diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index 171a086de81..20af344ea3c 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/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 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 fde779b3feb..20db54d6645 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,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 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 0c0952289e2..3fb6b12fa08 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16398 + "limit": 16396 }, "LIT011": { "limit": 5504