feat(proxy): add a CustomLogger hook that decides per request whether the semantic cache applies

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:

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,

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}

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.

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}