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>

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(

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 (

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