From dbe2fc909499c655ac899ecc7c26b686b7920078 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:17:47 -0700 Subject: [PATCH] fix(proxy): keep a rejected caller off the semantic cache on fallback hops and late batch rows A fallback entry's keys are copied over the request on that hop, so the check's no now also adds no-cache and no-store to the cache inside each fallback entry the request carries, in its own lists and in router_settings_override. Batch rows read the cache and the callbacks as each row runs, and the MCP sampling test runs through a real router. --- litellm/constants.py | 1 + litellm/proxy/batches_endpoints/endpoints.py | 4 +- .../litellm_executed_batches.py | 14 +- litellm/proxy/common_request_processing.py | 21 ++- .../proxy/common_utils/semantic_cache_gate.py | 49 ++++++- litellm/proxy/litellm_pre_call_utils.py | 6 +- litellm/proxy/mcp_server/sampling_handler.py | 20 +-- .../test_litellm_executed_batches.py | 28 +++- .../common_utils/test_semantic_cache_gate.py | 126 +++++++++++++----- .../test_mcp_sampling_completion_flow.py | 38 ++++-- .../proxy/test_common_request_processing.py | 16 ++- 11 files changed, 228 insertions(+), 95 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 2a20e812e68..496f82b3264 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -54,6 +54,7 @@ AZURE_DEFAULT_RESPONSES_API_VERSION: Final = str(os.getenv("AZURE_DEFAULT_RESPON AZURE_OPENAI_AUDIO_PROVIDERS: Final = frozenset({"azure", "azure_ai"}) ROUTER_MAX_FALLBACKS: Final = get_env_int("ROUTER_MAX_FALLBACKS", 5) ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS: Final = 2000 +REQUEST_FALLBACK_KEYS: Final = ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks") RUNTIME_UPDATABLE_ROUTER_SETTINGS: Final[frozenset[str]] = frozenset( { "routing_strategy_args", diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 5fbb240ea01..5ae1107e3b7 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -110,8 +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, + semantic_cache_gates=lambda: ProxyLogging.callback_capabilities().semantic_cache_gates, + semantic_cache=lambda: litellm.cache, ) diff --git a/litellm/proxy/batches_endpoints/litellm_executed_batches.py b/litellm/proxy/batches_endpoints/litellm_executed_batches.py index 1299c6dda6a..ffe803c24d8 100644 --- a/litellm/proxy/batches_endpoints/litellm_executed_batches.py +++ b/litellm/proxy/batches_endpoints/litellm_executed_batches.py @@ -29,7 +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.common_utils.semantic_cache_gate import gated_cache_overrides 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 @@ -426,8 +426,8 @@ class LiteLLMExecutedBatchRunner: batches: ManagedBatchRepository, proxy_logging_obj: ProxyLogging, general_settings: Mapping[str, object], - semantic_cache_gates: Sequence[CustomLogger], - semantic_cache: Cache | None, + semantic_cache_gates: Callable[[], Sequence[CustomLogger]], + semantic_cache: Callable[[], Cache | None], concurrency: int = LITELLM_EXECUTED_BATCH_CONCURRENCY, heartbeat_seconds: float = _HEARTBEAT_SECONDS, completion_window_seconds: float = _COMPLETION_WINDOW_SECONDS, @@ -654,14 +654,14 @@ class LiteLLMExecutedBatchRunner: request: Final = MappingProxyType( {**line.body, "model": run.model, "metadata": self._row_metadata(run), "disable_fallbacks": True} ) - cache_controls: Final = await gated_cache_controls( + cache_overrides: Final = await gated_cache_overrides( 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, + gates=self.semantic_cache_gates(), + cache=self.semantic_cache(), ) - params: Final = MappingProxyType({**request, **({} if cache_controls is None else {"cache": cache_controls})}) + params: Final = MappingProxyType({**request, **cache_overrides}) return _dump(await self._router_call(run.endpoint)(**params)) def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8b5c9a64a01..230d74bc731 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -114,7 +114,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.semantic_cache_gate import gated_cache_overrides from litellm.proxy.common_utils.sse_keepalive import ( SSE_COMMENT_PING_BYTES, SSE_STREAM_START_TAIL, @@ -2511,16 +2511,15 @@ 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 + cache_overrides: Final = await gated_cache_overrides( + 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, + ) + for key, value in cache_overrides.items(): + self.data[key] = value 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 index 8b45ce58031..b1c170ce45f 100644 --- a/litellm/proxy/common_utils/semantic_cache_gate.py +++ b/litellm/proxy/common_utils/semantic_cache_gate.py @@ -2,8 +2,11 @@ from collections.abc import Mapping, Sequence from types import MappingProxyType from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.caching.caching import Cache +from litellm.constants import REQUEST_FALLBACK_KEYS from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -22,6 +25,8 @@ SEMANTIC_CACHE_ROUTES: Final = frozenset( "asend_message", } ) +_FALLBACK_ENTRIES: Final = TypeAdapter(list[object]) +_FIELDS: Final = TypeAdapter(Mapping[str, object]) async def _gate_allows( @@ -51,16 +56,48 @@ async def _every_gate_allows( return True -async def gated_cache_controls( +def _skipping(controls: object) -> Mapping[str, object]: + return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS} + + +def _skipping_fallback_entry(entry: object) -> object: + if not isinstance(entry, Mapping): + return entry + fields: Final = _FIELDS.validate_python(entry) + if "cache" not in fields: + return fields + return {**fields, "cache": _skipping(fields["cache"])} + + +def _skipping_fallbacks(data: Mapping[str, object]) -> Mapping[str, Sequence[object]]: + return { + key: [_skipping_fallback_entry(entry) for entry in _FALLBACK_ENTRIES.validate_python(entries)] + for key in REQUEST_FALLBACK_KEYS + if isinstance(entries := data.get(key), list) + } + + +def _skipping_router_settings_override(data: Mapping[str, object]) -> Mapping[str, Mapping[str, object]]: + override: Final = data.get("router_settings_override") + if not isinstance(override, Mapping): + return {} + settings: Final = _FIELDS.validate_python(override) + return {"router_settings_override": {**settings, **_skipping_fallbacks(settings)}} + + +async def gated_cache_overrides( data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str, gates: Sequence[CustomLogger], cache: Cache | None, -) -> Mapping[str, object] | None: +) -> Mapping[str, object]: if not gates or cache is None or not cache.is_semantic_cache() or call_type not in SEMANTIC_CACHE_ROUTES: - return None + return {} 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} + return {} + return { + "cache": _skipping(data.get("cache")), + **_skipping_fallbacks(data), + **_skipping_router_settings_override(data), + } diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 24a31723faa..b22049690a1 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -39,6 +39,7 @@ from litellm.constants import ( MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, OTEL_SERVICE_NAME_METADATA_KEYS, PRE_CALL_EXECUTED_GUARDRAILS_KEY, + REQUEST_FALLBACK_KEYS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, ROUTING_REQUEST_TAGS_METADATA_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, @@ -3076,7 +3077,6 @@ async def add_litellm_data_to_request( return data -_REQUEST_FALLBACK_KEYS: Final = ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks") _FALLBACK_DISCOVERY_LIMIT: Final = 256 @@ -3150,7 +3150,7 @@ def _fallback_lists( ): if isinstance(router_list, list): lists.append(cast("Sequence[object]", router_list)) # cast-ok: fallback lists hold mixed entry shapes - for key in _REQUEST_FALLBACK_KEYS: + for key in REQUEST_FALLBACK_KEYS: entries: object = data.get(key) if isinstance(entries, list): lists.append(cast("Sequence[object]", entries)) # cast-ok: request fallback lists hold mixed entry shapes @@ -3403,7 +3403,7 @@ async def _resolve_user_provider_credentials_for_request( _fallback_entry_target_count(entry) for entry in cast("Sequence[object]", entries) # cast-ok: request fallback lists hold mixed entry shapes ) - for key in _REQUEST_FALLBACK_KEYS + for key in REQUEST_FALLBACK_KEYS if isinstance(entries := data.get(key), list) ) if request_fallback_targets > _FALLBACK_DISCOVERY_LIMIT: diff --git a/litellm/proxy/mcp_server/sampling_handler.py b/litellm/proxy/mcp_server/sampling_handler.py index 5af94a93871..efbfa152fa6 100644 --- a/litellm/proxy/mcp_server/sampling_handler.py +++ b/litellm/proxy/mcp_server/sampling_handler.py @@ -37,7 +37,7 @@ from fastapi import HTTPException from pydantic import TypeAdapter from litellm._logging import verbose_logger -from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_controls +from litellm.proxy.common_utils.semantic_cache_gate import gated_cache_overrides from litellm.proxy.utils import ProxyLogging # Guard imports that require the mcp package @@ -1127,14 +1127,16 @@ async def _run_guardrails_and_call_llm( import litellm - cache_controls: Final = await gated_cache_controls( - data=completion_kwargs, - user_api_key_dict=user_api_key_auth, - call_type="acompletion", - gates=ProxyLogging.callback_capabilities().semantic_cache_gates, - cache=litellm.cache, - ) - call_kwargs: Final = completion_kwargs if cache_controls is None else {**completion_kwargs, "cache": cache_controls} + call_kwargs: Final = { + **completion_kwargs, + **await gated_cache_overrides( + data=completion_kwargs, + user_api_key_dict=user_api_key_auth, + call_type="acompletion", + gates=ProxyLogging.callback_capabilities().semantic_cache_gates, + cache=litellm.cache, + ), + } try: from litellm.proxy.proxy_server import llm_router 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 1e62563f7e3..51a9e7c5eac 100644 --- a/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py +++ b/tests/unit/proxy/batches_endpoints/test_litellm_executed_batches.py @@ -383,8 +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, + semantic_cache_gates: Callable[[], Sequence[CustomLogger]] = lambda: (), + semantic_cache: Callable[[], Cache | None] = lambda: None, ) -> Harness: store = store_factory({INPUT_FILE_ID: managed_input_file()} if files is None else files) router = FakeRouter() @@ -722,6 +722,10 @@ class _TeamAllowlist(CustomLogger): return user_api_key_dict.team_id == self.team_id +def _semantic_cache() -> Cache: + return Cache(type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8) + + @pytest.mark.parametrize( ("allowed_team", "row_cache_controls"), [("team-1", None), ("team-other", {"no-cache": True, "no-store": True})], @@ -731,10 +735,8 @@ 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) + semantic_cache: Final = _semantic_cache() + harness: Final = make_runner(semantic_cache_gates=lambda: (gate,), semantic_cache=lambda: semantic_cache) await harness.create_and_finish() calls = harness.router.acompletion.await_args_list @@ -745,6 +747,20 @@ async def test_create_runs_the_semantic_cache_check_on_each_row( ] +async def test_rows_follow_the_semantic_cache_and_check_in_place_when_they_run_not_when_the_batch_was_created() -> None: + live_gates: Final[list[Sequence[CustomLogger]]] = [()] # mutable-ok: stands in for callbacks an admin adds later + live_cache: Final[list[Cache | None]] = [None] # mutable-ok: stands in for a cache an admin turns on later + harness: Final = make_runner(semantic_cache_gates=lambda: live_gates[-1], semantic_cache=lambda: live_cache[-1]) + await harness.create() + + live_gates.append((_TeamAllowlist("team-other"),)) + live_cache.append(_semantic_cache()) + await asyncio.gather(*list(litellm_executed_batches._RUNNING_BATCHES)) + + calls: Final = harness.router.acompletion.await_args_list + assert [call.kwargs.get("cache") for call in calls] == [{"no-cache": True, "no-store": True}] * 2 + + 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 eef95c5e589..8fe90c9d263 100644 --- a/tests/unit/proxy/common_utils/test_semantic_cache_gate.py +++ b/tests/unit/proxy/common_utils/test_semantic_cache_gate.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from typing import Final import pytest @@ -6,7 +7,8 @@ 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.common_utils.semantic_cache_gate import gated_cache_overrides +from litellm.proxy.route_llm_request import _merge_router_settings_override from litellm.proxy.utils import ProxyLogging from litellm.types.caching import ( DEFAULT_CACHING_SUPPORTED_CALL_TYPES, @@ -57,7 +59,7 @@ def _semantic_cache( ) -def _request(**cache_controls: object) -> dict[str, object]: +def _request(**cache_controls: object) -> Mapping[str, object]: return { "model": "gpt-4o-mini", "messages": [{"role": "user", "content": "What color is the sky?"}], @@ -70,23 +72,23 @@ def _key(team_id: str) -> UserAPIKeyAuth: @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}) +async def test_a_team_the_gate_rejects_skips_semantic_lookup_and_store_even_when_the_client_opts_in() -> None: + gate: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + data: Final = _request(**{"use-cache": True, "ttl": 60}) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( 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 result == {"cache": {"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) +async def test_a_team_the_gate_allows_keeps_its_own_cache_controls() -> None: + data: Final = _request(ttl=60) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=data, user_api_key_dict=_key(ALLOWED_TEAM), call_type="acompletion", @@ -94,33 +96,33 @@ async def test_a_team_the_gate_allows_keeps_its_own_cache_controls(): cache=_semantic_cache(), ) - assert result is None + assert result == {} @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() +async def test_the_gate_is_never_consulted_without_a_semantic_cache(cache: Cache | None) -> None: + gate: Final = _TeamAllowlist(frozenset()) + data: Final = _request() - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=cache ) - assert result is None + assert result == {} assert gate.seen == [] @pytest.mark.asyncio @pytest.mark.parametrize("route", ["acreate_batch", "avector_store_search", "aspeech", "atranscription", "arerank"]) -async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str): - gate = _TeamAllowlist(frozenset()) +async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cache(route: str) -> None: + gate: Final = _TeamAllowlist(frozenset()) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type=route, gates=(gate,), cache=_semantic_cache() ) - assert result is None + assert result == {} assert gate.seen == [] @@ -140,10 +142,10 @@ async def test_the_gate_is_never_consulted_for_a_route_that_cannot_reach_the_cac ) 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()) +) -> None: + gate: Final = _TeamAllowlist(frozenset()) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type=route, @@ -151,14 +153,14 @@ async def test_the_gate_covers_routes_that_reach_the_cache_through_a_bridge_what cache=_semantic_cache(supported_call_types=(served_call_type,)), ) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"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( +async def test_a_gate_that_fails_or_answers_anything_but_true_skips_the_semantic_cache(gate: CustomLogger) -> None: + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(ALLOWED_TEAM), call_type="aresponses", @@ -166,16 +168,16 @@ async def test_a_gate_that_fails_or_answers_anything_but_true_skips_the_semantic cache=_semantic_cache(), ) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"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()) +async def test_one_rejecting_gate_skips_the_semantic_cache_whatever_the_others_answer(deny_first: bool) -> None: + allow: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM})) + deny: Final = _TeamAllowlist(frozenset()) - result = await gated_cache_controls( + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(ALLOWED_TEAM), call_type="anthropic_messages", @@ -183,16 +185,16 @@ async def test_one_rejecting_gate_skips_the_semantic_cache_whatever_the_others_a cache=_semantic_cache(), ) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"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})) +async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: pytest.MonkeyPatch) -> None: + gate: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM})) monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), gate]) - gates = ProxyLogging.callback_capabilities().semantic_cache_gates - result = await gated_cache_controls( + gates: Final = ProxyLogging.callback_capabilities().semantic_cache_gates + result: Final = await gated_cache_overrides( data=_request(), user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", @@ -201,4 +203,54 @@ async def test_only_callbacks_that_override_the_hook_act_as_gates(monkeypatch: p ) assert gates == (gate,) - assert result == {"no-cache": True, "no-store": True} + assert result == {"cache": {"no-cache": True, "no-store": True}} + + +@pytest.mark.asyncio +async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_its_own_fallback_entries() -> None: + data: Final = { + **_request(ttl=60), + "fallbacks": [{"model": "gpt-4o", "cache": {"ttl": 5}}, "gpt-4.1-mini", {"model": "gpt-4.1"}], + "context_window_fallbacks": [{"gpt-4o-mini": ["gpt-4o"], "cache": {}}], + "content_policy_fallbacks": [{"model": "gpt-4o", "cache": None}], + } + + result: Final = await gated_cache_overrides( + data=data, + user_api_key_dict=_key(OTHER_TEAM), + call_type="acompletion", + gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),), + cache=_semantic_cache(), + ) + + assert result == { + "cache": {"ttl": 60, "no-cache": True, "no-store": True}, + "fallbacks": [ + {"model": "gpt-4o", "cache": {"ttl": 5, "no-cache": True, "no-store": True}}, + "gpt-4.1-mini", + {"model": "gpt-4.1"}, + ], + "context_window_fallbacks": [{"gpt-4o-mini": ["gpt-4o"], "cache": {"no-cache": True, "no-store": True}}], + "content_policy_fallbacks": [{"model": "gpt-4o", "cache": {"no-cache": True, "no-store": True}}], + } + + +@pytest.mark.asyncio +async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_router_settings_fallbacks() -> None: + data: Final = { + **_request(), + "router_settings_override": {"fallbacks": [{"model": "gpt-4o", "cache": {}}], "num_retries": 2}, + } + + overrides: Final = await gated_cache_overrides( + data=data, + user_api_key_dict=_key(OTHER_TEAM), + call_type="acompletion", + gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),), + cache=_semantic_cache(), + ) + routed: Final = {**data, **overrides} + _ = _merge_router_settings_override(routed) + + assert routed["fallbacks"] == [{"model": "gpt-4o", "cache": {"no-cache": True, "no-store": True}}] + assert routed["num_retries"] == 2 diff --git a/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py b/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py index 549fb99213e..908823f255f 100644 --- a/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py +++ b/tests/unit/proxy/mcp_server/test_mcp_sampling_completion_flow.py @@ -9,6 +9,7 @@ through the proxy router / guardrails, and the end-to-end from collections.abc import Mapping from types import SimpleNamespace +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -18,7 +19,9 @@ import litellm from litellm.caching.caching import Cache from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth +from litellm.router import Router from litellm.types.caching import LiteLLMCacheType +from litellm.types.utils import CallTypes from litellm.proxy.mcp_server.sampling_handler import ( _build_completion_kwargs, @@ -125,34 +128,47 @@ class TestRunGuardrailsAndCallLlm: assert result == "router-response" router.acompletion.assert_awaited_once() - async def test_should_skip_the_semantic_cache_for_a_key_the_semantic_cache_check_rejects(self): + async def test_should_skip_the_semantic_cache_for_a_key_the_semantic_cache_check_rejects(self) -> None: class _RejectEveryTeam(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.deployment_cache_controls: list[object] = [] + async def async_should_semantic_cache( self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str ) -> bool: return False - async def _cache_controls_the_router_received(**kwargs): - return kwargs.get("cache") + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.deployment_cache_controls.append(kwargs.get("cache")) + return None - semantic_cache = Cache( + gate: Final = _RejectEveryTeam() + router: Final = Router( + model_list=[ + { + "model_name": "gpt-4o", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-test", "mock_response": "hi"}, + } + ] + ) + semantic_cache: Final = Cache( type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8 ) with ( patch("litellm.proxy.proxy_server.proxy_logging_obj", None), - patch( - "litellm.proxy.proxy_server.llm_router", - SimpleNamespace(acompletion=_cache_controls_the_router_received), - ), - patch.object(litellm, "callbacks", [_RejectEveryTeam()]), + patch("litellm.proxy.proxy_server.llm_router", router), + patch.object(litellm, "callbacks", [gate]), patch.object(litellm, "cache", semantic_cache), ): - received = await _run_guardrails_and_call_llm( + await _run_guardrails_and_call_llm( completion_kwargs={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, user_api_key_auth=UserAPIKeyAuth(api_key="sk-test", team_id="team-other"), ) - assert received == {"no-cache": True, "no-store": True} + assert gate.deployment_cache_controls == [{"no-cache": True, "no-store": True}] async def test_should_propagate_guardrail_rejection(self): plo = MagicMock() diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 5263b0592bd..7e8db33095f 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -515,8 +515,11 @@ class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio @pytest.mark.parametrize( - ("gate_allows", "expected_cache_controls"), - [(True, {"ttl": 60}), (False, {"ttl": 60, "no-cache": True, "no-store": True})], + ("gate_allows", "expected_cache_controls", "expected_fallback_cache_controls"), + [ + (True, {"ttl": 60}, {}), + (False, {"ttl": 60, "no-cache": True, "no-store": True}, {"no-cache": True, "no-store": True}), + ], ids=["allowed", "rejected"], ) async def test_common_processing_pre_call_logic_applies_the_semantic_cache_gate( @@ -524,6 +527,7 @@ class TestProxyBaseLLMRequestProcessing: monkeypatch: pytest.MonkeyPatch, gate_allows: bool, expected_cache_controls: dict[str, object], + expected_fallback_cache_controls: dict[str, object], ) -> None: from litellm.caching.caching import Cache from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig @@ -552,7 +556,12 @@ class TestProxyBaseLLMRequestProcessing: ) returned_data, _ = await ProxyBaseLLMRequestProcessing( - data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}], "cache": {"ttl": 60}} + data={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "cache": {"ttl": 60}, + "fallbacks": [{"model": "gpt-4o", "cache": {}}], + } ).common_processing_pre_call_logic( request=request, general_settings={}, @@ -563,6 +572,7 @@ class TestProxyBaseLLMRequestProcessing: ) assert returned_data["cache"] == expected_cache_controls + assert returned_data["fallbacks"] == [{"model": "gpt-4o", "cache": expected_fallback_cache_controls}] @pytest.mark.asyncio @pytest.mark.parametrize("safe_memory_mode", [False, True])