From 971037577cff750b72e797a396d7c6f00a01e60a Mon Sep 17 00:00:00 2001 From: nuernber Date: Fri, 11 Sep 2026 13:03:18 -0700 Subject: [PATCH] feat(prompt-caching): include system parameter in cache affinity key Add prepend_system_prompt helper to merge system parameter into messages for cache key computation. Update async_filter_deployments to prepend system prompt before computing affinity and pass tools to async_get_model_id. Update async_log_success_event to prepend system prompt and pass tools to async_add_model_id. Add test coverage for system parameter in cache affinity. --- .../prompt_caching_deployment_check.py | 23 +++++++++++++++---- litellm/router_utils/prompt_caching_cache.py | 12 ++++++++++ .../test_prompt_caching_deployment_check.py | 21 +++++++++++++++++ type-discipline-budget.json | 2 +- 4 files changed, 52 insertions(+), 6 deletions(-) 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..da74788c92b 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,18 @@ 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(list[AllToolParamValues] | None, request_kwargs.get("tools")) + 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 +147,23 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + logged_messages: Final = PromptCachingCache.prepend_system_prompt( + cast(list[AllMessageValues], messages), + 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(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 20af344ea3c..c5c89afc6ae 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -139,6 +139,18 @@ class PromptCachingCache: return cacheable_prefix + @staticmethod + def prepend_system_prompt( + messages: list[AllMessageValues], + system: object | None, + ) -> list[AllMessageValues]: + if system is None: + return messages + return cast( + list[AllMessageValues], + [{"role": "system", "content": system}, *messages], + ) + @staticmethod def get_prompt_caching_ttl( messages: list[AllMessageValues] | None, 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 20db54d6645..eac5f5ec0ae 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,27 @@ def _messages(word_count: int) -> List[AllMessageValues]: ) +@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)), diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 3fb6b12fa08..d21f4882147 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16396 + "limit": 16394 }, "LIT011": { "limit": 5504