diff --git a/cmd/frontend/internal/dotcom/productsubscription/codygateway_dotcom_user.go b/cmd/frontend/internal/dotcom/productsubscription/codygateway_dotcom_user.go index c64fc2ddfb9..2a3aa7466d4 100644 --- a/cmd/frontend/internal/dotcom/productsubscription/codygateway_dotcom_user.go +++ b/cmd/frontend/internal/dotcom/productsubscription/codygateway_dotcom_user.go @@ -246,7 +246,7 @@ func getCompletionsRateLimit(ctx context.Context, db database.DB, userID int32, if featureflag.FromContext(ctx).GetBoolOr("rate-limits-exceeded-for-testing", false) { return licensing.CodyGatewayRateLimit{ // For this special tester user, just allow all models a Pro user can get. - AllowedModels: allowedModels(scope, true, false), + AllowedModels: allowedModels(scope, true), Limit: 1, IntervalSeconds: math.MaxInt32, }, graphqlbackend.CodyGatewayRateLimitSourceOverride, nil @@ -279,7 +279,7 @@ func getCompletionsRateLimit(ctx context.Context, db database.DB, userID int32, if err != nil { return licensing.CodyGatewayRateLimit{}, graphqlbackend.CodyGatewayRateLimitSourcePlan, errors.Wrap(err, "error fetching user's cody subscription") } - models := allowedModels(scope, true, featureflag.FromContext(ctx).GetBoolOr("cody-pro-gemini-enabled", false)) + models := allowedModels(scope, true) if limit == nil && cfg != nil { source = graphqlbackend.CodyGatewayRateLimitSourcePlan // Update the allowed models based on the user's plan. @@ -355,7 +355,7 @@ var allCodeCompletionModels = slices.Concat([]string{"anthropic/" + anthropic.Cl prefix("fireworks/", fireworks.FineTunedMixtralModelVariants), prefix("fireworks/", fireworks.FineTunedLlamaModelVariants)) -func allowedModels(scope types.CompletionsFeature, isProUser bool, googleGeminiEnabled bool) []string { +func allowedModels(scope types.CompletionsFeature, isProUser bool) []string { switch scope { case types.CompletionsFeatureChat: // When updating the below lists, make sure you also update `isAllowedCustomChatModel` in `chat.go` @@ -382,6 +382,8 @@ func allowedModels(scope types.CompletionsFeature, isProUser bool, googleGeminiE "openai/gpt-4o", "openai/gpt-4-turbo", "openai/gpt-4-turbo-preview", + "google/gemini-1.5-pro-latest", + "google/gemini-1.5-flash-latest", // Remove after the Claude 3 rollout is complete "anthropic/claude-2", @@ -392,9 +394,6 @@ func allowedModels(scope types.CompletionsFeature, isProUser bool, googleGeminiE "anthropic/claude-instant-v1", "anthropic/claude-instant-1", } - if googleGeminiEnabled { - chatModels = append(chatModels, "google/gemini-1.5-pro-latest", "google/gemini-1.5-flash-latest") - } return chatModels case types.CompletionsFeatureCode: diff --git a/internal/licensing/codygateway.go b/internal/licensing/codygateway.go index a6dd034e76d..10728570d2b 100644 --- a/internal/licensing/codygateway.go +++ b/internal/licensing/codygateway.go @@ -47,6 +47,9 @@ func NewCodyGatewayChatRateLimit(plan Plan, userCount *int) CodyGatewayRateLimit "openai/gpt-4o", "openai/gpt-4-turbo", "openai/gpt-4-turbo-preview", + + "google/gemini-1.5-pro-latest", + "google/gemini-1.5-flash-latest", } switch plan { // TODO: This is just an example for now. diff --git a/internal/licensing/codygateway_test.go b/internal/licensing/codygateway_test.go index 6bed8153fe2..8e74d7d413a 100644 --- a/internal/licensing/codygateway_test.go +++ b/internal/licensing/codygateway_test.go @@ -36,6 +36,8 @@ func TestNewCodyGatewayChatRateLimit(t *testing.T) { "openai/gpt-4o", "openai/gpt-4-turbo", "openai/gpt-4-turbo-preview", + "google/gemini-1.5-pro-latest", + "google/gemini-1.5-flash-latest", }, Limit: 2500, IntervalSeconds: 60 * 60 * 24, @@ -61,6 +63,8 @@ func TestNewCodyGatewayChatRateLimit(t *testing.T) { "openai/gpt-4o", "openai/gpt-4-turbo", "openai/gpt-4-turbo-preview", + "google/gemini-1.5-pro-latest", + "google/gemini-1.5-flash-latest", }, Limit: 50, IntervalSeconds: 60 * 60 * 24, @@ -86,6 +90,8 @@ func TestNewCodyGatewayChatRateLimit(t *testing.T) { "openai/gpt-4o", "openai/gpt-4-turbo", "openai/gpt-4-turbo-preview", + "google/gemini-1.5-pro-latest", + "google/gemini-1.5-flash-latest", }, Limit: 10, IntervalSeconds: 60 * 60 * 24,