Merge branch 'id-model-metadata-for-completions' into 'master'

Send model metadata for completion features

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



Merged-by: Eduardo Bonet <ebonet@gitlab.com>
Approved-by: Eduardo Bonet <ebonet@gitlab.com>
Reviewed-by: Igor Drozdov <idrozdov@gitlab.com>
Reviewed-by: Eduardo Bonet <ebonet@gitlab.com>
Co-authored-by: Igor Drozdov <idrozdov@gitlab.com>

ee/lib/gitlab/duo/chat/react_executor.rb

@@ -296,26 +296,7 @@ def unavailable_resources_params
         attr_reader :stream_response_handler
 
         def model_metadata_params
-          if chat_feature_setting&.self_hosted?
-
-            self_hosted_model = chat_feature_setting.self_hosted_model
-
-            return {
-              provider: :openai, # for self-hosted models we support Messages API format at the moment
-              name: self_hosted_model.model,
-              endpoint: self_hosted_model.endpoint,
-              api_key: self_hosted_model.api_token,
-              identifier: self_hosted_model.identifier
-            }
-          end
-
-          return unless ::Ai::AmazonQ.connected?
-
-          {
-            provider: :amazon_q,
-            name: :amazon_q,
-            role_arn: ::Ai::Setting.instance.amazon_q_role_arn
-          }
+          ::Gitlab::Llm::AiGateway::ModelMetadata.new(feature_setting: chat_feature_setting).to_params
         end
 
         def conversation

ee/lib/gitlab/llm/ai_gateway/completions/base.rb

@@ -44,16 +44,7 @@ def post_process(response)
           end
 
           def request!
-            ai_client = ::Gitlab::Llm::AiGateway::Client.new(user, service_name: service_name,
-              tracking_context: tracking_context)
-
-            request_body = { 'inputs' => inputs }
-            request_body['prompt_version'] = prompt_version unless prompt_version.nil?
-
-            response = ai_client.complete(
-              url: "#{::Gitlab::AiGateway.url}/v1/prompts/#{prompt_message.ai_action}",
-              body: request_body
-            )
+            response = perform_ai_gateway_request!
 
             return if response&.body.blank?
             return Gitlab::Json.parse(response.body) if response&.success?
@@ -67,6 +58,22 @@ def request!
             { 'detail' => DEFAULT_ERROR }
           end
 
+          def perform_ai_gateway_request!
+            ai_client = ::Gitlab::Llm::AiGateway::Client.new(user, service_name: service_name,
+              tracking_context: tracking_context)
+
+            request_body = { 'inputs' => inputs }
+            request_body['prompt_version'] = prompt_version unless prompt_version.nil?
+
+            model_metadata_params = ::Gitlab::Llm::AiGateway::ModelMetadata.new.to_params
+            request_body['model_metadata'] = model_metadata_params if model_metadata_params.present?
+
+            ai_client.complete(
+              url: "#{::Gitlab::AiGateway.url}/v1/prompts/#{prompt_message.ai_action}",
+              body: request_body
+            )
+          end
+
           def service_name
             prompt_message.ai_action.to_sym
           end

ee/lib/gitlab/llm/ai_gateway/completions/resolve_vulnerability.rb

@@ -53,12 +53,7 @@ def inputs
           private
 
           def request!
-            ai_client = ::Gitlab::Llm::AiGateway::Client.new(user, service_name: service_name,
-              tracking_context: tracking_context)
-            response = ai_client.complete(
-              url: "#{::Gitlab::AiGateway.url}/v1/prompts/#{prompt_message.ai_action}",
-              body: { 'inputs' => inputs }
-            )
+            response = perform_ai_gateway_request!
 
             raise EmptyResponseError, "Empty response from LLM" if response&.body.blank?
 

ee/lib/gitlab/llm/ai_gateway/model_metadata.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Llm
+    module AiGateway
+      class ModelMetadata
+        def initialize(feature_setting: nil)
+          @feature_setting = feature_setting
+        end
+
+        def to_params
+          return self_hosted_params if feature_setting&.self_hosted?
+
+          amazon_q_params if ::Ai::AmazonQ.connected?
+        end
+
+        private
+
+        attr_reader :feature_setting
+
+        def self_hosted_params
+          self_hosted_model = feature_setting.self_hosted_model
+
+          {
+            provider: self_hosted_model.provider,
+            name: self_hosted_model.model,
+            endpoint: self_hosted_model.endpoint,
+            api_key: self_hosted_model.api_token,
+            identifier: self_hosted_model.identifier
+          }
+        end
+
+        def amazon_q_params
+          {
+            provider: :amazon_q,
+            name: :amazon_q,
+            role_arn: ::Ai::Setting.instance.amazon_q_role_arn
+          }
+        end
+      end
+    end
+  end
+end

