scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

Merge 256e87789e into 2a9fad611f

Browse source
This commit is contained in:
tessellark 2026-11-04 01:31:59 +00:00 • committed by GitHub
commit 194c142b15
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 2063 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

203
litellm/litellm_core_utils/routing_common_utils.py
View file

@ -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`

118
litellm/router_python_core.py
View file

@ -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:

86
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

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

10
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py
View file

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

119
litellm/router_utils/pre_call_checks/model_rate_limit_check.py
View file

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

3
litellm/types/router.py
View file

@ -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 = [

2
tests/code_coverage_tests/router_code_coverage.py
View file

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

554
tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py
View file

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

844
tests/unit/router_utils/test_router_utils_common_utils.py
View file

@ -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(

188
tests/unit/test_router/test_enforce_model_rate_limits.py
View file

@ -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)