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

refactor(auth): decide the team member per-model budget fallback in AuthCore

Browse source 
The core now decides when an over-budget team member may fall back and walks
the router's fallback targets in order, skipping the requested model, through
its own TeamMemberStep and driver. The proxy handler only reads the targets
and vets each one, so _try_team_member_model_budget_fallback and its
per-model helper are gone

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Yujong Lee 2026-11-06 15:10:10 +09:00
parent 84851bc75f
commit 54122b4d89
7 changed files with 583 additions and 291 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/_auth/AGENTS.md
View file

@ -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

56
litellm/_auth/decision.py
View file

@ -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: ...

46
litellm/_auth/driver.py
View file

@ -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)

64
litellm/_auth/model_budgets.py
View file

@ -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),
)

10
litellm/_auth/python_core.py
View file

@ -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,

576
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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,

119
tests/unit/_auth/test_model_budgets.py
View file

@ -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