diff --git a/litellm/litellm_core_utils/routing_common_utils.py b/litellm/litellm_core_utils/routing_common_utils.py index a06860e7786..b46bdcff54a 100644 --- a/litellm/litellm_core_utils/routing_common_utils.py +++ b/litellm/litellm_core_utils/routing_common_utils.py @@ -1,10 +1,11 @@ import hashlib import json -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Final, TypeVar -from pydantic import TypeAdapter, ValidationError +from pydantic import BaseModel, TypeAdapter, ValidationError if TYPE_CHECKING: from litellm.types.llms.openai import OpenAIFileObject @@ -13,10 +14,12 @@ import litellm from litellm._logging import verbose_logger, verbose_router_logger from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.llms.sagemaker.common_utils import deployment_inference_component +from litellm.rust_bridge.public_call import optional_sequence from litellm.types.router import ( CredentialLiteLLMParams, Deployment, @@ -32,14 +35,31 @@ from litellm.types.utils import ( ) _V = TypeVar("_V") +_Picked: Final = TypeVar("_Picked") +_DeploymentT: Final = TypeVar("_DeploymentT", bound=Mapping[str, object]) MID_STREAM_CONTINUATION_KWARG: Final = "_mid_stream_continuation" ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors", MID_STREAM_CONTINUATION_KWARG}) _OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object]) +_GEMINI_SEARCH_TOOL_KEYS: Final = frozenset( + { + "googleSearch", + "google_search", + "googleSearchRetrieval", + "google_search_retrieval", + "enterpriseWebSearch", + "enterprise_web_search", + } +) +_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "google_search"}) +_OPENAI_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) _DEPLOYMENT_ROWS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) _DEPLOYMENT_IDS_ADAPTER: Final = TypeAdapter(tuple[str, ...]) +_RETRY_BREADCRUMBS_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) _FALLBACK_ENTRIES_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( tuple[Mapping[str, object], ...] ) @@ -408,10 +428,7 @@ def deployments_the_caller_may_use( return () -_WEB_SEARCH_TOOL_TYPES: Final = frozenset({"web_search", "web_search_preview"}) - - -def _deployment_supports_web_search(deployment: dict) -> bool: +def _deployment_supports_web_search(deployment: Mapping[str, object]) -> bool: """ Check if a deployment supports web search. @@ -423,12 +440,37 @@ def _deployment_supports_web_search(deployment: dict) -> bool: model_prices_and_context_window.json doesn't have supports_web_search tags on all models yet. TODO: backfill and add fallback. """ - model_info: Final = deployment.get("model_info", {}) + model_info: Final = deployment.get("model_info") + fields: Final = as_str_mapping(model_info.model_dump() if isinstance(model_info, BaseModel) else model_info) or {} + return bool(fields.get("supports_web_search", True)) - if "supports_web_search" in model_info: - return model_info["supports_web_search"] - return True +def _is_web_search_tool(tool: object) -> bool: + tool_fields: Final = as_str_mapping(tool) + if tool_fields is None: + return False + tool_type: Final = tool_fields.get("type") + if isinstance(tool_type, str) and (tool_type in _WEB_SEARCH_TOOL_TYPES or tool_type.startswith("web_search_")): + return True + return not _GEMINI_SEARCH_TOOL_KEYS.isdisjoint(tool_fields) and any( + tool_fields[key] is not None for key in _GEMINI_SEARCH_TOOL_KEYS.intersection(tool_fields) + ) + + +def _is_web_search_request(tools: object, web_search_options: object) -> bool: + return as_str_mapping(web_search_options) is not None or any( + _is_web_search_tool(tool) for tool in optional_sequence(tools) or () + ) + + +def web_search_capable_deployments( + deployments: Sequence[_DeploymentT], request_kwargs: Mapping[str, object] | None +) -> tuple[_DeploymentT, ...]: + if request_kwargs is None or not _is_web_search_request( + request_kwargs.get("tools"), request_kwargs.get("web_search_options") + ): + return () + return tuple(d for d in deployments if _deployment_supports_web_search(d)) def filter_web_search_deployments( @@ -449,7 +491,7 @@ def filter_web_search_deployments( is_web_search_request: Final = isinstance(tools, (list, tuple)) and any( isinstance(tool, dict) and isinstance(tool_type := tool.get("type"), str) - and tool_type in _WEB_SEARCH_TOOL_TYPES + and tool_type in _OPENAI_WEB_SEARCH_TOOL_TYPES for tool in tools ) if not is_web_search_request: @@ -457,11 +499,148 @@ def filter_web_search_deployments( # Filter out deployments that don't support web search final_deployments: Final = [d for d in healthy_deployments if _deployment_supports_web_search(d)] - if len(healthy_deployments) > 0 and len(final_deployments) == 0: + if healthy_deployments and not final_deployments: verbose_logger.warning("No deployments support web search for request") return final_deployments +def _deployment_id(deployment: Mapping[str, object]) -> object: + model_info: Final = as_str_mapping(deployment.get("model_info")) + return None if model_info is None else model_info.get("id") + + +def _retry_breadcrumbs(metadata: object) -> tuple[Mapping[str, object], ...]: + fields: Final = as_str_mapping(metadata) + breadcrumbs: Final = None if fields is None else fields.get("previous_models") + if breadcrumbs is None: + return () + try: + return _RETRY_BREADCRUMBS_ADAPTER.validate_python(breadcrumbs) + except ValidationError: + return () + + +def _retry_history_was_trimmed(metadata: object) -> bool: + fields: Final = as_str_mapping(metadata) + retry_count: Final = None if fields is None else fields.get("request_retry_count") + return type(retry_count) is int and retry_count > len(_retry_breadcrumbs(metadata)) + + +def _breadcrumb_deployment_ids(breadcrumb: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + deployment_id + for deployment_id in (breadcrumb.get("deployment_id"), breadcrumb.get("configured_deployment_id")) + if isinstance(deployment_id, str) + ) + + +def _deployment_ids_retried_in_request(request_kwargs: Mapping[str, object]) -> tuple[str, ...]: + breadcrumbs: Final = ( + *_retry_breadcrumbs(request_kwargs.get("metadata")), + *_retry_breadcrumbs(request_kwargs.get("litellm_metadata")), + ) + return tuple(chain.from_iterable(map(_breadcrumb_deployment_ids, breadcrumbs))) + + +def _web_search_tried_deployment_ids( + request_kwargs: Mapping[str, object] | None, at_rate_limit: frozenset[str] +) -> frozenset[str] | None: + if request_kwargs is None or not _is_web_search_request( + request_kwargs.get("tools"), request_kwargs.get("web_search_options") + ): + return None + if _retry_history_was_trimmed(request_kwargs.get("metadata")) or _retry_history_was_trimmed( + request_kwargs.get("litellm_metadata") + ): + verbose_router_logger.debug("This request's retry history was trimmed, so every deployment stays") + return None + return frozenset( + ( + *as_excluded_deployment_ids(request_kwargs.get("_excluded_deployment_ids")), + *as_retry_skipped_deployment_ids(request_kwargs.get("_retry_skipped_deployment_ids")), + *_deployment_ids_retried_in_request(request_kwargs), + *at_rate_limit, + ) + ) + + +def prefer_web_search_deployments( + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), +) -> list[_DeploymentT]: # mutable-ok: returns a mutable deployment list + tried: Final = _web_search_tried_deployment_ids(request_kwargs, at_rate_limit) + if tried is None: + return deployments + untried: Final = [d for d in deployments if _deployment_id(d) not in tried] + orders_with_an_untried_capable_deployment: Final = tuple( + litellm.utils.get_deployment_order(d) for d in untried if _deployment_supports_web_search(d) + ) + if not orders_with_an_untried_capable_deployment: + if not untried or not any(_deployment_supports_web_search(d) for d in deployments): + verbose_router_logger.debug("No untried deployment for this web search request can search, so all stay") + return deployments + if any(_deployment_supports_web_search(d) and _deployment_id(d) in at_rate_limit for d in deployments): + verbose_router_logger.debug("A deployment that can search is only at its limit, so all stay") + return deployments + verbose_router_logger.debug("Every deployment that can search was tried, so the untried ones stay") + return untried + return [ + d + for d in deployments + if litellm.utils.get_deployment_order(d) not in orders_with_an_untried_capable_deployment + or (_deployment_supports_web_search(d) and _deployment_id(d) not in tried) + ] + + +def web_search_preference_can_narrow(deployments: Sequence[_DeploymentT], request_kwargs: Mapping[str, object]) -> bool: + return 0 < len(web_search_capable_deployments(deployments, request_kwargs)) < len(deployments) + + +def web_search_alternatives_to_pin( + pinned: _DeploymentT, + deployments: Sequence[_DeploymentT], + request_kwargs: Mapping[str, object], + at_rate_limit: frozenset[str] = frozenset(), +) -> tuple[_DeploymentT, ...]: + tried: Final = _web_search_tried_deployment_ids(request_kwargs, at_rate_limit) + if tried is None or (_deployment_supports_web_search(pinned) and _deployment_id(pinned) not in tried): + return () + return tuple(d for d in deployments if _deployment_supports_web_search(d) and _deployment_id(d) not in tried) + + +def pick_preferring_web_search( + pick: Callable[[list[_DeploymentT]], _Picked | None], # mutable-ok: router deployment-list contract + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), +) -> _Picked | None: + preferred: Final = prefer_web_search_deployments(deployments, request_kwargs, at_rate_limit) + if len(preferred) == len(deployments): + return pick(deployments) + try: + from_preferred: Final = pick(preferred) + except (litellm.RateLimitError, ValueError): + return pick(deployments) + return from_preferred if from_preferred is not None else pick(deployments) + + +async def apick_preferring_web_search( + pick: Callable[[list[_DeploymentT]], Awaitable[_Picked | None]], # mutable-ok: router deployment-list contract + deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + request_kwargs: Mapping[str, object] | None, + at_rate_limit: frozenset[str] = frozenset(), +) -> _Picked | None: + preferred: Final = prefer_web_search_deployments(deployments, request_kwargs, at_rate_limit) + if len(preferred) == len(deployments): + return await pick(deployments) + try: + from_preferred: Final = await pick(preferred) + except (litellm.RateLimitError, ValueError): + return await pick(deployments) + return from_preferred if from_preferred is not None else await pick(deployments) + + # Credential params that only one provider family reads, paired with the providers # that read them. A deployment carrying them while resolving elsewhere is almost # always a missing route prefix: `model: claude-sonnet-5` with `aws_region_name` diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index 08f717ee81c..163f9185c74 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -126,6 +126,7 @@ from litellm.litellm_core_utils.clientside_credential_handler import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, coerce_token_limit, get_litellm_metadata_from_kwargs, get_metadata_variable_name_from_kwargs, @@ -187,6 +188,7 @@ from litellm.litellm_core_utils.request_timeout_resolver import ( ) from litellm.litellm_core_utils.routing_common_utils import ( MID_STREAM_CONTINUATION_KWARG, + apick_preferring_web_search, as_retry_skipped_deployment_ids, deployment_model_id_pin, deployments_the_caller_may_use, @@ -198,6 +200,8 @@ from litellm.litellm_core_utils.routing_common_utils import ( format_no_fallback_group_message, get_request_team_id, is_proxy_admin_request, + pick_preferring_web_search, + prefer_web_search_deployments, provider_for_generic_call, request_user_api_key_auth, resolve_model_group_alias, @@ -205,6 +209,7 @@ from litellm.litellm_core_utils.routing_common_utils import ( truncate_fallback_error_detail, warn_on_ignored_model_capability_flags, warn_on_provider_credential_mismatch, + web_search_capable_deployments, with_admin_api_key_secret_reference, without_model_capability_flags, without_router_only_kwargs, @@ -3278,6 +3283,8 @@ class PythonCore(DeprecatedRouterMembers): model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, + deployment_ids_at_limit=self._async_deployment_ids_at_rate_limit, + model_group_deployments=self._model_group_deployments, ) self.optional_callbacks.append(affinity_callback) self._host.callbacks.add_litellm_callback(affinity_callback) @@ -3325,6 +3332,8 @@ class PythonCore(DeprecatedRouterMembers): model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, + deployment_ids_at_limit=self._async_deployment_ids_at_rate_limit, + model_group_deployments=self._model_group_deployments, ) self.optional_callbacks.append(affinity_callback) self._host.callbacks.add_litellm_callback(affinity_callback) @@ -10017,6 +10026,7 @@ class PythonCore(DeprecatedRouterMembers): await self._repin_session_after_success( _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), + _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("optional_params")), ) # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits @@ -10042,7 +10052,9 @@ class PythonCore(DeprecatedRouterMembers): "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) - async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: + async def _repin_session_after_success( + self, litellm_params: Mapping[str, object], response: object, optional_params: object = None + ) -> None: if not self.session_affinity_repin_on_failover: return try: @@ -10068,7 +10080,12 @@ class PythonCore(DeprecatedRouterMembers): model_ids=list(registered_ids), parent_otel_span=None ) ) - await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) + await callback.async_repin_session_after_success( + litellm_params, + model_group, + registered_ids - cooldown_ids, + _LITELLM_PARAMS_ADAPTER.validate_python(optional_params or {}), + ) except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting verbose_router_logger.debug("Session re-pinning failed: %s", error) @@ -10396,14 +10413,15 @@ class PythonCore(DeprecatedRouterMembers): _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs model_group: Final = kwargs.get("model") - model_info: Final = request_metadata.get("model_info") - deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None + model_info: Final = as_str_mapping(request_metadata.get("model_info")) or {} + deployment_id: Final = model_info.get("id") + configured_deployment_id: Final = model_info.get("original_model_id") base_model: Final = get_base_model_from_litellm_call_metadata( metadata=request_metadata ) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None) deployment_model: Final = request_metadata.get("deployment") attempted_retries: Final = request_metadata.get("attempted_retries") - attempt_record: Final[RetryAttemptRecord] = { + attempt_fields: Final[RetryAttemptRecord] = { "model_group": model_group if isinstance(model_group, str) else None, "deployment_id": deployment_id if isinstance(deployment_id, str) else None, "model": deployment_model if isinstance(deployment_model, str) else None, @@ -10412,6 +10430,11 @@ class PythonCore(DeprecatedRouterMembers): "exception_string": str(e), "attempted_retries": attempted_retries if type(attempted_retries) is int else None, } + attempt_record: Final[RetryAttemptRecord] = ( + {**attempt_fields, "configured_deployment_id": configured_deployment_id} + if isinstance(configured_deployment_id, str) + else attempt_fields + ) earlier_breadcrumbs: Final = request_metadata.get("previous_models") kept_breadcrumbs: Final[tuple[object, ...]] = ( tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] @@ -10658,6 +10681,24 @@ class PythonCore(DeprecatedRouterMembers): healthy_deployments = self._filter_blocked_deployments(healthy_deployments) return healthy_deployments, _all_deployments + def _model_rate_limit_checks(self) -> tuple[ModelRateLimitingCheck, ...]: + return tuple(callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)) + + def _model_group_deployments(self, model_group: str) -> Sequence[Mapping[str, object]]: + return self._core_get_model_list(model_name=model_group) or () + + async def _async_deployment_ids_at_rate_limit( + self, deployments: Sequence[dict[str, object]], parent_otel_span: Span | None + ) -> frozenset[str]: + return frozenset( + chain.from_iterable( + [ + await check.async_deployment_ids_at_limit(deployments, parent_otel_span) + for check in self._model_rate_limit_checks() + ] + ) + ) + def routing_strategy_pre_call_checks(self, deployment: dict): """ Mimics 'async_routing_strategy_pre_call_checks' @@ -16230,26 +16271,39 @@ class PythonCore(DeprecatedRouterMembers): return healthy_deployments[0] start_time: Final = time.time() + routing_kwargs: Final = cast(Mapping[str, object], request_kwargs) # cast-ok: str-keyed request kwargs + web_search_capable: Final = web_search_capable_deployments( + request_kwargs=routing_kwargs, + deployments=cast(list[dict[str, object]], healthy_deployments), # cast-ok: str-keyed dicts + ) + at_rate_limit: Final = await self._async_deployment_ids_at_rate_limit(web_search_capable, parent_otel_span) if strategy == "simple-shuffle": shuffled: Final = simple_shuffle( resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, + healthy_deployments=prefer_web_search_deployments( + healthy_deployments, routing_kwargs, at_rate_limit + ), model=model, - request_kwargs=request_kwargs, + request_kwargs=routing_kwargs, ) phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) return shuffled with PrefetchedUsage.scoped( routing_read_batch.prefetched_usage if routing_read_batch is not None else None ): - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, + deployment: Final = await apick_preferring_web_search( + lambda candidates: self._select_deployment_async( + strategy=strategy, + selector=strategy_selector, + model=model, + healthy_deployments=candidates, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ), + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + routing_kwargs, + at_rate_limit, ) if deployment is None: exception: Final = await async_raise_no_deployment_exception( @@ -17277,25 +17331,39 @@ class PythonCore(DeprecatedRouterMembers): ), ) + web_search_capable: Final = web_search_capable_deployments( + request_kwargs=routing_kwargs, + deployments=cast(list[dict[str, object]], healthy_deployments), # cast-ok: str-keyed dicts + ) + at_rate_limit: Final = frozenset( + chain.from_iterable( + check.deployment_ids_at_limit(web_search_capable) for check in self._model_rate_limit_checks() + ) + ) if strategy == "simple-shuffle": # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm ############## Check 'weight' param set for weighted pick ################# shuffled: Final = simple_shuffle( resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, + healthy_deployments=prefer_web_search_deployments(healthy_deployments, routing_kwargs, at_rate_limit), model=model, - request_kwargs=request_kwargs, + request_kwargs=routing_kwargs, ) phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) return shuffled - deployment: Final = self._select_deployment_sync( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, + deployment: Final = pick_preferring_web_search( + lambda candidates: self._select_deployment_sync( + strategy=strategy, + selector=strategy_selector, + model=model, + healthy_deployments=candidates, + messages=messages, + input=input, + request_kwargs=request_kwargs, + ), + cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts + routing_kwargs, + at_rate_limit, ) if deployment is None: diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index dbf471d875c..dfb2fe7e5c6 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,8 +13,9 @@ where routing to a consistent deployment is still beneficial. """ import hashlib -from collections.abc import Callable, Mapping, Sequence -from typing import Final, cast +from collections.abc import Awaitable, Callable, Mapping, Sequence +from types import MappingProxyType +from typing import Final, TypeVar, cast from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict @@ -25,16 +26,24 @@ from litellm.caching.affinity_cache import ( ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, claim_affinity_pin_in_memory, + read_affinity_pin, set_local_affinity_pin, ) from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span -from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request +from litellm.litellm_core_utils.routing_common_utils import ( + deployment_ids_failed_in_request, + web_search_alternatives_to_pin, + web_search_capable_deployments, + web_search_preference_can_narrow, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes +_DeploymentT: Final = TypeVar("_DeploymentT", bound=dict[str, object]) + class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] @@ -110,6 +119,9 @@ class DeploymentAffinityCheck(CustomLogger): model_group_affinity_config: dict[str, list[str]] | None = None, is_priority_group: Callable[[str], bool] | None = None, session_affinity_repin_on_failover: bool = False, + deployment_ids_at_limit: Callable[[Sequence[dict[str, object]], Span | None], Awaitable[frozenset[str]]] + | None = None, + model_group_deployments: Callable[[str], Sequence[Mapping[str, object]]] | None = None, ): super().__init__() self.cache = cache @@ -120,6 +132,8 @@ class DeploymentAffinityCheck(CustomLogger): self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} self.is_priority_group = is_priority_group self.session_affinity_repin_on_failover = session_affinity_repin_on_failover + self.deployment_ids_at_limit = deployment_ids_at_limit + self.model_group_deployments = model_group_deployments def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -366,7 +380,7 @@ class DeploymentAffinityCheck(CustomLogger): return self._pinned_model_id(winner) @staticmethod - def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: + def _find_deployment_by_model_id(healthy_deployments: Sequence[_DeploymentT], model_id: str) -> _DeploymentT | None: for deployment in healthy_deployments: model_info = deployment.get("model_info") if not isinstance(model_info, dict): @@ -391,7 +405,7 @@ class DeploymentAffinityCheck(CustomLogger): 2. cached API-key deployment affinity """ request_kwargs = request_kwargs or {} - typed_healthy_deployments: Final = cast(list[dict], healthy_deployments) + typed_healthy_deployments: Final = cast(list[dict[str, object]], healthy_deployments) # cast-ok: router dicts if request_kwargs.get("_target_order") is not None: return typed_healthy_deployments @@ -463,17 +477,16 @@ class DeploymentAffinityCheck(CustomLogger): session_model_id = session_cache_result if session_model_id and session_model_id not in failed_deployment_ids: - session_deployment: Final = self._find_deployment_by_model_id( - healthy_deployments=typed_healthy_deployments, - model_id=session_model_id, + session_deployments: Final = await self._pinned_deployments( + typed_healthy_deployments, session_model_id, request_kwargs, parent_otel_span ) - if session_deployment is not None: + if session_deployments is not None: verbose_router_logger.debug( "DeploymentAffinityCheck: session-id affinity hit -> deployment=%s session_id=%s", session_model_id, session_id, ) - return [session_deployment] + return list(session_deployments) else: verbose_router_logger.debug( "DeploymentAffinityCheck: session-id pinned deployment=%s not found in healthy_deployments", @@ -500,11 +513,10 @@ class DeploymentAffinityCheck(CustomLogger): if not model_id or model_id in failed_deployment_ids: return typed_healthy_deployments - deployment = self._find_deployment_by_model_id( - healthy_deployments=typed_healthy_deployments, - model_id=model_id, + pinned_deployments: Final = await self._pinned_deployments( + typed_healthy_deployments, model_id, request_kwargs, parent_otel_span ) - if deployment is None: + if pinned_deployments is None: verbose_router_logger.debug( "DeploymentAffinityCheck: pinned deployment=%s not found in healthy_deployments", model_id, @@ -516,7 +528,34 @@ class DeploymentAffinityCheck(CustomLogger): model_id, self._shorten_for_logs(self._hash_user_key(user_key)), ) - return [deployment] + return list(pinned_deployments) + + async def _pinned_deployments( + self, + healthy_deployments: list[_DeploymentT], # mutable-ok: router deployment-list contract + model_id: str, + request_kwargs: Mapping[str, object], + parent_otel_span: Span | None, + ) -> tuple[_DeploymentT, ...] | None: + pinned: Final = self._find_deployment_by_model_id(healthy_deployments=healthy_deployments, model_id=model_id) + if pinned is None: + return None + at_rate_limit: Final = await self._web_search_deployment_ids_at_limit( + healthy_deployments, request_kwargs, parent_otel_span + ) + return (*web_search_alternatives_to_pin(pinned, healthy_deployments, request_kwargs, at_rate_limit), pinned) + + async def _web_search_deployment_ids_at_limit( + self, + deployments: Sequence[_DeploymentT], + request_kwargs: Mapping[str, object], + parent_otel_span: Span | None, + ) -> frozenset[str]: + if self.deployment_ids_at_limit is None: + return frozenset() + return await self.deployment_ids_at_limit( + web_search_capable_deployments(deployments, request_kwargs), parent_otel_span + ) async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: redis_cache: Final = self.cache.redis_cache @@ -560,6 +599,7 @@ class DeploymentAffinityCheck(CustomLogger): request_kwargs: Mapping[str, object], model_group: str, eligible_deployment_ids: frozenset[str], + request_params: Mapping[str, object] = MappingProxyType({}), ) -> None: marker_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs) if not self._get_effective_flags(model_group)[2] and marker_ttl is None: @@ -574,11 +614,17 @@ class DeploymentAffinityCheck(CustomLogger): cache_key: Final = self.get_session_affinity_cache_key( model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) ) + ttl_seconds: Final = marker_ttl if marker_ttl is not None else self.ttl_seconds + is_search_in_mixed_group: Final = self.model_group_deployments is not None and web_search_preference_can_narrow( + self.model_group_deployments(model_group), request_params + ) + if is_search_in_mixed_group and await read_affinity_pin(self.cache, cache_key, ttl_seconds) is None: + return await claim_affinity_pin( self.cache, cache_key, DeploymentAffinityCacheValue(model_id=str(model_id)), - marker_ttl if marker_ttl is not None else self.ttl_seconds, + ttl_seconds, eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids), preserve_legacy_model_ids=True, ) @@ -635,6 +681,14 @@ class DeploymentAffinityCheck(CustomLogger): if not ((enable_user_key and user_key is not None) or session_id is not None): return None + if self.model_group_deployments is not None and web_search_preference_can_narrow( + self.model_group_deployments(deployment_model_name), kwargs + ): + verbose_router_logger.debug( + "DeploymentAffinityCheck: web search request in a group with deployments that cannot search sets no pin" + ) + return None + model_info = kwargs.get("model_info") if not isinstance(model_info, dict): model_info = None diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index f887c4551f2..571c2d7f59d 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -156,6 +156,16 @@ def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | return itpm_key, otpm_key +def io_token_limits_by_cache_key( + deployment: dict[str, object], # mutable-ok: the io token readers take router deployment dicts + current_minute: str, +) -> tuple[tuple[str, int | None], ...]: + cache_keys: Final = _get_cache_keys(deployment, current_minute) + if cache_keys is None: + return () + return tuple(zip(cache_keys, get_deployment_io_token_limits(deployment), strict=True)) + + def _served_model(deployment: object) -> str: params: Final = deployment.get("litellm_params") if is_str_keyed_mapping(deployment) else None served: Final = params.get("model") if is_str_keyed_mapping(params) else None diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index f41d19cd62c..4dd98895e94 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -11,13 +11,16 @@ When a deployment sets both itpm/otpm and tpm/rpm, both are enforced. A warning is logged the first time such a deployment is seen. """ +import asyncio import contextlib -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timedelta +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal import httpx -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, ValidationError import litellm from litellm._internal_context import with_service_target @@ -31,6 +34,7 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( async_io_token_refund_failure, deployment_has_io_token_limits, get_io_token_rate_limit_request_kwargs, + io_token_limits_by_cache_key, io_token_pre_call_check, io_token_reconcile_success, io_token_refund_failure, @@ -116,6 +120,32 @@ class _DailyLimitDeployment(BaseModel): ) +def _deployment_model_id(deployment: Mapping[str, object]) -> object: + model_info: Final = deployment.get("model_info") + return model_info.get("id") if isinstance(model_info, Mapping) else None + + +@dataclass(frozen=True, slots=True) +class _LimitCounter: + model_id: str + key: str + limit: float + increments_first: bool + + def reached(self, usage: object) -> bool: + if self.increments_first: + return (usage if isinstance(usage, (int, float)) else 0) + 1 > self.limit + return isinstance(usage, (int, float)) and usage >= self.limit + + +def _counters_with_limits(model_id: str, counters: Sequence[tuple[str, object, bool]]) -> tuple[_LimitCounter, ...]: + return tuple( + _LimitCounter(model_id, key, limit, increments_first) + for key, limit, increments_first in counters + if isinstance(limit, (int, float)) + ) + + def _reset_offset_from_kwargs(kwargs: Mapping[str, object]) -> timedelta: litellm_params: Final = kwargs.get("litellm_params") if litellm_params is None: @@ -277,7 +307,7 @@ class ModelRateLimitingCheck(CustomLogger): return tpm_key, rpm_key @with_service_target(ROUTER_USAGE_TARGET) - def _get_current_tpm(self, tpm_key: str, tpm_limit: int) -> int | None: + def _get_current_tpm(self, tpm_key: str, tpm_limit: float) -> int | None: local_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): @@ -288,7 +318,7 @@ class ModelRateLimitingCheck(CustomLogger): return local_tpm @with_service_target(ROUTER_USAGE_TARGET) - async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: int, parent_otel_span: Span | None) -> int | None: + async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: float, parent_otel_span: Span | None) -> int | None: local_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) redis_cache: Final = self.dual_cache.redis_cache if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): @@ -298,6 +328,87 @@ class ModelRateLimitingCheck(CustomLogger): except RedisCircuitBreakerOpenError: return local_tpm + def _limit_counters( + self, + deployment: dict[str, object], # mutable-ok: the limit and key readers take router deployment dicts + now: datetime, + ) -> tuple[_LimitCounter, ...]: + model_id: Final = _deployment_model_id(deployment) + if not isinstance(model_id, str): + return () + tpm_limit, rpm_limit = self._get_deployment_limits(deployment) + current_minute: Final = now.strftime("%H-%M") + tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + minute_counters: Final = ( + (tpm_key, tpm_limit, False), + (rpm_key, rpm_limit, True), + *((key, limit, True) for key, limit in io_token_limits_by_cache_key(deployment, current_minute)), + ) + try: + daily: Final = _DailyLimitDeployment.model_validate(deployment) + except ValidationError: + return _counters_with_limits(model_id, minute_counters) + tpd_limit, rpd_limit = daily.limits() + tpd_key, rpd_key = daily.cache_keys(now) + return _counters_with_limits( + model_id, (*minute_counters, (tpd_key, tpd_limit, False), (rpd_key, rpd_limit, True)) + ) + + def _counters_of( + self, + deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts + ) -> tuple[_LimitCounter, ...]: + now: Final = self.clock() + return tuple(chain.from_iterable(self._limit_counters(deployment, now) for deployment in deployments)) + + def _local_usage(self, key: str) -> object: + return self.dual_cache.get_cache(key=key, local_only=True) + + def _redis_usages(self, keys: list[str]) -> Mapping[str, object] | None: + redis_cache: Final = self.dual_cache.redis_cache + if redis_cache is None or not keys: + return None + try: + return redis_cache.batch_get_cache(key_list=keys) + except RedisCircuitBreakerOpenError: + return None + + @with_service_target(ROUTER_USAGE_TARGET) + def _current_usages(self, counters: Sequence[_LimitCounter]) -> tuple[object, ...]: + local_usages: Final = tuple(self._local_usage(counter.key) for counter in counters) + redis_usages: Final = self._redis_usages( + [counter.key for counter, usage in zip(counters, local_usages, strict=True) if not counter.reached(usage)] + ) + if redis_usages is None: + return local_usages + return tuple( + usage if counter.reached(usage) else redis_usages.get(counter.key) + for counter, usage in zip(counters, local_usages, strict=True) + ) + + def deployment_ids_at_limit( + self, + deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts + ) -> frozenset[str]: + counters: Final = self._counters_of(deployments) + usages: Final = self._current_usages(counters) + return frozenset( + counter.model_id for counter, usage in zip(counters, usages, strict=True) if counter.reached(usage) + ) + + async def async_deployment_ids_at_limit( + self, + deployments: Sequence[dict[str, object]], # mutable-ok: _limit_counters reads router deployment dicts + parent_otel_span: Span | None, + ) -> frozenset[str]: + counters: Final = self._counters_of(deployments) + usages: Final = await asyncio.gather( + *(self._async_get_current_tpm(counter.key, counter.limit, parent_otel_span) for counter in counters) + ) + return frozenset( + counter.model_id for counter, usage in zip(counters, usages, strict=True) if counter.reached(usage) + ) + @with_service_target(ROUTER_USAGE_TARGET) def pre_call_check(self, deployment: dict) -> dict | None: """ diff --git a/litellm/types/router.py b/litellm/types/router.py index 4e3e4d92b3c..3eaed4ab055 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -39,7 +39,7 @@ from pydantic import ( field_validator, model_validator, ) -from typing_extensions import Protocol, ReadOnly, Required, TypedDict, runtime_checkable +from typing_extensions import NotRequired, Protocol, ReadOnly, Required, TypedDict, runtime_checkable from litellm._logging import verbose_logger from litellm._uuid import uuid @@ -1272,6 +1272,7 @@ class RetryAttemptRecord(TypedDict): exception_type: ReadOnly[str] exception_string: ReadOnly[str] attempted_retries: ReadOnly[int | None] + configured_deployment_id: ReadOnly[NotRequired[str]] VALID_LITELLM_ENVIRONMENTS = [ diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 57ae9f1905f..8803cd695e9 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -179,6 +179,8 @@ ignored_function_names = [ "_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py "_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py + "_model_rate_limit_checks", # Tested through the web search pin rpm tests in test_deployment_affinity_check.py + "_async_deployment_ids_at_rate_limit", # Tested through the web search pin rpm tests in test_deployment_affinity_check.py "_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py "priority_header_update", # Tested through the priority header tests in test_common_request_processing.py "_drops_params", # Read by split_request_priority and priority_header_update, tested through both diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index fa75f980cc0..6b3a651361a 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -2,6 +2,7 @@ import asyncio import itertools import json from collections.abc import Sequence +from datetime import datetime, timezone from typing import Final from unittest.mock import AsyncMock, patch @@ -1299,3 +1300,556 @@ async def test_previous_response_id_pin_holds_even_after_its_deployment_failed() }, parent_otel_span=None, ) == [healthy_deployments[0]] + + +def _mixed_web_search_router( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + orders: tuple[int | None, int | None] = (None, None), + search_rpm: int | None = None, + model_group_affinity_config: dict[str, list[str]] | None = None, # mutable-ok: Router signature +) -> litellm.Router: + no_search_order, search_order = orders + return litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + **({} if no_search_order is None else {"order": no_search_order}), + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + **({} if search_order is None else {"order": search_order}), + **({} if search_rpm is None else {"rpm": search_rpm}), + }, + "model_info": {"id": "search"}, + }, + ], + optional_pre_call_checks=optional_pre_call_checks, + model_group_affinity_config=model_group_affinity_config, + num_retries=0, + ) + + +_CALLER_METADATA: Final = {"user_api_key_hash": "caller-1", "session_id": "session-1"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "web_search_request", + [{"web_search_options": {}}, {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}], + ids=["web-search-options", "web-search-tool"], +) +@pytest.mark.parametrize( + ("optional_pre_call_checks", "pin_key"), + [ + ( + ["deployment_affinity"], + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ( + ["session_affinity"], + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ], + ids=["user-key-pin", "session-pin"], +) +async def test_pin_on_a_deployment_without_web_search_yields_to_a_web_search_request( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + pin_key: str, + web_search_request: dict[str, object], # mutable-ok: request kwargs +): + router: Final = _mixed_web_search_router(optional_pre_call_checks) + await router.cache.async_set_cache(pin_key, {"model_id": "no-search"}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + plain: Final = await router.acompletion(model="mixed", messages=messages, metadata=dict(_CALLER_METADATA)) + searching: Final = await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), **web_search_request + ) + + assert plain._hidden_params["model_id"] == "no-search" + assert searching._hidden_params["model_id"] == "search" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("orders", "searched_on"), + [((2, 1), "search"), ((None, 1), "search"), ((1, 2), "no-search"), ((1, None), "no-search")], + ids=["search-tier-first", "only-search-tier-ordered", "pinned-tier-first", "only-pinned-tier-ordered"], +) +@pytest.mark.parametrize( + ("optional_pre_call_checks", "pin_key"), + [ + ( + ["deployment_affinity"], + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ( + ["session_affinity"], + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ], + ids=["user-key-pin", "session-pin"], +) +async def test_pin_without_web_search_yields_to_a_web_search_request_only_when_search_sits_in_an_earlier_order_tier( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + pin_key: str, + orders: tuple[int | None, int | None], + searched_on: str, +): + router: Final = _mixed_web_search_router(optional_pre_call_checks, orders) + await router.cache.async_set_cache(pin_key, {"model_id": "no-search"}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + plain: Final = await router.acompletion(model="mixed", messages=messages, metadata=dict(_CALLER_METADATA)) + searching: Final = await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), web_search_options={} + ) + + assert plain._hidden_params["model_id"] == "no-search" + assert searching._hidden_params["model_id"] == searched_on + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "orders", [(None, 1), (2, 1), (None, None)], ids=["pin-unordered", "pin-in-later-tier", "same-tier"] +) +@pytest.mark.parametrize( + ("affinity_check", "per_group", "pin_key"), + [ + ( + "deployment_affinity", + False, + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ( + "session_affinity", + False, + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ( + "deployment_affinity", + True, + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + ), + ], + ids=["user-key-pin", "session-pin", "per-group-user-key-pin"], +) +async def test_pin_without_web_search_holds_for_a_web_search_request_once_the_earlier_search_tier_hits_its_rpm( + affinity_check: str, + per_group: bool, + pin_key: str, + orders: tuple[int | None, int | None], +): + router: Final = _mixed_web_search_router( + ["enforce_model_rate_limits"] if per_group else [affinity_check, "enforce_model_rate_limits"], + orders, + search_rpm=1, + model_group_affinity_config={"mixed": [affinity_check]} if per_group else None, + ) + await router.cache.async_set_cache(pin_key, {"model_id": "no-search"}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + with patch( + "litellm.router_utils.pre_call_checks.model_rate_limit_check.get_utc_datetime", + return_value=datetime(2026, 1, 1, 12, 30, 15, tzinfo=timezone.utc), + ): + served: Final = [ + await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), web_search_options={} + ) + for _ in range(4) + ] + + assert [response._hidden_params["model_id"] for response in served] == ["search", *["no-search"] * 3] + + +def _tiered_web_search_router( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + deployments: tuple[tuple[str, bool, int | None, int | None], ...], +) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + **({} if order is None else {"order": order}), + **({} if rpm is None else {"rpm": rpm}), + }, + "model_info": {"id": model_id, **({} if supports_web_search else {"supports_web_search": False})}, + } + for model_id, supports_web_search, order, rpm in deployments + ], + optional_pre_call_checks=optional_pre_call_checks, + num_retries=0, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("deployments", "pinned_id", "served_ids"), + [ + ( + (("search-early", True, 1, 1), ("no-search", False, None, None), ("search", True, None, None)), + "no-search", + ["search-early", *["search"] * 3], + ), + ( + (("search-early", True, 1, 1), ("no-search", False, 2, None), ("search", True, 2, None)), + "no-search", + ["search-early", *["search"] * 3], + ), + ( + (("no-search-early", False, 1, None), ("no-search", False, 2, None), ("search", True, 2, None)), + "no-search", + ["search"] * 4, + ), + ( + (("search-pinned", True, None, 1), ("search", True, None, None)), + "search-pinned", + ["search-pinned", *["search"] * 3], + ), + ], + ids=[ + "rate-limited-search-before-unordered-pin", + "rate-limited-search-before-pin-tier", + "flagged-before-pin-tier", + "pinned-search-at-its-rpm", + ], +) +@pytest.mark.parametrize( + ("affinity_check", "pin_key"), + [ + ("deployment_affinity", DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1")), + ( + "session_affinity", + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="mixed", session_id="session-1", user_key="caller-1" + ), + ), + ], + ids=["user-key-pin", "session-pin"], +) +async def test_web_search_request_leaves_a_pin_only_for_an_untried_search_deployment_under_its_limit( + affinity_check: str, + pin_key: str, + deployments: tuple[tuple[str, bool, int | None, int | None], ...], + pinned_id: str, + served_ids: list[str], # mutable-ok: expected id sequence +): + router: Final = _tiered_web_search_router([affinity_check, "enforce_model_rate_limits"], deployments) + await router.cache.async_set_cache(pin_key, {"model_id": pinned_id}, ttl=3600) + messages: Final = [{"role": "user", "content": "hi"}] + + with patch( + "litellm.router_utils.pre_call_checks.model_rate_limit_check.get_utc_datetime", + return_value=datetime(2026, 1, 1, 12, 30, 15, tzinfo=timezone.utc), + ): + served: Final = [ + await router.acompletion( + model="mixed", messages=messages, metadata=dict(_CALLER_METADATA), web_search_options={} + ) + for _ in range(4) + ] + + assert [response._hidden_params["model_id"] for response in served] == served_ids + + +@pytest.mark.asyncio +async def test_session_pin_on_a_later_order_tier_holds_for_a_web_search_request_when_it_can_search(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "tiered", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + "order": order, + }, + "model_info": {"id": model_id}, + } + for model_id, order in (("search-first-tier", 1), ("search-second-tier", 2)) + ], + optional_pre_call_checks=["session_affinity"], + ) + await router.cache.async_set_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key( + model_group="tiered", session_id="session-1", user_key="caller-1" + ), + {"model_id": "search-second-tier"}, + ttl=3600, + ) + + searching: Final = await router.acompletion( + model="tiered", + messages=[{"role": "user", "content": "hi"}], + metadata=dict(_CALLER_METADATA), + web_search_options={}, + ) + + assert searching._hidden_params["model_id"] == "search-second-tier" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_supported", "tried"), + [((False, False), ()), ((False, True), ("search",))], + ids=["no-deployment-can-search", "capable-deployment-already-tried"], +) +async def test_pin_holds_for_a_web_search_request_when_no_untried_deployment_can_search( + search_supported: tuple[bool, bool], tried: tuple[str, ...] +): + cache: Final = AsyncMock() + cache.async_get_cache = AsyncMock(return_value={"model_id": "no-search"}) + callback: Final = DeploymentAffinityCheck( + cache=cache, ttl_seconds=60, enable_user_key_affinity=True, enable_responses_api_affinity=False + ) + healthy_deployments: Final = [ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x"}, + "model_info": {"id": model_id, "supports_web_search": supported}, + } + for model_id, supported in zip(("no-search", "search"), search_supported) + ] + + filtered: Final = await callback.async_filter_deployments( + model="mixed", + healthy_deployments=healthy_deployments, + messages=None, + request_kwargs={ + "metadata": {"user_api_key_hash": "caller-1"}, + "web_search_options": {}, + "_excluded_deployment_ids": list(tried), + }, + ) + + assert [d["model_info"]["id"] for d in filtered] == ["no-search"] + + +@pytest.mark.asyncio +async def test_key_pin_holds_for_a_web_search_request_when_tag_filtering_removes_every_deployment_that_can_search(): + router: Final = litellm.Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + "tags": ["cheap"], + }, + "model_info": {"id": model_id, "supports_web_search": False}, + } + for model_id in ("flagged-a", "flagged-b") + ] + + [ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + "tags": ["premium"], + }, + "model_info": {"id": "search"}, + } + ], + optional_pre_call_checks=["deployment_affinity"], + enable_tag_filtering=True, + ) + await router.cache.async_set_cache( + DeploymentAffinityCheck.get_affinity_cache_key(model_group="mixed", user_key="caller-1"), + {"model_id": "flagged-a"}, + ttl=3600, + ) + + served_by: Final = { + ( + await router.acompletion( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + metadata={"user_api_key_hash": "caller-1", "tags": ["cheap"]}, + web_search_options={}, + ) + )._hidden_params["model_id"] + for _ in range(20) + } + + assert served_by == {"flagged-a"} + + + +def _web_search_group_router( + optional_pre_call_checks: list[str], # mutable-ok: Router signature + heavy_can_search: bool, + heavy_weight: int, + session_affinity_repin_on_failover: bool = False, +) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "group", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + "weight": heavy_weight, + }, + "model_info": {"id": "heavy", **({} if heavy_can_search else {"supports_web_search": False})}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "weight": 1}, + "model_info": {"id": "light"}, + }, + ], + optional_pre_call_checks=optional_pre_call_checks, + session_affinity_repin_on_failover=session_affinity_repin_on_failover, + num_retries=0, + ) + + +_AFFINITY_KINDS: Final = pytest.mark.parametrize( + "optional_pre_call_checks", [["deployment_affinity"], ["session_affinity"]], ids=["user-key-pin", "session-pin"] +) + + +async def _served_by(router: litellm.Router, caller: int, **web_search_request: object) -> object: + response: Final = await router.acompletion( + model="group", + messages=[{"role": "user", "content": "hi"}], + metadata={"user_api_key_hash": f"caller-{caller}", "session_id": f"session-{caller}"}, + **web_search_request, + ) + return response._hidden_params["model_id"] + + +@pytest.mark.asyncio +@_AFFINITY_KINDS +async def test_a_first_web_search_request_does_not_pin_plain_traffic_to_the_deployment_that_can_search( + optional_pre_call_checks: list[str], # mutable-ok: Router signature +): + router: Final = _web_search_group_router(optional_pre_call_checks, heavy_can_search=False, heavy_weight=10**6) + + searched_on: Final = await _served_by(router, 1, web_search_options={}) + plain_on: Final = {await _served_by(router, 1) for _ in range(5)} + + assert searched_on == "light" + assert plain_on == {"heavy"} + + +@pytest.mark.asyncio +@_AFFINITY_KINDS +async def test_a_first_web_search_request_still_pins_a_group_where_every_deployment_can_search( + optional_pre_call_checks: list[str], # mutable-ok: Router signature +): + router: Final = _web_search_group_router(optional_pre_call_checks, heavy_can_search=True, heavy_weight=1) + + for caller in range(10): + searched_on = await _served_by(router, caller, web_search_options={}) + assert {await _served_by(router, caller) for _ in range(3)} == {searched_on} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("optional_params", "session_pin"), + [ + ({"web_search_options": {}}, None), + ({"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, None), + ({}, {"model_id": "light"}), + ], + ids=["web-search-options", "web-search-tool", "plain"], +) +async def test_a_successful_web_search_request_does_not_repin_a_session_in_a_mixed_group( + optional_params: dict[str, object], # mutable-ok: logging kwargs shape + session_pin: dict[str, str] | None, # mutable-ok: cache value shape +): + router: Final = _web_search_group_router( + ["session_affinity"], heavy_can_search=False, heavy_weight=1, session_affinity_repin_on_failover=True + ) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + await router.deployment_callback_on_success( + { + "litellm_params": { + "model_info": {"id": "light"}, + "metadata": { + "model_group": "group", + "deployment_model_name": "group", + "session_id": "session-1", + "user_api_key_hash": "caller-1", + }, + }, + "optional_params": optional_params, + "standard_logging_object": {"total_tokens": 0}, + }, + {"choices": [{"message": {"content": "ok"}}]}, + timestamp, + timestamp, + ) + + assert ( + await router.cache.async_get_cache( + DeploymentAffinityCheck.get_session_affinity_cache_key("group", "session-1", "caller-1") + ) + == session_pin + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pinned_model_id", "session_pin"), + [("removed", {"model_id": "light"}), ("heavy", {"model_id": "heavy"})], + ids=["ineligible-pin-moves", "eligible-pin-stays"], +) +async def test_a_successful_web_search_request_moves_a_session_only_off_an_ineligible_pin_in_a_mixed_group( + pinned_model_id: str, + session_pin: dict[str, str], # mutable-ok: cache value shape +): + router: Final = _web_search_group_router( + ["session_affinity"], heavy_can_search=False, heavy_weight=1, session_affinity_repin_on_failover=True + ) + cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "session-1", "caller-1") + await router.cache.async_set_cache(cache_key, {"model_id": pinned_model_id}) + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + + await router.deployment_callback_on_success( + { + "litellm_params": { + "model_info": {"id": "light"}, + "metadata": { + "model_group": "group", + "deployment_model_name": "group", + "session_id": "session-1", + "user_api_key_hash": "caller-1", + }, + }, + "optional_params": {"web_search_options": {}}, + "standard_logging_object": {"total_tokens": 0}, + }, + {"choices": [{"message": {"content": "ok"}}]}, + timestamp, + timestamp, + ) + + assert await router.cache.async_get_cache(cache_key) == session_pin diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 6ec72868b6e..c73e354a8e2 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -1,5 +1,9 @@ +import asyncio import logging +from collections.abc import Mapping +from datetime import datetime, timezone from itertools import chain +from types import MappingProxyType from typing import Dict, Final, List, Optional, Union from unittest.mock import Mock @@ -7,6 +11,7 @@ import pytest import litellm from litellm import Router +from litellm.caching.dual_cache import DualCache from litellm.constants import ROUTER_FALLBACK_ERROR_DETAIL_MAX_CHARS from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.proxy._types import UserAPIKeyAuth @@ -15,6 +20,8 @@ from litellm.litellm_core_utils.routing_common_utils import ( add_model_file_id_mappings, filter_team_based_models, filter_web_search_deployments, + pick_preferring_web_search, + prefer_web_search_deployments, format_content_policy_no_fallback_note, format_fallback_outcome_message, provider_for_generic_call, @@ -24,6 +31,7 @@ from litellm.litellm_core_utils.routing_common_utils import ( PROVIDER_SCOPED_CREDENTIAL_PARAMS, warn_on_provider_credential_mismatch, ) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.types.router import Deployment, FallbackSkipReason, LiteLLM_Params, ModelInfo, SkippedFallbackTarget from litellm.types.utils import ProviderSpecificModelInfo @@ -238,6 +246,11 @@ class TestDeploymentSupportsWebSearch: deployment = {"litellm_params": {"model": "gpt-4"}} assert _deployment_supports_web_search(deployment) is True + @pytest.mark.parametrize("supports_web_search", [True, False]) + def test_model_info_object_setting_is_read(self, supports_web_search: bool) -> None: + deployment: Final = {"model_info": ModelInfo(id="1", supports_web_search=supports_web_search)} + assert _deployment_supports_web_search(deployment) is supports_web_search + class TestFilterWebSearchDeployments: """Test cases for filter_web_search_deployments function""" @@ -309,6 +322,44 @@ class TestFilterWebSearchDeployments: result_ids = [d["model_info"]["id"] for d in result] assert "deployment-3" not in result_ids + @pytest.mark.parametrize( + "tool", + [ + {"type": "web_search_20250305", "name": "web_search"}, + {"type": "web_search_20260209", "name": "web_search"}, + {"type": "web_search_preview_2025_03_11"}, + {"googleSearch": {}}, + {"google_search": {}}, + {"googleSearchRetrieval": {}}, + {"google_search_retrieval": {}}, + {"enterpriseWebSearch": {}}, + {"enterprise_web_search": {}}, + {"type": "google_search"}, + ], + ) + def test_native_web_search_tools_filter_unsupported( + self, sample_deployments: list[dict[str, object]], tool: dict[str, object] + ) -> None: + result: Final = prefer_web_search_deployments(sample_deployments, {"tools": [tool]}) + assert result == sample_deployments[:2] + + @pytest.mark.parametrize( + "tool", + [ + {"type": "web_fetch_20250910", "name": "web_fetch"}, + {"type": "function", "function": {"name": "web_search"}}, + {"type": "function", "name": "googleSearch"}, + {"urlContext": {}}, + {"googleSearch": None}, + {"google_search": None, "urlContext": {}}, + ], + ) + def test_non_web_search_tools_preserve_deployments( + self, sample_deployments: list[dict[str, object]], tool: dict[str, object] + ) -> None: + result: Final = prefer_web_search_deployments(sample_deployments, {"tools": [tool]}) + assert result == sample_deployments + def test_web_search_with_other_tools(self, sample_deployments): """Web search filtering works when mixed with other tools""" request_kwargs = { @@ -332,15 +383,44 @@ class TestFilterWebSearchDeployments: result = filter_web_search_deployments(deployments, request_kwargs) assert len(result) == 2 - def test_no_deployments_support_web_search(self): - """When no deployments support web search, all are filtered out""" - deployments = [ + @pytest.mark.parametrize("tool_type", ["web_search", "web_search_preview"]) + def test_no_deployment_supporting_openai_web_search_leaves_none(self, tool_type: str) -> None: + deployments: Final = [ {"model_info": {"id": "d1", "supports_web_search": False}}, {"model_info": {"id": "d2", "supports_web_search": False}}, ] - request_kwargs = {"tools": [{"type": "web_search"}]} - result = filter_web_search_deployments(deployments, request_kwargs) - assert len(result) == 0 + assert filter_web_search_deployments(deployments, {"tools": [{"type": tool_type}]}) == [] + + @pytest.mark.parametrize( + "request_kwargs", + [ + {"tools": [{"type": "web_search_20250305", "name": "web_search"}]}, + {"tools": [{"type": "web_search_2025_08_26"}]}, + {"tools": [{"googleSearch": {}}]}, + {"tools": [{"type": "web_search"}]}, + {"web_search_options": {}}, + ], + ) + def test_no_deployment_supporting_web_search_keeps_every_deployment(self, request_kwargs: dict[str, object]) -> None: + deployments: Final[list[dict[str, object]]] = [ + {"model_info": {"id": "d1", "supports_web_search": False}}, + {"model_info": {"id": "d2", "supports_web_search": False}}, + ] + assert prefer_web_search_deployments(deployments, request_kwargs) == deployments + + def test_a_tool_that_is_not_an_object_is_no_web_search_request( + self, sample_deployments: list[dict[str, object]] + ) -> None: + result: Final = prefer_web_search_deployments(sample_deployments, {"tools": ["web_search", None]}) + assert result == sample_deployments + + def test_web_search_options_filter_unsupported(self, sample_deployments: list[dict[str, object]]) -> None: + result: Final = prefer_web_search_deployments(sample_deployments, {"web_search_options": {}}) + assert result == sample_deployments[:2] + + def test_null_web_search_options_preserve_deployments(self, sample_deployments: list[dict[str, object]]) -> None: + result: Final = prefer_web_search_deployments(sample_deployments, {"web_search_options": None}) + assert result == sample_deployments def test_missing_config_defaults_to_supported(self): """Deployments without supports_web_search config default to True""" @@ -375,6 +455,758 @@ class TestFilterWebSearchDeployments: assert result == deployment +def _router_with_a_deployment_that_disables_web_search(routing_strategy: str) -> Router: + return Router( + model_list=[ + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key"}, + "model_info": {"id": "search"}, + }, + ], + routing_strategy=routing_strategy, + ) + + +_ANTHROPIC_WEB_SEARCH_REQUEST: Final = {"tools": [{"type": "web_search_20250305", "name": "web_search"}]} + + +@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing"]) +def test_sync_pick_skips_deployments_without_web_search(routing_strategy: str) -> None: + router: Final = _router_with_a_deployment_that_disables_web_search(routing_strategy) + picks: Final = { + router.get_available_deployment(model="search-model", request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST))[ + "model_info" + ]["id"] + for _ in range(20) + } + assert picks == {"search"} + + +def _router_with_a_group_that_disables_web_search(routing_strategy: str = "simple-shuffle") -> Router: + return Router( + model_list=[ + { + "model_name": "flagged-model", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "flagged"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "search-model", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "searched"}, + "model_info": {"id": "search"}, + }, + ], + fallbacks=[{"flagged-model": ["search-model"]}], + routing_strategy=routing_strategy, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("tools", "expected_content"), + [ + ([{"type": "web_search"}], "searched"), + ([{"type": "web_search_20250305", "name": "web_search"}], "flagged"), + ([{"googleSearch": dict[str, object]()}], "flagged"), + ], +) +async def test_web_search_to_a_group_without_web_search_routes_like_main( + tools: list[dict[str, object]], expected_content: str +) -> None: + response: Final = await _router_with_a_group_that_disables_web_search().acompletion( + model="flagged-model", messages=[{"role": "user", "content": "hi"}], tools=tools + ) + assert response.choices[0].message.content == expected_content + + +@pytest.mark.parametrize("routing_strategy", ["simple-shuffle", "usage-based-routing"]) +def test_sync_pick_keeps_a_group_without_web_search_for_openai_web_search(routing_strategy: str) -> None: + router: Final = _router_with_a_group_that_disables_web_search(routing_strategy) + deployment: Final = router.get_available_deployment( + model="flagged-model", request_kwargs={"tools": [{"type": "web_search"}]} + ) + assert deployment["model_info"]["id"] == "no-search" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "routing_strategy", ["usage-based-routing", "usage-based-routing-v2", "least-busy", "latency-based-routing"] +) +async def test_async_strategy_pick_skips_deployments_without_web_search(routing_strategy: str) -> None: + router: Final = _router_with_a_deployment_that_disables_web_search(routing_strategy) + picks: Final = { + ( + await router.async_get_available_deployment( + model="search-model", request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST) + ) + )["model_info"]["id"] + for _ in range(20) + } + assert picks == {"search"} + + +def _order_param(order: int | None) -> dict[str, int]: + return {} if order is None else {"order": order} + + +def _router_with_a_mixed_web_search_group( + enable_tag_filtering: bool, orders: tuple[int | None, int | None] = (None, None) +) -> Router: + return Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "tags": ["cheap"], + **_order_param(orders[0]), + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "tags": ["premium"], + **_order_param(orders[1]), + }, + "model_info": {"id": "search"}, + }, + ], + enable_tag_filtering=enable_tag_filtering, + ) + + +def _cool_down(router: Router, model_ids: tuple[str, ...]) -> None: + for model_id in model_ids: + router.cooldown_cache.add_deployment_to_cooldown( + model_id=model_id, original_exception=Exception("rate limited"), exception_status=429, cooldown_time=300 + ) + + +_LATER_FILTER_CASES: Final = pytest.mark.parametrize( + ("orders", "cooled_down", "request_kwargs", "expected_id"), + [ + ((None, None), (), _ANTHROPIC_WEB_SEARCH_REQUEST, "search"), + ((None, None), (), {"web_search_options": {}}, "search"), + ((1, 2), (), _ANTHROPIC_WEB_SEARCH_REQUEST, "no-search"), + ((1, 2), (), {"web_search_options": {}}, "no-search"), + ((None, None), ("search",), _ANTHROPIC_WEB_SEARCH_REQUEST, "no-search"), + ((None, None), ("search",), {"web_search_options": {}}, "no-search"), + ((None, None), (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_excluded_deployment_ids": ["search"]}, "no-search"), + ((None, None), (), {**_ANTHROPIC_WEB_SEARCH_REQUEST, "_retry_skipped_deployment_ids": ("search",)}, "no-search"), + ], + ids=[ + "same-tier", + "same-tier-options", + "lower-order-tier", + "lower-order-tier-options", + "cooldown", + "cooldown-options", + "weighted-failover", + "retry-skip", + ], +) + + +@pytest.mark.asyncio +@_LATER_FILTER_CASES +async def test_async_pick_prefers_web_search_until_later_filters_leave_none_that_can( + orders: tuple[int | None, int | None], + cooled_down: tuple[str, ...], + request_kwargs: dict[str, object], + expected_id: str, +) -> None: + router: Final = _router_with_a_mixed_web_search_group(False, orders) + _cool_down(router, cooled_down) + deployment: Final = await router.async_get_available_deployment( + model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(request_kwargs) + ) + assert deployment["model_info"]["id"] == expected_id + + +@_LATER_FILTER_CASES +def test_sync_pick_prefers_web_search_until_later_filters_leave_none_that_can( + orders: tuple[int | None, int | None], + cooled_down: tuple[str, ...], + request_kwargs: dict[str, object], + expected_id: str, +) -> None: + router: Final = _router_with_a_mixed_web_search_group(False, orders) + _cool_down(router, cooled_down) + deployment: Final = router.get_available_deployment( + model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(request_kwargs) + ) + assert deployment["model_info"]["id"] == expected_id + + +_WEB_SEARCH_REQUESTS: Final = pytest.mark.parametrize( + "web_search_request", + [_ANTHROPIC_WEB_SEARCH_REQUEST, {"web_search_options": {}}], + ids=["web-search-tool", "web-search-options"], +) + + +@pytest.mark.asyncio +@_WEB_SEARCH_REQUESTS +async def test_order_fallback_answers_a_web_search_request_from_the_next_tier_when_the_capable_one_fails( + web_search_request: dict[str, object], +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "ordered", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + "order": 1, + }, + "model_info": {"id": "search"}, + }, + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "order": 2}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=0, + ) + response: Final = await router.acompletion( + model="ordered", messages=[{"role": "user", "content": "hi"}], **web_search_request + ) + assert response._hidden_params["model_id"] == "no-search" + + +@pytest.mark.asyncio +@_WEB_SEARCH_REQUESTS +async def test_web_search_request_prefers_the_capable_deployment_inside_its_order_tier( + web_search_request: dict[str, object], +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "order": 1}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "order": 1}, + "model_info": {"id": "search"}, + }, + { + "model_name": "ordered", + "litellm_params": {"model": "openai/gpt-z", "api_key": "fake-key", "order": 2}, + "model_info": {"id": "later-no-search", "supports_web_search": False}, + }, + ], + ) + picks: Final = { + ( + await router.async_get_available_deployment( + model="ordered", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(web_search_request) + ) + )["model_info"]["id"] + for _ in range(20) + } + assert picks == {"search"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +async def test_web_search_retry_goes_to_an_untried_deployment_once_the_capable_deployment_failed_this_request( + metadata_field: str, +) -> None: + router: Final = _router_with_a_mixed_web_search_group(False) + picks: Final = { + ( + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={ + **_ANTHROPIC_WEB_SEARCH_REQUEST, + metadata_field: {"previous_models": ({"deployment_id": "search", "exception_type": "Timeout"},)}, + }, + ) + )["model_info"]["id"] + for _ in range(40) + } + assert picks == {"no-search"} + + + +@pytest.mark.parametrize( + ("deployments", "failed", "expected"), + [ + ( + ( + {"id": "search"}, + {"id": "flag-a", "supports_web_search": False}, + {"id": "flag-b", "supports_web_search": False}, + ), + ("search", "flag-a"), + ("flag-b",), + ), + ( + ({"id": "search"}, {"id": "flag-a", "supports_web_search": False}), + ("search", "flag-a"), + ("search", "flag-a"), + ), + ( + ({"id": "flag-a", "supports_web_search": False}, {"id": "flag-b", "supports_web_search": False}), + ("flag-a",), + ("flag-a", "flag-b"), + ), + ], + ids=["skips-every-tried-one", "keeps-all-once-all-were-tried", "no-deployment-can-search"], +) +def test_web_search_retry_prefers_untried_deployments_only_after_one_that_can_search_was_tried( + deployments: tuple[dict[str, object], ...], failed: tuple[str, ...], expected: tuple[str, ...] +) -> None: + request_kwargs: Final = { + **_ANTHROPIC_WEB_SEARCH_REQUEST, + "metadata": {"previous_models": tuple({"deployment_id": d, "exception_type": "Timeout"} for d in failed)}, + } + kept: Final = prefer_web_search_deployments([{"model_info": info} for info in deployments], request_kwargs) + assert tuple(d["model_info"]["id"] for d in kept) == expected + + +@pytest.mark.parametrize( + ("failed", "at_limit", "expected"), + [((), ("search",), ("search", "flag")), (("search",), (), ("flag",))], + ids=["at-its-limit-keeps-all", "failed-keeps-the-untried-one"], +) +def test_a_web_search_deployment_only_at_its_limit_leaves_every_deployment_to_the_weights_and_fallbacks( + failed: tuple[str, ...], at_limit: tuple[str, ...], expected: tuple[str, ...] +) -> None: + request_kwargs: Final = { + **_ANTHROPIC_WEB_SEARCH_REQUEST, + "metadata": {"previous_models": tuple({"deployment_id": d, "exception_type": "Timeout"} for d in failed)}, + } + infos: Final = ({"id": "search"}, {"id": "flag", "supports_web_search": False}) + kept: Final = prefer_web_search_deployments( + [{"model_info": info} for info in infos], request_kwargs, frozenset(at_limit) + ) + assert tuple(d["model_info"]["id"] for d in kept) == expected + +@pytest.mark.asyncio +@_WEB_SEARCH_REQUESTS +@pytest.mark.parametrize( + "caller_credentials", [{}, {"api_key": "caller-key"}], ids=["configured-credentials", "caller-credentials"] +) +@pytest.mark.parametrize( + "operator_model_info", [{}, {"original_model_id": "set-by-operator"}], ids=["plain", "operator-original-model-id"] +) +async def test_retry_after_a_retryable_failure_on_the_capable_deployment_reaches_the_other_one( + web_search_request: dict[str, object], + caller_credentials: dict[str, str], + operator_model_info: dict[str, str], +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + "weight": 1, + }, + "model_info": {"id": "search", **operator_model_info}, + }, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "from no-search", + "weight": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=1, + disable_cooldowns=True, + ) + response: Final = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **web_search_request, **caller_credentials + ) + assert response.choices[0].message.content == "from no-search" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("metadata_field", ["metadata", "litellm_metadata"]) +@pytest.mark.parametrize(("request_retry_count", "flagged_reachable"), [(1, False), (2, True)], ids=["kept", "trimmed"]) +async def test_web_search_preference_steps_aside_only_once_the_retry_history_was_trimmed( + metadata_field: str, request_retry_count: int, flagged_reachable: bool +) -> None: + router: Final = Router( + model_list=[ + *( + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("search-a", "search-b") + ), + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + "weight": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + ) + picks: Final = { + ( + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={ + **_ANTHROPIC_WEB_SEARCH_REQUEST, + metadata_field: { + "previous_models": ({"deployment_id": "search-a", "exception_type": "Timeout"},), + "request_retry_count": request_retry_count, + }, + }, + ) + )["model_info"]["id"] + for _ in range(20) + } + assert ("no-search" in picks) is flagged_reachable + + +@pytest.mark.asyncio +async def test_retries_past_the_kept_retry_history_reach_the_deployment_that_cannot_search() -> None: + capable_deployments: Final = [ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + }, + "model_info": {"id": f"search-{index}"}, + } + for index in range(5) + ] + router: Final = Router( + model_list=[ + *capable_deployments, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "from no-search", + "weight": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=5, + disable_cooldowns=True, + ) + response: Final = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + assert response.choices[0].message.content == "from no-search" + + +@pytest.mark.asyncio +async def test_retry_after_a_capable_deployment_fails_goes_to_another_capable_one_before_any_other() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": Exception("connection error"), + "weight": 10**6, + }, + "model_info": {"id": "dead-search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "from live-search"}, + "model_info": {"id": "live-search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "from no-search"}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + num_retries=1, + disable_cooldowns=True, + ) + replies: Final = tuple( + [ + ( + await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + ) + .choices[0] + .message.content + for _ in range(5) + ] + ) + assert replies == ("from live-search",) * 5 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("web_search_options", [False, "", 0, []], ids=["false", "empty-string", "zero", "list"]) +async def test_web_search_options_that_is_not_an_object_routes_like_a_plain_request( + web_search_options: object, +) -> None: + router: Final = _router_with_a_mixed_web_search_group(False) + picks: Final = { + ( + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"web_search_options": web_search_options}, + ) + )["model_info"]["id"] + for _ in range(40) + } + assert picks == {"search", "no-search"} + + +_SEARCH_RPM_OF_ONE: Final = ( + pytest.param({"litellm_params": {"rpm": 1}}, id="int-rpm-in-litellm-params"), + pytest.param({"model_info": {"rpm": 1.0}}, id="float-rpm-in-model-info"), + pytest.param({"model_info": {"rpm": True}}, id="bool-rpm-in-model-info"), +) + + +_ONE_MINUTE: Final = datetime(2026, 10, 28, 12, 0, 30, tzinfo=timezone.utc) + + +def _web_search_group_with_model_rate_limits( + search_rpm: Mapping[str, Mapping[str, object]] = MappingProxyType({"litellm_params": {"rpm": 1}}), +) -> Router: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-x", + "api_key": "fake-key", + "mock_response": "hi", + **search_rpm.get("litellm_params", {}), + }, + "model_info": {"id": "search", **search_rpm.get("model_info", {})}, + }, + { + "model_name": "mixed", + "litellm_params": { + "model": "openai/gpt-y", + "api_key": "fake-key", + "mock_response": "hi", + "rpm": 10**6, + }, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + ) + check: Final = ModelRateLimitingCheck(dual_cache=DualCache(), clock=lambda: _ONE_MINUTE) + router.optional_callbacks = [check] + litellm.logging_callback_manager.add_litellm_callback(check) + return router + + +@pytest.mark.asyncio +@pytest.mark.parametrize("search_rpm", _SEARCH_RPM_OF_ONE) +async def test_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm( + search_rpm: Mapping[str, Mapping[str, object]], +) -> None: + router: Final = _web_search_group_with_model_rate_limits(search_rpm) + served_by: Final = tuple( + [ + ( + await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + )._hidden_params["model_id"] + for _ in range(3) + ] + ) + assert served_by == ("search", "no-search", "no-search") + + +@pytest.mark.parametrize("search_rpm", _SEARCH_RPM_OF_ONE) +def test_sync_web_search_request_goes_to_the_other_deployment_once_the_capable_one_hits_its_enforced_rpm( + search_rpm: Mapping[str, Mapping[str, object]], +) -> None: + router: Final = _web_search_group_with_model_rate_limits(search_rpm) + served_by: Final = tuple( + router.completion( + model="mixed", messages=[{"role": "user", "content": "hi"}], num_retries=0, **_ANTHROPIC_WEB_SEARCH_REQUEST + )._hidden_params["model_id"] + for _ in range(3) + ) + assert served_by == ("search", "no-search", "no-search") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("zero_limit", ["rpm", "itpm", "otpm"]) +async def test_web_search_request_never_goes_to_a_capable_deployment_whose_enforced_limit_is_zero( + zero_limit: str, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", zero_limit: 0}, + "model_info": {"id": "search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "rpm": 1000}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + served_by: Final = tuple( + [ + ( + await router.acompletion( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + num_retries=0, + **_ANTHROPIC_WEB_SEARCH_REQUEST, + ) + )._hidden_params["model_id"] + for _ in range(5) + ] + ) + assert served_by == ("no-search",) * 5 + + +def _rate_limited_web_search_group(routing_strategy: str) -> Router: + return Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key", "mock_response": "hi", "tpm": 1}, + "model_info": {"id": "search"}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key", "mock_response": "hi", "rpm": 1000}, + "model_info": {"id": "no-search", "supports_web_search": False}, + }, + ], + routing_strategy=routing_strategy, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("routing_strategy", ["usage-based-routing-v2", "cost-based-routing"]) +async def test_web_search_request_goes_to_the_other_deployment_when_the_capable_one_has_no_tpm_room( + routing_strategy: str, +) -> None: + router: Final = _rate_limited_web_search_group(routing_strategy) + response: Final = await router.acompletion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + assert response._hidden_params["model_id"] == "no-search" + + +def test_sync_web_search_request_goes_to_the_other_deployment_when_the_capable_one_has_no_tpm_room() -> None: + router: Final = _rate_limited_web_search_group("usage-based-routing-v2") + response: Final = router.completion( + model="mixed", messages=[{"role": "user", "content": "hi"}], **_ANTHROPIC_WEB_SEARCH_REQUEST + ) + assert response._hidden_params["model_id"] == "no-search" + + +def test_sync_pick_falls_back_to_every_deployment_when_the_strategy_finds_none_that_can_search() -> None: + deployments: Final = [ + {"model_info": {"id": "search"}}, + {"model_info": {"id": "no-search", "supports_web_search": False}}, + ] + picked: Final = pick_preferring_web_search( + lambda candidates: next((d for d in candidates if d["model_info"]["id"] == "no-search"), None), + deployments, + _ANTHROPIC_WEB_SEARCH_REQUEST, + ) + assert picked == deployments[1] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("search_order", "no_search_order"), + [([1], [1]), ({"tier": 1}, None)], + ids=["list-order", "dict-order"], +) +async def test_web_search_request_routes_on_a_group_whose_order_is_not_a_number( + search_order: object, no_search_order: object +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-x", "api_key": "fake-key"}, + "model_info": {"id": "search", "order": search_order}, + }, + { + "model_name": "mixed", + "litellm_params": {"model": "openai/gpt-y", "api_key": "fake-key"}, + "model_info": { + "id": "no-search", + "supports_web_search": False, + **({} if no_search_order is None else {"order": no_search_order}), + }, + }, + ], + ) + deployment: Final = await router.async_get_available_deployment( + model="mixed", messages=[{"role": "user", "content": "hi"}], request_kwargs=dict(_ANTHROPIC_WEB_SEARCH_REQUEST) + ) + assert deployment["model_info"]["id"] == "search" + + +@pytest.mark.asyncio +async def test_tag_routing_to_a_deployment_without_web_search_keeps_it() -> None: + deployment: Final = await _router_with_a_mixed_web_search_group(True).async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"web_search_options": {}, "metadata": {"tags": ["cheap"]}}, + ) + assert deployment["model_info"]["id"] == "no-search" + + +@pytest.mark.asyncio +async def test_openai_web_search_still_finds_no_deployment_while_the_only_one_that_can_search_cools_down() -> None: + router: Final = _router_with_a_mixed_web_search_group(False) + _cool_down(router, ("search",)) + with pytest.raises(litellm.RouterRateLimitError): + await router.async_get_available_deployment( + model="mixed", + messages=[{"role": "user", "content": "hi"}], + request_kwargs={"tools": [{"type": "web_search"}]}, + ) + + def test_invalidate_model_group_info_cache(): """Test that _invalidate_model_group_info_cache clears the LRU cache.""" router = Router( diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index ff3ffc0db0c..6c452f193fb 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,6 +6,7 @@ regardless of the routing strategy being used. """ import asyncio +from collections.abc import Callable, Mapping from datetime import datetime, timedelta, timezone from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -17,6 +18,7 @@ import litellm from litellm import Router from litellm.caching.dual_cache import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import io_token_limits_by_cache_key from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) @@ -483,6 +485,21 @@ class Clock: return self.now +class _ReadsThatFinishOnceAllStarted: + def __init__(self, expected: int, usage: int): + self.expected = expected + self.usage = usage + self.started = 0 + self.all_started = asyncio.Event() + + async def async_get_cache(self, key: str, parent_otel_span: object = None) -> int: + self.started += 1 + if self.started == self.expected: + self.all_started.set() + await self.all_started.wait() + return self.usage + + def _pre_call(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> object: return check.pre_call_check(deployment) @@ -491,6 +508,51 @@ async def _async_pre_call(check: ModelRateLimitingCheck, deployment: dict[str, o return await check.async_pre_call_check(deployment) +def _rejects(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> bool: + try: + _pre_call(check, deployment) + except litellm.RateLimitError: + return True + return False + + +async def _async_rejects(check: ModelRateLimitingCheck, deployment: dict[str, object]) -> bool: + try: + await _async_pre_call(check, deployment) + except litellm.RateLimitError: + return True + return False + + +_FRESH_LIMITS: Final = ( + {"rpm": 0}, + {"rpm": 0.0}, + {"rpm": 0.5}, + {"rpm": False}, + {"rpd": 0}, + {"tpm": 0}, + {"tpd": 0}, + {"itpm": 0}, + {"otpm": 0}, + {"rpm": 1, "tpm": 1}, +) + + +def _redis_batch_answering(usage: int) -> Callable[[list[str]], dict[str, int]]: + def batch_get_cache(key_list: list[str]) -> dict[str, int]: + return dict.fromkeys(key_list, usage) + + return batch_get_cache + + +def _deployment_with(limits: Mapping[str, float]) -> dict[str, object]: + return { + "litellm_params": {"model": "gpt-4.1-mini", **limits}, + "model_info": {"id": "limit-test-id"}, + "model_name": "daily-model", + } + + async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, object]) -> None: await check.async_log_success_event(kwargs, None, None, None) @@ -539,6 +601,132 @@ class TestDailyLimits: assert "TPD limit=100, current usage=120" in str(exc_info.value) + @pytest.mark.asyncio + async def test_deployments_at_their_rpm_or_rpd_limit_are_reported_and_ones_below_are_not(self): + minute_spent = { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1}, + "model_info": {"id": "minute-test-id"}, + "model_name": "daily-model", + } + below_limits = { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 2, "rpd": 2}, + "model_info": {"id": "below-test-id"}, + "model_name": "daily-model", + } + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + await _async_pre_call(check, RPD_DEPLOYMENT) + await _async_pre_call(check, RPD_DEPLOYMENT) + await _async_pre_call(check, minute_spent) + await _async_pre_call(check, below_limits) + + assert await check.async_deployment_ids_at_limit( + [RPD_DEPLOYMENT, minute_spent, below_limits], None + ) == frozenset({"daily-test-id", "minute-test-id"}) + + @pytest.mark.asyncio + async def test_deployment_limit_reads_wait_on_redis_together_rather_than_one_by_one(self): + deployments = [ + { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 5}, + "model_info": {"id": f"concurrent-{index}"}, + "model_name": "daily-model", + } + for index in range(3) + ] + reads = _ReadsThatFinishOnceAllStarted(expected=len(deployments), usage=5) + redis_cache = MagicMock() + redis_cache.async_get_cache = reads.async_get_cache + check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + + at_limit = await asyncio.wait_for(check.async_deployment_ids_at_limit(deployments, None), timeout=5) + + assert at_limit == frozenset({"concurrent-0", "concurrent-1", "concurrent-2"}) + + def test_sync_deployment_limit_reads_reach_redis_in_one_round_trip(self): + deployments: list[dict[str, object]] = [ + { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 5, "rpd": 50}, + "model_info": {"id": f"batched-{index}"}, + "model_name": "daily-model", + } + for index in range(3) + ] + redis_cache = MagicMock() + redis_cache.batch_get_cache.side_effect = _redis_batch_answering(5) + check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + + at_limit = check.deployment_ids_at_limit(deployments) + + assert at_limit == frozenset({"batched-0", "batched-1", "batched-2"}) + assert redis_cache.batch_get_cache.call_count == 1 + assert len(redis_cache.batch_get_cache.call_args.kwargs["key_list"]) == 6 + + @pytest.mark.parametrize( + "redis_batch_get", + [_redis_batch_answering(0), RedisCircuitBreakerOpenError()], + ids=["redis-behind-this-worker", "redis-circuit-open"], + ) + def test_sync_a_deployment_this_worker_counted_to_its_limit_stays_at_its_limit( + self, redis_batch_get: Callable[[list[str]], dict[str, int]] | Exception + ): + spent: dict[str, object] = { + "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1}, + "model_info": {"id": "spent-test-id"}, + "model_name": "daily-model", + } + redis_cache = MagicMock() + redis_cache.batch_get_cache.side_effect = redis_batch_get + dual_cache = DualCache(redis_cache=redis_cache) + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) + _, rpm_key = check._get_cache_keys(spent, DAY_ONE.strftime("%H-%M")) + dual_cache.set_cache(key=rpm_key, value=1, local_only=True) + + assert check.deployment_ids_at_limit([spent, _deployment_with({"rpd": 2})]) == frozenset({"spent-test-id"}) + assert rpm_key not in redis_cache.batch_get_cache.call_args.kwargs["key_list"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("limit_kind", ["itpm", "otpm"]) + @pytest.mark.parametrize(("used", "reported"), [(99, False), (100, True)], ids=["room-left", "used-up"]) + async def test_a_deployment_is_reported_at_its_token_limit_once_no_token_is_left( + self, limit_kind: str, used: int, reported: bool + ): + deployment = _deployment_with({limit_kind: 100}) + dual_cache = DualCache() + check = ModelRateLimitingCheck(dual_cache=dual_cache, clock=Clock(DAY_ONE)) + token_key = next( + key + for key, limit in io_token_limits_by_cache_key(deployment, DAY_ONE.strftime("%H-%M")) + if limit is not None + ) + dual_cache.set_cache(key=token_key, value=used, local_only=True) + + assert ("limit-test-id" in check.deployment_ids_at_limit([deployment])) is reported + assert ("limit-test-id" in await check.async_deployment_ids_at_limit([deployment], None)) is reported + + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(f"{kind}={value!r}" for kind, value in limits.items())) + def test_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( + self, limits: Mapping[str, float] + ): + deployment = _deployment_with(limits) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + reported = "limit-test-id" in check.deployment_ids_at_limit([deployment]) + + assert reported == _rejects(check, deployment) + + @pytest.mark.asyncio + @pytest.mark.parametrize("limits", _FRESH_LIMITS, ids=lambda limits: ",".join(f"{kind}={value!r}" for kind, value in limits.items())) + async def test_async_a_deployment_is_reported_at_its_limit_exactly_when_its_next_request_is_rejected( + self, limits: Mapping[str, float] + ): + deployment = _deployment_with(limits) + check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + + reported = "limit-test-id" in await check.async_deployment_ids_at_limit([deployment], None) + + assert reported == await _async_rejects(check, deployment) + def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock)