feat(agents): use litellm in generic agents
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