From 03f58663804a6382d2fc709e5113c6833b95c812 Mon Sep 17 00:00:00 2001
From: Beatrix <68532117+abeatrix@users.noreply.github.com>
Date: Tue, 4 Jun 2024 07:15:35 -0700
Subject: [PATCH] Cody Gateway: remove feature flag requirement for Gemini
 models (#63055)

Cody Gateway: add google/gemini-1.5-pro-latest and google/gemini-1.5-flash-latest to allowed models

This change adds the Google Gemini 1.5 Pro and Flash models to the list of allowed models for Cody Gateway chat completions by default without requiring "cody-pro-gemini-enabled". This allows Cody Pro users to access these models.

The changes were made in the following files:
- `cmd/frontend/internal/dotcom/productsubscription/codygateway_dotcom_user.go`
- `internal/licensing/codygateway.go`
- `internal/licensing/codygateway_test.go`
---
 .../productsubscription/codygateway_dotcom_user.go    | 11 +++++------
 internal/licensing/codygateway.go                     |  3 +++
 internal/licensing/codygateway_test.go                |  6 ++++++
 3 files changed, 14 insertions(+), 6 deletions(-)

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,