From 7853a0c60744958ab99580a03d9af27f27fd451e Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 7 Aug 2026 08:23:57 +0000 Subject: [PATCH] fix(streaming): keep upstream usage when the usage chunk also carries choices Normalize provider SDK usage objects to litellm's Usage before they reach the stream aggregators, which rely on dict-style membership checks that only litellm's Usage supports. Fixes cached_tokens and provider prompt_tokens being dropped for OpenAI-compatible providers that send usage on a chunk with a non-empty choices array Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 4 ++ .../litellm_core_utils/streaming_handler.py | 18 +++++- .../test_streaming_handler.py | 56 +++++++++++++++++++ 3 files changed, 76 insertions(+), 2 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 3f967e29002..68d1dacf72e 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -5,6 +5,8 @@ from itertools import groupby from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Union, cast +from pydantic import BaseModel + from litellm._logging import verbose_logger from litellm.types.llms.openai import ( ChatCompletionAssistantContentValue, @@ -636,6 +638,8 @@ class ChunkProcessor: if isinstance(usage_chunk, dict): return Usage(**usage_chunk) + if isinstance(usage_chunk, BaseModel) and not isinstance(usage_chunk, Usage): + return Usage(**usage_chunk.model_dump()) return usage_chunk def _calculate_usage_per_chunk( diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 68465d06b15..13f02e943a0 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1450,8 +1450,9 @@ class CustomStreamWrapper: self.tool_call = True - if hasattr(chunk, "usage") and chunk.usage is not None: - model_response.usage = chunk.usage + chunk_usage: Final = getattr(chunk, "usage", None) + if chunk_usage is not None: + model_response.usage = _normalize_usage(chunk_usage) ## RETURN ARG result: Final = self.return_processed_chunk_logic( @@ -2243,6 +2244,19 @@ def _coerce_token_details( return details_type(**(raw if isinstance(raw, dict) else raw.model_dump())) +def _normalize_usage(usage: Usage | BaseModel | dict[str, Any]) -> Usage: + """ + Upstream usage can arrive as a provider SDK model (e.g. openai's CompletionUsage) or a raw + dict. Only litellm's Usage supports the `key in usage` membership checks the stream + aggregators rely on, so anything else silently reads as empty and the usage is dropped. + """ + if isinstance(usage, Usage): + return usage + if isinstance(usage, BaseModel): + return Usage(**usage.model_dump()) + return Usage(**usage) + + def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: """Assume most recent usage chunk has total usage uptil then.""" from litellm.litellm_core_utils.streaming_chunk_builder_utils import ( diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 5806b37539c..35579897857 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1604,6 +1604,62 @@ async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Loggin assert usage_chunks[-1].usage.cost == 0.00025 +@pytest.mark.asyncio +async def test_openai_compatible_usage_on_chunk_with_non_empty_choices( + logging_obj: Logging, +): + """Some OpenAI-compatible providers put the final usage on a chunk that still carries a + content-free `choices` entry. That usage arrives as openai's CompletionUsage, which does + not support the `key in usage` membership checks the aggregators use, so it used to read + as empty and prompt_tokens fell back to a local tokenizer estimate with cached_tokens + lost.""" + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk + from openai.types.completion_usage import CompletionUsage, PromptTokensDetails + + from litellm.utils import ModelResponseListIterator + + def _chunk(choices: list[dict], usage: CompletionUsage | None) -> ChatCompletionChunk: + return ChatCompletionChunk( + id="chatcmpl-x", + created=1742056047, + model="kimi", + object="chat.completion.chunk", + choices=choices, + usage=usage, + ) + + upstream_usage = CompletionUsage( + completion_tokens=32, + prompt_tokens=39779, + total_tokens=39811, + prompt_tokens_details=PromptTokensDetails(cached_tokens=39424), + ) + completion_stream = ModelResponseListIterator( + model_responses=[ + _chunk([{"index": 0, "delta": {"role": "assistant", "content": "OK"}}], None), + _chunk([{"index": 0, "delta": {}, "finish_reason": "stop"}], None), + _chunk([{"index": 0, "delta": {}}], upstream_usage), + ] + ) + response = CustomStreamWrapper( + completion_stream=completion_stream, + model="kimi", + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + + collected_chunks = [chunk async for chunk in response] + + usage = litellm.stream_chunk_builder( + collected_chunks, messages=[{"role": "user", "content": "ciao"}] + ).usage + assert usage.prompt_tokens == 39779 + assert usage.completion_tokens == 32 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 39424 + + def test_openrouter_streaming_cost_propagates_to_hidden_params(): """ Verify that provider-reported cost from usage.cost flows into