From d346d769f02811d5e325a2ac5f45ac54ff5593de Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 16 Oct 2026 23:28:11 -0700 Subject: [PATCH] fix(caching): keep prompt cache usage on streamed cache hits (#47073) * fix(caching): keep prompt cache usage on streamed cache hits A stream=true request served from the response cache rebuilt its usage from the prompt, completion and total counts only, so prompt_tokens_details and the cache read and creation counts were dropped on the replayed stream. /v1/messages clients routed to Bedrock Converse saw input_tokens equal to the whole prompt and no cache tokens on a hit. Both replays now build the usage through one helper that also restores prompt_tokens_details and the integer cache read and creation counts, skipping cached values that do not validate. * fix(caching): coerce replayed token counts the way the live usage does The replay accepted only exact ints for the token counts, so an integral float or numeric string in a cached entry became 0 where the previous replay kept it. Counts now go through pydantic int coercion and fall back only when they do not validate. --- .../convert_dict_to_response.py | 72 ++++++++++++---- .../test_convert_to_streaming_response.py | 83 +++++++++++++++++++ .../test_streaming_iterator_combined_chunk.py | 41 +++++++++ 3 files changed, 182 insertions(+), 14 deletions(-) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index f3c60679c81..8b94c34942c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -251,14 +251,7 @@ async def convert_to_streaming_response_async( setattr( model_response_object, "usage", - Usage( - completion_tokens=response_object["usage"].get("completion_tokens", 0), - prompt_tokens=response_object["usage"].get("prompt_tokens", 0), - total_tokens=response_object["usage"].get("total_tokens", 0), - completion_tokens_details=_cached_completion_tokens_details( - _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) - ), - ), + _replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])), ) if "id" in response_object: @@ -345,12 +338,10 @@ def convert_to_streaming_response( model_response_object.choices = choice_list if "usage" in response_object and response_object["usage"] is not None: - setattr(model_response_object, "usage", Usage()) - model_response_object.usage.completion_tokens = response_object["usage"].get("completion_tokens", 0) - model_response_object.usage.prompt_tokens = response_object["usage"].get("prompt_tokens", 0) - model_response_object.usage.total_tokens = response_object["usage"].get("total_tokens", 0) - model_response_object.usage.completion_tokens_details = _cached_completion_tokens_details( - _CACHED_USAGE_DETAILS.validate_python(response_object["usage"]) + setattr( + model_response_object, + "usage", + _replayed_usage(_CACHED_USAGE_DETAILS.validate_python(response_object["usage"])), ) if "id" in response_object: @@ -414,13 +405,28 @@ class _CachedMessageProviderFields(TypedDict, total=False): class _CachedUsageDetails(TypedDict, total=False): + completion_tokens: ReadOnly[object] + prompt_tokens: ReadOnly[object] + total_tokens: ReadOnly[object] completion_tokens_details: ReadOnly[object] + prompt_tokens_details: ReadOnly[object] + cache_read_input_tokens: ReadOnly[object] + cache_creation_input_tokens: ReadOnly[object] + + +class _PromptCacheTokenCounts(TypedDict, total=False): + cache_read_input_tokens: ReadOnly[int] + cache_creation_input_tokens: ReadOnly[int] _CACHED_MESSAGE_PROVIDER_FIELDS: Final = TypeAdapter(_CachedMessageProviderFields) _CACHED_USAGE_DETAILS: Final = TypeAdapter(_CachedUsageDetails) _FIELD_MAP: Final = TypeAdapter(dict[str, object]) _COMPLETION_TOKENS_DETAILS: Final = TypeAdapter(CompletionTokensDetailsWrapper) +_PROMPT_TOKENS_DETAILS: Final = TypeAdapter(PromptTokensDetailsWrapper) +_PROMPT_CACHE_TOKEN_COUNTS: Final = TypeAdapter(_PromptCacheTokenCounts) +_TOKEN_COUNT: Final = TypeAdapter(int) +_PROMPT_CACHE_TOKEN_FIELDS: Final = ("cache_read_input_tokens", "cache_creation_input_tokens") def _cached_provider_specific_fields(message: _CachedMessageProviderFields) -> Mapping[str, object] | None: @@ -438,6 +444,44 @@ def _cached_completion_tokens_details(usage: _CachedUsageDetails) -> CompletionT return None +def _cached_prompt_tokens_details(usage: _CachedUsageDetails) -> PromptTokensDetailsWrapper | None: + details: Final = usage.get("prompt_tokens_details") + if not isinstance(details, dict): + return None + try: + return _PROMPT_TOKENS_DETAILS.validate_python(details) + except ValidationError: + return None + + +def _cached_token_count(count: object) -> int | None: + try: + return _TOKEN_COUNT.validate_python(count) + except ValidationError: + return None + + +def _cached_prompt_cache_token_counts(usage: _CachedUsageDetails) -> _PromptCacheTokenCounts: + return _PROMPT_CACHE_TOKEN_COUNTS.validate_python( + { + field: count + for field in _PROMPT_CACHE_TOKEN_FIELDS + if (count := _cached_token_count(usage.get(field))) is not None + } + ) + + +def _replayed_usage(usage: _CachedUsageDetails) -> Usage: + return Usage( + completion_tokens=_cached_token_count(usage.get("completion_tokens")) or 0, + prompt_tokens=_cached_token_count(usage.get("prompt_tokens")) or 0, + total_tokens=_cached_token_count(usage.get("total_tokens")) or 0, + completion_tokens_details=_cached_completion_tokens_details(usage), + prompt_tokens_details=_cached_prompt_tokens_details(usage), + **_cached_prompt_cache_token_counts(usage), + ) + + def handle_invalid_parallel_tool_calls( tool_calls: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall], ) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] | None: diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 6b017f9ff79..348d0105a6c 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -155,6 +155,89 @@ async def test_a_replayed_stream_skips_cached_completion_tokens_details_that_do_ assert rebuilt.usage.completion_tokens == 7 +def _prompt_cached_payload(content: str, cache_read: int, cache_write: int) -> dict: + cached = _async_payload(content=content) + cached["usage"] = litellm.Usage( + prompt_tokens=42145, + completion_tokens=6, + total_tokens=42151, + cache_read_input_tokens=cache_read, + cache_creation_input_tokens=cache_write, + ).model_dump() + return cached + + +@pytest.mark.parametrize("content", ["", "Hello world! How are you?"], ids=["single_frame", "sliced"]) +@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_the_cached_prompt_cache_tokens(replay, cache_read, cache_write, content): + cached = _prompt_cached_payload(content, cache_read, cache_write) + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None) + + assert usage.prompt_tokens == 42145 + assert usage.prompt_tokens_details.cached_tokens == cache_read + assert usage.prompt_tokens_details.cache_write_tokens == cache_write + assert usage.cache_read_input_tokens == cache_read + assert usage.cache_creation_input_tokens == cache_write + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_cached_tokens_reported_only_in_the_prompt_tokens_details(replay): + cached = _async_payload() + cached["usage"]["prompt_tokens_details"] = {"cached_tokens": 4, "audio_tokens": 1} + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.usage.prompt_tokens_details.cached_tokens == 4 + assert rebuilt.usage.prompt_tokens_details.audio_tokens == 1 + + +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_keeps_integral_counts_cached_as_floats_or_numeric_strings(replay): + cached = _async_payload() + cached["usage"] = { + "prompt_tokens": 100.0, + "completion_tokens": "6", + "total_tokens": 106.0, + "cache_read_input_tokens": 90.0, + } + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + usage: Final = next(c.usage for c in chunks if getattr(c, "usage", None) is not None) + + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (100, 6, 106) + assert usage.cache_read_input_tokens == 90 + + +@pytest.mark.parametrize( + "field, value", + [ + ("prompt_tokens_details", {"cached_tokens": "abc"}), + ("prompt_tokens_details", [1, 2]), + ("cache_read_input_tokens", "abc"), + ("cache_creation_input_tokens", 1.5), + ], + ids=["details_text", "details_list", "read_text", "write_fraction"], +) +@pytest.mark.parametrize("replay", ["async", "sync"]) +@pytest.mark.asyncio +async def test_a_replayed_stream_skips_cached_prompt_cache_usage_that_does_not_validate(replay, field, value): + cached = _async_payload() + cached["usage"][field] = value + + chunks = await _collect_async(cached) if replay == "async" else list(convert_to_streaming_response(cached)) + rebuilt = litellm.stream_chunk_builder(chunks=chunks) + + assert rebuilt.choices[0].message.content == "Hello world! How are you?" + assert rebuilt.usage.prompt_tokens == 5 + + @pytest.mark.asyncio async def test_async_finish_reason_only_on_last_chunk(): chunks = await _collect_async(_async_payload()) diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py index b73df255cad..7400cd08dd3 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_streaming_iterator_combined_chunk.py @@ -18,6 +18,10 @@ import anthropic import httpx import pytest +from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response, + convert_to_streaming_response_async, +) from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import ( AnthropicStreamWrapper, _CombinedChunkSplitter, @@ -143,6 +147,43 @@ def test_delayed_usage_chunk_preserves_cache_tokens(): assert message_delta["usage"]["cache_creation_input_tokens"] == 20 +def _cached_bedrock_response(cache_read: int, cache_write: int) -> dict: + return { + "id": "chatcmpl-cached", + "object": "chat.completion", + "created": 1700000000, + "model": "us.anthropic.claude-opus-4-7", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "pong"}, "finish_reason": "stop"}], + "usage": Usage( + prompt_tokens=42145, + completion_tokens=6, + total_tokens=42151, + cache_read_input_tokens=cache_read, + cache_creation_input_tokens=cache_write, + ).model_dump(), + } + + +async def _replay_async(cached: dict) -> list[ModelResponseStream]: + return [chunk async for chunk in convert_to_streaming_response_async(cached)] + + +@pytest.mark.parametrize("cache_read, cache_write", [(42139, 0), (0, 42139)], ids=["read", "write"]) +@pytest.mark.parametrize("replay", ["async", "sync"]) +def test_a_cache_hit_replayed_as_a_messages_stream_keeps_the_prompt_cache_tokens(replay, cache_read, cache_write): + cached: Final = _cached_bedrock_response(cache_read, cache_write) + chunks: Final = ( + asyncio.run(_replay_async(cached)) if replay == "async" else list(convert_to_streaming_response(cached)) + ) + + events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="claude-opus-4-7")) + message_delta: Final = next(event for event in events if event.get("type") == "message_delta") + + assert message_delta["usage"]["input_tokens"] == 6 + assert message_delta["usage"].get("cache_read_input_tokens", 0) == cache_read + assert message_delta["usage"].get("cache_creation_input_tokens", 0) == cache_write + + def test_splitter_passes_through_non_combined_chunks(): """A chunk with content but no finish_reason is not split.""" chunk = ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(content="partial"), finish_reason=None)])