mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-24 00:52:24 +00:00
Merge pull request #42080 from BerriAI/litellm_prompt_caching_affinity_lookback
fix(router): keep prompt caching affinity when the breakpoint moves
This commit is contained in:
Mateo Wang
GitHub
commit
e8d97d381a
4 changed files with 482 additions and 138 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue