Merge 147e2452b6 into 1c61c2606e

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

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
 

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