Skip to content

Handle structured output from new prompts endpoint

Alejandro Rodríguezrequested to merge
prompt-token-usage into master

What does this MR do and why?

As part of gitlab-org/modelops/applied-ml/code-suggestions/ai-assist#1174 (closed) we introduced a new version of the prompts endpoint that also includes token usage metadata. This MR introduces a feature flag that switched between the old and the new endpoint. When the user has the expanded_ai_logging feature flag enabled, a log entry including the usage will be generated.

References

Screenshots or screen recordings

Before After

How to set up and validate locally

  • Run bundle exec 'rake gitlab:duo:setup' to configure all the necessary pre-requisites and feature flags
  • Invoke an AI action that relies on the prompts endpoint. For example, create a new issue (http://gdk.test:3000/gitlab-duo/test/-/issues/new) and use the "Generate Issue Description" button
  • On log/llm.log, search for a log entry with message "Received response from AI Gateway". It should also include token usage. For example:
{"severity":"INFO","time":"2025-06-27T19:55:40.058Z","correlation_id":"01JYSFBEGWM01JCH9NHK9ZP24X","meta.caller_id":"Llm::CompletionWorker","meta.feature_category":"ai_abstraction_layer","meta.organization_id":1,"meta.remote_ip":"127.0.0.1","meta.http_router_rule_action":"classify","meta.http_router_rule_type":"SESSION_PREFIX","meta.user":"root","meta.user_id":1,"meta.client_id":"user/1","meta.root_caller_id":"GraphqlController#execute","response_from_llm":{"content":"Test","usage":{"claude-3-5-sonnet-20241022":{"input_tokens":59,"output_tokens":4}}},"message":"Received response from AI Gateway","class":"Gitlab::Llm::AiGateway::Client","ai_event_name":"response_received","ai_component":"abstraction_layer"}

MR acceptance checklist

Evaluate this MR against the MR acceptance checklist. It helps you analyze changes to reduce risks in quality, performance, reliability, security, and maintainability.

Merge request reports