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] 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}