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