From 109268651b8dbd3cf974e4809e417542e0a399e2 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Thu, 15 May 2025 08:19:19 +0200
Subject: [PATCH 01/11] Add top_namespace_model_ref logic to request

Adjust model_params to have provider: 'gitlab'

Changelog: added
EE: true
---
 .../gitlab/llm/chain/requests/ai_gateway.rb   |  20 +++-
 ee/lib/gitlab/llm/completions/chat.rb         |   6 +-
 .../namespace_feature_settings.rb             |   9 ++
 .../llm/chain/requests/ai_gateway_spec.rb     | 103 ++++++++++++++++++
 .../lib/gitlab/llm/completions/chat_spec.rb   |   8 ++
 5 files changed, 143 insertions(+), 3 deletions(-)

diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index 1be68ab763dce7..b3577f59239b1d 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -22,15 +22,18 @@ class AiGateway < Base
           STOP_WORDS = ["\n\nHuman", "Observation:"].freeze
           DEFAULT_MAX_TOKENS = 4096
 
-          def initialize(user, service_name: :duo_chat, tracking_context: {})
+          def initialize(user, service_name: :duo_chat, tracking_context: {}, root_namespace: nil)
             @user = user
             @tracking_context = tracking_context
+            @root_namespace = root_namespace
             @ai_client = ::Gitlab::Llm::AiGateway::Client.new(user, service_name: processed_service_name(service_name),
               tracking_context: tracking_context)
           end
 
           def request(prompt, unit_primitive: nil)
             options = default_options.merge(prompt.fetch(:options, {}))
+            options[:unit_primitive] ||= unit_primitive
+
             return unless model_provider_valid?(options)
 
             response = ai_client.stream(
@@ -56,7 +59,7 @@ def request(prompt, unit_primitive: nil)
 
           private
 
-          attr_reader :user
+          attr_reader :user, :root_namespace
 
           def default_options
             {
@@ -82,6 +85,16 @@ def model_provider_valid?(options)
             provider(options)
           end
 
+          def top_namespace_model_ref(unit_primitive = nil)
+            return unless root_namespace&.root?
+            return unless Feature.enabled?(:ai_model_switching, root_namespace)
+
+            feature = unit_primitive ? :"duo_chat_#{unit_primitive}" : :duo_chat
+            setting = ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(root_namespace,
+              feature)
+            setting&.offered_model_ref
+          end
+
           def endpoint(unit_primitive, use_ai_gateway_agent_prompt)
             path =
               if use_ai_gateway_agent_prompt
@@ -148,6 +161,9 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
           end
 
           def model_params(options)
+            namespace_model = top_namespace_model_ref(options[:unit_primitive])
+            return { provider: 'gitlab', identifier: namespace_model } if namespace_model.present?
+
             if chat_feature_setting&.self_hosted?
               self_hosted_model = chat_feature_setting.self_hosted_model
 
diff --git a/ee/lib/gitlab/llm/completions/chat.rb b/ee/lib/gitlab/llm/completions/chat.rb
index c449e382c434e3..0d968779983211 100644
--- a/ee/lib/gitlab/llm/completions/chat.rb
+++ b/ee/lib/gitlab/llm/completions/chat.rb
@@ -53,7 +53,11 @@ def initialize(prompt_message, ai_prompt_class, options = {})
         end
 
         def ai_request
-          ::Gitlab::Llm::Chain::Requests::AiGateway.new(user, tracking_context: tracking_context)
+          ::Gitlab::Llm::Chain::Requests::AiGateway.new(
+            user,
+            tracking_context: tracking_context,
+            root_namespace: resource.try(:resource_parent)&.root_ancestor
+          )
         end
 
         def execute
diff --git a/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb b/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb
index 139eb0e7eb3570..b0f6aed696d687 100644
--- a/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb
+++ b/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb
@@ -14,6 +14,8 @@
         "models" => [
           { 'name' => 'Claude Sonnet 3.5', 'identifier' => 'claude_sonnet_3_5' },
           { 'name' => 'Claude Sonnet 3.7', 'identifier' => 'claude_sonnet_3_7' },
+          { 'name' => 'Claude Sonnet 3.7', 'identifier' => 'claude-3-7-sonnet-20250219' },
+          { 'name' => 'Claude 3.5 Sonnet', 'identifier' => 'claude-3-5-sonnet-20240620' },
           { 'name' => 'OpenAI Chat GPT 4o', 'identifier' => 'openai_chatgpt_4o' }
         ],
         "unit_primitives" => [
@@ -30,6 +32,13 @@
             "selectable_models" => %w[claude_sonnet_3_5 claude_sonnet_3_7 openai_chatgpt_4o],
             "beta_models" => [],
             "unit_primitives" => ["generate_code"]
+          },
+          {
+            "feature_setting" => "duo_chat_explain_code",
+            "default_model" => "claude-3-7-sonnet-20250219",
+            "selectable_models" => %w[claude-3-7-sonnet-20250219 claude-3-5-sonnet-20240620],
+            "beta_models" => [],
+            "unit_primitives" => ["explain_code"]
           }
         ]
       }
diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
index 752a41481818e5..18a61425172921 100644
--- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
@@ -367,6 +367,109 @@
         end
       end
     end
+
+    context 'when root_namespace is passed' do
+      let_it_be(:root_namespace) { create(:group) }
+      let(:tracking_context) { { action: 'chat', request_id: 'uuid' } }
+      let(:prompt) { { prompt: user_prompt, options: {} } }
+      let(:user_prompt) { "Some prompt" }
+      let(:ai_client) { double }
+      let(:logger) { instance_double(Gitlab::Llm::Logger) }
+      let(:response) { 'response from llm' }
+
+      before do
+        allow(Gitlab::Llm::Logger).to receive(:build).and_return(logger)
+        allow(logger).to receive(:conditional_info)
+        allow_next_instance_of(described_class) do |instance|
+          allow(instance).to receive(:ai_client).and_return(ai_client)
+        end
+      end
+
+      context 'and root_namespace is root and model switching is enabled' do
+        let(:unit_primitive) { :explain_code }
+        let(:feature) { :"duo_chat_#{unit_primitive}" }
+        let(:endpoint) { "#{described_class::BASE_ENDPOINT}/#{unit_primitive}" }
+        let(:url) { "#{::Gitlab::AiGateway.url}#{endpoint}" }
+        let(:model_ref) { 'claude-3-7-sonnet-20250219' }
+
+        before do
+          create(:ai_namespace_feature_setting, namespace: root_namespace,
+            feature: :duo_chat_explain_code, offered_model_ref: model_ref)
+
+          stub_feature_flags(ai_model_switching: true)
+        end
+
+        it 'uses model from namespace feature setting based on unit_primitive' do
+          expect(ai_client).to receive(:stream).with(
+            hash_including(
+              url: "https://cloud.gitlab.com/ai/v1/chat/explain_code",
+              body: hash_including(
+                prompt_components: array_including(
+                  hash_including(payload: hash_including(
+                    identifier: model_ref,
+                    provider: "gitlab"
+                  ))
+                )
+              )
+            )
+          ).and_return(response)
+
+          gateway = described_class.new(user, root_namespace: root_namespace, tracking_context: tracking_context)
+          expect(gateway.request(prompt, unit_primitive: unit_primitive)).to eq(response)
+        end
+      end
+
+      context 'and root_namespace is root but model switching is disabled' do
+        let(:unit_primitive) { nil }
+        let(:endpoint) { described_class::ENDPOINT }
+        let(:url) { "#{::Gitlab::AiGateway.url}#{endpoint}" }
+
+        before do
+          stub_feature_flags(ai_model_switching: false)
+        end
+
+        it 'falls back to default model' do
+          expect(ai_client).to receive(:stream).with(
+            hash_including(
+              url: url,
+              body: hash_including(
+                prompt_components: array_including(
+                  hash_including(payload: hash_including(model: described_class::CLAUDE_3_5_SONNET))
+                )
+              )
+            )
+          ).and_return(response)
+
+          gateway = described_class.new(user, root_namespace: root_namespace, tracking_context: tracking_context)
+          expect(gateway.request(prompt, unit_primitive: unit_primitive)).to eq(response)
+        end
+      end
+
+      context 'and root_namespace is not a root group' do
+        let(:subgroup) { create(:group, parent: root_namespace) }
+        let(:unit_primitive) { nil }
+        let(:endpoint) { described_class::ENDPOINT }
+        let(:url) { "#{::Gitlab::AiGateway.url}#{endpoint}" }
+
+        it 'does not use model switching even if feature flag is enabled' do
+          stub_feature_flags(ai_model_switching: true)
+
+          expect(ai_client).to receive(:stream).with(
+            hash_including(
+              url: url,
+              body: hash_including(
+                prompt_components: array_including(
+                  hash_including(payload: hash_including(model: described_class::CLAUDE_3_5_SONNET))
+                )
+              )
+            )
+          ).and_return(response)
+
+          gateway = described_class.new(user, root_namespace: subgroup, tracking_context: tracking_context)
+          expect(gateway.request(prompt, unit_primitive: unit_primitive)).to eq(response)
+        end
+      end
+    end
   end
   # rubocop:enable RSpec/MultipleMemoizedHelpers
 end
diff --git a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
index 070efbbb953031..f9d880e6c0bd5a 100644
--- a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
+++ b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
@@ -101,6 +101,14 @@
         stream_response_handler: stream_response_handler
       ]
 
+      expect(Gitlab::Llm::Chain::Requests::AiGateway).to receive(:new).with(
+        user,
+        hash_including(
+          tracking_context: anything,
+          root_namespace: expected_container&.root_ancestor
+        )
+      ).and_return(ai_request)
+
       expect_next_instance_of(::Gitlab::Duo::Chat::ReactExecutor, *expected_params) do |instance|
         expect(instance).to receive(:execute).and_return(answer)
       end
