scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(streaming): let a later usage event zero out stale cache counts (#40736) (#42330)

Browse source 
* fix(streaming): let a later usage event zero out stale cache counts (#40736)

The usage merger only replaced cache_creation_input_tokens and
cache_read_input_tokens with a positive value, so when Anthropic's
message_delta restated the cache block with a 0 write, the 58k write from
message_start survived next to the 58k read. prompt_tokens minus both cache
counts then went negative and the write was billed twice

A usage event that reports any prompt-side count is now authoritative for both
cache fields, zeros included. An event with no prompt-side counts, such as an
output-only message_delta, still leaves the earlier values alone

* test(streaming): cover an input-only message_delta and the derived uncached input in the cache count merger test

---------

Co-authored-by: muhammadwaqar12 <m_waqar@live.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: shrey kharbanda <shreshth@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-09-22 11:50:19 -07:00 • committed by GitHub
parent 4bcdaf3d4b
commit ee3f5a8bdf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 77 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -165,6 +165,14 @@ class _UsageSummary(TypedDict):
cost: float | None
def _reports_prompt_side_usage(usage_summary: "_UsageSummary") -> bool:
return (
(usage_summary["prompt_tokens"] or 0) > 0
or (usage_summary["cache_creation_input_tokens"] or 0) > 0
or (usage_summary["cache_read_input_tokens"] or 0) > 0
)
def capture_cache_creation_token_details(
prompt_tokens_details: PromptTokensDetailsWrapper | None,
current: CacheCreationTokenDetails | None,
@ -888,11 +896,11 @@ class ChunkProcessor:
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
_reports_prompt_side_usage(usage_chunk_dict) or cache_creation_input_tokens is None
):
cache_creation_input_tokens = usage_chunk_dict["cache_creation_input_tokens"]
if usage_chunk_dict["cache_read_input_tokens"] is not None and (
usage_chunk_dict["cache_read_input_tokens"] > 0 or cache_read_input_tokens is None
_reports_prompt_side_usage(usage_chunk_dict) or cache_read_input_tokens is None
):
cache_read_input_tokens = usage_chunk_dict["cache_read_input_tokens"]
if usage_chunk_dict["completion_tokens_details"] is not None:

67
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -1,5 +1,6 @@
from collections.abc import Mapping, Sequence
from typing import Final
from unittest.mock import MagicMock
import pytest
@ -7,6 +8,7 @@ import pytest
from litellm import ChatCompletionUsageBlock, stream_chunk_builder
from litellm.types.utils import GenericStreamingChunk
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.llms.anthropic.chat.handler import ModelResponseIterator
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
ChatCompletionMessageToolCall,
@ -1650,6 +1652,71 @@ def test_calculate_usage_falls_back_to_prompt_counter_when_mock_stream_has_no_ad
assert usage.prompt_tokens == 77
@pytest.mark.parametrize(
("message_delta_usage", "expected_cache_creation", "expected_cache_read"),
[
(
{
"input_tokens": 2,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 58352,
"output_tokens": 408,
},
0,
58352,
),
({"output_tokens": 408}, 58352, 0),
({"input_tokens": 2, "output_tokens": 408}, 58352, 0),
],
ids=["delta_restates_cache_counts", "delta_reports_output_only", "delta_reports_input_and_output_only"],
)
def test_anthropic_stream_usage_takes_cache_counts_from_last_event_that_reports_them(
message_delta_usage: Mapping[str, int], expected_cache_creation: int, expected_cache_read: int
) -> None:
iterator: Final = ModelResponseIterator(streaming_response=MagicMock(), sync_stream=True, json_mode=False)
events: Final = (
{
"type": "message_start",
"message": {
"id": "msg_1",
"type": "message",
"role": "assistant",
"model": "claude-sonnet-4-5",
"content": [],
"stop_reason": None,
"usage": {
"input_tokens": 2,
"cache_creation_input_tokens": 58352,
"cache_read_input_tokens": 0,
"output_tokens": 1,
},
},
},
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}},
{"type": "content_block_stop", "index": 0},
{"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": dict(message_delta_usage)},
{"type": "message_stop"},
)
response: Final = stream_chunk_builder(
chunks=[iterator.chunk_parser(event) for event in events],
messages=[{"role": "user", "content": "hi"}],
)
assert response.usage.cache_creation_input_tokens == expected_cache_creation
assert response.usage.cache_read_input_tokens == expected_cache_read
assert response.usage.prompt_tokens == 58354
assert response.usage.prompt_tokens_details.cache_creation_tokens == expected_cache_creation
assert response.usage.prompt_tokens_details.cached_tokens == expected_cache_read
assert (
response.usage.prompt_tokens
- response.usage.cache_read_input_tokens
- response.usage.cache_creation_input_tokens
== 2
)
_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}]),