Preserve provider usage in LLM results
Carry token usage and LiteLLM-reported response cost through Chat Completions and Responses, including terminal streaming usage chunks. Keep missing provider accounting absent instead of synthesizing values.
Alessandro committed
Aug 24, 2026 at 17:05 UTC
af56d211c09b93aaa40ddce6df761be95f7427f2
6 files changed
+105
-11
helpers/litellm_transport.py
+32
-9
@@ -214,10 +214,11 @@ class LiteLLMTransport:
214
while True:
215
try:
216
if self.policy.mode is TransportMode.CHAT_COMPLETIONS:
217
- parsed = ChatCompletionsTransport.parse(
218
- completion(**self._chat_request(stream=False))
217
+ raw_response = completion(**self._chat_request(stream=False))
218
+ parsed = ChatCompletionsTransport.parse(raw_response)
219
+ self.last_result = self._llm_result_from_chat(
220
+ parsed, raw_response
221
)
220
- self.last_result = self._llm_result_from_chat(parsed)
222
return parsed
223
request = self._responses_request(stream=False)
224
raw_response = responses(**request)
@@ -235,10 +236,13 @@ class LiteLLMTransport:
236
while True:
237
try:
238
if self.policy.mode is TransportMode.CHAT_COMPLETIONS:
238
- parsed = ChatCompletionsTransport.parse(
239
- await acompletion(**self._chat_request(stream=False))
239
+ raw_response = await acompletion(
240
+ **self._chat_request(stream=False)
241
+ )
242
+ parsed = ChatCompletionsTransport.parse(raw_response)
243
+ self.last_result = self._llm_result_from_chat(
244
+ parsed, raw_response
245
)
241
- self.last_result = self._llm_result_from_chat(parsed)
246
return parsed
247
request = self._responses_request(stream=False)
248
raw_response = await aresponses(**request)
@@ -400,10 +404,13 @@ class LiteLLMTransport:
404
**response_kwargs,
405
}
406
403
- def _llm_result_from_chat(self, parsed: ChatChunk) -> LLMResult:
407
+ def _llm_result_from_chat(
408
+ self, parsed: ChatChunk, response: Any = None
409
+ ) -> LLMResult:
410
return LLMResult.from_chat(
411
response=parsed["response_delta"],
412
reasoning=parsed["reasoning_delta"],
413
+ usage=_reported_usage(response),
414
input_items=ResponsesTransport.input_from_messages(self.messages),
415
output_items=parsed.get("_output_items"),
416
provider_model_key=self.model,
@@ -413,7 +420,7 @@ class LiteLLMTransport:
420
def _llm_result_from_response(
421
self, response: Any, request: dict[str, Any]
422
) -> LLMResult:
416
- return LLMResult.from_response(
423
+ result = LLMResult.from_response(
424
response,
425
input_items=_as_list(request.get("input")),
426
previous_response_id=str(request.get("previous_response_id") or ""),
@@ -422,6 +429,8 @@ class LiteLLMTransport:
429
state=self.last_request_state,
430
capability=self._capability_metadata(),
431
)
432
+ result.usage = _reported_usage(response)
433
+ return result
434
435
def _stream_result_from_parser(
436
self, parser: "ResponsesEventParser", request: dict[str, Any]
@@ -440,10 +449,11 @@ class LiteLLMTransport:
449
self, parser: "ChatCompletionsStreamParser"
450
) -> LLMResult | None:
451
output_items = parser.output_items()
443
- if not output_items:
452
+ if not output_items and not parser.usage:
453
return None
454
return LLMResult.from_chat(
455
response=parser.function_calls_text(),
456
+ usage=parser.usage,
457
input_items=ResponsesTransport.input_from_messages(self.messages),
458
output_items=output_items,
459
provider_model_key=self.model,
@@ -589,8 +599,11 @@ class ChatCompletionsStreamParser:
599
self.tool_calls: dict[str, dict[str, Any]] = {}
600
self.order: list[str] = []
601
self.emitted = False
602
+ self.usage: dict[str, Any] = {}
603
604
def parse(self, chunk: Any) -> ChatChunk:
605
+ if usage := _reported_usage(chunk):
606
+ self.usage.update(usage)
607
parsed = ChatCompletionsTransport.parse(chunk)
608
choice = _first_choice(chunk)
609
delta = _get_value(choice, "delta") or {}
@@ -1699,6 +1712,16 @@ def _object_to_dict(obj: Any) -> dict[str, Any]:
1712
return {}
1713
1714
1715
+def _reported_usage(response: Any) -> dict[str, Any]:
1716
+ usage = _object_to_dict(_get_value(response, "usage"))
1717
+ hidden = _object_to_dict(_get_value(response, "_hidden_params"))
1718
+ if usage.get("cost") is None:
1719
+ usage.pop("cost", None)
1720
+ if hidden.get("response_cost") is not None:
1721
+ usage["cost"] = hidden["response_cost"]
1722
+ return usage
1723
+
1724
+
1725
def _normalize_reasoning_effort(effort: Any) -> str | None:
1726
if isinstance(effort, str):
1727
normalized = effort.strip().lower()
helpers/litellm_transport.py.dox.md
+1
@@ -33,6 +33,7 @@
33
- Fall back to Chat Completions when a Responses endpoint fails before output with an endpoint-specific server error, proxy path-unavailable error, or LiteLLM proxy-extra import error.
34
- Fall back to Chat Completions when LiteLLM's Responses mock streaming path tries to JSON-decode a real SSE stream before any output.
35
- Preserve Chat Completions tool calls from both non-streaming responses and streaming deltas as canonical `LLMResult` function-call items.
36
+- Preserve provider usage and LiteLLM response cost for both transports only when the response or stream actually supplies them; do not synthesize unavailable provider accounting.
37
- Preserve Responses function calls collected from stream events when a terminal completed event omits them.
38
- Preserve provider-state metadata when Responses API calls succeed, and fall back to local replay when provider state is unsupported.
39
- Keep prompt-cache markers only for providers that accept them.
helpers/llm_result.py
+2
@@ -124,6 +124,7 @@ class LLMResult:
124
*,
125
response: str,
126
reasoning: str = "",
127
+ usage: dict[str, Any] | None = None,
128
input_items: list[dict[str, Any]] | None = None,
129
output_items: list[dict[str, Any]] | None = None,
130
provider_model_key: str = "",
@@ -160,6 +161,7 @@ class LLMResult:
161
provider_model_key=provider_model_key,
162
mode="chat_completions",
163
state="off",
164
+ usage=object_to_dict(usage or {}),
165
capability=dict(capability or {}),
166
)
167
if not result.response and result.function_calls:
helpers/llm_result.py.dox.md
+1
-1
@@ -19,7 +19,7 @@
19
20
- `LLMResult.metadata()` stores only durable provider state under `RESPONSE_METADATA_KEY`: response IDs, structured output items, provider/mode/state, usage, and capability data. Runtime prompt inputs, raw responses, and duplicated response/reasoning text are not persisted in history.
21
- `from_response(...)` must preserve provider `response_id`, `previous_response_id`, raw output items, usage, and capability metadata.
22
-- `from_chat(...)` must produce an equivalent chat-completions result with `mode="chat_completions"` and `state="off"`, preserving optional function-call output items when the chat transport supplies them.
22
+- `from_chat(...)` must produce an equivalent chat-completions result with `mode="chat_completions"` and `state="off"`, preserving optional function-call output items and provider usage when the chat transport supplies them.
23
- Function-call output items must preserve `call_id` and optional acknowledged safety checks.
24
- Argument parsing must tolerate JSON strings, dictionaries, and malformed values without throwing.
25
tests/test_responses_architecture.py
+37
-1
@@ -129,6 +129,37 @@ def test_history_migrates_legacy_ai_metadata_and_preserves_tool_inputs():
129
assert old.sequence == 0
130
131
132
+@pytest.mark.asyncio
133
+async def test_chat_completion_transport_preserves_reported_usage(monkeypatch):
134
+ async def fake_acompletion(**kwargs):
135
+ return {
136
+ "choices": [{"message": {"content": "done"}}],
137
+ "usage": {
138
+ "prompt_tokens": 120,
139
+ "completion_tokens": 8,
140
+ "total_tokens": 128,
141
+ },
142
+ "_hidden_params": {"response_cost": 0.0042},
143
+ }
144
+
145
+ monkeypatch.setattr(litellm_transport, "acompletion", fake_acompletion)
146
+ transport = litellm_transport.LiteLLMTransport(
147
+ model="custom/model",
148
+ messages=[{"role": "user", "content": "question"}],
149
+ kwargs={"a0_api_mode": "chat_completions"},
150
+ )
151
+
152
+ await transport.acomplete()
153
+
154
+ assert transport.last_result is not None
155
+ assert transport.last_result.usage == {
156
+ "prompt_tokens": 120,
157
+ "completion_tokens": 8,
158
+ "total_tokens": 128,
159
+ "cost": 0.0042,
160
+ }
161
+
162
+
163
def test_responses_provider_state_uses_previous_response_and_new_items():
164
new_items = [{"type": "function_call_output", "call_id": "call_1", "output": "done"}]
165
local_items = [{"role": "user", "content": "full replay"}]
@@ -362,6 +393,7 @@ async def test_unified_turn_waits_for_completed_native_responses_calls(monkeypat
393
"id": "resp_parallel",
394
"output": calls,
395
"usage": {"input_tokens": 10, "output_tokens": 5},
396
+ "_hidden_params": {"response_cost": 0.0012},
397
},
398
},
399
]
@@ -394,7 +426,11 @@ async def test_unified_turn_waits_for_completed_native_responses_calls(monkeypat
426
assert stream.closed is False
427
assert result.mode == "responses"
428
assert result.response_id == "resp_parallel"
397
- assert result.usage == {"input_tokens": 10, "output_tokens": 5}
429
+ assert result.usage == {
430
+ "input_tokens": 10,
431
+ "output_tokens": 5,
432
+ "cost": 0.0012,
433
+ }
434
assert [call.name for call in result.function_calls] == ["lookup", "summarize"]
435
assert json.loads(result.response) == {
436
"tool_name": "parallel_tool_calls",
tests/test_stream_tool_early_stop.py
+32
@@ -1527,6 +1527,38 @@ def test_chat_completions_stream_parser_reads_dumped_tool_calls():
1527
}
1528
1529
1530
+def test_chat_completions_stream_parser_preserves_optional_usage():
1531
+ parser = litellm_transport.ChatCompletionsStreamParser()
1532
+ parser.parse(
1533
+ {
1534
+ "choices": [],
1535
+ "usage": {"prompt_tokens": 240},
1536
+ "_hidden_params": {"response_cost": 0.0084},
1537
+ }
1538
+ )
1539
+ parser.parse(
1540
+ {
1541
+ "choices": [],
1542
+ "usage": {"completion_tokens": 16, "total_tokens": 256},
1543
+ }
1544
+ )
1545
+ transport = litellm_transport.LiteLLMTransport(
1546
+ model="custom/model",
1547
+ messages=[{"role": "user", "content": "question"}],
1548
+ kwargs={"a0_api_mode": "chat_completions"},
1549
+ )
1550
+
1551
+ result = transport._stream_result_from_chat_parser(parser)
1552
+
1553
+ assert result is not None
1554
+ assert result.usage == {
1555
+ "prompt_tokens": 240,
1556
+ "completion_tokens": 16,
1557
+ "total_tokens": 256,
1558
+ "cost": 0.0084,
1559
+ }
1560
+
1561
+
1562
@pytest.mark.asyncio
1563
async def test_unified_turn_preserves_chat_streaming_tool_calls(monkeypatch):
1564
async def fake_acompletion(*args, **kwargs):