Merge branch...

Merge branch '519374-create-routing-logic-in-ai-gateway-gecko-codestral-document-changes' into 'master' 

Respect Qwen opt out even when code gecko is disabled

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/183585



Merged-by: Vitali Tatarintev <vtatarintev@gitlab.com>
Approved-by: Vitali Tatarintev <vtatarintev@gitlab.com>
Co-authored-by: Allen Cook <acook@gitlab.com>

ee/lib/code_suggestions/model_details/code_completion.rb

@@ -15,6 +15,8 @@ def current_model
 
         return fireworks_qwen_2_5_model_details if use_fireworks_qwen_for_code_completions?
 
+        return codestral_2502_vertex_model_details if code_gecko_disabled? && !user_opted_in_to_qwen?
+
         # the default behavior is returning an empty hash
         # AI Gateway will fall back to the code-gecko model if model details are not provided
         {}
@@ -25,6 +27,10 @@ def saas_primary_model_class
 
         return CodeSuggestions::Prompts::CodeCompletion::FireworksQwen if use_fireworks_qwen_for_code_completions?
 
+        if code_gecko_disabled? && !user_opted_in_to_qwen?
+          return CodeSuggestions::Prompts::CodeCompletion::CodestralVertex
+        end
+
         CodeSuggestions::Prompts::CodeCompletion::VertexAi
       end
 
@@ -37,6 +43,13 @@ def fireworks_qwen_2_5_model_details
         }
       end
 
+      def codestral_2502_vertex_model_details
+        {
+          model_provider: CodeSuggestions::Prompts::CodeCompletion::CodestralVertex::MODEL_PROVIDER,
+          model_name: CodeSuggestions::Prompts::CodeCompletion::CodestralVertex::MODEL_NAME
+        }
+      end
+
       # We introduced an ops FF to allow organizations to opt out of Fireworks/Qwen.
       # On GitLab SaaS, Duo access is managed by top-level group,
       #   so we are checking the FF by the user's top-level group
@@ -48,7 +61,11 @@ def use_fireworks_qwen_for_code_completions?
         return false if Feature.disabled?(:fireworks_qwen_code_completion, current_user, type: :beta)
 
         # if the beta FF is enabled, proceed to check the ops FF
+        user_opted_in_to_qwen?
+      end
 
+      # check if user's groups or instance has opted out of Fireworks/Qwen
+      def user_opted_in_to_qwen?
         # on saas, check the user's groups
         return all_user_groups_opted_in_to_fireworks_qwen? if Gitlab.org_or_com? # rubocop: disable Gitlab/AvoidGitlabInstanceChecks -- see comment above method definition
 
@@ -56,6 +73,10 @@ def use_fireworks_qwen_for_code_completions?
         instance_opted_in_to_fireworks_qwen?
       end
 
+      def code_gecko_disabled?
+        Feature.enabled?(:disable_code_gecko_default, current_user)
+      end
+
       def all_user_groups_opted_in_to_fireworks_qwen?
         # while this has potential to be expensive in terms of Feature Flag checking,
         #   we are expecting that for most if not all users, only 1 group is providing them the Duo Access

ee/lib/code_suggestions/prompts/code_completion/codestral_vertex.rb

+# frozen_string_literal: true
+
+module CodeSuggestions
+  module Prompts
+    module CodeCompletion
+      class CodestralVertex < CodeSuggestions::Prompts::Base
+        GATEWAY_PROMPT_VERSION = 1
+        MODEL_NAME = 'codestral-2501'
+        MODEL_PROVIDER = 'vertex-ai'
+
+        def request_params
+          {
+            prompt_version: GATEWAY_PROMPT_VERSION,
+            model_name: MODEL_NAME,
+            model_provider: MODEL_PROVIDER
+          }
+        end
+      end
+    end
+  end
+end

ee/spec/lib/code_suggestions/model_details/code_completion_spec.rb

@@ -13,6 +13,7 @@
       before do
         stub_feature_flags(fireworks_qwen_code_completion: true)
         stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: false)
+        stub_feature_flags(disable_code_gecko_default: false)
       end
 
       context 'on GitLab self-managed' do
@@ -28,7 +29,16 @@
           it 'returns the codegecko model' do
             stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
 
