diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md index 04ffcebf3b6..eaa28600268 100644 --- a/litellm/_auth/AGENTS.md +++ b/litellm/_auth/AGENTS.md @@ -13,4 +13,5 @@ - A `Step` is either an `Outcome` or an effect. Each effect carries its typed `resume` continuation, and `driver.drive` answers effects through an `EffectHandler` until an outcome comes back. The proxy's handler (`_BuilderEffects` in `litellm/proxy/auth/user_api_key_auth.py`) owns all cache and DB IO, so both cores share it - Projection stays in the proxy. `_authenticated` in the builder module turns an `Allow` grant into `UserAPIKeyAuth`, and `litellm/proxy/auth/auth_denial.py` maps every `Denial` to the existing error contract - The key checks run in the core, in the order the proxy has always applied them: blocked key, model access, SCIM-deactivated owner, budget exemption, team member budget, expiry, budget alerts, the key's own max budget (or its throttle), the soft budget alert, per-model budgets, the team and project context, the proxy budget alert, then the key's token. A budget-exempt request skips everything from the team member budget to per-model budgets except expiry. `start_resolved_key` runs the same checks for keys the proxy resolved itself, such as internal jobs -- Per-model budgets are still one opaque effect (`EnforceModelBudgets`), because the key, team member and custom-auth paths share their fallback logic. Move that logic into the core for every path at once rather than forking it per path +- `model_budgets.py` decides the key and user per-model budgets: which models to check, when an over-budget key may fall back (only for the top-level request model, to the first `budget_fallbacks` entry it can afford and may call), and that the user check sees the rerouted model. The key chain runs it inline, custom auth runs it through `AuthCore.check_model_budgets`, and JWT runs the user half through `AuthCore.check_user_model_budgets`. Those two entry points end in `Allow(WithinModelBudgets())` and are driven by `driver.drive_model_budgets`, which only needs a `ModelBudgetHandler` +- The team, team member and end-user per-model budgets still run in the proxy after the core. The team member check has its own fallback walk over the router's fallbacks, which is the next thing to move diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py index 9e697de7fad..abd6b54d0fa 100644 --- a/litellm/_auth/decision.py +++ b/litellm/_auth/decision.py @@ -1,4 +1,4 @@ -from collections.abc import Callable +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime from typing import Protocol, TypeAlias @@ -45,7 +45,12 @@ class VirtualKey: budget_throttle_pct: float | None -Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey +@dataclass(frozen=True, slots=True) +class WithinModelBudgets: + pass + + +Grant: TypeAlias = NoMasterKey | UserAuthRoute | MasterKey | CachedAdminKey | VirtualKey | WithinModelBudgets @dataclass(frozen=True, slots=True) @@ -114,6 +119,15 @@ class KeyWithoutToken: pass +@dataclass(frozen=True, slots=True) +class ModelBudgetExceeded: + message: str + current_cost: float + max_budget: float + entity_type: str | None + entity_id: str | None + + Denial: TypeAlias = ( MissingKey | MalformedKey @@ -127,6 +141,7 @@ Denial: TypeAlias = ( | TeamMemberBudgetExceeded | KeyBudgetExceeded | KeyWithoutToken + | ModelBudgetExceeded ) @@ -181,6 +196,15 @@ class TeamMemberBudget: max_budget: float +@dataclass(frozen=True, slots=True) +class ModelBudgetScope: + models: tuple[str, ...] + rewritable_model: str | None + key_has_model_budgets: bool + key_budget_fallbacks: Mapping[str, tuple[str, ...]] + user_has_model_budgets: bool + + @dataclass(frozen=True, slots=True) class ResolveEndUser: prefetch_key: bool @@ -269,11 +293,6 @@ class AlertKeySoftBudget: resume: "Callable[[], Step]" -@dataclass(frozen=True, slots=True) -class EnforceModelBudgets: - resume: "Callable[[], Step]" - - @dataclass(frozen=True, slots=True) class LoadKeyContext: resume: "Callable[[], Step]" @@ -290,6 +309,35 @@ class AlertProxyBudget: resume: "Callable[[], Step]" +@dataclass(frozen=True, slots=True) +class ReadModelBudgetScope: + resume: "Callable[[ModelBudgetScope], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckKeyModelBudget: + model: str + resume: "Callable[[ModelBudgetExceeded | None], Step]" + + +@dataclass(frozen=True, slots=True) +class AuthorizeBudgetFallback: + model: str + resume: "Callable[[bool], Step]" + + +@dataclass(frozen=True, slots=True) +class RewriteRequestModel: + model: str + resume: "Callable[[], Step]" + + +@dataclass(frozen=True, slots=True) +class CheckUserModelBudget: + model: str + resume: "Callable[[ModelBudgetExceeded | None], Step]" + + IdentityEffect: TypeAlias = ( ResolveEndUser | ResolvePresentedKey @@ -309,11 +357,13 @@ BudgetEffect: TypeAlias = ( | ReadKeySpend | AlertKeyBudget | AlertKeySoftBudget - | EnforceModelBudgets | ReadProxySpend | AlertProxyBudget ) -Effect: TypeAlias = IdentityEffect | BudgetEffect +ModelBudgetEffect: TypeAlias = ( + ReadModelBudgetScope | CheckKeyModelBudget | AuthorizeBudgetFallback | RewriteRequestModel | CheckUserModelBudget +) +Effect: TypeAlias = IdentityEffect | BudgetEffect | ModelBudgetEffect Step: TypeAlias = Outcome | Effect @@ -321,3 +371,7 @@ class AuthCore(Protocol): def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... + + def check_model_budgets(self) -> Step: ... + + def check_user_model_budgets(self) -> Step: ... diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py index e709e095b2a..b014dba46f8 100644 --- a/litellm/_auth/driver.py +++ b/litellm/_auth/driver.py @@ -6,12 +6,14 @@ from litellm._auth.decision import ( AlertProxyBudget, AlertTeamMemberBudget, Allow, + AuthorizeBudgetFallback, BudgetEffect, CheckBudgetExemption, + CheckKeyModelBudget, + CheckUserModelBudget, Deny, Effect, EnforceModelAccess, - EnforceModelBudgets, EvictPresentedKey, IdentityEffect, InspectKey, @@ -19,22 +21,39 @@ from litellm._auth.decision import ( LoadKey, LoadKeyContext, LoadKeyOwner, + ModelBudgetEffect, + ModelBudgetExceeded, + ModelBudgetScope, Outcome, OwnerFacts, PrepareKey, PresentedKey, ReadKeySpend, + ReadModelBudgetScope, ReadProxySpend, ReadTeamMemberBudget, ResolveEndUser, ResolvePresentedKey, + RewriteRequestModel, SendKeyBudgetAlerts, Step, TeamMemberBudget, ) -class EffectHandler(Protocol): +class ModelBudgetHandler(Protocol): + async def read_model_budget_scope(self) -> ModelBudgetScope: ... + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + + async def authorize_budget_fallback(self, model: str) -> bool: ... + + async def rewrite_request_model(self, model: str) -> None: ... + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + + +class EffectHandler(ModelBudgetHandler, Protocol): async def resolve_end_user(self, prefetch_key: bool) -> None: ... async def resolve_presented_key(self, hashed_token: str) -> PresentedKey | None: ... @@ -65,8 +84,6 @@ class EffectHandler(Protocol): async def alert_key_soft_budget(self) -> None: ... - async def enforce_model_budgets(self) -> None: ... - async def load_key_context(self) -> None: ... async def read_proxy_spend(self) -> float | None: ... @@ -89,6 +106,14 @@ async def _answer(step: Effect, handler: EffectHandler) -> Step: | LoadKeyContext() ): return await _answer_identity(step, handler) + case ( + ReadModelBudgetScope() + | CheckKeyModelBudget() + | AuthorizeBudgetFallback() + | RewriteRequestModel() + | CheckUserModelBudget() + ): + return await _answer_model_budget(step, handler) case _: return await _answer_budget(step, handler) @@ -143,9 +168,6 @@ async def _answer_budget(step: BudgetEffect, handler: EffectHandler) -> Step: case AlertKeySoftBudget(resume=resume): await handler.alert_key_soft_budget() return resume() - case EnforceModelBudgets(resume=resume): - await handler.enforce_model_budgets() - return resume() case ReadProxySpend(resume=resume): return resume(await handler.read_proxy_spend()) case AlertProxyBudget(spend=spend, resume=resume): @@ -155,6 +177,39 @@ async def _answer_budget(step: BudgetEffect, handler: EffectHandler) -> Step: assert_never(step) +async def _answer_model_budget(step: ModelBudgetEffect, handler: ModelBudgetHandler) -> Step: + match step: + case ReadModelBudgetScope(resume=resume): + return resume(await handler.read_model_budget_scope()) + case CheckKeyModelBudget(model=model, resume=resume): + return resume(await handler.check_key_model_budget(model)) + case AuthorizeBudgetFallback(model=model, resume=resume): + return resume(await handler.authorize_budget_fallback(model)) + case RewriteRequestModel(model=model, resume=resume): + await handler.rewrite_request_model(model) + return resume() + case CheckUserModelBudget(model=model, resume=resume): + return resume(await handler.check_user_model_budget(model)) + case _: + assert_never(step) + + +async def drive_model_budgets(step: Step, handler: ModelBudgetHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case ( + ReadModelBudgetScope() + | CheckKeyModelBudget() + | AuthorizeBudgetFallback() + | RewriteRequestModel() + | CheckUserModelBudget() + ): + return await drive_model_budgets(await _answer_model_budget(step, handler), handler) + case _: + raise RuntimeError(f"{type(step).__name__} is not a model budget effect") + + async def drive(step: Step, handler: EffectHandler) -> Outcome: match step: case Allow() | Deny(): diff --git a/litellm/_auth/model_budgets.py b/litellm/_auth/model_budgets.py new file mode 100644 index 00000000000..9e5ad2d9f3d --- /dev/null +++ b/litellm/_auth/model_budgets.py @@ -0,0 +1,87 @@ +from collections.abc import Callable +from typing import Final + +from litellm._auth.decision import ( + AuthorizeBudgetFallback, + CheckKeyModelBudget, + CheckUserModelBudget, + Deny, + ModelBudgetExceeded, + ModelBudgetScope, + ReadModelBudgetScope, + RewriteRequestModel, + Step, +) + + +def check_model_budgets(then: Callable[[], Step]) -> Step: + return ReadModelBudgetScope(resume=lambda scope: _check_key_budgets(scope, then)) + + +def check_user_model_budgets(then: Callable[[], Step]) -> Step: + return ReadModelBudgetScope(resume=lambda scope: _check_user_budgets(scope, scope.models, then)) + + +def _check_key_budgets(scope: ModelBudgetScope, then: Callable[[], Step]) -> Step: + if not scope.key_has_model_budgets or not scope.models: + return _check_user_budgets(scope, scope.models, then) + return _check_key_budget_of( + scope, + scope.models, + scope.rewritable_model, + then=lambda: ReadModelBudgetScope(resume=lambda rescoped: _check_user_budgets(rescoped, rescoped.models, then)), + ) + + +def _check_key_budget_of( + scope: ModelBudgetScope, models: tuple[str, ...], rewritable_model: str | None, then: Callable[[], Step] +) -> Step: + if not models: + return then() + model: Final = models[0] + rest: Final = models[1:] + + def after(exceeded: ModelBudgetExceeded | None) -> Step: + if exceeded is None: + return _check_key_budget_of(scope, rest, rewritable_model, then) + if model != rewritable_model: + return Deny(exceeded) + return _first_fallback_within_budget( + scope.key_budget_fallbacks.get(model, ()), + found=lambda fallback: _reroute(fallback, exceeded, then=lambda: _check_key_budget_of(scope, rest, fallback, then)), + exhausted=lambda: Deny(exceeded), + ) + + return CheckKeyModelBudget(model=model, resume=after) + + +def _first_fallback_within_budget( + candidates: tuple[str, ...], found: Callable[[str], Step], exhausted: Callable[[], Step] +) -> Step: + if not candidates: + return exhausted() + candidate: Final = candidates[0] + return CheckKeyModelBudget( + model=candidate, + resume=lambda exceeded: ( + found(candidate) + if exceeded is None + else _first_fallback_within_budget(candidates[1:], found, exhausted) + ), + ) + + +def _reroute(fallback: str, exceeded: ModelBudgetExceeded, then: Callable[[], Step]) -> Step: + return AuthorizeBudgetFallback( + model=fallback, + resume=lambda authorized: RewriteRequestModel(model=fallback, resume=then) if authorized else Deny(exceeded), + ) + + +def _check_user_budgets(scope: ModelBudgetScope, models: tuple[str, ...], then: Callable[[], Step]) -> Step: + if not scope.user_has_model_budgets or not models: + return then() + return CheckUserModelBudget( + model=models[0], + resume=lambda exceeded: _check_user_budgets(scope, models[1:], then) if exceeded is None else Deny(exceeded), + ) diff --git a/litellm/_auth/python_core.py b/litellm/_auth/python_core.py index c57577c8227..a4c7463e71e 100644 --- a/litellm/_auth/python_core.py +++ b/litellm/_auth/python_core.py @@ -14,7 +14,6 @@ from litellm._auth.decision import ( CheckBudgetExemption, Deny, EnforceModelAccess, - EnforceModelBudgets, EvictPresentedKey, ExpiredKey, InspectKey, @@ -51,8 +50,10 @@ from litellm._auth.decision import ( UserAuthDisabled, UserAuthRoute, VirtualKey, + WithinModelBudgets, ) from litellm._auth.key import hash_token, is_master_key +from litellm._auth.model_budgets import check_model_budgets, check_user_model_budgets _USER_AUTH_ROUTE: Final = "/user/auth" _VIRTUAL_KEY_PREFIX: Final = "sk-" @@ -91,6 +92,12 @@ class PythonCore: def start_resolved_key(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: return _check_key(_KeyCheck(settings, now, facts.is_llm_api_route, key_label=facts.token or "")) + def check_model_budgets(self) -> Step: + return check_model_budgets(then=lambda: Allow(WithinModelBudgets())) + + def check_user_model_budgets(self) -> Step: + return check_user_model_budgets(then=lambda: Allow(WithinModelBudgets())) + def _after_presented_key( presented: PresentedKey | None, @@ -211,7 +218,7 @@ def _after_key_spend(check: _KeyCheck, key: KeyFacts, spend: float, max_budget: def _check_soft_budget(check: _KeyCheck, key: KeyFacts, budget_throttle_pct: float | None) -> Step: def model_budgets() -> Step: - return EnforceModelBudgets(resume=lambda: _load_context(check, key, budget_throttle_pct)) + return check_model_budgets(then=lambda: _load_context(check, key, budget_throttle_pct)) if key.soft_budget and key.spend >= key.soft_budget: return AlertKeySoftBudget(resume=model_budgets) diff --git a/litellm/proxy/auth/auth_denial.py b/litellm/proxy/auth/auth_denial.py index acbff4996a8..5b872561e6f 100644 --- a/litellm/proxy/auth/auth_denial.py +++ b/litellm/proxy/auth/auth_denial.py @@ -12,6 +12,7 @@ from litellm._auth.decision import ( MalformedKey, MasterKeyOnlyRoute, MissingKey, + ModelBudgetExceeded, NoDatabase, NotAVirtualKey, OwnerDeactivated, @@ -127,5 +128,19 @@ def denial_exception( return HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key, no token associated" ) + case ModelBudgetExceeded( + message=message, + current_cost=current_cost, + max_budget=max_budget, + entity_type=entity_type, + entity_id=entity_id, + ): + return litellm.BudgetExceededError( + current_cost=current_cost, + max_budget=max_budget, + message=message, + entity_type=entity_type, + entity_id=entity_id, + ) case _: assert_never(denial) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f58757b5d7e..dfb3d088794 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -9,7 +9,7 @@ Returns a UserAPIKeyAuth object if the API key is valid import asyncio import fnmatch -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from datetime import datetime, timezone from enum import Enum from types import MappingProxyType @@ -39,6 +39,8 @@ from litellm._auth.decision import ( Grant, KeyFacts, MasterKey, + ModelBudgetExceeded, + ModelBudgetScope, NoMasterKey, NonAdminKey, Outcome, @@ -46,11 +48,13 @@ from litellm._auth.decision import ( PresentedKey, RequestFacts, Settings, + Step, TeamMemberBudget, UserAuthRoute, VirtualKey, + WithinModelBudgets, ) -from litellm._auth.driver import drive +from litellm._auth.driver import drive, drive_model_budgets from litellm._auth.python_core import PythonCore from litellm._internal_context import service_target from litellm._logging import verbose_logger, verbose_proxy_logger @@ -395,13 +399,9 @@ def _budget_fallback_target_model(target: object) -> str | None: return model if isinstance(model, str) else None -class _KeyModelBudgetLimiter(Protocol): +class _KeyAndUserModelBudgetLimiter(Protocol): async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ... - async def get_fallback_model_within_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> str | None: ... - - -class _UserModelBudgetLimiter(Protocol): async def is_user_within_model_budget( self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None ) -> bool: ... @@ -528,28 +528,6 @@ async def _read_user_model_max_budget( return user_obj.model_max_budget if user_obj is not None else None -async def _check_user_model_budget( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _UserModelBudgetLimiter, - models: list[str], -) -> None: - """Enforce the internal user's own `model_max_budget` across the request's models. - - Separate from the key check: a user's per-model budget caps every key they - own, so a caller cannot escape it by minting another key. - """ - user_model_max_budget: Final = valid_token.user_model_max_budget - if valid_token.user_id is None or not isinstance(user_model_max_budget, Mapping) or not user_model_max_budget: - return - for model_name in models: - await model_max_budget_limiter.is_user_within_model_budget( - user_id=valid_token.user_id, - user_model_max_budget=user_model_max_budget, - model=model_name, - team_id=valid_token.team_id, - ) - - async def _check_team_model_budget( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _TeamModelBudgetLimiter, @@ -1049,62 +1027,118 @@ async def _try_team_member_model_budget_fallback( return False -async def _check_key_model_budget_with_fallback( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _KeyModelBudgetLimiter, - model_name: str, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None = None, - llm_router: litellm.Router | None = None, -) -> None: - """ - Enforce the key's per-model budget for `model_name`. If exceeded and the - key has a `budget_fallbacks` chain configured for `model_name`, reroute - the request to the first fallback model still within its own budget - instead of rejecting the request. - - The selected fallback is validated against the key's model-access - allowlist and the team's model restrictions so that budget_fallbacks - cannot bypass model authorization. The rewrite is persisted to the - parsed-body cache, Starlette's JSON cache (``request._json``), and - path parameters so that downstream handlers see the final model - regardless of whether they consume ``_read_request_body()``, - ``request.json()``, or the path ``model`` parameter. - - Fallback is only attempted when ``model_name`` matches the top-level - ``request_data["model"]``; models extracted from nested fields - (``session.model``, ``completion.model``, etc.) are not rewritable - and raise immediately. - - Raises: - BudgetExceededError: if `model_name` is over budget and no configured - fallback is within budget either (or the fallback is not authorized). - """ +async def _model_budget_verdict(check: Awaitable[bool]) -> ModelBudgetExceeded | None: try: - await model_max_budget_limiter.is_key_within_model_budget( - user_api_key_dict=valid_token, - model=model_name, - ) + with tracer.trace("litellm.proxy.auth.budget_checks"): + _ = await check except litellm.BudgetExceededError as e: - if request_data.get("model") != model_name: - raise e - fallback_model: Final = await model_max_budget_limiter.get_fallback_model_within_budget( - user_api_key_dict=valid_token, - model=model_name, + return ModelBudgetExceeded( + message=e.message, + current_cost=e.current_cost, + max_budget=e.max_budget, + entity_type=e.entity_type, + entity_id=e.entity_id, ) - if fallback_model is None: - raise e - if not await _is_fallback_model_authorized( - model=fallback_model, - valid_token=valid_token, - llm_model_list=llm_model_list, - llm_router=llm_router, - ): - raise - rewrite_request_model(request_data=request_data, request=request, model=fallback_model) - if "path_params" in request.scope: - request.scope["path_params"]["model"] = fallback_model # rebind-ok: route handlers use the fallback model + return None + + +class _ModelBudgetEffects: + def __init__( + self, + *, + request: Request, + request_data: dict[str, object], # mutable-ok: budget fallbacks rewrite the request model in place + valid_token: UserAPIKeyAuth, + route: str, + limiter: _KeyAndUserModelBudgetLimiter, + llm_router: litellm.Router | None, + llm_model_list: Sequence[object] | None, + ) -> None: + self._request: Final = request + self._request_data: Final = request_data + self._valid_token: Final = valid_token + self._route: Final = route + self._limiter: Final = limiter + self._llm_router: Final = llm_router + self._llm_model_list: Final = llm_model_list + + @property + def valid_token(self) -> UserAPIKeyAuth: + return self._valid_token + + @property + def route(self) -> str: + return self._route + + async def read_model_budget_scope(self) -> ModelBudgetScope: + requested: Final = self._request_data.get("model") + user_model_max_budget: Final = self._valid_token.user_model_max_budget + budget_fallbacks: Final[Mapping[str, list[str]]] = self._valid_token.budget_fallbacks or {} + return ModelBudgetScope( + models=tuple( + _get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=self._request_data, + route=self._route, + request=self._request, + llm_router=self._llm_router, + team_id=self._valid_token.team_id, + ) + ) + ), + rewritable_model=requested if isinstance(requested, str) else None, + key_has_model_budgets=bool(self._valid_token.model_max_budget) and self._valid_token.token is not None, + key_budget_fallbacks=MappingProxyType({model: tuple(chain) for model, chain in budget_fallbacks.items()}), + user_has_model_budgets=self._valid_token.user_id is not None + and isinstance(user_model_max_budget, Mapping) + and bool(user_model_max_budget), + ) + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + return await _model_budget_verdict( + self._limiter.is_key_within_model_budget(user_api_key_dict=self._valid_token, model=model) + ) + + async def authorize_budget_fallback(self, model: str) -> bool: + return await _is_fallback_model_authorized( + model=model, + valid_token=self._valid_token, + llm_model_list=self._llm_model_list, + llm_router=self._llm_router, + ) + + async def rewrite_request_model(self, model: str) -> None: + rewrite_request_model(request_data=self._request_data, request=self._request, model=model) + if "path_params" in self._request.scope: + self._request.scope["path_params"]["model"] = model + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + user_id: Final = self._valid_token.user_id + user_model_max_budget: Final = self._valid_token.user_model_max_budget + if user_id is None or not isinstance(user_model_max_budget, Mapping): + return None + return await _model_budget_verdict( + self._limiter.is_user_within_model_budget( + user_id=user_id, + user_model_max_budget=user_model_max_budget, + model=model, + team_id=self._valid_token.team_id, + ) + ) + + +async def _enforce_model_budgets(step: Step, effects: _ModelBudgetEffects) -> None: + match await drive_model_budgets(step, effects): + case Deny(denial=denial): + raise denial_exception( + denial, + token=effects.valid_token.token, + route=effects.route, + jwt_auth_enabled=False, + key=effects.valid_token, + ) + case Allow(): + return def _routing_selector_matches_claim( @@ -2622,11 +2656,33 @@ class _BuilderEffects: async def alert_key_soft_budget(self) -> None: send_soft_budget_alert(self.key, self._proxy_logging_obj, self.owner) - async def enforce_model_budgets(self) -> None: - with tracer.trace("litellm.proxy.auth.budget_checks"): - await _check_model_budgets( - request=self._request, request_data=self._request_data, valid_token=self.key, route=self._route - ) + def _model_budgets(self) -> _ModelBudgetEffects: + from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter + + return _ModelBudgetEffects( + request=self._request, + request_data=self._request_data, + valid_token=self.key, + route=self._route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), + ) + + async def read_model_budget_scope(self) -> ModelBudgetScope: + return await self._model_budgets().read_model_budget_scope() + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + return await self._model_budgets().check_key_model_budget(model) + + async def authorize_budget_fallback(self, model: str) -> bool: + return await self._model_budgets().authorize_budget_fallback(model) + + async def rewrite_request_model(self, model: str) -> None: + await self._model_budgets().rewrite_request_model(model) + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + return await self._model_budgets().check_user_model_budget(model) async def load_key_context(self) -> None: await _load_team_and_project( @@ -3037,17 +3093,16 @@ async def user_api_key_auth_builder( # zero-cost model would be refused here and allowed there, # while the log above claims all budget checks were skipped. if not skip_budget_checks: - await _check_user_model_budget( - valid_token=cast(UserAPIKeyAuth, valid_token), - model_max_budget_limiter=model_max_budget_limiter, - models=_get_model_names_for_budget_checks( - model=_get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) + await _enforce_model_budgets( + _AUTH_CORE.check_user_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=cast(UserAPIKeyAuth, valid_token), + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=_proxy_model_list(), ), ) @@ -3122,52 +3177,6 @@ async def user_api_key_auth_builder( _user_api_key_auth_builder: Final = user_api_key_auth_builder -async def _check_model_budgets( - request: Request, - request_data: dict[str, object], # mutable-ok: per-model budget fallbacks rewrite the request model in place - valid_token: UserAPIKeyAuth, - route: str, -) -> None: - from litellm.proxy.proxy_server import llm_router, model_max_budget_limiter, prisma_client - - def requested_models() -> list[str]: - return _get_model_names_for_budget_checks( - model=_get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - ) - - models: Final = requested_models() - checks_key_budgets: Final = ( - bool(valid_token.model_max_budget) - and prisma_client is not None - and bool(models) - and valid_token.token is not None - ) - for model_name in models if checks_key_budgets else (): - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, _proxy_model_list() - ), - llm_router=llm_router, - ) - # A budget fallback may have rewritten the request model, so the user check validates the final one. - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=requested_models() if checks_key_budgets else models, - ) - - async def _load_team_and_project( valid_token: UserAPIKeyAuth, prisma_client: PrismaClient | None, @@ -3317,6 +3326,8 @@ async def _granted( return _with_end_user(effects.key, effects.end_user, effects.parent_otel_span) case VirtualKey(budget_throttle_pct=budget_throttle_pct): return await _virtual_key_auth(effects, budget_throttle_pct, route, start_time) + case WithinModelBudgets(): + raise RuntimeError("the request auth core never ends on a model budget check") case _: assert_never(grant) @@ -4702,14 +4713,13 @@ async def _run_post_custom_auth_checks( llm_router=llm_router, ) - current_model = _get_model_from_request_context( + current_model: Final = _get_model_from_request_context( request_data=request_data, route=route, request=request, llm_router=llm_router, team_id=valid_token.team_id, ) - current_models = _get_model_names_for_budget_checks(model=current_model) # A zero-cost model cannot move any counter, so refusing it means refusing on # spend some other model accrued. The JWT and virtual-key paths already skip @@ -4721,41 +4731,6 @@ async def _run_post_custom_auth_checks( else False ) - # 3. Check key-level model_max_budget - max_budget_per_model: Final = valid_token.model_max_budget - if ( - not skip_budget_checks - and max_budget_per_model is not None - and isinstance(max_budget_per_model, dict) - and len(max_budget_per_model) > 0 - and current_models - and valid_token.token is not None - ): - for model_name in current_models: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - model_name=model_name, - request_data=request_data, - request=request, - llm_model_list=cast( # cast-ok: router model lists are deployment sequences - Sequence[object] | None, llm_model_list - ), - llm_router=llm_router, - ) - - # Recompute after a potential budget-fallback rewrite so - # the end-user check below validates the final model - current_model = _get_model_from_request_context( - request_data=request_data, - route=route, - request=request, - llm_router=llm_router, - team_id=valid_token.team_id, - ) - current_models = _get_model_names_for_budget_checks(model=current_model) - - # 3b. Attach and check the internal user's model_max_budget. # Custom auth builds its own token, so unlike the main path nothing has # loaded the user row yet. The attach is unconditional because the post-call # spend hook reads this field off the token: gating it on the same condition @@ -4769,14 +4744,33 @@ async def _run_post_custom_auth_checks( proxy_logging_obj=proxy_logging_obj, ) valid_token.user_model_max_budget = user_budget # rebind-ok: the spend hook reads it off this token - if not skip_budget_checks and current_models: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=current_models, + if not skip_budget_checks: + await _enforce_model_budgets( + _AUTH_CORE.check_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=valid_token, + route=route, + limiter=model_max_budget_limiter, + llm_router=llm_router, + llm_model_list=cast( # cast-ok: router model lists are deployment sequences + Sequence[object] | None, llm_model_list + ), + ), ) - # 4. Check end-user model_max_budget + # Read after the budget checks, which may have rerouted the request to a fallback model. + current_models: Final = _get_model_names_for_budget_checks( + model=_get_model_from_request_context( + request_data=request_data, + route=route, + request=request, + llm_router=llm_router, + team_id=valid_token.team_id, + ) + ) + end_user_mmb: Final = valid_token.end_user_model_max_budget if ( not skip_budget_checks diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 889fb9fb373..cc2e91cd93e 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -882,21 +882,6 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): ), ) - @with_service_target("model_budgets") - async def get_fallback_model_within_budget( - self, - user_api_key_dict: UserAPIKeyAuth, - model: str, - ) -> str | None: - budget_fallbacks: Final[dict[str, list[str]]] = user_api_key_dict.budget_fallbacks or {} - for fallback_model in budget_fallbacks.get(model, []): - try: - await self.is_key_within_model_budget(user_api_key_dict=user_api_key_dict, model=fallback_model) - return fallback_model - except litellm.BudgetExceededError: - continue - return None - @with_service_target("model_budgets") async def is_user_within_model_budget( self, diff --git a/tests/unit/_auth/test_model_budgets.py b/tests/unit/_auth/test_model_budgets.py new file mode 100644 index 00000000000..4370275a111 --- /dev/null +++ b/tests/unit/_auth/test_model_budgets.py @@ -0,0 +1,163 @@ +from dataclasses import dataclass, replace +from typing import Final + +import pytest + +from litellm._auth.decision import Allow, Deny, ModelBudgetExceeded, ModelBudgetScope, Outcome, Step, WithinModelBudgets +from litellm._auth.driver import drive_model_budgets +from litellm._auth.python_core import PythonCore + +Call = tuple[str, str | None] + +SCOPE: Final = ModelBudgetScope( + models=("gpt-4o",), + rewritable_model="gpt-4o", + key_has_model_budgets=True, + key_budget_fallbacks={"gpt-4o": ("gpt-4o-mini", "haiku")}, + user_has_model_budgets=True, +) +ALLOWED: Final = Allow(WithinModelBudgets()) + + +def _exceeded(entity_type: str, model: str) -> ModelBudgetExceeded: + return ModelBudgetExceeded( + message=f"{entity_type} over budget for {model}", + current_cost=2.0, + max_budget=1.0, + entity_type=entity_type, + entity_id="id", + ) + + +@dataclass +class FakeRequest: + scope: ModelBudgetScope + key_over: frozenset[str] = frozenset() + user_over: frozenset[str] = frozenset() + unauthorized: frozenset[str] = frozenset() + calls: tuple[Call, ...] = () + + def _record(self, name: str, model: str | None = None) -> None: + self.calls = (*self.calls, (name, model)) + + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read") + return self.scope + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("key", model) + return _exceeded("key", model) if model in self.key_over else None + + async def authorize_budget_fallback(self, model: str) -> bool: + self._record("authorize", model) + return model not in self.unauthorized + + async def rewrite_request_model(self, model: str) -> None: + self._record("rewrite", model) + self.scope = replace( + self.scope, + models=tuple(model if m == self.scope.rewritable_model else m for m in self.scope.models), + rewritable_model=model, + ) + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("user", model) + return _exceeded("user", model) if model in self.user_over else None + + +async def _run(step: Step, request: FakeRequest) -> tuple[Outcome, tuple[Call, ...]]: + outcome: Final = await drive_model_budgets(step, request) + return outcome, request.calls + + +@pytest.mark.asyncio +async def test_within_budget_checks_the_key_then_the_user_without_rerouting(): + assert await _run(PythonCore().check_model_budgets(), FakeRequest(SCOPE)) == ( + ALLOWED, + (("read", None), ("key", "gpt-4o"), ("read", None), ("user", "gpt-4o")), + ) + + +@pytest.mark.asyncio +async def test_over_budget_reroutes_to_the_first_fallback_within_budget_and_checks_the_user_on_it(): + request: Final = FakeRequest(SCOPE, key_over=frozenset({"gpt-4o", "gpt-4o-mini"})) + assert await _run(PythonCore().check_model_budgets(), request) == ( + ALLOWED, + ( + ("read", None), + ("key", "gpt-4o"), + ("key", "gpt-4o-mini"), + ("key", "haiku"), + ("authorize", "haiku"), + ("rewrite", "haiku"), + ("read", None), + ("user", "haiku"), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "key_over", "unauthorized"), + [ + (replace(SCOPE, rewritable_model=None), {"gpt-4o"}, set()), + (replace(SCOPE, rewritable_model="other"), {"gpt-4o"}, set()), + (replace(SCOPE, key_budget_fallbacks={}), {"gpt-4o"}, set()), + (SCOPE, {"gpt-4o", "gpt-4o-mini", "haiku"}, set()), + (SCOPE, {"gpt-4o"}, {"gpt-4o-mini"}), + ], + ids=["nested_model", "different_top_level_model", "no_fallbacks", "every_fallback_over", "fallback_unauthorized"], +) +async def test_an_unreroutable_over_budget_model_is_denied_with_its_own_budget_error( + scope: ModelBudgetScope, key_over: set[str], unauthorized: set[str] +): + request: Final = FakeRequest(scope, key_over=frozenset(key_over), unauthorized=frozenset(unauthorized)) + outcome, calls = await _run(PythonCore().check_model_budgets(), request) + assert outcome == Deny(_exceeded("key", "gpt-4o")) + assert not any(name in ("rewrite", "user") for name, _ in calls) + + +@pytest.mark.asyncio +async def test_a_later_model_is_checked_after_an_earlier_one_was_rerouted(): + scope: Final = replace(SCOPE, models=("gpt-4o", "embed")) + request: Final = FakeRequest(scope, key_over=frozenset({"gpt-4o", "embed"})) + outcome, calls = await _run(PythonCore().check_model_budgets(), request) + assert outcome == Deny(_exceeded("key", "embed")) + assert calls[-2:] == (("rewrite", "gpt-4o-mini"), ("key", "embed")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "expected_calls"), + [ + ( + replace(SCOPE, key_has_model_budgets=False), + (("read", None), ("user", "gpt-4o")), + ), + ( + replace(SCOPE, user_has_model_budgets=False), + (("read", None), ("key", "gpt-4o"), ("read", None)), + ), + (replace(SCOPE, models=()), (("read", None),)), + ], + ids=["no_key_budgets", "no_user_budgets", "no_models"], +) +async def test_only_budgets_that_exist_are_checked(scope: ModelBudgetScope, expected_calls: tuple[Call, ...]): + assert await _run(PythonCore().check_model_budgets(), FakeRequest(scope)) == (ALLOWED, expected_calls) + + +@pytest.mark.asyncio +async def test_a_user_over_budget_is_denied_on_every_path(): + over: Final = frozenset({"gpt-4o"}) + for step in (PythonCore().check_model_budgets(), PythonCore().check_user_model_budgets()): + outcome, _ = await _run(step, FakeRequest(SCOPE, user_over=over)) + assert outcome == Deny(_exceeded("user", "gpt-4o")) + + +@pytest.mark.asyncio +async def test_the_user_only_check_never_reads_key_budgets(): + request: Final = FakeRequest(SCOPE, key_over=frozenset({"gpt-4o"})) + assert await _run(PythonCore().check_user_model_budgets(), request) == ( + ALLOWED, + (("read", None), ("user", "gpt-4o")), + ) diff --git a/tests/unit/_auth/test_python_core.py b/tests/unit/_auth/test_python_core.py index d48d0545d0c..ae6b20b6b54 100644 --- a/tests/unit/_auth/test_python_core.py +++ b/tests/unit/_auth/test_python_core.py @@ -14,6 +14,8 @@ from litellm._auth.decision import ( KeyBlocked, KeyBudgetExceeded, KeyFacts, + ModelBudgetExceeded, + ModelBudgetScope, KeyWithoutToken, MalformedKey, MasterKey, @@ -67,10 +69,18 @@ CHECKED: Final = ( "enforce_model_access", "check_budget_exemption", "send_key_budget_alerts", - "enforce_model_budgets", + "read_model_budget_scope", "load_key_context", ) +NO_MODEL_BUDGETS: Final = ModelBudgetScope( + models=(), + rewritable_model=None, + key_has_model_budgets=False, + key_budget_fallbacks={}, + user_has_model_budgets=False, +) + Call = tuple[str, object] @@ -83,6 +93,8 @@ class RecordingHandler: team_member_budget: TeamMemberBudget | None = None key_spend: float = 0.0 proxy_spend: float | None = None + model_budget_scope: ModelBudgetScope = NO_MODEL_BUDGETS + user_model_budget_exceeded: ModelBudgetExceeded | None = None calls: tuple[Call, ...] = () def _record(self, name: str, payload: object = None) -> None: @@ -139,8 +151,24 @@ class RecordingHandler: async def alert_key_soft_budget(self) -> None: self._record("alert_key_soft_budget") - async def enforce_model_budgets(self) -> None: - self._record("enforce_model_budgets") + async def read_model_budget_scope(self) -> ModelBudgetScope: + self._record("read_model_budget_scope") + return self.model_budget_scope + + async def check_key_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("check_key_model_budget", model) + return None + + async def authorize_budget_fallback(self, model: str) -> bool: + self._record("authorize_budget_fallback", model) + return True + + async def rewrite_request_model(self, model: str) -> None: + self._record("rewrite_request_model", model) + + async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("check_user_model_budget", model) + return self.user_model_budget_exceeded async def load_key_context(self) -> None: self._record("load_key_context") @@ -411,7 +439,7 @@ async def test_key_budget_denies_or_throttles_at_the_limit( assert outcome == expected assert calls[3:6] == ("send_key_budget_alerts", "read_key_spend", "alert_key_budget") assert handler.calls[5] == ("alert_key_budget", spend) - assert ("enforce_model_budgets" in calls) == isinstance(expected, Allow) + assert ("read_model_budget_scope" in calls) == isinstance(expected, Allow) @pytest.mark.asyncio @@ -420,7 +448,7 @@ async def test_key_budget_is_enforced_only_on_llm_routes(): outcome, calls = await _check(key, is_llm_api_route=False, key_spend=2.0) assert outcome == _allowed() assert "read_key_spend" not in calls - assert "enforce_model_budgets" in calls + assert "read_model_budget_scope" in calls @pytest.mark.asyncio @@ -441,7 +469,7 @@ async def test_soft_budget_alerts_at_the_threshold_before_model_budgets( outcome, calls = await _check(replace(PLAIN_KEY, spend=spend, soft_budget=soft_budget)) assert outcome == _allowed() soft_alert: Final = ("alert_key_soft_budget",) if alerted else () - assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "enforce_model_budgets", "load_key_context") + assert calls[3:] == ("send_key_budget_alerts", *soft_alert, "read_model_budget_scope", "load_key_context") @pytest.mark.asyncio @@ -449,7 +477,7 @@ async def test_throttled_key_still_gets_soft_and_model_budget_checks(): key: Final = replace(PLAIN_KEY, max_budget=5.0, budget_throttle_pct=0.5, spend=6.0, soft_budget=1.0) outcome, calls = await _check(key, key_spend=6.0) assert outcome == _allowed(budget_throttle_pct=0.5) - assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "enforce_model_budgets", "load_key_context") + assert calls[-4:] == ("alert_key_budget", "alert_key_soft_budget", "read_model_budget_scope", "load_key_context") @pytest.mark.asyncio @@ -488,3 +516,16 @@ async def test_key_without_token_is_denied_after_its_context_loads(exempt: bool) outcome, calls = await _run(PythonCore().start_resolved_key(_facts("hash"), settings, NOW), handler) assert outcome == Deny(KeyWithoutToken()) assert calls[-3:] == ("load_key_context", "read_proxy_spend", "alert_proxy_budget") + + +@pytest.mark.asyncio +async def test_a_model_budget_denial_stops_the_key_chain_before_its_context_loads(): + exceeded: Final = ModelBudgetExceeded( + message="over", current_cost=2.0, max_budget=1.0, entity_type="user", entity_id="u-1" + ) + scope: Final = replace(NO_MODEL_BUDGETS, models=("gpt-4o",), user_has_model_budgets=True) + handler: Final = RecordingHandler(model_budget_scope=scope, user_model_budget_exceeded=exceeded) + outcome, calls = await _run(PythonCore().start_resolved_key(_facts("hash"), SETTINGS, NOW), handler) + assert outcome == Deny(exceeded) + assert calls[-2:] == ("read_model_budget_scope", "check_user_model_budget") + diff --git a/tests/unit/proxy/auth/test_user_api_key_auth.py b/tests/unit/proxy/auth/test_user_api_key_auth.py index efa0ee912d8..e8ac0e96760 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth.py @@ -15,7 +15,10 @@ from litellm._logging import verbose_proxy_logger import json import logging import litellm +from litellm._auth.python_core import PythonCore from litellm.proxy.auth.user_api_key_auth import ( + _enforce_model_budgets, + _ModelBudgetEffects, user_api_key_auth, UserAPIKeyAuth, ) @@ -25,6 +28,23 @@ from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles from tests._master_key import MASTER_KEY +async def _check_user_model_budgets(valid_token: UserAPIKeyAuth, limiter: object, model: str) -> None: + request: Final = MagicMock() + request.scope = {} + await _enforce_model_budgets( + PythonCore().check_user_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data={"model": model}, + valid_token=valid_token, + route="/chat/completions", + limiter=limiter, + llm_router=None, + llm_model_list=None, + ), + ) + + class Request: def __init__(self, client_ip: Optional[str] = None, headers: Optional[dict] = None): self.client = MagicMock() @@ -1381,11 +1401,12 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ stored on LiteLLM_UserTable, accepted by /user/new and /user/update, and read by nothing, so a user-level per-model budget never blocked anything. """ - from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget - calls = [] class _Limiter: + async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: + return True + async def is_user_within_model_budget( self, user_id: str, user_model_max_budget: Mapping[str, object], model: str, team_id: str | None = None ) -> bool: @@ -1397,11 +1418,7 @@ async def test_check_user_model_budget(user_id, user_model_max_budget, expected_ user_id=user_id, user_model_max_budget=user_model_max_budget, ) - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=_Limiter(), - models=["gpt-4"], - ) + await _check_user_model_budgets(valid_token, _Limiter(), model="gpt-4") assert len(calls) == expected_calls if expected_calls: assert calls[0] == ("u-1", user_model_max_budget, "gpt-4") @@ -1543,7 +1560,6 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove nothing ever reads it, which looks enforced and is not. """ from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth - from litellm.proxy.auth.user_api_key_auth import _check_user_model_budget from litellm.proxy.hooks.model_max_budget_limiter import model_budget_spend_cache_key from litellm.proxy.proxy_server import model_max_budget_limiter @@ -1571,18 +1587,11 @@ async def test_jwt_user_model_budget_is_enforced_before_the_jwt_path_returns(ove if expect_refusal: with pytest.raises(litellm.BudgetExceededError) as exc: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=[model], - ) + await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) assert exc.value.entity_type == Litellm_EntityType.USER.value + assert exc.value.entity_id == user_id else: - await _check_user_model_budget( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=[model], - ) + await _check_user_model_budgets(valid_token, model_max_budget_limiter, model=model) def test_jwt_path_enforces_the_user_model_budget_before_returning(): @@ -1609,7 +1618,7 @@ def test_jwt_path_enforces_the_user_model_budget_before_returning(): if isinstance(child, ast.Call): fn = child.func name = getattr(fn, "id", None) or getattr(fn, "attr", None) - if name == "_check_user_model_budget": + if name == "_enforce_model_budgets": seen_check.append(child.lineno) return seen_check @@ -1801,13 +1810,13 @@ def test_zero_cost_models_skip_the_user_budget_check_on_every_path(): calls = [ node for node in ast.walk(tree) - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_check_user_model_budget" + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "_enforce_model_budgets" ] assert len(calls) == 1, f"expected the JWT call site, found {len(calls)}" unguarded = [c for c in calls if not guarded_by_skip(tree, c)] assert not unguarded, ( - f"{len(unguarded)} _check_user_model_budget call(s) run even when " + f"{len(unguarded)} _enforce_model_budgets call(s) run even when " "skip_budget_checks is set, so a zero-cost model is enforced on one auth path and not the other" ) @@ -1835,8 +1844,7 @@ def test_custom_auth_also_skips_budget_checks_for_zero_cost_models(): assert "skip_budget_checks" in src, "the custom-auth path never computes the zero-cost skip flag" budget_calls = ( - "_check_key_model_budget_with_fallback", - "_check_user_model_budget", + "_enforce_model_budgets", "is_end_user_within_model_budget", ) @@ -2107,7 +2115,7 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods( with ( patch("litellm.proxy.proxy_server.premium_user", True), patch("litellm.proxy.auth.handle_jwt.JWTAuthManager.auth_builder", return_value=jwt_response), - patch("litellm.proxy.auth.user_api_key_auth._check_user_model_budget", user_model_budget_check), + patch("litellm.proxy.auth.user_api_key_auth._enforce_model_budgets", user_model_budget_check), ): result: Final = await user_api_key_auth_builder( request=mock_request, diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index f71f8fabc28..883ad06c172 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -68,10 +68,12 @@ from litellm.proxy.auth.auth_checks import ( cache_key_object, jwt_key_mapping_cache_key, ) +from litellm._auth.python_core import PythonCore from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model from litellm.proxy.auth.user_api_key_auth import ( - _check_key_model_budget_with_fallback, + _enforce_model_budgets, + _ModelBudgetEffects, _check_team_member_model_budget_with_fallback, _ensure_litellm_received_at_on_request_state, _ensure_parent_otel_span_on_request_state, @@ -9587,10 +9589,43 @@ class TestJWTAuthUserEmail: assert result.user_email == "admin@example.com" -class TestCheckKeyModelBudgetWithFallback: - """`_check_key_model_budget_with_fallback` must reroute a request to the - first configured `budget_fallbacks` entry still within its own budget, - and only raise `BudgetExceededError` when no fallback is available.""" +def _key_over_budget_for(*models: str) -> AsyncMock: + async def is_key_within_model_budget(user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: + if model in models: + raise litellm.BudgetExceededError(current_cost=10, max_budget=5, message=f"over budget for {model}") + return True + + limiter: Final = AsyncMock() + limiter.is_key_within_model_budget.side_effect = is_key_within_model_budget + return limiter + + +async def _check_key_model_budgets( + valid_token: UserAPIKeyAuth, + limiter: object, + request_data: dict[str, object], + request: object, + llm_router: litellm.Router | None = None, +) -> None: + await _enforce_model_budgets( + PythonCore().check_model_budgets(), + _ModelBudgetEffects( + request=request, + request_data=request_data, + valid_token=valid_token, + route="/chat/completions", + limiter=limiter, + llm_router=llm_router, + llm_model_list=None, + ), + ) + + +_KEY_MODEL_BUDGET: Final = {"gpt-4o": {"budget_limit": 5.0, "time_period": "1d"}} + + +class TestKeyModelBudgetFallback: + """An over-budget key reroutes to the first `budget_fallbacks` entry it may call and can still afford.""" def _make_request(self): request = MagicMock() @@ -9599,162 +9634,64 @@ class TestCheckKeyModelBudgetWithFallback: @pytest.mark.asyncio async def test_within_budget_does_not_reroute(self): - valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) - limiter = AsyncMock() - limiter.is_key_within_model_budget.return_value = True + valid_token = UserAPIKeyAuth( + token="test-key", model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]} + ) request_data = {"model": "gpt-4o"} request = self._make_request() - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for(), request_data, request) assert request_data["model"] == "gpt-4o" - limiter.get_fallback_model_within_budget.assert_not_awaited() assert "parsed_body" not in request.scope @pytest.mark.asyncio - async def test_exceeded_budget_reroutes_to_fallback(self): + async def test_exceeded_budget_reroutes_to_fallback_visible_to_a_later_body_read(self): valid_token = UserAPIKeyAuth( token="test-key", + model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini", "claude-haiku"]}, ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for("gpt-4o"), request_data, request) assert request_data["model"] == "gpt-4o-mini" - limiter.get_fallback_model_within_budget.assert_awaited_once_with(user_api_key_dict=valid_token, model="gpt-4o") - # the rerouted model must be visible to a later, separate - # `_read_request_body` call on the same `request` (route handlers - # re-parse the body from this cache instead of reusing the dict). - cached_keys, cached_body = request.scope["parsed_body"] + _, cached_body = request.scope["parsed_body"] assert cached_body["model"] == "gpt-4o-mini" @pytest.mark.asyncio - async def test_raises_when_every_fallback_also_exceeded(self): - valid_token = UserAPIKeyAuth(token="test-key", budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = None + @pytest.mark.parametrize( + ("token_fields", "over_budget"), + [ + ({"budget_fallbacks": {"gpt-4o": ["gpt-4o-mini"]}}, ("gpt-4o", "gpt-4o-mini")), + ({"models": ["gpt-4o"], "budget_fallbacks": {"gpt-4o": ["restricted-model"]}}, ("gpt-4o",)), + ( + { + "models": ["gpt-4o", "restricted-model"], + "team_id": "team-1", + "team_models": ["gpt-4o"], + "budget_fallbacks": {"gpt-4o": ["restricted-model"]}, + }, + ("gpt-4o",), + ), + ], + ids=["every_fallback_over_budget", "fallback_outside_key_models", "fallback_outside_team_models"], + ) + async def test_raises_the_original_budget_error_when_no_fallback_is_usable( + self, token_fields: dict[str, object], over_budget: tuple[str, ...] + ): + valid_token = UserAPIKeyAuth(token="test-key", model_max_budget=_KEY_MODEL_BUDGET, **token_fields) request_data = {"model": "gpt-4o"} - request = self._make_request() with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, + await _check_key_model_budgets( + valid_token, _key_over_budget_for(*over_budget), request_data, self._make_request() ) - assert exc_info.value is original_error - assert request_data["model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_raises_when_fallback_not_authorized(self): - """If the fallback model is within budget but the key is not allowed - to call it, the original BudgetExceededError must be raised instead - of rerouting to an unauthorized model.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o"], - budget_fallbacks={"gpt-4o": ["restricted-model"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = "restricted-model" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert exc_info.value is original_error - assert request_data["model"] == "gpt-4o" - - @pytest.mark.asyncio - async def test_reroute_succeeds_when_fallback_is_authorized(self): - """If the fallback model is within budget AND authorized, the request - must be rerouted to it.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "gpt-4o-mini"], - budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, - ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert request_data["model"] == "gpt-4o-mini" - - @pytest.mark.asyncio - async def test_raises_when_fallback_blocked_by_team_models(self): - """If the key allows the fallback but the team does not, the original - BudgetExceededError must be raised.""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "restricted-model"], - team_id="team-1", - team_models=["gpt-4o"], - budget_fallbacks={"gpt-4o": ["restricted-model"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - limiter.get_fallback_model_within_budget.return_value = "restricted-model" - request_data = {"model": "gpt-4o"} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) - - assert exc_info.value is original_error + assert exc_info.value.message == "over budget for gpt-4o" + assert (exc_info.value.current_cost, exc_info.value.max_budget) == (10, 5) assert request_data["model"] == "gpt-4o" @pytest.mark.asyncio @@ -9765,56 +9702,18 @@ class TestCheckKeyModelBudgetWithFallback: valid_token = UserAPIKeyAuth( token="test-key", models=["gpt-4o", "gpt-4o-mini"], + model_max_budget=_KEY_MODEL_BUDGET, budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, ) - limiter = AsyncMock() - limiter.is_key_within_model_budget.side_effect = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.get_fallback_model_within_budget.return_value = "gpt-4o-mini" request_data = {"model": "gpt-4o"} request = self._make_request() request.scope["path_params"] = {"model": "gpt-4o"} - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + await _check_key_model_budgets(valid_token, _key_over_budget_for("gpt-4o"), request_data, request) assert request_data["model"] == "gpt-4o-mini" assert request.scope["path_params"]["model"] == "gpt-4o-mini" - @pytest.mark.asyncio - async def test_raises_when_model_from_nested_field(self): - """Budget fallback must not attempt rewrite when the checked model - came from a nested field (e.g. session.model) rather than the - top-level request_data['model'].""" - valid_token = UserAPIKeyAuth( - token="test-key", - models=["gpt-4o", "gpt-4o-mini"], - budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, - ) - limiter = AsyncMock() - original_error = litellm.BudgetExceededError(current_cost=10, max_budget=5) - limiter.is_key_within_model_budget.side_effect = original_error - request_data = {"session": {"model": "gpt-4o"}} - request = self._make_request() - - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) - - assert exc_info.value is original_error - assert "model" not in request_data - class TestCheckTeamMemberModelBudgetWithFallback: @pytest.fixture(autouse=True) @@ -10833,13 +10732,7 @@ class TestCheckTeamMemberModelBudgetWithFallback: request_data: Final = {"model": "gpt-4o"} request: Final = self._request("gpt-4o") - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - ) + await _check_key_model_budgets(valid_token, limiter, request_data, request) assert request_data["model"] == "gpt-4o-mini" diff --git a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py index 21cbe8667ff..55beaf5254d 100644 --- a/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py @@ -510,70 +510,6 @@ async def test_model_budget_limiter_without_redis_tracks_spend_locally_and_queue assert limiter.redis_increment_operation_queue == [] -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_none_without_fallbacks( - budget_limiter, -): - user_api_key = UserAPIKeyAuth(token="test-key", budget_fallbacks={}) - assert await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") is None - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_first_within_budget( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={"gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}}, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=1.0): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result == "gpt-4o-mini" - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_skips_exhausted_fallback( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={ - "gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}, - "claude-haiku": {"budget_limit": 100.0, "time_period": "1d"}, - }, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - - async def _spend_for_model(entity_type, entity_id, model, resolved): - return 150.0 if resolved.budget_model == "gpt-4o-mini" else 1.0 - - with patch.object( - budget_limiter, - "_get_spend_for_model_budget", - side_effect=_spend_for_model, - ): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result == "claude-haiku" - - -@pytest.mark.asyncio -async def test_get_fallback_model_within_budget_returns_none_when_chain_exhausted( - budget_limiter, -): - user_api_key = UserAPIKeyAuth( - token="test-key", - model_max_budget={ - "gpt-4o-mini": {"budget_limit": 100.0, "time_period": "1d"}, - "claude-haiku": {"budget_limit": 100.0, "time_period": "1d"}, - }, - budget_fallbacks={"gpt-4": ["gpt-4o-mini", "claude-haiku"]}, - ) - with patch.object(budget_limiter, "_get_spend_for_model_budget", return_value=150.0): - result = await budget_limiter.get_fallback_model_within_budget(user_api_key, "gpt-4") - assert result is None - - @pytest.mark.asyncio async def test_async_log_success_event_skips_redis_push_without_redis(budget_limiter): """When dual_cache has no Redis backend, do not await _push_in_memory_increments_to_redis."""