-- 
GitLab


From 27b35840697b1fc4c6f35a3c8ed0f6c79e6c90d9 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Thu, 15 May 2025 12:57:29 +0200
Subject: [PATCH 02/11] Adjust model_params and model_metadata

---
 .../gitlab/llm/chain/requests/ai_gateway.rb   | 26 +++++-
 .../llm/chain/requests/ai_gateway_spec.rb     | 79 ++++++++++++++-----
 2 files changed, 83 insertions(+), 22 deletions(-)

diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index b3577f59239b1d..06d80e819d7be2 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -151,6 +151,11 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
             model_family = model_metadata_params && model_metadata_params[:name]
             default_version = ::Gitlab::Llm::PromptVersions.version_for_prompt("chat/#{unit_primitive}", model_family)
 
+            if top_namespace_model_selected?(unit_primitive)
+              apply_model_selection_metadata(params, unit_primitive,
+                prompt_version)
+            end
+
             params[:prompt_version] = if feature_setting&.self_hosted? || ::Ai::AmazonQ.connected?
                                         default_version
                                       else
@@ -161,8 +166,7 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
           end
 
           def model_params(options)
-            namespace_model = top_namespace_model_ref(options[:unit_primitive])
-            return { provider: 'gitlab', identifier: namespace_model } if namespace_model.present?
+            unit_primitive = options[:unit_primitive]
 
             if chat_feature_setting&.self_hosted?
               self_hosted_model = chat_feature_setting.self_hosted_model
@@ -174,6 +178,11 @@ def model_params(options)
                 model_api_key: self_hosted_model.api_token,
                 model_identifier: self_hosted_model.identifier
               }
+            elsif top_namespace_model_selected?(unit_primitive)
+              {
+                provider: 'gitlab',
+                identifier: top_namespace_model_ref(unit_primitive)
+              }
             else
               {
                 provider: provider(options),
@@ -215,6 +224,19 @@ def chat_feature_setting(unit_primitive: nil)
             feature_setting
           end
 
+          def top_namespace_model_selected?(unit_primitive)
+            top_namespace_model_ref(unit_primitive).present?
+          end
+
+          def apply_model_selection_metadata(params, unit_primitive, _prompt_version)
+            identifier = top_namespace_model_ref(unit_primitive)
+
+            params[:model_metadata] = {
+              provider: 'gitlab',
+              id: identifier
+            }
+          end
+
           def processed_service_name(service_name)
             return service_name unless service_name == :duo_chat
             return service_name unless chat_feature_setting&.self_hosted?
diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
index 18a61425172921..50d3b5164e5de7 100644
--- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
@@ -371,11 +371,10 @@
     context 'when root_namespace is passed' do
       let_it_be(:root_namespace) { create(:group) }
       let(:tracking_context) { { action: 'chat', request_id: 'uuid' } }
-      let(:prompt) { { prompt: user_prompt, options: {} } }
       let(:user_prompt) { "Some prompt" }
-      let(:ai_client) { double }
-      let(:logger) { instance_double(Gitlab::Llm::Logger) }
       let(:response) { 'response from llm' }
+      let(:logger) { instance_double(Gitlab::Llm::Logger) }
+      let(:ai_client) { double }
 
       before do
         allow(Gitlab::Llm::Logger).to receive(:build).and_return(logger)
@@ -385,29 +384,28 @@
         end
       end
 
-      context 'and root_namespace is root and model switching is enabled' do
+      context 'when model switching is enabled and namespace is root' do
         let(:unit_primitive) { :explain_code }
-        let(:feature) { :"duo_chat_#{unit_primitive}" }
-        let(:endpoint) { "#{described_class::BASE_ENDPOINT}/#{unit_primitive}" }
-        let(:url) { "#{::Gitlab::AiGateway.url}#{endpoint}" }
         let(:model_ref) { 'claude-3-7-sonnet-20250219' }
+        let(:prompt) { { prompt: user_prompt, options: {} } }
 
         before do
           create(:ai_namespace_feature_setting, namespace: root_namespace,
-            feature: :duo_chat_explain_code, offered_model_ref: model_ref)
-
+            feature: :"duo_chat_#{unit_primitive}", offered_model_ref: model_ref)
           stub_feature_flags(ai_model_switching: true)
         end
 
