From a64febb3e7d32aef73dc7b90665c33d0b65b4805 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:01:08 -0700 Subject: [PATCH] fix(streaming): keep an explicit provider prompt_tokens=0 or completion_tokens=0 in streamed usage The stream chunk builder started its per-chunk accumulators at 0 and adopted only nonzero counts, then fell back to litellm's tokenizer whenever the accumulated value was falsy, so a provider that reported an explicit 0 for prompt or completion tokens was billed the estimate instead. The accumulators now start at None, a usage chunk that reports a count marks it reported (a later chunk's 0 never replaces a reported nonzero), and the estimate only runs when no chunk reported the count. The Anthropic message_start cursor reset now yields None so the estimate still covers a cancelled stream, and Ollama chat streaming only attaches usage on the done chunk when both counts are present instead of inventing 0/0 on every chunk --- .../streaming_chunk_builder_utils.py | 46 ++++++----- litellm/llms/ollama/chat/transformation.py | 22 +++-- .../streaming_chunk_builder_utils.py | 4 +- .../test_streaming_chunk_builder_cursor.py | 26 +++--- .../test_streaming_chunk_builder_utils.py | 80 +++++++++++++++++++ .../ollama/test_ollama_chat_transformation.py | 47 ++++++++++- 6 files changed, 183 insertions(+), 42 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index fcd55c844c6..025db65a7ce 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -839,15 +839,17 @@ class ChunkProcessor: UsagePerChunk, ) - # # Update usage information if needed - prompt_tokens = 0 - completion_tokens = 0 + # None means no usage chunk reported the count, which is the only case + # calculate_usage() estimates with the tokenizer. An explicit provider 0 + # is a reported count and stays 0; a reported count is never replaced by + # a later chunk's 0 (Ollama sends 0/0 on every chunk before the done one). + prompt_tokens: int | None = None + completion_tokens: int | None = None # Anthropic's `message_start` SSE event carries usage.output_tokens=1 as a # cursor/placeholder; the real value only arrives in `message_delta`. - # If a stream is cancelled before `message_delta` lands, the last-wins - # accumulator below leaves completion_tokens stuck at 1 — which then - # bypasses the `completion_tokens or token_counter(...)` fallback in - # calculate_usage() because 1 is truthy. Count the completion-bearing + # If a stream is cancelled before `message_delta` lands, the accumulator + # below leaves completion_tokens stuck at 1, a reported count that + # calculate_usage() would keep. Count the completion-bearing # usage events so `_reset_anthropic_cursor_completion_tokens` can tell a # legitimate single-token reply (Anthropic emits 1 in BOTH message_start # AND message_delta, so >=2 events is positive evidence message_delta @@ -875,10 +877,15 @@ class ChunkProcessor: if usage_chunk is not None: usage_chunk_dict = self._usage_chunk_calculation_helper(usage_chunk) - if usage_chunk_dict["prompt_tokens"] is not None and usage_chunk_dict["prompt_tokens"] > 0: + if usage_chunk_dict["prompt_tokens"] is not None and ( + usage_chunk_dict["prompt_tokens"] > 0 or prompt_tokens is None + ): prompt_tokens = usage_chunk_dict["prompt_tokens"] - if usage_chunk_dict["completion_tokens"] is not None and usage_chunk_dict["completion_tokens"] > 0: + if usage_chunk_dict["completion_tokens"] is not None and ( + usage_chunk_dict["completion_tokens"] > 0 or completion_tokens is None + ): completion_tokens = usage_chunk_dict["completion_tokens"] + if usage_chunk_dict["completion_tokens"] is not None and usage_chunk_dict["completion_tokens"] > 0: completion_usage_updates += 1 if usage_chunk_dict["cache_creation_input_tokens"] is not None and ( usage_chunk_dict["cache_creation_input_tokens"] > 0 or cache_creation_input_tokens is None @@ -995,10 +1002,10 @@ class ChunkProcessor: @staticmethod def _reset_anthropic_cursor_completion_tokens( chunks: Sequence["_UsageBearingChunk | ModelResponse"], - completion_tokens: int, + completion_tokens: int | None, completion_usage_updates: int, - ) -> int: - """Reset a stale Anthropic ``message_start`` cursor placeholder to 0. + ) -> int | None: + """Reset a stale Anthropic ``message_start`` cursor placeholder to unreported. See the ``completion_usage_updates`` comment in ``_calculate_usage_per_chunk``. The accumulated value is NOT a stale @@ -1006,8 +1013,8 @@ class ChunkProcessor: carried a ``finish_reason`` (positive evidence ``message_delta`` arrived). Otherwise the only completion update we ever saw was the Anthropic ``message_start`` cursor, a small placeholder whose magnitude - varies per request (1 and 8 both observed live), so reset to 0 and let - ``calculate_usage()``'s ``or token_counter(...)`` fallback estimate from + varies per request (1 and 8 both observed live), so reset to None and let + ``calculate_usage()``'s ``token_counter(...)`` fallback estimate from the actually-received text and reasoning instead. Gated on ``custom_llm_provider == "anthropic"`` so the heuristic (which encodes Anthropic's specific message_start SSE shape) does not silently affect @@ -1028,7 +1035,7 @@ class ChunkProcessor: custom_llm_provider = hp.get("custom_llm_provider") if custom_llm_provider == "anthropic": - return 0 + return None return completion_tokens def calculate_usage( @@ -1063,15 +1070,18 @@ class ChunkProcessor: cost: Final[float | None] = calculated_usage_per_chunk["cost"] try: - returned_usage.prompt_tokens = prompt_tokens or ( - count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages) + returned_usage.prompt_tokens = ( + prompt_tokens + if prompt_tokens is not None + else (count_prompt_tokens() if count_prompt_tokens else token_counter(model=model, messages=messages)) ) except Exception: # don't allow this failing to block a complete streaming response from being returned print_verbose("token_counter failed, assuming prompt tokens is 0") returned_usage.prompt_tokens = 0 returned_usage.completion_tokens = ( completion_tokens - or ( + if completion_tokens is not None + else ( token_counter( model=model, text=completion_output, diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 181894646e3..d39b4049903 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -1,6 +1,6 @@ import json import time -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, cast from httpx._models import Headers, Response @@ -420,6 +420,18 @@ class OllamaChatConfig(BaseConfig): ) +def _done_chunk_usage(chunk: Mapping[str, object]) -> ChatCompletionUsageBlock | None: + prompt_eval_count: Final = chunk.get("prompt_eval_count") + eval_count: Final = chunk.get("eval_count") + if chunk.get("done") is not True or not isinstance(prompt_eval_count, int) or not isinstance(eval_count, int): + return None + return ChatCompletionUsageBlock( + prompt_tokens=prompt_eval_count, + completion_tokens=eval_count, + total_tokens=prompt_eval_count + eval_count, + ) + + class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): started_reasoning_content: bool = False finished_reasoning_content: bool = False @@ -528,17 +540,11 @@ class OllamaChatCompletionResponseIterator(BaseModelResponseIterator): ) ] - usage: Final = ChatCompletionUsageBlock( - prompt_tokens=chunk.get("prompt_eval_count", 0), - completion_tokens=chunk.get("eval_count", 0), - total_tokens=chunk.get("prompt_eval_count", 0) + chunk.get("eval_count", 0), - ) - return ModelResponseStream( id=str(uuid.uuid4()), object="chat.completion.chunk", created=int(time.time()), # ollama created_at is in UTC - usage=usage, + usage=_done_chunk_usage(chunk), model=chunk["model"], choices=choices, ) diff --git a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py index f981089d370..3deb307881f 100644 --- a/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/types/litellm_core_utils/streaming_chunk_builder_utils.py @@ -4,8 +4,8 @@ from ..utils import CompletionTokensDetails, PromptTokensDetailsWrapper, ServerT class UsagePerChunk(TypedDict): - prompt_tokens: int - completion_tokens: int + prompt_tokens: ReadOnly[int | None] + completion_tokens: ReadOnly[int | None] cache_creation_input_tokens: int | None cache_read_input_tokens: int | None server_tool_use: ServerToolUse | None diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py index 8617c5b81e8..b4a1f7733e7 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py @@ -15,7 +15,7 @@ calculate_usage() never fires, and the request is billed for 1 output token even when several thousand tokens of text were actually streamed. These tests pin the post-fix behavior: completion_tokens should reset -to 0 when the only update we saw was the cursor, allowing the +to None when the only update we saw was the cursor, allowing the text-based fallback to estimate from the real completion text. """ @@ -63,10 +63,10 @@ def _make_chunk( class TestAnthropicCursorBug: """The core regression: completion_tokens=1 cursor must not leak through.""" - def test_only_message_start_cursor_resets_completion_to_zero(self): + def test_only_message_start_cursor_resets_completion_to_unreported(self): """ Stream cancelled before message_delta — only the message_start cursor - (output_tokens=1) was seen. Per-chunk accumulator must reset to 0 so + (output_tokens=1) was seen. Per-chunk accumulator must reset to None so token_counter fallback can estimate from completion text. """ # Anthropic message_start: input_tokens accurate, output_tokens=1 cursor @@ -83,11 +83,11 @@ class TestAnthropicCursorBug: result = processor._calculate_usage_per_chunk(chunks=chunks) assert result["prompt_tokens"] == 1024 - # The cursor value of 1 must NOT leak through — should be reset to 0 + # The cursor value of 1 must NOT leak through — should be reset to None # so the text-based fallback estimates the real completion length. - assert result["completion_tokens"] == 0, ( + assert result["completion_tokens"] is None, ( "completion_tokens=1 from message_start cursor leaked through. " - "Should reset to 0 when only cursor was seen, so token_counter " + "Should reset to None when only cursor was seen, so token_counter " "fallback in calculate_usage() can estimate from completion text." ) @@ -233,10 +233,10 @@ class TestAnthropicCursorBug: result = processor._calculate_usage_per_chunk(chunks=chunks) assert result["cache_read_input_tokens"] == 4096 - assert result["completion_tokens"] == 0, ( + assert result["completion_tokens"] is None, ( "cache chunks alone don't count as completion progress — only " "completion_tokens > 0 in a usage event proves real output happened. " - "Reset to 0 forces token_counter fallback." + "Reset to None forces token_counter fallback." ) @pytest.mark.parametrize("placeholder", [1, 3, 8]) @@ -326,7 +326,7 @@ class TestAnthropicCursorBug: ] processor = ChunkProcessor(chunks=chunks, messages=[]) result = processor._calculate_usage_per_chunk(chunks=chunks) - assert result["completion_tokens"] == 0 + assert result["completion_tokens"] is None assert result["completion_tokens_details"] is None def test_estimated_reasoning_is_capped_to_trusted_completion_total(self): @@ -403,11 +403,11 @@ class TestNonAnthropicStreamingIntact: result = processor._calculate_usage_per_chunk(chunks=chunks) assert result["completion_tokens"] == 5 - def test_no_usage_chunks_leaves_zero(self): - """Stream with zero usage info → completion_tokens stays 0 + def test_no_usage_chunks_leaves_unreported(self): + """Stream with zero usage info → both counts stay None (token_counter fallback will handle it).""" chunks = [_make_chunk(content="hi"), _make_chunk(content=" there")] processor = ChunkProcessor(chunks=chunks, messages=[]) result = processor._calculate_usage_per_chunk(chunks=chunks) - assert result["prompt_tokens"] == 0 - assert result["completion_tokens"] == 0 + assert result["prompt_tokens"] is None + assert result["completion_tokens"] is None diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 8e7ed52fade..5d75c6699cf 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1648,3 +1648,83 @@ def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_ad ) assert usage.prompt_tokens == 77 + + +_ZERO_USAGE_TEXT_CHUNKS: Final = ( + _openai_chunk(choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hi"}, "finish_reason": None}]), + _openai_chunk(choices=[{"index": 0, "delta": {"content": " there"}, "finish_reason": None}]), + _openai_chunk(choices=[{"index": 0, "delta": {}, "finish_reason": "stop"}]), +) + + +@pytest.mark.parametrize( + "reported", + [ + pytest.param({"prompt_tokens": 0, "completion_tokens": 17, "total_tokens": 17}, id="zero_prompt"), + pytest.param({"prompt_tokens": 5, "completion_tokens": 0, "total_tokens": 5}, id="zero_completion"), + pytest.param({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, id="all_zero"), + ], +) +def test_calculate_usage_keeps_an_explicit_provider_zero(reported: Mapping[str, int]) -> None: + chunks: Final = [*_ZERO_USAGE_TEXT_CHUNKS, _openai_chunk(choices=[], usage=reported)] + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="Hi there", + messages=[{"role": "user", "content": "hi"}], + count_prompt_tokens=lambda: 999, + ) + + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + reported["prompt_tokens"], + reported["completion_tokens"], + reported["prompt_tokens"] + reported["completion_tokens"], + ) + + +def test_stream_chunk_builder_keeps_an_explicit_zero_prompt_count_end_to_end() -> None: + reported: Final = {"prompt_tokens": 0, "completion_tokens": 17, "total_tokens": 17} + chunks: Final = [*_ZERO_USAGE_TEXT_CHUNKS, _openai_chunk(choices=[], usage=reported)] + + response: Final = stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "hi"}]) + + assert response is not None + assert (response.usage.prompt_tokens, response.usage.completion_tokens, response.usage.total_tokens) == (0, 17, 17) + + +def test_calculate_usage_estimates_only_when_no_chunk_reported_usage() -> None: + chunks: Final = list(_ZERO_USAGE_TEXT_CHUNKS) + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="Hi there", + count_prompt_tokens=lambda: 77, + ) + + assert usage.prompt_tokens == 77 + assert usage.completion_tokens > 0 + assert usage.total_tokens == 77 + usage.completion_tokens + + +def test_calculate_usage_keeps_a_reported_count_over_a_later_chunks_zero() -> None: + chunks: Final = [ + _openai_chunk( + choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hi"}, "finish_reason": None}], + usage={"prompt_tokens": 5, "completion_tokens": 0, "total_tokens": 5}, + ), + _openai_chunk( + choices=[{"index": 0, "delta": {"content": " there"}, "finish_reason": "stop"}], + usage={"prompt_tokens": 0, "completion_tokens": 17, "total_tokens": 17}, + ), + ] + + usage: Final = ChunkProcessor(chunks=chunks).calculate_usage( + chunks=chunks, + model="gpt-5.4-mini", + completion_output="Hi there", + count_prompt_tokens=lambda: 999, + ) + + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (5, 17, 22) diff --git a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py index 25f9645faa0..c7fba21d222 100644 --- a/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/test_litellm/llms/ollama/test_ollama_chat_transformation.py @@ -22,7 +22,7 @@ import json from unittest.mock import MagicMock import litellm -from litellm.types.utils import Choices, Message, ModelResponse +from litellm.types.utils import Choices, Message, ModelResponse, ModelResponseStream class TestEvent(BaseModel): @@ -944,3 +944,48 @@ class TestOllamaToolCallTransformation: assert tool_msg["content"] == "Sunny, 72°F" assert "tool_call_id" in tool_msg, "tool_call_id must be forwarded to Ollama" assert tool_msg["tool_call_id"] == "call_abc123" + + +class TestOllamaStreamingUsage: + @staticmethod + def _parse(chunk: dict) -> ModelResponseStream: + iterator = OllamaChatCompletionResponseIterator(streaming_response=iter([]), sync_stream=True) + return iterator.chunk_parser(chunk) + + def test_done_chunk_reports_the_counts_ollama_sent(self): + result = self._parse( + { + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + "prompt_eval_count": 100, + "eval_count": 50, + } + ) + + assert result.usage is not None + assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (100, 50, 150) + + def test_done_chunk_without_counts_reports_no_usage_instead_of_zeros(self): + result = self._parse( + { + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": ""}, + "done": True, + "done_reason": "stop", + } + ) + + assert result.usage is None + + def test_chunk_before_done_reports_no_usage(self): + result = self._parse( + { + "model": "qwen3:0.6b", + "message": {"role": "assistant", "content": "Hi"}, + "done": False, + } + ) + + assert result.usage is None