From cb70bc560e1bbd8a1d619fb5dead7f2f0c333e74 Mon Sep 17 00:00:00 2001 From: Ambuj Upadhyay Date: Wed, 28 Oct 2026 03:51:11 +0530 Subject: [PATCH] fix(streaming): dict-usage arm drops prompt_tokens_details/completion_tokens_details its sibling arms preserve (#36678) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(streaming): dict-usage arm no longer drops token detail fields In _dispatch_provider_chunk's openai/azure else-branch, a plain-dict usage was rebuilt into litellm.Usage from only prompt_tokens / completion_tokens / total_tokens. prompt_tokens_details (cached_tokens), completion_tokens_details (reasoning_tokens) and cache token counts were silently dropped — while both sibling arms (Usage passthrough, BaseModel via Usage(**model_dump())) preserve them. Every in-repo chunk producer currently normalizes usage to Usage/BaseModel before this point, so the dict arm is reachable today only via chunks litellm did not construct (e.g. a custom streaming_decoder feeding CustomStreamWrapper). When it does fire on a usage-only final chunk (stream_options.include_usage), the stripped Usage is exactly what the stream consumer receives — the early-return path has no later aggregation to repair it. Fix: pass the full dict through Usage(**usage_dict). Usage.__init__ coerces nested detail dicts into their wrapper types and tolerates extra keys. Side effect: a legitimate 0 count is no longer coerced to None by the old 'or None' expressions. Co-Authored-By: Claude Fable 5 * chore: drop explanatory source comment per repo comment policy Rationale lives in the PR body and the regression tests. Co-Authored-By: Claude Fable 5 * fix(streaming): type the dict-usage unpack for the basedpyright gate Usage(**response_obj["usage"]) unpacks a dict whose value type pyright cannot see, and the delta gate reports that once per typed constructor parameter: five, against the three .get() calls the line replaced. Moving the unpack behind _usage_from_payload keeps the same Usage(**payload) call and is reported once, at the call site, so the rule count ends below its merge base --------- Co-authored-by: Claude Fable 5 --- .../litellm_core_utils/streaming_handler.py | 11 ++- .../test_streaming_handler.py | 99 +++++++++++++++++++ 2 files changed, 105 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8fb602ebeb0..b64e146ca08 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1589,11 +1589,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, - ), + _usage_from_payload(response_obj["usage"]), ) elif isinstance(response_obj["usage"], Usage): setattr( @@ -2662,6 +2658,11 @@ class CustomStreamWrapper: _strip_sse_data_from_chunk = strip_sse_data_from_chunk +def _usage_from_payload(payload: Mapping[str, Any]) -> Usage: + """Build ``Usage`` from a wire-shaped usage dict, forwarding every key.""" + return litellm.Usage(**payload) + + def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: for key in keys: value = getattr(details, key, None) diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index cc1fa1b8803..6d700bae514 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7844,3 +7844,102 @@ async def test_leaving_an_async_with_block_closes_the_provider_stream(): async with _wrapper_over(stream): pass assert stream.close_calls == 1 + + +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