ee/lib/gitlab/llm/chain/requests/ai_gateway.rb

@@ -133,25 +133,9 @@ def request_body_agent(inputs:, unit_primitive: nil, prompt_version: nil)
 
             feature_setting = chat_feature_setting(unit_primitive: unit_primitive)
 
-            if feature_setting&.self_hosted?
-              self_hosted_model = feature_setting.self_hosted_model
-
-              params[:model_metadata] = {
-                provider: self_hosted_model.provider,
-                name: self_hosted_model.model,
-                endpoint: self_hosted_model.endpoint,
-                api_key: self_hosted_model.api_token,
-                identifier: self_hosted_model.identifier
-              }
-            end
-
-            if ::Ai::AmazonQ.connected?
-              params[:model_metadata] = {
-                provider: :amazon_q,
-                name: :amazon_q,
-                role_arn: ::Ai::Setting.instance.amazon_q_role_arn
-              }
-            end
+            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?
 
             params
           end

ee/spec/lib/gitlab/llm/ai_gateway/completions/base_spec.rb

@@ -97,6 +97,7 @@
       let(:processed_response) { { 'detail' => 'An unexpected error has occurred.' } }
 
       before do
+        allow(::Ai::AmazonQ).to receive(:connected?).and_return(false)
         allow(Gitlab::Json).to receive(:parse).and_raise(StandardError)
       end
 
@@ -143,5 +144,33 @@
         is_expected.to be(result)
       end
     end
+
+    context 'when model_metadata is provided' do
+      it 'includes model_metadata in the request body' do
+        params = {
+          provider: 'provider',
+          name: 'model',
+          endpoint: 'http://example.com',
+          identifier: 'identifier'
+        }
+
+        expect_next_instance_of(::Gitlab::Llm::AiGateway::ModelMetadata) do |instance|
+          expect(instance).to receive(:to_params).and_return(params)
+        end
+
+        expect(client).to receive(:complete).with(
+          url: "#{Gitlab::AiGateway.url}/v1/prompts/#{ai_action}",
+          body: { 'inputs' => inputs, 'model_metadata' => params }
+        ).and_return(http_response)
+
+        expect(response_modifier_class).to receive(:new).with(processed_response)
+          .and_return(response_modifier)
+        expect(Gitlab::Llm::GraphqlSubscriptionResponseService).to receive(:new)
+          .with(user, resource, response_modifier, options: response_options).and_return(response_service)
+        expect(response_service).to receive(:execute).and_return(result)
+
+        is_expected.to be(result)
+      end
+    end
   end
 end

ee/spec/lib/gitlab/llm/ai_gateway/model_metadata_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Llm::AiGateway::ModelMetadata, feature_category: :ai_abstraction_layer do
+  let_it_be(:self_hosted_model) { create(:ai_self_hosted_model) }
+  let_it_be(:feature_setting) { create(:ai_feature_setting, self_hosted_model: self_hosted_model) }
+
+  describe '#to_params' do
+    context 'when feature_setting is self-hosted' do
+      it 'returns self-hosted params' do
+        model_metadata = described_class.new(feature_setting: feature_setting)
+
+        expect(model_metadata.to_params).to eq({
+          provider: self_hosted_model.provider,
+          name: self_hosted_model.model,
+          endpoint: self_hosted_model.endpoint,
+          api_key: self_hosted_model.api_token,
+          identifier: self_hosted_model.identifier
+        })
+      end
+    end
+
+    context 'when feature_setting is not self-hosted and Ai::AmazonQ is connected' do
+      before do
+        allow(::Ai::AmazonQ).to receive(:connected?).and_return(true)
+        ::Ai::Setting.instance.update!(amazon_q_role_arn: "role-arn")
+      end
+
+      it 'returns amazon_q params' do
+        model_metadata = described_class.new
+
+        expect(model_metadata.to_params).to eq({
+          provider: :amazon_q,
+          name: :amazon_q,
+          role_arn: "role-arn"
+        })
+      end
+    end
+
+    context 'when feature_setting is not self-hosted and Ai::AmazonQ is not connected' do
+      let(:self_hosted) { false }
+
+      before do
+        allow(::Ai::AmazonQ).to receive(:connected?).and_return(false)
+      end
+
+      it 'returns nil' do
+        model_metadata = described_class.new
+
+        expect(model_metadata.to_params).to be_nil
+      end
+    end
+  end
+end