Skip to content

feat(agents): use litellm in generic agents

Igor Drozdov requested to merge id-use-litellm-in-generic-agents into main

What does this merge request do and why?

Use litellm in generic agents. It'll allow us to easily extend the number of supported providers via configuration rather than code changes

Related issue: Prompt Migration: Prepare AI Gateway to be exte... (#524 - closed)

How to set up and validate locally

Use Use AIGW endpoint for generate description (gitlab-org/gitlab!152429) MR and verify that the request is sent to Anthropic correctly.

Notes:

  • Had to introduce ai_gateway/agents/generate_issue_description in AI Gateway project
  • Had to extend the list of headers here with 'X-Gitlab-Unit-Primitive' => unit_primitive
GitLab Rails Diff
diff --git a/ee/lib/gitlab/llm/ai_gateway/client.rb b/ee/lib/gitlab/llm/ai_gateway/client.rb
index bea7e0be556a..f4901e105729 100644
--- a/ee/lib/gitlab/llm/ai_gateway/client.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/client.rb
@@ -12,8 +12,9 @@ class Client

         ConnectionError = Class.new(StandardError)

-        def initialize(user, service_name:, tracking_context: {})
+        def initialize(user, service_name:, tracking_context: {}, unit_primitive: nil)
           @user = user
+          @unit_primitive = unit_primitive
           @access_token = ::CloudConnector::AvailableServices.find_by_name(service_name).access_token(user)
           @tracking_context = tracking_context
           @logger = Gitlab::Llm::Logger.build
@@ -58,7 +59,7 @@ def stream(endpoint:, body:, timeout: DEFAULT_TIMEOUT)

         private

-        attr_reader :user, :access_token, :logger, :tracking_context
+        attr_reader :user, :access_token, :logger, :tracking_context, :unit_primitive

         def perform_completion_request(endpoint:, body:, timeout:, stream:)
           logger.info_or_debug(user, message: "Performing request to AI Gateway", body: body, timeout: timeout,
@@ -85,6 +86,7 @@ def request_headers
             'X-Gitlab-Authentication-Type' => 'oidc',
             'Authorization' => "Bearer #{access_token}",
             'Content-Type' => 'application/json',
+            'X-Gitlab-Unit-Primitive' => unit_primitive,
             'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id
           }.merge(cloud_connector_headers(user))
         end
diff --git a/ee/lib/gitlab/llm/ai_gateway/completions/base.rb b/ee/lib/gitlab/llm/ai_gateway/completions/base.rb
index 14daf48bfd2c..26f189f45aeb 100644
--- a/ee/lib/gitlab/llm/ai_gateway/completions/base.rb
+++ b/ee/lib/gitlab/llm/ai_gateway/completions/base.rb
@@ -36,6 +36,7 @@ def unit_primitive

           def request!
             ai_client = ::Gitlab::Llm::AiGateway::Client.new(user, service_name: service_name,
+              unit_primitive: unit_primitive.to_s,
               tracking_context: tracking_context)
             ai_client.complete(
               endpoint: '/v1/agents/invoke',
Edited by Igor Drozdov

Merge request reports