Skip to content
This repository was archived by the owner on Sep 30, 2024. It is now read-only.

Commit 4984274

Browse files
chrsmithchrsmith
authored
Refactor the 'getModel' callbacks into their own file (#63359)
Minor refactoring to the `cmd/frontend/internal/completions` package. When we call `newCompletionsHandler` one of the parameters is a function named `getModel`. This is called to determine which LLM model should be use to respond to the incoming completion request. And we have two implementations of this function, one for code completions and another for chats. This PR just moves the logic for those two implementations into their own file (`get_model.go`), along with a couple of functions for which they were the only caller. There were two minor functionality changes I made, which I'll call out in comments on this PR. ## Why? As we rework how LLM models and associated configuration flows throughout the backend, updating these functions will be a bit easier if they are pulled out like this rather than being defined inline like they are today. Also, pretty much any place in the codebase where we have hard-coded an LLM model is "on notice" and should instead be driven entirely by some sort of global configuration file. (Since this is one of the key problems server-side LLM config is trying to solve.) ## Test plan NA, no functional changes. Relying on CI/CD and linter. ## Changelog NA
1 parent 4613452 commit 4984274

File tree

5 files changed

+161
-138
lines changed

5 files changed

+161
-138
lines changed

‎cmd/frontend/internal/httpapi/completions/BUILD.bazel‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ go_library(
66
srcs = [
77
"chat.go",
88
"codecompletion.go",
9+
"get_model.go",
910
"handler.go",
1011
"limiter.go",
1112
"observability.go",

‎cmd/frontend/internal/httpapi/completions/chat.go‎

Lines changed: 1 addition & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -5,17 +5,9 @@ import (
55

66
"net/http"
77

8-
"github.com/sourcegraph/sourcegraph/cmd/frontend/internal/cody"
9-
sgactor "github.com/sourcegraph/sourcegraph/internal/actor"
10-
"github.com/sourcegraph/sourcegraph/internal/dotcom"
11-
128
"github.com/sourcegraph/log"
139

14-
"github.com/sourcegraph/sourcegraph/internal/completions/client/anthropic"
15-
"github.com/sourcegraph/sourcegraph/internal/completions/client/fireworks"
16-
"github.com/sourcegraph/sourcegraph/internal/completions/client/google"
1710
"github.com/sourcegraph/sourcegraph/internal/completions/types"
18-
"github.com/sourcegraph/sourcegraph/internal/conf/conftypes"
1911
"github.com/sourcegraph/sourcegraph/internal/database"
2012
"github.com/sourcegraph/sourcegraph/internal/redispool"
2113
"github.com/sourcegraph/sourcegraph/internal/telemetry/telemetryrecorder"
@@ -43,79 +35,5 @@ func NewChatCompletionsStreamHandler(logger log.Logger, db database.DB) http.Han
4335
types.CompletionsFeatureChat,
4436
rl,
4537
"chat",
46-
func(ctx context.Context, requestParams types.CodyCompletionRequestParameters, c *conftypes.CompletionsConfig) (string, error) {
47-
// Allow a number of additional models on Dotcom
48-
if dotcom.SourcegraphDotComMode() {
49-
actor := sgactor.FromContext(ctx)
50-
user, err := actor.User(ctx, db.Users())
51-
if err != nil {
52-
return "", err
53-
}
54-
55-
subscription, err := cody.SubscriptionForUser(ctx, db, *user)
56-
if err != nil {
57-
return "", err
58-
}
59-
60-
if isAllowedCustomChatModel(requestParams.Model, subscription.ApplyProRateLimits) {
61-
return requestParams.Model, nil
62-
}
63-
}
64-
// No user defined models for now.
65-
if requestParams.Fast {
66-
return c.FastChatModel, nil
67-
}
68-
return c.ChatModel, nil
69-
},
70-
)
71-
}
72-
73-
// We only allow dotcom clients to select a custom chat model and maintain an allowlist for which
74-
// custom values we support
75-
func isAllowedCustomChatModel(model string, isProUser bool) bool {
76-
// When updating these two lists, make sure you also update `allowedModels` in codygateway_dotcom_user.go.
77-
if isProUser {
78-
switch model {
79-
case
80-
"anthropic/" + anthropic.Claude3Haiku,
81-
"anthropic/" + anthropic.Claude3Sonnet,
82-
"anthropic/" + anthropic.Claude3Opus,
83-
"fireworks/" + fireworks.Mixtral8x7bInstruct,
84-
"fireworks/" + fireworks.Mixtral8x22Instruct,
85-
"openai/gpt-3.5-turbo",
86-
"openai/gpt-4o",
87-
"openai/gpt-4-turbo",
88-
"openai/gpt-4-turbo-preview",
89-
"google/" + google.Gemini15FlashLatest,
90-
"google/" + google.Gemini15ProLatest,
91-
"google/" + google.GeminiProLatest,
92-
"google/" + google.Gemini15Flash,
93-
"google/" + google.Gemini15Pro,
94-
"google/" + google.GeminiPro,
95-
96-
// Remove after the Claude 3 rollout is complete
97-
"anthropic/claude-2",
98-
"anthropic/claude-2.0",
99-
"anthropic/claude-2.1",
100-
"anthropic/claude-instant-1.2-cyan",
101-
"anthropic/claude-instant-1.2",
102-
"anthropic/claude-instant-v1",
103-
"anthropic/claude-instant-1":
104-
return true
105-
}
106-
} else {
107-
switch model {
108-
case
109-
"anthropic/" + anthropic.Claude3Haiku,
110-
"anthropic/" + anthropic.Claude3Sonnet,
111-
// Remove after the Claude 3 rollout is complete
112-
"anthropic/claude-2",
113-
"anthropic/claude-2.0",
114-
"anthropic/claude-instant-v1",
115-
"anthropic/claude-instant-1":
116-
return true
117-
}
118-
}
119-
120-
return false
38+
getChatModelFn(db))
12139
}

‎cmd/frontend/internal/httpapi/completions/codecompletion.go‎

Lines changed: 1 addition & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,15 @@
11
package completions
22

33
import (
4-
"context"
54
"net/http"
65

76
"github.com/sourcegraph/log"
87

9-
"github.com/sourcegraph/sourcegraph/internal/completions/client/fireworks"
10-
"github.com/sourcegraph/sourcegraph/internal/completions/client/google"
118
"github.com/sourcegraph/sourcegraph/internal/completions/types"
12-
"github.com/sourcegraph/sourcegraph/internal/conf/conftypes"
139
"github.com/sourcegraph/sourcegraph/internal/database"
1410
"github.com/sourcegraph/sourcegraph/internal/guardrails"
1511
"github.com/sourcegraph/sourcegraph/internal/redispool"
1612
"github.com/sourcegraph/sourcegraph/internal/telemetry/telemetryrecorder"
17-
"github.com/sourcegraph/sourcegraph/lib/errors"
1813
)
1914

2015
// NewCodeCompletionsHandler is an http handler which sends back code completion results.
@@ -31,54 +26,5 @@ func NewCodeCompletionsHandler(logger log.Logger, db database.DB, test guardrail
3126
types.CompletionsFeatureCode,
3227
rl,
3328
"code",
34-
func(_ context.Context, requestParams types.CodyCompletionRequestParameters, c *conftypes.CompletionsConfig) (string, error) {
35-
customModel := allowedCustomModel(requestParams.Model)
36-
if customModel != "" {
37-
return customModel, nil
38-
}
39-
if requestParams.Model != "" {
40-
return "", errors.Newf("Unsupported code completion model %q", requestParams.Model)
41-
}
42-
return c.CompletionModel, nil
43-
},
44-
)
45-
}
46-
47-
func allowedCustomModel(model string) string {
48-
switch model {
49-
case "fireworks/starcoder",
50-
"fireworks/starcoder-16b",
51-
"fireworks/starcoder-7b",
52-
"fireworks/starcoder2-15b",
53-
"fireworks/starcoder2-7b",
54-
"fireworks/" + fireworks.Starcoder16b,
55-
"fireworks/" + fireworks.Starcoder7b,
56-
"fireworks/" + fireworks.Llama27bCode,
57-
"fireworks/" + fireworks.Llama213bCode,
58-
"fireworks/" + fireworks.Llama213bCodeInstruct,
59-
"fireworks/" + fireworks.Llama234bCodeInstruct,
60-
"fireworks/" + fireworks.Mistral7bInstruct,
61-
"fireworks/" + fireworks.FineTunedFIMVariant1,
62-
"fireworks/" + fireworks.FineTunedFIMVariant2,
63-
"fireworks/" + fireworks.FineTunedFIMVariant3,
64-
"fireworks/" + fireworks.FineTunedFIMVariant4,
65-
"fireworks/" + fireworks.FineTunedFIMLangSpecificMixtral,
66-
"fireworks/" + fireworks.DeepseekCoder1p3b,
67-
"fireworks/" + fireworks.DeepseekCoder7b,
68-
"anthropic/claude-instant-1.2",
69-
"anthropic/claude-3-haiku-20240307",
70-
// Deprecated model identifiers
71-
"anthropic/claude-instant-v1",
72-
"anthropic/claude-instant-1",
73-
"anthropic/claude-instant-1.2-cyan",
74-
"google/" + google.Gemini15Flash,
75-
"google/" + google.GeminiPro,
76-
"fireworks/accounts/sourcegraph/models/starcoder-7b",
77-
"fireworks/accounts/sourcegraph/models/starcoder-16b",
78-
"fireworks/accounts/fireworks/models/starcoder-3b-w8a16",
79-
"fireworks/accounts/fireworks/models/starcoder-1b-w8a16":
80-
return model
81-
}
82-
83-
return ""
29+
getCodeCompletionModelFn())
8430
}

