From b67da11a5c34139c718f9de0de5c92d3ee8cadc8 Mon Sep 17 00:00:00 2001
From: Manoj M J <mmj@gitlab.com>
Date: Fri, 18 Jul 2025 17:15:41 +0200
Subject: [PATCH] Hybrid support for custom models

---
 ee/app/models/ai/feature_setting.rb           | 16 ++++++++++++-
 .../concerns/ai/feature_configurable.rb       |  5 ++++
 .../model_selection/features_configurable.rb  |  6 +++--
 ee/lib/api/code_suggestions.rb                |  2 +-
 ee/lib/code_suggestions/model_details/base.rb |  4 ++++
 .../model_details/code_completion.rb          |  3 +++
 .../model_switching/ai_gateway.rb             |  4 +++-
 .../concerns/gitlab_default_params.rb         | 24 +++++++++++++++++++
 ee/lib/code_suggestions/tasks/base.rb         |  2 +-
 .../code_suggestions/tasks/code_completion.rb |  2 +-
 ee/lib/gitlab/ai_gateway.rb                   |  7 +++---
 ee/lib/gitlab/duo/chat/react_executor.rb      |  5 +++-
 ee/lib/gitlab/duo/chat/step_executor.rb       |  9 ++++---
 .../llm/ai_gateway/code_suggestions_client.rb |  6 ++++-
 ee/lib/gitlab/llm/ai_gateway/docs_client.rb   |  8 ++++++-
 ...lection_feature_setting_shared_examples.rb |  2 +-
 16 files changed, 88 insertions(+), 17 deletions(-)
 create mode 100644 ee/lib/code_suggestions/prompts/code_completion/model_switching/concerns/gitlab_default_params.rb

diff --git a/ee/app/models/ai/feature_setting.rb b/ee/app/models/ai/feature_setting.rb
index 85a4bddf654eba..ae4d7095aced0f 100644
--- a/ee/app/models/ai/feature_setting.rb
+++ b/ee/app/models/ai/feature_setting.rb
@@ -113,7 +113,9 @@ def provider_title
     end
 
     def base_url
-      Gitlab::AiGateway.url if self_hosted?
+      return Gitlab::AiGateway.url if self_hosted?
+
+      Gitlab::AiGateway.cloud_connector_url if vendored?
     end
 
     def compatible_self_hosted_models
@@ -143,6 +145,7 @@ def ready_for_request?
     end
 
     def model_metadata_params
