From 172e16acde90355b23668633b85efa675f6d5cc6 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Thu, 12 Nov 2026 21:57:18 -0800 Subject: [PATCH] feat(proxy): router_settings.fallback_on_saturation walks key and team fallbacks before a model saturates, replacing team model_preferences (#51439) * feat(proxy): fallback_on_saturation walks key and team fallbacks before a model saturates, replacing team_metadata.model_preferences Team router_settings.fallbacks now fire on proxy-side rate limits: the pre-call fallback walk reads router settings Key > Team > Global instead of the key's column alone, and skips a fallback the caller may not call. A key or team with router_settings.fallback_on_saturation: true walks that same fallback list as soon as the requested model reaches saturation_threshold in dynamic_rate_limiter_v3, and goes back to the requested model when no fallback admits the request. Responses carry x-litellm-fallback-reason: rate_limit:model_saturation_fallback. Removes the team_metadata.model_preferences steering and the x-litellm-model-steered-from header from #51350, which did the same job through a second knob. * fix(proxy): strict bool for fallback_on_saturation, drop a stale import and settle the pre-call tests fallback_on_saturation only turns on for a real boolean true, so a string like "yes" in router_settings no longer enables it. The saturation unit test reads saturation through the limiter's cache, so it asserts the threshold was reached instead of an exact count. The response polling test now hands the pre-call path an awaitable router_settings lookup, as the pre-call fallback walk reads Key > Team settings before the first call. * fix(proxy): pre-call fallbacks pass the budget and access-group checks a direct request does A fallback the pre-call walk picks now has to clear every per-model budget layer and the budgets of the access groups that grant it, and its spend is charged to those groups instead of the requested model's. At saturation a fallback that fails for any reason counts as a miss and the request goes back to the requested model. The client body can no longer opt itself into the saturation walk through router_settings_override, only the saved key or team router_settings can. The integration test for a forbidden fallback waits for the preferred model's own 429 so it cannot pass before saturation. * refactor(proxy): one helper decides which refused fallbacks end the pre-call walk _pre_call_with_fallbacks had grown past the strict C901 ceiling, so the two except clauses fold into one and _fallback_refusal_ends_the_walk carries the rule: at saturation every refusal is a miss, otherwise only a non-rate-limit failure or a hard per-model 429 reaches the caller. * fix(proxy): keep the caller's request object through the pre-call fallback walk Copying the request to drop a client fallback_on_saturation opt-in handed the hooks a different dict than the one the caller holds, so header metadata written into one never reached the other (test_add_litellm_metadata_from_request_headers). drop_client_saturation_opt_in now edits the request in place. * fix(proxy): pre-call fallbacks keep the pricing and budget-reservation rules a direct request passes A fallback is skipped when unpriced models are blocked and it has no pricing, when a zero-cost request would land on a priced model, and when the budget auth reserved for the requested model does not fit it (higher per-token prices, or reserved against an access group that does not grant it). * refactor(proxy): read fallback prices through the router registry entry point router_surface_check only lets the router surface shrink, so the per-token price lookup goes through registry().group() instead of the deprecated get_model_group_info. * fix(proxy): a fallback's group budget check skips the groups the request already holds The auth-time reservation sits in the same counter the plain budget read uses, so a fallback in a group the request already reserved against was refused by its own hold once the group was within one estimate of max_budget. Groups without a hold keep the plain read. * test(dynamic_rate_limiter_v3): the throttle metric test from main uses fallback_on_saturation instead of the removed steering A saturation handoff carries no rate_limited_by reason, so it is never counted as a throttle. * fix(proxy): the saturation handoff only fires under the fallback walk _pre_call_with_fallbacks arms the request before the hooks run, so routes that call the hooks directly (batches, fine-tuning, polling, images) keep admitting or refusing an opted-in key on a saturated model as before instead of surfacing the handoff 429. * style(proxy): keep the cast-ok reason within the line length * fix(proxy): never invent a metadata bucket when arming the saturation fallback walk, take main's limiter metric tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): arm the saturation fallback walk in the request context, refuse fallbacks for agent-owned keys Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep the saturation walk state out of the body, check the end user's allowlist on fallbacks, release a refused fallback's rate limit hold Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): only load the end user for fallback checks when a database is connected Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): grant pre-call fallbacks on the model an alias serves and release the rate limited model's hold before the walk Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): check a pre-call fallback's per-model budgets on its listed name and its served group, as auth does Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert the listed fallback name's budget blocks an aliased fallback, not the order of budget reads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): release the dynamic limiter's reservation when the fallback walk leaves a model, and replay an exhausted walk with the requested model The dynamic rate limiter admits through its own inner v3 limiter, so the walk's hold release never reached its team-share reservation. It now exposes async_post_call_failure_hook, which refunds that reservation when the proxy refuses a request after admission, and _release_pre_call_hold calls it next to the registered limiter. An exhausted saturation walk replayed with the already aliased model name, so chained model_alias_map entries moved the request on to a further alias. It now replays with the model argument the first attempt saw. The hold-release test drives the real v3 limiter with one slot per key instead of recording hook calls. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): type the dynamic limiter's failure hook request dict for the type-discipline gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 416 +++++- .../common_utils/proxy_rate_limit_error.py | 88 +- .../proxy/hooks/dynamic_rate_limiter_v3.py | 469 +------ .../hooks/parallel_request_limiter_v3.py | 1 - litellm/types/router.py | 1 + .../src/lib/http/schema.d.ts | 2 + ...ring.py => test_fallback_on_saturation.py} | 77 +- .../test_proxy_rate_limit_error.py | 70 + .../hooks/test_dynamic_rate_limiter_v3.py | 437 ++---- .../proxy/test_common_request_processing.py | 1166 ++++++++++++++++- .../test_response_polling_pre_call_checks.py | 2 +- 11 files changed, 1848 insertions(+), 881 deletions(-) rename tests/integration/routing/{test_team_model_preference_steering.py => test_fallback_on_saturation.py} (57%) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3efdbfd30e2..47c807f4f47 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -5,7 +5,17 @@ import json import logging import math import time -from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Coroutine, Iterator, Mapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Coroutine, + Iterator, + Mapping, + MutableMapping, + Sequence, +) from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache @@ -153,7 +163,11 @@ from litellm.proxy.common_utils.sse_keepalive import ( ) from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.guardrails.auto_router_compression import arm_pre_call as _arm_auto_router_compression -from litellm.proxy.hooks.parallel_request_limiter_v3 import get_or_create_request_stash +from litellm.proxy.hooks.dynamic_rate_limiter_v3 import PROXY_DynamicRateLimitHandlerV3 +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + PROXY_MaxParallelRequestsHandler_v3, + get_or_create_request_stash, +) from litellm.proxy.hooks.proxy_track_cost_callback import failure_row_response_cost from litellm.proxy.native_compaction import with_proxy_compaction_executor from litellm.proxy.route_llm_request import ( @@ -1022,6 +1036,7 @@ async def hierarchical_router_settings_updates( _REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None) +_FALLBACK_LIST: Final[TypeAdapter[list[object]]] = TypeAdapter(list[object], config=ConfigDict(strict=True)) _DECODED_JSON: Final = TypeAdapter(object) _ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -2266,6 +2281,265 @@ def _timing_values( return getattr(logging_obj, "response_timing_metrics", None) or {} +def _fallback_refusal_ends_the_walk( + exc: Exception, at_saturation: bool, general_settings: Mapping[str, object] +) -> bool: + """At saturation the requested model still has room, so a fallback refused for any reason is a miss. Otherwise a + fallback's own non-rate-limit failure, or a hard per-model 429, is the answer the caller gets.""" + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, + ProxyRateLimitError, + per_model_rate_limits_disable_fallbacks, + ) + + if at_saturation: + return False + if not isinstance(exc, ProxyRateLimitError): + return True + return exc.descriptor_key in PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS and per_model_rate_limits_disable_fallbacks( + general_settings + ) + + +async def _fallback_access_groups( + fallback_model: str, user_api_key_dict: UserAPIKeyAuth, llm_router: Router, proxy_logging_obj: ProxyLogging +) -> tuple[str, ...]: + from litellm.proxy.auth.auth_checks import ( + collect_matched_model_access_groups, + get_project_object, + get_team_membership, + get_team_object, + ) + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + return () + team_id: Final = user_api_key_dict.team_id + user_id: Final = user_api_key_dict.user_id + project_id: Final = user_api_key_dict.project_id + 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 + ) + return await collect_matched_model_access_groups( + model=fallback_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, + ) + + +_PER_TOKEN_PRICES: Final[TypeAdapter[tuple[float, float]]] = TypeAdapter(tuple[float, float]) +_RESERVATION_ENTRIES: Final[TypeAdapter[tuple[Mapping[str, object], ...]]] = TypeAdapter( + tuple[Mapping[str, object], ...] +) +_MODEL_ACCESS_GROUP_ENTITY: Final = "Model access group" + + +def _per_token_prices(llm_router: Router, model: str) -> tuple[float, float] | None: + info: Final = llm_router.registry().group(model) + if info is None: + return None + try: + return _PER_TOKEN_PRICES.validate_python((info.input_cost_per_token, info.output_cost_per_token)) + except ValidationError: + return None + + +def _reserved_access_groups(user_api_key_dict: UserAPIKeyAuth) -> frozenset[str] | None: + """The model access groups the request's ``budget_reservation`` already holds, None when its entries are + unreadable.""" + reservation: Final = user_api_key_dict.budget_reservation + if reservation is None: + return frozenset() + try: + entries: Final = _RESERVATION_ENTRIES.validate_python(reservation.get("entries") or ()) + except ValidationError: + return None + return frozenset( + str(entry.get("entity_id")) for entry in entries if entry.get("entity_type") == _MODEL_ACCESS_GROUP_ENTITY + ) + + +def _reservation_covers( + user_api_key_dict: UserAPIKeyAuth, + access_groups: Sequence[str], + held: frozenset[str] | None, + requested_model: str, + fallback_model: str, + llm_router: Router, +) -> bool: + """Whether the budget auth reserved for ``requested_model`` still fits ``fallback_model``: the fallback costs no + more per token than the model the estimate was made for, and every access group ``held`` by the reservation + grants the fallback, so settlement never charges a group that did not serve.""" + if user_api_key_dict.budget_reservation is None: + return True + if held is None: + return False + requested: Final = _per_token_prices(llm_router, requested_model) + fallback: Final = _per_token_prices(llm_router, fallback_model) + if requested is None or fallback is None or fallback[0] > requested[0] or fallback[1] > requested[1]: + return False + return held <= frozenset(access_groups) + + +def _fallback_pricing_admits(requested_model: str, fallback_model: str, llm_router: Router) -> bool: + """A zero-cost ``requested_model`` waived the aggregate budget checks at auth, so only another zero-cost model may + take its place, and an unpriced fallback is refused while ``block_requests_for_models_without_pricing`` is on.""" + from litellm.proxy.auth.auth_checks import ( + _unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # the pre-call walk shares the auth-time fallback contract + ) + + if is_model_cost_zero(model=requested_model, llm_router=llm_router) and not is_model_cost_zero( + model=fallback_model, llm_router=llm_router + ): + return False + return not ( + litellm.block_requests_for_models_without_pricing + and _unpriced_models_in_request(model=fallback_model, llm_router=llm_router) + ) + + +async def _release_pre_call_hold( + proxy_logging_obj: ProxyLogging, + request_data: Mapping[str, object], + refusal: Exception, + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Give back the slot and tokens the rate limiter holds for a model the walk is leaving: the requested model a + later hook rate limited after the limiter admitted it, or a fallback the proxy refused after admitting it. The + walk's next attempt overwrites the request's stash, so a hold not released here stays charged until the window + resets. The request has not failed, so only the limiters' releases run, not the failure logging and alerting of + ``post_call_failure_hook``. The dynamic rate limiter admits through its own inner limiter, so its team-share + reservation is released alongside the registered limiter's slot and tokens.""" + limiter: Final = proxy_logging_obj.get_proxy_hook("parallel_request_limiter") + dynamic_limiters: Final = litellm.logging_callback_manager.get_custom_loggers_for_type( + PROXY_DynamicRateLimitHandlerV3 + ) + for holder in (limiter, *dynamic_limiters): + if not isinstance(holder, (PROXY_MaxParallelRequestsHandler_v3, PROXY_DynamicRateLimitHandlerV3)): + continue + await holder.async_post_call_failure_hook( # pyright: ignore[reportUnknownMemberType] # the limiter hooks take the proxy's bare request dict + request_data=dict(request_data), original_exception=refusal, user_api_key_dict=user_api_key_dict + ) + + +async def _fallback_grant( + fallback_model: str, + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router | None, + proxy_logging_obj: ProxyLogging, + requested_model: str, + listed_as: str | None = None, +) -> list[str] | None: + """The model access groups a pre-call fallback to ``fallback_model`` charges, or None when the caller may not use + it. The caller passes the group the fallback is served by, through every alias rewrite, and the name the fallback + list gave it as ``listed_as``, so an alias never skips the budgets of the model that answers nor those set on the + listed name: an agent-owned key never leaves the model it asked for, as at auth, then the key, team and project allowlists + while ``enforce_fallback_model_access`` is on, the end user's allowlist, the pricing rules, every per-model budget layer, the budgets of the access groups that grant the model, and the budget auth reserved for + ``requested_model``, the same checks a direct request passes at auth.""" + from litellm.proxy.auth.auth_checks import ( + _model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # the pre-call walk shares the auth-time fallback contract + can_customer_access_model, + get_end_user_object, + ) + from litellm.proxy.auth.fallback_model_access import is_model_authorized_for_token, router_fallback_access_check + from litellm.proxy.auth.user_api_key_auth import ( + _is_within_every_model_budget, # pyright: ignore[reportPrivateUsage] # the pre-call walk shares the auth-time fallback contract + ) + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client, user_api_key_cache + + if llm_router is None: + return [] + if user_api_key_dict.agent_id is not None: + return None + if router_fallback_access_check.is_enforced() and not await is_model_authorized_for_token( + model=fallback_model, valid_token=user_api_key_dict, llm_router=llm_router + ): + return None + if not _fallback_pricing_admits(requested_model, fallback_model, llm_router): + return None + try: + end_user: Final = ( + await get_end_user_object( + end_user_id=user_api_key_dict.end_user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_api_key_dict.end_user_id is not None and prisma_client is not None + else None + ) + if end_user is not None and end_user.models: + can_customer_access_model( + model=fallback_model, end_user_object=end_user, llm_router=llm_router, valid_token=user_api_key_dict + ) + if not is_model_cost_zero(model=fallback_model, llm_router=llm_router): + budget_names: Final = tuple( + dict.fromkeys( + (listed_as or fallback_model, fallback_model, llm_router.registry().resolve(fallback_model)) + ) + ) + for budget_name in budget_names: + if not await _is_within_every_model_budget( + valid_token=user_api_key_dict, model_max_budget_limiter=model_max_budget_limiter, model=budget_name + ): + return None + access_groups: Final = await _fallback_access_groups( + fallback_model, user_api_key_dict, llm_router, proxy_logging_obj + ) + held: Final = _reserved_access_groups(user_api_key_dict) + await _model_access_group_max_budget_check( + matched_model_access_groups=tuple(group for group in access_groups if group not in (held or ())), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if not _reservation_covers(user_api_key_dict, access_groups, held, requested_model, fallback_model, llm_router): + return None + except Exception as e: # noqa: BLE001 # fail closed: a grant the walk cannot settle must not admit the fallback + verbose_proxy_logger.warning( + "Skipping fallback to model=%s: %s", + fallback_model.replace("\r", "").replace("\n", ""), + e, + ) + return None + return list(access_groups) + + async def _team_member_within_fallback_model_budget( user_api_key_dict: UserAPIKeyAuth, model: str, @@ -2918,6 +3192,12 @@ class ProxyBaseLLMRequestProcessing: PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, ProxyRateLimitError, RateLimitFallback, + arm_saturation_fallback_walk, + disarm_saturation_fallback_walk, + drop_client_saturation_opt_in, + fallback_on_saturation_enabled, + is_saturation_fallback, + mark_saturation_fallbacks_exhausted, per_model_rate_limits_disable_fallbacks, with_rate_limit_fallback, ) @@ -2932,12 +3212,26 @@ class ProxyBaseLLMRequestProcessing: and per_model_rate_limits_disable_fallbacks(general_settings_view) ) - configured_fallbacks: Final = ( - self._configured_fallbacks(llm_router=llm_router, user_api_key_dict=user_api_key_dict) + router_settings: Final = ( + await self._hierarchical_router_settings(user_api_key_dict, proxy_config, proxy_logging_obj) if llm_router is not None else None ) - pristine: Final = independent_snapshot(self.data) if configured_fallbacks else None + configured_fallbacks: Final = ( + self._configured_fallbacks(llm_router=llm_router, router_settings=router_settings) + if llm_router is not None + else None + ) + scrubbed: Final[dict] = self.data # mutable-ok: caller-defined request-payload shape + payload: Final = cast(MutableMapping[str, object], scrubbed) # cast-ok: untyped caller payload + drop_client_saturation_opt_in(payload) + walk: Final = arm_saturation_fallback_walk() + self.data = scrubbed + pristine: Final = ( + independent_snapshot(scrubbed) + if configured_fallbacks or fallback_on_saturation_enabled(router_settings) + else None + ) try: return await self.common_processing_pre_call_logic( @@ -2959,20 +3253,17 @@ class ProxyBaseLLMRequestProcessing: except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data original_model: Final = rate_limited_data.get("model") - if ( - pristine is None - or not configured_fallbacks - or rate_limited_data.get("disable_fallbacks") - or not isinstance(original_model, str) - or is_hard_per_model_limit(original_exc) - ): + # The limiter raises this one while the requested model still has room, so a walk that admits nothing + # ends on that model instead of a 429. + at_saturation: Final = is_saturation_fallback(original_exc) + if pristine is None or not isinstance(original_model, str) or is_hard_per_model_limit(original_exc): raise - - fallback_models: Final = self._resolve_fallback_models( - model=original_model, - fallbacks=configured_fallbacks, + fallback_models: Final = ( + self._resolve_fallback_models(model=original_model, fallbacks=configured_fallbacks) + if configured_fallbacks and not rate_limited_data.get("disable_fallbacks") + else None ) - if not fallback_models: + if not fallback_models and not at_saturation: raise verbose_proxy_logger.info( @@ -2988,11 +3279,26 @@ class ProxyBaseLLMRequestProcessing: if isinstance(raw_router_settings, Mapping) else None ) + await _release_pre_call_hold( + proxy_logging_obj, _REQUEST_FIELDS.validate_python(rate_limited_data), original_exc, user_api_key_dict + ) attempt_numbers: Final[Iterator[int]] = itertools.count(1) + requested_groups: Final = user_api_key_dict.matched_model_access_groups try: - for fallback_model in fallback_models: + for fallback_model in fallback_models or (): if fallback_model == original_model: continue + if ( + grant := await _fallback_grant( + pre_call_served_model(fallback_model, user_api_key_dict, group_alias), + user_api_key_dict, + llm_router, + proxy_logging_obj, + original_model, + listed_as=fallback_model, + ) + ) is None: + continue if not await _team_member_within_fallback_model_budget( user_api_key_dict=user_api_key_dict, model=fallback_model, @@ -3012,6 +3318,7 @@ class ProxyBaseLLMRequestProcessing: ): continue attempt = next(attempt_numbers) + user_api_key_dict.matched_model_access_groups = grant # rebind-ok: bill the fallback's groups admitted: tuple[ Mapping[str, object], LiteLLMLoggingObj ] = await self.common_processing_pre_call_logic( @@ -3031,9 +3338,16 @@ class ProxyBaseLLMRequestProcessing: llm_router=llm_router, rate_limited_model=original_model, ) - except ProxyRateLimitError as fallback_exc: - if is_hard_per_model_limit(fallback_exc): + except Exception as fallback_exc: # noqa: BLE001 # the helper decides which refusals end the walk + user_api_key_dict.matched_model_access_groups = requested_groups # rebind-ok: undo the fallback + if _fallback_refusal_ends_the_walk(fallback_exc, at_saturation, general_settings_view): raise + await _release_pre_call_hold( + proxy_logging_obj, + self.data, # pyright: ignore[reportUnknownMemberType] # caller-defined request-payload shape + fallback_exc, + user_api_key_dict, + ) continue else: admitted_data, admitted_logging_obj = admitted @@ -3050,19 +3364,65 @@ class ProxyBaseLLMRequestProcessing: ) self.data = stamped return stamped, admitted_logging_obj + if not at_saturation: + self.data = rate_limited_data + raise original_exc + verbose_proxy_logger.info( + "No fallback of model=%s is below the saturation threshold, serving the request on it", + original_model.replace("\r", "").replace("\n", ""), + ) + mark_saturation_fallbacks_exhausted() + self.data = independent_snapshot(dict(_REQUEST_FIELDS.validate_python(pristine))) + return await self.common_processing_pre_call_logic( + request=request, + general_settings=general_settings, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=user_api_key_dict, + version=version, + proxy_config=proxy_config, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + model=model, # the requested name as the first attempt saw it, so its aliases apply once, not again + route_type=route_type, # pyright: ignore[reportArgumentType] # the proxy hands its route literal over as str, as the two calls above take it + llm_router=llm_router, + ) except BaseException: + user_api_key_dict.matched_model_access_groups = requested_groups # rebind-ok: undo the fallback self.data = rate_limited_data raise - - self.data = rate_limited_data - raise original_exc + finally: + disarm_saturation_fallback_walk(walk) @staticmethod - def _configured_fallbacks(llm_router: Router, user_api_key_dict: UserAPIKeyAuth) -> list | None: - key_router_settings: Final = user_api_key_dict.router_settings - key_fallbacks: Final = key_router_settings.get("fallbacks") if isinstance(key_router_settings, dict) else None - fallbacks: Final = key_fallbacks if key_fallbacks is not None else llm_router.settings().fallbacks - return fallbacks if isinstance(fallbacks, list) and fallbacks else None + async def _hierarchical_router_settings( + user_api_key_dict: UserAPIKeyAuth, proxy_config: ProxyConfig, proxy_logging_obj: ProxyLogging + ) -> Mapping[str, object] | None: + """The key's ``router_settings``, else its team's, read through the same cached lookup the request body + gets them from, so the pre-call fallback walk and the router agree on Key > Team > Global.""" + from litellm.proxy.proxy_server import prisma_client + + return _ROUTER_SETTINGS_ADAPTER.validate_python( + await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ) + + @staticmethod + def _configured_fallbacks(llm_router: Router, router_settings: Mapping[str, object] | None) -> list[object] | None: + settings_fallbacks: Final = router_settings.get("fallbacks") if router_settings is not None else None + fallbacks: Final[object] = ( + settings_fallbacks if settings_fallbacks is not None else llm_router.settings().fallbacks + ) + try: + configured: Final = _FALLBACK_LIST.validate_python(fallbacks) + except ValidationError: + return None + return configured or None @staticmethod def _resolve_fallback_models(model: str, fallbacks: list) -> list[str] | None: diff --git a/litellm/proxy/common_utils/proxy_rate_limit_error.py b/litellm/proxy/common_utils/proxy_rate_limit_error.py index 32daf68e910..82d7f08d8ac 100644 --- a/litellm/proxy/common_utils/proxy_rate_limit_error.py +++ b/litellm/proxy/common_utils/proxy_rate_limit_error.py @@ -36,10 +36,11 @@ This module provides a single proxy-side error class that: """ import json -from collections.abc import Mapping +from collections.abc import Mapping, MutableMapping +from contextvars import ContextVar, Token from dataclasses import dataclass from types import MappingProxyType -from typing import Final +from typing import Final, Literal, TypeAlias from fastapi import HTTPException from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError @@ -78,6 +79,87 @@ def per_model_rate_limits_disable_fallbacks(general_settings: Mapping[str, objec return False +FALLBACK_ON_SATURATION_SETTING: Final = "fallback_on_saturation" +SATURATION_FALLBACK_DESCRIPTOR_KEY: Final = "model_saturation_fallback" +_FALLBACK_ON_SATURATION_FLAG: Final[TypeAdapter[bool | None]] = TypeAdapter(bool | None) + + +def fallback_on_saturation_enabled(router_settings: object) -> bool: + """Whether the key's or team's ``router_settings`` ask for fallbacks once the requested model reaches + ``saturation_threshold``, before the dynamic rate limiter admits or refuses the request on that model.""" + raw_value: Final = _metadata_bucket(router_settings).get(FALLBACK_ON_SATURATION_SETTING) + try: + return _FALLBACK_ON_SATURATION_FLAG.validate_python(raw_value, strict=True) is True + except ValidationError: + verbose_proxy_logger.warning( + "router_settings.%s=%r is not a boolean, treating it as disabled", FALLBACK_ON_SATURATION_SETTING, raw_value + ) + return False + + +def drop_client_saturation_opt_in( + request_data: MutableMapping[str, object], # mutable-ok: in-place edit keeps the caller's request object +) -> None: + """Drop ``fallback_on_saturation`` from the client's ``router_settings_override``. Only the key's or team's saved + ``router_settings`` turn the saturation walk on, and the proxy writes those over this field before the limiter + reads it.""" + override: Final = _metadata_bucket(request_data.get("router_settings_override")) + if FALLBACK_ON_SATURATION_SETTING not in override: + return + request_data["router_settings_override"] = { # rebind-ok: the request object must stay the caller's + key: value for key, value in override.items() if key != FALLBACK_ON_SATURATION_SETTING + } + + +_SaturationWalkState: TypeAlias = Literal["off", "armed", "exhausted"] +_SATURATION_FALLBACK_WALK: Final[ContextVar[_SaturationWalkState]] = ContextVar( + "litellm_saturation_fallback_walk", default="off" +) + + +def saturation_fallback_walk_armed() -> bool: + """Whether a caller that catches the handoff (``_pre_call_with_fallbacks``) is running this request's hooks and + has not walked the fallbacks yet.""" + return _SATURATION_FALLBACK_WALK.get() == "armed" + + +def saturation_fallbacks_exhausted() -> bool: + """Whether this request already tried every fallback at saturation and is back on the model it asked for.""" + return _SATURATION_FALLBACK_WALK.get() == "exhausted" + + +def arm_saturation_fallback_walk() -> Token[_SaturationWalkState]: + """Mark the current request as running under the fallback walk, so the limiter may hand a saturated model off to + it. The state lives in the request's context, never in its body, so nothing reaches the provider and no client + field can arm, skip or exhaust the walk. Routes that call the hooks directly never arm it and keep admitting or + refusing as before.""" + return _SATURATION_FALLBACK_WALK.set("armed") + + +def mark_saturation_fallbacks_exhausted() -> None: + _SATURATION_FALLBACK_WALK.set("exhausted") + + +def disarm_saturation_fallback_walk(token: Token[_SaturationWalkState]) -> None: + _SATURATION_FALLBACK_WALK.reset(token) + + +def fallback_on_saturation_requested(request_data: Mapping[str, object]) -> bool: + """Whether the dynamic rate limiter should hand a saturated model's request to the fallback walk instead of + admitting or refusing it: the key or team opted in through ``router_settings_override`` and the fallback walk + is the caller and has not already run for this request.""" + return ( + fallback_on_saturation_enabled(request_data.get("router_settings_override")) + and saturation_fallback_walk_armed() + ) + + +def is_saturation_fallback(exc: "ProxyRateLimitError") -> bool: + """Whether ``exc`` is the limiter asking for the fallback walk rather than refusing the request. The model it + names still has room, so when no fallback admits the request it goes back to that model instead of a 429.""" + return exc.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY + + RATE_LIMIT_FALLBACK_METADATA_KEY: Final = "rate_limit_fallback" RATE_LIMIT_FALLBACK_REASON_HEADER: Final = "x-litellm-fallback-reason" ATTEMPTED_FALLBACKS_HEADER: Final = "x-litellm-attempted-fallbacks" @@ -204,7 +286,7 @@ def _coerce_message(detail: object) -> str: try: return json.dumps(detail) except (TypeError, ValueError): - return str(detail) + return f"{detail}" return str(detail) diff --git a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 44fc44c337d..798b1e5cc74 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -2,16 +2,15 @@ Dynamic rate limiter v3 - Saturation-aware priority-based rate limiting """ -import json import os -from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal from fastapi import HTTPException -from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, ValidationError from redis.exceptions import RedisError import litellm @@ -27,11 +26,12 @@ 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.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, ProxyRateLimitError, + fallback_on_saturation_requested, map_v3_rate_limit_type, ) from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # legacy module exports @@ -56,7 +56,6 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( # noqa: F401 # l claim_request_stash_for_data, get_or_create_request_stash, get_request_stash, - get_request_stash_for_call, request_stash_for_callback, ) from litellm.proxy.hooks.priority_class_share import ( @@ -73,7 +72,7 @@ from litellm.proxy.hooks.rate_limiter_utils import ( resolve_llm_provider_for_rate_limit, socket_session_usage, ) -from litellm.proxy.utils import InternalUsageCache, check_and_merge_model_level_guardrails +from litellm.proxy.utils import InternalUsageCache 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 @@ -83,273 +82,6 @@ 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.customer_hierarchy_auth import check_customer_hierarchy_model_access - 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 - ) - if user_api_key_dict.customer_hierarchy is not None: - check_customer_hierarchy_model_access(user_api_key_dict.customer_hierarchy, model, llm_router) - 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.""" @@ -572,10 +304,8 @@ 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) self.active_teams = ActiveTeamRegistry(internal_usage_cache) @@ -1639,11 +1369,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): 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: - admitted: Final = await self._admit_steered_or_preferred( - model, requested_group_info, steered, user_api_key_dict, data, call_type - ) + await self._admit(model, requested_group_info, user_api_key_dict, data, call_type) except ProxyRateLimitError as refusal: raise self._throttled(refusal, model, user_api_key_dict) from None except HTTPException: @@ -1652,34 +1379,6 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): verbose_proxy_logger.error("Error in dynamic rate limiter: %s, allowing request", e) # 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, - data: Mapping[str, object], - call_type: str | None, - ) -> "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, data, call_type) - 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, data, call_type) return None async def _admit( @@ -1690,7 +1389,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): data: Mapping[str, object], call_type: str | None, ) -> None: - """Read saturation, then the three-phase rate limit check, raising the 429 when ``model`` is out of capacity.""" + """Read saturation, then the three-phase rate limit check, raising the 429 when ``model`` is out of capacity. + A key or team with ``router_settings.fallback_on_saturation`` is sent to its fallbacks first once ``model`` is + at or above ``saturation_threshold``, and comes back here only when none of them admitted the request.""" 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 @@ -1698,6 +1399,8 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): f"[Dynamic Rate Limiter] Model={model}, Saturation={saturation:.1%}, " f"Threshold={saturation_threshold:.1%}, Priority={priority}" ) + if saturation >= saturation_threshold and fallback_on_saturation_requested(data): + raise self._saturation_fallback_error(model, saturation, priority) await self._check_rate_limits( model=model, model_group_info=model_group_info, @@ -1708,104 +1411,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): call_type=call_type, ) - 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, + def _saturation_fallback_error(self, model: str, saturation: float, priority: str | None) -> ProxyRateLimitError: + resolved_model, llm_provider = resolve_llm_provider_for_rate_limit(model) + return ProxyRateLimitError( + detail={ + "error": f"Model {model} is at {saturation:.0%} of its capacity, at or above the saturation " + f"threshold of {_get_priority_settings().saturation_threshold:.0%}, trying fallbacks first. " + f"Priority: {priority}" + }, + headers={"retry-after": str(self.v3_limiter.window_size), "x-litellm-priority": priority or "default"}, + model=resolved_model, + llm_provider=llm_provider, + descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY, ) - 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( @@ -1842,21 +1460,6 @@ 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: Mapping[str, object], response_obj: object, start_time: object, end_time: object @@ -1986,6 +1589,30 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): verbose_proxy_logger.exception("Error in dynamic rate limiter success event: %s", e) @with_service_target(RATE_LIMITS_TARGET) + async def async_post_call_failure_hook( + self, + request_data: dict[str, object], # mutable-ok: the base hook declares this parameter as a dict + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> None: + """Give back the team-share tokens held at admission when the proxy refuses the request before the call runs, + as a later hook or the pre-call fallback walk does. ``async_log_failure_event`` is a completion callback and + never fires for those refusals, so without this the hold stays charged until the window resets. Idempotent: + the first release zeroes the stash's reservation, so a settlement that still follows charges only usage.""" + stash: Final = get_request_stash() + if stash is None: + return + refund: Final = _team_model_settlement(stash, None, 0, self.v3_limiter.window_size) + stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_cap_reserved_tokens = 0 + stash.borrowing_team_model_reserved_window = None + stash.borrowing_team_model_cap_reserved_window = None + if refund: + await self.v3_limiter.async_increment_reservation_aware_tokens( + pipeline_operations=refund, parent_otel_span=user_api_key_dict.parent_otel_span + ) + async def async_log_failure_event( self, kwargs: object, response_obj: object, start_time: object, end_time: object ) -> None: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b8eeaeab217..a6697004fc7 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -620,7 +620,6 @@ 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) diff --git a/litellm/types/router.py b/litellm/types/router.py index 6cc38a30d6f..868e0022871 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -237,6 +237,7 @@ class UpdateRouterConfig(LiteLLMBaseModel): tag_routing_prefix: str | None = None optional_pre_call_checks: OptionalPreCallChecks | None = None fallback_on_unsupported_params: bool | None = None + fallback_on_saturation: bool | None = None model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index c6a636d8c50..5b240a6fe3b 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -51867,6 +51867,8 @@ export interface components { cooldown_time?: number | null; /** Enable Tag Filtering */ enable_tag_filtering?: boolean | null; + /** Fallback On Saturation */ + fallback_on_saturation?: boolean | null; /** Fallback On Unsupported Params */ fallback_on_unsupported_params?: boolean | null; /** Fallbacks */ diff --git a/tests/integration/routing/test_team_model_preference_steering.py b/tests/integration/routing/test_fallback_on_saturation.py similarity index 57% rename from tests/integration/routing/test_team_model_preference_steering.py rename to tests/integration/routing/test_fallback_on_saturation.py index c99cfcf660b..edfbbac22e3 100644 --- a/tests/integration/routing/test_team_model_preference_steering.py +++ b/tests/integration/routing/test_fallback_on_saturation.py @@ -13,9 +13,10 @@ from integration._support.wire import Reply, Request, wire_server OPENAI_MODEL: Final = "gpt-5.6" PREFERRED_RPM: Final = 2 -ALTERNATE_RPM: Final = 1000 +FALLBACK_RPM: Final = 1000 SATURATION_THRESHOLD: Final = 0.5 -STEERED_FROM_HEADER: Final = "x-litellm-model-steered-from" +FALLBACK_REASON_HEADER: Final = "x-litellm-fallback-reason" +SATURATION_REASON: Final = "rate_limit:model_saturation_fallback" def _reply() -> bytes: @@ -25,15 +26,15 @@ def _reply() -> bytes: "object": "chat.completion", "created": 1700000000, "model": OPENAI_MODEL, - "choices": [{"index": 0, "message": {"role": "assistant", "content": "steered"}, "finish_reason": "stop"}], + "choices": [{"index": 0, "message": {"role": "assistant", "content": "served"}, "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") +def saturation_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[Gateway, str]]: + directory: Final = tmp_path_factory.mktemp("fallback-on-saturation") base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) configuration: Final = { **base, @@ -46,7 +47,7 @@ def steering_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[tuple[G }, }, } - path: Final = directory / "model_preference_steering.yaml" + path: Final = directory / "fallback_on_saturation.yaml" path.write_text(yaml.safe_dump(configuration)) def respond(request: Request) -> Reply: @@ -73,7 +74,7 @@ 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}"}]}, + {"model": model, "messages": [{"role": "user", "content": f"fallback on saturation {uuid.uuid4().hex}"}]}, key=key, ) @@ -84,61 +85,69 @@ def _spend_row(owned: Gateway, response: httpx.Response) -> list[dict[str, objec ) -def test_flexible_team_is_steered_to_its_permitted_alternate_once_the_preferred_model_saturates( - steering_proxy: tuple[Gateway, str], +def test_a_team_with_fallback_on_saturation_is_served_by_its_fallback_once_the_model_saturates( + saturation_proxy: tuple[Gateway, str], ) -> None: - owned, upstream = steering_proxy + owned, upstream = saturation_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])) + fallback, fallback_id = _deployment(owned, scenario, upstream, FALLBACK_RPM) + rigid_key: Final = scenario.key(team_id=scenario.team(models=[preferred, fallback])) flexible_key: Final = scenario.key( team_id=scenario.team( - models=[preferred, alternate], - metadata={"model_preferences": [{"model": preferred, "alternates": [alternate]}]}, + models=[preferred, fallback], + router_settings={"fallbacks": [{preferred: [fallback]}], "fallback_on_saturation": True}, ) ) 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 + assert FALLBACK_REASON_HEADER not in warmup.headers - steered: Final = eventually( + fell_back: Final = eventually( lambda: _chat(owned, flexible_key, preferred), - lambda response: STEERED_FROM_HEADER in response.headers, + lambda response: response.headers.get(FALLBACK_REASON_HEADER) == SATURATION_REASON, 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 + assert fell_back.status_code == 200, fell_back.text + assert fell_back.headers[FALLBACK_REASON_HEADER] == SATURATION_REASON + assert fell_back.headers["x-litellm-attempted-fallbacks"] == "1" + assert fell_back.headers["x-litellm-model-id"] == fallback_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 + assert FALLBACK_REASON_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 + charged: Final = eventually(lambda: _spend_row(owned, fell_back), lambda rows: len(rows) == 1, seconds=30) + assert charged[0]["model_group"] == fallback -def test_alternate_the_team_may_not_call_is_never_used(steering_proxy: tuple[Gateway, str]) -> None: - owned, upstream = steering_proxy +def test_a_fallback_the_team_may_not_call_is_never_used(saturation_proxy: tuple[Gateway, str]) -> None: + owned, upstream = saturation_proxy with owned.scenario() as scenario: - preferred, _ = _deployment(owned, scenario, upstream, PREFERRED_RPM) - forbidden, _ = _deployment(owned, scenario, upstream, ALTERNATE_RPM) + preferred, preferred_id = _deployment(owned, scenario, upstream, PREFERRED_RPM) + forbidden, _ = _deployment(owned, scenario, upstream, FALLBACK_RPM) key: Final = scenario.key( team_id=scenario.team( models=[preferred], - metadata={"model_preferences": [{"model": preferred, "alternates": [forbidden]}]}, + router_settings={"fallbacks": [{preferred: [forbidden]}], "fallback_on_saturation": True}, ) ) - 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 + def served_by_preferred_or_refused() -> httpx.Response: + response: Final = _chat(owned, key, preferred) + assert FALLBACK_REASON_HEADER not in response.headers, response.headers + if response.status_code == 200: + assert response.headers["x-litellm-model-id"] == preferred_id + return response + + refused: Final = eventually( + served_by_preferred_or_refused, + lambda response: response.status_code != 200, + seconds=15, + return_last_on_timeout=True, ) - assert second.status_code == 429, second.text - assert STEERED_FROM_HEADER not in second.headers + assert refused.status_code == 429, refused.text diff --git a/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py b/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py index a7628b72341..e1978b84e96 100644 --- a/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py +++ b/tests/unit/proxy/common_utils/test_proxy_rate_limit_error.py @@ -3,9 +3,18 @@ from typing import Final import pytest from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, RateLimitFallback, + arm_saturation_fallback_walk, + disarm_saturation_fallback_walk, + fallback_on_saturation_enabled, + fallback_on_saturation_requested, + is_saturation_fallback, + mark_saturation_fallbacks_exhausted, rate_limit_fallback_from_request_data, rate_limit_fallback_headers, + saturation_fallbacks_exhausted, ) STAMP: Final = {"original_model_group": "opus", "descriptor_key": "model_per_end_user"} @@ -56,3 +65,64 @@ def test_rate_limit_fallback_headers_count_every_candidate_a_limit_rejected( ) assert headers["x-litellm-attempted-fallbacks"] == expected + + +@pytest.mark.parametrize( + ("router_settings", "expected"), + ( + ({"fallback_on_saturation": True}, True), + ({"fallback_on_saturation": False}, False), + ({"fallback_on_saturation": "yes"}, False), + ({"fallbacks": [{"luna": ["opus"]}]}, False), + (None, False), + ("fallback_on_saturation", False), + ), +) +def test_fallback_on_saturation_is_enabled_only_by_a_true_router_setting( + router_settings: object, expected: bool +) -> None: + assert fallback_on_saturation_enabled(router_settings) is expected + + +def test_a_request_asks_for_saturation_fallbacks_only_while_the_walk_is_armed_and_until_it_has_run() -> None: + request: Final = {"model": "luna", "router_settings_override": {"fallback_on_saturation": True}} + assert not fallback_on_saturation_requested(request) + + walk: Final = arm_saturation_fallback_walk() + try: + assert fallback_on_saturation_requested(request) + assert not saturation_fallbacks_exhausted() + mark_saturation_fallbacks_exhausted() + assert saturation_fallbacks_exhausted() + assert not fallback_on_saturation_requested(request) + finally: + disarm_saturation_fallback_walk(walk) + assert not saturation_fallbacks_exhausted() + assert not fallback_on_saturation_requested(request) + + +def test_client_metadata_never_arms_skips_or_exhausts_the_walk() -> None: + forged: Final = { + "model": "luna", + "router_settings_override": {"fallback_on_saturation": True}, + "metadata": {"saturation_fallback_walk": True, "saturation_fallbacks_exhausted": True}, + "litellm_metadata": {"saturation_fallback_walk": True, "saturation_fallbacks_exhausted": True}, + } + assert not fallback_on_saturation_requested(forged) + assert not saturation_fallbacks_exhausted() + walk: Final = arm_saturation_fallback_walk() + try: + assert fallback_on_saturation_requested(forged) + finally: + disarm_saturation_fallback_walk(walk) + + +def test_a_request_without_the_setting_never_asks_for_saturation_fallbacks() -> None: + assert not fallback_on_saturation_requested({"model": "luna"}) + assert not fallback_on_saturation_requested({"model": "luna", "router_settings_override": {"fallbacks": []}}) + + +def test_only_the_saturation_descriptor_marks_a_fallback_walk_the_model_can_still_serve() -> None: + assert is_saturation_fallback(ProxyRateLimitError(detail="x", descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY)) + assert not is_saturation_fallback(ProxyRateLimitError(detail="x", descriptor_key="model_per_team")) + assert not is_saturation_fallback(ProxyRateLimitError(detail="x")) 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 1c71e918cf5..87c88ae98f6 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -6,9 +6,8 @@ Core tests to validate that priority weights are respected (0.9/0.1) instead of import asyncio import time -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime, timedelta -from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, patch @@ -17,11 +16,9 @@ from prometheus_client import CollectorRegistry 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, _PROXY_DynamicRateLimitHandlerV3, # pyright: ignore[reportPrivateUsage] # the public name is a Final alias, not a type ) @@ -1837,317 +1834,92 @@ 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 _saturation_fallback_handler(monkeypatch: pytest.MonkeyPatch) -> _PROXY_DynamicRateLimitHandlerV3: + return _borrowing_handler(monkeypatch, borrows=False, saturation_threshold=0.5, rpm=10) -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} - ) +def _interactive_key(api_key: str, team_id: str) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key=api_key, team_id=team_id, metadata={"priority": "interactive"}) -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 -@pytest.mark.parametrize("team_allowed_models, granted", [(("luna", "opus"), True), (("luna",), False)]) -async def test_an_alternate_is_granted_only_when_every_customer_level_permits_it( - monkeypatch: pytest.MonkeyPatch, team_allowed_models: tuple[str, ...], granted: bool -): - from litellm.proxy import proxy_server - from litellm.proxy.auth import auth_checks, user_api_key_auth - from litellm.proxy.hooks.dynamic_rate_limiter_v3 import alternate_grant_for - from litellm.types.proxy.auth.customer_hierarchy import CustomerHierarchyAuth, CustomerScopeAuth - - monkeypatch.setattr(proxy_server, "llm_router", _steering_router("luna", "opus")) - monkeypatch.setattr(proxy_server, "prisma_client", None) - monkeypatch.setattr(user_api_key_auth, "_is_fallback_model_authorized", AsyncMock(return_value=True)) - monkeypatch.setattr(user_api_key_auth, "_is_within_every_model_budget", AsyncMock(return_value=True)) - monkeypatch.setattr(auth_checks, "collect_matched_model_access_groups", AsyncMock(return_value=())) - monkeypatch.setattr(auth_checks, "_model_access_group_max_budget_check", AsyncMock(return_value=None)) - customer: Final = UserAPIKeyAuth( - customer_hierarchy=CustomerHierarchyAuth( - customer_member_id="member-1", - customer_team_id="team-1", - customer_org_id="org-1", - scopes=( - CustomerScopeAuth(key="customer_org", entity_id="org-1"), - CustomerScopeAuth(key="customer_team", entity_id="team-1", allowed_models=team_allowed_models), - CustomerScopeAuth(key="customer_member", entity_id="member-1", allowed_models=("luna", "opus")), - ), - ) - ) - - assert (await alternate_grant_for(customer, "opus") is not None) is granted - - -@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}, +def _opted_in(model: str = "shared-model", **extra: object) -> dict[str, object]: + return { + "model": model, + "router_settings_override": {"fallback_on_saturation": True}, + **extra, } - 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.fixture +def armed_walk() -> Iterator[None]: + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + arm_saturation_fallback_walk, + disarm_saturation_fallback_walk, ) - -@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 + walk: Final = arm_saturation_fallback_walk() + yield + disarm_saturation_fallback_walk(walk) @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 test_an_opted_in_key_on_a_route_without_the_fallback_walk_is_admitted_on_the_saturated_model( + monkeypatch: pytest.MonkeyPatch, +): + handler = _saturation_fallback_handler(monkeypatch) + key: Final = _interactive_key("sk-flexible", "team-a") + await _saturate_to_threshold(handler, key) + direct: Final = {"model": "shared-model", "router_settings_override": {"fallback_on_saturation": True}} - async def broken(user_api_key_dict: UserAPIKeyAuth, model: str) -> AlternateGrant | None: - raise RuntimeError("membership lookup is down") + assert await handler.async_pre_call_hook(key, DualCache(), direct, "acompletion") is None - 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 - ) +async def _saturate_to_threshold(handler: _PROXY_DynamicRateLimitHandlerV3, key: UserAPIKeyAuth) -> None: + for _ in range(5): + assert await handler.async_pre_call_hook(key, DualCache(), {"model": "shared-model"}, "acompletion") is None + group: Final = handler.llm_router.registry().group("shared-model") + assert group is not None + assert await handler._check_model_saturation("shared-model", group) >= 0.5 + + +@pytest.mark.asyncio +async def test_a_key_with_fallback_on_saturation_is_sent_to_its_fallbacks_once_the_model_saturates( + monkeypatch: pytest.MonkeyPatch, armed_walk: None +): + from litellm.proxy.common_utils.proxy_rate_limit_error import SATURATION_FALLBACK_DESCRIPTOR_KEY + + handler = _saturation_fallback_handler(monkeypatch) + key: Final = _interactive_key("sk-flexible", "team-a") + assert await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion") is None + await _saturate_to_threshold(handler, key) + + with pytest.raises(ProxyRateLimitError) as raised: + await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion") + assert raised.value.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY + assert "trying fallbacks first" in str(raised.value.detail) + + +@pytest.mark.asyncio +async def test_a_key_without_fallback_on_saturation_is_admitted_on_the_saturated_model(monkeypatch: pytest.MonkeyPatch): + handler = _saturation_fallback_handler(monkeypatch) + key: Final = _interactive_key("sk-rigid", "team-b") + await _saturate_to_threshold(handler, key) + + assert await handler.async_pre_call_hook(key, DualCache(), {"model": "shared-model"}, "acompletion") is None + + +@pytest.mark.asyncio +async def test_a_request_back_from_an_exhausted_fallback_walk_is_admitted_on_the_saturated_model( + monkeypatch: pytest.MonkeyPatch, armed_walk: None +): + from litellm.proxy.common_utils.proxy_rate_limit_error import mark_saturation_fallbacks_exhausted + + handler = _saturation_fallback_handler(monkeypatch) + key: Final = _interactive_key("sk-flexible", "team-a") + await _saturate_to_threshold(handler, key) + + mark_saturation_fallbacks_exhausted() + assert await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion") is None def _borrowing_handler( @@ -2586,6 +2358,23 @@ async def test_a_reservation_at_saturation_is_refunded_when_the_call_fails(monke assert await _team_tokens_remaining(handler) == 100 +@pytest.mark.asyncio +async def test_a_reservation_at_saturation_is_refunded_when_the_proxy_refuses_the_request_before_the_call( + monkeypatch: pytest.MonkeyPatch, +): + """A guardrail or the fallback walk turning the request away after admission gives the team its tokens back, and + a second release changes nothing.""" + handler = _borrowing_handler(monkeypatch, borrows=True) + key: Final = _team_key() + await _check(handler, key, saturation=0.5, data=_prompt(words=30)) + assert await _team_tokens_remaining(handler) < 70 + for _ in range(2): + await handler.async_post_call_failure_hook( + request_data=_prompt(words=30), original_exception=Exception("guardrail refused"), user_api_key_dict=key + ) + assert await _team_tokens_remaining(handler) == 100 + + @pytest.mark.asyncio async def test_a_fallback_target_is_refused_once_it_saturates_and_the_teams_share_on_it_is_used( monkeypatch: pytest.MonkeyPatch, @@ -2806,7 +2595,7 @@ async def test_a_served_fallbacks_hold_is_not_refunded_again_under_the_name_the_ @pytest.mark.asyncio -async def test_a_steered_requests_images_are_counted_for_the_model_that_serves_it(monkeypatch: pytest.MonkeyPatch): +async def test_a_request_served_by_another_model_counts_its_images_for_that_model(monkeypatch: pytest.MonkeyPatch): from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) @@ -3710,40 +3499,20 @@ async def test_a_class_share_refusal_is_counted_under_the_teams_priority(monkeyp @pytest.mark.asyncio -async def test_an_alternates_refusal_is_not_counted_while_the_preferred_model_admits(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "priority_reservation", None) - monkeypatch.setattr(litellm.priority_reservation_settings, "saturation_threshold", 0.5) - monkeypatch.setattr( - litellm.priority_reservation_settings, "team_model_limit_borrows_idle_capacity_from_other_teams", True - ) - registry = _prometheus_sink(monkeypatch) - handler: Final[_PROXY_DynamicRateLimitHandlerV3] = _steering_handler() - router: Final = _steering_router("luna", "opus", preferred_rpm=4) - monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) - handler.update_variables(llm_router=router) - flexible: Final = UserAPIKeyAuth( - team_id="flexible-team", - team_models=["luna", "opus"], - team_metadata={ - "model_preferences": [{"model": "luna", "alternates": ["opus"]}], - "model_rpm_limit": {"opus": 1, "luna": 1}, - "team_model_limit_max_borrow_multiple": 1, - }, - ) - await _saturate(handler, "luna", requests=2) - assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") == { - "model": "opus" - } - assert await handler.async_pre_call_hook(flexible, DualCache(), {"model": "luna"}, "completion") is None - assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team") is None +async def test_a_saturation_fallback_handoff_is_not_counted_as_a_throttle( + monkeypatch: pytest.MonkeyPatch, armed_walk: None +): + from litellm.proxy.common_utils.proxy_rate_limit_error import SATURATION_FALLBACK_DESCRIPTOR_KEY - refusal: Final = await _refused_at_the_hook(handler, flexible, model="luna") - assert refusal.headers is not None - assert refusal.headers["rate_limited_by"] == "model_per_team_saturated" - assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team", model_group="luna") == 1 - assert ( - _sample(registry, "litellm_dynamic_limiter_throttled_total", team="flexible-team", model_group="opus") is None - ) + handler = _saturation_fallback_handler(monkeypatch) + registry = _prometheus_sink(monkeypatch) + key: Final = _interactive_key("sk-flexible", "team-a") + await _saturate_to_threshold(handler, key) + + with pytest.raises(ProxyRateLimitError) as handed_off: + await handler.async_pre_call_hook(key, DualCache(), _opted_in(), "acompletion") + assert handed_off.value.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY + assert _sample(registry, "litellm_dynamic_limiter_throttled_total", team="team-a") is None @pytest.mark.asyncio diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index bbc77b8645b..1959f81cef8 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -2,16 +2,14 @@ import asyncio import contextvars import copy import datetime -import time import gzip import json +import time from collections.abc import AsyncIterable, Awaitable, Coroutine, Mapping from types import MappingProxyType, SimpleNamespace from typing import AsyncGenerator, Callable, Final, Iterator, Literal, Optional, Sequence, TypeVar -from urllib.parse import unquote_plus from unittest.mock import AsyncMock, MagicMock, patch - -from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs +from urllib.parse import unquote_plus import anthropic import httpx @@ -23,29 +21,39 @@ from fastapi.responses import JSONResponse, StreamingResponse import litellm from litellm._uuid import uuid -from litellm.exceptions import APIError, InternalServerError, MidStreamFallbackError -from litellm.litellm_core_utils.litellm_logging import Logging from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame from litellm.caching.caching import DualCache -from litellm.litellm_core_utils.bug_report import ( - DISABLE_ENV_VAR, - ISSUE_URL_BASE, - bug_report_notice, - build_bug_report, -) -from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, RETURN_RAW_MODEL_NAME_METADATA_KEY, STREAM_SSE_KEEPALIVE_PING_BYTES, ) -from litellm.integrations.custom_logger import CustomLogger +from litellm.exceptions import APIError, InternalServerError, MidStreamFallbackError from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.opentelemetry import UserAPIKeyAuth -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.litellm_core_utils.add_retry_fallback_headers import prepare_response_for_header_attachment +from litellm.litellm_core_utils.bug_report import ( + DISABLE_ENV_VAR, + ISSUE_URL_BASE, + bug_report_notice, + build_bug_report, +) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.litellm_core_utils.total_deadline import current_total_deadline +from litellm.llms.anthropic.chat.transformation import AnthropicConfig +from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException +from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth +from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck from litellm.proxy.common_request_processing import ( + OPENAI_FIRST_CHUNK_FRAMING, + RESPONSES_FIRST_CHUNK_FRAMING, + RESPONSES_LATE_SSE_FRAMING, + CostBreakdownHeaderValues, + FirstChunkSseFraming, ProxyBaseLLMRequestProcessing, ProxyConfig, _await_llm_call_cancelling_on_disconnect, @@ -53,39 +61,29 @@ from litellm.proxy.common_request_processing import ( _buffer_first_chunk_honoring_disconnect, _cancel_llm_call_on_client_disconnect, _ClientDisconnectedBeforeFirstChunk, - attach_guardrail_information, _extract_error_from_sse_chunk, - FirstChunkSseFraming, _get_cost_breakdown_from_logging_obj, - CostBreakdownHeaderValues, _has_attribute_error_in_chain, - OPENAI_FIRST_CHUNK_FRAMING, - RESPONSES_FIRST_CHUNK_FRAMING, - include_guardrail_response_requested, - is_azure_model_router_request, - RESPONSES_LATE_SSE_FRAMING, - open_sse_before_first_byte, - resolve_litellm_call_id, - ttft_keepalive_interval, _override_openai_response_model, _parse_event_data_for_error, _resolve_per_request_model_group_alias, - should_return_raw_model_name, _sse_error_frames, _UpstreamClosingStreamingResponse, + attach_guardrail_information, create_response, + include_guardrail_response_requested, + is_azure_model_router_request, + open_sse_before_first_byte, + resolve_litellm_call_id, + should_return_raw_model_name, sse_error_payload, + ttft_keepalive_interval, ) from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.dd_span_tagger import DDSpanTagger -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.llms.anthropic.chat.transformation import AnthropicConfig -from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException -from litellm.litellm_core_utils.total_deadline import current_total_deadline -from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth -from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, model_budget_spend_cache_key, @@ -94,10 +92,11 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.proxy.utils import ProxyLogging from litellm.responses.mcp.parked_stream_dispatch import IN_FLIGHT_ROUND from litellm.router import Router -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic +from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices from litellm.utils import ModelResponseListIterator -from litellm.litellm_core_utils.add_retry_fallback_headers import prepare_response_for_header_attachment from tests._support.clock_loops import AutojumpClockLoop +from tests.unit.proxy.router_entry_point_stubs import route_entry_points_to_stubs def test_attach_guardrail_information_copies_recorded_entries_onto_model_response(): @@ -2632,10 +2631,11 @@ class TestCommonRequestProcessingHelpers: async def test_serialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" + import json as _json + from litellm.proxy.common_request_processing import ( serialize_http_exception_detail, ) - import json as _json assert serialize_http_exception_detail("plain") == ("plain", None) @@ -6569,6 +6569,7 @@ class TestAllmPassthroughRoutePostCallGuardrails: def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: import json import struct + from botocore.eventstream import crc32 as esm_crc32 payload_bytes = json.dumps(payload, separators=(",", ":")).encode() @@ -8064,11 +8065,22 @@ class TestCostHeadersForCallsPricedAtZero: assert breakdown.output_cost == 0.0 +def _proxy_config(router_settings: Mapping[str, object] | None = None) -> MagicMock: + """A proxy config whose Key > Team router_settings read resolves to ``router_settings``.""" + return MagicMock(get_hierarchical_router_settings=AsyncMock(return_value=router_settings)) + + class TestPreCallWithFallbacksOnLocalRateLimit: + @pytest.fixture(autouse=True) + def _every_key_may_call_every_fallback(self, monkeypatch): + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "enforce_fallback_model_access", False) + @pytest.mark.asyncio async def test_fallback_triggered_on_local_rate_limit(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError primary_model = "gpt-4" fallback_model = "gpt-3.5-turbo" @@ -8101,9 +8113,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: request=MagicMock(), general_settings={}, proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8157,7 +8169,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: proxy_logging_obj=MagicMock(), user_api_key_dict=ProxyUserAPIKeyAuth(), version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8172,8 +8184,8 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @pytest.mark.asyncio async def test_raises_when_no_fallbacks_configured(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) @@ -8196,9 +8208,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: request=MagicMock(), general_settings={}, proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8211,8 +8223,8 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @pytest.mark.asyncio async def test_raises_when_all_fallbacks_also_rate_limited(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) @@ -8235,9 +8247,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: request=MagicMock(), general_settings={}, proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8252,8 +8264,8 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @pytest.mark.asyncio async def test_fallback_uses_key_level_router_settings(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) @@ -8268,7 +8280,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: mock_router = route_entry_points_to_stubs(MagicMock()) mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] - user_api_key_dict = MagicMock() + user_api_key_dict = MagicMock(agent_id=None, end_user_id=None) user_api_key_dict.router_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} with patch.object( @@ -8282,7 +8294,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: proxy_logging_obj=MagicMock(), user_api_key_dict=user_api_key_dict, version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(user_api_key_dict.router_settings), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8296,9 +8308,1045 @@ class TestPreCallWithFallbacksOnLocalRateLimit: assert processor.data["model"] == "claude-3-haiku" @pytest.mark.asyncio - async def test_disable_fallbacks_flag_respected(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + async def test_fallback_uses_team_level_router_settings_when_the_key_has_none(self): from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + team_settings = {"fallbacks": [{"gpt-4": ["claude-3-haiku"]}]} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config(team_settings), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_a_fallback_the_key_may_not_call_is_skipped(self, monkeypatch): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def only_haiku(*, model, valid_token, llm_router): + return model == "claude-3-haiku" + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setitem(proxy_server.general_settings, "enforce_fallback_model_access", True) + monkeypatch.setattr("litellm.proxy.auth.fallback_model_access.is_model_authorized_for_token", only_haiku) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_an_agent_owned_key_never_leaves_the_model_it_asked_for(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["claude-3-haiku"]}] + + with ( + patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic), + pytest.raises(ProxyRateLimitError), + ): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id="agent-1"), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + @pytest.mark.asyncio + async def test_a_fallback_the_end_user_may_not_call_is_skipped(self, monkeypatch): + from litellm.proxy._types import LiteLLM_EndUserTable + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def customer(**kwargs): + return LiteLLM_EndUserTable(user_id="cust-1", models=["gpt-3.5-turbo"], blocked=False) + + monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_end_user_object", customer) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock()) + router = litellm.Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}} + for name in ("gpt-4", "claude-3-haiku", "gpt-3.5-turbo") + ], + fallbacks=[{"gpt-4": ["claude-3-haiku", "gpt-3.5-turbo"]}], + ) + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(end_user_id="cust-1"), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=router, + ) + + assert processor.data["model"] == "gpt-3.5-turbo" + + @pytest.mark.asyncio + async def test_an_end_user_without_a_database_keeps_its_fallbacks(self, monkeypatch): + from litellm.proxy import proxy_server + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + monkeypatch.setattr(proxy_server, "prisma_client", None) + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "user": "cust-1"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(end_user_id="cust-1"), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_the_rate_limited_models_hold_and_a_refused_fallbacks_hold_are_released_before_the_walk_moves_on( + self, monkeypatch + ): + """A key allowed one request at a time is admitted on the model the walk ends on, so the slots the limiter + held for the saturated model and for the refused fallback must have come back. Afterwards only the served + request holds the key's slot: one release frees it.""" + from fastapi import HTTPException + + from litellm.caching.caching import DualCache + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + saturation_fallbacks_exhausted, + ) + from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 + from litellm.proxy.utils import InternalUsageCache, hash_token + + cache = DualCache() + limiter = PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=InternalUsageCache(cache)) + key = ProxyKey(api_key=hash_token("sk-one-at-a-time"), max_parallel_requests=1) + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + refusal = HTTPException(status_code=400, detail="guardrail refused the fallback") + + async def admit_then_decide(**kwargs): + await limiter.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data=processor.data, call_type="acompletion" + ) + if processor.data.get("model") == "gpt-4" and not saturation_fallbacks_exhausted(): + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + if processor.data.get("model") == "gpt-3.5-turbo": + raise refusal + return processor.data, MagicMock() + + async def every_model(*, model, valid_token, llm_router): + return True + + monkeypatch.setattr("litellm.proxy.auth.fallback_model_access.is_model_authorized_for_token", every_model) + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook.return_value = limiter + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [] + settings = {"fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}], "fallback_on_saturation": True} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=admit_then_decide): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=proxy_logging_obj, + user_api_key_dict=key, + version=None, + proxy_config=_proxy_config(settings), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "gpt-4" + with pytest.raises(HTTPException): + await limiter.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data={"model": "gpt-4"}, call_type="acompletion" + ) + await limiter.async_post_call_failure_hook( + request_data=data, original_exception=Exception("client went away"), user_api_key_dict=key + ) + await limiter.async_pre_call_hook( + user_api_key_dict=key, cache=cache, data={"model": "gpt-4"}, call_type="acompletion" + ) + + @pytest.mark.asyncio + async def test_an_exhausted_saturation_walk_hands_the_request_back_to_the_model_its_first_attempt_resolved( + self, + ): + """With aliases gpt-4 -> gpt-4-pool and gpt-4-pool -> gpt-4-overflow, the request the walk hands back is + served by gpt-4-pool as the first attempt was, not aliased a second time on to gpt-4-overflow.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + saturation_fallbacks_exhausted, + ) + + aliases = {"gpt-4": "gpt-4-pool", "gpt-4-pool": "gpt-4-overflow"} + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def alias_then_saturate(*, model, **kwargs): + processor.data["model"] = aliases.get(model, model) + if not saturation_fallbacks_exhausted(): + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [] + settings = {"fallbacks": [{"gpt-4-pool": ["gpt-3.5-turbo"]}], "fallback_on_saturation": True} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=alias_then_saturate): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config(settings), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "gpt-4-pool" + + @pytest.mark.asyncio + async def test_a_fallback_over_its_model_budget_is_skipped(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def cheap_model_budget_left(*, valid_token, model_max_budget_limiter, model): + return model != "gpt-3.5-turbo" + + monkeypatch.setattr( + "litellm.proxy.auth.user_api_key_auth._is_within_every_model_budget", cheap_model_budget_left + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_a_fallback_aliased_to_a_model_over_its_budget_is_skipped(self, monkeypatch): + """The walk grants the group an alias is served by, so ``fallbacks: [{A: [B]}]`` with + ``model_group_alias: {B: C}`` checks C's budget as well as B's, as the auth-time fallback walk does.""" + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4", "router_settings_override": {"model_group_alias": {"gpt-3.5-turbo": "pricey"}}} + ) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + budget_checked: list[str] = [] + + async def served_model_over_budget(*, valid_token, model_max_budget_limiter, model): + budget_checked.append(model) + return model != "pricey" + + monkeypatch.setattr( + "litellm.proxy.auth.user_api_key_auth._is_within_every_model_budget", served_model_over_budget + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == "claude-3-haiku" + assert "pricey" in budget_checked + + @pytest.mark.asyncio + async def test_a_fallback_whose_listed_name_is_over_budget_is_skipped_even_when_its_alias_target_has_room( + self, monkeypatch + ): + """A per-model budget set on the fallback list's name still applies when an alias serves that name through + another group, as it does at auth.""" + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4", "router_settings_override": {"model_group_alias": {"gpt-3.5-turbo": "roomy"}}} + ) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def listed_name_over_budget(*, valid_token, model_max_budget_limiter, model): + return model != "gpt-3.5-turbo" + + monkeypatch.setattr( + "litellm.proxy.auth.user_api_key_auth._is_within_every_model_budget", listed_name_over_budget + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert processor.data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("unpriced_models_blocked", "served_by"), [(True, "claude-3-haiku"), (False, "gpt-3.5-turbo")] + ) + async def test_an_unpriced_fallback_is_skipped_only_while_unpriced_models_are_blocked( + self, monkeypatch, unpriced_models_blocked, served_by + ): + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", unpriced_models_blocked) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.model_has_no_cost_mapping", + lambda *, model, llm_router: model == "gpt-3.5-turbo", + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == served_by + + @pytest.mark.asyncio + async def test_a_fallback_priced_above_the_reserved_estimate_is_skipped(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.types.router import ModelGroupInfo + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + prices = {"gpt-4": (1e-5, 3e-5), "gpt-3.5-turbo": (1e-5, 4e-5), "claude-3-haiku": (1e-6, 1e-6)} + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + def priced_group(model_group, **kwargs): + input_cost, output_cost = prices[model_group] + return ModelGroupInfo( + model_group=model_group, + providers=["openai"], + input_cost_per_token=input_cost, + output_cost_per_token=output_cost, + ) + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + mock_router.registry.return_value.group = priced_group + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(budget_reservation={"entries": []}), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_a_fallback_outside_the_reserved_access_groups_is_skipped(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.types.router import ModelGroupInfo + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def only_haiku_shares_the_reserved_group(*, model, **kwargs): + return ("primary-group",) if model == "claude-3-haiku" else ("other-group",) + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.collect_matched_model_access_groups", only_haiku_shares_the_reserved_group + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + mock_router.registry.return_value.group = lambda model_group, **kwargs: ModelGroupInfo( + model_group=model_group, providers=["openai"], input_cost_per_token=1e-6, output_cost_per_token=1e-6 + ) + reserved_against_primary = {"entries": [{"entity_type": "Model access group", "entity_id": "primary-group"}]} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(budget_reservation=reserved_against_primary), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "claude-3-haiku" + + @pytest.mark.asyncio + async def test_a_fallback_in_a_group_the_request_already_holds_is_not_refused_by_its_own_hold(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + from litellm.types.router import ModelGroupInfo + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def every_model_in_the_shared_group(*, model, **kwargs): + return ("shared-group",) + + async def the_hold_fills_the_shared_group(*, matched_model_access_groups, **kwargs): + if "shared-group" in matched_model_access_groups: + raise RuntimeError("shared-group spend plus the request's own hold reached max_budget") + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.collect_matched_model_access_groups", every_model_in_the_shared_group + ) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks._model_access_group_max_budget_check", the_hold_fills_the_shared_group + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + mock_router.registry.return_value.group = lambda model_group, **kwargs: ModelGroupInfo( + model_group=model_group, providers=["openai"], input_cost_per_token=1e-6, output_cost_per_token=1e-6 + ) + held_at_auth = {"entries": [{"entity_type": "Model access group", "entity_id": "shared-group"}]} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=ProxyKey(budget_reservation=held_at_auth), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "gpt-3.5-turbo" + + @pytest.mark.asyncio + async def test_a_served_fallback_charges_the_access_groups_that_grant_it(self, monkeypatch): + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + key: Final = ProxyKey(matched_model_access_groups=["primary-group"]) + charged_when_served = [] + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + charged_when_served.append(list(key.matched_model_access_groups or [])) + return processor.data, MagicMock() + + async def backup_group_grants_the_fallback(*, model, **kwargs): + return ("backup-group",) if model == "gpt-3.5-turbo" else () + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.collect_matched_model_access_groups", backup_group_grants_the_fallback + ) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=key, + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "gpt-3.5-turbo" + assert charged_when_served == [["backup-group"]] + + @pytest.mark.asyncio + async def test_a_fallback_over_its_access_group_budget_is_skipped(self, monkeypatch): + from litellm.exceptions import BudgetExceededError + from litellm.proxy._types import UserAPIKeyAuth as ProxyKey + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + key: Final = ProxyKey(matched_model_access_groups=["primary-group"]) + + async def mock_pre_call_logic(**kwargs): + if processor.data.get("model") == "gpt-4": + raise ProxyRateLimitError(detail="TPM limit exceeded", headers={"retry-after": "30"}) + return processor.data, MagicMock() + + async def one_group_per_fallback(*, model, **kwargs): + return (f"{model}-group",) + + async def haiku_group_is_spent(*, matched_model_access_groups, **kwargs): + if "gpt-3.5-turbo-group" in matched_model_access_groups: + raise BudgetExceededError(current_cost=2.0, max_budget=1.0) + + import litellm.proxy.proxy_server as proxy_server + + monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + monkeypatch.setattr( + "litellm.proxy.auth.auth_checks.collect_matched_model_access_groups", one_group_per_fallback + ) + monkeypatch.setattr("litellm.proxy.auth.auth_checks._model_access_group_max_budget_check", haiku_group_is_spent) + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=key, + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert data["model"] == "claude-3-haiku" + assert key.matched_model_access_groups == ["claude-3-haiku-group"] + + @pytest.mark.asyncio + async def test_a_client_body_cannot_opt_into_saturation_fallbacks(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import fallback_on_saturation_requested + + processor = ProxyBaseLLMRequestProcessing( + data={"model": "gpt-4", "router_settings_override": {"fallback_on_saturation": True, "num_retries": 1}} + ) + seen_by_the_limiter = [] + + async def mock_pre_call_logic(**kwargs): + seen_by_the_limiter.append(fallback_on_saturation_requested(processor.data)) + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config(), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert seen_by_the_limiter == [False] + assert data["router_settings_override"] == {"num_retries": 1} + + @pytest.mark.asyncio + async def test_fallback_on_saturation_moves_past_a_fallback_that_refuses_the_request(self): + from fastapi import HTTPException + + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + saturation_fallbacks_exhausted, + ) + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + attempted = [] + + async def mock_pre_call_logic(**kwargs): + attempted.append(processor.data.get("model")) + if processor.data.get("model") == "gpt-4" and not saturation_fallbacks_exhausted(): + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + if processor.data.get("model") == "gpt-3.5-turbo": + raise HTTPException(status_code=400, detail="guardrail refused the fallback") + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [] + settings = {"fallbacks": [{"gpt-4": ["gpt-3.5-turbo"]}], "fallback_on_saturation": True} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config(settings), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert attempted == ["gpt-4", "gpt-3.5-turbo", "gpt-4"] + assert data["model"] == "gpt-4" + + @pytest.mark.asyncio + async def test_fallback_on_saturation_serves_the_first_fallback_with_room(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + rate_limit_fallback_from_request_data, + ) + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + attempted = [] + + async def mock_pre_call_logic(**kwargs): + attempted.append(processor.data.get("model")) + if processor.data.get("model") in ("gpt-4", "gpt-3.5-turbo"): + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + return processor.data, MagicMock() + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [] + settings = {"fallbacks": [{"gpt-4": ["gpt-3.5-turbo", "claude-3-haiku"]}], "fallback_on_saturation": True} + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config(settings), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert attempted == ["gpt-4", "gpt-3.5-turbo", "claude-3-haiku"] + assert data["model"] == "claude-3-haiku" + stamp = rate_limit_fallback_from_request_data(data) + assert stamp is not None + assert stamp.descriptor_key == SATURATION_FALLBACK_DESCRIPTOR_KEY + assert stamp.original_model_group == "gpt-4" + assert stamp.attempts == 2 + + @pytest.mark.asyncio + async def test_fallback_on_saturation_serves_the_requested_model_when_no_fallback_has_room(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + rate_limit_fallback_from_request_data, + saturation_fallbacks_exhausted, + ) + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "metadata": {"tags": ["t"]}}) + attempted = [] + + async def mock_pre_call_logic(**kwargs): + attempted.append(processor.data.get("model")) + if saturation_fallbacks_exhausted(): + return processor.data, MagicMock() + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config({"fallback_on_saturation": True}), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert attempted == ["gpt-4", "gpt-3.5-turbo", "gpt-4"] + assert data["model"] == "gpt-4" + assert data["metadata"]["tags"] == ["t"] + assert rate_limit_fallback_from_request_data(data) is None + + @pytest.mark.asyncio + async def test_fallback_on_saturation_without_a_fallback_list_serves_the_requested_model(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + saturation_fallbacks_exhausted, + ) + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + attempted = [] + + async def mock_pre_call_logic(**kwargs): + attempted.append(processor.data.get("model")) + if saturation_fallbacks_exhausted(): + return processor.data, MagicMock() + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + data, _ = await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config({"fallback_on_saturation": True}), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + assert attempted == ["gpt-4", "gpt-4"] + assert data["model"] == "gpt-4" + + @pytest.mark.asyncio + async def test_a_real_429_on_the_requested_model_after_saturation_fallbacks_is_returned(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ( + SATURATION_FALLBACK_DESCRIPTOR_KEY, + ProxyRateLimitError, + saturation_fallbacks_exhausted, + ) + + processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4"}) + + async def mock_pre_call_logic(**kwargs): + if saturation_fallbacks_exhausted(): + raise ProxyRateLimitError(detail="gpt-4 is out of capacity", headers={"retry-after": "30"}) + raise ProxyRateLimitError( + detail="at saturation", headers={}, descriptor_key=SATURATION_FALLBACK_DESCRIPTOR_KEY + ) + + mock_router = route_entry_points_to_stubs(MagicMock()) + mock_router.fallbacks = [{"gpt-4": ["gpt-3.5-turbo"]}] + + with patch.object(processor, "common_processing_pre_call_logic", side_effect=mock_pre_call_logic): + with pytest.raises(ProxyRateLimitError, match="out of capacity"): + await processor._pre_call_with_fallbacks( + request=MagicMock(), + general_settings={}, + proxy_logging_obj=MagicMock(), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), + version=None, + proxy_config=_proxy_config({"fallback_on_saturation": True}), + user_model=None, + user_temperature=None, + user_request_timeout=None, + user_max_tokens=None, + user_api_base=None, + model="gpt-4", + route_type="acompletion", + llm_router=mock_router, + ) + + @pytest.mark.asyncio + async def test_disable_fallbacks_flag_respected(self): + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError processor = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4", "disable_fallbacks": True}) @@ -8321,9 +9369,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: request=MagicMock(), general_settings={}, proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8336,8 +9384,8 @@ class TestPreCallWithFallbacksOnLocalRateLimit: @pytest.mark.asyncio async def test_model_restored_on_non_rate_limit_exception(self): - from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError primary_model = "gpt-4" @@ -8365,9 +9413,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit: request=MagicMock(), general_settings={}, proxy_logging_obj=MagicMock(), - user_api_key_dict=MagicMock(router_settings=None), + user_api_key_dict=MagicMock(router_settings=None, agent_id=None, end_user_id=None), version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -8472,7 +9520,7 @@ class TestPreCallWithFallbacksOnLocalRateLimit: proxy_logging_obj=MagicMock(), user_api_key_dict=user_api_key_dict, version=None, - proxy_config=MagicMock(), + proxy_config=_proxy_config(), user_model=None, user_temperature=None, user_request_timeout=None, @@ -13088,8 +14136,8 @@ class TestResponseCostFromLoggingObjForBridgedMessages: return logging_obj def test_recomputed_cost_prices_the_reasoning_tokens_of_the_bridged_chat_completion(self): - from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage from litellm.llms.anthropic.messages.adapters.transformation import LiteLLMAnthropicMessagesAdapter + from litellm.types.utils import CompletionTokensDetailsWrapper, ModelResponse, Usage model = "dashscope/qwen-turbo" pricing = litellm.get_model_info(model) diff --git a/tests/unit/proxy/test_response_polling_pre_call_checks.py b/tests/unit/proxy/test_response_polling_pre_call_checks.py index 4228c1ff8f8..2c1dc9f6ae5 100644 --- a/tests/unit/proxy/test_response_polling_pre_call_checks.py +++ b/tests/unit/proxy/test_response_polling_pre_call_checks.py @@ -92,7 +92,7 @@ class TestSkipPreCallLogic: proxy_logging_obj=mock_proxy_logging, llm_router=MagicMock(), general_settings={}, - proxy_config=MagicMock(), + proxy_config=MagicMock(get_hierarchical_router_settings=AsyncMock(return_value=None)), ) except Exception: pass