fix(proxy): apply the semantic cache veto to grouped and nested fallback targets and forward only skip controls through generateContent

This commit is contained in:
mateo-berri 2026-10-29 00:11:00 -07:00
parent dc951d73ce
commit 543c9917aa
7 changed files with 87 additions and 26 deletions

View file

@ -16,6 +16,7 @@ import traceback
from collections.abc import Generator, Mapping
from contextlib import contextmanager
from enum import Enum
from types import MappingProxyType
from typing import Any, Final, Literal, cast
from pydantic import BaseModel, TypeAdapter
@ -43,6 +44,15 @@ from .s3_cache import S3Cache
RESPONSE_CACHE_TARGET: Final = "llm_response"
_ALWAYS_KEYED_PROVIDER_PARAMS: Final = frozenset({"continuation_token"})
SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True})
def requested_skip_controls(controls: object) -> Mapping[str, object]:
if not isinstance(controls, Mapping):
return {}
return {key: value for key, value in SKIP_CACHE_CONTROLS.items() if controls.get(key) is value}
_DECODED_LITERAL: Final = TypeAdapter(object)

View file

@ -4,6 +4,7 @@ from typing import Final, cast
from openai.types.chat import ChatCompletionStreamOptionsParam
import litellm
from litellm.caching.caching import requested_skip_controls
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.google_genai.adapters import GenerateContentCompletionKwargs
from litellm.types.router import GenericLiteLLMParams
@ -69,8 +70,9 @@ 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 and litellm.cache is not None and litellm.cache.is_semantic_cache():
completion_kwargs["cache"] = extra_kwargs["cache"]
skip_controls: Final = requested_skip_controls(extra_kwargs.get("cache"))
if skip_controls and litellm.cache is not None and litellm.cache.is_semantic_cache():
completion_kwargs["cache"] = skip_controls
if stream:
completion_kwargs["stream"] = stream

View file

@ -449,11 +449,13 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
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.
"""Runs once per request on a route the proxy's semantic cache serves, after the pre-call hooks,
even when the cache would skip that request anyway, and once per Responses WebSocket when it opens.
It gets the authenticated key, the request body, and the route's internal call type, such as
`acompletion`, `aresponses`, `_aresponses_websocket`, or `asend_message`. 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, and nothing times it out.
"""
return True

View file