-        it 'uses model from namespace feature setting based on unit_primitive' do
+        it 'uses namespace model ref in prompt body' do
+          url = "#{::Gitlab::AiGateway.url}#{described_class::BASE_ENDPOINT}/#{unit_primitive}"
+
           expect(ai_client).to receive(:stream).with(
             hash_including(
-              url: "https://cloud.gitlab.com/ai/v1/chat/explain_code",
+              url: url,
               body: hash_including(
                 prompt_components: array_including(
                   hash_including(payload: hash_including(
-                    identifier: model_ref,
-                    provider: "gitlab"
+                    provider: "gitlab",
+                    identifier: model_ref
                   ))
                 )
               )
@@ -419,16 +417,54 @@
         end
       end
 
-      context 'and root_namespace is root but model switching is disabled' do
+      context 'when using agent prompt with model switching' do
+        let(:unit_primitive) { :explain_code }
+        let(:model_ref) { 'claude-3-7-sonnet-20250219' }
+        let(:prompt) do
+          {
+            options: {
+              use_ai_gateway_agent_prompt: true,
+              inputs: { a: 1 }
+            }
+          }
+        end
+
+        before do
+          create(:ai_namespace_feature_setting, namespace: root_namespace,
+            feature: :"duo_chat_#{unit_primitive}", offered_model_ref: model_ref)
+          stub_feature_flags(ai_model_switching: true)
+        end
+
+        it 'sends agent request with model_metadata' do
+          url = "#{::Gitlab::AiGateway.url}#{described_class::BASE_PROMPTS_CHAT_ENDPOINT}/#{unit_primitive}"
+
+          expect(ai_client).to receive(:stream).with(
+            hash_including(
+              url: url,
+              body: hash_including(
+                inputs: { a: 1 },
+                model_metadata: { provider: 'gitlab', id: model_ref },
+                prompt_version: a_kind_of(String)
+              )
+            )
+          ).and_return(response)
+
+          gateway = described_class.new(user, root_namespace: root_namespace, tracking_context: tracking_context)
+          expect(gateway.request(prompt, unit_primitive: unit_primitive)).to eq(response)
+        end
+      end
+
+      context 'when model switching is disabled' do
         let(:unit_primitive) { nil }
-        let(:endpoint) { described_class::ENDPOINT }
-        let(:url) { "#{::Gitlab::AiGateway.url}#{endpoint}" }
+        let(:prompt) { { prompt: user_prompt, options: {} } }
 
         before do
           stub_feature_flags(ai_model_switching: false)
         end
 
         it 'falls back to default model' do
+          url = "#{::Gitlab::AiGateway.url}#{described_class::ENDPOINT}"
+
           expect(ai_client).to receive(:stream).with(
             hash_including(
               url: url,
@@ -445,14 +481,17 @@
         end
       end
 
-      context 'and root_namespace is not a root group' do
+      context 'when root_namespace is not root' do
         let(:subgroup) { create(:group, parent: root_namespace) }
         let(:unit_primitive) { nil }
-        let(:endpoint) { described_class::ENDPOINT }
-        let(:url) { "#{::Gitlab::AiGateway.url}#{endpoint}" }
+        let(:prompt) { { prompt: user_prompt, options: {} } }
 
-        it 'does not use model switching even if feature flag is enabled' do
+        before do
           stub_feature_flags(ai_model_switching: true)
+        end
+
+        it 'ignores model switching and uses default model' do
+          url = "#{::Gitlab::AiGateway.url}#{described_class::ENDPOINT}"
 
           expect(ai_client).to receive(:stream).with(
             hash_including(
-- 
GitLab


From 0983931e3bd2a46bb6e369ec012c9afc6a74234e Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Thu, 15 May 2025 14:40:05 +0200
Subject: [PATCH 03/11] Change model_metadata to use identifier not id

---
 ee/lib/gitlab/llm/chain/requests/ai_gateway.rb           | 2 +-
 ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index 06d80e819d7be2..2dc736bf4b041d 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -233,7 +233,7 @@ def apply_model_selection_metadata(params, unit_primitive, _prompt_version)
 
             params[:model_metadata] = {
               provider: 'gitlab',
-              id: identifier
+              identifier: identifier
             }
           end
 
diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
index 50d3b5164e5de7..af681564ffb9c0 100644
--- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
@@ -443,7 +443,7 @@
               url: url,
               body: hash_including(
                 inputs: { a: 1 },
-                model_metadata: { provider: 'gitlab', id: model_ref },
+                model_metadata: { provider: 'gitlab', identifier: model_ref },
                 prompt_version: a_kind_of(String)
               )
             )
-- 
GitLab


From fc97f2be2234b2a1d3ab74600f77de09cef513c9 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Thu, 15 May 2025 15:40:19 +0200
Subject: [PATCH 04/11] Fallback to feature_setting if no identifier

---
 .../gitlab/llm/chain/requests/ai_gateway.rb   | 34 +++++++++++++------
 .../llm/chain/requests/ai_gateway_spec.rb     | 13 +++----
 2 files changed, 31 insertions(+), 16 deletions(-)

diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index 2dc736bf4b041d..bc3e046601871d 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -178,11 +178,19 @@ def model_params(options)
                 model_api_key: self_hosted_model.api_token,
                 model_identifier: self_hosted_model.identifier
               }
-            elsif top_namespace_model_selected?(unit_primitive)
-              {
-                provider: 'gitlab',
-                identifier: top_namespace_model_ref(unit_primitive)
-              }
+            elsif Feature.enabled?(:ai_model_switching, root_namespace)
+              if top_namespace_model_selected?(unit_primitive)
+                {
+                  provider: 'gitlab',
+                  identifier: top_namespace_model_ref(unit_primitive)
+                }
+              else
+                {
+                  provider: 'gitlab',
+                  feature_setting: top_namespace_model_feature(unit_primitive)
+                }
+              end
+
             else
               {
                 provider: provider(options),
@@ -191,6 +199,11 @@ def model_params(options)
             end
           end
 
+          def top_namespace_model_feature(unit_primitive)
+            feature = unit_primitive ? :"duo_chat_#{unit_primitive}" : :duo_chat
+            feature.to_s
+          end
+
           def payload_params(options)
             allowed_params = ALLOWED_PARAMS.fetch(provider(options))
             params = options.slice(*allowed_params)
@@ -230,11 +243,12 @@ def top_namespace_model_selected?(unit_primitive)
 
           def apply_model_selection_metadata(params, unit_primitive, _prompt_version)
             identifier = top_namespace_model_ref(unit_primitive)
-
-            params[:model_metadata] = {
-              provider: 'gitlab',
-              identifier: identifier
-            }
+            params[:model_metadata] = if identifier.present?
+                                        { provider: 'gitlab', identifier: identifier }
+                                      else
+                                        { provider: 'gitlab',
+                                          feature_setting: top_namespace_model_feature(unit_primitive) }
+                                      end
           end
 
           def processed_service_name(service_name)
diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
index af681564ffb9c0..521a43c7b54ccc 100644
--- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
@@ -93,6 +93,7 @@
       allow(Gitlab::Llm::Logger).to receive(:build).and_return(logger)
       allow(logger).to receive(:conditional_info)
       allow(instance).to receive(:ai_client).and_return(ai_client)
+      stub_feature_flags(ai_model_switching: false)
     end
 
     shared_examples 'performing request to the AI Gateway' do
@@ -390,9 +391,9 @@
         let(:prompt) { { prompt: user_prompt, options: {} } }
 
         before do
+          stub_feature_flags(ai_model_switching: true)
           create(:ai_namespace_feature_setting, namespace: root_namespace,
             feature: :"duo_chat_#{unit_primitive}", offered_model_ref: model_ref)
-          stub_feature_flags(ai_model_switching: true)
         end
 
         it 'uses namespace model ref in prompt body' do
@@ -430,9 +431,9 @@
         end
 
         before do
+          stub_feature_flags(ai_model_switching: true)
           create(:ai_namespace_feature_setting, namespace: root_namespace,
             feature: :"duo_chat_#{unit_primitive}", offered_model_ref: model_ref)
-          stub_feature_flags(ai_model_switching: true)
         end
 
         it 'sends agent request with model_metadata' do
@@ -483,22 +484,22 @@
 
       context 'when root_namespace is not root' do
         let(:subgroup) { create(:group, parent: root_namespace) }
-        let(:unit_primitive) { nil }
+        let(:unit_primitive) { 'write_tests' }
         let(:prompt) { { prompt: user_prompt, options: {} } }
 
         before do
           stub_feature_flags(ai_model_switching: true)
         end
 
-        it 'ignores model switching and uses default model' do
-          url = "#{::Gitlab::AiGateway.url}#{described_class::ENDPOINT}"
+        it 'ignores model switching and uses default feature_setting' do
+          url = "#{::Gitlab::AiGateway.url}/v1/chat/write_tests"
 
           expect(ai_client).to receive(:stream).with(
             hash_including(
               url: url,
               body: hash_including(
                 prompt_components: array_including(
-                  hash_including(payload: hash_including(model: described_class::CLAUDE_3_5_SONNET))
+                  hash_including(payload: hash_including(feature_setting: 'duo_chat_write_tests', provider: "gitlab"))
                 )
               )
             )
-- 
GitLab


From 90249cda1cb4116f739ec55c931f53ea1469de17 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Fri, 16 May 2025 09:19:55 +0200
Subject: [PATCH 05/11] Always pass both feature setting and identifier

---
 .../gitlab/llm/chain/requests/ai_gateway.rb   | 35 ++++-------
 .../llm/chain/requests/ai_gateway_spec.rb     | 61 ++++++++++++++++---
 2 files changed, 65 insertions(+), 31 deletions(-)

diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index bc3e046601871d..d7b49ebc22afa0 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -151,10 +151,7 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
             model_family = model_metadata_params && model_metadata_params[:name]
             default_version = ::Gitlab::Llm::PromptVersions.version_for_prompt("chat/#{unit_primitive}", model_family)
 
-            if top_namespace_model_selected?(unit_primitive)
-              apply_model_selection_metadata(params, unit_primitive,
-                prompt_version)
-            end
+            apply_model_selection_metadata(params, unit_primitive) if top_namespace_model_selected?(unit_primitive)
 
             params[:prompt_version] = if feature_setting&.self_hosted? || ::Ai::AmazonQ.connected?
                                         default_version
@@ -179,17 +176,7 @@ def model_params(options)
                 model_identifier: self_hosted_model.identifier
               }
             elsif Feature.enabled?(:ai_model_switching, root_namespace)
-              if top_namespace_model_selected?(unit_primitive)
-                {
-                  provider: 'gitlab',
-                  identifier: top_namespace_model_ref(unit_primitive)
-                }
-              else
-                {
-                  provider: 'gitlab',
-                  feature_setting: top_namespace_model_feature(unit_primitive)
-                }
-              end
+              model_selection_metadata(unit_primitive)
 
             else
               {
@@ -241,14 +228,16 @@ def top_namespace_model_selected?(unit_primitive)
             top_namespace_model_ref(unit_primitive).present?
           end
 
-          def apply_model_selection_metadata(params, unit_primitive, _prompt_version)
-            identifier = top_namespace_model_ref(unit_primitive)
-            params[:model_metadata] = if identifier.present?
-                                        { provider: 'gitlab', identifier: identifier }
-                                      else
-                                        { provider: 'gitlab',
-                                          feature_setting: top_namespace_model_feature(unit_primitive) }
-                                      end
+          def model_selection_metadata(unit_primitive)
+            {
+              provider: 'gitlab',
+              identifier: top_namespace_model_ref(unit_primitive),
+              feature_setting: top_namespace_model_feature(unit_primitive)
+            }
+          end
+
+          def apply_model_selection_metadata(params, unit_primitive)
+            params[:model_metadata] = model_selection_metadata(unit_primitive)
           end
 
           def processed_service_name(service_name)
diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
index 521a43c7b54ccc..443bf3d71264a3 100644
--- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
@@ -385,7 +385,7 @@
         end
       end
 
-      context 'when model switching is enabled and namespace is root' do
+      context 'when model switching is enabled and model is selected' do
         let(:unit_primitive) { :explain_code }
         let(:model_ref) { 'claude-3-7-sonnet-20250219' }
         let(:prompt) { { prompt: user_prompt, options: {} } }
@@ -396,7 +396,7 @@
             feature: :"duo_chat_#{unit_primitive}", offered_model_ref: model_ref)
         end
 
-        it 'uses namespace model ref in prompt body' do
+        it 'sends identifier and feature_setting' do
           url = "#{::Gitlab::AiGateway.url}#{described_class::BASE_ENDPOINT}/#{unit_primitive}"
 
           expect(ai_client).to receive(:stream).with(
@@ -406,6 +406,7 @@
                 prompt_components: array_including(
                   hash_including(payload: hash_including(
                     provider: "gitlab",
+                    feature_setting: 'duo_chat_explain_code',
                     identifier: model_ref
                   ))
                 )
@@ -418,6 +419,39 @@
         end
       end
 
+      context 'when model switching is enabled and model is NOT selected' do
+        let(:unit_primitive) { :explain_code }
+        let(:prompt) { { prompt: user_prompt, options: {} } }
+
+        before do
+          stub_feature_flags(ai_model_switching: true)
+          create(:ai_namespace_feature_setting, namespace: root_namespace,
+            feature: :"duo_chat_#{unit_primitive}", offered_model_ref: nil)
+        end
+
+        it 'sends only feature_setting (identifier is nil)' do
+          url = "#{::Gitlab::AiGateway.url}#{described_class::BASE_ENDPOINT}/#{unit_primitive}"
+
+          expect(ai_client).to receive(:stream).with(
+            hash_including(
+              url: url,
+              body: hash_including(
+                prompt_components: array_including(
+                  hash_including(payload: hash_including(
+                    provider: "gitlab",
+                    feature_setting: 'duo_chat_explain_code',
+                    identifier: nil
+                  ))
+                )
+              )
+            )
+          ).and_return(response)
+
+          gateway = described_class.new(user, root_namespace: root_namespace, tracking_context: tracking_context)
+          expect(gateway.request(prompt, unit_primitive: unit_primitive)).to eq(response)
+        end
+      end
+
       context 'when using agent prompt with model switching' do
         let(:unit_primitive) { :explain_code }
         let(:model_ref) { 'claude-3-7-sonnet-20250219' }
@@ -436,7 +470,7 @@
             feature: :"duo_chat_#{unit_primitive}", offered_model_ref: model_ref)
         end
 
-        it 'sends agent request with model_metadata' do
+        it 'sends model_metadata with identifier and feature_setting' do
           url = "#{::Gitlab::AiGateway.url}#{described_class::BASE_PROMPTS_CHAT_ENDPOINT}/#{unit_primitive}"
 
           expect(ai_client).to receive(:stream).with(
@@ -444,7 +478,11 @@
               url: url,
               body: hash_including(
                 inputs: { a: 1 },
-                model_metadata: { provider: 'gitlab', identifier: model_ref },
+                model_metadata: {
+                  provider: 'gitlab',
+                  feature_setting: 'duo_chat_explain_code',
+                  identifier: model_ref
+                },
                 prompt_version: a_kind_of(String)
               )
             )
@@ -463,7 +501,7 @@
           stub_feature_flags(ai_model_switching: false)
         end
 
-        it 'falls back to default model' do
+        it 'uses default classic model' do
           url = "#{::Gitlab::AiGateway.url}#{described_class::ENDPOINT}"
 
           expect(ai_client).to receive(:stream).with(
@@ -471,7 +509,10 @@
               url: url,
               body: hash_including(
                 prompt_components: array_including(
-                  hash_including(payload: hash_including(model: described_class::CLAUDE_3_5_SONNET))
+                  hash_including(payload: hash_including(
+                    provider: :anthropic,
+                    model: described_class::CLAUDE_3_5_SONNET
+                  ))
                 )
               )
             )
@@ -491,7 +532,7 @@
           stub_feature_flags(ai_model_switching: true)
         end
 
-        it 'ignores model switching and uses default feature_setting' do
+        it 'ignores model switching and sends only feature_setting with nil identifier' do
           url = "#{::Gitlab::AiGateway.url}/v1/chat/write_tests"
 
           expect(ai_client).to receive(:stream).with(
@@ -499,7 +540,11 @@
               url: url,
               body: hash_including(
                 prompt_components: array_including(
-                  hash_including(payload: hash_including(feature_setting: 'duo_chat_write_tests', provider: "gitlab"))
+                  hash_including(payload: hash_including(
+                    provider: "gitlab",
+                    feature_setting: 'duo_chat_write_tests',
+                    identifier: nil
+                  ))
                 )
               )
             )
-- 
GitLab


From 8580d49468b7ec4035467028b6d3efc38badd591 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Mon, 19 May 2025 10:24:36 +0200
Subject: [PATCH 06/11] Move logic to ModelMetadata for setting check

---
 .../gitlab/llm/ai_gateway/model_metadata.rb   | 10 +++++++
 .../gitlab/llm/chain/requests/ai_gateway.rb   | 30 +++++++++++--------
 .../namespace_feature_settings.rb             |  7 +++++
 .../llm/ai_gateway/model_metadata_spec.rb     | 24 +++++++++++++++
 .../llm/chain/requests/ai_gateway_spec.rb     |  7 ++---
 5 files changed, 61 insertions(+), 17 deletions(-)

diff --git a/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb b/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb
index 52d78e73d9a4c0..5219299e1e77de 100644
--- a/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb
@@ -9,6 +9,8 @@ def initialize(feature_setting: nil)
         end
 
         def to_params
+          return namespace_settings_params if feature_setting.is_a?(::Ai::ModelSelection::NamespaceFeatureSetting)
+
           return self_hosted_params if feature_setting&.self_hosted?
 
           amazon_q_params if ::Ai::AmazonQ.connected?
@@ -32,6 +34,14 @@ def self_hosted_params
 
         attr_reader :feature_setting
 
+        def namespace_settings_params
+          {
+            provider: 'gitlab',
+            identifier: feature_setting.offered_model_ref,
+            feature_setting: feature_setting.feature
+          }
+        end
+
         def amazon_q_params
           {
             provider: :amazon_q,
diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index d7b49ebc22afa0..4d70b4b610109e 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -175,9 +175,12 @@ def model_params(options)
                 model_api_key: self_hosted_model.api_token,
                 model_identifier: self_hosted_model.identifier
               }
-            elsif Feature.enabled?(:ai_model_switching, root_namespace)
-              model_selection_metadata(unit_primitive)
+            elsif root_namespace&.root? && Feature.enabled?(:ai_model_switching, root_namespace)
+              feature = top_namespace_model_feature(unit_primitive)
+              setting = ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(root_namespace,
+                feature)
 
+              ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: setting).to_params
             else
               {
                 provider: provider(options),
@@ -187,8 +190,7 @@ def model_params(options)
           end
 
           def top_namespace_model_feature(unit_primitive)
-            feature = unit_primitive ? :"duo_chat_#{unit_primitive}" : :duo_chat
-            feature.to_s
+            unit_primitive ? :"duo_chat_#{unit_primitive}" : "duo_chat"
           end
 
           def payload_params(options)
@@ -228,16 +230,18 @@ def top_namespace_model_selected?(unit_primitive)
             top_namespace_model_ref(unit_primitive).present?
           end
 
-          def model_selection_metadata(unit_primitive)
-            {
-              provider: 'gitlab',
-              identifier: top_namespace_model_ref(unit_primitive),
-              feature_setting: top_namespace_model_feature(unit_primitive)
-            }
-          end
-
           def apply_model_selection_metadata(params, unit_primitive)
-            params[:model_metadata] = model_selection_metadata(unit_primitive)
+            feature = top_namespace_model_feature(unit_primitive)
+
+            return unless root_namespace&.root? && Feature.enabled?(:ai_model_switching, root_namespace)
+
+            setting = ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(
+              root_namespace, feature
+            )
+
+            model_metadata = ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: setting).to_params
+
+            params[:model_metadata] = model_metadata if model_metadata
           end
 
           def processed_service_name(service_name)
diff --git a/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb b/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb
index b0f6aed696d687..1f879db71972aa 100644
--- a/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb
+++ b/ee/spec/factories/ai/model_selection/namespace_feature_settings.rb
@@ -39,6 +39,13 @@
             "selectable_models" => %w[claude-3-7-sonnet-20250219 claude-3-5-sonnet-20240620],
             "beta_models" => [],
             "unit_primitives" => ["explain_code"]
+          },
+          {
+            "feature_setting" => "duo_chat",
+            "default_model" => "claude-3-7-sonnet-20250219",
+            "selectable_models" => %w[claude-3-7-sonnet-20250219 claude_3_5_sonnet_20240620],
+            "beta_models" => [],
+            "unit_primitives" => %w[ask_build ask_commit]
           }
         ]
       }
diff --git a/ee/spec/lib/gitlab/llm/ai_gateway/model_metadata_spec.rb b/ee/spec/lib/gitlab/llm/ai_gateway/model_metadata_spec.rb
index 19f052c3647ec3..82f5d81cd4fc39 100644
--- a/ee/spec/lib/gitlab/llm/ai_gateway/model_metadata_spec.rb
+++ b/ee/spec/lib/gitlab/llm/ai_gateway/model_metadata_spec.rb
@@ -50,6 +50,30 @@
 
       it { is_expected.to be_nil }
     end
+
+    context 'when feature_setting is a NamespaceFeatureSetting' do
+      let_it_be(:root_namespace) { create(:group) }
+
+      let(:feature_name) { 'duo_chat' }
+      let(:model_ref) { 'claude-3-7-sonnet-20250219' }
+
+      let(:feature_setting) do
+        create(
+          :ai_namespace_feature_setting,
+          namespace: root_namespace,
+          feature: feature_name,
+          offered_model_ref: model_ref
+        )
+      end
+
+      it 'returns the correct namespace_settings_params' do
+        is_expected.to eq({
+          provider: 'gitlab',
+          identifier: model_ref,
+          feature_setting: 'duo_chat'
+        })
+      end
+    end
   end
 
   describe '#self_hosted_params' do
diff --git a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
index 443bf3d71264a3..93e1bdc289a525 100644
--- a/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
+++ b/ee/spec/lib/gitlab/llm/chain/requests/ai_gateway_spec.rb
@@ -532,7 +532,7 @@
           stub_feature_flags(ai_model_switching: true)
         end
 
-        it 'ignores model switching and sends only feature_setting with nil identifier' do
+        it 'ignores model switching' do
           url = "#{::Gitlab::AiGateway.url}/v1/chat/write_tests"
 
           expect(ai_client).to receive(:stream).with(
@@ -541,9 +541,8 @@
               body: hash_including(
                 prompt_components: array_including(
                   hash_including(payload: hash_including(
-                    provider: "gitlab",
-                    feature_setting: 'duo_chat_write_tests',
-                    identifier: nil
+                    provider: :anthropic,
+                    model: 'claude-3-5-sonnet-20240620'
                   ))
                 )
               )
-- 
GitLab


From 06664a88ccaeffd029a15f7419dbf3fac65d8609 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Tue, 20 May 2025 09:39:40 +0200
Subject: [PATCH 07/11] Add Duo chat root_namespace settings support

---
 doc/api/graphql/reference/_index.md           |  1 +
 .../ai/graphql/chat.mutation.graphql          |  2 +
 .../ai/tanuki_bot/components/app.vue          |  6 ++
 .../assets/javascripts/ai/tanuki_bot/index.js |  3 +-
 ee/app/graphql/mutations/ai/action.rb         |  6 +-
 .../views/layouts/_tanuki_bot_chat.html.haml  |  3 +-
 ee/lib/gitlab/duo/chat/react_executor.rb      |  9 +-
 .../gitlab/llm/chain/requests/ai_gateway.rb   |  4 +-
 ee/lib/gitlab/llm/completions/chat.rb         |  9 +-
 ee/lib/gitlab/llm/tanuki_bot.rb               |  5 ++
 ee/spec/graphql/mutations/ai/action_spec.rb   | 85 +++++++++++++++++++
 .../gitlab/duo/chat/react_executor_spec.rb    | 41 +++++++++
 .../lib/gitlab/llm/completions/chat_spec.rb   | 65 +++++++++++++-
 ee/spec/lib/gitlab/llm/tanuki_bot_spec.rb     | 34 ++++++++
 .../_tanuki_bot_chat.html.haml_spec.rb        |  9 +-
 15 files changed, 271 insertions(+), 11 deletions(-)

diff --git a/doc/api/graphql/reference/_index.md b/doc/api/graphql/reference/_index.md
index 67f9bb9c57fa85..25f211c217e3e5 100644
--- a/doc/api/graphql/reference/_index.md
+++ b/doc/api/graphql/reference/_index.md
@@ -2068,6 +2068,7 @@ Input type: `AiActionInput`
 | <a id="mutationaiactionplatformorigin"></a>`platformOrigin` | [`String`](#string) | Specifies the origin platform of the request. |
 | <a id="mutationaiactionprojectid"></a>`projectId` | [`ProjectID`](#projectid) | Global ID of the project the user is acting on. |
 | <a id="mutationaiactionresolvevulnerability"></a>`resolveVulnerability` | [`AiResolveVulnerabilityInput`](#airesolvevulnerabilityinput) | Input for resolve_vulnerability AI action. |
+| <a id="mutationaiactionrootnamespaceid"></a>`rootNamespaceId` | [`NamespaceID`](#namespaceid) | Global ID of the top-level namespace the user is acting on. |
 | <a id="mutationaiactionsummarizecomments"></a>`summarizeComments` | [`AiSummarizeCommentsInput`](#aisummarizecommentsinput) | Input for summarize_comments AI action. |
 | <a id="mutationaiactionsummarizenewmergerequest"></a>`summarizeNewMergeRequest` | [`AiSummarizeNewMergeRequestInput`](#aisummarizenewmergerequestinput) | Input for summarize_new_merge_request AI action. |
 | <a id="mutationaiactionsummarizereview"></a>`summarizeReview` | [`AiSummarizeReviewInput`](#aisummarizereviewinput) | Input for summarize_review AI action. |
diff --git a/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql b/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
index 43618561a556d9..86cd40a13a0a7f 100644
--- a/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
+++ b/ee/app/assets/javascripts/ai/graphql/chat.mutation.graphql
@@ -7,6 +7,7 @@ mutation chat(
   $currentFileContext: AiCurrentFileInput
   $conversationType: AiConversationsThreadsConversationType
   $threadId: AiConversationThreadID
+  $rootNamespaceId: NamespaceID
 ) {
   aiAction(
     input: {
@@ -20,6 +21,7 @@ mutation chat(
       clientSubscriptionId: $clientSubscriptionId
       conversationType: $conversationType
       threadId: $threadId
+      rootNamespaceId: $rootNamespaceId
     }
   ) {
     requestId
diff --git a/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue b/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
index 382c0c846c2174..bc832c8a89b57c 100644
--- a/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
+++ b/ee/app/assets/javascripts/ai/tanuki_bot/components/app.vue
@@ -75,6 +75,11 @@ export default {
       required: false,
       default: null,
     },
+    rootNamespaceId: {
+      type: String,
+      required: false,
+      default: null,
+    },
     chatTitle: {
       type: String,
       required: false,
@@ -354,6 +359,7 @@ export default {
         projectId: this.projectId,
         threadId: this.activeThread,
         conversationType: MULTI_THREADED_CONVERSATION_TYPE,
+        rootNamespaceId: this.rootNamespaceId,
         ...variables,
       };
 
diff --git a/ee/app/assets/javascripts/ai/tanuki_bot/index.js b/ee/app/assets/javascripts/ai/tanuki_bot/index.js
index 2f84399235ad91..b45f5e4cb9c8f6 100644
--- a/ee/app/assets/javascripts/ai/tanuki_bot/index.js
+++ b/ee/app/assets/javascripts/ai/tanuki_bot/index.js
@@ -27,7 +27,7 @@ export const initTanukiBotChatDrawer = () => {
     });
   }
 
-  const { userId, resourceId, projectId, chatTitle } = el.dataset;
+  const { userId, resourceId, projectId, chatTitle, rootNamespaceId } = el.dataset;
 
   return new Vue({
     el,
@@ -40,6 +40,7 @@ export const initTanukiBotChatDrawer = () => {
           resourceId,
           projectId,
           chatTitle,
+          rootNamespaceId,
         },
       });
     },
diff --git a/ee/app/graphql/mutations/ai/action.rb b/ee/app/graphql/mutations/ai/action.rb
index 9c9d2797b58d90..3570755fe75e49 100644
--- a/ee/app/graphql/mutations/ai/action.rb
+++ b/ee/app/graphql/mutations/ai/action.rb
@@ -28,6 +28,10 @@ class Action < BaseMutation
         required: false,
         description: "Global ID of the project the user is acting on."
 
+      argument :root_namespace_id, ::Types::GlobalIDType[::Namespace],
+        required: false,
+        description: "Global ID of the top-level namespace the user is acting on."
+
       argument :conversation_type, Types::Ai::Conversations::Threads::ConversationTypeEnum,
         required: false,
         description: 'Conversation type of the thread.'
@@ -174,7 +178,7 @@ def authorized_resource?(object)
 
       def extract_method_params!(attributes)
         options = attributes.extract!(:client_subscription_id, :platform_origin, :project_id,
-          :conversation_type, :thread_id)
+          :conversation_type, :thread_id, :root_namespace_id)
         methods = methods(attributes.transform_values(&:to_h))
 
         # At this point, we only have one method since we filtered it in `#ready?`
diff --git a/ee/app/views/layouts/_tanuki_bot_chat.html.haml b/ee/app/views/layouts/_tanuki_bot_chat.html.haml
index db5282683471ef..96a515c682e098 100644
--- a/ee/app/views/layouts/_tanuki_bot_chat.html.haml
+++ b/ee/app/views/layouts/_tanuki_bot_chat.html.haml
@@ -1,6 +1,7 @@
 - return unless ::Gitlab::Llm::TanukiBot.enabled_for?(user: current_user, container: nil)
 - resource_id = Gitlab::Llm::TanukiBot.resource_id
 - project_id = Gitlab::Llm::TanukiBot.project_id
+- root_namespace_id = Gitlab::Llm::TanukiBot.root_namespace_id
 - chat_title = ::Ai::AmazonQ.enabled? ? s_('GitLab Duo Chat with Amazon Q') : s_('GitLab Duo Chat')
 
-#js-tanuki-bot-chat-app{ data: { user_id: current_user.to_global_id, resource_id: resource_id, project_id: project_id, chat_title: chat_title } }
+#js-tanuki-bot-chat-app{ data: { user_id: current_user.to_global_id, resource_id: resource_id, project_id: project_id, root_namespace_id: root_namespace_id, chat_title: chat_title } }
diff --git a/ee/lib/gitlab/duo/chat/react_executor.rb b/ee/lib/gitlab/duo/chat/react_executor.rb
index 1fa7f2db80b93f..4e90b128d97b3d 100644
--- a/ee/lib/gitlab/duo/chat/react_executor.rb
+++ b/ee/lib/gitlab/duo/chat/react_executor.rb
@@ -339,7 +339,14 @@ def current_blob
         strong_memoize_attr :current_blob
 
         def chat_feature_setting
-          ::Ai::FeatureSetting.find_by_feature(:duo_chat)
+          root_namespace = context&.ai_request&.root_namespace
+
+          if Feature.enabled?(:ai_model_switching, root_namespace)
+            ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(root_namespace,
+              :duo_chat)
+          else
+            ::Ai::FeatureSetting.find_by_feature(:duo_chat)
+          end
         end
 
         def record_first_token_apex
diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index 4d70b4b610109e..f01e3b7c48e849 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -11,7 +11,7 @@ class AiGateway < Base
           include ::Gitlab::Llm::Concerns::AllowedParams
           include ::Gitlab::Llm::Concerns::EventTracking
 
-          attr_reader :ai_client, :tracking_context
+          attr_reader :ai_client, :tracking_context, :root_namespace
 
           ENDPOINT = '/v1/chat/agent'
           BASE_ENDPOINT = '/v1/chat'
@@ -59,7 +59,7 @@ def request(prompt, unit_primitive: nil)
 
           private
 
-          attr_reader :user, :root_namespace
+          attr_reader :user
 
           def default_options
             {
diff --git a/ee/lib/gitlab/llm/completions/chat.rb b/ee/lib/gitlab/llm/completions/chat.rb
index 0d968779983211..ff4298737c0cb6 100644
--- a/ee/lib/gitlab/llm/completions/chat.rb
+++ b/ee/lib/gitlab/llm/completions/chat.rb
@@ -56,7 +56,7 @@ def ai_request
           ::Gitlab::Llm::Chain::Requests::AiGateway.new(
             user,
             tracking_context: tracking_context,
-            root_namespace: resource.try(:resource_parent)&.root_ancestor
+            root_namespace: resource.try(:resource_parent)&.root_ancestor || find_root_namespace
           )
         end
 
@@ -103,6 +103,13 @@ def tools
           TOOLS
         end
 
+        def find_root_namespace
+          return unless options[:root_namespace_id]
+
+          root_namespace_id = GlobalID.parse(options[:root_namespace_id])
+          ::Gitlab::Graphql::Lazy.force(GitlabSchema.find_by_gid(root_namespace_id))
+        end
+
         def response_post_processing
           return if Rails.env.development?
           return unless Gitlab::Saas.feature_available?(:duo_chat_categorize_question)
diff --git a/ee/lib/gitlab/llm/tanuki_bot.rb b/ee/lib/gitlab/llm/tanuki_bot.rb
index 0770ec25bfcb64..9acc8b572896e2 100644
--- a/ee/lib/gitlab/llm/tanuki_bot.rb
+++ b/ee/lib/gitlab/llm/tanuki_bot.rb
@@ -51,6 +51,11 @@ def self.project_id
         project_path = Gitlab::ApplicationContext.current_context_attribute(:project).presence
         Project.find_by_full_path(project_path).try(:to_global_id) if project_path
       end
+
+      def self.root_namespace_id
+        namespace_path = Gitlab::ApplicationContext.current_context_attribute(:root_namespace).presence
+        Group.find_by_full_path(namespace_path).try(:to_global_id) if namespace_path
+      end
     end
   end
 end
diff --git a/ee/spec/graphql/mutations/ai/action_spec.rb b/ee/spec/graphql/mutations/ai/action_spec.rb
index 170f5756709577..a56f61132806ce 100644
--- a/ee/spec/graphql/mutations/ai/action_spec.rb
+++ b/ee/spec/graphql/mutations/ai/action_spec.rb
@@ -476,5 +476,90 @@
 
       it_behaves_like 'an AI action'
     end
+
+    context 'when root_namespace_id is specified' do
+      let_it_be(:group) { create(:group) }
+
+      let(:input) { { chat: { resource_id: resource_id }, root_namespace_id: group.to_global_id } }
+
+      before_all do
+        resource.project.add_developer(user)
+        group.add_developer(user)
+      end
+
+      before do
+        allow(Ability).to receive(:allowed?).and_call_original
+        allow(Ability)
+          .to receive(:allowed?)
+          .with(user, :read_namespace, group)
+          .and_return(true)
+      end
+
+      it 'passes the root_namespace_id to the service' do
+        expect_next_instance_of(
+          Llm::ExecuteMethodService,
+          user,
+          resource,
+          :chat,
+          hash_including(root_namespace_id: kind_of(GlobalID))
+        ) do |svc|
+          expect(svc)
+            .to receive(:execute)
+            .and_return(ServiceResponse.success(
+              payload: {
+                ai_message: build(:ai_message, request_id: request_id)
+              }))
+        end
+
+        result = subject
+        expect(result[:errors]).to be_empty
+        expect(result[:request_id]).to eq(request_id)
+      end
+
+      context 'when resource is null' do
+        let(:resource_id) { nil }
+
+        it 'passes the root_namespace_id to the service' do
+          expect_next_instance_of(
+            Llm::ExecuteMethodService,
+            user,
+            nil,
+            :chat,
+            hash_including(root_namespace_id: kind_of(GlobalID))
+          ) do |svc|
+            expect(svc)
+              .to receive(:execute)
+              .and_return(ServiceResponse.success(
+                payload: {
+                  ai_message: build(:ai_message, request_id: request_id)
+                }))
+          end
+
+          result = subject
+          expect(result[:errors]).to be_empty
+          expect(result[:request_id]).to eq(request_id)
+        end
+      end
+
+      context 'when service returns an error' do
+        it 'returns the error message' do
+          expect_next_instance_of(
+            Llm::ExecuteMethodService,
+            user,
+            resource,
+            :chat,
+            hash_including(root_namespace_id: kind_of(GlobalID))
+          ) do |svc|
+            expect(svc)
+              .to receive(:execute)
+              .and_return(ServiceResponse.error(message: 'error'))
+          end
+
+          result = subject
+          expect(result[:errors]).to eq(['error'])
+          expect(result[:request_id]).to be_nil
+        end
+      end
+    end
   end
 end
diff --git a/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb b/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
index 9a7da107954866..d90474f3e10f79 100644
--- a/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
+++ b/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
@@ -568,6 +568,47 @@ def expect_sli_error(failed)
       end
     end
 
+    context 'when Duo chat model is selected at namespace level' do
+      let_it_be(:root_namespace) { create(:group) }
+      let(:ai_request) { instance_double(::Gitlab::Llm::Chain::Requests::AiGateway, root_namespace: root_namespace) }
+      let(:context) do
+        Gitlab::Llm::Chain::GitlabContext.new(
+          current_user: user,
+          container: nil,
+          resource: resource,
+          ai_request: ai_request,
+          extra_resource: extra_resource,
+          started_at: started_at_timestamp,
+          current_file: current_file,
+          agent_version: nil,
+          additional_context: additional_context
+        )
+      end
+
+      let_it_be(:feature_setting) do
+        create(:ai_namespace_feature_setting,
+          namespace: root_namespace,
+          feature: :duo_chat,
+          offered_model_ref: 'claude-3-7-sonnet-20250219')
+      end
+
+      it 'sends the namespace model metadata' do
+        params = step_params
+        params[:model_metadata] = {
+          provider: 'gitlab',
+          identifier: 'claude-3-7-sonnet-20250219',
+          feature_setting: 'duo_chat'
+        }
+
+        expect_next_instance_of(Gitlab::Duo::Chat::StepExecutor) do |react_agent|
+          expect(react_agent).to receive(:step).with(params)
+            .and_yield(action_event).and_return([action_event])
+        end
+
+        agent.execute
+      end
+    end
+
     context 'when amazon q is connected' do
       let_it_be(:add_on_purchase) { create(:gitlab_subscription_add_on_purchase, :duo_amazon_q) }
 
diff --git a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
index f9d880e6c0bd5a..489ac0bfb13912 100644
--- a/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
+++ b/ee/spec/lib/gitlab/llm/completions/chat_spec.rb
@@ -218,12 +218,71 @@
   end
 
   describe '.initialize' do
-    subject { described_class.new(prompt_message, nil, **options) }
+    let_it_be(:root_group) { create(:group) }
+    let_it_be(:subgroup) { create(:group, parent: root_group) }
+    let_it_be(:project) { create(:project, namespace: subgroup) }
+    let_it_be(:issue) { create(:issue, project: project) }
+    let_it_be(:user) { create(:user) }
+
+    let(:resource) { issue }
+    let(:prompt_message) do
+      build(:ai_chat_message, user: user, resource: resource, request_id: 'uuid', content: content, thread: thread)
+    end
+
+    let(:options) do
+      {
+        current_file: nil,
+        additional_context: [],
+        started_at: Time.current,
+        agent_version_id: nil,
+        extra_resource: {},
+        tracking_context: {},
+        resource: resource,
+        user: user
+      }
+    end
+
+    subject(:instance) do
+      described_class.new(prompt_message, nil, **options)
+    end
 
     it 'trims additional context' do
-      expect(::CodeSuggestions::Context).to receive(:new).with(additional_context).and_call_original
+      expect(::CodeSuggestions::Context).to receive(:new).with([]).and_call_original
+      instance
+    end
 
-      subject
+    it 'sets root_namespace correctly on ai_request' do
+      expect(instance.send(:context).ai_request.root_namespace).to eq(root_group)
+    end
+
+    context 'when resource has no resource_parent, fallback is used via root_namespace_global_id' do
+      let_it_be(:fallback_namespace) { create(:group) }
+
+      let(:prompt_message) do
+        build(:ai_chat_message, user: user, resource: user, request_id: 'uuid', content: content, thread: thread)
+      end
+
+      let(:options) do
+        {
+          current_file: nil,
+          additional_context: [],
+          started_at: Time.current,
+          agent_version_id: nil,
+          extra_resource: {},
+          tracking_context: {},
+          resource: user,
+          user: user,
+          root_namespace_id: fallback_namespace.to_global_id.to_s
+        }
+      end
+
+      subject(:instance) do
+        described_class.new(prompt_message, nil, **options)
+      end
+
+      it 'uses root_namespace_global_id fallback when resource_parent is nil' do
+        expect(instance.send(:context).ai_request.root_namespace).to eq(fallback_namespace)
+      end
     end
   end
 
diff --git a/ee/spec/lib/gitlab/llm/tanuki_bot_spec.rb b/ee/spec/lib/gitlab/llm/tanuki_bot_spec.rb
index 22b8552bb0a55e..a782fc6f34b900 100644
--- a/ee/spec/lib/gitlab/llm/tanuki_bot_spec.rb
+++ b/ee/spec/lib/gitlab/llm/tanuki_bot_spec.rb
@@ -237,4 +237,38 @@
       end
     end
   end
+
+  describe '.root_namespace_id' do
+    let_it_be(:group) { create(:group) }
+
+    context 'with current context including root_namespace' do
+      let(:result) do
+        ::Gitlab::ApplicationContext.with_raw_context(root_namespace: group.full_path) do
+          described_class.root_namespace_id
+        end
+      end
+
+      it 'returns the global ID of the root namespace when found' do
+        expect(result).to eq(group.to_global_id)
+      end
+    end
+
+    context 'when root_namespace is not found' do
+      let(:result) do
+        ::Gitlab::ApplicationContext.with_raw_context(root_namespace: 'non_existent_namespace') do
+          described_class.root_namespace_id
+        end
+      end
+
+      it 'returns nil' do
+        expect(result).to be_nil
+      end
+    end
+
+    context 'when root_namespace is not present in the context' do
+      it 'returns nil' do
+        expect(described_class.root_namespace_id).to be_nil
+      end
+    end
+  end
 end
diff --git a/ee/spec/views/layouts/_tanuki_bot_chat.html.haml_spec.rb b/ee/spec/views/layouts/_tanuki_bot_chat.html.haml_spec.rb
index 5b76bd6e8602db..c3ec7b334d34fa 100644
--- a/ee/spec/views/layouts/_tanuki_bot_chat.html.haml_spec.rb
+++ b/ee/spec/views/layouts/_tanuki_bot_chat.html.haml_spec.rb
@@ -10,10 +10,17 @@
     allow(::Gitlab::Llm::TanukiBot).to receive_messages(
       enabled_for?: true,
       resource_id: 'test_resource_id',
-      project_id: 'test_project_id'
+      project_id: 'test_project_id',
+      root_namespace_id: 'test_root_namespace_id'
     )
   end
 
+  it 'includes the root_namespace_id in the data attributes' do
+    render
+
+    expect(rendered).to have_css("#js-tanuki-bot-chat-app[data-root-namespace-id='test_root_namespace_id']")
+  end
+
   context 'when AmazonQ is enabled' do
     before do
       allow(::Ai::AmazonQ).to receive(:enabled?).and_return(true)
-- 
GitLab


From 579cf8f89aac036ba9356a53ab5e01ff721987ab Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Tue, 20 May 2025 10:24:00 +0200
Subject: [PATCH 08/11] Fix failing FE and BE specs

Add find_or_initialize_by_feature coverage
Add rootNamespaceId case
---
 .../namespace_feature_setting.rb              |  2 +-
 .../ai/tanuki_bot/components/app_spec.js      | 30 +++++++++++++++++++
 .../gitlab/duo/chat/react_executor_spec.rb    |  4 +++
 .../namespace_feature_setting_spec.rb         |  7 +++++
 4 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/ee/app/models/ai/model_selection/namespace_feature_setting.rb b/ee/app/models/ai/model_selection/namespace_feature_setting.rb
index 28397999174899..97936ae3e8bdf4 100644
--- a/ee/app/models/ai/model_selection/namespace_feature_setting.rb
+++ b/ee/app/models/ai/model_selection/namespace_feature_setting.rb
@@ -17,7 +17,7 @@ class NamespaceFeatureSetting < ApplicationRecord
       scope :for_namespace, ->(namespace_id) { where(namespace_id: namespace_id) }
 
       def self.find_or_initialize_by_feature(namespace, feature)
-        return unless ::Feature.enabled?(:ai_model_switching, namespace)
+        return unless namespace.present? && ::Feature.enabled?(:ai_model_switching, namespace)
         return unless namespace.root?
 
         find_or_initialize_by(namespace_id: namespace.id, feature: feature)
diff --git a/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js b/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js
index 80e1ef655db97c..089e659287d7f2 100644
--- a/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js
+++ b/ee/spec/frontend/ai/tanuki_bot/components/app_spec.js
@@ -328,6 +328,7 @@ describeSkipVue3(skipReason, () => {
         resourceId: 'command::1',
         projectId: null,
         conversationType: 'DUO_CHAT',
+        rootNamespaceId: null,
         threadId: undefined,
       });
     });
@@ -398,6 +399,7 @@ describeSkipVue3(skipReason, () => {
           resourceId: MOCK_USER_ID,
           projectId: 'project-123',
           conversationType: 'DUO_CHAT',
+          rootNamespaceId: null,
           threadId: undefined,
         });
       });
@@ -417,6 +419,7 @@ describeSkipVue3(skipReason, () => {
           resourceId: MOCK_RESOURCE_ID,
           projectId: null,
           conversationType: 'DUO_CHAT',
+          rootNamespaceId: null,
           threadId: undefined,
         });
       });
@@ -464,6 +467,7 @@ describeSkipVue3(skipReason, () => {
             clientSubscriptionId: '123',
             projectId: null,
             conversationType: 'DUO_CHAT',
+            rootNamespaceId: null,
             threadId: undefined,
           });
         });
@@ -850,6 +854,7 @@ describeSkipVue3(skipReason, () => {
           resourceId: 'gid://gitlab/Issue/1',
           projectId: null,
           conversationType: 'DUO_CHAT',
+          rootNamespaceId: null,
         });
       });
 
@@ -869,9 +874,34 @@ describeSkipVue3(skipReason, () => {
           resourceId: MOCK_RESOURCE_ID,
           projectId: null,
           conversationType: 'DUO_CHAT',
+          rootNamespaceId: null,
           threadId: mockThreadId,
         });
       });
+      it('passes rootNamespaceId when provided in props', async () => {
+        // Create component with rootNamespaceId in props
+        createComponent({
+          propsData: {
+            userId: MOCK_USER_ID,
+            resourceId: MOCK_RESOURCE_ID,
+            rootNamespaceId: 'namespace-123',
+          },
+        });
+
+        findDuoChat().vm.$emit('send-chat-prompt', MOCK_USER_MESSAGE.content);
+
+        await nextTick();
+
+        expect(chatMutationHandlerMock).toHaveBeenCalledWith({
+          clientSubscriptionId: '123',
+          question: MOCK_USER_MESSAGE.content,
+          resourceId: MOCK_RESOURCE_ID,
+          projectId: null,
+          conversationType: 'DUO_CHAT',
+          rootNamespaceId: 'namespace-123',
+          threadId: undefined,
+        });
+      });
     });
 
     describe('thread handling', () => {
diff --git a/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb b/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
index d90474f3e10f79..3c39b566089834 100644
--- a/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
+++ b/ee/spec/lib/gitlab/duo/chat/react_executor_spec.rb
@@ -549,6 +549,10 @@ def expect_sli_error(failed)
       let_it_be(:self_hosted_model) { create(:ai_self_hosted_model, api_token: 'test_token') }
       let_it_be(:ai_feature) { create(:ai_feature_setting, self_hosted_model: self_hosted_model, feature: :duo_chat) }
 
+      before do
+        stub_feature_flags(ai_model_switching: false)
+      end
+
       it 'sends the self-hosted model metadata' do
         params = step_params
         params[:model_metadata] = {
diff --git a/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb b/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb
index 21145b387f06bc..fcb2496039baed 100644
--- a/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb
+++ b/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb
@@ -28,6 +28,13 @@
     let(:existing_feature) { ai_feature_setting.feature.to_sym }
     let(:new_feature_enum) { :code_completions }
 
+    context 'when namespace is nil' do
+      it 'returns nil' do
+        result = described_class.find_or_initialize_by_feature(nil, existing_feature)
+        expect(result).to be_nil
+      end
+    end
+
     it 'returns existing setting when one exists for the feature' do
       ai_feature_setting.save!
       result = described_class.find_or_initialize_by_feature(group, existing_feature)
-- 
GitLab


From 82fb20ce980e51da0934d7cd1a046f38c924735e Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Wed, 21 May 2025 10:13:39 +0200
Subject: [PATCH 09/11] Refactor Chain::Requests::AiGateway methods

Add find_by_feature method and guards
Change the model params usage
---
 .../namespace_feature_setting.rb              |  7 ++
 ee/app/workers/llm/completion_worker.rb       |  2 +-
 ee/lib/gitlab/duo/chat/react_executor.rb      |  5 +-
 .../gitlab/llm/chain/requests/ai_gateway.rb   | 97 ++++++++-----------
 .../namespace_feature_setting_spec.rb         | 56 ++++++++++-
 5 files changed, 103 insertions(+), 64 deletions(-)

diff --git a/ee/app/models/ai/model_selection/namespace_feature_setting.rb b/ee/app/models/ai/model_selection/namespace_feature_setting.rb
index 97936ae3e8bdf4..ae69932a350b6a 100644
--- a/ee/app/models/ai/model_selection/namespace_feature_setting.rb
+++ b/ee/app/models/ai/model_selection/namespace_feature_setting.rb
@@ -23,6 +23,13 @@ def self.find_or_initialize_by_feature(namespace, feature)
         find_or_initialize_by(namespace_id: namespace.id, feature: feature)
       end
 
+      def self.find_by_feature(namespace, feature)
+        return unless namespace.present? && ::Feature.enabled?(:ai_model_switching, namespace)
+        return unless namespace.root?
+
+        find_by(namespace: namespace, feature: feature)
+      end
+
       def model_selection_scope
         namespace
       end
diff --git a/ee/app/workers/llm/completion_worker.rb b/ee/app/workers/llm/completion_worker.rb
index 0cde546cf402c8..1397000cb082e3 100644
--- a/ee/app/workers/llm/completion_worker.rb
+++ b/ee/app/workers/llm/completion_worker.rb
@@ -45,7 +45,7 @@ def perform_for(message, options = {})
         # set SESSION_ID_HASH_KEY to ensure inside Sidekiq `Gitlab::Session.current` is not nil
         with_ip_address_state.set(
           Gitlab::SidekiqMiddleware::SetSession::Server::SESSION_ID_HASH_KEY => ::Gitlab::Session.session_id_for_worker
-        ).perform_async(serialize_message(message), options)
+        ).perform_inline(serialize_message(message), options)
       end
 
       def resource(message_hash)
diff --git a/ee/lib/gitlab/duo/chat/react_executor.rb b/ee/lib/gitlab/duo/chat/react_executor.rb
index 4e90b128d97b3d..a1df84bbf3b1dd 100644
--- a/ee/lib/gitlab/duo/chat/react_executor.rb
+++ b/ee/lib/gitlab/duo/chat/react_executor.rb
@@ -339,11 +339,10 @@ def current_blob
         strong_memoize_attr :current_blob
 
         def chat_feature_setting
-          root_namespace = context&.ai_request&.root_namespace
+          root_namespace = context.ai_request&.root_namespace
 
           if Feature.enabled?(:ai_model_switching, root_namespace)
-            ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(root_namespace,
-              :duo_chat)
+            ::Ai::ModelSelection::NamespaceFeatureSetting.find_by_feature(root_namespace, :duo_chat)
           else
             ::Ai::FeatureSetting.find_by_feature(:duo_chat)
           end
diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index f01e3b7c48e849..e47b0e16d47a74 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -32,8 +32,6 @@ def initialize(user, service_name: :duo_chat, tracking_context: {}, root_namespa
 
           def request(prompt, unit_primitive: nil)
             options = default_options.merge(prompt.fetch(:options, {}))
-            options[:unit_primitive] ||= unit_primitive
-
             return unless model_provider_valid?(options)
 
             response = ai_client.stream(
@@ -85,16 +83,6 @@ def model_provider_valid?(options)
             provider(options)
           end
 
-          def top_namespace_model_ref(unit_primitive = nil)
-            return unless root_namespace&.root?
-            return unless Feature.enabled?(:ai_model_switching, root_namespace)
-
-            feature = unit_primitive ? :"duo_chat_#{unit_primitive}" : :duo_chat
-            setting = ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(root_namespace,
-              feature)
-            setting&.offered_model_ref
-          end
-
           def endpoint(unit_primitive, use_ai_gateway_agent_prompt)
             path =
               if use_ai_gateway_agent_prompt
@@ -116,11 +104,11 @@ def body(prompt, options, unit_primitive: nil)
               request_body_agent(inputs: options[:inputs], unit_primitive: unit_primitive,
                 prompt_version: options[:prompt_version])
             else
-              request_body(prompt: prompt[:prompt], options: options)
+              request_body(prompt: prompt[:prompt], options: options, unit_primitive: unit_primitive)
             end
           end
 
-          def request_body(prompt:, options: {})
+          def request_body(prompt:, options: {}, unit_primitive: nil)
             {
               prompt_components: [{
                 type: DEFAULT_TYPE,
@@ -130,7 +118,7 @@ def request_body(prompt:, options: {})
                 },
                 payload: {
                   content: prompt
-                }.merge(payload_params(options)).merge(model_params(options))
+                }.merge(payload_params(options)).merge(model_params(options, unit_primitive))
               }],
               stream: true
             }
@@ -142,18 +130,17 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
               inputs: inputs
             }
 
-            feature_setting = chat_feature_setting(unit_primitive: unit_primitive)
+            feature_setting = namespace_feature_setting(unit_primitive) ||
+              chat_feature_setting(unit_primitive: unit_primitive)
 
-            model_metadata_params =
-              ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: feature_setting).to_params
-            params[:model_metadata] = model_metadata_params if model_metadata_params.present?
+            model_metadata = model_metadata(feature_setting)
+            params[:model_metadata] = model_metadata if model_metadata.present?
 
