mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-10 08:28:25 +00:00
Merge 543c9917aa into d037d86a94
This commit is contained in:
commit
78437fc6c3
19 changed files with 836 additions and 15 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)
|
||||
|
||||
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
108
litellm/proxy/common_utils/semantic_cache_gate.py
Normal file
108
litellm/proxy/common_utils/semantic_cache_gate.py
Normal 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),
|
||||
}
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
295
tests/unit/proxy/common_utils/test_semantic_cache_gate.py
Normal file
295
tests/unit/proxy/common_utils/test_semantic_cache_gate.py
Normal 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
|
||||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue