diff --git a/litellm/_auth/AGENTS.md b/litellm/_auth/AGENTS.md index eaa28600268..0c78570ae64 100644 --- a/litellm/_auth/AGENTS.md +++ b/litellm/_auth/AGENTS.md @@ -14,4 +14,5 @@ - 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 - `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 +- The team member per-model budget runs through `AuthCore.check_team_member_model_budgets`, which has its own step type (`TeamMemberStep`) and driver (`drive_team_member_model_budgets`). The core decides when an over-budget member may fall back and walks the router's fallback targets in order, skipping the requested model. The proxy handler (`_TeamMemberModelBudgetEffects`) reads those targets and vets each one (pricing, customer access, every per-model budget, model access and access group budgets) +- The team and end-user per-model budgets still run in the proxy after the core, as plain loops with no fallback diff --git a/litellm/_auth/decision.py b/litellm/_auth/decision.py index abd6b54d0fa..b002d28b701 100644 --- a/litellm/_auth/decision.py +++ b/litellm/_auth/decision.py @@ -205,6 +205,21 @@ class ModelBudgetScope: user_has_model_budgets: bool +@dataclass(frozen=True, slots=True) +class TeamMemberModelBudgetScope: + models: tuple[str, ...] + rewritable_model: str | None + has_model_budgets: bool + can_fall_back: bool + + +@dataclass(frozen=True, slots=True) +class FallbackTarget: + requested: str + model: str + served_model: str + + @dataclass(frozen=True, slots=True) class ResolveEndUser: prefetch_key: bool @@ -367,6 +382,45 @@ Effect: TypeAlias = IdentityEffect | BudgetEffect | ModelBudgetEffect Step: TypeAlias = Outcome | Effect +@dataclass(frozen=True, slots=True) +class ReadTeamMemberModelBudgetScope: + resume: "Callable[[TeamMemberModelBudgetScope], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class CheckTeamMemberModelBudget: + model: str + resume: "Callable[[ModelBudgetExceeded | None], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class ReadTeamMemberFallbacks: + model: str + resume: "Callable[[tuple[FallbackTarget, ...]], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class VetTeamMemberFallback: + target: FallbackTarget + resume: "Callable[[bool], TeamMemberStep]" + + +@dataclass(frozen=True, slots=True) +class RerouteTeamMemberRequest: + model: str + resume: "Callable[[], TeamMemberStep]" + + +TeamMemberEffect: TypeAlias = ( + ReadTeamMemberModelBudgetScope + | CheckTeamMemberModelBudget + | ReadTeamMemberFallbacks + | VetTeamMemberFallback + | RerouteTeamMemberRequest +) +TeamMemberStep: TypeAlias = Outcome | TeamMemberEffect + + class AuthCore(Protocol): def start(self, facts: RequestFacts, settings: Settings, now: datetime) -> Step: ... @@ -375,3 +429,5 @@ class AuthCore(Protocol): def check_model_budgets(self) -> Step: ... def check_user_model_budgets(self) -> Step: ... + + def check_team_member_model_budgets(self) -> TeamMemberStep: ... diff --git a/litellm/_auth/driver.py b/litellm/_auth/driver.py index b014dba46f8..e447d7c0d04 100644 --- a/litellm/_auth/driver.py +++ b/litellm/_auth/driver.py @@ -10,11 +10,13 @@ from litellm._auth.decision import ( BudgetEffect, CheckBudgetExemption, CheckKeyModelBudget, + CheckTeamMemberModelBudget, CheckUserModelBudget, Deny, Effect, EnforceModelAccess, EvictPresentedKey, + FallbackTarget, IdentityEffect, InspectKey, KeyFacts, @@ -32,12 +34,19 @@ from litellm._auth.decision import ( ReadModelBudgetScope, ReadProxySpend, ReadTeamMemberBudget, + ReadTeamMemberFallbacks, + ReadTeamMemberModelBudgetScope, + RerouteTeamMemberRequest, ResolveEndUser, ResolvePresentedKey, RewriteRequestModel, SendKeyBudgetAlerts, Step, TeamMemberBudget, + TeamMemberEffect, + TeamMemberModelBudgetScope, + TeamMemberStep, + VetTeamMemberFallback, ) @@ -53,6 +62,18 @@ class ModelBudgetHandler(Protocol): async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... +class TeamMemberModelBudgetHandler(Protocol): + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: ... + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: ... + + async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: ... + + async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: ... + + async def reroute_team_member_request(self, model: str) -> None: ... + + class EffectHandler(ModelBudgetHandler, Protocol): async def resolve_end_user(self, prefetch_key: bool) -> None: ... @@ -216,3 +237,28 @@ async def drive(step: Step, handler: EffectHandler) -> Outcome: return step case _: return await drive(await _answer(step, handler), handler) + + +async def _answer_team_member(step: TeamMemberEffect, handler: TeamMemberModelBudgetHandler) -> TeamMemberStep: + match step: + case ReadTeamMemberModelBudgetScope(resume=resume): + return resume(await handler.read_team_member_model_budget_scope()) + case CheckTeamMemberModelBudget(model=model, resume=resume): + return resume(await handler.check_team_member_model_budget(model)) + case ReadTeamMemberFallbacks(model=model, resume=resume): + return resume(await handler.read_team_member_fallbacks(model)) + case VetTeamMemberFallback(target=target, resume=resume): + return resume(await handler.vet_team_member_fallback(target)) + case RerouteTeamMemberRequest(model=model, resume=resume): + await handler.reroute_team_member_request(model) + return resume() + case _: + assert_never(step) + + +async def drive_team_member_model_budgets(step: TeamMemberStep, handler: TeamMemberModelBudgetHandler) -> Outcome: + match step: + case Allow() | Deny(): + return step + case _: + return await drive_team_member_model_budgets(await _answer_team_member(step, handler), handler) diff --git a/litellm/_auth/model_budgets.py b/litellm/_auth/model_budgets.py index 9e5ad2d9f3d..7bf2e109287 100644 --- a/litellm/_auth/model_budgets.py +++ b/litellm/_auth/model_budgets.py @@ -2,15 +2,25 @@ from collections.abc import Callable from typing import Final from litellm._auth.decision import ( + Allow, AuthorizeBudgetFallback, CheckKeyModelBudget, + CheckTeamMemberModelBudget, CheckUserModelBudget, Deny, + FallbackTarget, ModelBudgetExceeded, ModelBudgetScope, ReadModelBudgetScope, + ReadTeamMemberFallbacks, + ReadTeamMemberModelBudgetScope, + RerouteTeamMemberRequest, RewriteRequestModel, Step, + TeamMemberModelBudgetScope, + TeamMemberStep, + VetTeamMemberFallback, + WithinModelBudgets, ) @@ -85,3 +95,57 @@ def _check_user_budgets(scope: ModelBudgetScope, models: tuple[str, ...], then: model=models[0], resume=lambda exceeded: _check_user_budgets(scope, models[1:], then) if exceeded is None else Deny(exceeded), ) + + +def check_team_member_model_budgets() -> TeamMemberStep: + return ReadTeamMemberModelBudgetScope( + resume=lambda scope: ( + _check_team_member_budget_of(scope, scope.models, scope.rewritable_model) + if scope.has_model_budgets + else Allow(WithinModelBudgets()) + ) + ) + + +def _check_team_member_budget_of( + scope: TeamMemberModelBudgetScope, models: tuple[str, ...], rewritable_model: str | None +) -> TeamMemberStep: + if not models: + return Allow(WithinModelBudgets()) + model: Final = models[0] + rest: Final = models[1:] + + def reroute(target: FallbackTarget) -> TeamMemberStep: + return RerouteTeamMemberRequest( + model=target.model, resume=lambda: _check_team_member_budget_of(scope, rest, target.model) + ) + + def after(exceeded: ModelBudgetExceeded | None) -> TeamMemberStep: + if exceeded is None: + return _check_team_member_budget_of(scope, rest, rewritable_model) + if model != rewritable_model or not scope.can_fall_back: + return Deny(exceeded) + return ReadTeamMemberFallbacks( + model=model, + resume=lambda targets: _first_usable_fallback( + tuple(target for target in targets if target.model != model), + found=reroute, + exhausted=lambda: Deny(exceeded), + ), + ) + + return CheckTeamMemberModelBudget(model=model, resume=after) + + +def _first_usable_fallback( + targets: tuple[FallbackTarget, ...], + found: Callable[[FallbackTarget], TeamMemberStep], + exhausted: Callable[[], TeamMemberStep], +) -> TeamMemberStep: + if not targets: + return exhausted() + target: Final = targets[0] + return VetTeamMemberFallback( + target=target, + resume=lambda usable: found(target) if usable else _first_usable_fallback(targets[1:], found, exhausted), + ) diff --git a/litellm/_auth/python_core.py b/litellm/_auth/python_core.py index a4c7463e71e..c7d6210bdc6 100644 --- a/litellm/_auth/python_core.py +++ b/litellm/_auth/python_core.py @@ -47,13 +47,18 @@ from litellm._auth.decision import ( Step, TeamMemberBudget, TeamMemberBudgetExceeded, + TeamMemberStep, 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 +from litellm._auth.model_budgets import ( + check_model_budgets, + check_team_member_model_budgets, + check_user_model_budgets, +) _USER_AUTH_ROUTE: Final = "/user/auth" _VIRTUAL_KEY_PREFIX: Final = "sk-" @@ -98,6 +103,9 @@ class PythonCore: def check_user_model_budgets(self) -> Step: return check_user_model_budgets(then=lambda: Allow(WithinModelBudgets())) + def check_team_member_model_budgets(self) -> TeamMemberStep: + return check_team_member_model_budgets() + def _after_presented_key( presented: PresentedKey | None, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index dfb3d088794..ec253a04a4c 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -36,6 +36,7 @@ from litellm._auth.decision import ( AuthCore, CachedAdminKey, Deny, + FallbackTarget, Grant, KeyFacts, MasterKey, @@ -50,11 +51,12 @@ from litellm._auth.decision import ( Settings, Step, TeamMemberBudget, + TeamMemberModelBudgetScope, UserAuthRoute, VirtualKey, WithinModelBudgets, ) -from litellm._auth.driver import drive, drive_model_budgets +from litellm._auth.driver import drive, drive_model_budgets, drive_team_member_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 @@ -744,289 +746,6 @@ async def _is_within_all_fallback_model_budgets( return True -async def _check_team_member_model_budget_with_fallback( - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - models: Sequence[str], - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router | None, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> None: - team_member_model_max_budget: Final = valid_token.team_member_model_max_budget - if ( - valid_token.user_id is None - or valid_token.team_id is None - or not isinstance(team_member_model_max_budget, Mapping) - or not team_member_model_max_budget - ): - return - - user_id: Final = valid_token.user_id - team_id: Final = valid_token.team_id - for model_name in models: - await _check_team_member_model_budget_for_model( - user_id=user_id, - team_id=team_id, - team_member_model_max_budget=team_member_model_max_budget, - model_name=model_name, - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - request_data=request_data, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - project_object=project_object, - team_object=team_object, - end_user_object=end_user_object, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - - -async def _check_team_member_model_budget_for_model( - user_id: str, - team_id: str, - team_member_model_max_budget: Mapping[str, object], - model_name: str, - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router | None, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> None: - try: - await model_max_budget_limiter.is_team_member_within_model_budget( - user_id=user_id, - team_id=team_id, - team_member_model_max_budget=team_member_model_max_budget, - model=model_name, - ) - except litellm.BudgetExceededError: - if request_data.get("model") != model_name or llm_router is None or fallbacks_disabled(request_data): - raise - fallback_succeeded: Final = await _try_team_member_model_budget_fallback( - user_id=user_id, - team_id=team_id, - model_name=model_name, - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - request_data=request_data, - request=request, - llm_model_list=llm_model_list, - llm_router=llm_router, - project_object=project_object, - team_object=team_object, - end_user_object=end_user_object, - route=route, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, - ) - if not fallback_succeeded: - raise - - -async def _try_team_member_model_budget_fallback( - user_id: str, - team_id: str, - model_name: str, - valid_token: UserAPIKeyAuth, - model_max_budget_limiter: _EveryModelBudgetLimiter, - request_data: dict[str, object], # mutable-ok: fallback rewrites the request body - request: Request, - llm_model_list: Sequence[object] | None, - llm_router: litellm.Router, - project_object: LiteLLM_ProjectTableCachedObj | None = None, - team_object: LiteLLM_TeamTable | None = None, - end_user_object: LiteLLM_EndUserTable | None = None, - route: str = "", - prisma_client: PrismaClient | None = None, - user_api_key_cache: UserApiKeyCache | None = None, - proxy_logging_obj: ProxyLogging | None = None, -) -> bool: - if request.headers.get("content-type", "").lower().startswith( - "multipart/form-data" - ) or references_provider_scoped_resource(request_data): - return False - - from litellm.proxy import proxy_server - from litellm.proxy.proxy_server import proxy_config - - resolved_prisma_client: Final = prisma_client if prisma_client is not None else proxy_server.prisma_client - resolved_user_api_key_cache: Final = ( - user_api_key_cache if user_api_key_cache is not None else proxy_server.user_api_key_cache - ) - resolved_proxy_logging_obj: Final = ( - proxy_logging_obj if proxy_logging_obj is not None else proxy_server.proxy_logging_obj - ) - - try: - router_settings: Final[dict[str, object] | None] = ( - cast( # cast-ok: proxy settings return an unparameterized dict - dict[str, object] | None, - await proxy_config.get_hierarchical_router_settings( - user_api_key_dict=valid_token, - prisma_client=resolved_prisma_client, - proxy_logging_obj=resolved_proxy_logging_obj, - ), - ) - ) - except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. - return False - - configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None - router_fallbacks: Final[Sequence[object] | None] = ( - cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below - if isinstance(configured_fallbacks, list) - else _router_fallbacks(llm_router) - ) - if router_fallbacks is None: - return False - - fallback_models, _ = request_fallback_chain(list(router_fallbacks), request_data, model_name) - fallback_targets: Final[Sequence[object]] = fallback_models or () - fallback_candidates: Final[tuple[str, ...]] = tuple( - model for target in fallback_targets if (model := _budget_fallback_target_model(target)) is not None - ) - if not fallback_candidates: - return False - - model_group_aliases: Final[object] = ( - router_settings.get("model_group_alias") if router_settings is not None else None - ) - key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime - object, valid_token.aliases - ) - resolved_fallback_candidates: Final[tuple[tuple[str, str], ...]] = tuple( - ( - candidate, - target if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None else candidate, - ) - for candidate in fallback_candidates - ) - fallback_candidates_with_served_models: Final = ( - ( - fallback_candidate, - fallback_model, - resolve_served_model( - fallback_model, - litellm.model_alias_map, - key_aliases, - model_group_aliases, - ), - ) - for fallback_candidate, fallback_model in resolved_fallback_candidates - ) - - try: - team_membership: Final = await get_team_membership( - user_id=user_id, - team_id=team_id, - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - proxy_logging_obj=resolved_proxy_logging_obj, - ) - except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback - return False - - for fallback_candidate, fallback_model, served_model in fallback_candidates_with_served_models: - if fallback_model == model_name: - continue - try: - if ( - litellm.block_requests_for_models_without_pricing - and RouteChecks.is_llm_api_route(route=route) - and _unpriced_models_in_request(model=served_model, llm_router=llm_router) - ): - continue - if end_user_object is not None and end_user_object.models: - can_customer_access_model( - model=fallback_model, - end_user_object=end_user_object, - llm_router=llm_router, - valid_token=valid_token, - ) - if not is_model_cost_zero( - model=served_model, llm_router=llm_router - ) and not await _is_within_all_fallback_model_budgets( - valid_token=valid_token, - model_max_budget_limiter=model_max_budget_limiter, - models=tuple( - dict.fromkeys( - ( - fallback_candidate, - fallback_model, - served_model, - llm_router.registry().resolve(served_model), - ) - ) - ), - ): - continue - if not await _is_fallback_model_authorized( - model=fallback_model, - valid_token=valid_token, - llm_model_list=llm_model_list, - llm_router=llm_router, - prisma_client=prisma_client, - team_object=team_object, - team_membership=team_membership, - project_object=project_object, - refuse_agent_keys=True, - ): - continue - valid_token.matched_model_access_groups = [] # rebind-ok: clear prior groups before evaluating this candidate - matched_groups = await stamp_matched_model_access_groups( - model=fallback_model, - valid_token=valid_token, - team_object=team_object, - project_object=project_object, - llm_router=llm_router, - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - proxy_logging_obj=resolved_proxy_logging_obj, - team_membership=team_membership, - team_membership_loaded=True, - ) - valid_token.matched_model_access_groups = list(matched_groups) # rebind-ok: spend uses the candidate groups - if matched_groups: - await _model_access_group_max_budget_check( - matched_model_access_groups=list(matched_groups), - prisma_client=resolved_prisma_client, - user_api_key_cache=resolved_user_api_key_cache, - ) - except Exception: # noqa: BLE001 # A candidate check failure must not escape the fail-closed walk. - verbose_proxy_logger.warning( - "Skipping fallback candidate %s after a budget or authorization check error", - fallback_model.replace("\r", "").replace("\n", ""), - exc_info=True, - ) - continue - 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 True - return False - - async def _model_budget_verdict(check: Awaitable[bool]) -> ModelBudgetExceeded | None: try: with tracer.trace("litellm.proxy.auth.budget_checks"): @@ -1108,9 +827,7 @@ class _ModelBudgetEffects: ) 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 + _reroute_request(self._request_data, self._request, model) async def check_user_model_budget(self, model: str) -> ModelBudgetExceeded | None: user_id: Final = self._valid_token.user_id @@ -1141,6 +858,291 @@ async def _enforce_model_budgets(step: Step, effects: _ModelBudgetEffects) -> No return +def _reroute_request( + request_data: dict[str, object], # mutable-ok: the fallback model replaces the requested one in place + request: Request, + model: str, +) -> None: + rewrite_request_model(request_data=request_data, request=request, model=model) + if "path_params" in request.scope: + request.scope["path_params"]["model"] = model # rebind-ok: route handlers use the fallback model + + +def _budget_fallback_candidates(targets: Sequence[object]) -> tuple[str, ...]: + return tuple(model for target in targets if (model := _budget_fallback_target_model(target)) is not None) + + +class _TeamMemberModelBudgetEffects: + def __init__( + self, + *, + valid_token: UserAPIKeyAuth, + limiter: _EveryModelBudgetLimiter, + models: Sequence[str], + request_data: dict[str, object], # mutable-ok: fallback rewrites the request body + request: Request, + llm_model_list: Sequence[object] | None, + llm_router: litellm.Router | None, + project_object: LiteLLM_ProjectTableCachedObj | None, + team_object: LiteLLM_TeamTable | None, + end_user_object: LiteLLM_EndUserTable | None, + route: str, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache | None, + proxy_logging_obj: ProxyLogging | None, + ) -> None: + self._valid_token: Final = valid_token + self._limiter: Final = limiter + self._models: Final = tuple(models) + self._request_data: Final = request_data + self._request: Final = request + self._llm_model_list: Final = llm_model_list + self._llm_router: Final = llm_router + self._project_object: Final = project_object + self._team_object: Final = team_object + self._end_user_object: Final = end_user_object + self._route: Final = route + self._prisma_client: Final = prisma_client + self._user_api_key_cache: Final = user_api_key_cache + self._proxy_logging_obj: Final = proxy_logging_obj + + def _resolved_dependencies(self) -> tuple[PrismaClient | None, UserApiKeyCache, ProxyLogging]: + from litellm.proxy import proxy_server + + return ( + self._prisma_client if self._prisma_client is not None else proxy_server.prisma_client, + self._user_api_key_cache if self._user_api_key_cache is not None else proxy_server.user_api_key_cache, + self._proxy_logging_obj if self._proxy_logging_obj is not None else proxy_server.proxy_logging_obj, + ) + + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: + budget: Final = self._valid_token.team_member_model_max_budget + requested: Final = self._request_data.get("model") + return TeamMemberModelBudgetScope( + models=self._models, + rewritable_model=requested if isinstance(requested, str) else None, + has_model_budgets=self._valid_token.user_id is not None + and self._valid_token.team_id is not None + and isinstance(budget, Mapping) + and bool(budget), + can_fall_back=self._llm_router is not None and not fallbacks_disabled(self._request_data), + ) + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: + user_id: Final = self._valid_token.user_id + team_id: Final = self._valid_token.team_id + budget: Final = self._valid_token.team_member_model_max_budget + if user_id is None or team_id is None or not isinstance(budget, Mapping): + return None + return await _model_budget_verdict( + self._limiter.is_team_member_within_model_budget( + user_id=user_id, team_id=team_id, team_member_model_max_budget=budget, model=model + ) + ) + + async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: + llm_router: Final = self._llm_router + if ( + llm_router is None + or self._request.headers.get("content-type", "").lower().startswith("multipart/form-data") + or references_provider_scoped_resource(self._request_data) + ): + return () + from litellm.proxy.proxy_server import proxy_config + + prisma_client, _, proxy_logging_obj = self._resolved_dependencies() + try: + router_settings: Final[dict[str, object] | None] = ( + cast( # cast-ok: proxy settings return an unparameterized dict + dict[str, object] | None, + await proxy_config.get_hierarchical_router_settings( + user_api_key_dict=self._valid_token, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ), + ) + ) + except Exception: # noqa: BLE001 # Settings lookup failures must not authorize a fallback. + return () + + configured_fallbacks: Final[object] = router_settings.get("fallbacks") if router_settings is not None else None + router_fallbacks: Final[Sequence[object] | None] = ( + cast(Sequence[object], configured_fallbacks) # cast-ok: guarded by the list check below + if isinstance(configured_fallbacks, list) + else _router_fallbacks(llm_router) + ) + if router_fallbacks is None: + return () + fallback_models, _ = request_fallback_chain(list(router_fallbacks), self._request_data, model) + model_group_aliases: Final[object] = ( + router_settings.get("model_group_alias") if router_settings is not None else None + ) + key_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime + object, self._valid_token.aliases + ) + resolved: Final[tuple[tuple[str, str], ...]] = tuple( + ( + candidate, + target if (target := resolve_model_group_alias(model_group_aliases, candidate)) is not None else candidate, + ) + for candidate in _budget_fallback_candidates(fallback_models or ()) + ) + return tuple( + FallbackTarget( + requested=candidate, + model=fallback_model, + served_model=resolve_served_model( + fallback_model, litellm.model_alias_map, key_aliases, model_group_aliases + ), + ) + for candidate, fallback_model in resolved + ) + + async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: + llm_router: Final = self._llm_router + user_id: Final = self._valid_token.user_id + team_id: Final = self._valid_token.team_id + if llm_router is None or user_id is None or team_id is None: + return False + prisma_client, user_api_key_cache, proxy_logging_obj = self._resolved_dependencies() + try: + 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, + ) + except Exception: # noqa: BLE001 # membership lookup failures must not authorize a fallback + return False + try: + return await self._is_usable_fallback( + target, llm_router, team_membership, prisma_client, user_api_key_cache, proxy_logging_obj + ) + except Exception: # noqa: BLE001 # A candidate check failure must not escape the fail-closed walk. + verbose_proxy_logger.warning( + "Skipping fallback candidate %s after a budget or authorization check error", + target.model.replace("\r", "").replace("\n", ""), + exc_info=True, + ) + return False + + async def _is_usable_fallback( + self, + target: FallbackTarget, + llm_router: litellm.Router, + team_membership: LiteLLM_TeamMembership | None, + prisma_client: PrismaClient | None, + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + ) -> bool: + if ( + litellm.block_requests_for_models_without_pricing + and RouteChecks.is_llm_api_route(route=self._route) + and _unpriced_models_in_request(model=target.served_model, llm_router=llm_router) + ): + return False + if self._end_user_object is not None and self._end_user_object.models: + can_customer_access_model( + model=target.model, + end_user_object=self._end_user_object, + llm_router=llm_router, + valid_token=self._valid_token, + ) + if not is_model_cost_zero( + model=target.served_model, llm_router=llm_router + ) and not await _is_within_all_fallback_model_budgets( + valid_token=self._valid_token, + model_max_budget_limiter=self._limiter, + models=tuple( + dict.fromkeys( + ( + target.requested, + target.model, + target.served_model, + llm_router.registry().resolve(target.served_model), + ) + ) + ), + ): + return False + if not await _is_fallback_model_authorized( + model=target.model, + valid_token=self._valid_token, + llm_model_list=self._llm_model_list, + llm_router=llm_router, + prisma_client=self._prisma_client, + team_object=self._team_object, + team_membership=team_membership, + project_object=self._project_object, + refuse_agent_keys=True, + ): + return False + self._valid_token.matched_model_access_groups = [] + matched_groups: Final = await stamp_matched_model_access_groups( + model=target.model, + valid_token=self._valid_token, + team_object=self._team_object, + project_object=self._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, + ) + self._valid_token.matched_model_access_groups = list(matched_groups) + if matched_groups: + await _model_access_group_max_budget_check( + matched_model_access_groups=list(matched_groups), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + return True + + async def reroute_team_member_request(self, model: str) -> None: + _reroute_request(self._request_data, self._request, model) + + +async def _check_team_member_model_budget_with_fallback( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _EveryModelBudgetLimiter, + models: Sequence[str], + request_data: dict[str, object], # mutable-ok: fallback rewrites the request body + request: Request, + llm_model_list: Sequence[object] | None, + llm_router: litellm.Router | None, + project_object: LiteLLM_ProjectTableCachedObj | None = None, + team_object: LiteLLM_TeamTable | None = None, + end_user_object: LiteLLM_EndUserTable | None = None, + route: str = "", + prisma_client: PrismaClient | None = None, + user_api_key_cache: UserApiKeyCache | None = None, + proxy_logging_obj: ProxyLogging | None = None, +) -> None: + effects: Final = _TeamMemberModelBudgetEffects( + valid_token=valid_token, + limiter=model_max_budget_limiter, + models=models, + request_data=request_data, + request=request, + llm_model_list=llm_model_list, + llm_router=llm_router, + project_object=project_object, + team_object=team_object, + end_user_object=end_user_object, + route=route, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + match await drive_team_member_model_budgets(_AUTH_CORE.check_team_member_model_budgets(), effects): + case Deny(denial=denial): + raise denial_exception(denial, token=valid_token.token, route=route, jwt_auth_enabled=False, key=valid_token) + case Allow(): + return + + def _routing_selector_matches_claim( selector_value: object, claim_value: object, diff --git a/tests/unit/_auth/test_model_budgets.py b/tests/unit/_auth/test_model_budgets.py index 4370275a111..f5ad09a68c9 100644 --- a/tests/unit/_auth/test_model_budgets.py +++ b/tests/unit/_auth/test_model_budgets.py @@ -3,8 +3,18 @@ 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.decision import ( + Allow, + Deny, + FallbackTarget, + ModelBudgetExceeded, + ModelBudgetScope, + Outcome, + Step, + TeamMemberModelBudgetScope, + WithinModelBudgets, +) +from litellm._auth.driver import drive_model_budgets, drive_team_member_model_budgets from litellm._auth.python_core import PythonCore Call = tuple[str, str | None] @@ -161,3 +171,108 @@ async def test_the_user_only_check_never_reads_key_budgets(): ALLOWED, (("read", None), ("user", "gpt-4o")), ) + + +MEMBER_SCOPE: Final = TeamMemberModelBudgetScope( + models=("gpt-4o",), rewritable_model="gpt-4o", has_model_budgets=True, can_fall_back=True +) + + +def _target(model: str) -> FallbackTarget: + return FallbackTarget(requested=model, model=model, served_model=model) + + +@dataclass +class FakeTeamMemberRequest: + scope: TeamMemberModelBudgetScope + over: frozenset[str] = frozenset() + fallbacks: tuple[str, ...] = ("gpt-4o", "unusable", "gpt-4o-mini", "haiku") + unusable: frozenset[str] = frozenset({"unusable"}) + calls: tuple[Call, ...] = () + + def _record(self, name: str, model: str | None = None) -> None: + self.calls = (*self.calls, (name, model)) + + async def read_team_member_model_budget_scope(self) -> TeamMemberModelBudgetScope: + self._record("read") + return self.scope + + async def check_team_member_model_budget(self, model: str) -> ModelBudgetExceeded | None: + self._record("member", model) + return _exceeded("member", model) if model in self.over else None + + async def read_team_member_fallbacks(self, model: str) -> tuple[FallbackTarget, ...]: + self._record("fallbacks", model) + return tuple(_target(fallback) for fallback in self.fallbacks) + + async def vet_team_member_fallback(self, target: FallbackTarget) -> bool: + self._record("vet", target.model) + return target.model not in self.unusable + + async def reroute_team_member_request(self, model: str) -> None: + self._record("reroute", model) + + +async def _run_member(request: FakeTeamMemberRequest) -> tuple[Outcome, tuple[Call, ...]]: + outcome: Final = await drive_team_member_model_budgets(PythonCore().check_team_member_model_budgets(), request) + return outcome, request.calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "expected_calls"), + [ + (MEMBER_SCOPE, (("read", None), ("member", "gpt-4o"))), + (replace(MEMBER_SCOPE, has_model_budgets=False), (("read", None),)), + (replace(MEMBER_SCOPE, models=()), (("read", None),)), + ], + ids=["within_budget", "no_member_budgets", "no_models"], +) +async def test_a_team_member_within_budget_is_allowed( + scope: TeamMemberModelBudgetScope, expected_calls: tuple[Call, ...] +): + assert await _run_member(FakeTeamMemberRequest(scope)) == (ALLOWED, expected_calls) + + +@pytest.mark.asyncio +async def test_an_over_budget_team_member_reroutes_to_the_first_usable_fallback_other_than_itself(): + request: Final = FakeTeamMemberRequest(replace(MEMBER_SCOPE, models=("gpt-4o", "embed")), over=frozenset({"gpt-4o"})) + assert await _run_member(request) == ( + ALLOWED, + ( + ("read", None), + ("member", "gpt-4o"), + ("fallbacks", "gpt-4o"), + ("vet", "unusable"), + ("vet", "gpt-4o-mini"), + ("reroute", "gpt-4o-mini"), + ("member", "embed"), + ), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("scope", "unusable"), + [ + (replace(MEMBER_SCOPE, rewritable_model=None), frozenset()), + (replace(MEMBER_SCOPE, can_fall_back=False), frozenset()), + (MEMBER_SCOPE, frozenset({"unusable", "gpt-4o-mini", "haiku"})), + ], + ids=["nested_model", "fallbacks_unavailable", "no_usable_fallback"], +) +async def test_an_over_budget_team_member_without_a_usable_fallback_is_denied_with_its_own_error( + scope: TeamMemberModelBudgetScope, unusable: frozenset[str] +): + outcome, calls = await _run_member(FakeTeamMemberRequest(scope, over=frozenset({"gpt-4o"}), unusable=unusable)) + assert outcome == Deny(_exceeded("member", "gpt-4o")) + assert not any(name == "reroute" for name, _ in calls) + + +@pytest.mark.asyncio +async def test_a_later_model_is_not_rerouted_once_the_request_model_changed(): + scope: Final = replace(MEMBER_SCOPE, models=("gpt-4o", "gpt-4o")) + request: Final = FakeTeamMemberRequest(scope, over=frozenset({"gpt-4o"})) + outcome, calls = await _run_member(request) + assert outcome == Deny(_exceeded("member", "gpt-4o")) + assert [name for name, _ in calls].count("reroute") == 1