scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(router): keep prompt caching affinity when the breakpoint moves

Browse source 
The prompt_caching pre-call check keyed a deployment pin on a hash of the
whole cacheable prefix, cache_control markers included. Agent clients
such as Claude Code move the marker to the newest user turn on every
request, so the key changed every turn, the pin never matched, and a
multi-turn session drifted across deployments and lost its provider
cache.

Hash the prefix per content block with the markers stripped, chained so
every block position has a key, and write the pin at the breakpoint
block. Lookup walks back over the last PROMPT_CACHE_LOOKBACK_POSITIONS
positions (a run of tool_use or tool_result blocks counting as one), the
same window the provider probes for a cached prefix, in one batch cache
read. Both sides hash the prefix after base64 truncation so a request
carrying raw image bytes derives the keys the success event stored.
This commit is contained in:
mateo-berri 2026-09-19 19:52:22 -07:00
parent d1773d96e9
commit 517fff5bb7
3 changed files with 450 additions and 76 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/constants.py
View file

@ -399,6 +399,9 @@ 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)
) # default ratio of tokens to trim from the end of a prompt

250
litellm/router_utils/prompt_caching_cache.py
View file

@ -4,12 +4,21 @@ 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,
truncate_base64_in_messages_async,
)
from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam
if TYPE_CHECKING:
@ -28,10 +37,100 @@ 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)))
).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:
@ -140,114 +239,123 @@ 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,
)
),
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 _positions_of(
_PREFIX_ADAPTER.validate_python(
to_jsonable_python(
await truncate_base64_in_messages_async(PromptCachingCache.extract_cacheable_prefix(messages)),
serialize_unknown=True,
)
),
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
)
)
)

273
tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py
View file

@ -1,5 +1,6 @@
import asyncio
import copy
import functools
from typing import List, cast
import pytest
@ -7,7 +8,7 @@ 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 (
@ -30,7 +31,6 @@ def _local_model_cost_map_autouse(local_model_cost_map):
yield
def _deployments(*models: str) -> List[dict]:
return [
{
@ -84,7 +84,9 @@ def test_write_gate_is_what_prevents_a_pin_below_the_model_minimum():
"""
messages = _messages(word_count=1400)
token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True)
token_count = token_counter(
messages=messages, model="anthropic/claude-opus-4-5", use_default_image_token_count=True
)
assert 1024 < token_count < 4096
assert is_prompt_caching_valid_prompt(model="anthropic/claude-opus-4-5", messages=messages) is False
@ -110,7 +112,9 @@ async def test_async_filter_deployments_does_not_narrow_prompt_below_model_minim
deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
messages = _messages(word_count=1400)
token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True)
token_count = token_counter(
messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True
)
assert DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT < token_count < OPUS_4_6_MIN_TOKENS
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
@ -136,7 +140,9 @@ async def test_async_filter_deployments_narrows_prompt_above_model_minimum():
deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6")
messages = _messages(word_count=5000)
token_count = token_counter(messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True)
token_count = token_counter(
messages=messages, model="anthropic/claude-opus-4-6", use_default_image_token_count=True
)
assert token_count > OPUS_4_6_MIN_TOKENS
await PromptCachingCache(cache=cache).async_add_model_id(model_id="dep-2", messages=messages, tools=None)
@ -539,3 +545,260 @@ 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():
"""Walking back must still bind every block's content, or an edited conversation pins to a stale cache."""
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("next")]},
)
assert await prompt_cache.async_get_model_id(messages=edited, tools=None) is None
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 ("dep-1", "dep-2", "dep-3")
],
optional_pre_call_checks=["prompt_caching"],
)
user_turns = [LONG_PROMPT, *(f"follow-up {number}" for number in range(1, 6))]
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)