diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 1e0b778d244..d9641a6306b 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2336,6 +2336,9 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + if self.stream_options is None: + usage: Final = calculate_total_usage(chunks=self.chunks) + processed_chunk._hidden_params["usage"] = usage # pyright: ignore[reportPrivateUsage] # sync parity # see sync __next__'s sibling branch: deliberately do NOT restore # here - this chunk is still this call's own data, and restoring # before returning it would corrupt the caller's own log 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 7f54fbfb4c2..7caf401ce38 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4692,3 +4692,40 @@ async def test_async_stream_assembled_response_keeps_vertex_traffic_type(logging assembled = litellm.stream_chunk_builder(chunks=received, messages=[{"role": "user", "content": "hi"}]) assert assembled is not None assert assembled._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND_FLEX" + + +@pytest.mark.asyncio +async def test_async_fake_stream_final_chunk_carries_hidden_usage(logging_obj: Logging): + from litellm.llms.base_llm.base_model_iterator import MockResponseIterator + from litellm.types.utils import ModelResponse + + model_response = ModelResponse( + id="chatcmpl-fake-stream", + model="my-random-model", + choices=[ + { + "index": 0, + "message": {"role": "assistant", "content": "hello world"}, + "finish_reason": "stop", + } + ], + ) + model_response.usage = Usage(prompt_tokens=1234, completion_tokens=7, total_tokens=1241) + + wrapper = CustomStreamWrapper( + completion_stream=MockResponseIterator(model_response=model_response), + model="my-random-model", + custom_llm_provider="anthropic", + logging_obj=logging_obj, + ) + + final_chunk = None + async for chunk in wrapper: + final_chunk = chunk + + assert final_chunk is not None + hidden_usage = final_chunk._hidden_params.get("usage") + assert hidden_usage is not None + assert hidden_usage.prompt_tokens == 1234 + assert hidden_usage.completion_tokens == 7 + assert hidden_usage.total_tokens == 1241