diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 68465d06b15..5643e2df17e 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 + raw_usage: Final = getattr(chunk, "usage", None) + if isinstance(raw_usage, (Usage, BaseModel, dict)): + model_response.usage = _normalize_usage(raw_usage) ## RETURN ARG result: Final = self.return_processed_chunk_logic( @@ -2243,6 +2244,14 @@ def _coerce_token_details( return details_type(**(raw if isinstance(raw, dict) else raw.model_dump())) +def _normalize_usage(usage: Usage | BaseModel | dict[str, object]) -> Usage: + 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 ( @@ -2259,7 +2268,10 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: for chunk in chunks: if "usage" in chunk and chunk["usage"] is not None: - usage = chunk["usage"] + raw_usage = chunk["usage"] + if not isinstance(raw_usage, (Usage, BaseModel, dict)): + continue + usage = _normalize_usage(raw_usage) latest_usage_chunk = usage if "prompt_tokens" in usage: prompt_tokens = usage.get("prompt_tokens", 0) or 0 @@ -2286,11 +2298,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage: ) if latest_usage_chunk is not None: - latest_cost: Final = ( - latest_usage_chunk.get("cost") - if isinstance(latest_usage_chunk, dict) - else getattr(latest_usage_chunk, "cost", None) - ) + latest_cost: Final = getattr(latest_usage_chunk, "cost", None) if latest_cost is not None: returned_usage_chunk.cost = latest_cost 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..125bcfb83d1 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -1428,6 +1428,85 @@ def test_calculate_total_usage_with_cost(): assert usage.completion_tokens == 5 +def test_openai_choice_usage_chunk_normalizes_sdk_usage( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + from openai.types.chat import ChatCompletionChunk + from openai.types.chat.chat_completion_chunk import Choice, ChoiceDelta + from openai.types.completion_usage import CompletionUsage, PromptTokensDetails + + initialized_custom_stream_wrapper.model = "openai/glm-5.2" + initialized_custom_stream_wrapper.custom_llm_provider = "openai" + initialized_custom_stream_wrapper.stream_options = {"include_usage": True} + chunk = ChatCompletionChunk( + id="chatcmpl-cached-usage", + choices=[ + Choice( + delta=ChoiceDelta(content="", role="assistant"), + finish_reason="stop", + index=0, + logprobs=None, + ) + ], + created=1745513206, + model="glm-5.2", + object="chat.completion.chunk", + usage=CompletionUsage( + completion_tokens=100, + prompt_tokens=1000, + total_tokens=1100, + prompt_tokens_details=PromptTokensDetails( + audio_tokens=0, + cached_tokens=600, + ), + ), + ) + + processed_chunk = initialized_custom_stream_wrapper.chunk_creator(chunk=chunk) + + assert processed_chunk is not None + assert isinstance(processed_chunk.usage, Usage) + assembled_response = litellm.stream_chunk_builder( + chunks=[processed_chunk], + messages=[{"role": "user", "content": "cached prompt"}], + ) + assert assembled_response is not None + assert assembled_response.usage.prompt_tokens == 1000 + assert assembled_response.usage.completion_tokens == 100 + assert assembled_response.usage.prompt_tokens_details is not None + assert assembled_response.usage.prompt_tokens_details.cached_tokens == 600 + + +def test_calculate_total_usage_preserves_token_details(): + from openai.types.completion_usage import CompletionUsage, PromptTokensDetails + + from litellm.litellm_core_utils.streaming_handler import calculate_total_usage + + raw_usage = CompletionUsage( + completion_tokens=100, + prompt_tokens=1000, + total_tokens=1100, + prompt_tokens_details=PromptTokensDetails( + audio_tokens=0, + cached_tokens=600, + ), + ) + chunk = ModelResponseStream( + id="chatcmpl-cached-usage", + created=1745513206, + model="openai/glm-5.2", + choices=[], + ) + chunk.usage = raw_usage + + usage = calculate_total_usage([chunk]) + + assert usage.prompt_tokens == 1000 + assert usage.completion_tokens == 100 + assert usage.prompt_tokens_details is not None + assert usage.prompt_tokens_details.cached_tokens == 600 + + def test_calculate_total_usage_with_dict_usage_cost(): """Regression: dict-shaped `usage` with a `cost` key must still surface provider cost even though `hasattr` on a dict does not consult its keys."""