From 17fc935ed60c5e55c00b4da27a55abc8ea1a00ec Mon Sep 17 00:00:00 2001 From: nuernber Date: Fri, 11 Sep 2026 13:28:57 -0700 Subject: [PATCH] feat(prompt-caching): use shortest TTL for mixed cache control blocks and exclude trailing uncached tools from affinity key Extract cacheable_tools helper to slice tools at the last cache_control breakpoint. Update get_prompt_caching_ttl to use all() instead of any() so mixed TTL values default to the shorter 5 minute affinity. Update get_prompt_caching_cache_key to serialize only cacheable tools. Add test coverage for prepend_system_prompt preserving messages without system parameter, tool affinity ignoring trailing --- litellm/router_utils/prompt_caching_cache.py | 25 +++++++-- .../test_prompt_caching_deployment_check.py | 51 ++++++++++++++++++- type-discipline-budget.json | 2 +- 3 files changed, 73 insertions(+), 5 deletions(-) diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index c5c89afc6ae..e8ef3d0d3fd 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -139,6 +139,22 @@ 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(tools[index].get("cache_control"), dict) + and tools[index]["cache_control"].get("type") == "ephemeral" + ), + None, + ) + # Match the provider prefix exactly instead of pinning on uncached trailing tools + return tools[: cacheable_tool_index + 1] if cacheable_tool_index is not None else [] + @staticmethod def prepend_system_prompt( messages: list[AllMessageValues], @@ -164,6 +180,7 @@ class PromptCachingCache: cacheable_prefix: list[AllMessageValues], tools: list[ChatCompletionToolParam] | None, ) -> int: + cacheable_tools: Final = PromptCachingCache.extract_cacheable_tools(tools or []) cache_control_values: Final = tuple( cache_control for message in cacheable_prefix @@ -176,8 +193,9 @@ class PromptCachingCache: ), ) if isinstance(cache_control, dict) and cache_control.get("type") == "ephemeral" - ) + tuple(tool.get("cache_control") for tool in (tools or ()) if isinstance(tool.get("cache_control"), dict)) - return 3600 if any(value.get("ttl") == "1h" for value in cache_control_values) else 300 + ) + tuple(tool.get("cache_control") for tool in cacheable_tools if isinstance(tool.get("cache_control"), dict)) + # 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( @@ -205,7 +223,8 @@ class PromptCachingCache: 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 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 eac5f5ec0ae..cabd41c1943 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 @@ -61,6 +61,29 @@ 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 + ) + + @pytest.mark.asyncio async def test_system_parameter_is_part_of_prompt_cache_affinity(): cache = DualCache() @@ -96,6 +119,24 @@ def test_prompt_caching_affinity_ttl_matches_cache_control(ttl: str | None, expe assert PromptCachingCache.get_prompt_caching_ttl(messages) == expected_affinity_ttl +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() @@ -141,7 +182,15 @@ 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 = _messages(word_count=1400) + messages = cast( + List[AllMessageValues], + [ + { + "role": "user", + "content": [{"type": "text", "text": "cached", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + } + ], + ) tools = cast( list[ChatCompletionToolParam], [ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index d21f4882147..9456cf2327e 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16394 + "limit": 16392 }, "LIT011": { "limit": 5504