-            model_family = model_metadata_params && model_metadata_params[:name]
+            model_family = model_metadata && model_metadata[:name]
             default_version = ::Gitlab::Llm::PromptVersions.version_for_prompt("chat/#{unit_primitive}", model_family)
 
-            apply_model_selection_metadata(params, unit_primitive) if top_namespace_model_selected?(unit_primitive)
-
-            params[:prompt_version] = if feature_setting&.self_hosted? || ::Ai::AmazonQ.connected?
+            is_self_hosted = feature_setting.is_a?(::Ai::FeatureSetting) && feature_setting.self_hosted?
+            params[:prompt_version] = if is_self_hosted || ::Ai::AmazonQ.connected?
                                         default_version
                                       else
                                         prompt_version || default_version
@@ -162,35 +149,47 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
             params
           end
 
-          def model_params(options)
-            unit_primitive = options[:unit_primitive]
+          def model_metadata(feature_setting)
+            ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: feature_setting).to_params
+          end
 
-            if chat_feature_setting&.self_hosted?
-              self_hosted_model = chat_feature_setting.self_hosted_model
+          def namespace_feature_setting(unit_primitive)
+            feature = unit_primitive ? "duo_chat_#{unit_primitive}" : "duo_chat"
+            ::Ai::ModelSelection::NamespaceFeatureSetting.find_by_feature(root_namespace, feature)
+          end
+
+          def model_params(options, unit_primitive = nil)
+            unit_primitive ||= options[:unit_primitive]
+            feature_setting = namespace_feature_setting(unit_primitive) ||
+              chat_feature_setting(unit_primitive: unit_primitive)
 
