From 3ffe6272c96c08f54f972ef43a2541d73222f2ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 20:22:28 -0700 Subject: [PATCH] fix(router): hash the prompt caching affinity prefix off the event loop Offload the per-block hashing through offload_token_count on both the pre-call read and the success-event write, hash raw bytes as base64 instead of raising, drop the unused serialize_object helper, and bind the chained digest, the message envelope, and the bytes path in the regression tests --- litellm/constants.py | 2 - litellm/router_utils/prompt_caching_cache.py | 38 +++------------ .../test_router_prompt_caching.py | 48 ------------------- .../test_prompt_caching_deployment_check.py | 40 ++++++++++++++-- 4 files changed, 43 insertions(+), 85 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index e4576ad4d5c..215f25bccd1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -399,8 +399,6 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = ( if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT ) -# Anthropic checks at most 20 block positions behind a breakpoint for a cached prefix, a run of tool_use -# or tool_result blocks counting as one position, so deployment affinity probes the same window PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20 DEFAULT_TRIM_RATIO: Final = float( os.getenv("DEFAULT_TRIM_RATIO", 0.75) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 0b784e1fa91..78fc5e3fe6d 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -15,10 +15,8 @@ from typing_extensions import TypedDict from litellm.caching.caching import DualCache from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS -from litellm.litellm_core_utils.logging_utils import ( - truncate_base64_in_messages, - truncate_base64_in_messages_async, -) +from litellm.litellm_core_utils.logging_utils import truncate_base64_in_messages +from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam if TYPE_CHECKING: @@ -88,7 +86,9 @@ def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes: if tools is None: return hashlib.sha256(b"").digest() return hashlib.sha256( - _canonical_bytes(_TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True))) + _canonical_bytes( + _TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64")) + ) ).digest() @@ -132,23 +132,6 @@ class PromptCachingCache: def __init__(self, cache: DualCache): self.cache = cache - @staticmethod - def serialize_object(obj: Any) -> object: - """Helper function to serialize Pydantic objects, dictionaries, or fallback to string.""" - if hasattr(obj, "dict"): - # If the object is a Pydantic model, use its `dict()` method - return obj.dict() - elif isinstance(obj, dict): - # If the object is a dictionary, serialize it with sorted keys - return json.dumps(obj, sort_keys=True, separators=(",", ":")) # Standardize serialization - - elif isinstance(obj, list): - # Serialize lists by ensuring each element is handled properly - return [PromptCachingCache.serialize_object(item) for item in obj] - elif isinstance(obj, (int, float, bool)): - return obj # Keep primitive types as-is - return str(obj) - @staticmethod def extract_cacheable_prefix( messages: list[AllMessageValues], @@ -263,6 +246,7 @@ class PromptCachingCache: to_jsonable_python( truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)), serialize_unknown=True, + bytes_mode="base64", ) ), tools, @@ -275,15 +259,7 @@ class PromptCachingCache: ) -> tuple[PrefixPosition, ...]: if not messages: return () - return _positions_of( - _PREFIX_ADAPTER.validate_python( - to_jsonable_python( - await truncate_base64_in_messages_async(PromptCachingCache.extract_cacheable_prefix(messages)), - serialize_unknown=True, - ) - ), - tools, - ) + return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools) @staticmethod def get_prompt_caching_cache_key( diff --git a/tests/router_unit_tests/test_router_prompt_caching.py b/tests/router_unit_tests/test_router_prompt_caching.py index 5c36c30e818..879264ca502 100644 --- a/tests/router_unit_tests/test_router_prompt_caching.py +++ b/tests/router_unit_tests/test_router_prompt_caching.py @@ -11,57 +11,9 @@ from unittest.mock import patch, MagicMock, AsyncMock from create_mock_standard_logging_payload import create_standard_logging_payload from litellm.types.utils import StandardLoggingPayload import unittest -from pydantic import BaseModel from litellm.router_utils.prompt_caching_cache import PromptCachingCache -class ExampleModel(BaseModel): - field1: str - field2: int - - -def test_serialize_pydantic_object(): - model = ExampleModel(field1="value", field2=42) - serialized = PromptCachingCache.serialize_object(model) - assert serialized == {"field1": "value", "field2": 42} - - -def test_serialize_dict(): - obj = {"b": 2, "a": 1} - serialized = PromptCachingCache.serialize_object(obj) - assert serialized == '{"a":1,"b":2}' # JSON string with sorted keys - - -def test_serialize_nested_dict(): - obj = {"z": {"b": 2, "a": 1}, "x": [1, 2, {"c": 3}]} - serialized = PromptCachingCache.serialize_object(obj) - expected = '{"x":[1,2,{"c":3}],"z":{"a":1,"b":2}}' # JSON string with sorted keys - assert serialized == expected - - -def test_serialize_list(): - obj = ["item1", {"a": 1, "b": 2}, 42] - serialized = PromptCachingCache.serialize_object(obj) - expected = ["item1", '{"a":1,"b":2}', 42] - assert serialized == expected - - -def test_serialize_fallback(): - obj = 12345 # Simple non-serializable object - serialized = PromptCachingCache.serialize_object(obj) - assert serialized == 12345 - - -def test_serialize_non_serializable(): - class CustomClass: - def __str__(self): - return "custom_object" - - obj = CustomClass() - serialized = PromptCachingCache.serialize_object(obj) - assert serialized == "custom_object" # Fallback to string conversion - - @pytest.mark.asyncio async def test_router_prompt_caching_same_cacheable_prefix_routes_to_same_deployment(): """ diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index d0a9223dfa7..ad92f442a6e 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -708,7 +708,10 @@ async def test_a_run_of_tool_blocks_counts_as_one_lookback_position(): @pytest.mark.asyncio async def test_an_edited_earlier_block_does_not_inherit_the_pin(): - """Walking back must still bind every block's content, or an edited conversation pins to a stale cache.""" + """ + Every key must bind the whole prefix before its block, not the block alone, or a conversation + that repeats a pinned block after an edit walks back onto a cache the provider no longer holds. + """ prompt_cache = PromptCachingCache(cache=DualCache()) await prompt_cache.async_add_model_id( model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("original")]}), tools=None @@ -716,12 +719,41 @@ async def test_an_edited_earlier_block_does_not_inherit_the_pin(): edited = _turn( {"role": "user", "content": [_text("edited")]}, {"role": "assistant", "content": "ok"}, - {"role": "user", "content": [_marked("next")]}, + {"role": "user", "content": [_marked("original")]}, ) assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None +@pytest.mark.asyncio +async def test_swapped_roles_do_not_inherit_the_pin(): + """The message envelope is part of what the provider caches, so the same blocks under other roles key apart.""" + prompt_cache = PromptCachingCache(cache=DualCache()) + pinned = _turn( + {"role": "user", "content": [_text("question")]}, + {"role": "assistant", "content": [_marked("answer")]}, + ) + swapped = _turn( + {"role": "assistant", "content": [_text("question")]}, + {"role": "user", "content": [_marked("answer")]}, + ) + await prompt_cache.async_add_model_id(model_id="dep-1", messages=pinned, tools=None) + + assert await prompt_cache.async_get_model_id(messages=pinned, tools=None) == {"model_id": "dep-1"} + assert await prompt_cache.async_get_model_id(messages=swapped, tools=None) is None + + +@pytest.mark.asyncio +async def test_raw_bytes_in_a_block_hash_instead_of_failing_the_request(): + """A block carrying raw bytes must key like any other block rather than raising out of the router filter.""" + prompt_cache = PromptCachingCache(cache=DualCache()) + binary_block = {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": b"\xff\xfe"}} + turn = _turn({"role": "user", "content": [binary_block, _marked("describe")]}) + await prompt_cache.async_add_model_id(model_id="dep-1", messages=turn, tools=None) + + assert await prompt_cache.async_get_model_id(messages=turn, tools=None) == {"model_id": "dep-1"} + + class _BrokenBatchReadCache(DualCache): async def async_batch_get_cache(self, keys, parent_otel_span=None, local_only=False, **kwargs): return None @@ -786,11 +818,11 @@ async def test_claude_code_style_session_stays_on_one_deployment_across_turns(lo "litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"}, "model_info": {"id": model_id}, } - for model_id in ("dep-1", "dep-2", "dep-3") + for model_id in (f"dep-{number}" for number in range(1, 7)) ], optional_pre_call_checks=["prompt_caching"], ) - user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 6))] + user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 9))] history: List[AllMessageValues] = [] served: List[str] = [] for text in user_turns: