This commit is contained in:
m6atlas 2026-10-29 03:59:58 -07:00 • committed by GitHub
commit 78437fc6c3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 836 additions and 15 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)
@ -58,7 +68,7 @@ def response_cache_phase(operation: Literal["get", "set"]) -> Generator[None]:
yield
def print_verbose(print_statement):
def print_verbose(print_statement: object) -> None:
try:
verbose_logger.debug(print_statement)
if litellm.set_verbose:
@ -353,7 +363,7 @@ class Cache:
)
_SEMANTIC_CACHE_END_USER_SCOPE_FIELD: Final = "user_api_key_end_user_id"
def _is_semantic_cache(self) -> bool:
def is_semantic_cache(self) -> bool:
return self.type in (
LiteLLMCacheType.REDIS_SEMANTIC,
LiteLLMCacheType.QDRANT_SEMANTIC,
@ -395,7 +405,7 @@ class Cache:
return preset_cache_key
combined_kwargs: Final = ModelParamHelper.get_all_llm_api_params()
is_semantic_cache: Final = self._is_semantic_cache()
is_semantic_cache: Final = self.is_semantic_cache()
scope_excluded_params: Final = self._SEMANTIC_CACHE_SCOPE_EXCLUDED_PARAMS if is_semantic_cache else frozenset()
for param in kwargs:
if param in scope_excluded_params:
@ -963,7 +973,7 @@ class Cache:
except Exception as e:
self._log_add_cache_failure(e)
def should_use_cache(self, **kwargs):
def should_use_cache(self, **kwargs: object) -> bool:
"""
Returns true if we should use the cache for LLM API calls

View file

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

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,6 +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"]
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

@ -443,6 +443,22 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
"""
return model_names
async def async_should_semantic_cache(
self,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: str,
) -> bool:
"""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
async def async_post_call_response_headers_hook(
self,
data: dict,

View file

@ -23,6 +23,7 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
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
@ -182,6 +183,16 @@ async def _resolve_backend_auth(
)
def _request_cache_controls(
agent_params: Mapping[str, object], data: Mapping[str, object], cache: Cache | None
) -> Mapping[str, object]:
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}}
def _forwarding_headers(
caller_identity: Mapping[str, str],
request_data: Mapping[str, object],
@ -842,7 +853,10 @@ 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(backend_auth.litellm_params, data, litellm.cache),
}
# Merge agent-level guardrails into data so post_call_success_hook and
# _handle_stream_message both pick them up. A2A agents use model

View file

@ -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=lambda: ProxyLogging.callback_capabilities().semantic_cache_gates,
semantic_cache=lambda: litellm.cache,
)

View file

@ -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_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
@ -423,6 +426,8 @@ class LiteLLMExecutedBatchRunner:
batches: ManagedBatchRepository,
proxy_logging_obj: ProxyLogging,
general_settings: Mapping[str, object],
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,
@ -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_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(),
)
params: Final = MappingProxyType({**request, **cache_overrides})
return _dump(await self._router_call(run.endpoint)(**params))
def _router_call(self, endpoint: BatchEndpoint) -> _RouterCall:

View file

@ -115,6 +115,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_overrides
from litellm.proxy.common_utils.sse_keepalive import (
SSE_COMMENT_PING_BYTES,
SSE_STREAM_START_TAIL,
@ -2562,6 +2563,15 @@ class ProxyBaseLLMRequestProcessing:
call_type=route_type,
skip_guardrails=skip_guardrails,
)
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,

View file

@ -0,0 +1,108 @@
from collections.abc import Mapping, Sequence
from typing import Final
from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
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
SEMANTIC_CACHE_ROUTES: Final = frozenset(
{
"acompletion",
"atext_completion",
"aembedding",
"aresponses",
"_aresponses_websocket",
"anthropic_messages",
"agenerate_content",
"agenerate_content_stream",
"acreate_interaction",
"asend_message",
}
)
_FALLBACK_ENTRIES: Final = TypeAdapter(list[object])
_FIELDS: Final = TypeAdapter(Mapping[str, object])
async def _gate_allows(
gate: CustomLogger,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: str,
) -> bool:
try:
return await gate.async_should_semantic_cache(user_api_key_dict, data, call_type) is True
except Exception: # noqa: BLE001 # operator hooks can raise anything, and every failure skips the cache
verbose_proxy_logger.exception(
"%s.async_should_semantic_cache raised, so this request skips the semantic cache", type(gate).__name__
)
return False
async def _every_gate_allows(
gates: Sequence[CustomLogger],
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: str,
) -> bool:
for gate in gates:
if not await _gate_allows(gate, user_api_key_dict, data, call_type):
return False
return True
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)
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]]:
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]:
if not gates or cache is None or not cache.is_semantic_cache() or call_type not in SEMANTIC_CACHE_ROUTES:
return {}
if await _every_gate_allows(gates, user_api_key_dict, data, call_type):
return {}
return {
"cache": _skipping(data.get("cache")),
**_skipping_fallbacks(data),
**_skipping_router_settings_override(data),
}

View file

@ -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,
@ -3075,7 +3076,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
@ -3149,7 +3149,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
@ -3402,7 +3402,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:

View file

@ -37,6 +37,8 @@ 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_overrides
from litellm.proxy.utils import ProxyLogging
# Guard imports that require the mcp package
try:
@ -1125,14 +1127,24 @@ async def _run_guardrails_and_call_llm(
import litellm
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
if llm_router is not None:
return await _AcompletionCall(fn=llm_router.acompletion).fn(**completion_kwargs)
return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs)
return await _AcompletionCall(fn=llm_router.acompletion).fn(**call_kwargs)
return await _AcompletionCall(fn=litellm.acompletion).fn(**call_kwargs)
except ImportError:
return await _AcompletionCall(fn=litellm.acompletion).fn(**completion_kwargs)
return await _AcompletionCall(fn=litellm.acompletion).fn(**call_kwargs)
async def handle_sampling_create_message(

View file

@ -1654,6 +1654,7 @@ class _CallbackCapabilities:
# every string entry in ``litellm.callbacks``.
resolved_callbacks: tuple[object, ...] = field(default_factory=tuple)
listed_models_filters: tuple[CustomLogger, ...] = field(default_factory=tuple)
semantic_cache_gates: tuple[CustomLogger, ...] = field(default_factory=tuple)
def _overrides_hook(callback: CustomLogger, hook_name: str) -> bool:
@ -3436,6 +3437,9 @@ class ProxyLogging:
listed_models_filters=tuple(
callback for callback in resolved_callbacks if _overrides_hook(callback, "async_filter_listed_models")
),
semantic_cache_gates=tuple(
callback for callback in resolved_callbacks if _overrides_hook(callback, "async_should_semantic_cache")
),
)
# Limit cache to handle test churn without leaking; production
# callback lists are stable so this rarely grows past 1 entry.

View file

@ -22,6 +22,7 @@ class GenerateContentCompletionKwargs(TypedDict, total=False):
stream: bool
metadata: dict[str, object]
litellm_metadata: ReadOnly[Mapping[str, object]]
cache: ReadOnly[Mapping[str, object]]
stream_options: ReadOnly[ChatCompletionStreamOptionsParam]
extra_headers: dict[str, str] | None
proxy_server_request: dict[str, object] | None

View file

@ -2,12 +2,14 @@
"""
Test to verify the Google GenAI generate_content adapter functionality
"""
import asyncio
import json
import unittest
from collections.abc import Mapping
from typing import Final
import pytest
import respx
import litellm
from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter
@ -1820,3 +1822,83 @@ async def test_generate_content_sends_response_schema_and_tool_parameters_to_the
}
assert response["candidates"][0]["content"]["parts"] == [{"text": '{"park_name": "EPCOT"}'}]
assert "text" not in response
@pytest.mark.asyncio
@pytest.mark.parametrize(
("first_cache_controls", "second_cache_controls"),
[({"no-store": True}, None), (None, {"no-cache": True})],
ids=["no-store", "no-cache"],
)
async def test_generate_content_bridge_leaves_an_exact_match_cache_as_it_was(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
first_cache_controls: Mapping[str, bool] | None,
second_cache_controls: Mapping[str, bool] | None,
) -> None:
import httpx
from litellm.caching.caching import Cache
from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
from litellm.types.caching import LiteLLMCacheType
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
monkeypatch.setattr(litellm, "cache", Cache(type=LiteLLMCacheType.LOCAL))
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-capital",
"object": "chat.completion",
"created": 1234567890,
"model": "gpt-4.1",
"choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Paris"}}],
"usage": {"prompt_tokens": 10, "completion_tokens": 1, "total_tokens": 11},
},
)
)
request: Final = {
"model": "openai/gpt-4.1",
"contents": [{"role": "user", "parts": [{"text": "What is the capital of France?"}]}],
"api_key": "sk-test",
}
await agenerate_content(**request, **({} if first_cache_controls is None else {"cache": first_cache_controls}))
await asyncio.gather(*_PENDING_CACHE_WRITES)
await agenerate_content(**request, **({} if second_cache_controls is None else {"cache": second_cache_controls}))
assert route.call_count == 1
@pytest.mark.parametrize(
("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_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
from litellm.types.caching import LiteLLMCacheType
cache_type: Final = LiteLLMCacheType.VALKEY_SEMANTIC if semantic else LiteLLMCacheType.LOCAL
monkeypatch.setattr(
litellm, "cache", Cache(type=cache_type, host="localhost", port="6379", similarity_threshold=0.8)
)
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": request_cache},
)
assert completion_kwargs.get("cache") == expected_cache

View file

@ -19,8 +19,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]]]
@ -44,6 +47,7 @@ def _sse_upstream_response(*lines: str, first_line_delay: float = 0.0) -> httpx.
class CapturedAgentCall:
request_id: object
agent_extra_headers: dict[str, str] | None
litellm_params: Mapping[str, object] | None
@pytest.mark.asyncio
@ -398,7 +402,12 @@ def _make_agent_mock(url: str = "http://backend-agent:10001") -> MagicMock:
return agent
def _make_request_mock(method: str, params: Mapping[str, object], request_id: object = "req-1") -> MagicMock:
def _make_request_mock(
method: str,
params: Mapping[str, object],
request_id: object = "req-1",
cache: Mapping[str, object] | None = None,
) -> MagicMock:
req = MagicMock()
req.headers = {}
req.json = AsyncMock(
@ -407,6 +416,7 @@ def _make_request_mock(method: str, params: Mapping[str, object], request_id: ob
"id": request_id,
"method": method,
"params": params,
**({} if cache is None else {"cache": cache}),
}
)
return req
@ -627,7 +637,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
@ -641,6 +655,67 @@ 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):
agent = _make_agent_mock()
agent.litellm_params = {"cache": {"ttl": 60}}
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,
agent=agent,
)
assert (captured.litellm_params or {}).get("cache") == {"ttl": 60, "no-cache": True, "no-store": True}
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
@pytest.mark.parametrize(
("cache_type", "expected_agent_cache"),
[
(LiteLLMCacheType.VALKEY_SEMANTIC, {"namespace": "team-a", "no-store": True}),
(LiteLLMCacheType.LOCAL, {"namespace": "team-a"}),
],
ids=["semantic", "exact-match"],
)
async def test_message_methods_forward_only_the_callers_semantic_cache_opt_outs_to_the_agent_call(
method: str, cache_type: LiteLLMCacheType, expected_agent_cache: Mapping[str, object]
) -> None:
agent: Final = _make_agent_mock()
agent.litellm_params = {"cache": {"namespace": "team-a"}}
mock_request: Final = _make_request_mock(
method, _HELLO_MESSAGE_PARAMS, cache={"namespace": "team-b", "use-cache": True, "ttl": 1, "no-store": True}
)
cache: Final = Cache(type=cache_type, host="localhost", port="6379", similarity_threshold=0.8)
with patch.object(litellm, "cache", cache):
captured: Final = await _invoke_message_method(
method,
mock_request,
UserAPIKeyAuth(api_key="sk-test", team_id="team-a"),
_add_proxy_data,
agent=agent,
)
assert (captured.litellm_params or {}).get("cache") == expected_agent_cache
@pytest.mark.asyncio
@pytest.mark.parametrize("method", ["message/send", "message/stream"])
async def test_message_methods_forward_caller_identity_headers(method: str):

View file

@ -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: 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()
@ -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,58 @@ 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
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})],
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: 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
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_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",

View file

@ -0,0 +1,295 @@
from collections.abc import Mapping
from typing import Final
import pytest
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_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,
CachingSupportedCallTypes,
LiteLLMCacheType,
)
ALLOWED_TEAM = "team-eval"
OTHER_TEAM = "team-other"
class _TeamAllowlist(CustomLogger):
def __init__(self, teams: frozenset[str]) -> None:
super().__init__()
self.teams = teams
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("model"), call_type))
return user_api_key_dict.team_id in self.teams
class _RaisingGate(CustomLogger):
async def async_should_semantic_cache(
self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str
) -> bool:
raise RuntimeError("allowlist service down")
class _TruthyGate(CustomLogger):
async def async_should_semantic_cache( # pyright: ignore[reportIncompatibleMethodOverride] # operator hooks are untyped Python
self, user_api_key_dict: UserAPIKeyAuth, data: Mapping[str, object], call_type: str
) -> str:
return "yes"
def _semantic_cache(
supported_call_types: tuple[CachingSupportedCallTypes, ...] = DEFAULT_CACHING_SUPPORTED_CALL_TYPES,
) -> Cache:
return Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
host="localhost",
port="6379",
similarity_threshold=0.8,
supported_call_types=list(supported_call_types),
)
def _request(**cache_controls: object) -> Mapping[str, object]:
return {
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "What color is the sky?"}],
**({"cache": cache_controls} if cache_controls else {}),
}
def _key(team_id: str) -> UserAPIKeyAuth:
return UserAPIKeyAuth(api_key="sk-hashed", team_id=team_id)
@pytest.mark.asyncio
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: Final = await gated_cache_overrides(
data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=_semantic_cache()
)
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() -> None:
data: Final = _request(ttl=60)
result: Final = await gated_cache_overrides(
data=data,
user_api_key_dict=_key(ALLOWED_TEAM),
call_type="acompletion",
gates=(_TeamAllowlist(frozenset({ALLOWED_TEAM})),),
cache=_semantic_cache(),
)
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) -> None:
gate: Final = _TeamAllowlist(frozenset())
data: Final = _request()
result: Final = await gated_cache_overrides(
data=data, user_api_key_dict=_key(OTHER_TEAM), call_type="acompletion", gates=(gate,), cache=cache
)
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) -> None:
gate: Final = _TeamAllowlist(frozenset())
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 == {}
assert gate.seen == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
("route", "served_call_type"),
[
("aresponses", "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_the_cache_through_a_bridge_whatever_call_types_it_serves(
route: str, served_call_type: CachingSupportedCallTypes
) -> None:
gate: Final = _TeamAllowlist(frozenset())
result: Final = await gated_cache_overrides(
data=_request(),
user_api_key_dict=_key(OTHER_TEAM),
call_type=route,
gates=(gate,),
cache=_semantic_cache(supported_call_types=(served_call_type,)),
)
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) -> None:
result: Final = await gated_cache_overrides(
data=_request(),
user_api_key_dict=_key(ALLOWED_TEAM),
call_type="aresponses",
gates=(gate,),
cache=_semantic_cache(),
)
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) -> None:
allow: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM}))
deny: Final = _TeamAllowlist(frozenset())
result: Final = await gated_cache_overrides(
data=_request(),
user_api_key_dict=_key(ALLOWED_TEAM),
call_type="anthropic_messages",
gates=(deny, allow) if deny_first else (allow, deny),
cache=_semantic_cache(),
)
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) -> None:
gate: Final = _TeamAllowlist(frozenset({ALLOWED_TEAM}))
monkeypatch.setattr(litellm, "callbacks", [CustomLogger(), gate])
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",
gates=gates,
cache=_semantic_cache(),
)
assert gates == (gate,)
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_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 = {
**_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

View file

@ -7,12 +7,22 @@ through the proxy router / guardrails, and the end-to-end
`handle_sampling_create_message` success and error-propagation behaviour.
"""
from collections.abc import Mapping
from types import SimpleNamespace
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from mcp.types import CreateMessageResult, ErrorData
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,
_run_guardrails_and_call_llm,
@ -118,6 +128,48 @@ 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) -> 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 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
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", router),
patch.object(litellm, "callbacks", [gate]),
patch.object(litellm, "cache", semantic_cache),
):
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 gate.deployment_cache_controls == [{"no-cache": True, "no-store": True}]
async def test_should_propagate_guardrail_rejection(self):
plo = MagicMock()
plo.pre_call_hook = AsyncMock(side_effect=ValueError("blocked by guardrail"))

View file

@ -514,6 +514,67 @@ class TestProxyBaseLLMRequestProcessing:
assert exc_info.value.param == "model"
add_litellm_data_to_request.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("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(
self,
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
from litellm.types.caching import LiteLLMCacheType
class _Gate(CustomLogger):
async def async_should_semantic_cache(
self, user_api_key_dict: ProxyUserAPIKeyAuth, data: Mapping[str, object], call_type: str
) -> bool:
return gate_allows
monkeypatch.setattr(litellm, "callbacks", [_Gate()])
monkeypatch.setattr(
litellm,
"cache",
Cache(type=LiteLLMCacheType.VALKEY_SEMANTIC, host="localhost", port="6379", similarity_threshold=0.8),
)
request: Final = Request(
{
"type": "http",
"method": "POST",
"path": "/v1/chat/completions",
"headers": [(b"content-type", b"application/json")],
"query_string": b"",
}
)
returned_data, _ = await ProxyBaseLLMRequestProcessing(
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={},
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-hashed", team_id="team-other"),
proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()),
proxy_config=RealProxyConfig(),
route_type="acompletion",
)
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(
("path", "route_type", "has_router", "priority_header", "body_priority", "expected_priority"),