From 901b08f524966e370a2779357c11b22bfe5b8a92 Mon Sep 17 00:00:00 2001 From: Mohammad Javad Naderi Date: Sun, 16 Aug 2026 14:18:54 +0330 Subject: [PATCH 1/3] fix(streaming): preserve usage details in usage-only chunks --- .../litellm_core_utils/streaming_handler.py | 6 +-- .../test_streaming_handler.py | 43 +++++++++++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 99b1c1a2ab7..07f979a063f 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1451,11 +1451,7 @@ class CustomStreamWrapper: setattr( model_response, "usage", - litellm.Usage( - prompt_tokens=response_obj["usage"].get("prompt_tokens", None) or None, - completion_tokens=response_obj["usage"].get("completion_tokens", None) or None, - total_tokens=response_obj["usage"].get("total_tokens", None) or None, - ), + litellm.Usage(**response_obj["usage"]), ) elif isinstance(response_obj["usage"], Usage): setattr( 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 101935cac0a..b7091bec3c8 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -2,6 +2,7 @@ import json import os import sys import time +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -618,6 +619,48 @@ def test_streaming_handler_with_stop_chunk( assert returned_chunk is None +def test_usage_only_openai_chunk_preserves_usage_details( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + initialized_custom_stream_wrapper.custom_llm_provider = "openai" + initialized_custom_stream_wrapper.stream_options = {"include_usage": True} + chunk = SimpleNamespace( + id="chatcmpl-test", + model="glm-4.7", + choices=[], + usage={ + "completion_tokens": 28, + "prompt_tokens": 3721, + "total_tokens": 3749, + "completion_tokens_details": { + "accepted_prediction_tokens": None, + "audio_tokens": 0, + "reasoning_tokens": 0, + "rejected_prediction_tokens": None, + "image_tokens": 0, + }, + "prompt_tokens_details": { + "audio_tokens": 0, + "cache_write_tokens": 0, + "cached_tokens": 3237, + "video_tokens": 0, + }, + "cost": 0.0003777566, + "is_byok": False, + "cost_details": { + "upstream_inference_cost": 0.0003777566, + "upstream_inference_prompt_cost": 0.0003494206, + "upstream_inference_completions_cost": 0.000028336, + }, + }, + ) + + response = initialized_custom_stream_wrapper.chunk_creator(chunk) + + assert response.usage.prompt_tokens_details.cached_tokens == 3237 + assert response.usage.completion_tokens_details.reasoning_tokens == 0 + + def test_finish_reason_chunk_preserves_non_openai_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 58d7a321fa5520f6bb2501052b8002e6ce9f34a4 Mon Sep 17 00:00:00 2001 From: Mohammad Javad Naderi Date: Sun, 16 Aug 2026 16:36:12 +0330 Subject: [PATCH 2/3] fix(streaming): preserve OpenAI SDK usage in final chunks --- .../litellm_core_utils/streaming_handler.py | 6 ++- .../test_streaming_handler.py | 42 +++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 07f979a063f..861dccd0099 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1570,7 +1570,11 @@ class CustomStreamWrapper: self.tool_call = True if hasattr(chunk, "usage") and chunk.usage is not None: - model_response.usage = chunk.usage + model_response.usage = ( + litellm.Usage(**chunk.usage.model_dump()) + if isinstance(chunk.usage, BaseModel) + else chunk.usage + ) ## RETURN ARG result: Final = self.return_processed_chunk_logic( 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 b7091bec3c8..a09feddf6f0 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -661,6 +661,48 @@ def test_usage_only_openai_chunk_preserves_usage_details( assert response.usage.completion_tokens_details.reasoning_tokens == 0 +def test_openai_sdk_usage_is_preserved_in_stream_assembly(): + from openai.types.chat.chat_completion_chunk import ChatCompletionChunk + + stream = CustomStreamWrapper( + completion_stream=None, + model="z-ai/glm-5.2", + logging_obj=SimpleNamespace( + model_call_details={"litellm_params": {}}, + messages=[{"role": "user", "content": "Reply with exactly: ok"}], + stream_options={"include_usage": True}, + ), + custom_llm_provider="openai", + ) + chunk = ChatCompletionChunk.model_validate( + { + "id": "gen-test", + "object": "chat.completion.chunk", + "created": 0, + "model": "z-ai/glm-5.2", + "choices": [{"index": 0, "delta": {}, "finish_reason": "length"}], + "usage": { + "prompt_tokens": 17, + "completion_tokens": 4, + "total_tokens": 21, + "prompt_tokens_details": {"cached_tokens": 15}, + "completion_tokens_details": {"reasoning_tokens": 4}, + }, + } + ) + + assembled = litellm.stream_chunk_builder( + chunks=[stream.chunk_creator(chunk)], messages=stream.messages + ) + + assert assembled is not None + assert assembled.usage.prompt_tokens == 17 + assert assembled.usage.completion_tokens == 4 + assert assembled.usage.total_tokens == 21 + assert assembled.usage.prompt_tokens_details.cached_tokens == 15 + assert assembled.usage.completion_tokens_details.reasoning_tokens == 4 + + def test_finish_reason_chunk_preserves_non_openai_attributes( initialized_custom_stream_wrapper: CustomStreamWrapper, ): From 678e89fc2ec9094481b1de9f7616ea1ca5aafa2f Mon Sep 17 00:00:00 2001 From: Mohammad Javad Naderi Date: Sun, 16 Aug 2026 16:40:47 +0330 Subject: [PATCH 3/3] style: format streaming usage normalization --- litellm/litellm_core_utils/streaming_handler.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 861dccd0099..959cdd27791 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1571,9 +1571,7 @@ class CustomStreamWrapper: if hasattr(chunk, "usage") and chunk.usage is not None: model_response.usage = ( - litellm.Usage(**chunk.usage.model_dump()) - if isinstance(chunk.usage, BaseModel) - else chunk.usage + litellm.Usage(**chunk.usage.model_dump()) if isinstance(chunk.usage, BaseModel) else chunk.usage ) ## RETURN ARG