-              {
+            # Handle self-hosted model settings
+            if feature_setting.is_a?(::Ai::FeatureSetting) && feature_setting.self_hosted?
+              self_hosted_model = feature_setting.self_hosted_model
+              return {
                 provider: :litellm,
                 model: self_hosted_model.model,
                 model_endpoint: self_hosted_model.endpoint,
                 model_api_key: self_hosted_model.api_token,
                 model_identifier: self_hosted_model.identifier
               }
-            elsif root_namespace&.root? && Feature.enabled?(:ai_model_switching, root_namespace)
-              feature = top_namespace_model_feature(unit_primitive)
-              setting = ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(root_namespace,
-                feature)
+            end
 
-              ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: setting).to_params
-            else
-              {
-                provider: provider(options),
-                model: model(options)
+            # Handle namespace feature settings
+            if feature_setting.is_a?(::Ai::ModelSelection::NamespaceFeatureSetting)
+              return {
+                provider: "gitlab",
+                feature_setting: feature_setting.feature,
+                identifier: feature_setting.offered_model_ref
               }
+
             end
-          end
 
-          def top_namespace_model_feature(unit_primitive)
-            unit_primitive ? :"duo_chat_#{unit_primitive}" : "duo_chat"
+            # Default model parameters
+            {
+              provider: provider(options),
+              model: model(options)
+            }
           end
 
           def payload_params(options)
@@ -226,24 +225,6 @@ def chat_feature_setting(unit_primitive: nil)
             feature_setting
           end
 
-          def top_namespace_model_selected?(unit_primitive)
-            top_namespace_model_ref(unit_primitive).present?
-          end
-
-          def apply_model_selection_metadata(params, unit_primitive)
-            feature = top_namespace_model_feature(unit_primitive)
-
-            return unless root_namespace&.root? && Feature.enabled?(:ai_model_switching, root_namespace)
-
-            setting = ::Ai::ModelSelection::NamespaceFeatureSetting.find_or_initialize_by_feature(
-              root_namespace, feature
-            )
-
-            model_metadata = ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: setting).to_params
-
-            params[:model_metadata] = model_metadata if model_metadata
-          end
-
           def processed_service_name(service_name)
             return service_name unless service_name == :duo_chat
             return service_name unless chat_feature_setting&.self_hosted?