-            expect(actual_result).to eq(expected_codegecko_result)
+            expect(actual_result).to eq(expected_default_result)
+          end
+
+          context 'when code gecko is disabled' do
+            it 'returns codestral on vertex' do
+              stub_feature_flags(disable_code_gecko_default: true)
+              stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
+
+              expect(actual_result).to eq(expected_codestral_result)
+            end
           end
         end
       end
@@ -61,7 +71,15 @@
           end
 
           it 'returns the code gecko model' do
-            expect(actual_result).to eq(expected_codegecko_result)
+            expect(actual_result).to eq(expected_default_result)
+          end
+
+          context 'when code gecko is disabled' do
+            it 'returns codestral on vertex' do
+              stub_feature_flags(disable_code_gecko_default: true)
+
+              expect(actual_result).to eq(expected_codestral_result)
+            end
           end
         end
 
@@ -80,10 +98,11 @@
     context 'when Fireworks/Qwen beta FF is disabled' do
       before do
         stub_feature_flags(fireworks_qwen_code_completion: false)
+        stub_feature_flags(disable_code_gecko_default: false)
       end
 
-      it 'returns the codegecko model' do
-        expect(actual_result).to eq(expected_codegecko_result)
+      it 'returns the default model' do
+        expect(actual_result).to eq(expected_default_result)
       end
     end
 
@@ -111,7 +130,14 @@
         }
       end
 
-      let(:expected_codegecko_result) { {} }
+      let(:expected_codestral_result) do
+        {
+          model_provider: 'vertex-ai',
+          model_name: 'codestral-2501'
+        }
+      end
+
+      let(:expected_default_result) { {} }
 
       let(:expected_self_hosted_model_result) { {} }
     end
@@ -125,7 +151,11 @@
         CodeSuggestions::Prompts::CodeCompletion::FireworksQwen
       end
 
-      let(:expected_codegecko_result) do
+      let(:expected_codestral_result) do
+        CodeSuggestions::Prompts::CodeCompletion::CodestralVertex
+      end
+
+      let(:expected_default_result) do
         CodeSuggestions::Prompts::CodeCompletion::VertexAi
       end
 

ee/spec/lib/code_suggestions/tasks/code_completion_spec.rb

@@ -128,6 +128,13 @@
         )
       end
 
+      let(:request_body_for_vertrex_codestral) do
+        request_body_without_model_details.merge(
+          "model_name" => "codestral-2501",
+          "model_provider" => "vertex-ai"
+        )
+      end
+
       context 'on GitLab self-managed' do
         before do
           allow(Gitlab).to receive(:org_or_com?).and_return(false)
@@ -141,10 +148,22 @@
           it_behaves_like 'code suggestion task' do
             before do
               stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
+              stub_feature_flags(disable_code_gecko_default: false)
             end
 
             let(:expected_body) { request_body_without_model_details }
           end
+
+          context "when code gecko is disabled" do
+            before do
+              stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: true)
+              stub_feature_flags(disable_code_gecko_default: true)
+            end
+
+            it_behaves_like 'code suggestion task' do
+              let(:expected_body) { request_body_for_vertrex_codestral }
+            end
+          end
         end
       end
 
@@ -173,11 +192,22 @@
           before do
             # opt out for group2
             stub_feature_flags(code_completion_model_opt_out_from_fireworks_qwen: group2)
+            stub_feature_flags(disable_code_gecko_default: false)
           end
 
           it_behaves_like 'code suggestion task' do
             let(:expected_body) { request_body_without_model_details }
           end
+
+          context "when code gecko is disabled" do
+            before do
+              stub_feature_flags(disable_code_gecko_default: true)
+            end
+
+            it_behaves_like 'code suggestion task' do
+              let(:expected_body) { request_body_for_vertrex_codestral }
+            end
+          end
         end
       end
     end
@@ -294,6 +324,7 @@
     context 'when amazon_q_chat_and_code_suggestions is disabled' do
       before do
         stub_feature_flags(amazon_q_chat_and_code_suggestions: false)
+        stub_feature_flags(disable_code_gecko_default: false)
       end
 
       it_behaves_like 'code suggestion task' do