From 77a3d061be9d4ccc3f6b2c13d4640ce3b68be370 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:24:51 -0700 Subject: [PATCH 1/8] feat(proxy): add a CustomLogger hook that decides per request whether the semantic cache applies --- litellm/caching/caching.py | 4 +- litellm/integrations/custom_logger.py | 14 ++ .../proxy/common_utils/semantic_cache_gate.py | 65 ++++++ litellm/proxy/utils.py | 4 + .../common_utils/test_semantic_cache_gate.py | 201 ++++++++++++++++++ 5 files changed, 286 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/common_utils/semantic_cache_gate.py create mode 100644 tests/unit/proxy/common_utils/test_semantic_cache_gate.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..8bbfed260c7 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -353,7 +353,7 @@ class Cache: ) _SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id" - def _is_semantic_cache(self) -> bool: + def is_semantic_cache(self) -> bool: return self.type in ( LiteLLMCacheType.REDIS_SEMANTIC, LiteLLMCacheType.QDRANT_SEMANTIC, @@ -395,7 +395,7 @@ class Cache: return preset_cache_key combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params() - is_semantic_cache: Final = self._is_semantic_cache() + is_semantic_cache: Final = self.is_semantic_cache() scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset() for param in kwargs: if param in scope_excluded_params: diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..4f3801d6e9b 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -443,6 +443,20 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ return model_names + async def async_should_semantic_cache( + self, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str, + ) -> bool: + """Runs once per request the proxy's semantic cache would serve, after the pre-call hooks, + with the authenticated key and the request body. Return True to let the request read from + and write to the semantic cache. Any other return value, or an exception, skips both for that + request. Every callback that overrides this hook must return True for the request to be cached. + Exact-match cache types never call it. + """ + return True + async def async_post_call_response_headers_hook( self, data: dict, diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py new file mode 100644 index 00000000000..217fe596e23 --- /dev/null +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -0,0 +1,65 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + +SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True}) +CACHED_CALL_TYPES_BY_ROUTE: Final = MappingProxyType( + { + "aresponses": frozenset({"aresponses", "responses", "acompletion"}), + "anthropic_messages": frozenset({"anthropic_messages", "aanthropic_messages", "acompletion"}), + "agenerate_content": frozenset({"acompletion"}), + "agenerate_content_stream": frozenset({"acompletion"}), + } +) + + +def _cache_serves(cache: Cache, call_type: str) -> bool: + cached_call_types: Final = CACHED_CALL_TYPES_BY_ROUTE.get(call_type, frozenset({call_type})) + return not cached_call_types.isdisjoint(cache.supported_call_types or ()) + + +async def _gate_allows( + gate: CustomLogger, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str, +) -> bool: + try: + return await gate.async_should_semantic_cache(user_api_key_dict, data, call_type) is True + except Exception: # noqa: BLE001 # operator hooks can raise anything, and every failure skips the cache + verbose_proxy_logger.exception( + "%s.async_should_semantic_cache raised, so this request skips the semantic cache", type(gate).__name__ + ) + return False + + +async def _every_gate_allows( + gates: Sequence[CustomLogger], + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str, +) -> bool: + for gate in gates: + if not await _gate_allows(gate, user_api_key_dict, data, call_type): + return False + return True + + +async def gated_cache_controls( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + call_type: str, + gates: Sequence[CustomLogger], + cache: Cache | None, +) -> Mapping[str, object] | None: + if not gates or cache is None or not cache.is_semantic_cache() or not _cache_serves(cache, call_type): + return None + if await _every_gate_allows(gates, user_api_key_dict, data, call_type): + return None + controls: Final = data.get("cache") + return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ea1044d4a55..1966c32fb07 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1496,6 +1496,7 @@ class _CallbackCapabilities: # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) + semantic_cache_gates: tuple[CustomLogger, ...] = field(default_factory=tuple) def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: @@ -3275,6 +3276,9 @@ class ProxyLogging: listed_models_filters=tuple( callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") ), + semantic_cache_gates=tuple( + callback for callback in resolved_callbacks if _overrides_hook(callback, "async_should_semantic_cache") + ), ) # Limit cache to handle test churn without leaking; production # callback lists are stable so this rarely grows past 1 entry. diff --git a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py new file mode 100644 index 00000000000..d99067a1b0e --- /dev/null +++ b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py @@ -0,0 +1,201 @@ +from collections.abc import Mapping + +import pytest + +import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.utils import ProxyLogging +from litellm.types.caching import ( + DEFAULT_CACHING_SUPPORTED_CALL_TYPES, + CachingSupportedCallTypes, + LiteLLMCacheType, +) + +ALLOWED_TEAM = "team-eval" +OTHER_TEAM = "team-other" + + +class _TeamAllowlist(CustomLogger): + def __init__(self, teams: frozenset[str]) -> None: + super().__init__() + self.teams = teams + self.seen: list[tuple[str | None, object, str]] = [] + + async def async_should_semantic_cache( + self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> bool: + self.seen.append((user_api_key_dict.team_id, data.get("model"), call_type)) + return user_api_key_dict.team_id in self.teams + + +class _RaisingGate(CustomLogger): + async def async_should_semantic_cache( + self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> bool: + raise RuntimeError("allowlist service down") + + +class _TruthyGate(CustomLogger): + async def async_should_semantic_cache( # pyright: ignore[reportIncompatibleMethodOverride] # operator hooks are untyped Python + self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> str: + return "yes" + + +def _semantic_cache( + supported_call_types: tuple[CachingSupportedCallTypes, ...] = DEFAULT_CACHING_SUPPORTED_CALL_TYPES, +) -> Cache: + return Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + host="localhost", + port="6379", + similarity_threshold=0.8, + supported_call_types=list(supported_call_types), + ) + + +def _request(**cache_controls: object) -> dict[str, object]: + return { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "What color is the sky?"}], + **({"cache": cache_controls} if cache_controls else {}), + } + + +def _key(team_id: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-hashed", team_id=team_id) + + +@pytest.mark.asyncio +async def test_a_team_the_gate_rejects_skips_semantic_lookup_and_store_even_when_the_client_opts_in(): + gate = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + data = _request(**{"use-cache": True, "ttl": 60}) + + result = await gated_cache_controls( + data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=_semantic_cache() + ) + + assert result == {"use-cache": True, "ttl": 60, "no-cache": True, "no-store": True} + assert gate.seen == [(OTHER_TEAM, "gpt-4o-mini", "acompletion")] + + +@pytest.mark.asyncio +async def test_a_team_the_gate_allows_keeps_its_own_cache_controls(): + data = _request(ttl=60) + + result = await gated_cache_controls( + data=data, + user_api_key_dict=_key(ALLOWED_TEAM), + call_type="acompletion", + gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),), + cache=_semantic_cache(), + ) + + assert result is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cache", [None, Cache(type=LiteLLMCacheType.LOCAL)], ids=["no-cache", "exact-match-cache"]) +async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache | None): + gate = _TeamAllowlist(frozenset()) + data = _request() + + result = await gated_cache_controls( + data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=cache + ) + + assert result is None + assert gate.seen == [] + + +@pytest.mark.asyncio +async def test_the_gate_is_never_consulted_for_a_call_type_the_cache_does_not_serve(): + gate = _TeamAllowlist(frozenset()) + data = _request() + + result = await gated_cache_controls( + data=data, + user_api_key_dict=_key(OTHER_TEAM), + call_type="aembedding", + gates=(gate,), + cache=_semantic_cache(supported_call_types=("acompletion",)), + ) + + assert result is None + assert gate.seen == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("route", "served_call_type"), + [ + ("aresponses", "responses"), + ("aresponses", "acompletion"), + ("anthropic_messages", "acompletion"), + ("agenerate_content", "acompletion"), + ("agenerate_content_stream", "acompletion"), + ], +) +async def test_the_gate_covers_routes_that_reach_a_served_call_type_through_an_alias_or_a_bridge( + route: str, served_call_type: CachingSupportedCallTypes +): + gate = _TeamAllowlist(frozenset()) + + result = await gated_cache_controls( + data=_request(), + user_api_key_dict=_key(OTHER_TEAM), + call_type=route, + gates=(gate,), + cache=_semantic_cache(supported_call_types=(served_call_type,)), + ) + + assert result == {"no-cache": True, "no-store": True} + assert gate.seen == [(OTHER_TEAM, "gpt-4o-mini", route)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("gate", [_RaisingGate(), _TruthyGate()], ids=["raises", "returns-non-bool"]) +async def test_a_gate_that_fails_or_answers_anything_but_true_skips_the_semantic_cache(gate: CustomLogger): + result = await gated_cache_controls( + data=_request(), + user_api_key_dict=_key(ALLOWED_TEAM), + call_type="aresponses", + gates=(gate,), + cache=_semantic_cache(), + ) + + assert result == {"no-cache": True, "no-store": True} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("deny_first", [True, False]) +async def test_one_rejecting_gate_skips_the_semantic_cache_whatever_the_others_answer(deny_first: bool): + allow = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + deny = _TeamAllowlist(frozenset()) + + result = await gated_cache_controls( + data=_request(), + user_api_key_dict=_key(ALLOWED_TEAM), + call_type="anthropic_messages", + gates=(deny, allow) if deny_first else (allow, deny), + cache=_semantic_cache(), + ) + + assert result == {"no-cache": True, "no-store": True} + + +@pytest.mark.asyncio +async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: pytest.MonkeyPatch): + gate = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), gate]) + + gates = ProxyLogging.callback_capabilities().semantic_cache_gates + result = await gated_cache_controls( + data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=gates, cache=_semantic_cache() + ) + + assert gates == (gate,) + assert result == {"no-cache": True, "no-store": True} From 1198cbcbb8aa718f376312c0445fc3cc015033e4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:16:47 -0700 Subject: [PATCH 2/8] test(proxy): hold the semantic cache gate to its regression tests until the code lands in the next commit --- litellm/caching/caching.py | 4 +- litellm/integrations/custom_logger.py | 14 ---- .../proxy/common_utils/semantic_cache_gate.py | 65 ------------------- litellm/proxy/utils.py | 4 -- .../google_genai/test_google_genai_adapter.py | 38 +++++++++++ 5 files changed, 40 insertions(+), 85 deletions(-) delete mode 100644 litellm/proxy/common_utils/semantic_cache_gate.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 8bbfed260c7..5f5b8210096 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -353,7 +353,7 @@ class Cache: ) _SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id" - def is_semantic_cache(self) -> bool: + def _is_semantic_cache(self) -> bool: return self.type in ( LiteLLMCacheType.REDIS_SEMANTIC, LiteLLMCacheType.QDRANT_SEMANTIC, @@ -395,7 +395,7 @@ class Cache: return preset_cache_key combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params() - is_semantic_cache: Final = self.is_semantic_cache() + is_semantic_cache: Final = self._is_semantic_cache() scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset() for param in kwargs: if param in scope_excluded_params: diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 4f3801d6e9b..a4de1be1842 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -443,20 +443,6 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ return model_names - async def async_should_semantic_cache( - self, - user_api_key_dict: UserAPIKeyAuth, - data: Mapping[str, object], - call_type: str, - ) -> bool: - """Runs once per request the proxy's semantic cache would serve, after the pre-call hooks, - with the authenticated key and the request body. Return True to let the request read from - and write to the semantic cache. Any other return value, or an exception, skips both for that - request. Every callback that overrides this hook must return True for the request to be cached. - Exact-match cache types never call it. - """ - return True - async def async_post_call_response_headers_hook( self, data: dict, diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py deleted file mode 100644 index 217fe596e23..00000000000 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ /dev/null @@ -1,65 +0,0 @@ -from collections.abc import Mapping, Sequence -from types import MappingProxyType -from typing import Final - -from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import Cache -from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth - -SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True}) -CACHED_CALL_TYPES_BY_ROUTE: Final = MappingProxyType( - { - "aresponses": frozenset({"aresponses", "responses", "acompletion"}), - "anthropic_messages": frozenset({"anthropic_messages", "aanthropic_messages", "acompletion"}), - "agenerate_content": frozenset({"acompletion"}), - "agenerate_content_stream": frozenset({"acompletion"}), - } -) - - -def _cache_serves(cache: Cache, call_type: str) -> bool: - cached_call_types: Final = CACHED_CALL_TYPES_BY_ROUTE.get(call_type, frozenset({call_type})) - return not cached_call_types.isdisjoint(cache.supported_call_types or ()) - - -async def _gate_allows( - gate: CustomLogger, - user_api_key_dict: UserAPIKeyAuth, - data: Mapping[str, object], - call_type: str, -) -> bool: - try: - return await gate.async_should_semantic_cache(user_api_key_dict, data, call_type) is True - except Exception: # noqa: BLE001 # operator hooks can raise anything, and every failure skips the cache - verbose_proxy_logger.exception( - "%s.async_should_semantic_cache raised, so this request skips the semantic cache", type(gate).__name__ - ) - return False - - -async def _every_gate_allows( - gates: Sequence[CustomLogger], - user_api_key_dict: UserAPIKeyAuth, - data: Mapping[str, object], - call_type: str, -) -> bool: - for gate in gates: - if not await _gate_allows(gate, user_api_key_dict, data, call_type): - return False - return True - - -async def gated_cache_controls( - data: Mapping[str, object], - user_api_key_dict: UserAPIKeyAuth, - call_type: str, - gates: Sequence[CustomLogger], - cache: Cache | None, -) -> Mapping[str, object] | None: - if not gates or cache is None or not cache.is_semantic_cache() or not _cache_serves(cache, call_type): - return None - if await _every_gate_allows(gates, user_api_key_dict, data, call_type): - return None - controls: Final = data.get("cache") - return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1966c32fb07..ea1044d4a55 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1496,7 +1496,6 @@ class _CallbackCapabilities: # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) - semantic_cache_gates: tuple[CustomLogger, ...] = field(default_factory=tuple) def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: @@ -3276,9 +3275,6 @@ class ProxyLogging: listed_models_filters=tuple( callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") ), - semantic_cache_gates=tuple( - callback for callback in resolved_callbacks if _overrides_hook(callback, "async_should_semantic_cache") - ), ) # Limit cache to handle test churn without leaking; production # callback lists are stable so this rarely grows past 1 entry. diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 5e521506699..dbcc80a5d25 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -2,6 +2,7 @@ """ Test to verify the Google GenAI generate_content adapter functionality """ +import asyncio import json import unittest @@ -1641,3 +1642,40 @@ async def test_generate_content_sends_response_schema_and_tool_parameters_to_the } assert response["candidates"][0]["content"]["parts"] == [{"text": '{"park_name": "EPCOT"}'}] assert "text" not in response + + +@pytest.mark.asyncio +async def test_generate_content_bridge_honors_the_request_cache_controls(respx_mock, monkeypatch): + import httpx + + from litellm.caching.caching import Cache + from litellm.types.caching import LiteLLMCacheType + + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) + route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-capital", + "object": "chat.completion", + "created": 1234567890, + "model": "gpt-4.1", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Paris"}}], + "usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11}, + }, + ) + ) + request = { + "model": "openai/gpt-4.1", + "contents": [{"role": "user", "parts": [{"text": "What is the capital of France?"}]}], + "cache": {"no-cache": True, "no-store": True}, + "api_key": "sk-test", + } + + await agenerate_content(**request) + await asyncio.sleep(0.5) + await agenerate_content(**request) + + assert route.call_count == 2 + From 6517331b1df485614ff025b11a73c3e37c566dea Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 20:53:25 -0700 Subject: [PATCH 3/8] feat(proxy): call the semantic cache check on proxy requests and forward cache controls through the generateContent bridge --- litellm/caching/caching.py | 8 +-- litellm/google_genai/adapters/handler.py | 2 + litellm/integrations/custom_logger.py | 14 ++++ litellm/proxy/common_request_processing.py | 11 ++++ .../proxy/common_utils/semantic_cache_gate.py | 65 +++++++++++++++++++ litellm/proxy/utils.py | 4 ++ litellm/types/google_genai/adapters.py | 1 + .../proxy/test_common_request_processing.py | 57 ++++++++++++++++ 8 files changed, 158 insertions(+), 4 deletions(-) create mode 100644 litellm/proxy/common_utils/semantic_cache_gate.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 5f5b8210096..17cf2c98da0 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -58,7 +58,7 @@ def response_cache_phase(operation: Literal["get", "set"]) -> Generator[None]: yield -def print_verbose(print_statement): +def print_verbose(print_statement: object) -> None: try: verbose_logger.debug(print_statement) if litellm.set_verbose: @@ -353,7 +353,7 @@ class Cache: ) _SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id" - def _is_semantic_cache(self) -> bool: + def is_semantic_cache(self) -> bool: return self.type in ( LiteLLMCacheType.REDIS_SEMANTIC, LiteLLMCacheType.QDRANT_SEMANTIC, @@ -395,7 +395,7 @@ class Cache: return preset_cache_key combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params() - is_semantic_cache: Final = self._is_semantic_cache() + is_semantic_cache: Final = self.is_semantic_cache() scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset() for param in kwargs: if param in scope_excluded_params: @@ -963,7 +963,7 @@ class Cache: except Exception as e: self._log_add_cache_failure(e) - def should_use_cache(self, **kwargs): + def should_use_cache(self, **kwargs: object) -> bool: """ Returns true if we should use the cache for LLM API calls diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index 7caa0c62571..e386008f5d3 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -69,6 +69,8 @@ class GenerateContentToCompletionHandler: completion_kwargs["proxy_server_request"] = extra_kwargs["proxy_server_request"] if "litellm_metadata" in extra_kwargs: completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"] + if "cache" in extra_kwargs: + completion_kwargs["cache"] = extra_kwargs["cache"] if stream: completion_kwargs["stream"] = stream diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index a4de1be1842..4f3801d6e9b 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -443,6 +443,20 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac """ return model_names + async def async_should_semantic_cache( + self, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str, + ) -> bool: + """Runs once per request the proxy's semantic cache would serve, after the pre-call hooks, + with the authenticated key and the request body. Return True to let the request read from + and write to the semantic cache. Any other return value, or an exception, skips both for that + request. Every callback that overrides this hook must return True for the request to be cached. + Exact-match cache types never call it. + """ + return True + async def async_post_call_response_headers_hook( self, data: dict, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 59db23820fc..b98a8f2184b 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -107,6 +107,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_param, openai_error_type, ) +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -2472,6 +2473,16 @@ class ProxyBaseLLMRequestProcessing: call_type=route_type, skip_guardrails=skip_guardrails, ) + if ( + cache_controls := await gated_cache_controls( + data=self.data, + user_api_key_dict=user_api_key_dict, + call_type=route_type, + gates=ProxyLogging.callback_capabilities().semantic_cache_gates, + cache=litellm.cache, + ) + ) is not None: + self.data["cache"] = cache_controls await _enforce_guardrail_added_tag_budgets( data=self.data, tags_before_guardrails=self._tags_before_guardrails, diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py new file mode 100644 index 00000000000..217fe596e23 --- /dev/null +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -0,0 +1,65 @@ +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + +SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True}) +CACHED_CALL_TYPES_BY_ROUTE: Final = MappingProxyType( + { + "aresponses": frozenset({"aresponses", "responses", "acompletion"}), + "anthropic_messages": frozenset({"anthropic_messages", "aanthropic_messages", "acompletion"}), + "agenerate_content": frozenset({"acompletion"}), + "agenerate_content_stream": frozenset({"acompletion"}), + } +) + + +def _cache_serves(cache: Cache, call_type: str) -> bool: + cached_call_types: Final = CACHED_CALL_TYPES_BY_ROUTE.get(call_type, frozenset({call_type})) + return not cached_call_types.isdisjoint(cache.supported_call_types or ()) + + +async def _gate_allows( + gate: CustomLogger, + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str, +) -> bool: + try: + return await gate.async_should_semantic_cache(user_api_key_dict, data, call_type) is True + except Exception: # noqa: BLE001 # operator hooks can raise anything, and every failure skips the cache + verbose_proxy_logger.exception( + "%s.async_should_semantic_cache raised, so this request skips the semantic cache", type(gate).__name__ + ) + return False + + +async def _every_gate_allows( + gates: Sequence[CustomLogger], + user_api_key_dict: UserAPIKeyAuth, + data: Mapping[str, object], + call_type: str, +) -> bool: + for gate in gates: + if not await _gate_allows(gate, user_api_key_dict, data, call_type): + return False + return True + + +async def gated_cache_controls( + data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + call_type: str, + gates: Sequence[CustomLogger], + cache: Cache | None, +) -> Mapping[str, object] | None: + if not gates or cache is None or not cache.is_semantic_cache() or not _cache_serves(cache, call_type): + return None + if await _every_gate_allows(gates, user_api_key_dict, data, call_type): + return None + controls: Final = data.get("cache") + return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ea1044d4a55..1966c32fb07 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1496,6 +1496,7 @@ class _CallbackCapabilities: # every string entry in ``litellm.callbacks``. resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple) + semantic_cache_gates: tuple[CustomLogger, ...] = field(default_factory=tuple) def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool: @@ -3275,6 +3276,9 @@ class ProxyLogging: listed_models_filters=tuple( callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models") ), + semantic_cache_gates=tuple( + callback for callback in resolved_callbacks if _overrides_hook(callback, "async_should_semantic_cache") + ), ) # Limit cache to handle test churn without leaking; production # callback lists are stable so this rarely grows past 1 entry. diff --git a/litellm/types/google_genai/adapters.py b/litellm/types/google_genai/adapters.py index f03c67aa01a..61e760d21c9 100644 --- a/litellm/types/google_genai/adapters.py +++ b/litellm/types/google_genai/adapters.py @@ -22,6 +22,7 @@ class GenerateContentCompletionKwargs(TypedDict, total=False): stream: bool metadata: dict[str, object] litellm_metadata: ReadOnly[Mapping[str, object]] + cache: ReadOnly[Mapping[str, object]] stream_options: ReadOnly[ChatCompletionStreamOptionsParam] extra_headers: dict[str, str] | None proxy_server_request: dict[str, object] | None diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index d742956bb59..5dee9b6c614 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -512,6 +512,63 @@ class TestProxyBaseLLMRequestProcessing: assert exc_info.value.param == "model" add_litellm_data_to_request.assert_not_awaited() + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("gate_allows", "expected_cache_controls"), + [(True, {"ttl": 60}), (False, {"ttl": 60, "no-cache": True, "no-store": True})], + ids=["allowed", "rejected"], + ) + async def test_common_processing_pre_call_logic_applies_the_semantic_cache_gate( + self, + monkeypatch: pytest.MonkeyPatch, + gate_allows: bool, + expected_cache_controls: dict[str, object], + ) -> None: + from litellm.caching.caching import Cache + from litellm.types.caching import LiteLLMCacheType + + class _Gate(CustomLogger): + async def async_should_semantic_cache( + self, user_api_key_dict: ProxyUserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> bool: + return gate_allows + + async def passthrough_pre_call_hook( + user_api_key_dict: ProxyUserAPIKeyAuth, + data: dict[str, object], + call_type: str, + skip_guardrails: bool = False, + ) -> dict[str, object]: + return data + + monkeypatch.setattr(litellm, "callbacks", [_Gate()]) + monkeypatch.setattr( + litellm, + "cache", + Cache(type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8), + ) + body: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "add_litellm_data_to_request", + AsyncMock(return_value={**body, "metadata": {}, "proxy_server_request": {"body": body}}), + ) + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + mock_proxy_logging_obj: Final = MagicMock(spec=ProxyLogging) + mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) + + returned_data, _ = await ProxyBaseLLMRequestProcessing(data={}).common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-hashed", team_id="team-other"), + proxy_logging_obj=mock_proxy_logging_obj, + proxy_config=MagicMock(spec=ProxyConfig), + route_type="acompletion", + ) + + assert returned_data["cache"] == expected_cache_controls + @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True]) @pytest.mark.parametrize( From 6984196c14355f50ffd6e05f46a98abb994999ed Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:31:29 -0700 Subject: [PATCH 4/8] fix(proxy): run the semantic cache check on every route that can reach the cache, including A2A, WebSocket, interactions and batch rows --- .../proxy/agent_endpoints/a2a_endpoints.py | 7 ++- litellm/proxy/batches_endpoints/endpoints.py | 2 + .../litellm_executed_batches.py | 17 +++++++- .../proxy/common_utils/semantic_cache_gate.py | 25 ++++++----- .../agent_endpoints/test_a2a_endpoints.py | 33 +++++++++++++- .../test_litellm_executed_batches.py | 43 +++++++++++++++++++ .../common_utils/test_semantic_cache_gate.py | 25 ++++++----- 7 files changed, 127 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c11d296837e..f883d49532a 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -180,6 +180,11 @@ async def _resolve_backend_auth( ) +def _request_cache_controls(data: Mapping[str, object]) -> Mapping[str, object]: + cache_controls: Final = data.get("cache") + return {"cache": cache_controls} if isinstance(cache_controls, Mapping) else {} + + def _forwarding_headers( caller_identity: Mapping[str, str], request_data: Mapping[str, object], @@ -838,7 +843,7 @@ async def invoke_agent_a2a( ), backend_auth_header=backend_auth.header, ) - forwarded_litellm_params: Final = backend_auth.litellm_params + forwarded_litellm_params: Final = {**backend_auth.litellm_params, **_request_cache_controls(data)} # Merge agent-level guardrails into data so post_call_success_hook and # _handle_stream_message both pick them up. A2A agents use model diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 763737191ef..5fbb240ea01 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -110,6 +110,8 @@ def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyL batches=ManagedBatchRepository(prisma_client), proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, + semantic_cache_gates=ProxyLogging.callback_capabilities().semantic_cache_gates, + semantic_cache=litellm.cache, ) diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 098dd2f9816..1299c6dda6a 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -18,7 +18,9 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid as uuid_module +from litellm.caching.caching import Cache from litellm.constants import LITELLM_EXECUTED_BATCH_CONCURRENCY +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME from litellm.llms.base_llm.files.storage_backend import BaseFileStorageBackend @@ -27,6 +29,7 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.models.managed_files import LiteLLM_ManagedFileTable from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import is_request_body_safe +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.openai_files_endpoints.common_utils import LITELLM_EXECUTED_BATCH_ID_PREFIX from litellm.proxy.openai_files_endpoints.storage_backend_service import StorageBackendFileService @@ -423,6 +426,8 @@ class LiteLLMExecutedBatchRunner: batches: ManagedBatchRepository, proxy_logging_obj: ProxyLogging, general_settings: Mapping[str, object], + semantic_cache_gates: Sequence[CustomLogger], + semantic_cache: Cache | None, concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY, heartbeat_seconds: float = _HEARTBEAT_SECONDS, completion_window_seconds: float = _COMPLETION_WINDOW_SECONDS, @@ -435,6 +440,8 @@ class LiteLLMExecutedBatchRunner: self.batches = batches self.proxy_logging_obj = proxy_logging_obj self.general_settings = general_settings + self.semantic_cache_gates = semantic_cache_gates + self.semantic_cache = semantic_cache self.concurrency = concurrency self.heartbeat_seconds = heartbeat_seconds self.completion_window_seconds = completion_window_seconds @@ -644,9 +651,17 @@ class LiteLLMExecutedBatchRunner: return RowOutcome(custom_id=line.custom_id, status_code=200, body=body, succeeded=True) async def _dispatch(self, run: _BatchRun, line: BatchInputLine) -> Mapping[str, object]: - params: Final = MappingProxyType( + request: Final = MappingProxyType( {**line.body, "model": run.model, "metadata": self._row_metadata(run), "disable_fallbacks": True} ) + cache_controls: Final = await gated_cache_controls( + data=request, + user_api_key_dict=run.user_api_key_dict, + call_type=_ROUTER_METHODS[run.endpoint], + gates=self.semantic_cache_gates, + cache=self.semantic_cache, + ) + params: Final = MappingProxyType({**request, **({} if cache_controls is None else {"cache": cache_controls})}) return _dump(await self._router_call(run.endpoint)(**params)) def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall: diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py index 217fe596e23..b080c166f2d 100644 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -8,21 +8,24 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True}) -CACHED_CALL_TYPES_BY_ROUTE: Final = MappingProxyType( +SEMANTIC_CACHE_ROUTES: Final = frozenset( { - "aresponses": frozenset({"aresponses", "responses", "acompletion"}), - "anthropic_messages": frozenset({"anthropic_messages", "aanthropic_messages", "acompletion"}), - "agenerate_content": frozenset({"acompletion"}), - "agenerate_content_stream": frozenset({"acompletion"}), + "acompletion", + "atext_completion", + "aembedding", + "atranscription", + "arerank", + "aresponses", + "_aresponses_websocket", + "anthropic_messages", + "agenerate_content", + "agenerate_content_stream", + "acreate_interaction", + "asend_message", } ) -def _cache_serves(cache: Cache, call_type: str) -> bool: - cached_call_types: Final = CACHED_CALL_TYPES_BY_ROUTE.get(call_type, frozenset({call_type})) - return not cached_call_types.isdisjoint(cache.supported_call_types or ()) - - async def _gate_allows( gate: CustomLogger, user_api_key_dict: UserAPIKeyAuth, @@ -57,7 +60,7 @@ async def gated_cache_controls( gates: Sequence[CustomLogger], cache: Cache | None, ) -> Mapping[str, object] | None: - if not gates or cache is None or not cache.is_semantic_cache() or not _cache_serves(cache, call_type): + if not gates or cache is None or not cache.is_semantic_cache() or call_type not in SEMANTIC_CACHE_ROUTES: return None if await _every_gate_allows(gates, user_api_key_dict, data, call_type): return None diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index a2448694dd7..223b247a779 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -18,8 +18,11 @@ import pytest import respx import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.types.agents import AgentCaller +from litellm.types.caching import LiteLLMCacheType AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] @@ -28,6 +31,7 @@ AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] class CapturedAgentCall: request_id: object agent_extra_headers: dict[str, str] | None + litellm_params: Mapping[str, object] | None @pytest.mark.asyncio @@ -611,7 +615,11 @@ async def _invoke_message_method( kwargs: Final = downstream.call_args.kwargs request_id: Final = kwargs["request"].__dict__["id"] if is_send else kwargs["request_id"] - return CapturedAgentCall(request_id=request_id, agent_extra_headers=kwargs.get("agent_extra_headers")) + return CapturedAgentCall( + request_id=request_id, + agent_extra_headers=kwargs.get("agent_extra_headers"), + litellm_params=kwargs.get("litellm_params"), + ) @pytest.mark.asyncio @@ -625,6 +633,29 @@ async def test_message_methods_preserve_numeric_zero_request_id(method: str): assert captured.request_id == 0 +class _RejectEveryTeam(CustomLogger): + async def async_should_semantic_cache( + self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> bool: + return False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent_call(method: str): + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + semantic_cache = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 + ) + + with patch.object(litellm, "callbacks", [_RejectEveryTeam()]), patch.object(litellm, "cache", semantic_cache): + captured = await _invoke_message_method( + method, mock_request, UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), _add_proxy_data + ) + + assert (captured.litellm_params or {}).get("cache") == {"no-cache": True, "no-store": True} + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_forward_caller_identity_headers(method: str): diff --git a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py index f5c656c485c..1e62563f7e3 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -13,6 +13,8 @@ import pytest from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles from openai.types.batch_request_counts import BatchRequestCounts +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger from litellm.models.managed_files import LiteLLM_ManagedFileTable from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.batches_endpoints import litellm_executed_batches @@ -38,6 +40,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.repositories.managed_batch_repository import ManagedBatchRepository from litellm.router import Router +from litellm.types.caching import LiteLLMCacheType from litellm.types.llms.openai import LiteLLMBatchCreateRequest, OpenAIFileObject, OpenAIFilesPurpose from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, ModelResponse, SpecialEnums @@ -380,6 +383,8 @@ def make_runner( heartbeat_seconds: float = 30.0, completion_window_seconds: float = 24 * 60 * 60, user: UserAPIKeyAuth | None = None, + semantic_cache_gates: Sequence[CustomLogger] = (), + semantic_cache: Cache | None = None, ) -> Harness: store = store_factory({INPUT_FILE_ID: managed_input_file()} if files is None else files) router = FakeRouter() @@ -397,6 +402,8 @@ def make_runner( batches=ManagedBatchRepository(prisma), proxy_logging_obj=MagicMock(spec=ProxyLogging), general_settings=general_settings, + semantic_cache_gates=semantic_cache_gates, + semantic_cache=semantic_cache, concurrency=concurrency, heartbeat_seconds=heartbeat_seconds, completion_window_seconds=completion_window_seconds, @@ -702,6 +709,42 @@ async def test_create_dispatches_each_row_with_the_batch_model_and_the_key_metad assert metadata["user_api_key_user_email"] == "user@example.com" +class _TeamAllowlist(CustomLogger): + def __init__(self, team_id: str) -> None: + super().__init__() + self.team_id = team_id + self.seen: list[tuple[str | None, object, str]] = [] + + async def async_should_semantic_cache( + self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> bool: + self.seen.append((user_api_key_dict.team_id, data.get("messages"), call_type)) + return user_api_key_dict.team_id == self.team_id + + +@pytest.mark.parametrize( + ("allowed_team", "row_cache_controls"), + [("team-1", None), ("team-other", {"no-cache": True, "no-store": True})], + ids=["allowed-team", "rejected-team"], +) +async def test_create_runs_the_semantic_cache_check_on_each_row( + allowed_team: str, row_cache_controls: Mapping[str, object] | None +) -> None: + gate = _TeamAllowlist(allowed_team) + semantic_cache = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 + ) + harness = make_runner(semantic_cache_gates=(gate,), semantic_cache=semantic_cache) + await harness.create_and_finish() + + calls = harness.router.acompletion.await_args_list + assert [call.kwargs.get("cache") for call in calls] == [row_cache_controls, row_cache_controls] + assert sorted(gate.seen, key=str) == [ + ("team-1", [{"role": "user", "content": "hi 1"}], "acompletion"), + ("team-1", [{"role": "user", "content": "hi 2"}], "acompletion"), + ] + + async def test_create_dispatches_each_row_with_the_key_team_and_project_tags_as_inherited_tags() -> None: user = UserAPIKeyAuth( api_key="sk-batch-key", diff --git a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py index d99067a1b0e..9c95d28973d 100644 --- a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py +++ b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py @@ -112,16 +112,12 @@ async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache @pytest.mark.asyncio -async def test_the_gate_is_never_consulted_for_a_call_type_the_cache_does_not_serve(): +@pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech"]) +async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str): gate = _TeamAllowlist(frozenset()) - data = _request() result = await gated_cache_controls( - data=data, - user_api_key_dict=_key(OTHER_TEAM), - call_type="aembedding", - gates=(gate,), - cache=_semantic_cache(supported_call_types=("acompletion",)), + data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type=route, gates=(gate,), cache=_semantic_cache() ) assert result is None @@ -132,14 +128,17 @@ async def test_the_gate_is_never_consulted_for_a_call_type_the_cache_does_not_se @pytest.mark.parametrize( ("route", "served_call_type"), [ - ("aresponses", "responses"), ("aresponses", "acompletion"), - ("anthropic_messages", "acompletion"), + ("acompletion", "aresponses"), + ("anthropic_messages", "aresponses"), ("agenerate_content", "acompletion"), ("agenerate_content_stream", "acompletion"), + ("acreate_interaction", "aresponses"), + ("_aresponses_websocket", "aresponses"), + ("asend_message", "acompletion"), ], ) -async def test_the_gate_covers_routes_that_reach_a_served_call_type_through_an_alias_or_a_bridge( +async def test_the_gate_covers_routes_that_reach_the_cache_through_a_bridge_whatever_call_types_it_serves( route: str, served_call_type: CachingSupportedCallTypes ): gate = _TeamAllowlist(frozenset()) @@ -194,7 +193,11 @@ async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: p gates = ProxyLogging.callback_capabilities().semantic_cache_gates result = await gated_cache_controls( - data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=gates, cache=_semantic_cache() + data=_request(), + user_api_key_dict=_key(OTHER_TEAM), + call_type="acompletion", + gates=gates, + cache=_semantic_cache(), ) assert gates == (gate,) From 34bbeaaa381292b1c4aa125757a768783c63556b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:54:16 -0700 Subject: [PATCH 5/8] fix(proxy): run the semantic cache check on MCP sampling, merge A2A cache controls with the agent's own, and drop routes that never reach the cache --- .../mcp_server/sampling_handler.py | 16 +++++++-- .../proxy/agent_endpoints/a2a_endpoints.py | 14 +++++--- .../proxy/common_utils/semantic_cache_gate.py | 2 -- .../test_mcp_sampling_completion_flow.py | 36 +++++++++++++++++++ .../agent_endpoints/test_a2a_endpoints.py | 10 ++++-- .../common_utils/test_semantic_cache_gate.py | 2 +- 6 files changed, 68 insertions(+), 12 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py index e1469f2fff7..64f5530fa5c 100644 --- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py +++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py @@ -37,6 +37,8 @@ from fastapi import HTTPException from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.utils import ProxyLogging # Guard imports that require the mcp package try: @@ -1125,14 +1127,22 @@ async def _run_guardrails_and_call_llm( import litellm + cache_controls: Final = await gated_cache_controls( + data=completion_kwargs, + user_api_key_dict=user_api_key_auth, + call_type="acompletion", + gates=ProxyLogging.callback_capabilities().semantic_cache_gates, + cache=litellm.cache, + ) + call_kwargs: Final = completion_kwargs if cache_controls is None else {**completion_kwargs, "cache": cache_controls} try: from litellm.proxy.proxy_server import llm_router if llm_router is not None: - return await _AcompletionCall(fn=llm_router.acompletion).fn(**completion_kwargs) - return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs) + return await _AcompletionCall(fn=llm_router.acompletion).fn(**call_kwargs) + return await _AcompletionCall(fn=litellm.acompletion).fn(**call_kwargs) except ImportError: - return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs) + return await _AcompletionCall(fn=litellm.acompletion).fn(**call_kwargs) async def handle_sampling_create_message( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index f883d49532a..6a0f4462160 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -180,9 +180,12 @@ async def _resolve_backend_auth( ) -def _request_cache_controls(data: Mapping[str, object]) -> Mapping[str, object]: - cache_controls: Final = data.get("cache") - return {"cache": cache_controls} if isinstance(cache_controls, Mapping) else {} +def _request_cache_controls(agent_params: Mapping[str, object], data: Mapping[str, object]) -> Mapping[str, object]: + request_cache: Final = data.get("cache") + if not isinstance(request_cache, Mapping): + return {} + agent_cache: Final = agent_params.get("cache") + return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **request_cache}} def _forwarding_headers( @@ -843,7 +846,10 @@ async def invoke_agent_a2a( ), backend_auth_header=backend_auth.header, ) - forwarded_litellm_params: Final = {**backend_auth.litellm_params, **_request_cache_controls(data)} + forwarded_litellm_params: Final = { + **backend_auth.litellm_params, + **_request_cache_controls(backend_auth.litellm_params, data), + } # Merge agent-level guardrails into data so post_call_success_hook and # _handle_stream_message both pick them up. A2A agents use model diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py index b080c166f2d..8b45ce58031 100644 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -13,8 +13,6 @@ SEMANTIC_CACHE_ROUTES: Final = frozenset( "acompletion", "atext_completion", "aembedding", - "atranscription", - "arerank", "aresponses", "_aresponses_websocket", "anthropic_messages", diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py index d17b407a1be..30b9617c42d 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py +++ b/tests/unit/proxy/_experimental/mcp_server/test_mcp_sampling_completion_flow.py @@ -7,6 +7,7 @@ through the proxy router / guardrails, and the end-to-end `handle_sampling_create_message` success and error-propagation behaviour. """ +from collections.abc import Mapping from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -14,6 +15,12 @@ import pytest from mcp.types import CreateMessageResult, ErrorData +import litellm +from litellm.caching.caching import Cache +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.caching import LiteLLMCacheType + from litellm.proxy._experimental.mcp_server.sampling_handler import ( _build_completion_kwargs, _run_guardrails_and_call_llm, @@ -119,6 +126,35 @@ class TestRunGuardrailsAndCallLlm: assert result == "router-response" router.acompletion.assert_awaited_once() + async def test_should_skip_the_semantic_cache_for_a_key_the_semantic_cache_check_rejects(self): + class _RejectEveryTeam(CustomLogger): + async def async_should_semantic_cache( + self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str + ) -> bool: + return False + + async def _cache_controls_the_router_received(**kwargs): + return kwargs.get("cache") + + semantic_cache = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 + ) + with ( + patch("litellm.proxy.proxy_server.proxy_logging_obj", None), + patch( + "litellm.proxy.proxy_server.llm_router", + SimpleNamespace(acompletion=_cache_controls_the_router_received), + ), + patch.object(litellm, "callbacks", [_RejectEveryTeam()]), + patch.object(litellm, "cache", semantic_cache), + ): + received = await _run_guardrails_and_call_llm( + completion_kwargs={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), + ) + + assert received == {"no-cache": True, "no-store": True} + async def test_should_propagate_guardrail_rejection(self): plo = MagicMock() plo.pre_call_hook = AsyncMock(side_effect=ValueError("blocked by guardrail")) diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 223b247a779..011df4f8daf 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -643,6 +643,8 @@ class _RejectEveryTeam(CustomLogger): @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent_call(method: str): + agent = _make_agent_mock() + agent.litellm_params = {"cache": {"ttl": 60}} mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) semantic_cache = Cache( type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 @@ -650,10 +652,14 @@ async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent with patch.object(litellm, "callbacks", [_RejectEveryTeam()]), patch.object(litellm, "cache", semantic_cache): captured = await _invoke_message_method( - method, mock_request, UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), _add_proxy_data + method, + mock_request, + UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), + _add_proxy_data, + agent=agent, ) - assert (captured.litellm_params or {}).get("cache") == {"no-cache": True, "no-store": True} + assert (captured.litellm_params or {}).get("cache") == {"ttl": 60, "no-cache": True, "no-store": True} @pytest.mark.asyncio diff --git a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py index 9c95d28973d..eef95c5e589 100644 --- a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py +++ b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py @@ -112,7 +112,7 @@ async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache @pytest.mark.asyncio -@pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech"]) +@pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech", "atranscription", "arerank"]) async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str): gate = _TeamAllowlist(frozenset()) From d0a9e2181c1063aecb60a80889522804c6602a48 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:08:31 -0700 Subject: [PATCH 6/8] fix(caching): forward bridged cache controls only to a semantic cache, and only a caller's opt-outs on A2A The generateContent and A2A completion bridges forwarded the request's cache controls for every cache type, so an exact-match cache started honoring them. Both now forward only when the proxy's cache is semantic. The A2A bridge merged every request control over the agent's own cache setting, so a caller could name another agent's namespace and read its entries. It now forwards only a no-cache or no-store the request sets. The generateContent cache test waits on the pending cache writes instead of sleeping, and the pre-call gate test runs the real request helpers. --- litellm/google_genai/adapters/handler.py | 2 +- .../proxy/agent_endpoints/a2a_endpoints.py | 15 +++-- .../google_genai/test_google_genai_adapter.py | 55 ++++++++++++++++--- .../agent_endpoints/test_a2a_endpoints.py | 40 +++++++++++++- .../proxy/test_common_request_processing.py | 36 +++++------- 5 files changed, 113 insertions(+), 35 deletions(-) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index e386008f5d3..f0c8fbab106 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -69,7 +69,7 @@ class GenerateContentToCompletionHandler: completion_kwargs["proxy_server_request"] = extra_kwargs["proxy_server_request"] if "litellm_metadata" in extra_kwargs: completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"] - if "cache" in extra_kwargs: + if "cache" in extra_kwargs and litellm.cache is not None and litellm.cache.is_semantic_cache(): completion_kwargs["cache"] = extra_kwargs["cache"] if stream: diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 6a0f4462160..20cc9844920 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -23,6 +23,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.caching import Cache from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth @@ -40,6 +41,7 @@ from litellm.proxy.agent_endpoints.databricks_oauth import ( ) from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.semantic_cache_gate import SKIP_CACHE_CONTROLS from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING, coerce_keepalive_interval, @@ -180,12 +182,17 @@ async def _resolve_backend_auth( ) -def _request_cache_controls(agent_params: Mapping[str, object], data: Mapping[str, object]) -> Mapping[str, object]: +def _request_cache_controls( + agent_params: Mapping[str, object], data: Mapping[str, object], cache: Cache | None +) -> Mapping[str, object]: request_cache: Final = data.get("cache") - if not isinstance(request_cache, Mapping): + if cache is None or not cache.is_semantic_cache() or not isinstance(request_cache, Mapping): + return {} + bypass: Final = {key: value for key, value in SKIP_CACHE_CONTROLS.items() if request_cache.get(key) is value} + if not bypass: return {} agent_cache: Final = agent_params.get("cache") - return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **request_cache}} + return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **bypass}} def _forwarding_headers( @@ -848,7 +855,7 @@ async def invoke_agent_a2a( ) forwarded_litellm_params: Final = { **backend_auth.litellm_params, - **_request_cache_controls(backend_auth.litellm_params, data), + **_request_cache_controls(backend_auth.litellm_params, data, litellm.cache), } # Merge agent-level guardrails into data so post_call_success_hook and diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index dbcc80a5d25..4add065b36f 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -5,8 +5,11 @@ Test to verify the Google GenAI generate_content adapter functionality import asyncio import json import unittest +from collections.abc import Mapping +from typing import Final import pytest +import respx from litellm.google_genai.main import agenerate_content @@ -1645,15 +1648,26 @@ async def test_generate_content_sends_response_schema_and_tool_parameters_to_the @pytest.mark.asyncio -async def test_generate_content_bridge_honors_the_request_cache_controls(respx_mock, monkeypatch): +@pytest.mark.parametrize( + ("first_cache_controls", "second_cache_controls"), + [({"no-store": True}, None), (None, {"no-cache": True})], + ids=["no-store", "no-cache"], +) +async def test_generate_content_bridge_leaves_an_exact_match_cache_as_it_was( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + first_cache_controls: Mapping[str, bool] | None, + second_cache_controls: Mapping[str, bool] | None, +) -> None: import httpx from litellm.caching.caching import Cache + from litellm.caching.caching_handler import _PENDING_CACHE_WRITES from litellm.types.caching import LiteLLMCacheType monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL)) - route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( 200, json={ @@ -1666,16 +1680,41 @@ async def test_generate_content_bridge_honors_the_request_cache_controls(respx_m }, ) ) - request = { + request: Final = { "model": "openai/gpt-4.1", "contents": [{"role": "user", "parts": [{"text": "What is the capital of France?"}]}], - "cache": {"no-cache": True, "no-store": True}, "api_key": "sk-test", } - await agenerate_content(**request) - await asyncio.sleep(0.5) - await agenerate_content(**request) + await agenerate_content(**request, **({} if first_cache_controls is None else {"cache": first_cache_controls})) + await asyncio.gather(*_PENDING_CACHE_WRITES) + await agenerate_content(**request, **({} if second_cache_controls is None else {"cache": second_cache_controls})) - assert route.call_count == 2 + assert route.call_count == 1 + + +@pytest.mark.parametrize( + ("semantic", "expected_cache"), + [(True, {"no-cache": True, "no-store": True}), (False, None)], + ids=["semantic", "exact-match"], +) +def test_generate_content_bridge_forwards_the_request_cache_controls_only_to_a_semantic_cache( + monkeypatch: pytest.MonkeyPatch, semantic: bool, expected_cache: Mapping[str, bool] | None +) -> None: + from litellm.caching.caching import Cache + from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler + from litellm.types.caching import LiteLLMCacheType + + cache_type: Final = LiteLLMCacheType.VALKEY_SEMANTIC if semantic else LiteLLMCacheType.LOCAL + monkeypatch.setattr( + litellm, "cache", Cache(type=cache_type, host="localhost", port="6379", similarity_threshold=0.8) + ) + + completion_kwargs: Final = GenerateContentToCompletionHandler._prepare_completion_kwargs( + model="openai/gpt-4.1", + contents=[{"role": "user", "parts": [{"text": "What is the capital of France?"}]}], + extra_kwargs={"cache": {"no-cache": True, "no-store": True}}, + ) + + assert completion_kwargs.get("cache") == expected_cache diff --git a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py index 011df4f8daf..c7c501149f6 100644 --- a/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/unit/proxy/agent_endpoints/test_a2a_endpoints.py @@ -386,7 +386,12 @@ def _make_agent_mock(url: str = "http://backend-agent:10001") -> MagicMock: return agent -def _make_request_mock(method: str, params: Mapping[str, object], request_id: object = "req-1") -> MagicMock: +def _make_request_mock( + method: str, + params: Mapping[str, object], + request_id: object = "req-1", + cache: Mapping[str, object] | None = None, +) -> MagicMock: req = MagicMock() req.headers = {} req.json = AsyncMock( @@ -395,6 +400,7 @@ def _make_request_mock(method: str, params: Mapping[str, object], request_id: ob "id": request_id, "method": method, "params": params, + **({} if cache is None else {"cache": cache}), } ) return req @@ -662,6 +668,38 @@ async def test_message_methods_hand_the_semantic_cache_check_result_to_the_agent assert (captured.litellm_params or {}).get("cache") == {"ttl": 60, "no-cache": True, "no-store": True} +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +@pytest.mark.parametrize( + ("cache_type", "expected_agent_cache"), + [ + (LiteLLMCacheType.VALKEY_SEMANTIC, {"namespace": "team-a", "no-store": True}), + (LiteLLMCacheType.LOCAL, {"namespace": "team-a"}), + ], + ids=["semantic", "exact-match"], +) +async def test_message_methods_forward_only_the_callers_semantic_cache_opt_outs_to_the_agent_call( + method: str, cache_type: LiteLLMCacheType, expected_agent_cache: Mapping[str, object] +) -> None: + agent: Final = _make_agent_mock() + agent.litellm_params = {"cache": {"namespace": "team-a"}} + mock_request: Final = _make_request_mock( + method, _HELLO_MESSAGE_PARAMS, cache={"namespace": "team-b", "use-cache": True, "ttl": 1, "no-store": True} + ) + cache: Final = Cache(type=cache_type, host="localhost", port="6379", similarity_threshold=0.8) + + with patch.object(litellm, "cache", cache): + captured: Final = await _invoke_message_method( + method, + mock_request, + UserAPIKeyAuth(api_key="sk-test", team_id="team-a"), + _add_proxy_data, + agent=agent, + ) + + assert (captured.litellm_params or {}).get("cache") == expected_agent_cache + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_forward_caller_identity_headers(method: str): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 700c4958a69..8e52a263bc5 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -525,6 +525,7 @@ class TestProxyBaseLLMRequestProcessing: expected_cache_controls: dict[str, object], ) -> None: from litellm.caching.caching import Cache + from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig from litellm.types.caching import LiteLLMCacheType class _Gate(CustomLogger): @@ -533,37 +534,30 @@ class TestProxyBaseLLMRequestProcessing: ) -> bool: return gate_allows - async def passthrough_pre_call_hook( - user_api_key_dict: ProxyUserAPIKeyAuth, - data: dict[str, object], - call_type: str, - skip_guardrails: bool = False, - ) -> dict[str, object]: - return data - monkeypatch.setattr(litellm, "callbacks", [_Gate()]) monkeypatch.setattr( litellm, "cache", Cache(type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8), ) - body: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} - monkeypatch.setattr( - litellm.proxy.common_request_processing, - "add_litellm_data_to_request", - AsyncMock(return_value={**body, "metadata": {}, "proxy_server_request": {"body": body}}), + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + } ) - mock_request: Final = MagicMock(spec=Request) - mock_request.headers = {} - mock_proxy_logging_obj: Final = MagicMock(spec=ProxyLogging) - mock_proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=passthrough_pre_call_hook) - returned_data, _ = await ProxyBaseLLMRequestProcessing(data={}).common_processing_pre_call_logic( - request=mock_request, + returned_data, _ = await ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} + ).common_processing_pre_call_logic( + request=request, general_settings={}, user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-hashed", team_id="team-other"), - proxy_logging_obj=mock_proxy_logging_obj, - proxy_config=MagicMock(spec=ProxyConfig), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + proxy_config=RealProxyConfig(), route_type="acompletion", ) From dbe2fc909499c655ac899ecc7c26b686b7920078 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:17:47 -0700 Subject: [PATCH 7/8] fix(proxy): keep a rejected caller off the semantic cache on fallback hops and late batch rows A fallback entry's keys are copied over the request on that hop, so the check's no now also adds no-cache and no-store to the cache inside each fallback entry the request carries, in its own lists and in router_settings_override. Batch rows read the cache and the callbacks as each row runs, and the MCP sampling test runs through a real router. --- litellm/constants.py | 1 + litellm/proxy/batches_endpoints/endpoints.py | 4 +- .../litellm_executed_batches.py | 14 +- litellm/proxy/common_request_processing.py | 21 ++- .../proxy/common_utils/semantic_cache_gate.py | 49 ++++++- litellm/proxy/litellm_pre_call_utils.py | 6 +- litellm/proxy/mcp_server/sampling_handler.py | 20 +-- .../test_litellm_executed_batches.py | 28 +++- .../common_utils/test_semantic_cache_gate.py | 126 +++++++++++++----- .../test_mcp_sampling_completion_flow.py | 38 ++++-- .../proxy/test_common_request_processing.py | 16 ++- 11 files changed, 228 insertions(+), 95 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2a20e812e68..496f82b3264 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -54,6 +54,7 @@ AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPON AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"}) ROUTER_MAX_FALLBACKS: Final = get_env_int("ROUTER_MAX_FALLBACKS", 5) ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 +REQUEST_FALLBACK_KEYS: Final = ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks") RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( { "routing_strategy_args", diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5fbb240ea01..5ae1107e3b7 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -110,8 +110,8 @@ def _litellm_executed_batch_runner(llm_router: Router, proxy_logging_obj: ProxyL batches=ManagedBatchRepository(prisma_client), proxy_logging_obj=proxy_logging_obj, general_settings=general_settings, - semantic_cache_gates=ProxyLogging.callback_capabilities().semantic_cache_gates, - semantic_cache=litellm.cache, + semantic_cache_gates=lambda: ProxyLogging.callback_capabilities().semantic_cache_gates, + semantic_cache=lambda: litellm.cache, ) diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 1299c6dda6a..ffe803c24d8 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.models.managed_files import LiteLLM_ManagedFileTable from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_utils import is_request_body_safe -from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_overrides from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.openai_files_endpoints.common_utils import LITELLM_EXECUTED_BATCH_ID_PREFIX from litellm.proxy.openai_files_endpoints.storage_backend_service import StorageBackendFileService @@ -426,8 +426,8 @@ class LiteLLMExecutedBatchRunner: batches: ManagedBatchRepository, proxy_logging_obj: ProxyLogging, general_settings: Mapping[str, object], - semantic_cache_gates: Sequence[CustomLogger], - semantic_cache: Cache | None, + semantic_cache_gates: Callable[[], Sequence[CustomLogger]], + semantic_cache: Callable[[], Cache | None], concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY, heartbeat_seconds: float = _HEARTBEAT_SECONDS, completion_window_seconds: float = _COMPLETION_WINDOW_SECONDS, @@ -654,14 +654,14 @@ class LiteLLMExecutedBatchRunner: request: Final = MappingProxyType( {**line.body, "model": run.model, "metadata": self._row_metadata(run), "disable_fallbacks": True} ) - cache_controls: Final = await gated_cache_controls( + cache_overrides: Final = await gated_cache_overrides( data=request, user_api_key_dict=run.user_api_key_dict, call_type=_ROUTER_METHODS[run.endpoint], - gates=self.semantic_cache_gates, - cache=self.semantic_cache, + gates=self.semantic_cache_gates(), + cache=self.semantic_cache(), ) - params: Final = MappingProxyType({**request, **({} if cache_controls is None else {"cache": cache_controls})}) + params: Final = MappingProxyType({**request, **cache_overrides}) return _dump(await self._router_call(run.endpoint)(**params)) def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8b5c9a64a01..230d74bc731 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -114,7 +114,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_param, openai_error_type, ) -from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_overrides from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -2511,16 +2511,15 @@ class ProxyBaseLLMRequestProcessing: call_type=route_type, skip_guardrails=skip_guardrails, ) - if ( - cache_controls := await gated_cache_controls( - data=self.data, - user_api_key_dict=user_api_key_dict, - call_type=route_type, - gates=ProxyLogging.callback_capabilities().semantic_cache_gates, - cache=litellm.cache, - ) - ) is not None: - self.data["cache"] = cache_controls + cache_overrides: Final = await gated_cache_overrides( + data=self.data, + user_api_key_dict=user_api_key_dict, + call_type=route_type, + gates=ProxyLogging.callback_capabilities().semantic_cache_gates, + cache=litellm.cache, + ) + for key, value in cache_overrides.items(): + self.data[key] = value await _enforce_guardrail_added_tag_budgets( data=self.data, tags_before_guardrails=self._tags_before_guardrails, diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py index 8b45ce58031..b1c170ce45f 100644 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -2,8 +2,11 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.caching.caching import Cache +from litellm.constants import REQUEST_FALLBACK_KEYS from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -22,6 +25,8 @@ SEMANTIC_CACHE_ROUTES: Final = frozenset( "asend_message", } ) +_FALLBACK_ENTRIES: Final = TypeAdapter(list[object]) +_FIELDS: Final = TypeAdapter(Mapping[str, object]) async def _gate_allows( @@ -51,16 +56,48 @@ async def _every_gate_allows( return True -async def gated_cache_controls( +def _skipping(controls: object) -> Mapping[str, object]: + return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} + + +def _skipping_fallback_entry(entry: object) -> object: + if not isinstance(entry, Mapping): + return entry + fields: Final = _FIELDS.validate_python(entry) + if "cache" not in fields: + return fields + return {**fields, "cache": _skipping(fields["cache"])} + + +def _skipping_fallbacks(data: Mapping[str, object]) -> Mapping[str, Sequence[object]]: + return { + key: [_skipping_fallback_entry(entry) for entry in _FALLBACK_ENTRIES.validate_python(entries)] + for key in REQUEST_FALLBACK_KEYS + if isinstance(entries := data.get(key), list) + } + + +def _skipping_router_settings_override(data: Mapping[str, object]) -> Mapping[str, Mapping[str, object]]: + override: Final = data.get("router_settings_override") + if not isinstance(override, Mapping): + return {} + settings: Final = _FIELDS.validate_python(override) + return {"router_settings_override": {**settings, **_skipping_fallbacks(settings)}} + + +async def gated_cache_overrides( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str, gates: Sequence[CustomLogger], cache: Cache | None, -) -> Mapping[str, object] | None: +) -> Mapping[str, object]: if not gates or cache is None or not cache.is_semantic_cache() or call_type not in SEMANTIC_CACHE_ROUTES: - return None + return {} if await _every_gate_allows(gates, user_api_key_dict, data, call_type): - return None - controls: Final = data.get("cache") - return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} + return {} + return { + "cache": _skipping(data.get("cache")), + **_skipping_fallbacks(data), + **_skipping_router_settings_override(data), + } diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 24a31723faa..b22049690a1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -39,6 +39,7 @@ from litellm.constants import ( MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + REQUEST_FALLBACK_KEYS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -3076,7 +3077,6 @@ async def add_litellm_data_to_request( return data -_REQUEST_FALLBACK_KEYS: Final = ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks") _FALLBACK_DISCOVERY_LIMIT: Final = 256 @@ -3150,7 +3150,7 @@ def _fallback_lists( ): if isinstance(router_list, list): lists.append(cast("Sequence[object]", router_list)) # cast-ok: fallback lists hold mixed entry shapes - for key in _REQUEST_FALLBACK_KEYS: + for key in REQUEST_FALLBACK_KEYS: entries: object = data.get(key) if isinstance(entries, list): lists.append(cast("Sequence[object]", entries)) # cast-ok: request fallback lists hold mixed entry shapes @@ -3403,7 +3403,7 @@ async def _resolve_user_provider_credentials_for_request( _fallback_entry_target_count(entry) for entry in cast("Sequence[object]", entries) # cast-ok: request fallback lists hold mixed entry shapes ) - for key in _REQUEST_FALLBACK_KEYS + for key in REQUEST_FALLBACK_KEYS if isinstance(entries := data.get(key), list) ) if request_fallback_targets > _FALLBACK_DISCOVERY_LIMIT: diff --git a/litellm/proxy/mcp_server/sampling_handler.py b/litellm/proxy/mcp_server/sampling_handler.py index 5af94a93871..efbfa152fa6 100644 --- a/litellm/proxy/mcp_server/sampling_handler.py +++ b/litellm/proxy/mcp_server/sampling_handler.py @@ -37,7 +37,7 @@ from fastapi import HTTPException from pydantic import TypeAdapter from litellm._logging import verbose_logger -from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_overrides from litellm.proxy.utils import ProxyLogging # Guard imports that require the mcp package @@ -1127,14 +1127,16 @@ async def _run_guardrails_and_call_llm( import litellm - cache_controls: Final = await gated_cache_controls( - data=completion_kwargs, - user_api_key_dict=user_api_key_auth, - call_type="acompletion", - gates=ProxyLogging.callback_capabilities().semantic_cache_gates, - cache=litellm.cache, - ) - call_kwargs: Final = completion_kwargs if cache_controls is None else {**completion_kwargs, "cache": cache_controls} + call_kwargs: Final = { + **completion_kwargs, + **await gated_cache_overrides( + data=completion_kwargs, + user_api_key_dict=user_api_key_auth, + call_type="acompletion", + gates=ProxyLogging.callback_capabilities().semantic_cache_gates, + cache=litellm.cache, + ), + } try: from litellm.proxy.proxy_server import llm_router diff --git a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py index 1e62563f7e3..51a9e7c5eac 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -383,8 +383,8 @@ def make_runner( heartbeat_seconds: float = 30.0, completion_window_seconds: float = 24 * 60 * 60, user: UserAPIKeyAuth | None = None, - semantic_cache_gates: Sequence[CustomLogger] = (), - semantic_cache: Cache | None = None, + semantic_cache_gates: Callable[[], Sequence[CustomLogger]] = lambda: (), + semantic_cache: Callable[[], Cache | None] = lambda: None, ) -> Harness: store = store_factory({INPUT_FILE_ID: managed_input_file()} if files is None else files) router = FakeRouter() @@ -722,6 +722,10 @@ class _TeamAllowlist(CustomLogger): return user_api_key_dict.team_id == self.team_id +def _semantic_cache() -> Cache: + return Cache(type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8) + + @pytest.mark.parametrize( ("allowed_team", "row_cache_controls"), [("team-1", None), ("team-other", {"no-cache": True, "no-store": True})], @@ -731,10 +735,8 @@ async def test_create_runs_the_semantic_cache_check_on_each_row( allowed_team: str, row_cache_controls: Mapping[str, object] | None ) -> None: gate = _TeamAllowlist(allowed_team) - semantic_cache = Cache( - type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 - ) - harness = make_runner(semantic_cache_gates=(gate,), semantic_cache=semantic_cache) + semantic_cache: Final = _semantic_cache() + harness: Final = make_runner(semantic_cache_gates=lambda: (gate,), semantic_cache=lambda: semantic_cache) await harness.create_and_finish() calls = harness.router.acompletion.await_args_list @@ -745,6 +747,20 @@ async def test_create_runs_the_semantic_cache_check_on_each_row( ] +async def test_rows_follow_the_semantic_cache_and_check_in_place_when_they_run_not_when_the_batch_was_created() -> None: + live_gates: Final[list[Sequence[CustomLogger]]] = [()] # mutable-ok: stands in for callbacks an admin adds later + live_cache: Final[list[Cache | None]] = [None] # mutable-ok: stands in for a cache an admin turns on later + harness: Final = make_runner(semantic_cache_gates=lambda: live_gates[-1], semantic_cache=lambda: live_cache[-1]) + await harness.create() + + live_gates.append((_TeamAllowlist("team-other"),)) + live_cache.append(_semantic_cache()) + await asyncio.gather(*list(litellm_executed_batches._RUNNING_BATCHES)) + + calls: Final = harness.router.acompletion.await_args_list + assert [call.kwargs.get("cache") for call in calls] == [{"no-cache": True, "no-store": True}] * 2 + + async def test_create_dispatches_each_row_with_the_key_team_and_project_tags_as_inherited_tags() -> None: user = UserAPIKeyAuth( api_key="sk-batch-key", diff --git a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py index eef95c5e589..8fe90c9d263 100644 --- a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py +++ b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from typing import Final import pytest @@ -6,7 +7,8 @@ import litellm from litellm.caching.caching import Cache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_overrides +from litellm.proxy.route_llm_request import _merge_router_settings_override from litellm.proxy.utils import ProxyLogging from litellm.types.caching import ( DEFAULT_CACHING_SUPPORTED_CALL_TYPES, @@ -57,7 +59,7 @@ def _semantic_cache( ) -def _request(**cache_controls: object) -> dict[str, object]: +def _request(**cache_controls: object) -> Mapping[str, object]: return { "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "What color is the sky?"}], @@ -70,23 +72,23 @@ def _key(team_id: str) -> UserAPIKeyAuth: @pytest.mark.asyncio -async def test_a_team_the_gate_rejects_skips_semantic_lookup_and_store_even_when_the_client_opts_in(): - gate = _TeamAllowlist(frozenset({ALLOWED_TEAM})) - data = _request(**{"use-cache": True, "ttl": 60}) +async def test_a_team_the_gate_rejects_skips_semantic_lookup_and_store_even_when_the_client_opts_in() -> None: + gate: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + data: Final = _request(**{"use-cache": True, "ttl": 60}) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=_semantic_cache() ) - assert result == {"use-cache": True, "ttl": 60, "no-cache": True, "no-store": True} + assert result == {"cache": {"use-cache": True, "ttl": 60, "no-cache": True, "no-store": True}} assert gate.seen == [(OTHER_TEAM, "gpt-4o-mini", "acompletion")] @pytest.mark.asyncio -async def test_a_team_the_gate_allows_keeps_its_own_cache_controls(): - data = _request(ttl=60) +async def test_a_team_the_gate_allows_keeps_its_own_cache_controls() -> None: + data: Final = _request(ttl=60) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=data, user_api_key_dict=_key(ALLOWED_TEAM), call_type="acompletion", @@ -94,33 +96,33 @@ async def test_a_team_the_gate_allows_keeps_its_own_cache_controls(): cache=_semantic_cache(), ) - assert result is None + assert result == {} @pytest.mark.asyncio @pytest.mark.parametrize("cache", [None, Cache(type=LiteLLMCacheType.LOCAL)], ids=["no-cache", "exact-match-cache"]) -async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache | None): - gate = _TeamAllowlist(frozenset()) - data = _request() +async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache | None) -> None: + gate: Final = _TeamAllowlist(frozenset()) + data: Final = _request() - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=cache ) - assert result is None + assert result == {} assert gate.seen == [] @pytest.mark.asyncio @pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech", "atranscription", "arerank"]) -async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str): - gate = _TeamAllowlist(frozenset()) +async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str) -> None: + gate: Final = _TeamAllowlist(frozenset()) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type=route, gates=(gate,), cache=_semantic_cache() ) - assert result is None + assert result == {} assert gate.seen == [] @@ -140,10 +142,10 @@ async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cac ) async def test_the_gate_covers_routes_that_reach_the_cache_through_a_bridge_whatever_call_types_it_serves( route: str, served_call_type: CachingSupportedCallTypes -): - gate = _TeamAllowlist(frozenset()) +) -> None: + gate: Final = _TeamAllowlist(frozenset()) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type=route, @@ -151,14 +153,14 @@ async def test_the_gate_covers_routes_that_reach_the_cache_through_a_bridge_what cache=_semantic_cache(supported_call_types=(served_call_type,)), ) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"no-cache": True, "no-store": True}} assert gate.seen == [(OTHER_TEAM, "gpt-4o-mini", route)] @pytest.mark.asyncio @pytest.mark.parametrize("gate", [_RaisingGate(), _TruthyGate()], ids=["raises", "returns-non-bool"]) -async def test_a_gate_that_fails_or_answers_anything_but_true_skips_the_semantic_cache(gate: CustomLogger): - result = await gated_cache_controls( +async def test_a_gate_that_fails_or_answers_anything_but_true_skips_the_semantic_cache(gate: CustomLogger) -> None: + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(ALLOWED_TEAM), call_type="aresponses", @@ -166,16 +168,16 @@ async def test_a_gate_that_fails_or_answers_anything_but_true_skips_the_semantic cache=_semantic_cache(), ) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"no-cache": True, "no-store": True}} @pytest.mark.asyncio @pytest.mark.parametrize("deny_first", [True, False]) -async def test_one_rejecting_gate_skips_the_semantic_cache_whatever_the_others_answer(deny_first: bool): - allow = _TeamAllowlist(frozenset({ALLOWED_TEAM})) - deny = _TeamAllowlist(frozenset()) +async def test_one_rejecting_gate_skips_the_semantic_cache_whatever_the_others_answer(deny_first: bool) -> None: + allow: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + deny: Final = _TeamAllowlist(frozenset()) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(ALLOWED_TEAM), call_type="anthropic_messages", @@ -183,16 +185,16 @@ async def test_one_rejecting_gate_skips_the_semantic_cache_whatever_the_others_a cache=_semantic_cache(), ) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"no-cache": True, "no-store": True}} @pytest.mark.asyncio -async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: pytest.MonkeyPatch): - gate = _TeamAllowlist(frozenset({ALLOWED_TEAM})) +async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: pytest.MonkeyPatch) -> None: + gate: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM})) monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), gate]) - gates = ProxyLogging.callback_capabilities().semantic_cache_gates - result = await gated_cache_controls( + gates: Final = ProxyLogging.callback_capabilities().semantic_cache_gates + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", @@ -201,4 +203,54 @@ async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: p ) assert gates == (gate,) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"no-cache": True, "no-store": True}} + + +@pytest.mark.asyncio +async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_its_own_fallback_entries() -> None: + data: Final = { + **_request(ttl=60), + "fallbacks": [{"model": "gpt-4o", "cache": {"ttl": 5}}, "gpt-4.1-mini", {"model": "gpt-4.1"}], + "context_window_fallbacks": [{"gpt-4o-mini": ["gpt-4o"], "cache": {}}], + "content_policy_fallbacks": [{"model": "gpt-4o", "cache": None}], + } + + result: Final = await gated_cache_overrides( + data=data, + user_api_key_dict=_key(OTHER_TEAM), + call_type="acompletion", + gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),), + cache=_semantic_cache(), + ) + + assert result == { + "cache": {"ttl": 60, "no-cache": True, "no-store": True}, + "fallbacks": [ + {"model": "gpt-4o", "cache": {"ttl": 5, "no-cache": True, "no-store": True}}, + "gpt-4.1-mini", + {"model": "gpt-4.1"}, + ], + "context_window_fallbacks": [{"gpt-4o-mini": ["gpt-4o"], "cache": {"no-cache": True, "no-store": True}}], + "content_policy_fallbacks": [{"model": "gpt-4o", "cache": {"no-cache": True, "no-store": True}}], + } + + +@pytest.mark.asyncio +async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_router_settings_fallbacks() -> None: + data: Final = { + **_request(), + "router_settings_override": {"fallbacks": [{"model": "gpt-4o", "cache": {}}], "num_retries": 2}, + } + + overrides: Final = await gated_cache_overrides( + data=data, + user_api_key_dict=_key(OTHER_TEAM), + call_type="acompletion", + gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),), + cache=_semantic_cache(), + ) + routed: Final = {**data, **overrides} + _ = _merge_router_settings_override(routed) + + assert routed["fallbacks"] == [{"model": "gpt-4o", "cache": {"no-cache": True, "no-store": True}}] + assert routed["num_retries"] == 2 diff --git a/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py b/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py index 549fb99213e..908823f255f 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py +++ b/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py @@ -9,6 +9,7 @@ through the proxy router / guardrails, and the end-to-end from collections.abc import Mapping from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -18,7 +19,9 @@ import litellm from litellm.caching.caching import Cache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth +from litellm.router import Router from litellm.types.caching import LiteLLMCacheType +from litellm.types.utils import CallTypes from litellm.proxy.mcp_server.sampling_handler import ( _build_completion_kwargs, @@ -125,34 +128,47 @@ class TestRunGuardrailsAndCallLlm: assert result == "router-response" router.acompletion.assert_awaited_once() - async def test_should_skip_the_semantic_cache_for_a_key_the_semantic_cache_check_rejects(self): + async def test_should_skip_the_semantic_cache_for_a_key_the_semantic_cache_check_rejects(self) -> None: class _RejectEveryTeam(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.deployment_cache_controls: list[object] = [] + async def async_should_semantic_cache( self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str ) -> bool: return False - async def _cache_controls_the_router_received(**kwargs): - return kwargs.get("cache") + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.deployment_cache_controls.append(kwargs.get("cache")) + return None - semantic_cache = Cache( + gate: Final = _RejectEveryTeam() + router: Final = Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + semantic_cache: Final = Cache( type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 ) with ( patch("litellm.proxy.proxy_server.proxy_logging_obj", None), - patch( - "litellm.proxy.proxy_server.llm_router", - SimpleNamespace(acompletion=_cache_controls_the_router_received), - ), - patch.object(litellm, "callbacks", [_RejectEveryTeam()]), + patch("litellm.proxy.proxy_server.llm_router", router), + patch.object(litellm, "callbacks", [gate]), patch.object(litellm, "cache", semantic_cache), ): - received = await _run_guardrails_and_call_llm( + await _run_guardrails_and_call_llm( completion_kwargs={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), ) - assert received == {"no-cache": True, "no-store": True} + assert gate.deployment_cache_controls == [{"no-cache": True, "no-store": True}] async def test_should_propagate_guardrail_rejection(self): plo = MagicMock() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5263b0592bd..7e8db33095f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -515,8 +515,11 @@ class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio @pytest.mark.parametrize( - ("gate_allows", "expected_cache_controls"), - [(True, {"ttl": 60}), (False, {"ttl": 60, "no-cache": True, "no-store": True})], + ("gate_allows", "expected_cache_controls", "expected_fallback_cache_controls"), + [ + (True, {"ttl": 60}, {}), + (False, {"ttl": 60, "no-cache": True, "no-store": True}, {"no-cache": True, "no-store": True}), + ], ids=["allowed", "rejected"], ) async def test_common_processing_pre_call_logic_applies_the_semantic_cache_gate( @@ -524,6 +527,7 @@ class TestProxyBaseLLMRequestProcessing: monkeypatch: pytest.MonkeyPatch, gate_allows: bool, expected_cache_controls: dict[str, object], + expected_fallback_cache_controls: dict[str, object], ) -> None: from litellm.caching.caching import Cache from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig @@ -552,7 +556,12 @@ class TestProxyBaseLLMRequestProcessing: ) returned_data, _ = await ProxyBaseLLMRequestProcessing( - data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "cache": {"ttl": 60}, + "fallbacks": [{"model": "gpt-4o", "cache": {}}], + } ).common_processing_pre_call_logic( request=request, general_settings={}, @@ -563,6 +572,7 @@ class TestProxyBaseLLMRequestProcessing: ) assert returned_data["cache"] == expected_cache_controls + assert returned_data["fallbacks"] == [{"model": "gpt-4o", "cache": expected_fallback_cache_controls}] @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True]) From 543c9917aa286cda7d64f3051b8187ac9befca8a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 00:11:00 -0700 Subject: [PATCH 8/8] fix(proxy): apply the semantic cache veto to grouped and nested fallback targets and forward only skip controls through generateContent --- litellm/caching/caching.py | 10 +++++ litellm/google_genai/adapters/handler.py | 6 ++- litellm/integrations/custom_logger.py | 12 +++--- .../proxy/agent_endpoints/a2a_endpoints.py | 10 ++--- .../proxy/common_utils/semantic_cache_gate.py | 17 +++++--- .../google_genai/test_google_genai_adapter.py | 19 ++++++--- .../common_utils/test_semantic_cache_gate.py | 39 +++++++++++++++++++ 7 files changed, 87 insertions(+), 26 deletions(-) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 17cf2c98da0..2f196f721d7 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -16,6 +16,7 @@ import traceback from collections.abc import Generator, Mapping from contextlib import contextmanager from enum import Enum +from types import MappingProxyType from typing import Any, Final, Literal, cast from pydantic import BaseModel, TypeAdapter @@ -43,6 +44,15 @@ from .s3_cache import S3Cache RESPONSE_CACHE_TARGET: Final = "llm_response" _ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"}) +SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True}) + + +def requested_skip_controls(controls: object) -> Mapping[str, object]: + if not isinstance(controls, Mapping): + return {} + return {key: value for key, value in SKIP_CACHE_CONTROLS.items() if controls.get(key) is value} + + _DECODED_LITERAL: Final = TypeAdapter(object) diff --git a/litellm/google_genai/adapters/handler.py b/litellm/google_genai/adapters/handler.py index f0c8fbab106..78c8f8b3765 100644 --- a/litellm/google_genai/adapters/handler.py +++ b/litellm/google_genai/adapters/handler.py @@ -4,6 +4,7 @@ from typing import Final, cast from openai.types.chat import ChatCompletionStreamOptionsParam import litellm +from litellm.caching.caching import requested_skip_controls from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.google_genai.adapters import GenerateContentCompletionKwargs from litellm.types.router import GenericLiteLLMParams @@ -69,8 +70,9 @@ class GenerateContentToCompletionHandler: completion_kwargs["proxy_server_request"] = extra_kwargs["proxy_server_request"] if "litellm_metadata" in extra_kwargs: completion_kwargs["litellm_metadata"] = extra_kwargs["litellm_metadata"] - if "cache" in extra_kwargs and litellm.cache is not None and litellm.cache.is_semantic_cache(): - completion_kwargs["cache"] = extra_kwargs["cache"] + skip_controls: Final = requested_skip_controls(extra_kwargs.get("cache")) + if skip_controls and litellm.cache is not None and litellm.cache.is_semantic_cache(): + completion_kwargs["cache"] = skip_controls if stream: completion_kwargs["stream"] = stream diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 4f3801d6e9b..14bb39879ea 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -449,11 +449,13 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac data: Mapping[str, object], call_type: str, ) -> bool: - """Runs once per request the proxy's semantic cache would serve, after the pre-call hooks, - with the authenticated key and the request body. Return True to let the request read from - and write to the semantic cache. Any other return value, or an exception, skips both for that - request. Every callback that overrides this hook must return True for the request to be cached. - Exact-match cache types never call it. + """Runs once per request on a route the proxy's semantic cache serves, after the pre-call hooks, + even when the cache would skip that request anyway, and once per Responses WebSocket when it opens. + It gets the authenticated key, the request body, and the route's internal call type, such as + `acompletion`, `aresponses`, `_aresponses_websocket`, or `asend_message`. Return True to let the + request read from and write to the semantic cache. Any other return value, or an exception, skips + both for that request. Every callback that overrides this hook must return True for the request to + be cached. Exact-match cache types never call it, and nothing times it out. """ return True diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index c3766c6b9f0..eca1463b66a 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -23,7 +23,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import Cache +from litellm.caching.caching import Cache, requested_skip_controls from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header @@ -43,7 +43,6 @@ from litellm.proxy.agent_endpoints.databricks_oauth import ( ) from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_utils.semantic_cache_gate import SKIP_CACHE_CONTROLS from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING, coerce_keepalive_interval, @@ -187,11 +186,8 @@ async def _resolve_backend_auth( def _request_cache_controls( agent_params: Mapping[str, object], data: Mapping[str, object], cache: Cache | None ) -> Mapping[str, object]: - request_cache: Final = data.get("cache") - if cache is None or not cache.is_semantic_cache() or not isinstance(request_cache, Mapping): - return {} - bypass: Final = {key: value for key, value in SKIP_CACHE_CONTROLS.items() if request_cache.get(key) is value} - if not bypass: + bypass: Final = requested_skip_controls(data.get("cache")) + if cache is None or not cache.is_semantic_cache() or not bypass: return {} agent_cache: Final = agent_params.get("cache") return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **bypass}} diff --git a/litellm/proxy/common_utils/semantic_cache_gate.py b/litellm/proxy/common_utils/semantic_cache_gate.py index b1c170ce45f..afd10fbe0a1 100644 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -1,16 +1,14 @@ from collections.abc import Mapping, Sequence -from types import MappingProxyType from typing import Final from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger -from litellm.caching.caching import Cache +from litellm.caching.caching import SKIP_CACHE_CONTROLS, Cache from litellm.constants import REQUEST_FALLBACK_KEYS from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth -SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True}) SEMANTIC_CACHE_ROUTES: Final = frozenset( { "acompletion", @@ -60,13 +58,20 @@ def _skipping(controls: object) -> Mapping[str, object]: return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} +def _skipping_chain(value: object) -> object: + if not isinstance(value, list): + return value + return [_skipping_fallback_entry(target) for target in _FALLBACK_ENTRIES.validate_python(value)] + + def _skipping_fallback_entry(entry: object) -> object: if not isinstance(entry, Mapping): return entry fields: Final = _FIELDS.validate_python(entry) - if "cache" not in fields: - return fields - return {**fields, "cache": _skipping(fields["cache"])} + return { + **{key: _skipping_chain(value) for key, value in fields.items()}, + **({"cache": _skipping(fields["cache"])} if "cache" in fields else {}), + } def _skipping_fallbacks(data: Mapping[str, object]) -> Mapping[str, Sequence[object]]: diff --git a/tests/unit/google_genai/test_google_genai_adapter.py b/tests/unit/google_genai/test_google_genai_adapter.py index 9fc4e84bdcb..f918203e432 100644 --- a/tests/unit/google_genai/test_google_genai_adapter.py +++ b/tests/unit/google_genai/test_google_genai_adapter.py @@ -1859,12 +1859,19 @@ async def test_generate_content_bridge_leaves_an_exact_match_cache_as_it_was( @pytest.mark.parametrize( - ("semantic", "expected_cache"), - [(True, {"no-cache": True, "no-store": True}), (False, None)], - ids=["semantic", "exact-match"], + ("semantic", "request_cache", "expected_cache"), + [ + (True, {"no-cache": True, "no-store": True}, {"no-cache": True, "no-store": True}), + (True, {"use-cache": True, "ttl": 5, "namespace": "team-a", "no-store": True}, {"no-store": True}), + (True, {"use-cache": True}, None), + (True, "true", None), + (True, ["no-cache"], None), + (False, {"no-cache": True, "no-store": True}, None), + ], + ids=["semantic", "only-skips-forwarded", "opt-in-dropped", "string", "list", "exact-match"], ) -def test_generate_content_bridge_forwards_the_request_cache_controls_only_to_a_semantic_cache( - monkeypatch: pytest.MonkeyPatch, semantic: bool, expected_cache: Mapping[str, bool] | None +def test_generate_content_bridge_forwards_only_skip_controls_and_only_to_a_semantic_cache( + monkeypatch: pytest.MonkeyPatch, semantic: bool, request_cache: object, expected_cache: Mapping[str, bool] | None ) -> None: from litellm.caching.caching import Cache from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler @@ -1878,7 +1885,7 @@ def test_generate_content_bridge_forwards_the_request_cache_controls_only_to_a_s completion_kwargs: Final = GenerateContentToCompletionHandler._prepare_completion_kwargs( model="openai/gpt-4.1", contents=[{"role": "user", "parts": [{"text": "What is the capital of France?"}]}], - extra_kwargs={"cache": {"no-cache": True, "no-store": True}}, + extra_kwargs={"cache": request_cache}, ) assert completion_kwargs.get("cache") == expected_cache diff --git a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py index 8fe90c9d263..fb5660b71e1 100644 --- a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py +++ b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py @@ -235,6 +235,45 @@ async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_its_own_ } +@pytest.mark.asyncio +async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_grouped_or_nested_fallback_targets() -> None: + skip: Final = {"no-cache": True, "no-store": True} + data: Final = { + **_request(), + "fallbacks": [ + { + "gpt-4o-mini": [ + "gpt-4.1-mini", + {"model": "gpt-4o", "cache": {}}, + {"model": "gpt-4.1", "fallbacks": [{"model": "gpt-5", "cache": {"ttl": 5}}]}, + ] + } + ], + "router_settings_override": {"content_policy_fallbacks": [{"gpt-4o-mini": [{"model": "gpt-4o", "cache": {}}]}]}, + } + + overrides: Final = await gated_cache_overrides( + data=data, + user_api_key_dict=_key(OTHER_TEAM), + call_type="acompletion", + gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),), + cache=_semantic_cache(), + ) + + assert overrides["fallbacks"] == [ + { + "gpt-4o-mini": [ + "gpt-4.1-mini", + {"model": "gpt-4o", "cache": skip}, + {"model": "gpt-4.1", "fallbacks": [{"model": "gpt-5", "cache": {"ttl": 5, **skip}}]}, + ] + } + ] + assert overrides["router_settings_override"] == { + "content_policy_fallbacks": [{"gpt-4o-mini": [{"model": "gpt-4o", "cache": skip}]}] + } + + @pytest.mark.asyncio async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_router_settings_fallbacks() -> None: data: Final = {