diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 3654d0b9176..bd20fbf3ab3 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -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: """ diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 6b53e952f98..4cbe08d0042 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -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): diff --git a/tests/integration/routing/test_team_model_preference_steering.py b/tests/integration/routing/test_team_model_preference_steering.py new file mode 100644 index 00000000000..c99cfcf660b --- /dev/null +++ b/tests/integration/routing/test_team_model_preference_steering.py @@ -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 diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index a8265680856..0bf0ea4d6ad 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -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,