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] 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,)