fix(proxy): run the semantic cache check on every route that can reach the cache, including A2A, WebSocket, interactions and batch rows

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

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

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:

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

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):

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",

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