diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 3565872f09e..651b169a9b2 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2326,7 +2326,7 @@ class CustomStreamWrapper: return if self.model: partial_response.model = self.model - zero_fill_missing_cache_usage_fields(usage) + backfill_missing_cache_usage_fields(usage) self.logging_obj.model_call_details["combined_usage_object"] = usage self.logging_obj.model_call_details["response_cost"] = ( self.logging_obj._response_cost_calculator(result=partial_response) or 0.0 @@ -2447,13 +2447,33 @@ class CustomStreamWrapper: return chunk -def zero_fill_missing_cache_usage_fields(usage: Usage) -> None: - if getattr(usage, "cache_creation_input_tokens", None) is None: - usage.cache_creation_input_tokens = 0 # rebind-ok: in-place zero-fill is the contract +def _cache_token_count(details: PromptTokensDetailsWrapper | None, keys: tuple[str, ...]) -> int: + for key in keys: + value = getattr(details, key, None) + if isinstance(value, int) and not isinstance(value, bool) and value: + return value + return 0 + + +def backfill_missing_cache_usage_fields(usage: Usage) -> None: + """Give partial-stream usage the same cache fields a complete stream reports. + + Carries OpenAI-style ``prompt_tokens_details`` counts up to the Anthropic-style + top-level keys, defaulting to zero. It must carry the real count rather than a + flat zero: downstream readers treat these keys as authoritative once present and + skip their own normalization, so a zero here would overwrite a real cache read. + """ + details: Final = usage.prompt_tokens_details if getattr(usage, "cache_read_input_tokens", None) is None: - usage.cache_read_input_tokens = 0 # rebind-ok: in-place zero-fill is the contract + usage.cache_read_input_tokens = _cache_token_count( # rebind-ok: in-place backfill is the contract + details, ("cached_tokens",) + ) + if getattr(usage, "cache_creation_input_tokens", None) is None: + usage.cache_creation_input_tokens = _cache_token_count( # rebind-ok: in-place backfill is the contract + details, ("cache_write_tokens", "cache_creation_tokens") + ) if usage.prompt_tokens_details is None: - usage.prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=0) # rebind-ok: zero-fill in place + usage.prompt_tokens_details = PromptTokensDetailsWrapper(cached_tokens=0) # rebind-ok: backfill in place _TokenDetails = TypeVar("_TokenDetails", PromptTokensDetailsWrapper, CompletionTokensDetailsWrapper) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d430d776350..72cc298d37d 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -44,7 +44,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.streaming_handler import ( - zero_fill_missing_cache_usage_fields, + backfill_missing_cache_usage_fields, ) from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model @@ -332,7 +332,7 @@ async def _bill_partial_streamed_spend_on_disconnect(request_data: dict, respons partial_response.model = wrapper_model partial_usage: Final = getattr(partial_response, "usage", None) if isinstance(partial_usage, Usage): - zero_fill_missing_cache_usage_fields(partial_usage) + backfill_missing_cache_usage_fields(partial_usage) try: await logging_obj.dispatch_success_handlers( partial_response, 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 a550160cd73..9f04d63b6ad 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -3459,7 +3459,7 @@ def test_record_partial_usage_for_failure_counts_prompt_tokens_from_request_mess assert stashed.prompt_tokens > 0 -def test_record_partial_usage_for_failure_zero_fills_missing_cache_fields(): +def test_record_partial_usage_for_failure_backfills_missing_cache_fields(): wrapper, logging_obj = _wrapper_with_partial_chunks(chunk_model="gpt-4o-mini") wrapper._record_partial_usage_for_failure() @@ -3471,6 +3471,24 @@ def test_record_partial_usage_for_failure_zero_fills_missing_cache_fields(): assert stashed.prompt_tokens_details.cached_tokens == 0 +def test_record_partial_usage_for_failure_carries_up_openai_style_cached_tokens(): + recovered = Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=500), + ) + wrapper, logging_obj = _wrapper_with_partial_chunks( + chunk_model="gpt-4o-mini", usage=recovered + ) + + wrapper._record_partial_usage_for_failure() + + stashed = logging_obj.model_call_details["combined_usage_object"] + assert stashed.cache_read_input_tokens == 500 + assert stashed.cache_creation_input_tokens == 0 + + def test_record_partial_usage_for_failure_keeps_cache_values_recovered_from_chunks(): recovered = Usage( prompt_tokens=40, diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 609d0dfe1b6..4d78bf164b1 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -5555,7 +5555,7 @@ class TestStreamingClientDisconnectBilling: assert standard_logging_object["response_cost"] > 0.0 @pytest.mark.asyncio - async def test_disconnect_billing_zero_fills_missing_cache_fields(self): + async def test_disconnect_billing_backfills_missing_cache_fields(self): event = await self._bill_and_collect_success_event() usage = event["response_obj"].usage @@ -5564,6 +5564,48 @@ class TestStreamingClientDisconnectBilling: assert usage.prompt_tokens_details is not None assert usage.prompt_tokens_details.cached_tokens == 0 + @pytest.mark.asyncio + async def test_disconnect_billing_carries_up_openai_style_cached_tokens(self): + from litellm.types.utils import ( + Delta, + ModelResponseStream, + PromptTokensDetailsWrapper, + StreamingChoices, + Usage, + ) + + def append_openai_style_cached_usage_chunk(response): + response.chunks.append( + ModelResponseStream( + id=response.chunks[0].id, + model="gpt-4o-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=" and more", role="assistant"), + ) + ], + usage=Usage( + prompt_tokens=1000, + completion_tokens=10, + total_tokens=1010, + prompt_tokens_details=PromptTokensDetailsWrapper( + cached_tokens=500 + ), + ), + ) + ) + + event = await self._bill_and_collect_success_event( + append_openai_style_cached_usage_chunk + ) + + usage = event["response_obj"].usage + assert getattr(usage, "cache_read_input_tokens", None) == 500 + assert getattr(usage, "cache_creation_input_tokens", None) == 0 + @pytest.mark.asyncio async def test_disconnect_billing_keeps_cache_values_recovered_from_chunks(self): from litellm.types.utils import (