scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-04 02:31:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 71

fix(streaming): estimate interrupted Anthropic stream usage from reasoning_content

Browse source 
Interrupted Anthropic streams that die before message_delta were billed at
the message_start placeholder (any value above 1 was trusted) or at 0 when
the partial response was reasoning-only, because the token_counter fallback
only looked at visible text. Reset the placeholder whenever no finish_reason
or second usage event arrived, fold the already-counted reasoning tokens into
the fallback estimate, and drop the stale completion_tokens_details so cost
is computed from the recovered count

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-16 22:49:25 +00:00
parent 9357491f91
commit 4b70696afa
2 changed files with 132 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

70
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -5,6 +5,7 @@ from itertools import groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast
from pydantic import BaseModel
from typing_extensions import ReadOnly, Required
from litellm._logging import verbose_logger
@ -148,6 +149,7 @@ class _ToolCallChunk(TypedDict):
class _UsageBearingChunk(TypedDict, total=False):
usage: Usage | None
_hidden_params: Mapping[str, str]
choices: ReadOnly[Sequence[StreamingChoices | Mapping[str, object]]]
class _UsageSummary(TypedDict):
@ -921,21 +923,22 @@ class ChunkProcessor:
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
recovered_completion_tokens: Final = self._reset_anthropic_cursor_completion_tokens(
chunks=chunks,
completion_tokens=completion_tokens,
completion_usage_updates=completion_usage_updates,
)
cursor_was_reset: Final = recovered_completion_tokens != completion_tokens
return UsagePerChunk(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
completion_tokens=recovered_completion_tokens,
cache_creation_input_tokens=cache_creation_input_tokens,
cache_read_input_tokens=cache_read_input_tokens,
server_tool_use=server_tool_use,
web_search_requests=web_search_requests,
google_maps_grounding_requests=google_maps_grounding_requests,
completion_tokens_details=completion_tokens_details,
completion_tokens_details=None if cursor_was_reset else completion_tokens_details,
prompt_tokens_details=prompt_tokens_details,
cost=cost,
inference_geo=self._last_provider_pricing_field(chunks, "inference_geo"),
@ -960,6 +963,32 @@ class ChunkProcessor:
]
return values[-1] if values else None
@staticmethod
def _finish_reason_of_choice(choice: object) -> str | None:
match choice:
case StreamingChoices(finish_reason=reason) | Choices(finish_reason=reason):
return reason
case {"finish_reason": str() as reason}:
return reason
case _:
return None
@staticmethod
def _chunk_choices(chunk: "_UsageBearingChunk | BaseModel") -> Sequence[object]:
if isinstance(chunk, dict):
return chunk.get("choices", ())
if isinstance(chunk, (ModelResponse, ModelResponseStream)):
return chunk.choices
return ()
@staticmethod
def _saw_finish_reason(chunks: Sequence["_UsageBearingChunk | ModelResponse"]) -> bool:
return any(
ChunkProcessor._finish_reason_of_choice(choice) is not None
for chunk in chunks
for choice in ChunkProcessor._chunk_choices(chunk)
)
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: Sequence["_UsageBearingChunk | ModelResponse"],
@ -970,18 +999,18 @@ class ChunkProcessor:
See the ``completion_usage_updates`` comment in
``_calculate_usage_per_chunk``. The accumulated value is NOT a stale
cursor when either it is > 1 (definitely not a placeholder) or we saw
>= 2 completion-bearing usage events (positive evidence ``message_delta``
arrived). Otherwise — the only completion update we ever saw was the
Anthropic ``message_start`` cursor (=1) — reset to 0 so
``calculate_usage()``'s ``or token_counter(text=...)`` fallback estimates
from the actually-received completion text instead of trusting the
placeholder. Gated on ``custom_llm_provider == "anthropic"`` so the
heuristic (which encodes Anthropic's specific message_start SSE shape)
does not silently affect other providers that may legitimately report
``completion_tokens=1`` from a single usage event.
cursor when we saw >= 2 completion-bearing usage events or any chunk
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
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
other providers that legitimately report usage from a single event.
"""
saw_non_cursor_completion: Final = completion_tokens > 1 or completion_usage_updates >= 2
saw_non_cursor_completion: Final = completion_usage_updates >= 2 or ChunkProcessor._saw_finish_reason(chunks)
if saw_non_cursor_completion:
return completion_tokens
@ -995,7 +1024,7 @@ class ChunkProcessor:
if isinstance(hp, dict):
custom_llm_provider = hp.get("custom_llm_provider")
if custom_llm_provider == "anthropic" and completion_tokens == 1:
if custom_llm_provider == "anthropic":
return 0
return completion_tokens
@ -1039,10 +1068,13 @@ class ChunkProcessor:
returned_usage.prompt_tokens = 0
returned_usage.completion_tokens = (
completion_tokens
or token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
or (
token_counter(
model=model,
text=completion_output,
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
)
+ (reasoning_tokens or 0)
)
)
returned_usage.total_tokens = returned_usage.prompt_tokens + returned_usage.completion_tokens

86
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_cursor.py
View file

@ -22,9 +22,10 @@ text-based fallback to estimate from the real completion text.
import pytest
import litellm
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
Delta,
ModelResponseStream,
StreamingChoices,
@ -35,6 +36,7 @@ from litellm.types.utils import (
def _make_chunk(
*,
content: str = "",
reasoning_content: str | None = None,
usage: Usage = None,
finish_reason: str = None,
custom_llm_provider: str = "anthropic",
@ -48,7 +50,7 @@ def _make_chunk(
StreamingChoices(
finish_reason=finish_reason,
index=0,
delta=Delta(content=content, role="assistant"),
delta=Delta(content=content, role="assistant", reasoning_content=reasoning_content),
)
],
usage=usage,
@ -253,6 +255,79 @@ class TestAnthropicCursorBug:
"Reset to 0 forces token_counter fallback."
)
@pytest.mark.parametrize("placeholder", [1, 3, 8])
def test_interrupted_reasoning_only_stream_estimates_from_reasoning(self, placeholder: int):
"""
message_start placeholders are not always 1 (live Anthropic streams
have been observed sending 1 and 8 for the same prompt), and a thinking
model cut off before message_delta has streamed only reasoning_content.
The recovered usage, including the completion_tokens_details the cost
calculator bills from, must come from that reasoning rather than from
the placeholder.
"""
message_start = _make_chunk(
usage=Usage(
prompt_tokens=100,
completion_tokens=placeholder,
total_tokens=100 + placeholder,
completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=0, text_tokens=placeholder),
)
)
reasoning_text = "Let me work through the scheduling constraints step by step. " * 40
reasoning_chunks = [
_make_chunk(reasoning_content=reasoning_text[i : i + 50]) for i in range(0, len(reasoning_text), 50)
]
response = litellm.stream_chunk_builder(
chunks=[message_start, *reasoning_chunks],
messages=[{"role": "user", "content": "Plan the schedule."}],
)
assert response.choices[0].message.reasoning_content == reasoning_text
reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
assert reasoning_tokens > placeholder
assert response.usage.completion_tokens == reasoning_tokens, (
f"Expected completion_tokens to be the reasoning estimate, got "
f"completion_tokens={response.usage.completion_tokens} reasoning_tokens={reasoning_tokens}"
)
assert response.usage.total_tokens == response.usage.prompt_tokens + reasoning_tokens
details = response.usage.completion_tokens_details
assert (details.text_tokens or 0) + details.reasoning_tokens == response.usage.completion_tokens
def test_fallback_counts_reasoning_and_text_together(self):
"""
With no usable provider count, the estimate covers everything the
provider generated: reasoning_content plus visible text, not text alone.
"""
reasoning = "First I should check whether the input is sorted. " * 10
text = "The list is already sorted, so no work is needed."
chunks = [_make_chunk(reasoning_content=reasoning), _make_chunk(content=text)]
response = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": "Sort it."}])
text_only = litellm.token_counter(model="claude-sonnet-4-6", text=text, count_response_tokens=True)
reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens
assert reasoning_tokens > 0
assert response.usage.completion_tokens == text_only + reasoning_tokens
def test_lone_usage_event_with_finish_reason_is_trusted(self):
"""
Guardrails rebuild responses from the chunks yielded to the client,
which excludes the un-yielded message_start. A finished stream then has
exactly one usage event (message_delta) and it must be kept as-is.
"""
chunks = [
_make_chunk(content="Yes, "),
_make_chunk(content="that works."),
_make_chunk(
usage=Usage(prompt_tokens=20, completion_tokens=5, total_tokens=25),
finish_reason="stop",
),
]
processor = ChunkProcessor(chunks=chunks, messages=[])
result = processor._calculate_usage_per_chunk(chunks=chunks)
assert result["completion_tokens"] == 5
class TestProviderGuard:
"""Class A: the cursor-reset heuristic must NOT silently affect non-Anthropic
@ -297,11 +372,12 @@ class TestNonAnthropicStreamingIntact:
"""Make sure providers without cursor pattern still work."""
def test_completion_tokens_above_one_never_resets(self):
"""Any chunk reporting completion_tokens > 1 sets saw_non_cursor
and prevents the reset."""
"""A non-Anthropic provider reporting completion_tokens > 1 from a
single usage event keeps that value."""
chunks = [
_make_chunk(
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
custom_llm_provider="openai",
),
]
processor = ChunkProcessor(chunks=chunks, messages=[])