scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 175

feat(dynamic_rate_limiter_v3): steer flexible teams to a permitted alternate when their preferred model is saturated (#51350)

Browse source 
* feat(router): steer flexible teams to a permitted alternate model when their preferred model is saturated

Teams opt in with metadata.model_preferences: [{"model": "luna", "alternates": ["opus"]}].
The dynamic rate limiter v3 pre-call hook reads the saturation cache for the preferred model
and, once it is at or above saturation_threshold, rewrites the request to the first alternate
the team may call that is below the threshold. The response carries
x-litellm-model-steered-from with the preferred model. Rate limits and spend are charged to
the model that served. Teams without model_preferences are untouched.

* fix(dynamic_rate_limiter_v3): type the hook request dicts so the steering rewrite passes the lint gates

* fix(dynamic_rate_limiter_v3): authorize alternates against the key and its model budget, send the steering header on streams

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(dynamic_rate_limiter_v3): enforce member and project allowlists on alternates, skip the steering header for non Latin-1 names

* fix(dynamic_rate_limiter_v3): check every budget layer for an alternate, fall back to the preferred model when the alternate refuses, merge the served model's guardrails

* fix(dynamic_rate_limiter_v3): never steer agent keys, type the steering unit tests

* fix(dynamic_rate_limiter_v3): hold alternates to customer grants, access group budgets, guardrails and the zero-cost waiver

An alternate now passes every check a budget fallback passes. Customer model grants are enforced, the access groups that grant the alternate are budget checked and stamped on the token and request metadata so spend charges them, a zero-cost preferred model that waived the aggregate budget checks is never steered to a priced alternate, and an alternate whose deployments add model-level guardrails the request did not run is skipped.

* fix(dynamic_rate_limiter_v3): stamp the alternate's groups on every bucket and skip alternates the request cannot safely take

The served model's access groups now land on every metadata bucket and the auth copy inside it even when the preferred model had none, so group spend is charged. An alternate is skipped when a budget reservation already holds a different group, when it does not serve the request's endpoint, when unpriced models are blocked and it has no price, or when its guardrails differ from the request's by name or settings. A failing alternate check or alternate admission leaves the request on the preferred model instead of failing it.

* fix(dynamic_rate_limiter_v3): keep a reserved budget off pricier alternates

While the request holds a budget reservation sized for the preferred model, an alternate that costs more per token is skipped, so fail-closed budget enforcement never admits a call the reservation does not cover.

* fix(dynamic_rate_limiter_v3): an unpriced alternate never takes a reserved budget

While a budget reservation is held, an alternate with no input or output price is skipped rather than treated as free.

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu 2026-11-12 01:10:23 -08:00 • committed by GitHub
parent 1677c2250e
commit 52d1c4cd80
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 872 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

472
litellm/proxy/hooks/dynamic_rate_limiter_v3.py
View file

@ -2,14 +2,15 @@
Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting
"""
import json
import os
from collections.abc import Callable, Mapping
from collections.abc import Awaitable, Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from typing import TYPE_CHECKING, Final, Literal
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
import litellm
from litellm import Router
@ -22,6 +23,7 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import (
ensure_response_additional_headers,
response_has_hidden_params,
)
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import (
ProxyRateLimitError,
@ -40,12 +42,13 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l
borrowing_team_model_limits,
claim_request_stash_for_data,
get_or_create_request_stash,
get_request_stash_for_call,
)
from litellm.proxy.hooks.rate_limiter_utils import (
convert_priority_to_percent,
resolve_llm_provider_for_rate_limit,
)
from litellm.proxy.utils import InternalUsageCache
from litellm.proxy.utils import InternalUsageCache, check_and_merge_model_level_guardrails
from litellm.redis.stores.rate_limits import RATE_LIMITS_TARGET, CounterIncrement, RateLimitCounter, RateLimitScope
from litellm.types.router import ModelGroupInfo
from litellm.types.utils import CallTypesLiteral, LLMResponseTypes, Usage
@ -54,6 +57,270 @@ if TYPE_CHECKING:
from litellm.types.utils import PriorityReservationSettings
STEERED_FROM_HEADER: Final = "x-litellm-model-steered-from"
@dataclass(frozen=True, slots=True)
class AlternateGrant:
"""What an authorized alternate carries onto the request: the budgeted access groups that granted it."""
access_groups: tuple[str, ...] = ()
AlternateAuthorizer: TypeAlias = Callable[[UserAPIKeyAuth, str], Awaitable[AlternateGrant | None]]
class ModelPreference(BaseModel):
model_config = ConfigDict(frozen=True, extra="ignore")
model: str
alternates: tuple[str, ...] = ()
MODEL_PREFERENCES: Final = TypeAdapter(tuple[ModelPreference, ...])
@dataclass(frozen=True, slots=True)
class SteeredModel:
model: str
group_info: ModelGroupInfo
access_groups: tuple[str, ...] = ()
async def alternate_grant_for(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None:
"""The grant a budget fallback to ``model`` would get: the key, team, member, project and customer allowlists,
every model budget layer and the budgets of the access groups that authorize it, else None."""
from litellm.exceptions import BudgetExceededError
from litellm.proxy.auth.auth_checks import (
_model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract
can_customer_access_model,
collect_matched_model_access_groups,
get_end_user_object,
get_project_object,
get_team_membership,
get_team_object,
)
from litellm.proxy.auth.user_api_key_auth import (
_is_fallback_model_authorized, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract
_is_within_every_model_budget, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract
)
from litellm.proxy.proxy_server import (
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
llm_router,
model_max_budget_limiter,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
user_id: Final = user_api_key_dict.user_id
team_id: Final = user_api_key_dict.team_id
project_id: Final = user_api_key_dict.project_id
end_user_id: Final = user_api_key_dict.end_user_id
try:
team_object: Final = (
await get_team_object(
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if team_id is not None
else None
)
team_membership: Final = (
await get_team_membership(
user_id=user_id,
team_id=team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if user_id is not None and team_id is not None
else None
)
project_object: Final = (
await get_project_object(
project_id=project_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if project_id is not None
else None
)
end_user_object: Final = (
await get_end_user_object(
end_user_id=end_user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
if end_user_id is not None
else None
)
if end_user_object is not None and end_user_object.models:
can_customer_access_model(
model=model, end_user_object=end_user_object, llm_router=llm_router, valid_token=user_api_key_dict
)
except Exception: # noqa: BLE001 # a failed team, member, project or customer lookup must not authorize an alternate
return None
if not await _is_fallback_model_authorized(
model=model,
valid_token=user_api_key_dict,
llm_model_list=llm_model_list, # pyright: ignore[reportUnknownArgumentType] # proxy global has untyped deployment entries
llm_router=llm_router,
prisma_client=prisma_client,
team_membership=team_membership,
project_object=project_object,
refuse_agent_keys=True,
team_object=team_object,
):
return None
if not await _is_within_every_model_budget(
valid_token=user_api_key_dict, model_max_budget_limiter=model_max_budget_limiter, model=model
):
return None
try:
access_groups: Final = await collect_matched_model_access_groups(
model=model,
valid_token=user_api_key_dict,
team_object=team_object,
project_object=project_object,
llm_router=llm_router,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
team_membership=team_membership,
team_membership_loaded=True,
)
await _model_access_group_max_budget_check(
matched_model_access_groups=access_groups,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
except BudgetExceededError:
return None
except Exception: # noqa: BLE001 # an unreadable group budget must not authorize an alternate
return None
return AlternateGrant(access_groups=access_groups)
REQUEST_BUCKET: Final = TypeAdapter(dict[str, object])
REQUEST_ENTRIES: Final = TypeAdapter(tuple[object, ...])
RESERVATION_ENTRIES: Final = TypeAdapter(tuple[Mapping[str, object], ...])
MODEL_ACCESS_GROUP_ENTITY: Final = "Model access group"
def _metadata_bucket(data: Mapping[str, object], key: str) -> Mapping[str, object] | None:
"""The request's ``key`` metadata bucket as a typed dict, or None when it is absent or not a mapping."""
value: Final = data.get(key)
if not isinstance(value, Mapping):
return None
try:
return REQUEST_BUCKET.validate_python(value)
except ValidationError:
return None
def _guardrail_entries(data: Mapping[str, object]) -> frozenset[str]:
"""The request's guardrail entries, settings included, in a form two lists can be compared by."""
metadata: Final = _metadata_bucket(data, "metadata")
guardrails: Final = metadata.get("guardrails") if metadata is not None else None
entries: Final = (
REQUEST_ENTRIES.validate_python(guardrails)
if isinstance(guardrails, list)
else (guardrails,)
if guardrails
else ()
)
return frozenset(
entry if isinstance(entry, str) else json.dumps(entry, sort_keys=True, default=str) for entry in entries
)
def _auth_copy_with_groups(bucket: Mapping[str, object], access_groups: tuple[str, ...]) -> dict[str, object]:
auth: Final = bucket.get("user_api_key_auth")
if isinstance(auth, UserAPIKeyAuth):
return {"user_api_key_auth": auth.model_copy(update={"matched_model_access_groups": list(access_groups)})}
if isinstance(auth, Mapping):
return {
"user_api_key_auth": {
**REQUEST_BUCKET.validate_python(auth),
"matched_model_access_groups": list(access_groups),
}
}
return {}
def _with_access_groups(data: Mapping[str, object], access_groups: tuple[str, ...]) -> dict[str, object]:
"""``data`` with the served model's access groups stamped on every metadata bucket, the auth copy in it
included, so spend attribution charges the groups that granted the served model and nothing else."""
stamped: Final = {
key: {
**bucket,
MODEL_ACCESS_GROUP_METADATA_KEY: list(access_groups),
**_auth_copy_with_groups(bucket, access_groups),
}
for key in ("metadata", "litellm_metadata")
if (bucket := _metadata_bucket(data, key)) is not None
}
return {**data, **stamped}
def _priced_within(preferred: ModelGroupInfo, alternate: ModelGroupInfo) -> bool:
"""Whether ``alternate`` costs no more per token than ``preferred``, so a budget reserved for the preferred
model's estimate still covers the alternate."""
if alternate.input_cost_per_token is None or alternate.output_cost_per_token is None:
return False
return alternate.input_cost_per_token <= (preferred.input_cost_per_token or 0.0) and (
alternate.output_cost_per_token <= (preferred.output_cost_per_token or 0.0)
)
def _reservation_covers(
user_api_key_dict: UserAPIKeyAuth,
access_groups: tuple[str, ...],
preferred: ModelGroupInfo,
alternate: ModelGroupInfo,
) -> bool:
"""Whether the budget the request already reserved still fits the alternate: the alternate costs no more per
token than the preferred model the estimate was made for, and every model access group reserved against also
grants the alternate, so settlement never charges a group that did not serve."""
reservation: Final = user_api_key_dict.budget_reservation
if reservation is None:
return True
if not _priced_within(preferred, alternate):
return False
entries: Final = reservation.get("entries")
if not isinstance(entries, list):
return True
try:
reserved: Final = frozenset(
entry.get("entity_id")
for entry in RESERVATION_ENTRIES.validate_python(entries)
if entry.get("entity_type") == MODEL_ACCESS_GROUP_ENTITY
)
except ValidationError:
return False
return reserved <= frozenset(access_groups)
def alternates_for(user_api_key_dict: UserAPIKeyAuth, model: str) -> tuple[str, ...]:
"""The team's opt-in alternates for ``model`` from ``team_metadata.model_preferences``, or none."""
raw: Final = (user_api_key_dict.team_metadata or {}).get("model_preferences")
if raw is None:
return ()
try:
preferences: Final = MODEL_PREFERENCES.validate_python(raw)
except ValidationError as error:
verbose_proxy_logger.warning(
"Ignoring malformed team model_preferences for team %s: %s", user_api_key_dict.team_id, error
)
return ()
return next((preference.alternates for preference in preferences if preference.model == model), ())
class _SettlementMetadata(BaseModel):
"""The two fields of a request's standard logging metadata that token settlement reads."""
@ -155,8 +422,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
self,
internal_usage_cache: DualCache,
time_provider: Callable[[], datetime] | None = None,
is_model_authorized: AlternateAuthorizer = alternate_grant_for,
):
self.internal_usage_cache = InternalUsageCache(dual_cache=internal_usage_cache)
self._is_model_authorized = is_model_authorized
self.v3_limiter = PROXY_MaxParallelRequestsHandler_v3(self.internal_usage_cache, time_provider=time_provider)
def update_variables(self, llm_router: Router) -> None:
@ -781,7 +1050,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
data: Mapping[str, object],
call_type: CallTypesLiteral,
) -> Exception | str | dict | None:
"""
@ -823,38 +1092,20 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
claim_request_stash_for_data(data)
model: Final = data["model"]
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
if not isinstance(model, str):
return None
# Get model configuration
model_group_info: Final[ModelGroupInfo | None] = self.llm_router.registry().group(model)
if model_group_info is None:
requested_group_info: Final[ModelGroupInfo | None] = self.llm_router.registry().group(model)
if requested_group_info is None:
verbose_proxy_logger.debug("No model group info for %s, allowing request", model)
return None
steered: Final = await self._steer_to_alternate(model, requested_group_info, user_api_key_dict, data, call_type)
try:
# STEP 1: Check current saturation level
saturation: Final = await self._check_model_saturation(model, model_group_info)
saturation_threshold: Final = _get_priority_settings().saturation_threshold
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={priority}"
admitted: Final = await self._admit_steered_or_preferred(
model, requested_group_info, steered, user_api_key_dict
)
# STEP 2: Check rate limits in THREE phases
# Phase 1: Read-only check of ALL limits (no increments)
# Phase 2: Decide which limits to enforce (based on saturation)
# Phase 3: Increment ALL counters only if request will be allowed
# This prevents partial increments and ensures accurate tracking
await self._check_rate_limits(
model=model,
model_group_info=model_group_info,
user_api_key_dict=user_api_key_dict,
priority=priority,
saturation=saturation,
)
except HTTPException:
raise
except Exception as e:
@ -862,11 +1113,155 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
# Fail open on unexpected errors
return None
if admitted is None:
return None
get_or_create_request_stash().steered_from_model = model
groups: Final = list(admitted.access_groups)
user_api_key_dict.matched_model_access_groups = groups # rebind-ok: spend charges the served model's groups
return _with_access_groups(dict(data, model=admitted.model), admitted.access_groups)
async def _admit_steered_or_preferred(
self,
model: str,
requested_group_info: ModelGroupInfo,
steered: "SteeredModel | None",
user_api_key_dict: UserAPIKeyAuth,
) -> "SteeredModel | None":
"""The alternate once it admits the request, else None after the preferred model's own check."""
if steered is not None:
try:
await self._admit(steered.model, steered.group_info, user_api_key_dict)
except Exception: # noqa: BLE001 # the alternate's refusal or error keeps the request on its own model
verbose_proxy_logger.info(
"Alternate %s refused team %s, staying on %s", steered.model, user_api_key_dict.team_id, model
)
else:
return steered
await self._admit(model, requested_group_info, user_api_key_dict)
return None
async def _admit(self, model: str, model_group_info: ModelGroupInfo, user_api_key_dict: UserAPIKeyAuth) -> None:
"""Read saturation, then the three-phase rate limit check, raising the 429 when ``model`` is out of capacity."""
priority: Final = self._get_priority_from_user_api_key_dict(user_api_key_dict=user_api_key_dict)
saturation: Final = await self._check_model_saturation(model, model_group_info)
saturation_threshold: Final = _get_priority_settings().saturation_threshold
verbose_proxy_logger.debug(
f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, "
f"Threshold={saturation_threshold:.1%}, Priority={priority}"
)
await self._check_rate_limits(
model=model,
model_group_info=model_group_info,
user_api_key_dict=user_api_key_dict,
priority=priority,
saturation=saturation,
)
async def _steer_to_alternate(
self,
model: str,
model_group_info: ModelGroupInfo,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: CallTypesLiteral,
) -> "SteeredModel | None":
"""First team alternate below ``saturation_threshold`` once ``model`` is at or above it, else None. A
broken alternate check is logged and leaves the request on ``model``, as the limiter fails open elsewhere."""
try:
return await self._pick_alternate(model, model_group_info, user_api_key_dict, data, call_type)
except Exception: # noqa: BLE001 # a failing alternate check must not fail a request that can stay put
verbose_proxy_logger.warning(
"Alternate check failed for team %s, staying on %s", user_api_key_dict.team_id, model, exc_info=True
)
return None
async def _pick_alternate(
self,
model: str,
model_group_info: ModelGroupInfo,
user_api_key_dict: UserAPIKeyAuth,
data: Mapping[str, object],
call_type: CallTypesLiteral,
) -> "SteeredModel | None":
alternates: Final = alternates_for(user_api_key_dict, model)
if not alternates:
return None
threshold: Final = _get_priority_settings().saturation_threshold
if await self._check_model_saturation(model, model_group_info) < threshold:
return None
for alternate in alternates:
if alternate == model or not self._steerable(model, alternate, data, call_type, user_api_key_dict.team_id):
continue
if (
candidate := await self._idle_permitted_alternate(
alternate, user_api_key_dict, threshold, model_group_info
)
) is not None:
verbose_proxy_logger.info("Steering team %s from %s to %s", user_api_key_dict.team_id, model, alternate)
return candidate
return None
def _steerable(
self, model: str, alternate: str, data: Mapping[str, object], call_type: CallTypesLiteral, team_id: str | None
) -> bool:
"""Whether ``alternate`` may serve a request that already passed auth and the guardrail walk as ``model``: a
zero-cost ``model`` waived the aggregate budget checks, so only another zero-cost alternate may take it, an
unpriced alternate is blocked when unpriced models are, the alternate has to serve this endpoint, and an
alternate with a model-level guardrail the request has not run would reach it unchecked."""
from litellm.proxy.auth.auth_checks import (
_unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # steering shares the fallback contract
is_model_cost_zero,
)
from litellm.proxy.route_llm_request import endpoint_unsupported_by_model
if is_model_cost_zero(model=model, llm_router=self.llm_router) and not is_model_cost_zero(
model=alternate, llm_router=self.llm_router
):
verbose_proxy_logger.info("Not steering zero-cost %s to priced %s", model, alternate)
return False
if litellm.block_requests_for_models_without_pricing and _unpriced_models_in_request(
model=alternate, llm_router=self.llm_router
):
verbose_proxy_logger.info("Not steering %s to unpriced %s", model, alternate)
return False
if endpoint_unsupported_by_model(call_type, dict(data, model=alternate), self.llm_router, team_id):
verbose_proxy_logger.info(
"Not steering %s to %s, the alternate does not serve %s", model, alternate, call_type
)
return False
probe: Final = REQUEST_BUCKET.validate_python(
check_and_merge_model_level_guardrails(
data=dict(data, model=alternate, metadata=_metadata_bucket(data, "metadata") or {}),
llm_router=self.llm_router,
trust_client_model_info=False,
)
)
if _guardrail_entries(probe) != _guardrail_entries(data):
verbose_proxy_logger.info(
"Not steering %s to %s, the alternate adds model-level guardrails", model, alternate
)
return False
return True
async def _idle_permitted_alternate(
self, alternate: str, user_api_key_dict: UserAPIKeyAuth, threshold: float, preferred: ModelGroupInfo
) -> "SteeredModel | None":
group_info: Final = self.llm_router.registry().group(alternate)
if group_info is None:
return None
grant: Final = await self._is_model_authorized(user_api_key_dict, alternate)
if grant is None or not _reservation_covers(user_api_key_dict, grant.access_groups, preferred, group_info):
return None
if await self._check_model_saturation(alternate, group_info) >= threshold:
return None
return SteeredModel(model=alternate, group_info=group_info, access_groups=grant.access_groups)
@with_service_target(RATE_LIMITS_TARGET)
async def async_post_call_success_hook(
self, data: dict, user_api_key_dict: UserAPIKeyAuth, response
self,
data: dict[str, object], # mutable-ok: the CustomLogger contract hands the v3 limiter a mutable dict
user_api_key_dict: UserAPIKeyAuth,
response,
) -> LLMResponseTypes:
"""
Post-call hook to add rate limit headers to response.
@ -896,6 +1291,21 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger):
verbose_proxy_logger.exception("Error in dynamic rate limiter v3 post-call hook: %s", e)
return response
async def async_post_call_response_headers_hook(
self,
data: Mapping[str, object],
user_api_key_dict: UserAPIKeyAuth,
response: object,
request_headers: Mapping[str, str] | None = None,
litellm_call_info: Mapping[str, object] | None = None,
) -> dict[str, str] | None: # mutable-ok: the CustomLogger contract returns a dict
"""Names the model the caller asked for when steering served another, on streamed and buffered replies."""
call_id: Final = data.get("litellm_call_id")
stash: Final = get_request_stash_for_call(call_id if isinstance(call_id, str) else None)
if stash is None or stash.steered_from_model is None or not _is_latin1_encodable(stash.steered_from_model):
return None
return {STEERED_FROM_HEADER: stash.steered_from_model}
@with_service_target(RATE_LIMITS_TARGET)
async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time) -> None:
"""

3
litellm/proxy/hooks/parallel_request_limiter_v3.py
View file

@ -602,6 +602,7 @@ class RequestRateLimiterStash:
reserved_model: RateLimitedModel | None = None
reserved_scopes: frozenset[tuple[str, str]] = field(default_factory=frozenset)
fallback_models: frozenset[str] = field(default_factory=frozenset)
steered_from_model: str | None = None
ptu_ceiling: PTUTeamCeiling | None = None
ptu_reserved_tokens: int = 0
ptu_reserved_window_identities: frozenset[ReservationWindowIdentity] = field(default_factory=frozenset)
@ -706,7 +707,7 @@ def get_or_create_request_stash() -> RequestRateLimiterStash:
return stash
def claim_request_stash_for_data(data: dict) -> RequestRateLimiterStash:
def claim_request_stash_for_data(data: Mapping[str, object]) -> RequestRateLimiterStash:
stash: Final = get_or_create_request_stash()
owner_call_id: Final = data.get("litellm_call_id")
if isinstance(owner_call_id, str):

144
tests/integration/routing/test_team_model_preference_steering.py Normal file
View file

@ -0,0 +1,144 @@
import json
import uuid
from collections.abc import Iterator
from pathlib import Path
from typing import Final
import httpx
import pytest
import yaml
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
from integration._support.process import owned_proxy
from integration._support.wire import Reply, Request, wire_server
OPENAI_MODEL: Final = "gpt-5.6"
PREFERRED_RPM: Final = 2
ALTERNATE_RPM: Final = 1000
SATURATION_THRESHOLD: Final = 0.5
STEERED_FROM_HEADER: Final = "x-litellm-model-steered-from"
def _reply() -> bytes:
return json.dumps(
{
"id": f"chatcmpl_{uuid.uuid4().hex}",
"object": "chat.completion",
"created": 1700000000,
"model": OPENAI_MODEL,
"choices": [{"index": 0, "message": {"role": "assistant", "content": "steered"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30},
}
).encode()
@pytest.fixture(scope="module")
def steering_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]:
directory: Final = tmp_path_factory.mktemp("model-preference-steering")
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
configuration: Final = {
**base,
"litellm_settings": {
**base["litellm_settings"],
"callbacks": ["dynamic_rate_limiter_v3"],
"priority_reservation_settings": {
"saturation_threshold": SATURATION_THRESHOLD,
"saturation_check_cache_ttl": 1,
},
},
}
path: Final = directory / "model_preference_steering.yaml"
path.write_text(yaml.safe_dump(configuration))
def respond(request: Request) -> Reply:
assert request.method == "POST" and request.target == "/v1/chat/completions", request.target
return Reply(body=_reply())
with (
gateway_from_environment() as rig,
wire_server(respond) as wire,
owned_proxy(rig, directory, {}, config=path) as owned,
):
yield owned, wire.url
def _deployment(owned: Gateway, scenario: Scenario, upstream: str, rpm: int) -> tuple[str, str]:
name: Final = scenario.model(
model=f"openai/{OPENAI_MODEL}", api_base=f"{upstream}/v1", api_key="synthetic-key", rpm=rpm
)
rows: Final = owned.rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_name = %s', (name,))
return name, str(rows[0]["model_id"])
def _chat(owned: Gateway, key: str, model: str) -> httpx.Response:
return owned.request(
"POST",
"/v1/chat/completions",
{"model": model, "messages": [{"role": "user", "content": f"model preference steering {uuid.uuid4().hex}"}]},
key=key,
)
def _spend_row(owned: Gateway, response: httpx.Response) -> list[dict[str, object]]:
return owned.rows(
'SELECT model_group FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (str(response.json()["id"]),)
)
def test_flexible_team_is_steered_to_its_permitted_alternate_once_the_preferred_model_saturates(
steering_proxy: tuple[Gateway, str],
) -> None:
owned, upstream = steering_proxy
with owned.scenario() as scenario:
preferred, preferred_id = _deployment(owned, scenario, upstream, PREFERRED_RPM)
alternate, alternate_id = _deployment(owned, scenario, upstream, ALTERNATE_RPM)
rigid_key: Final = scenario.key(team_id=scenario.team(models=[preferred, alternate]))
flexible_key: Final = scenario.key(
team_id=scenario.team(
models=[preferred, alternate],
metadata={"model_preferences": [{"model": preferred, "alternates": [alternate]}]},
)
)
warmup: Final = _chat(owned, rigid_key, preferred)
assert warmup.status_code == 200, warmup.text
assert warmup.headers["x-litellm-model-id"] == preferred_id
assert STEERED_FROM_HEADER not in warmup.headers
steered: Final = eventually(
lambda: _chat(owned, flexible_key, preferred),
lambda response: STEERED_FROM_HEADER in response.headers,
seconds=15,
return_last_on_timeout=True,
)
assert steered.status_code == 200, steered.text
assert steered.headers[STEERED_FROM_HEADER] == preferred
assert steered.headers["x-litellm-model-id"] == alternate_id
rigid: Final = _chat(owned, rigid_key, preferred)
assert rigid.status_code == 200, rigid.text
assert rigid.headers["x-litellm-model-id"] == preferred_id
assert STEERED_FROM_HEADER not in rigid.headers
charged: Final = eventually(lambda: _spend_row(owned, steered), lambda rows: len(rows) == 1, seconds=30)
assert charged[0]["model_group"] == alternate
def test_alternate_the_team_may_not_call_is_never_used(steering_proxy: tuple[Gateway, str]) -> None:
owned, upstream = steering_proxy
with owned.scenario() as scenario:
preferred, _ = _deployment(owned, scenario, upstream, PREFERRED_RPM)
forbidden, _ = _deployment(owned, scenario, upstream, ALTERNATE_RPM)
key: Final = scenario.key(
team_id=scenario.team(
models=[preferred],
metadata={"model_preferences": [{"model": preferred, "alternates": [forbidden]}]},
)
)
first: Final = _chat(owned, key, preferred)
assert first.status_code == 200, first.text
second: Final = eventually(
lambda: _chat(owned, key, preferred), lambda response: response.status_code != 200, seconds=15
)
assert second.status_code == 429, second.text
assert STEERED_FROM_HEADER not in second.headers

285
tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py
View file

@ -6,7 +6,9 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of
import asyncio
import time
from collections.abc import Mapping
from datetime import datetime, timedelta
from types import MappingProxyType
from typing import Final
from unittest.mock import AsyncMock, patch
@ -14,9 +16,11 @@ import pytest
import litellm
from litellm import DualCache, Router
from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
AlternateGrant,
PROXY_DynamicRateLimitHandlerV3,
)
from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
@ -1824,6 +1828,287 @@ async def test_post_call_success_hook_priority_header_is_always_http_encodable(t
assert http_response.headers["x-litellm-rate-limiter-version"] == "v3"
def _steering_router(
preferred: str,
alternate: str,
preferred_rpm: int = 2,
alternate_rpm: int = 1000,
extra_params: Mapping[str, Mapping[str, object]] = MappingProxyType({}),
model_info: Mapping[str, Mapping[str, object]] = MappingProxyType({}),
) -> Router:
return Router(
model_list=[
{
"model_name": name,
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "test-key",
"api_base": "test-base",
"rpm": rpm,
**extra_params.get(name, {}),
},
"model_info": dict(model_info.get(name, {})),
}
for name, rpm in ((preferred, preferred_rpm), (alternate, alternate_rpm))
]
)
def _team_user(team_models: list[str], preferences: object) -> UserAPIKeyAuth:
return UserAPIKeyAuth(
team_id="flexible-team", team_models=team_models, team_metadata={"model_preferences": preferences}
)
async def _team_allowlist(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None:
return AlternateGrant() if model in user_api_key_dict.team_models else None
async def _pooled_team_allowlist(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None:
return AlternateGrant(access_groups=(f"{model}-pool",)) if model in user_api_key_dict.team_models else None
def _steering_handler() -> DynamicRateLimitHandler:
return DynamicRateLimitHandler(internal_usage_cache=DualCache(), is_model_authorized=_team_allowlist)
async def _saturate(handler: DynamicRateLimitHandler, model: str, requests: int = 1) -> None:
for _ in range(requests):
assert await handler.async_pre_call_hook(UserAPIKeyAuth(), DualCache(), {"model": model}, "completion") is None
@pytest.mark.asyncio
async def test_flexible_team_is_steered_to_its_alternate_once_the_preferred_model_saturates(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = _steering_handler()
handler.update_variables(llm_router=_steering_router("luna", "opus"))
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None
await _saturate(handler, "luna")
steered: Final = await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna", "n": 1}, "completion")
assert steered == {"model": "opus", "n": 1}
headers: Final = await handler.async_post_call_response_headers_hook(
data={"model": "opus", "stream": True}, user_api_key_dict=flexible, response=None
)
assert headers == {"x-litellm-model-steered-from": "luna"}
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash
get_or_create_request_stash().steered_from_model = "luna\u2728"
assert (
await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None
)
@pytest.mark.asyncio
async def test_teams_without_a_permitted_idle_alternate_stay_on_the_preferred_model(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = _steering_handler()
handler.update_variables(llm_router=_steering_router("luna", "opus", preferred_rpm=10, alternate_rpm=1))
await _saturate(handler, "luna", requests=5)
await _saturate(handler, "opus")
forbidden: Final = _team_user(["luna"], [{"model": "luna", "alternates": ["opus"]}])
unknown: Final = _team_user([], [{"model": "luna", "alternates": ["missing-model"]}])
malformed: Final = _team_user([], {"luna": "opus"})
both_saturated: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
for team in (forbidden, unknown, malformed, both_saturated):
assert await handler.async_pre_call_hook(team, DualCache(), {"model": "luna"}, "completion") is None
assert (
await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=team, response=None) is None
)
@pytest.mark.asyncio
async def test_a_zero_cost_preferred_model_is_never_steered_to_a_priced_alternate(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = _steering_handler()
free_luna: Final = {"luna": {"input_cost_per_token": 0, "output_cost_per_token": 0}}
handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=free_luna))
await _saturate(handler, "luna")
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None
assert (
await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None
)
@pytest.mark.asyncio
async def test_an_alternate_that_adds_model_level_guardrails_is_never_used(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = _steering_handler()
guarded_opus: Final = {"opus": {"guardrails": ["opus-only-input-scan"]}}
handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=guarded_opus))
await _saturate(handler, "luna")
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None
assert (
await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None
)
already_scanned: Final = {"model": "luna", "metadata": {"guardrails": ["opus-only-input-scan"]}}
steered: Final = await handler.async_pre_call_hook(flexible, DualCache(), already_scanned, "completion")
assert isinstance(steered, dict)
assert steered["model"] == "opus"
assert steered["metadata"]["guardrails"] == ["opus-only-input-scan"]
configured: Final = _steering_handler()
configured_opus: Final = {"opus": {"guardrails": [{"opus-only-input-scan": {"extra_body": {"strict": True}}}]}}
configured.update_variables(llm_router=_steering_router("luna", "opus", extra_params=configured_opus))
await _saturate(configured, "luna")
plain_scan: Final = {"model": "luna", "metadata": {"guardrails": ["opus-only-input-scan"]}}
assert await configured.async_pre_call_hook(flexible, DualCache(), plain_scan, "completion") is None
@pytest.mark.asyncio
async def test_a_steered_request_is_charged_to_the_groups_that_granted_the_alternate(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = DynamicRateLimitHandler(
internal_usage_cache=DualCache(), is_model_authorized=_pooled_team_allowlist
)
handler.update_variables(llm_router=_steering_router("luna", "opus"))
await _saturate(handler, "luna")
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
auth_copy: Final = flexible.model_copy()
data: Final = {"model": "luna", "metadata": {"user_api_key_auth": auth_copy, "tags": ["t"]}}
steered: Final = await handler.async_pre_call_hook(flexible, DualCache(), data, "completion")
assert isinstance(steered, dict)
assert steered["model"] == "opus"
assert steered["metadata"]["tags"] == ["t"]
assert steered["metadata"][MODEL_ACCESS_GROUP_METADATA_KEY] == ["opus-pool"]
assert steered["metadata"]["user_api_key_auth"].matched_model_access_groups == ["opus-pool"]
assert flexible.matched_model_access_groups == ["opus-pool"]
assert auth_copy.matched_model_access_groups is None
@pytest.mark.asyncio
async def test_an_alternate_granted_by_a_group_the_request_did_not_reserve_against_is_never_used(
monkeypatch: pytest.MonkeyPatch,
):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = DynamicRateLimitHandler(
internal_usage_cache=DualCache(), is_model_authorized=_pooled_team_allowlist
)
handler.update_variables(llm_router=_steering_router("luna", "opus"))
await _saturate(handler, "luna")
def reserved_against(group: str) -> UserAPIKeyAuth:
user: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
user.budget_reservation = {
"entries": [
{
"counter_key": f"spend:model_access_group:{group}",
"entity_type": "Model access group",
"entity_id": group,
"reserved_cost": 0.01,
"applied_adjustment": 0.0,
}
]
}
return user
assert (
await handler.async_pre_call_hook(reserved_against("luna-pool"), DualCache(), {"model": "luna"}, "completion")
is None
)
steered: Final = await handler.async_pre_call_hook(
reserved_against("opus-pool"), DualCache(), {"model": "luna"}, "completion"
)
assert isinstance(steered, dict)
assert steered["model"] == "opus"
@pytest.mark.asyncio
async def test_a_pricier_alternate_is_never_used_while_a_budget_reservation_is_held(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = _steering_handler()
pricier_opus: Final = {
"luna": {"input_cost_per_token": 0.000001, "output_cost_per_token": 0.000002},
"opus": {"input_cost_per_token": 0.000005, "output_cost_per_token": 0.000002},
}
handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=pricier_opus))
await _saturate(handler, "luna")
reserved: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
reserved.budget_reservation = {"reserved_cost": 0.0001, "entries": []}
assert await handler.async_pre_call_hook(reserved, DualCache(), {"model": "luna"}, "completion") is None
unreserved: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
steered: Final = await handler.async_pre_call_hook(unreserved, DualCache(), {"model": "luna"}, "completion")
assert isinstance(steered, dict)
assert steered["model"] == "opus"
unpriced: Final = _steering_handler()
unpriced_opus: Final = {"opus": {"model": "openai/no-such-model-xyz"}}
unpriced.update_variables(llm_router=_steering_router("luna", "opus", extra_params=unpriced_opus))
await _saturate(unpriced, "luna")
assert await unpriced.async_pre_call_hook(reserved, DualCache(), {"model": "luna"}, "completion") is None
@pytest.mark.asyncio
async def test_an_alternate_that_does_not_serve_the_endpoint_is_never_used(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
handler: Final = _steering_handler()
embeddings_only: Final = {"opus": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}}
handler.update_variables(llm_router=_steering_router("luna", "opus", model_info=embeddings_only))
await _saturate(handler, "luna")
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "acompletion") is None
assert (
await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None
)
@pytest.mark.asyncio
async def test_an_unpriced_alternate_is_never_used_while_unpriced_models_are_blocked(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True)
handler: Final = _steering_handler()
unpriced_opus: Final = {"opus": {"model": "openai/no-such-model-xyz"}}
handler.update_variables(llm_router=_steering_router("luna", "opus", extra_params=unpriced_opus))
await _saturate(handler, "luna")
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None
@pytest.mark.asyncio
async def test_a_broken_alternate_check_keeps_the_request_on_the_preferred_model(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "priority_reservation", None)
monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5)
async def broken(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None:
raise RuntimeError("membership lookup is down")
handler: Final = DynamicRateLimitHandler(internal_usage_cache=DualCache(), is_model_authorized=broken)
handler.update_variables(llm_router=_steering_router("luna", "opus"))
await _saturate(handler, "luna")
flexible: Final = _team_user(["luna", "opus"], [{"model": "luna", "alternates": ["opus"]}])
assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None
assert (
await handler.async_post_call_response_headers_hook(data={}, user_api_key_dict=flexible, response=None) is None
)
def _borrowing_handler(
monkeypatch: pytest.MonkeyPatch,
borrows: bool,