Log AI Gateway token usage per model in llm.log
What does this MR do and why?
The Received response from AI Gateway entry in llm.log reads the AI Gateway's token counts from usage.input_tokens and usage.output_tokens as flat keys, but /v1/prompts answers with a string and /v2/prompts nests the counts under each model's name, so those fields are never filled.
This MR logs one token_usage field instead, with one entry for each model:
"token_usage": [{ "model_name": "claude-sonnet-4-5", "input_tokens": 100, "output_tokens": 42 }]The keys are the same whichever models answer, as the Log Field Standards review rule asks of new fields, and match the names the AI Gateway logs for these counts. The rule also asks for new fields in Labkit::Fields, which has no token or model field yet. Adding one is a change to the LabKit spec that the field standardisation design describes, outside this repository.
A streamed /v2/prompts/chat answer arrives as newline-delimited frames, and several frames do not parse as one document. PromptResponseParser now keeps the usage the last frame reports, and Chain::Requests::AiGateway logs it. The client reads usage only from a response that is not streamed, so each answer's usage is logged once.
flowchart LR
A["POST /v2/prompts/:prompt_id"] --> B["Client#35;log_response_received"]
C["POST /v2/prompts/chat, streamed"] --> B
C --> D["PromptResponseParser#35;usage"]
D --> E["Chain::Requests::AiGateway#35;request"]
B -->|"not streamed"| F["token_usage_fields"]
E --> F
F --> G["llm.log response_received<br/>token_usage, one entry per model"]Danger's AI logging warning: both changed entries use log_conditional_info, which drops extra fields unless expanded AI logging is on, and token_usage holds only model names and token counts.
No changelog: only /v2/prompts responses carry usage, and ai_prompts_v2 is disabled by default (https://docs.gitlab.com/development/feature_flags/#changelog).
References
- Related to #499617
- !235406 (merged) added the two fields
- gitlab-org/modelops/applied-ml/code-suggestions/ai-assist!7341 (merged) fixes the placeholder model name, such as
gpt-3.5-turbo, that the AI Gateway reports for a self-hosted custom model called without streaming
How to set up and validate locally
- In a Rails console,
Feature.enable(:expanded_ai_logging)andFeature.enable(:ai_prompts_v2, User.find_by_username('root')). - In Duo Chat, with the Agentic toggle off,
/explaina code snippet. Inlog/llm.log,Made request to AI Clientcarriestoken_usageand that request'sReceived response from AI Gatewaydoes not. - Generate a commit message in a merge request. Its
Received response from AI Gatewayentry carriestoken_usage.
MR acceptance checklist
Evaluate this MR against the MR acceptance checklist.
An AI coding assistant wrote this change and ran the changed specs in a disposable development environment.