diff --git a/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb b/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb
index fcb2496039baed..ebdcda05119417 100644
--- a/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb
+++ b/ee/spec/models/ai/model_selection/namespace_feature_setting_spec.rb
@@ -38,14 +38,12 @@
     it 'returns existing setting when one exists for the feature' do
       ai_feature_setting.save!
       result = described_class.find_or_initialize_by_feature(group, existing_feature)
-
       expect(result).to eq(ai_feature_setting)
     end
 
     it 'initializes a new setting when none exists for the feature' do
       new_feature = :code_completions
       result = described_class.find_or_initialize_by_feature(group, new_feature_enum)
-
       expect(result).to be_a(described_class)
       expect(result).to be_new_record
       expect(result.namespace).to eq(group)
@@ -64,6 +62,60 @@
     end
   end
 
+  describe '.find_by_feature' do
+    let(:feature_name) { "duo_chat" }
+    let(:offered_model_ref) { "claude-3-7-sonnet-20250219" }
+
+    subject(:ai_feature_setting) do
+      create(:ai_namespace_feature_setting,
+        namespace: group,
+        feature: feature_name,
+        offered_model_ref: offered_model_ref)
+    end
+
+    before do
+      ai_feature_setting
+    end
+
+    context 'when namespace is nil' do
+      it 'returns nil' do
+        result = described_class.find_by_feature(nil, feature_name)
+        expect(result).to be_nil
+      end
+    end
+
+    context 'when namespace is not a root namespace' do
+      let(:subgroup) { create(:group, parent: group) }
+
+      it 'returns nil' do
+        result = described_class.find_by_feature(subgroup, feature_name)
+        expect(result).to be_nil
+      end
+    end
+
+    it 'returns existing setting when one exists for the feature' do
+      result = described_class.find_by_feature(group, feature_name)
+      expect(result).to eq(ai_feature_setting)
+    end
+
+    it 'returns nil when no setting exists for the feature' do
+      non_existent_feature = 'code_generations'
+      result = described_class.find_by_feature(group, non_existent_feature)
+      expect(result).to be_nil
+    end
+
+    context 'when the feature is not enabled' do
+      let(:ff_enabled) { false }
+
+      subject(:ai_feature_setting) { build(:ai_namespace_feature_setting) }
+
+      it 'returns nil' do
+        result = described_class.find_by_feature(group, feature_name)
+        expect(result).to be_nil
+      end
+    end
+  end
+
   it_behaves_like 'model selection feature setting', scope_class_name: 'Group'
 
   describe 'validations' do
