diff --git a/litellm/constants.py b/litellm/constants.py index 95ccad93284..72495b389d7 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -402,6 +402,7 @@ MINIMUM_PROMPT_CACHE_TOKEN_COUNT: Final = ( if MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE is not None else DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT ) +PROMPT_CACHE_LOOKBACK_POSITIONS: Final = 20 DEFAULT_TRIM_RATIO: Final = float( os.getenv("DEFAULT_TRIM_RATIO", 0.75) ) # default ratio of tokens to trim from the end of a prompt diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 39708e168f5..78fc5e3fe6d 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -4,12 +4,19 @@ Wrapper around router cache. Meant to store model id when prompt caching support import hashlib import json +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from itertools import accumulate from typing import TYPE_CHECKING, Any, Final, cast +from pydantic import JsonValue, TypeAdapter +from pydantic_core import to_jsonable_python from typing_extensions import TypedDict from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache +from litellm.constants import PROMPT_CACHE_LOOKBACK_POSITIONS +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: @@ -28,27 +35,102 @@ class PromptCachingCacheValue(TypedDict): model_id: str +PROMPT_CACHE_PIN_TTL_SECONDS: Final = 300 +_TOOL_RUN_BLOCK_TYPES: Final = frozenset({"tool_use", "tool_result"}) +_PREFIX_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) +_TOOLS_ADAPTER: Final = TypeAdapter(tuple[JsonValue, ...]) +_PINS_ADAPTER: Final[TypeAdapter[tuple[JsonValue, ...] | None]] = TypeAdapter(tuple[JsonValue, ...] | None) + + +@dataclass(frozen=True, slots=True) +class PrefixPosition: + cache_key: str + position: int + + +def _sorted_pairs(pairs: Iterable[tuple[str, JsonValue]]) -> tuple[tuple[str, JsonValue], ...]: + return tuple(sorted(pairs, key=lambda pair: pair[0])) + + +def _canonical_bytes(value: object) -> bytes: + return json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + + +def _block_unit( + envelope: tuple[tuple[str, JsonValue], ...], message_run_type: str | None, block: JsonValue +) -> tuple[bytes, str | None]: + if not isinstance(block, dict): + return _canonical_bytes((envelope, block)), message_run_type + block_type: Final = block.get("type") + block_run_type: Final = block_type if isinstance(block_type, str) and block_type in _TOOL_RUN_BLOCK_TYPES else None + stripped: Final = _sorted_pairs(item for item in block.items() if item[0] != "cache_control") + return _canonical_bytes((envelope, stripped)), message_run_type or block_run_type + + +def _message_units(message: Mapping[str, JsonValue]) -> tuple[tuple[bytes, str | None], ...]: + envelope: Final = _sorted_pairs(item for item in message.items() if item[0] not in ("content", "cache_control")) + message_run_type: Final = "tool_result" if message.get("role") == "tool" else None + content: Final = message.get("content") + if isinstance(content, list) and content: + return tuple(_block_unit(envelope, message_run_type, block) for block in content) + if isinstance(content, str) and content: + return ((_canonical_bytes((envelope, (("text", content), ("type", "text")))), message_run_type),) + return ((_canonical_bytes((envelope, None)), message_run_type),) + + +def _chain_digest(digest: bytes, unit: bytes) -> bytes: + return hashlib.sha256(digest + unit).digest() + + +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, bytes_mode="base64")) + ) + ).digest() + + +def _positions_of( + prefix: tuple[Mapping[str, JsonValue], ...], tools: Sequence[ChatCompletionToolParam] | None +) -> tuple[PrefixPosition, ...]: + units: Final = tuple(unit for message in prefix for unit in _message_units(message)) + digests: Final = tuple(accumulate((unit_bytes for unit_bytes, _ in units), _chain_digest, initial=_seed(tools)))[1:] + run_types: Final = tuple(run_type for _, run_type in units) + positions: Final = accumulate( + 0 if run_type is not None and run_type == previous else 1 + for run_type, previous in zip(run_types, (None, *run_types[:-1])) + ) + return tuple( + PrefixPosition(cache_key=f"deployment:{digest.hex()}:prompt_caching", position=position) + for digest, position in zip(digests, positions) + ) + + +def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]: + if not positions: + return () + oldest_probed_position: Final = positions[-1].position - PROMPT_CACHE_LOOKBACK_POSITIONS + return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position) + + +def _pinned_value(value: JsonValue) -> PromptCachingCacheValue | None: + if not isinstance(value, dict): + return None + model_id: Final = value.get("model_id") + return PromptCachingCacheValue(model_id=model_id) if isinstance(model_id, str) else None + + +def _first_pin(values: tuple[JsonValue, ...] | None) -> PromptCachingCacheValue | None: + if values is None: + return None + return next((pin for pin in map(_pinned_value, values) if pin is not None), None) + + class PromptCachingCache: def __init__(self, cache: DualCache): self.cache = cache - self.in_memory_cache = InMemoryCache() - - @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( @@ -140,114 +222,116 @@ class PromptCachingCache: return cacheable_prefix @staticmethod - def get_prompt_caching_cache_key( + def prefix_positions( messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, - ) -> str | None: - if messages is None and tools is None: - return None + tools: Sequence[ChatCompletionToolParam] | None, + ) -> tuple[PrefixPosition, ...]: + """ + One cache key per content block of the cacheable prefix, oldest block first. - # Extract cacheable prefix from messages (only include up to last cache_control block) - cacheable_messages = None - if messages is not None: - cacheable_messages = PromptCachingCache.extract_cacheable_prefix(messages) - # If no cacheable prefix found, return None (can't cache) - if not cacheable_messages: - return None + Each key hashes the prefix content up to and including that block, with cache_control markers + left out, so the key of a block is the same whichever turn's breakpoint the prefix ends at. + String content hashes like a single text block, which is how the provider treats it and how + Claude Code re-sends a previously marked message. `position` counts a run of consecutive + tool_use (or tool_result) blocks as one, matching the provider's lookback window. - # Use serialize_object for consistent and stable serialization - data_to_hash: Final = {} - if cacheable_messages is not None: - serialized_messages: Final = PromptCachingCache.serialize_object(cacheable_messages) - data_to_hash["messages"] = serialized_messages - if tools is not None: - serialized_tools: Final = PromptCachingCache.serialize_object(tools) - data_to_hash["tools"] = serialized_tools - - # Combine serialized data into a single string - data_to_hash_str: Final = json.dumps( - data_to_hash, - sort_keys=True, - separators=(",", ":"), + The prefix is hashed in the shape the success event sees it, with long base64 data URIs + already replaced by their size placeholder, so a request carrying the raw image bytes + derives the same keys the write side stored. + """ + if not messages: + return () + return _positions_of( + _PREFIX_ADAPTER.validate_python( + to_jsonable_python( + truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)), + serialize_unknown=True, + bytes_mode="base64", + ) + ), + tools, ) - # Create a hash of the serialized data for a stable cache key - hashed_data: Final = hashlib.sha256(data_to_hash_str.encode()).hexdigest() - return f"deployment:{hashed_data}:prompt_caching" + @staticmethod + async def async_prefix_positions( + messages: list[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + ) -> tuple[PrefixPosition, ...]: + if not messages: + return () + return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools) + + @staticmethod + def get_prompt_caching_cache_key( + messages: list[AllMessageValues] | None, + tools: Sequence[ChatCompletionToolParam] | None, + ) -> str | None: + positions: Final = PromptCachingCache.prefix_positions(messages, tools) + return positions[-1].cache_key if positions else None def add_model_id( self, model_id: str, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - if messages is None and tools is None: - return - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - # If no cacheable prefix found, don't cache (can't generate cache key) if cache_key is None: return - self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300) - return + self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=PROMPT_CACHE_PIN_TTL_SECONDS) async def async_add_model_id( self, model_id: str, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> None: - if messages is None and tools is None: - return - - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - # If no cacheable prefix found, don't cache (can't generate cache key) - if cache_key is None: + positions: Final = await PromptCachingCache.async_prefix_positions(messages, tools) + if not positions: return await self.cache.async_set_cache( - cache_key, + positions[-1].cache_key, PromptCachingCacheValue(model_id=model_id), - ttl=300, # store for 5 minutes + ttl=PROMPT_CACHE_PIN_TTL_SECONDS, ) - return async def async_get_model_id( self, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> PromptCachingCacheValue | None: """ - Get model ID from cache using the cacheable prefix. - - The cache key is based on the cacheable prefix (everything up to and including - the last cache_control block), so requests with the same cacheable prefix but - different user messages will have the same cache key. + Find the deployment that last served this prefix, walking back from the breakpoint the + same way the provider cache does, so a breakpoint that moved forward since the last + turn still lands on the deployment whose cache holds the earlier prefix. """ - if messages is None and tools is None: + cache_keys: Final = _lookback_keys(await PromptCachingCache.async_prefix_positions(messages, tools)) + if not cache_keys: return None - # Generate cache key using cacheable prefix - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - if cache_key is None: - return None - - # Perform cache lookup - cache_result: Final = await self.cache.async_get_cache(key=cache_key) - return cache_result + return _first_pin( + _PINS_ADAPTER.validate_python( + await self.cache.async_batch_get_cache( + keys=list(cache_keys), # mutable-ok: DualCache.async_batch_get_cache only takes a list + ) + ) + ) def get_model_id( self, messages: list[AllMessageValues] | None, - tools: list[ChatCompletionToolParam] | None, + tools: Sequence[ChatCompletionToolParam] | None, ) -> PromptCachingCacheValue | None: - if messages is None and tools is None: + cache_keys: Final = _lookback_keys(PromptCachingCache.prefix_positions(messages, tools)) + if not cache_keys: return None - cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, tools) - # If no cacheable prefix found, return None (can't cache) - if cache_key is None: - return None - - return self.cache.get_cache(cache_key) + return _first_pin( + _PINS_ADAPTER.validate_python( + self.cache.batch_get_cache( + keys=list(cache_keys), # mutable-ok: DualCache.batch_get_cache only takes a list + ) + ) + ) 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/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index 323abd68c59..a7006c62438 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -1,12 +1,13 @@ import asyncio import copy -from typing import cast +import functools +from typing import Final, cast import pytest import litellm from litellm.caching.dual_cache import DualCache -from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT +from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( @@ -19,6 +20,23 @@ from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_p MODEL_GROUP_ALIAS = "my-claude-group" OPUS_4_6_MIN_TOKENS = 4096 +CALLBACK_REGISTRIES: Final = ( + "input_callback", + "success_callback", + "failure_callback", + "_async_success_callback", + "_async_failure_callback", + "callbacks", +) + + +@pytest.fixture(autouse=True) +def _fresh_callback_registries(monkeypatch): + """`litellm.logging_callback_manager` keeps one callback per class, so a + `PromptCachingDeploymentCheck` or `_SentMessagesCapture` left behind by an + earlier test would swallow the next test's success events.""" + for registry in CALLBACK_REGISTRIES: + monkeypatch.setattr(litellm, registry, []) @pytest.fixture @@ -604,3 +622,292 @@ async def test_async_log_success_event_counts_the_prompt_off_the_event_loop(): "model_id": "dep-1" } assert_loop_stayed_free(took, lags) + + +LONG_PROMPT = "word " * 3000 +ONE_PIXEL_PNG = ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +def _turn(*messages: dict) -> list[AllMessageValues]: + return cast(list[AllMessageValues], list(messages)) + + +def _text(text: str) -> dict: + return {"type": "text", "text": text} + + +def _marked(text: str) -> dict: + return {"type": "text", "text": text, "cache_control": {"type": "ephemeral"}} + + +@pytest.mark.asyncio +async def test_pin_survives_the_breakpoint_moving_to_the_next_turn(): + """ + The regression. Claude Code marks only the newest user message each turn, so the last breakpoint + moves forward every turn. The key hashed the prefix up to that moving breakpoint, markers + included, so no turn after the first ever found the pin the previous turn wrote, and a + multi-deployment group re-rolled the deployment mid-session, paying a cache write on a + deployment whose provider cache held nothing of the conversation. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + turn_one = _turn({"role": "user", "content": [_marked(LONG_PROMPT)]}) + turn_two = _turn( + {"role": "user", "content": [_text(LONG_PROMPT)]}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": [_marked("next")]}, + ) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=turn_one, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[1]] + + +@pytest.mark.asyncio +async def test_pin_survives_the_marked_message_coming_back_as_string_content(): + """ + Claude Code sends the message that carries a breakpoint as a one-block content list and re-sends + it next turn as plain string content once the marker has moved on. The provider caches both + shapes identically, so the key has to as well, or the walk-back never lands on the turn-one write. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + turn_one = _turn( + {"role": "system", "content": [_marked(LONG_PROMPT)]}, + {"role": "user", "content": [_marked("hello")]}, + ) + turn_two = _turn( + {"role": "system", "content": LONG_PROMPT}, + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": "hi"}, + {"role": "user", "content": [_marked("again")]}, + ) + + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-1", messages=turn_one, tools=None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[0]] + + +@pytest.mark.asyncio +async def test_lookback_stops_where_the_provider_cache_stops(): + """ + Anthropic finds a cached prefix at most PROMPT_CACHE_LOOKBACK_POSITIONS block positions behind a + breakpoint, the breakpoint block included. Probing further would pin to a deployment whose cache + the provider will not consult, and probing less would drop pins the provider still honors. + """ + prompt_cache = PromptCachingCache(cache=DualCache()) + await prompt_cache.async_add_model_id( + model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("block 0")]}), tools=None + ) + + def turn_with_blocks_after(count: int) -> list[AllMessageValues]: + later = [_text(f"block {index}") for index in range(1, count)] + [_marked(f"block {count}")] + return _turn({"role": "user", "content": [_text("block 0"), *later]}) + + inside_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS - 1) + past_window = turn_with_blocks_after(PROMPT_CACHE_LOOKBACK_POSITIONS) + + assert await prompt_cache.async_get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"} + assert prompt_cache.get_model_id(messages=inside_window, tools=None) == {"model_id": "dep-1"} + assert await prompt_cache.async_get_model_id(messages=past_window, tools=None) is None + assert prompt_cache.get_model_id(messages=past_window, tools=None) is None + + +@pytest.mark.asyncio +async def test_a_run_of_tool_blocks_counts_as_one_lookback_position(): + """ + The provider counts consecutive tool_use blocks as one lookback position, and consecutive + tool_result blocks as one, in both the Anthropic and the OpenAI message shapes. An agent turn that + fans out into many tool calls would otherwise push the previous breakpoint out of the window + after a single turn, which is exactly when the conversation is longest and the cache matters most. + """ + prompt_cache = PromptCachingCache(cache=DualCache()) + await prompt_cache.async_add_model_id( + model_id="dep-1", messages=_turn({"role": "user", "content": [_marked("task")]}), tools=None + ) + fan_out = PROMPT_CACHE_LOOKBACK_POSITIONS + 5 + + def anthropic_shaped(tool_use_type: str, tool_result_type: str) -> list[AllMessageValues]: + return _turn( + {"role": "user", "content": [_text("task")]}, + { + "role": "assistant", + "content": [ + {"type": tool_use_type, "id": f"call-{index}", "name": "read", "input": {"index": index}} + for index in range(fan_out) + ], + }, + { + "role": "user", + "content": [ + *( + {"type": tool_result_type, "tool_use_id": f"call-{index}", "content": "ok"} + for index in range(fan_out) + ), + _marked("continue"), + ], + }, + ) + + openai_shaped = _turn( + {"role": "user", "content": [_text("task")]}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": f"call-{index}", "type": "function", "function": {"name": "read", "arguments": "{}"}} + for index in range(fan_out) + ], + }, + *({"role": "tool", "tool_call_id": f"call-{index}", "content": "ok"} for index in range(fan_out)), + {"role": "user", "content": [_marked("continue")]}, + ) + + assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("tool_use", "tool_result"), tools=None) == { + "model_id": "dep-1" + } + assert await prompt_cache.async_get_model_id(messages=openai_shaped, tools=None) == {"model_id": "dep-1"} + assert await prompt_cache.async_get_model_id(messages=anthropic_shaped("text", "text"), tools=None) is None + + +@pytest.mark.asyncio +async def test_an_edited_earlier_block_does_not_inherit_the_pin(): + """ + 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 + ) + edited = _turn( + {"role": "user", "content": [_text("edited")]}, + {"role": "assistant", "content": "ok"}, + {"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 + + +@pytest.mark.asyncio +async def test_a_failed_batch_read_pins_nothing(): + """DualCache answers None rather than a list when the batch read raises, and routing must fall through.""" + prompt_cache = PromptCachingCache(cache=_BrokenBatchReadCache()) + + assert ( + await prompt_cache.async_get_model_id(messages=_turn({"role": "user", "content": [_marked("x")]}), tools=None) + is None + ) + + +@pytest.mark.asyncio +async def test_pin_matches_when_the_success_event_truncated_an_image_payload(monkeypatch, local_model_cost_map): + """ + The success event only ever sees the standard logging payload, whose long base64 data URIs are + replaced by size placeholders, while routing sees the raw request. Hashing the raw bytes on the + read side would key every image-carrying session past its own pin. + """ + capture = _SentMessagesCapture() + monkeypatch.setattr(litellm, "callbacks", [capture]) + image = {"type": "image_url", "image_url": {"url": ONE_PIXEL_PNG}} + turn_one = _turn({"role": "user", "content": [image, _marked(LONG_PROMPT)]}) + + await litellm.acompletion( + model=AUTO_CACHING_MODEL, messages=copy.deepcopy(turn_one), mock_response="ok", api_key="sk-fake" + ) + logged = await _eventually(lambda: capture.messages) + assert logged is not None + assert logged != turn_one + + cache = DualCache() + await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=logged, tools=None) + turn_two = _turn( + {"role": "user", "content": [image, _text(LONG_PROMPT)]}, + {"role": "assistant", "content": "ok"}, + {"role": "user", "content": [_marked("next")]}, + ) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + + filtered = await PromptCachingDeploymentCheck(cache=cache).async_filter_deployments( + model=MODEL_GROUP_ALIAS, healthy_deployments=deployments, messages=turn_two + ) + + assert filtered == [deployments[1]] + + +@pytest.mark.asyncio +async def test_claude_code_style_session_stays_on_one_deployment_across_turns(local_model_cost_map): + """ + End to end over the router with a client that marks only the newest user message each turn, the + way Claude Code does. Every turn has to land on the deployment that served the first one. + """ + router = litellm.Router( + model_list=[ + { + "model_name": MODEL_GROUP_ALIAS, + "litellm_params": {"model": AUTO_CACHING_MODEL, "api_key": "sk-fake"}, + "model_info": {"id": model_id}, + } + 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, 9))] + history: list[AllMessageValues] = [] + served: list[str] = [] + for text in user_turns: + request = cast(list[AllMessageValues], [*history, {"role": "user", "content": [_marked(text)]}]) + response = await router.acompletion(model=MODEL_GROUP_ALIAS, messages=request, mock_response="ok") + served.append(response._hidden_params["model_id"]) + pin_key = PromptCachingCache.get_prompt_caching_cache_key(request, None) + assert await _eventually(functools.partial(router.cache.get_cache, key=pin_key)) is not None + history = [*history, {"role": "user", "content": [_text(text)]}, {"role": "assistant", "content": "ok"}] + + assert served == [served[0]] * len(user_turns)