refactor(prompt-caching): extract TTL and cache key logic from prefix computation

Split get_prompt_caching_ttl into get_prompt_caching_ttl_from_prefix to separate prefix extraction from TTL calculation. Split get_prompt_caching_cache_key into get_prompt_caching_cache_key_from_prefix similarly. Update add_model_id and async_add_model_id to extract cacheable prefix once and pass it to both cache key and TTL methods. Check tools for cache_control TTL values in addition to messages. Add test coverage for sync

litellm/router_utils/prompt_caching_cache.py

@@ -140,13 +140,20 @@ class PromptCachingCache:
         return cacheable_prefix
 
     @staticmethod
-    def get_prompt_caching_ttl(messages: list[AllMessageValues] | None) -> int:
-        if messages is None:
-            return 300
+    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 []
+        return PromptCachingCache.get_prompt_caching_ttl_from_prefix(cacheable_prefix, tools)
 
-        cacheable_prefix: Final = PromptCachingCache.extract_cacheable_prefix(messages)
-        cache_control_ttls: Final = tuple(
-            cache_control.get("ttl")
+    @staticmethod
+    def get_prompt_caching_ttl_from_prefix(
+        cacheable_prefix: list[AllMessageValues],
+        tools: list[ChatCompletionToolParam] | None,
+    ) -> int:
+        cache_control_values: Final = tuple(
+            cache_control
             for message in cacheable_prefix
             for cache_control in (
                 message.get("cache_control"),
@@ -157,24 +164,28 @@ class PromptCachingCache:
                 ),
             )
             if isinstance(cache_control, dict) and cache_control.get("type") == "ephemeral"
-        )
-        return 3600 if "1h" in cache_control_ttls else 300
+        ) + 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
 
     @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:
-            return 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)
 
-        # 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
+    @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
 
         # Use serialize_object for consistent and stable serialization
         data_to_hash: Final = {}
@@ -205,15 +216,17 @@ 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=PromptCachingCache.get_prompt_caching_ttl(messages),
+            ttl=PromptCachingCache.get_prompt_caching_ttl_from_prefix(cacheable_prefix or [], tools),
         )
         return
 
@@ -226,15 +239,17 @@ 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=PromptCachingCache.get_prompt_caching_ttl(messages),
+            ttl=PromptCachingCache.get_prompt_caching_ttl_from_prefix(cacheable_prefix or [], tools),
         )
         return
 

tests/test_litellm/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py

@@ -1,7 +1,7 @@
 import asyncio
 import copy
 from typing import List, cast
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 
@@ -16,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"
@@ -75,6 +75,20 @@ 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_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()
@@ -101,6 +115,28 @@ async def test_async_add_model_id_uses_one_hour_affinity_ttl():
     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 = _messages(word_count=1400)
+    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

type-discipline-budget.json

@@ -27,7 +27,7 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16398
+    "limit": 16396
   },
   "LIT011": {
     "limit": 5504