-- 
GitLab


From d11ec9f1f3a676bb6bbaa53b4303d953c75f740d Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Wed, 21 May 2025 11:12:07 +0200
Subject: [PATCH 10/11] Remove perform_inline in CompletionWorker left over

---
 ee/app/workers/llm/completion_worker.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ee/app/workers/llm/completion_worker.rb b/ee/app/workers/llm/completion_worker.rb
index 1397000cb082e3..0cde546cf402c8 100644
--- a/ee/app/workers/llm/completion_worker.rb
+++ b/ee/app/workers/llm/completion_worker.rb
@@ -45,7 +45,7 @@ def perform_for(message, options = {})
         # set SESSION_ID_HASH_KEY to ensure inside Sidekiq `Gitlab::Session.current` is not nil
         with_ip_address_state.set(
           Gitlab::SidekiqMiddleware::SetSession::Server::SESSION_ID_HASH_KEY => ::Gitlab::Session.session_id_for_worker
-        ).perform_inline(serialize_message(message), options)
+        ).perform_async(serialize_message(message), options)
       end
 
       def resource(message_hash)
-- 
GitLab


From 26f97b1b2defcb0b5641c0a67fa00d011bd7c348 Mon Sep 17 00:00:00 2001
From: Mohamed Hamda <mhamda@gitlab.com>
Date: Thu, 22 May 2025 10:38:52 +0200
Subject: [PATCH 11/11] Use provider from feature level

---
 ee/lib/gitlab/llm/ai_gateway/model_metadata.rb | 2 +-
 ee/lib/gitlab/llm/chain/requests/ai_gateway.rb | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb b/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb
index 5219299e1e77de..580956e1843227 100644
--- a/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/model_metadata.rb
@@ -36,7 +36,7 @@ def self_hosted_params
 
         def namespace_settings_params
           {
-            provider: 'gitlab',
+            provider: feature_setting.provider,
             identifier: feature_setting.offered_model_ref,
             feature_setting: feature_setting.feature
           }
diff --git a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
index e47b0e16d47a74..49eb4cce76ccde 100644
--- a/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
+++ b/ee/lib/gitlab/llm/chain/requests/ai_gateway.rb
@@ -178,7 +178,7 @@ def model_params(options, unit_primitive = nil)
             # Handle namespace feature settings
             if feature_setting.is_a?(::Ai::ModelSelection::NamespaceFeatureSetting)
               return {
-                provider: "gitlab",
+                provider: feature_setting.provider,
                 feature_setting: feature_setting.feature,
                 identifier: feature_setting.offered_model_ref
               }
-- 
GitLab