diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 6b1ada07987..3f967e29002 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -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], diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 8c67bfd97cd..a8d2da5e3b2 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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, ) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 36d0eed2a59..5806b37539c 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -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