diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 025db65a7ce..aa4e0cf5495 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -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: diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py index 5d75c6699cf..af763da2d87 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -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}]),