scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(streaming): carry Anthropic cache-creation TTL split through fallback usage reassembly

Browse source
This commit is contained in:
mateo-berri 2026-08-05 02:06:52 -07:00
parent 74f1b934be
commit a9902fcdb5
3 changed files with 98 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

62
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -39,6 +39,34 @@ if TYPE_CHECKING:
)
def capture_cache_creation_token_details(
prompt_tokens_details: PromptTokensDetailsWrapper | None,
current: CacheCreationTokenDetails | None,
) -> CacheCreationTokenDetails | None:
incoming: Final = cast(
CacheCreationTokenDetails | None,
getattr(prompt_tokens_details, "cache_creation_token_details", None),
)
if incoming is not None:
return incoming
return current
def attach_cache_creation_token_details(
prompt_tokens_details: PromptTokensDetailsWrapper | None,
cache_creation_token_details: CacheCreationTokenDetails | None,
) -> PromptTokensDetailsWrapper | None:
if prompt_tokens_details is None or cache_creation_token_details is None:
return prompt_tokens_details
existing: Final = cast(
CacheCreationTokenDetails | None,
getattr(prompt_tokens_details, "cache_creation_token_details", None),
)
if existing is not None:
return prompt_tokens_details
return prompt_tokens_details.model_copy(update={"cache_creation_token_details": cache_creation_token_details})
class ChunkProcessor:
def __init__(self, chunks: list, messages: list | None = None):
self.chunks = self._sort_chunks(chunks)
@ -701,16 +729,14 @@ class ChunkProcessor:
or prompt_tokens_details
)
cache_creation_token_details = self._capture_cache_creation_token_details(
cache_creation_token_details = capture_cache_creation_token_details(
prompt_tokens_details, cache_creation_token_details
)
if usage_chunk_dict["cost"] is not None:
cost = usage_chunk_dict["cost"]
prompt_tokens_details = self._attach_cache_creation_token_details(
prompt_tokens_details, cache_creation_token_details
)
prompt_tokens_details = attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details)
completion_tokens = self._reset_anthropic_cursor_completion_tokens(
chunks=chunks,
@ -730,34 +756,6 @@ class ChunkProcessor:
cost=cost,
)
@staticmethod
def _capture_cache_creation_token_details(
prompt_tokens_details: PromptTokensDetailsWrapper | None,
current: CacheCreationTokenDetails | None,
) -> CacheCreationTokenDetails | None:
incoming: Final = cast(
CacheCreationTokenDetails | None,
getattr(prompt_tokens_details, "cache_creation_token_details", None),
)
if incoming is not None:
return incoming
return current
@staticmethod
def _attach_cache_creation_token_details(
prompt_tokens_details: PromptTokensDetailsWrapper | None,
cache_creation_token_details: CacheCreationTokenDetails | None,
) -> PromptTokensDetailsWrapper | None:
if prompt_tokens_details is None or cache_creation_token_details is None:
return prompt_tokens_details
existing: Final = cast(
CacheCreationTokenDetails | None,
getattr(prompt_tokens_details, "cache_creation_token_details", None),
)
if existing is not None:
return prompt_tokens_details
return prompt_tokens_details.model_copy(update={"cache_creation_token_details": cache_creation_token_details})
@staticmethod
def _reset_anthropic_cursor_completion_tokens(
chunks: list[dict[str, Any] | ModelResponse],

18
litellm/litellm_core_utils/streaming_handler.py
View file

@ -25,6 +25,7 @@ from litellm.litellm_core_utils.thread_pool_executor import executor
from litellm.types.llms.openai import OpenAIChatCompletionChunk
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import (
CacheCreationTokenDetails,
CompletionTokensDetailsWrapper,
Delta,
LlmProviders,
@ -2251,11 +2252,17 @@ def _coerce_token_details(
def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
"""Assume most recent usage chunk has total usage uptil then."""
from litellm.litellm_core_utils.streaming_chunk_builder_utils import (
attach_cache_creation_token_details,
capture_cache_creation_token_details,
)
prompt_tokens: int = 0
completion_tokens: int = 0
latest_usage_chunk = None
prompt_tokens_details: PromptTokensDetailsWrapper | None = None
completion_tokens_details: CompletionTokensDetailsWrapper | None = None
cache_creation_token_details: CacheCreationTokenDetails | None = None
for chunk in chunks:
if "usage" in chunk and chunk["usage"] is not None:
@ -2265,10 +2272,13 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
prompt_tokens = usage.get("prompt_tokens", 0) or 0
if "completion_tokens" in usage:
completion_tokens = usage.get("completion_tokens", 0) or 0
prompt_tokens_details = (
_coerce_token_details(usage, "prompt_tokens_details", PromptTokensDetailsWrapper)
or prompt_tokens_details
incoming_prompt_tokens_details = _coerce_token_details(
usage, "prompt_tokens_details", PromptTokensDetailsWrapper
)
cache_creation_token_details = capture_cache_creation_token_details(
incoming_prompt_tokens_details, cache_creation_token_details
)
prompt_tokens_details = incoming_prompt_tokens_details or prompt_tokens_details
completion_tokens_details = (
_coerce_token_details(usage, "completion_tokens_details", CompletionTokensDetailsWrapper)
or completion_tokens_details
@ -2278,7 +2288,7 @@ def calculate_total_usage(chunks: list[ModelResponse]) -> Usage:
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=prompt_tokens + completion_tokens,
prompt_tokens_details=prompt_tokens_details,
prompt_tokens_details=attach_cache_creation_token_details(prompt_tokens_details, cache_creation_token_details),
completion_tokens_details=completion_tokens_details,
)

54
tests/test_litellm/litellm_core_utils/test_streaming_handler.py
View file

@ -1492,6 +1492,60 @@ def test_calculate_total_usage_preserves_prompt_cache_token_details():
assert usage.completion_tokens_details.reasoning_tokens == 2
def test_calculate_total_usage_preserves_anthropic_cache_creation_ttl_breakdown():
"""Anthropic sends the 5m/1h cache-write split only on `message_start`; the later
`message_delta` repeats the flat count without the split. Losing it here bills 1h
cache writes at the cheaper 5m rate."""
from litellm.litellm_core_utils.streaming_handler import calculate_total_usage
from litellm.types.utils import CacheCreationTokenDetails
message_start_chunk = ModelResponseStream(
id="chatcmpl-1",
created=1745513206,
model="claude-sonnet-5",
choices=[
StreamingChoices(finish_reason=None, index=0, delta=Delta(content="Hi"))
],
usage=Usage(
prompt_tokens=120,
completion_tokens=1,
total_tokens=121,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=0,
cache_creation_tokens=100,
cache_creation_token_details=CacheCreationTokenDetails(
ephemeral_5m_input_tokens=20, ephemeral_1h_input_tokens=80
),
),
),
)
message_delta_chunk = ModelResponseStream(
id="chatcmpl-1",
created=1745513207,
model="claude-sonnet-5",
choices=[
StreamingChoices(finish_reason="stop", index=0, delta=Delta(content=""))
],
usage=Usage(
prompt_tokens=120,
completion_tokens=4,
total_tokens=124,
prompt_tokens_details=PromptTokensDetailsWrapper(
cached_tokens=0, cache_creation_tokens=100
),
),
)
usage = calculate_total_usage([message_start_chunk, message_delta_chunk])
assert usage.prompt_tokens_details is not None
assert usage.prompt_tokens_details.cache_creation_tokens == 100
ttl_breakdown = usage.prompt_tokens_details.cache_creation_token_details
assert ttl_breakdown is not None
assert ttl_breakdown.ephemeral_5m_input_tokens == 20
assert ttl_breakdown.ephemeral_1h_input_tokens == 80
@pytest.mark.asyncio
async def test_openrouter_streaming_cost_after_finish_reason(logging_obj: Logging):
from litellm.utils import ModelResponseListIterator