‎cmd/frontend/internal/httpapi/completions/get_model.go‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
package completions
2+
3+
import (
4+
"context"
5+
6+
"github.com/sourcegraph/sourcegraph/cmd/frontend/internal/cody"
7+
sgactor "github.com/sourcegraph/sourcegraph/internal/actor"
8+
"github.com/sourcegraph/sourcegraph/internal/dotcom"
9+
"github.com/sourcegraph/sourcegraph/lib/errors"
10+
11+
"github.com/sourcegraph/sourcegraph/internal/completions/client/anthropic"
12+
"github.com/sourcegraph/sourcegraph/internal/completions/client/fireworks"
13+
"github.com/sourcegraph/sourcegraph/internal/completions/client/google"
14+
"github.com/sourcegraph/sourcegraph/internal/completions/types"
15+
"github.com/sourcegraph/sourcegraph/internal/conf/conftypes"
16+
"github.com/sourcegraph/sourcegraph/internal/database"
17+
)
18+
19+
// getModelFn is the thunk used to return the LLM model we should use for processing
20+
// the supplied completion request. Depending on the incomming request, site config,
21+
// feature used, etc. it could be any number of things.
22+
type getModelFn func(ctx context.Context, requestParams types.CodyCompletionRequestParameters, c *conftypes.CompletionsConfig) (string, error)
23+
24+
func getCodeCompletionModelFn() getModelFn {
25+
return func(_ context.Context, requestParams types.CodyCompletionRequestParameters, c *conftypes.CompletionsConfig) (string, error) {
26+
// For code completions, we only allow certain models to be used.
27+
// (Regardless of if the user is on Cody Free, Pro, or Enterprise.)
28+
if requestParams.Model != "" {
29+
if isAllowedCodeCompletionModel(requestParams.Model) {
30+
return requestParams.Model, nil
31+
}
32+
return "", errors.Newf("unsupported code completion model %q", requestParams.Model)
33+
}
34+
return c.CompletionModel, nil
35+
}
36+
}
37+
38+
func getChatModelFn(db database.DB) getModelFn {
39+
return func(ctx context.Context, requestParams types.CodyCompletionRequestParameters, c *conftypes.CompletionsConfig) (string, error) {
40+
// If running on dotcom, i.e. using Cody Free/Cody Pro, then a number
41+
// of models are available depending on the caller's subscription status.
42+
if dotcom.SourcegraphDotComMode() {
43+
actor := sgactor.FromContext(ctx)
44+
user, err := actor.User(ctx, db.Users())
45+
if err != nil {
46+
return "", err
47+
}
48+
49+
subscription, err := cody.SubscriptionForUser(ctx, db, *user)
50+
if err != nil {
51+
return "", err
52+
}
53+
54+
if isAllowedCustomChatModel(requestParams.Model, subscription.ApplyProRateLimits) {
55+
return requestParams.Model, nil
56+
}
57+
}
58+
59+
// For any other Sourcegraph instance, i.e. using Cody Enterprise,
60+
// we just use the configured "chat" or "fastChat" model.
61+
if requestParams.Fast {
62+
return c.FastChatModel, nil
63+
}
64+
return c.ChatModel, nil
65+
}
66+
}
67+
68+
func isAllowedCodeCompletionModel(model string) bool {
69+
switch model {
70+
case "fireworks/starcoder",
71+
"fireworks/starcoder-16b",
72+
"fireworks/starcoder-7b",
73+
"fireworks/starcoder2-15b",
74+
"fireworks/starcoder2-7b",
75+
"fireworks/" + fireworks.Starcoder16b,
76+
"fireworks/" + fireworks.Starcoder7b,
77+
"fireworks/" + fireworks.Llama27bCode,
78+
"fireworks/" + fireworks.Llama213bCode,
79+
"fireworks/" + fireworks.Llama213bCodeInstruct,
80+
"fireworks/" + fireworks.Llama234bCodeInstruct,
81+
"fireworks/" + fireworks.Mistral7bInstruct,
82+
"fireworks/" + fireworks.FineTunedFIMVariant1,
83+
"fireworks/" + fireworks.FineTunedFIMVariant2,
84+
"fireworks/" + fireworks.FineTunedFIMVariant3,
85+
"fireworks/" + fireworks.FineTunedFIMVariant4,
86+
"fireworks/" + fireworks.FineTunedFIMLangSpecificMixtral,
87+
"fireworks/" + fireworks.DeepseekCoder1p3b,
88+
"fireworks/" + fireworks.DeepseekCoder7b,
89+
"anthropic/claude-instant-1.2",
90+
"anthropic/claude-3-haiku-20240307",
91+
// Deprecated model identifiers
92+
"anthropic/claude-instant-v1",
93+
"anthropic/claude-instant-1",
94+
"anthropic/claude-instant-1.2-cyan",
95+
"google/" + google.Gemini15Flash,
96+
"google/" + google.GeminiPro,
97+
"fireworks/accounts/sourcegraph/models/starcoder-7b",
98+
"fireworks/accounts/sourcegraph/models/starcoder-16b",
99+
"fireworks/accounts/fireworks/models/starcoder-3b-w8a16",
100+
"fireworks/accounts/fireworks/models/starcoder-1b-w8a16":
101+
return true
102+
}
103+
104+
return false
105+
}
106+
107+
// We only allow dotcom clients to select a custom chat model and maintain an allowlist for which
108+
// custom values we support
109+
func isAllowedCustomChatModel(model string, isProUser bool) bool {
110+
// When updating these two lists, make sure you also update `allowedModels` in codygateway_dotcom_user.go.
111+
if isProUser {
112+
switch model {
113+
case
114+
"anthropic/" + anthropic.Claude3Haiku,
115+
"anthropic/" + anthropic.Claude3Sonnet,
116+
"anthropic/" + anthropic.Claude3Opus,
117+
"fireworks/" + fireworks.Mixtral8x7bInstruct,
118+
"fireworks/" + fireworks.Mixtral8x22Instruct,
119+
"openai/gpt-3.5-turbo",
120+
"openai/gpt-4o",
121+
"openai/gpt-4-turbo",
122+
"openai/gpt-4-turbo-preview",
123+
"google/" + google.Gemini15FlashLatest,
124+
"google/" + google.Gemini15ProLatest,
125+
"google/" + google.GeminiProLatest,
126+
"google/" + google.Gemini15Flash,
127+
"google/" + google.Gemini15Pro,
128+
"google/" + google.GeminiPro,
129+
130+
// Remove after the Claude 3 rollout is complete
131+
"anthropic/claude-2",
132+
"anthropic/claude-2.0",
133+
"anthropic/claude-2.1",
134+
"anthropic/claude-instant-1.2-cyan",
135+
"anthropic/claude-instant-1.2",
136+
"anthropic/claude-instant-v1",
137+
"anthropic/claude-instant-1":
138+
return true
139+
}
140+
} else {
141+
switch model {
142+
case
143+
"anthropic/" + anthropic.Claude3Haiku,
144+
"anthropic/" + anthropic.Claude3Sonnet,
145+
// Remove after the Claude 3 rollout is complete
146+
"anthropic/claude-2",
147+
"anthropic/claude-2.0",
148+
"anthropic/claude-instant-v1",
149+
"anthropic/claude-instant-1":
150+
return true
151+
}
152+
}
153+
154+
return false
155+
}

‎cmd/frontend/internal/httpapi/completions/handler.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,9 @@ func newCompletionsHandler(
115115
requestParams.Model, err = getModel(ctx, requestParams, completionsConfig)
116116
requestParams.User = completionsConfig.User
117117
if err != nil {
118+
// NOTE: We return the raw error to the user assuming that it contains relevant
119+
// user-facing diagnostic information, and doesn't leak any internal details.
120+
logger.Info("error fetching model", log.Error(err))
118121
http.Error(w, err.Error(), http.StatusBadRequest)
119122
return
120123
}

0 commit comments

Comments
 (0)