+      return params_as_if_gitlab_default_model if vendored?
       return unless ready_for_request?
 
       {
@@ -155,6 +158,7 @@ def model_metadata_params
     end
 
     def model_request_params
+      return params_as_if_gitlab_default_model if vendored?
       return unless ready_for_request?
 
       {
@@ -165,5 +169,15 @@ def model_request_params
         model_identifier: self_hosted_model.identifier
       }
     end
+
+    private
+
+    def params_as_if_gitlab_default_model
+      {
+        provider: MODEL_PROVIDER,
+        feature_setting: feature,
+        identifier: ''
+      }
+    end
   end
 end
diff --git a/ee/app/models/concerns/ai/feature_configurable.rb b/ee/app/models/concerns/ai/feature_configurable.rb
index e34b2bc4ea7c6d..79c2009a6946e3 100644
--- a/ee/app/models/concerns/ai/feature_configurable.rb
+++ b/ee/app/models/concerns/ai/feature_configurable.rb
@@ -6,6 +6,7 @@ module FeatureConfigurable
 
     FEATURE_METADATA_PATH = Rails.root.join('ee/lib/gitlab/ai/feature_settings/feature_metadata.yml')
     FEATURE_METADATA = YAML.load_file(FEATURE_METADATA_PATH)
+    MODEL_PROVIDER = "gitlab"
 
     FeatureMetadata = Struct.new(
       :title,
@@ -24,6 +25,10 @@ def disabled?
       raise NotImplementedError, '#disabled? method must be implemented'
     end
 
+    def vendored?
+      raise NotImplementedError, '#vendored? method must be implemented'
+    end
+
     def model_metadata_params
       raise NotImplementedError, '#model_metadata_params method must be implemented'
     end
diff --git a/ee/app/models/concerns/ai/model_selection/features_configurable.rb b/ee/app/models/concerns/ai/model_selection/features_configurable.rb
index a00f215e238054..8b67c743d25e87 100644
--- a/ee/app/models/concerns/ai/model_selection/features_configurable.rb
+++ b/ee/app/models/concerns/ai/model_selection/features_configurable.rb
@@ -6,8 +6,6 @@ module FeaturesConfigurable
       extend ActiveSupport::Concern
       include Ai::FeatureConfigurable
 
-      MODEL_PROVIDER = "gitlab"
-
       FEATURES = {
         code_generations: 0,
         code_completions: 1,
@@ -109,6 +107,10 @@ def disabled?
           false
         end
 
+        def vendored?
+          false
+        end
+
         def provider
           MODEL_PROVIDER
         end
diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb
index 176f41d705630e..a2928fafb13a40 100644
--- a/ee/lib/api/code_suggestions.rb
+++ b/ee/lib/api/code_suggestions.rb
@@ -233,7 +233,7 @@ def model_prompt_cache_enabled?(project_path)
           details_hash = completion_model_details.current_model
 
           access = {
-            base_url: ::Gitlab::AiGateway.url,
+            base_url: completion_model_details.base_url,
             # for development purposes we just return instance JWT, this should not be used in production
             # until we generate a short-term token for user
             # https://gitlab.com/gitlab-org/modelops/applied-ml/code-suggestions/ai-assist/-/issues/429
diff --git a/ee/lib/code_suggestions/model_details/base.rb b/ee/lib/code_suggestions/model_details/base.rb
index 39c4fc0f2e63d8..c225a6bbcc0f90 100644
--- a/ee/lib/code_suggestions/model_details/base.rb
+++ b/ee/lib/code_suggestions/model_details/base.rb
@@ -46,6 +46,10 @@ def self_hosted?
         !!feature_setting&.self_hosted?
       end
 
+      def vendored?
+        !!feature_setting&.vendored?
+      end
+
       def namespace_feature_setting?
         feature_setting.is_a?(::Ai::ModelSelection::NamespaceFeatureSetting)
       end
diff --git a/ee/lib/code_suggestions/model_details/code_completion.rb b/ee/lib/code_suggestions/model_details/code_completion.rb
index 51afa97e021149..0fc8c51d85dc21 100644
--- a/ee/lib/code_suggestions/model_details/code_completion.rb
+++ b/ee/lib/code_suggestions/model_details/code_completion.rb
@@ -3,6 +3,8 @@
 module CodeSuggestions
   module ModelDetails
     class CodeCompletion < Base
+      include CodeSuggestions::Prompts::CodeCompletion::ModelSwitching::Concerns::GitlabDefaultParams
+
       FEATURE_SETTING_NAME = 'code_completions'
 
       def initialize(current_user:, root_namespace: nil)
@@ -21,6 +23,7 @@ def initialize(current_user:, root_namespace: nil)
       def current_model
         # if self-hosted, the model details are provided by the client
         return {} if self_hosted?
+        return code_completion_params_as_if_gitlab_default if vendored?
 
         return vertex_codestral_2501_model_details if code_completion_opt_out_fireworks?
 
diff --git a/ee/lib/code_suggestions/prompts/code_completion/model_switching/ai_gateway.rb b/ee/lib/code_suggestions/prompts/code_completion/model_switching/ai_gateway.rb
index 261a05e4bd199a..798343384c3549 100644
--- a/ee/lib/code_suggestions/prompts/code_completion/model_switching/ai_gateway.rb
+++ b/ee/lib/code_suggestions/prompts/code_completion/model_switching/ai_gateway.rb
@@ -6,10 +6,10 @@ module CodeCompletion
       module ModelSwitching
         class AiGateway < CodeSuggestions::Prompts::Base
           include CodeSuggestions::Prompts::CodeCompletion::Anthropic::Concerns::Prompt
+          include CodeSuggestions::Prompts::CodeCompletion::ModelSwitching::Concerns::GitlabDefaultParams
           include Gitlab::Loggable
 
           GATEWAY_PROMPT_VERSION = 3
-          MODEL_PROVIDER = 'gitlab'
           GITLAB_PROVIDED_CLAUDE_SONNET_35_MODEL_NAME = 'claude_3_5_sonnet_20240620'
           GITLAB_PROVIDED_ANTHROPIC_MODELS_FOR_CODE_COMPLETION = [
             'claude_sonnet_3_7_20250219',
@@ -22,6 +22,8 @@ def initialize(params, current_user, feature_setting, user_group_with_claude_cod
           end
 
           def request_params
+            return code_completion_params_as_if_gitlab_default if feature_setting&.vendored?
+
             model_name = determine_model_name
 
             {
diff --git a/ee/lib/code_suggestions/prompts/code_completion/model_switching/concerns/gitlab_default_params.rb b/ee/lib/code_suggestions/prompts/code_completion/model_switching/concerns/gitlab_default_params.rb
new file mode 100644
index 00000000000000..dc2fb655bd7cec
--- /dev/null
+++ b/ee/lib/code_suggestions/prompts/code_completion/model_switching/concerns/gitlab_default_params.rb
@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+module CodeSuggestions
+  module Prompts
+    module CodeCompletion
+      module ModelSwitching
+        module Concerns
+          module GitlabDefaultParams
+            MODEL_PROVIDER = 'gitlab'
+
+            private
+
+            def code_completion_params_as_if_gitlab_default
+              {
+                model_name: '',
+                model_provider: MODEL_PROVIDER
+              }
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/ee/lib/code_suggestions/tasks/base.rb b/ee/lib/code_suggestions/tasks/base.rb
index 07c78a50561c98..7429a84dc38da7 100644
--- a/ee/lib/code_suggestions/tasks/base.rb
+++ b/ee/lib/code_suggestions/tasks/base.rb
@@ -6,7 +6,7 @@ class Base
       AI_GATEWAY_CONTENT_SIZE = 100_000
 
       delegate :base_url, :self_hosted?, :feature_setting, :feature_name, :feature_disabled?, :licensed_feature,
-        :namespace_feature_setting?, to: :model_details
+        :namespace_feature_setting?, :vendored?, to: :model_details
 
       def initialize(current_user:, params: {}, unsafe_passthrough_params: {}, client: nil)
         @params = params
diff --git a/ee/lib/code_suggestions/tasks/code_completion.rb b/ee/lib/code_suggestions/tasks/code_completion.rb
index 93e6a3e8abcc6d..39ac67290a1810 100644
--- a/ee/lib/code_suggestions/tasks/code_completion.rb
+++ b/ee/lib/code_suggestions/tasks/code_completion.rb
@@ -32,7 +32,7 @@ def prompt
           amazon_q_prompt
         elsif self_hosted?
           self_hosted_prompt
-        elsif use_model_switching?
+        elsif use_model_switching? || vendored?
           model_switching_ai_gateway_prompt
         else
           saas_prompt
diff --git a/ee/lib/gitlab/ai_gateway.rb b/ee/lib/gitlab/ai_gateway.rb
index d8e2a41501ee1d..0c25e7fc591d05 100644
--- a/ee/lib/gitlab/ai_gateway.rb
+++ b/ee/lib/gitlab/ai_gateway.rb
@@ -17,10 +17,11 @@ def self.cloud_connector_url
       "#{::CloudConnector::Config.base_url}/ai"
     end
 
-    def self.access_token_url
-      base_url = self_hosted_url || "#{::CloudConnector::Config.base_url}/auth"
+    def self.access_token_url(code_completion_feature_setting)
+      default_base_url = self_hosted_url || "#{::CloudConnector::Config.base_url}/auth"
+      return "#{default_base_url}/v1/code/user_access_token" unless code_completion_feature_setting&.vendored?
 
-      "#{base_url}/v1/code/user_access_token"
+      "#{::CloudConnector::Config.base_url}/auth/v1/code/user_access_token"
     end
 
     def self.self_hosted_url
diff --git a/ee/lib/gitlab/duo/chat/react_executor.rb b/ee/lib/gitlab/duo/chat/react_executor.rb
index f8f6c5d24f3c86..4d5720de31f6a4 100644
--- a/ee/lib/gitlab/duo/chat/react_executor.rb
+++ b/ee/lib/gitlab/duo/chat/react_executor.rb
@@ -226,7 +226,10 @@ def process_unknown(events)
         end
 
         def step_executor
-          @step_executor ||= Gitlab::Duo::Chat::StepExecutor.new(context.current_user)
+          @step_executor ||= Gitlab::Duo::Chat::StepExecutor.new(
+            context.current_user,
+            feature_setting: chat_feature_setting
+          )
         end
 
         def step_forward
diff --git a/ee/lib/gitlab/duo/chat/step_executor.rb b/ee/lib/gitlab/duo/chat/step_executor.rb
index 7cedd5acafb81a..a0703869acbac7 100644
--- a/ee/lib/gitlab/duo/chat/step_executor.rb
+++ b/ee/lib/gitlab/duo/chat/step_executor.rb
@@ -17,10 +17,11 @@ class StepExecutor
 
         attr_reader :agent_steps
 
-        def initialize(user)
+        def initialize(user, feature_setting: nil)
           @user = user
           @agent_steps = []
           @event_parser = AgentEventParser.new
+          @feature_setting = feature_setting
         end
 
         def step(params)
@@ -64,7 +65,7 @@ def update_observation(observation)
 
         private
 
-        attr_reader :user, :event_parser
+        attr_reader :user, :event_parser, :feature_setting
 
         def perform_agent_request(params)
           log_conditional_info(user, message: "Request to v2/chat/agent",
@@ -79,8 +80,10 @@ def perform_agent_request(params)
           # Ref: https://github.com/ruby/net-protocol/blob/master/lib/net/protocol.rb#L214
           buffer = ""
 
+          url = feature_setting&.base_url || Gitlab::AiGateway.url
+
           response = Gitlab::HTTP.post(
-            "#{Gitlab::AiGateway.url}#{CHAT_V2_ENDPOINT}",
+            "#{url}#{CHAT_V2_ENDPOINT}",
             headers: Gitlab::AiGateway.headers(user: user, service: service),
             body: params.to_json,
             timeout: DEFAULT_TIMEOUT,
diff --git a/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb b/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
index 8626ca48c0e5e2..08d07b7784b0fa 100644
--- a/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/code_suggestions_client.rb
@@ -78,7 +78,7 @@ def direct_access_token
           )
 
           response = Gitlab::HTTP.post(
-            Gitlab::AiGateway.access_token_url,
+            Gitlab::AiGateway.access_token_url(code_completion_feature_setting),
             headers: Gitlab::AiGateway.headers(user: user, service: service),
             body: nil,
             timeout: DEFAULT_TIMEOUT,
@@ -109,6 +109,10 @@ def direct_access_token
 
         attr_reader :user
 
+        def code_completion_feature_setting
+          ::Ai::FeatureSetting.find_by_feature(:code_completions)
+        end
+
         def error(message)
           {
             message: message,
diff --git a/ee/lib/gitlab/llm/ai_gateway/docs_client.rb b/ee/lib/gitlab/llm/ai_gateway/docs_client.rb
index 970540ef11c881..e16d88b87a1e21 100644
--- a/ee/lib/gitlab/llm/ai_gateway/docs_client.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/docs_client.rb
@@ -33,8 +33,10 @@ def perform_search_request(query:, options:)
             options: options)
           timeout = options.delete(:timeout) || DEFAULT_TIMEOUT
 
+          url = feature_setting&.base_url || Gitlab::AiGateway.url
+
           response = Gitlab::HTTP.post(
-            "#{Gitlab::AiGateway.url}/v1/search/gitlab-docs",
+            "#{url}/v1/search/gitlab-docs",
             headers: Gitlab::AiGateway.headers(user: user, service: service),
             body: request_body(query: query).to_json,
             timeout: timeout,
@@ -50,6 +52,10 @@ def perform_search_request(query:, options:)
           response
         end
 
+        def feature_setting
+          ::Ai::FeatureSetting.find_by_feature(:duo_chat)
+        end
+
         def service
           ::CloudConnector::AvailableServices.find_by_name(:duo_chat)
         end
diff --git a/ee/spec/support/shared_examples/ai/model_selection_feature_setting_shared_examples.rb b/ee/spec/support/shared_examples/ai/model_selection_feature_setting_shared_examples.rb
index 2669fd1e76c77c..a836c214dd4518 100644
--- a/ee/spec/support/shared_examples/ai/model_selection_feature_setting_shared_examples.rb
+++ b/ee/spec/support/shared_examples/ai/model_selection_feature_setting_shared_examples.rb
@@ -123,7 +123,7 @@
 
     describe '#provider' do
       it 'returns the MODEL_PROVIDER' do
-        expect(ai_feature_setting.provider).to eq(described_class::MODEL_PROVIDER)
+        expect(ai_feature_setting.provider).to eq(::Ai::FeatureConfigurable::MODEL_PROVIDER)
       end
     end
 
-- 
GitLab