@ -23,7 +23,7 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import Cache
from litellm.caching.caching import Cache, requested_skip_controls
from litellm.litellm_core_utils.core_helpers import as_str_mapping
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header
@ -43,7 +43,6 @@ from litellm.proxy.agent_endpoints.databricks_oauth import (
)
from litellm.proxy.agent_endpoints.utils import merge_agent_headers
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.semantic_cache_gate import SKIP_CACHE_CONTROLS
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING,
coerce_keepalive_interval,
@ -187,11 +186,8 @@ async def _resolve_backend_auth(
def _request_cache_controls(
agent_params: Mapping[str, object], data: Mapping[str, object], cache: Cache | None
) -> Mapping[str, object]:
request_cache: Final = data.get("cache")
if cache is None or not cache.is_semantic_cache() or not isinstance(request_cache, Mapping):
return {}
bypass: Final = {key: value for key, value in SKIP_CACHE_CONTROLS.items() if request_cache.get(key) is value}
if not bypass:
bypass: Final = requested_skip_controls(data.get("cache"))
if cache is None or not cache.is_semantic_cache() or not bypass:
return {}
agent_cache: Final = agent_params.get("cache")
return {"cache": {**(agent_cache if isinstance(agent_cache, Mapping) else {}), **bypass}}

View file

@ -1,16 +1,14 @@
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.caching.caching import SKIP_CACHE_CONTROLS, Cache
from litellm.constants import REQUEST_FALLBACK_KEYS
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy._types import UserAPIKeyAuth
SKIP_CACHE_CONTROLS: Final = MappingProxyType({"no-cache": True, "no-store": True})
SEMANTIC_CACHE_ROUTES: Final = frozenset(
{
"acompletion",
@ -60,13 +58,20 @@ def _skipping(controls: object) -> Mapping[str, object]:
return {**(controls if isinstance(controls, Mapping) else {}), **SKIP_CACHE_CONTROLS}
def _skipping_chain(value: object) -> object:
if not isinstance(value, list):
return value
return [_skipping_fallback_entry(target) for target in _FALLBACK_ENTRIES.validate_python(value)]
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"])}
return {
**{key: _skipping_chain(value) for key, value in fields.items()},
**({"cache": _skipping(fields["cache"])} if "cache" in fields else {}),
}
def _skipping_fallbacks(data: Mapping[str, object]) -> Mapping[str, Sequence[object]]:

View file

@ -1859,12 +1859,19 @@ async def test_generate_content_bridge_leaves_an_exact_match_cache_as_it_was(
@pytest.mark.parametrize(
("semantic", "expected_cache"),
[(True, {"no-cache": True, "no-store": True}), (False, None)],
ids=["semantic", "exact-match"],
("semantic", "request_cache", "expected_cache"),
[
(True, {"no-cache": True, "no-store": True}, {"no-cache": True, "no-store": True}),
(True, {"use-cache": True, "ttl": 5, "namespace": "team-a", "no-store": True}, {"no-store": True}),
(True, {"use-cache": True}, None),
(True, "true", None),
(True, ["no-cache"], None),
(False, {"no-cache": True, "no-store": True}, None),
],
ids=["semantic", "only-skips-forwarded", "opt-in-dropped", "string", "list", "exact-match"],
)
def test_generate_content_bridge_forwards_the_request_cache_controls_only_to_a_semantic_cache(
monkeypatch: pytest.MonkeyPatch, semantic: bool, expected_cache: Mapping[str, bool] | None
def test_generate_content_bridge_forwards_only_skip_controls_and_only_to_a_semantic_cache(
monkeypatch: pytest.MonkeyPatch, semantic: bool, request_cache: object, expected_cache: Mapping[str, bool] | None
) -> None:
from litellm.caching.caching import Cache
from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler
@ -1878,7 +1885,7 @@ def test_generate_content_bridge_forwards_the_request_cache_controls_only_to_a_s
completion_kwargs: Final = GenerateContentToCompletionHandler._prepare_completion_kwargs(
model="openai/gpt-4.1",
contents=[{"role": "user", "parts": [{"text": "What is the capital of France?"}]}],
extra_kwargs={"cache": {"no-cache": True, "no-store": True}},
extra_kwargs={"cache": request_cache},
)
assert completion_kwargs.get("cache") == expected_cache

View file

@ -235,6 +235,45 @@ async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_its_own_
}
@pytest.mark.asyncio
async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_grouped_or_nested_fallback_targets() -> None:
skip: Final = {"no-cache": True, "no-store": True}
data: Final = {
**_request(),
"fallbacks": [
{
"gpt-4o-mini": [
"gpt-4.1-mini",
{"model": "gpt-4o", "cache": {}},
{"model": "gpt-4.1", "fallbacks": [{"model": "gpt-5", "cache": {"ttl": 5}}]},
]
}
],
"router_settings_override": {"content_policy_fallbacks": [{"gpt-4o-mini": [{"model": "gpt-4o", "cache": {}}]}]},
}
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(),
)
assert overrides["fallbacks"] == [
{
"gpt-4o-mini": [
"gpt-4.1-mini",
{"model": "gpt-4o", "cache": skip},
{"model": "gpt-4.1", "fallbacks": [{"model": "gpt-5", "cache": {"ttl": 5, **skip}}]},
]
}
]
assert overrides["router_settings_override"] == {
"content_policy_fallbacks": [{"gpt-4o-mini": [{"model": "gpt-4o", "cache": skip}]}]
}
@pytest.mark.asyncio
async def test_a_rejected_team_cannot_reopen_the_semantic_cache_through_router_settings_fallbacks() -> None:
data: Final = {