mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
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:
parent
dc951d73ce
commit
543c9917aa
7 changed files with 87 additions and 26 deletions
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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}}
|
||||
|
|
|
|||
|
|
@ -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]]:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue