diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 0589e290b47..a628e0194a8 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -74,7 +74,7 @@ class PromptCachingDeploymentCheck(CustomLogger): ## AUTO PROMPT CACHING - the breakpoints this request will carry are injected inside ## `litellm.acompletion`, after a deployment has been picked, so the affinity key has to ## be derived from the messages as they will be sent, not as they arrive here. - affinity_messages: Final = AnthropicCacheControlHook.messages_with_default_injections( + injected_messages: Final = AnthropicCacheControlHook.messages_with_default_injections( messages=cast(list[AllMessageValues], messages), models=( deployment["litellm_params"]["model"] @@ -93,10 +93,20 @@ class PromptCachingDeploymentCheck(CustomLogger): ), request_kwargs=request_kwargs, ) + affinity_messages: Final = PromptCachingCache.prepend_system_prompt( + injected_messages, + request_kwargs.get("system") if request_kwargs is not None else None, + ) model_id_dict: Final = await prompt_cache.async_get_model_id( messages=affinity_messages, - tools=None, + tools=( + cast( # cast-ok: request kwargs are untyped + list[AllToolParamValues] | None, request_kwargs.get("tools") + ) # cast-ok: request kwargs are untyped + if request_kwargs is not None + else None + ), ) if model_id_dict is not None: model_id: Final = model_id_dict["model_id"] @@ -139,18 +149,25 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + logged_messages: Final = PromptCachingCache.prepend_system_prompt( + cast(list[AllMessageValues], messages), # cast-ok: standard logging payload is partially typed + kwargs.get("system"), + ) + ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, - messages=cast(list[AllMessageValues], messages), + messages=logged_messages, ): cache: Final = PromptCachingCache( cache=self.cache, ) await cache.async_add_model_id( model_id=model_id, - messages=messages, - tools=None, # [TODO]: add tools once standard_logging_object supports it + messages=logged_messages, + tools=cast( # cast-ok: callback kwargs are untyped + list[AllToolParamValues] | None, kwargs.get("tools") + ), ) return diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 39708e168f5..77ee9a6d386 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -10,7 +10,7 @@ from typing_extensions import TypedDict from litellm.caching.caching import DualCache from litellm.caching.in_memory_cache import InMemoryCache -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam +from litellm.types.llms.openai import AllMessageValues, ChatCompletionCachedContent, ChatCompletionToolParam if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -139,29 +139,119 @@ class PromptCachingCache: return cacheable_prefix + @staticmethod + def extract_cacheable_tools( + tools: list[ChatCompletionToolParam], + ) -> list[ChatCompletionToolParam]: + cacheable_tool_index: Final = next( + ( + index + for index in range(len(tools) - 1, -1, -1) + if isinstance(cache_control := tools[index].get("cache_control"), dict) + and cache_control.get("type") == "ephemeral" + ), + None, + ) + # A breakpoint truncates trailing, uncached tools out of the key. With no breakpoint, + # every tool still precedes whatever the real provider prefix caches on, so the full + # list has to stay in the key or distinct tool sets would collide on the same key. + return tools[: cacheable_tool_index + 1] if cacheable_tool_index is not None else tools[:] + + @staticmethod + def prepend_system_prompt( + messages: list[AllMessageValues], + system: object | None, + ) -> list[AllMessageValues]: + if system is None: + return messages + # Standard logging already prepends a string system prompt onto `messages` before this + # runs, so re-prepending here would double it up and produce a different affinity key + # than the one computed from the raw request messages. + if messages and messages[0].get("role") == "system" and messages[0].get("content") == system: + return messages + return cast( # cast-ok: system content is validated by the provider payload + list[AllMessageValues], + [{"role": "system", "content": system}, *messages], # mutable-ok: cast target requires a concrete list + ) + + @staticmethod + def get_prompt_caching_ttl( + messages: list[AllMessageValues] | None, + tools: list[ChatCompletionToolParam] | None = None, + ) -> int: + cacheable_prefix: Final = ( + PromptCachingCache.extract_cacheable_prefix(messages) + if messages is not None + else [] # mutable-ok: TTL helper requires a concrete list + ) + return PromptCachingCache.get_prompt_caching_ttl_from_prefix(cacheable_prefix, tools) + + @staticmethod + def _ephemeral_cache_controls( + message: AllMessageValues, + ) -> tuple[ChatCompletionCachedContent | dict[str, object], ...]: + content: Final = message.get("content") + content_cache_controls: Final = ( + tuple(content_block.get("cache_control") for content_block in content if isinstance(content_block, dict)) + if isinstance(content, list) + else () + ) + return tuple( + cache_control + for cache_control in (message.get("cache_control"), *content_cache_controls) + if isinstance(cache_control, dict) and cache_control.get("type") == "ephemeral" + ) + + @staticmethod + def get_prompt_caching_ttl_from_prefix( + cacheable_prefix: list[AllMessageValues], + tools: list[ChatCompletionToolParam] | None, + ) -> int: + cacheable_tools: Final = PromptCachingCache.extract_cacheable_tools( + tools or [] # mutable-ok: tool API requires a concrete list + ) + cache_control_values: Final = tuple( + cache_control + for message in cacheable_prefix + for cache_control in PromptCachingCache._ephemeral_cache_controls(message) + ) + tuple( + cache_control + for tool in cacheable_tools + if isinstance(cache_control := tool.get("cache_control"), dict) and cache_control.get("type") == "ephemeral" + ) + # Prefer the shortest provider lifetime so affinity never outlives a cached segment + return 3600 if cache_control_values and all(value.get("ttl") == "1h" for value in cache_control_values) else 300 + @staticmethod def get_prompt_caching_cache_key( messages: list[AllMessageValues] | None, tools: list[ChatCompletionToolParam] | None, ) -> str | None: - if messages is None and tools is None: + cacheable_messages: Final = ( + PromptCachingCache.extract_cacheable_prefix(messages) if messages is not None else None + ) + return PromptCachingCache.get_prompt_caching_cache_key_from_prefix(cacheable_messages, tools) + + @staticmethod + def get_prompt_caching_cache_key_from_prefix( + cacheable_messages: list[AllMessageValues] | None, + tools: list[ChatCompletionToolParam] | None, + ) -> str | None: + if cacheable_messages is None and tools is None: + return None + if cacheable_messages is not None and not cacheable_messages: return None - # 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 - + # Only exact cached-prefix matches can establish deployment affinity. + # Partial matches must fall back to normal routing. # 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) + cacheable_tools: Final = PromptCachingCache.extract_cacheable_tools(tools) + serialized_tools: Final = PromptCachingCache.serialize_object(cacheable_tools) data_to_hash["tools"] = serialized_tools # Combine serialized data into a single string @@ -184,12 +274,21 @@ class PromptCachingCache: 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) + cacheable_prefix: Final = ( + PromptCachingCache.extract_cacheable_prefix(messages) if messages is not None else None + ) + cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key_from_prefix(cacheable_prefix, tools) if cache_key is None: return - self.cache.set_cache(cache_key, PromptCachingCacheValue(model_id=model_id), ttl=300) + self.cache.set_cache( + cache_key, + PromptCachingCacheValue(model_id=model_id), + ttl=PromptCachingCache.get_prompt_caching_ttl_from_prefix( + cacheable_prefix or [], # mutable-ok: TTL helper requires a concrete list + tools, + ), + ) return async def async_add_model_id( @@ -201,15 +300,20 @@ class PromptCachingCache: 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) + cacheable_prefix: Final = ( + PromptCachingCache.extract_cacheable_prefix(messages) if messages is not None else None + ) + cache_key: Final = PromptCachingCache.get_prompt_caching_cache_key_from_prefix(cacheable_prefix, tools) if cache_key is None: return await self.cache.async_set_cache( cache_key, PromptCachingCacheValue(model_id=model_id), - ttl=300, # store for 5 minutes + ttl=PromptCachingCache.get_prompt_caching_ttl_from_prefix( + cacheable_prefix or [], # mutable-ok: TTL helper requires a concrete list + tools, + ), ) return 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 333e7b2ff31..7fb7907c230 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 @@ -1,6 +1,7 @@ import asyncio import copy from typing import List, cast +from unittest.mock import AsyncMock, Mock import pytest @@ -15,7 +16,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import _get_min_token_count_for_deployments, ) from litellm.router_utils.prompt_caching_cache import PromptCachingCache -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt, token_counter MODEL_GROUP_ALIAS = "my-claude-group" @@ -60,6 +61,222 @@ def _messages(word_count: int) -> List[AllMessageValues]: ) +def test_prepend_system_prompt_without_system_preserves_messages(): + messages = _messages(word_count=10) + + assert PromptCachingCache.prepend_system_prompt(messages, None) is messages + + +def test_tool_affinity_ignores_tools_after_cache_breakpoint(): + messages = _messages(word_count=10) + cached_tools = cast( + list[ChatCompletionToolParam], + [ + {"type": "function", "function": {"name": "stable", "parameters": {}}, "cache_control": {"type": "ephemeral"}}, + {"type": "function", "function": {"name": "first-trailing", "parameters": {}}}, + ], + ) + changed_trailing_tools = [*cached_tools[:-1], {"type": "function", "function": {"name": "second-trailing", "parameters": {}}}] + + assert PromptCachingCache.extract_cacheable_tools(cached_tools) == cached_tools[:1] + assert PromptCachingCache.get_prompt_caching_cache_key(messages, cached_tools) == PromptCachingCache.get_prompt_caching_cache_key( + messages, changed_trailing_tools + ) + + +def test_tool_affinity_key_differs_for_uncached_tools_with_no_breakpoint(): + """No tool carries a cache breakpoint, but the tools still precede whatever the messages-side + breakpoint caches on the provider, so distinct tool lists must not collide on the same key.""" + messages = _messages(word_count=10) + tools_a = cast(list[ChatCompletionToolParam], [{"type": "function", "function": {"name": "alpha", "parameters": {}}}]) + tools_b = cast(list[ChatCompletionToolParam], [{"type": "function", "function": {"name": "beta", "parameters": {}}}]) + + assert PromptCachingCache.extract_cacheable_tools(tools_a) == tools_a + assert PromptCachingCache.get_prompt_caching_cache_key(messages, tools_a) != PromptCachingCache.get_prompt_caching_cache_key( + messages, tools_b + ) + + +def test_prepend_system_prompt_does_not_duplicate_an_already_logged_system_message(): + """Standard logging's append_system_prompt_messages already prepends a string system prompt + onto `messages` before async_log_success_event runs, so prepend_system_prompt must not add it + a second time or the write-path affinity key would never match the read-path key.""" + system = "You are a helpful assistant" + messages = cast(List[AllMessageValues], [{"role": "system", "content": system}, {"role": "user", "content": "hi"}]) + + assert PromptCachingCache.prepend_system_prompt(messages, system) == messages + + +@pytest.mark.asyncio +async def test_string_system_parameter_does_not_double_up_the_affinity_key(): + """Regression: kwargs["system"] as a string is already baked into `messages` by standard + logging before async_log_success_event runs. Re-prepending it there produced a different + (duplicated) affinity key than async_filter_deployments computes from the raw request + messages, so a pin written on success was never read back.""" + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") + messages = _messages(word_count=5000) + system = "cached system prompt" + standard_logging_object = { + "call_type": "acompletion", + "model": "anthropic/claude-opus-4-6", + "messages": [{"role": "system", "content": system}, *messages], + "model_id": "dep-2", + } + + await check.async_log_success_event( + kwargs={"standard_logging_object": standard_logging_object, "system": system}, + response_obj=None, + start_time=None, + end_time=None, + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"system": system}, + ) + + assert filtered == [deployments[1]] + + +@pytest.mark.asyncio +async def test_system_parameter_is_part_of_prompt_cache_affinity(): + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments("anthropic/claude-opus-4-6", "anthropic/claude-opus-4-6") + messages = _messages(word_count=5000) + system = [{"type": "text", "text": "system", "cache_control": {"type": "ephemeral", "ttl": "1h"}}] + cached_messages = PromptCachingCache.prepend_system_prompt(messages, system) + + await PromptCachingCache(cache=cache).async_add_model_id("dep-2", cached_messages, None) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=messages, + request_kwargs={"system": system}, + ) + + assert filtered == [deployments[1]] + + +@pytest.mark.parametrize( + ("ttl", "expected_affinity_ttl"), + ((None, 300), ("5m", 300), ("1h", 3600)), +) +def test_prompt_caching_affinity_ttl_matches_cache_control(ttl: str | None, expected_affinity_ttl: int): + cache_control: dict[str, str] = {"type": "ephemeral", **({"ttl": ttl} if ttl is not None else {})} + messages = cast( + List[AllMessageValues], + [{"role": "system", "content": [{"type": "text", "text": "cached", "cache_control": cache_control}]}], + ) + + assert PromptCachingCache.get_prompt_caching_ttl(messages) == expected_affinity_ttl + + +def test_message_level_cache_control_on_string_content_sets_affinity_ttl(): + """cache_control can sit as a sibling of a string `content` field rather than inside a + content-block list, e.g. {"role": "system", "content": "...", "cache_control": {...}}. + get_prompt_caching_ttl must still pick it up.""" + messages = cast( + List[AllMessageValues], + [{"role": "system", "content": "cached", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + ) + + assert PromptCachingCache.get_prompt_caching_ttl(messages) == 3600 + + +def test_mixed_cache_ttls_use_the_shortest_affinity_ttl(): + messages = cast( + List[AllMessageValues], + [ + { + "role": "system", + "content": [{"type": "text", "text": "system", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + }, + { + "role": "user", + "content": [{"type": "text", "text": "user", "cache_control": {"type": "ephemeral", "ttl": "5m"}}], + }, + ], + ) + + assert PromptCachingCache.get_prompt_caching_ttl(messages) == 300 + + +def test_add_model_id_uses_one_hour_affinity_ttl(): + cache = DualCache() + set_cache = Mock() + cache.set_cache = set_cache + messages = cast( + List[AllMessageValues], + [{"role": "system", "content": [{"type": "text", "text": "cached", "cache_control": {"type": "ephemeral", "ttl": "1h"}}]}], + ) + + PromptCachingCache(cache=cache).add_model_id("dep-1", messages, None) + + assert set_cache.call_args.kwargs["ttl"] == 3600 + + +@pytest.mark.asyncio +async def test_async_add_model_id_uses_one_hour_affinity_ttl(): + cache = DualCache() + async_set_cache = AsyncMock() + cache.async_set_cache = async_set_cache + messages = cast( + List[AllMessageValues], + [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "cached", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ], + ) + + await PromptCachingCache(cache=cache).async_add_model_id("dep-1", messages, None) + + assert async_set_cache.call_args.kwargs["ttl"] == 3600 + + +@pytest.mark.asyncio +async def test_async_add_model_id_uses_one_hour_tool_affinity_ttl(): + cache = DualCache() + async_set_cache = AsyncMock() + cache.async_set_cache = async_set_cache + messages = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [{"type": "text", "text": "cached", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + ) + tools = cast( + list[ChatCompletionToolParam], + [ + { + "type": "function", + "function": {"name": "lookup", "parameters": {}}, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + ) + + await PromptCachingCache(cache=cache).async_add_model_id("dep-1", messages, tools) + + assert async_set_cache.call_args.kwargs["ttl"] == 3600 + + def test_get_min_token_count_for_deployments_takes_min_across_mixed_group(): """ A group may legally mix models whose real minimums differ, and one gate decides for every