Merge pull request #42323 from BerriAI/litellm_stream_usage_keep_explicit_zero

fix(streaming): keep an explicit provider prompt_tokens=0 or completion_tokens=0 in streamed usage

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,

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,
             )

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

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

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)

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