diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f6340426c1b..ab614e0a3ed 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1481,11 +1481,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 b5e33a4e421..8d8e4ae3be5 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -6,7 +6,7 @@ import pytest import asyncio import traceback -from typing import Optional +from typing import Final, Optional import litellm from litellm import verbose_logger @@ -4460,3 +4460,102 @@ def test_handle_stream_fallback_error_restores_context_only_after_exception_mapp finally: trace_id_var.set("") session_id_var.set("") + + +def _final_chunk_with_dict_usage(choices) -> ModelResponseStream: + """A stream chunk whose ``usage`` is a plain dict. + + ``ModelResponseStream.__init__`` coerces a dict ``usage`` kwarg into + ``litellm.Usage``, so the dict is attached after construction — the same + shape a custom ``streaming_decoder`` (or any third-party iterator that + bypasses litellm chunk construction) delivers to ``chunk_creator``. + """ + chunk: Final = ModelResponseStream( + id="chatcmpl-dict-usage-test", + created=1754900000, + model="gpt-4o", + object="chat.completion.chunk", + choices=choices, + ) + chunk.usage = { # mutable-ok: wire-shaped usage payload under test + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": {"cached_tokens": 80}, # mutable-ok: wire payload + "completion_tokens_details": {"reasoning_tokens": 5}, # mutable-ok: wire payload + "cache_creation_input_tokens": 42, + } + return chunk + + +def test_dispatch_provider_chunk_dict_usage_preserves_token_details( + logging_obj: Logging, +): + """The openai else-branch of ``_dispatch_provider_chunk`` handles three + usage shapes: ``Usage`` (passed through), ``BaseModel`` + (``Usage(**model_dump())``), and plain ``dict``. The dict arm used to + rebuild ``litellm.Usage`` from only prompt/completion/total tokens, + silently dropping ``prompt_tokens_details`` (cached_tokens), + ``completion_tokens_details`` (reasoning_tokens), and cache token counts + that both sibling arms preserve. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage( + choices=[ + StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=None)) + ] # mutable-ok: chunk payload + ) + model_response: Final = wrapper.model_response_creator() + wrapper._dispatch_provider_chunk( + chunk=chunk, + model_response=model_response, + completion_obj={"content": ""}, # mutable-ok: dispatch scratch dict + ) + + usage: Final = model_response.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens == 100 + assert usage.completion_tokens == 10 + assert usage.total_tokens == 110 + assert usage.prompt_tokens_details is not None, ( + "dict-usage arm dropped prompt_tokens_details (its Usage/BaseModel sibling arms preserve it)" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5 + assert getattr(usage, "cache_creation_input_tokens", None) == 42 + + +def test_chunk_creator_usage_only_chunk_with_dict_usage_keeps_details( + logging_obj: Logging, +): + """End-to-end through ``chunk_creator``: a usage-only final chunk (empty + ``choices``, ``stream_options.include_usage`` set) takes the early-return + path, so whatever the dict-usage arm wrote is exactly what the consumer + receives — there is no later aggregation to repair it. + """ + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="gpt-4o", + logging_obj=logging_obj, + custom_llm_provider="openai", + stream_options={"include_usage": True}, # mutable-ok: constructor payload + ) + chunk: Final = _final_chunk_with_dict_usage(choices=[]) # mutable-ok: empty choices list + result: Final = wrapper.chunk_creator(chunk=chunk) + + assert result is not None + usage: Final = result.usage + assert isinstance(usage, Usage) + assert usage.prompt_tokens_details is not None, ( + "usage-only early return delivered a Usage stripped of prompt_tokens_details to the stream consumer" + ) + assert usage.prompt_tokens_details.cached_tokens == 80 + assert usage.completion_tokens_details is not None + assert usage.completion_tokens_details.reasoning_tokens == 5