feat(proxy): per-model team member budgets with router fallback (#43887)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-10-10 20:33:55 +00:00 • committed by GitHub
parent 7383f59ad2
commit 5325c4a236
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
41 changed files with 5039 additions and 245 deletions

View file

@ -84,6 +84,7 @@ _PROPAGATED_METADATA_KEYS: Final = (
"user_api_key_model_max_budget",
"user_api_key_team_model_max_budget",
"user_api_key_user_model_max_budget",
"user_api_key_team_member_model_max_budget",
"user_api_key_end_user_model_max_budget",
"litellm_call_id",
"litellm_parent_otel_span",
@ -125,6 +126,23 @@ class _SummaryAcompletion(Protocol):
) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ...
class _SummaryModelBudgetLimiter(Protocol):
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
async def is_end_user_within_model_budget(
self,
end_user_id: str,
end_user_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
class _CreateRateLimitDescriptors(Protocol):
def __call__(
self,
@ -392,6 +410,42 @@ async def _check_summary_model_access(
return True
async def _check_summary_team_member_model_budget(
user_api_key_auth: "UserAPIKeyAuth",
model_max_budget_limiter: _SummaryModelBudgetLimiter,
summary_model: str,
) -> bool:
team_member_model_max_budget: Final = user_api_key_auth.team_member_model_max_budget
user_id: Final = user_api_key_auth.user_id
team_id: Final = user_api_key_auth.team_id
if (
not isinstance(team_member_model_max_budget, Mapping)
or not team_member_model_max_budget
or not isinstance(user_id, str)
or not isinstance(team_id, str)
):
return True
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=summary_model,
)
except litellm.BudgetExceededError:
return False
except Exception as e: # noqa: BLE001 # unexpected budget-check errors must deny compaction
verbose_logger.warning(
"compact_20260112: unexpected error during team member model-budget "
"check for summary_model=%s; denying: %s",
summary_model,
e,
)
return False
return True
async def _check_summary_model_budget(
user_api_key_auth: Optional["UserAPIKeyAuth"],
summary_model: str,
@ -407,9 +461,9 @@ async def _check_summary_model_budget(
per-model budget is configured.
Every scope is checked because the summary's spend is charged to every
scope: this file propagates the key, team, user and end-user budgets into the
subrequest's metadata, so skipping one of them would let compaction
increment a counter it can never be refused by.
scope: this file propagates the key, team, user, team-member and end-user
budgets into the subrequest's metadata, so skipping one of them would let
compaction increment a counter it can never be refused by.
"""
if user_api_key_auth is None:
return True
@ -497,7 +551,11 @@ async def _check_summary_model_budget(
)
return False
return True
return await _check_summary_team_member_model_budget(
user_api_key_auth=user_api_key_auth,
model_max_budget_limiter=model_max_budget_limiter,
summary_model=summary_model,
)
def _without_parallel_request_gauges(

View file

@ -2106,6 +2106,7 @@ class NewTeamRequest(TeamBase):
team_member_tpm_limit: int | None = None # allow user to set TPM limit for all team members
team_member_key_duration: str | None = None # e.g. "1d", "1w", "1m"
team_member_budget_duration: str | None = None # e.g. "30d", "1mo"
team_member_model_max_budget: GenericBudgetConfigType | None = None
allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None
enforced_batch_output_expires_after: dict | None = None
enforced_file_expires_after: dict | None = None
@ -2165,6 +2166,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
require_trace_id: bool | None = None
team_member_budget: float | None = None
team_member_budget_duration: str | None = None
team_member_model_max_budget: GenericBudgetConfigType | None = None
team_member_rpm_limit: int | None = None
team_member_tpm_limit: int | None = None
team_member_key_duration: str | None = None
@ -3363,6 +3365,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob
# and validating it here would make one malformed row fail auth outright.
# resolve_model_budget validates the single entry a request actually needs.
user_model_max_budget: Mapping[str, object] | None = None
team_member_model_max_budget: Mapping[str, object] | None = Field(default=None, exclude=True)
request_route: str | None = None
is_session_token: bool = False
# Server-only marker set exclusively by the MCP gateway admission path
@ -4651,6 +4654,13 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest):
default=None,
description="UTC expiry for temp_budget_increase",
)
model_max_budget: GenericBudgetConfigType | None = Field(
default=None,
description=(
"Per-model spend caps for this team member, each with its own budget_duration. "
"Overrides the team's default per-model member budget. Pass an empty dict to fall back to the team default."
),
)
@model_validator(mode="after")
def validate_temp_budget(self) -> "TeamMemberUpdateRequest":
@ -4669,6 +4679,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse):
allowed_models: list[str] | None = None
temp_budget_increase: float | None = None
temp_budget_expiry: datetime | None = None
model_max_budget: GenericBudgetConfigType | None = None
class TeamModelAddRequest(LiteLLMBaseModel):

View file

@ -1739,6 +1739,7 @@ async def get_team_member_default_budget(
budget_id: str,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
raise_on_lookup_error: bool = False,
) -> LiteLLM_BudgetTable | None:
"""
Fetches the team-level default per-member budget referenced by team.metadata["team_member_budget_id"].
@ -1773,6 +1774,8 @@ async def get_team_member_default_budget(
)
except Exception:
verbose_proxy_logger.exception("Error fetching team-default member budget %s", budget_id)
if raise_on_lookup_error:
raise
return None
if budget_record is None:

View file

@ -11,13 +11,13 @@ attempted, mirroring `fallback_model_access.py`. It deliberately leaves the prim
a zero-cost model is never blocked by budget, and only the paid fallback is refused. On by default;
set `general_settings.enforce_fallback_budget: false` to restore the unguarded behaviour.
Scope: the key's and the user's `max_budget`. Not covered yet, and each needs a read-only evaluation
path before it can be: team, team-member, end-user, org, global and per-model budgets, whose
auth-path functions enforce rather than report (they raise), so reusing them would fire threshold
alerts and take spend reservations for a target that is then skipped; and the key's rolling
`budget_limits` windows, whose accumulated spend lives only in per-window counters
(`spend:key:{token}:window:{budget_duration}`), so enforcing them means more counter reads on the
fallback path rather than reusing state auth already loaded.
Scope: the key's and the user's `max_budget`, plus the team member's per-model caps. Not covered yet,
and each needs a read-only evaluation path before it can be: team and team-member overall, end-user,
org, global and other per-model budgets, whose auth-path functions enforce rather than report (they
raise), so reusing them would fire threshold alerts and take spend reservations for a target that is
then skipped; and the key's rolling `budget_limits` windows, whose accumulated spend lives only in
per-window counters (`spend:key:{token}:window:{budget_duration}`), so enforcing them means more
counter reads on the fallback path rather than reusing state auth already loaded.
Two known limitations of that narrow scope, both shared with `fallback_model_access.py`:
@ -33,10 +33,11 @@ Two known limitations of that narrow scope, both shared with `fallback_model_acc
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from typing import Final
from typing import Final, Protocol
from pydantic import ValidationError
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module exports
@ -55,6 +56,16 @@ class _FallbackBudgetSettings(LiteLLMBaseModel):
enforce_fallback_budget: bool = True
class _TeamMemberModelBudgetLimiter(Protocol):
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None:
try:
return _RequestMetadata.model_validate(metadata).user_api_key_auth
@ -103,7 +114,13 @@ async def _counter_spend(counter_key: str, fallback_spend: float, max_budget: fl
)
async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKeyAuth, llm_router: Router) -> bool:
async def is_token_within_budget_for_model(
*,
model: str,
valid_token: UserAPIKeyAuth,
llm_router: Router,
team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None,
) -> bool:
"""
True when the key and the user behind it can still pay for `model`.
@ -139,6 +156,23 @@ async def is_token_within_budget_for_model(*, model: str, valid_token: UserAPIKe
if user_spend >= user_budget:
return False
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
if (
team_member_model_budget_limiter is not None
and team_member_model_max_budget
and isinstance(valid_token.user_id, str)
and isinstance(valid_token.team_id, str)
):
try:
await team_member_model_budget_limiter.is_team_member_within_model_budget(
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_member_model_max_budget=team_member_model_max_budget,
model=model,
)
except litellm.BudgetExceededError:
return False
return True
@ -151,6 +185,7 @@ class RouterFallbackBudgetCheck:
"""
is_enforced: Callable[[], bool]
team_member_model_budget_limiter: _TeamMemberModelBudgetLimiter | None = None
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: Router) -> bool:
if not self.is_enforced():
@ -158,11 +193,20 @@ class RouterFallbackBudgetCheck:
valid_token: Final = _user_api_key_auth_from_request(request_kwargs)
if valid_token is None:
return True
try:
return await is_token_within_budget_for_model(model=model, valid_token=valid_token, llm_router=llm_router)
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller
verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
return False
budget_models: Final = tuple(dict.fromkeys((model, llm_router.routable_model_group(model))))
for budget_model in budget_models:
try:
if not await is_token_within_budget_for_model(
model=budget_model,
valid_token=valid_token,
llm_router=llm_router,
team_member_model_budget_limiter=self.team_member_model_budget_limiter,
):
return False
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the caller
verbose_proxy_logger.warning("Skipping fallback to model=%s: budget lookup failed: %s", model, e)
return False
return True
router_fallback_budget_check: Final = RouterFallbackBudgetCheck(is_enforced=_enforced_by_general_settings)

View file

@ -11,12 +11,11 @@ import asyncio
import fnmatch
import re
import secrets
from collections.abc import Mapping
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from typing import Final, NamedTuple, Protocol, Union, cast
import fastapi
import orjson
from fastapi import HTTPException, Request, WebSocket, status
from fastapi.security.api_key import APIKeyHeader
from starlette.exceptions import WebSocketException
@ -51,13 +50,18 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
_get_user_role, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_is_model_cost_zero, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_is_user_proxy_admin, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_model_access_group_max_budget_check, # pyright: ignore[reportPrivateUsage] # reuse common-check group budget enforcement
_team_member_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_unpriced_models_in_request, # pyright: ignore[reportPrivateUsage] # reuse common-check pricing enforcement
_virtual_key_max_budget_alert_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_virtual_key_max_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
_virtual_key_soft_budget_check, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
cache_key_object,
can_customer_access_model,
can_key_call_model,
can_object_call_model,
can_project_access_model,
can_team_access_model,
check_end_user_budget,
common_checks,
delete_cache_key_object,
@ -67,6 +71,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
get_object_permission,
get_org_object_for_request,
get_project_object,
get_team_member_default_budget,
get_team_membership,
get_team_object,
get_user_object,
@ -76,9 +81,12 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
is_user_proxy_admin,
is_valid_fallback_model,
jwt_key_mapping_cache_key,
key_access_group_grants_model,
key_model_aliases_for_auth_check,
resolve_and_validate_end_user_id,
resolve_default_end_user_budget,
route_skips_budget_checks,
stamp_matched_model_access_groups,
team_member_max_budget_alert_check,
virtual_key_max_budget_alert_check,
virtual_key_max_budget_check,
@ -169,7 +177,13 @@ from litellm.proxy.utils import (
)
from litellm.repositories.table_repositories import JWTKeyMappingRepository, TeamMembershipRepository
from litellm.repositories.verification_token_repository import VerificationTokenRepository
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
from litellm.router_utils.fallback_event_handlers import (
PROVIDER_SCOPED_RESOURCE_KEYS,
fallback_lookup_groups,
fallbacks_disabled_for_request,
get_fallback_model_group_for_lookup_groups,
)
from litellm.secret_managers.main import get_secret_bool
from litellm.types.services import ServiceTypes
@ -332,6 +346,18 @@ def _get_model_names_for_budget_checks(
return model
def _budget_fallback_target_model(target: object) -> str | None:
if isinstance(target, str):
return target
if not isinstance(target, Mapping):
return None
target_mapping: Final = cast(Mapping[str, object], target) # cast-ok: isinstance cannot type the container's items
if set(target_mapping) != {"model"}:
return None
model: Final = target_mapping.get("model")
return model if isinstance(model, str) else None
class _KeyModelBudgetLimiter(Protocol):
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
@ -354,6 +380,34 @@ class _TeamModelBudgetLimiter(Protocol):
) -> bool: ...
class _EveryModelBudgetLimiter(Protocol):
async def is_key_within_model_budget(self, user_api_key_dict: UserAPIKeyAuth, model: str) -> bool: ...
async def is_user_within_model_budget(
self, user_id: str, user_model_max_budget: Mapping[str, object], model: str
) -> bool: ...
async def is_end_user_within_model_budget(
self, end_user_id: str, end_user_model_max_budget: Mapping[str, object], model: str
) -> bool: ...
async def is_team_within_model_budget(
self,
team_id: str,
team_model_max_budget: Mapping[str, object],
key_model_max_budget: Mapping[str, object] | None,
model: str,
) -> bool: ...
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool: ...
class _TokenTeamModels(Protocol):
@property
def team_models(self) -> list[str]: ...
@ -374,6 +428,30 @@ def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
return valid_token.team_models
class _TokenModelBudgetFields(Protocol):
@property
def model_max_budget(self) -> Mapping[str, object] | None: ...
@property
def end_user_model_max_budget(self) -> Mapping[str, object] | None: ...
@property
def team_model_aliases(self) -> dict[str, str] | None: ... # mutable-ok: mirrors _can_object_call_model's contract
def _token_model_budget_fields(valid_token: _TokenModelBudgetFields) -> _TokenModelBudgetFields:
return valid_token
class _RouterFallbacks(Protocol):
@property
def fallbacks(self) -> Sequence[object] | None: ...
def _router_fallbacks(router: _RouterFallbacks) -> Sequence[object] | None:
return router.fallbacks
async def _read_user_model_max_budget(
user_id: str | None,
prisma_client: PrismaClient | None,
@ -443,13 +521,455 @@ async def _check_team_model_budget(
)
async def _is_within_every_model_budget(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
model: str,
) -> bool:
budget_fields: Final = _token_model_budget_fields(valid_token)
key_model_max_budget: Final = budget_fields.model_max_budget
user_model_max_budget: Final = valid_token.user_model_max_budget
end_user_model_max_budget: Final = budget_fields.end_user_model_max_budget
team_model_max_budget: Final = valid_token.team_model_max_budget
team_member_model_max_budget: Final = valid_token.team_member_model_max_budget
try:
if valid_token.token is not None and isinstance(key_model_max_budget, Mapping) and key_model_max_budget:
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=valid_token,
model=model,
)
if valid_token.user_id is not None and isinstance(user_model_max_budget, Mapping) and user_model_max_budget:
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,
)
if (
valid_token.end_user_id is not None
and isinstance(end_user_model_max_budget, Mapping)
and end_user_model_max_budget
):
await model_max_budget_limiter.is_end_user_within_model_budget(
end_user_id=valid_token.end_user_id,
end_user_model_max_budget=end_user_model_max_budget,
model=model,
)
if valid_token.team_id is not None and isinstance(team_model_max_budget, Mapping) and team_model_max_budget:
await model_max_budget_limiter.is_team_within_model_budget(
team_id=valid_token.team_id,
team_model_max_budget=team_model_max_budget,
key_model_max_budget=key_model_max_budget,
model=model,
)
if (
valid_token.user_id is not None
and valid_token.team_id is not None
and isinstance(team_member_model_max_budget, Mapping)
and team_member_model_max_budget
):
await model_max_budget_limiter.is_team_member_within_model_budget(
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_member_model_max_budget=team_member_model_max_budget,
model=model,
)
except litellm.BudgetExceededError:
return False
return True
async def _is_fallback_model_authorized(
model: str,
valid_token: UserAPIKeyAuth,
llm_model_list: Sequence[object] | None,
llm_router: litellm.Router | None,
prisma_client: PrismaClient | None = None,
team_membership: LiteLLM_TeamMembership | None = None,
project_object: LiteLLM_ProjectTableCachedObj | None = None,
refuse_agent_keys: bool = False,
team_object: LiteLLM_TeamTable | None = None,
) -> bool:
if refuse_agent_keys and valid_token.agent_id is not None:
return False
try:
await can_key_call_model(
model=model,
llm_model_list=llm_model_list,
valid_token=valid_token,
llm_router=llm_router,
prisma_client=prisma_client,
)
if team_object is not None:
try:
await can_team_access_model(
model=model,
team_object=team_object,
llm_router=llm_router,
team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
prisma_client=prisma_client,
)
except ProxyException as team_denial:
if team_denial.type != ProxyErrorTypes.team_model_access_denied:
raise
if not await key_access_group_grants_model(
model=model,
valid_token=valid_token,
team_object=team_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
):
raise
else:
team_models: Final = _token_team_models(valid_token)
if team_models:
can_object_call_model(
model=model,
llm_router=llm_router,
models=team_models,
team_model_aliases=_token_model_budget_fields(valid_token).team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
member_allowed_models: Final = (
team_membership.litellm_budget_table.allowed_models
if team_membership is not None and team_membership.litellm_budget_table is not None
else None
)
if member_allowed_models:
can_object_call_model(
model=model,
llm_router=llm_router,
models=member_allowed_models,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
if project_object is not None and len(project_object.models) > 0:
can_project_access_model(
model=model,
project_object=project_object,
llm_router=llm_router,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
)
except ProxyException:
return False
return True
async def _is_within_all_fallback_model_budgets(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _EveryModelBudgetLimiter,
models: Sequence[str],
) -> bool:
for model in models:
if not await _is_within_every_model_budget(
valid_token=valid_token,
model_max_budget_limiter=model_max_budget_limiter,
model=model,
):
return False
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_for_request(kwargs=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 any(
request_data.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS
):
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, _ = get_fallback_model_group_for_lookup_groups(
fallbacks=list(router_fallbacks),
lookup_groups=fallback_lookup_groups(kwargs=request_data, model_group=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.routable_model_group(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 _check_key_model_budget_with_fallback(
valid_token: UserAPIKeyAuth,
model_max_budget_limiter: _KeyModelBudgetLimiter,
model_name: str,
request_data: dict,
request_data: dict[str, object], # mutable-ok: fallback rewrites the request body
request: Request,
llm_model_list: list | None = None,
llm_model_list: Sequence[object] | None = None,
llm_router: litellm.Router | None = None,
) -> None:
"""
@ -489,32 +1009,16 @@ async def _check_key_model_budget_with_fallback(
)
if fallback_model is None:
raise e
try:
await can_key_call_model(
model=fallback_model,
llm_model_list=llm_model_list,
valid_token=valid_token,
llm_router=llm_router,
)
if valid_token.team_models:
can_object_call_model(
model=fallback_model,
llm_router=llm_router,
models=valid_token.team_models,
team_model_aliases=valid_token.team_model_aliases,
team_id=valid_token.team_id,
key_model_aliases=key_model_aliases_for_auth_check(valid_token),
object_type="team",
)
except ProxyException:
raise e
request_data["model"] = fallback_model
safe_set_request_parsed_body(request=request, parsed_body=request_data)
request._json = request_data # pyright: ignore[reportPrivateUsage] # Starlette JSON cache
request._body = orjson.dumps(request_data) # pyright: ignore[reportPrivateUsage] # Starlette body cache
path_params: Final = request.scope.get("path_params")
if isinstance(path_params, dict) and "model" in path_params:
path_params["model"] = fallback_model
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
def _get_bearer_token_or_received_api_key(api_key: str) -> str:
@ -2543,7 +3047,9 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
model_name=model_name,
request_data=request_data,
request=request,
llm_model_list=llm_model_list,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
llm_router=llm_router,
)
@ -2836,6 +3342,79 @@ def is_no_auth_dev_mode(master_key: str | None, general_settings: Mapping[str, o
)
async def _load_team_member_default_model_budget(
user_api_key_auth_obj: UserAPIKeyAuth,
team_object: LiteLLM_TeamTable | None,
prisma_client: PrismaClient | None,
user_api_key_cache: UserApiKeyCache,
skip_budget_checks: bool,
proxy_logging_obj: ProxyLogging | None,
) -> None:
team_metadata_value: Final[object] = getattr(team_object, "metadata", None) if team_object is not None else None
team_metadata: Final[Mapping[str, object] | None] = (
cast( # cast-ok: team metadata is a string-keyed JSON object
Mapping[str, object], team_metadata_value
)
if isinstance(team_metadata_value, Mapping)
else None
)
team_member_budget_id: Final[object | None] = (
team_metadata.get("team_member_budget_id") if isinstance(team_metadata, Mapping) else None
)
if (
not isinstance(user_api_key_auth_obj.user_id, str)
or not isinstance(user_api_key_auth_obj.team_id, str)
or prisma_client is None
):
return
try:
membership: Final = await get_team_membership(
user_id=user_api_key_auth_obj.user_id,
team_id=user_api_key_auth_obj.team_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
member_budget_row: Final = membership.litellm_budget_table if membership is not None else None
member_budget_value: Final[object] = (
getattr(member_budget_row, "model_max_budget", None) if member_budget_row is not None else None
)
if (
membership is not None
and membership.budget_id is not None
and membership.budget_id != team_member_budget_id
and isinstance(member_budget_value, Mapping)
and member_budget_value
):
user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries member budget downstream
member_budget_value
)
return
if not isinstance(team_member_budget_id, str):
return
team_member_default_budget: Final = await get_team_member_default_budget(
budget_id=team_member_budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
raise_on_lookup_error=True,
)
except Exception:
if not skip_budget_checks:
raise
verbose_proxy_logger.warning("Skipping team-member model budget lookup after a lookup error", exc_info=True)
return
default_budget_value: Final[object] = (
getattr(team_member_default_budget, "model_max_budget", None)
if team_member_default_budget is not None
else None
)
user_api_key_auth_obj.team_member_model_max_budget = ( # rebind-ok: auth context carries default budget downstream
default_budget_value if isinstance(default_budget_value, Mapping) else None
)
@tracer.wrap()
async def run_centralized_common_checks(
user_api_key_auth_obj: UserAPIKeyAuth,
@ -2864,6 +3443,7 @@ async def run_centralized_common_checks(
from litellm.proxy.proxy_server import (
general_settings,
litellm_proxy_admin_name,
llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
llm_router,
master_key,
model_max_budget_limiter,
@ -3124,6 +3704,17 @@ async def run_centralized_common_checks(
llm_router=llm_router,
team_id=user_api_key_auth_obj.team_id,
)
typed_team_object: Final[LiteLLM_TeamTable | None] = (
team_object if isinstance(team_object, LiteLLM_TeamTable) else None
)
await _load_team_member_default_model_budget(
user_api_key_auth_obj=user_api_key_auth_obj,
team_object=typed_team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
skip_budget_checks=skip_budget_checks or route_skips_budget_checks(route=route),
proxy_logging_obj=proxy_logging_obj,
)
# Pin the metadata variable name (litellm_metadata vs metadata) before
# any tag merge runs. Without this, header tags from
@ -3164,19 +3755,42 @@ async def run_centralized_common_checks(
project_object=project_object,
)
if not skip_budget_checks:
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=user_api_key_auth_obj.team_id,
)
)
await _check_team_model_budget(
valid_token=user_api_key_auth_obj,
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=user_api_key_auth_obj.team_id,
)
),
models=models,
)
typed_project_object: Final[LiteLLM_ProjectTableCachedObj | None] = (
project_object if isinstance(project_object, LiteLLM_ProjectTableCachedObj) else None
)
if not route_skips_budget_checks(route=route):
await _check_team_member_model_budget_with_fallback(
valid_token=user_api_key_auth_obj,
model_max_budget_limiter=model_max_budget_limiter,
models=models,
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,
project_object=typed_project_object,
team_object=typed_team_object,
end_user_object=end_user_object if isinstance(end_user_object, LiteLLM_EndUserTable) else None,
route=route,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
await _reserve_budget_after_common_checks(
user_api_key_auth_obj=user_api_key_auth_obj,
@ -3184,7 +3798,10 @@ async def run_centralized_common_checks(
request_data=request_data,
route=route,
llm_router=llm_router,
team_object=team_object,
team_object=cast( # cast-ok: preserve plain team-table objects through reservation
LiteLLM_TeamTableCachedObj | None,
team_object,
),
user_object=user_object,
end_user_id=end_user_id,
end_user_object=end_user_object,
@ -3855,13 +4472,13 @@ async def enforce_key_and_fallback_model_access(
)
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access
_enforce_key_and_fallback_model_access: Final = enforce_key_and_fallback_model_access # pyright: ignore[reportUnknownVariableType] # existing tests use this alias
async def _run_post_custom_auth_checks(
valid_token: UserAPIKeyAuth,
request: Request,
request_data: dict,
request_data: dict[str, object],
route: str,
parent_otel_span: Span | None,
) -> UserAPIKeyAuth:
@ -3956,7 +4573,9 @@ async def _run_post_custom_auth_checks(
model_name=model_name,
request_data=request_data,
request=request,
llm_model_list=llm_model_list,
llm_model_list=cast( # cast-ok: router model lists are deployment sequences
Sequence[object] | None, llm_model_list
),
llm_router=llm_router,
)

View file

@ -83,6 +83,7 @@ from litellm.litellm_core_utils.streaming_handler import (
from litellm.proxy._types import LiteLLMRoutes, ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.auth.auth_checks import (
can_key_call_resolved_model,
is_model_cost_zero,
request_skips_budget_checks,
tag_max_budget_check_for_tags,
)
@ -127,7 +128,7 @@ from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
)
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.router_utils.common_utils import resolve_model_group_alias
from litellm.router_utils.common_utils import resolve_model_group_alias, resolve_served_model
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.router import RouterRateLimitError
from litellm.types.router_weights import validate_router_weights
@ -250,6 +251,7 @@ from litellm.proxy.litellm_pre_call_utils import (
add_litellm_data_to_request,
refresh_proxy_server_request_body_snapshot,
reject_url_valued_destination,
sanitize_for_log,
)
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.types.utils import (
@ -737,6 +739,7 @@ async def _resolve_per_request_model_group_alias(
_REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None)
_ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
def _request_model(data: Mapping[str, object]) -> str | list[str] | None:
@ -1733,6 +1736,69 @@ def _timing_values(
return getattr(logging_obj, "response_timing_metrics", None) or {}
async def _team_member_within_fallback_model_budget(
user_api_key_dict: UserAPIKeyAuth,
model: str,
llm_router: Router | None,
proxy_config: ProxyConfig | None,
proxy_logging_obj: ProxyLogging,
) -> bool:
member_caps: Final = user_api_key_dict.team_member_model_max_budget
user_id: Final = user_api_key_dict.user_id
team_id: Final = user_api_key_dict.team_id
if not isinstance(member_caps, Mapping) or not member_caps or user_id is None or team_id is None:
return True
from litellm.proxy.proxy_server import model_max_budget_limiter
try:
from litellm.proxy.proxy_server import prisma_client
router_settings: Final = (
_ROUTER_SETTINGS_ADAPTER.validate_python(
await proxy_config.get_hierarchical_router_settings( # pyright: ignore[reportUnknownMemberType] # legacy bare-dict return
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
)
if proxy_config is not None and llm_router is not None
else None
)
raw_aliases: Final[object] = cast( # cast-ok: auth aliases are dynamically typed at runtime
object, user_api_key_dict.aliases
)
served_model: Final = resolve_served_model(
model,
litellm.model_alias_map,
raw_aliases,
router_settings.get("model_group_alias") if router_settings is not None else None,
)
if is_model_cost_zero(model=served_model, llm_router=llm_router):
return True
budget_models: Final = (
tuple(dict.fromkeys((model, served_model, llm_router.routable_model_group(served_model))))
if llm_router is not None
else tuple(dict.fromkeys((model, served_model)))
)
for budget_model in budget_models:
await model_max_budget_limiter.is_team_member_within_model_budget(
user_id=user_id,
team_id=team_id,
team_member_model_max_budget=member_caps,
model=budget_model,
)
except litellm.BudgetExceededError:
return False
except Exception as e: # noqa: BLE001 # fail closed: a spend lookup failure must not bill the member
verbose_proxy_logger.warning(
"Skipping rate limit fallback to model=%s: member budget lookup failed: %s",
sanitize_for_log(model),
sanitize_for_log(e),
)
return False
return True
class ProxyBaseLLMRequestProcessing:
def __init__(self, data: dict):
self.data = data
@ -2371,6 +2437,14 @@ class ProxyBaseLLMRequestProcessing:
for fallback_model in fallback_models:
if fallback_model == original_model:
continue
if not await _team_member_within_fallback_model_budget(
user_api_key_dict=user_api_key_dict,
model=fallback_model,
llm_router=llm_router,
proxy_config=proxy_config,
proxy_logging_obj=proxy_logging_obj,
):
continue
self.data = independent_snapshot(pristine)
self.data["model"] = fallback_model
try:

View file

@ -4,7 +4,7 @@ import time
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from types import MappingProxyType
from typing import Final
from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing
from openai.types import Batch
@ -25,6 +25,7 @@ VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX: Final = "virtual_key_spend"
END_USER_SPEND_CACHE_KEY_PREFIX: Final = "end_user_model_spend"
USER_SPEND_CACHE_KEY_PREFIX: Final = "user_model_spend"
TEAM_SPEND_CACHE_KEY_PREFIX: Final = "team_model_spend"
TEAM_MEMBER_SPEND_CACHE_KEY_PREFIX: Final = "team_member_model_spend"
_SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
{
@ -32,6 +33,7 @@ _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType(
Litellm_EntityType.USER: USER_SPEND_CACHE_KEY_PREFIX,
Litellm_EntityType.END_USER: END_USER_SPEND_CACHE_KEY_PREFIX,
Litellm_EntityType.TEAM: TEAM_SPEND_CACHE_KEY_PREFIX,
Litellm_EntityType.TEAM_MEMBER: TEAM_MEMBER_SPEND_CACHE_KEY_PREFIX,
}
)
@ -45,6 +47,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType(
Litellm_EntityType.USER: "user_model_budget_start_time",
Litellm_EntityType.END_USER: "end_user_budget_start_time",
Litellm_EntityType.TEAM: "team_model_budget_start_time",
Litellm_EntityType.TEAM_MEMBER: "team_member_model_budget_start_time",
}
)
@ -63,6 +66,10 @@ class ResolvedModelBudget:
budget_config: BudgetConfig
def team_member_budget_entity_id(user_id: str, team_id: str) -> str:
return f"{len(user_id)}:{user_id}:{team_id}"
def model_budget_spend_cache_key(
entity_type: Litellm_EntityType,
entity_id: str | None,
@ -290,6 +297,18 @@ def _resolve_entity_model_budgets(
)
def _metadata_value(metadata: object, key: str) -> object | None:
metadata_mapping: Final[Mapping[str, object]] = cast( # cast-ok: callback metadata has string keys
Mapping[str, object], metadata
)
return metadata_mapping.get(key)
def _metadata_string_value(metadata: object, key: str) -> str | None:
value: Final = _metadata_value(metadata, key)
return value if isinstance(value, str) else None
class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
"""
Handles budgets for model + virtual key
@ -410,6 +429,24 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}",
)
@with_service_target("model_budgets")
async def is_team_member_within_model_budget(
self,
user_id: str,
team_id: str,
team_member_model_max_budget: Mapping[str, object],
model: str,
) -> bool:
return await self._is_entity_within_model_budget(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=team_member_budget_entity_id(user_id=user_id, team_id=team_id),
model_max_budget=team_member_model_max_budget,
model=model,
exceeded_message=(
f"LiteLLM Team Member: user={user_id}, team={team_id}, exceeded budget for model={model}"
),
)
async def _is_entity_within_model_budget(
self,
entity_type: Litellm_EntityType,
@ -523,6 +560,19 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
response_cost: Final[float] = standard_logging_payload.get("response_cost", 0)
key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget")
team_member_user_id: Final = _metadata_string_value(
payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
"user_api_key_user_id",
)
team_member_team_id: Final = _metadata_string_value(
payload_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
"user_api_key_team_id",
)
team_member_entity_id: Final = (
team_member_budget_entity_id(user_id=team_member_user_id, team_id=team_member_team_id)
if team_member_user_id is not None and team_member_team_id is not None
else None
)
entity_budgets: Final = (
(
Litellm_EntityType.KEY,
@ -553,6 +603,14 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
standard_logging_payload.get("end_user") or payload_metadata.get("user_api_key_end_user_id"),
_metadata.get("user_api_key_end_user_model_max_budget"),
),
(
Litellm_EntityType.TEAM_MEMBER,
team_member_entity_id,
_metadata_value(
_metadata, # pyright: ignore[reportUnknownArgumentType] # callback metadata is partially untyped
"user_api_key_team_member_model_max_budget",
),
),
)
resolved_budgets: Final = _resolve_entity_model_budgets(model=model, entity_budgets=entity_budgets)

View file

@ -2508,6 +2508,9 @@ async def add_litellm_data_to_request(
data[_metadata_variable_name]["user_api_key_user_max_budget"] = user_api_key_dict.user_max_budget
user_model_budget: Final = user_api_key_dict.user_model_max_budget
data[_metadata_variable_name]["user_api_key_user_model_max_budget"] = user_model_budget # rebind-ok: out-param
data[_metadata_variable_name][ # rebind-ok: adds member budget for spend tracking
"user_api_key_team_member_model_max_budget"
] = user_api_key_dict.team_member_model_max_budget
data[_metadata_variable_name].update(carried_budget_metadata(user_api_key_dict))
data[_metadata_variable_name]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)

View file

@ -19,6 +19,8 @@ from typing import Final
from fastapi import APIRouter, Depends, HTTPException
from litellm._internal_context import service_target
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import *
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
@ -141,7 +143,7 @@ async def update_budget(
- model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}}
- budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset.
"""
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client
from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client, user_api_key_cache
if prisma_client is None:
raise HTTPException(
@ -199,6 +201,11 @@ async def update_budget(
where={"budget_id": budget_obj.budget_id},
data=budget_obj_jsonified,
)
with service_target("team_member_budgets"):
try:
await user_api_key_cache.async_delete_cache(key=f"team_member_default_budget:{budget_obj.budget_id}")
except Exception: # noqa: BLE001 # Cache failure must not fail a committed budget update.
verbose_proxy_logger.warning("Failed to evict the team member default budget cache", exc_info=True)
return response

View file

@ -1,3 +1,4 @@
import json
import math
from collections.abc import Mapping
from types import MappingProxyType
@ -580,21 +581,31 @@ MEMBER_BUDGET_PATCH_FIELDS: Final = MappingProxyType(
"allowed_models": "allowed_models",
"temp_budget_increase": "temp_budget_increase",
"temp_budget_expiry": "temp_budget_expiry",
"model_max_budget": "model_max_budget",
}
)
def _prisma_value(value: object) -> object:
_JSON_BUDGET_COLUMNS: Final = frozenset({"model_max_budget"})
def _prisma_value(column: str, value: object) -> object:
if column in _JSON_BUDGET_COLUMNS:
return json.dumps(value if value is not None else {})
return list(value) if isinstance(value, tuple) else value
def _prisma_payload(data: Mapping[str, object]) -> dict[str, object]:
return {column: _prisma_value(column, value) for column, value in data.items()}
def member_budget_patch(source: BaseModel) -> Mapping[str, object]:
"""Map the per-member limit fields a request actually set to their budget-table
columns (merge-patch: a sent value updates, an explicit null clears, an absent
field is left untouched)."""
provided: Final = source.model_dump(exclude_unset=True)
return {
column: _prisma_value(provided[request_field])
column: provided[request_field]
for request_field, column in MEMBER_BUDGET_PATCH_FIELDS.items()
if request_field in provided
}
@ -603,7 +614,7 @@ def member_budget_patch(source: BaseModel) -> Mapping[str, object]:
def _is_set_budget_value(value: object) -> bool:
if value is None:
return False
if isinstance(value, list) and len(value) == 0:
if isinstance(value, (list, dict)) and len(value) == 0:
return False
return True
@ -645,10 +656,26 @@ async def upsert_budget_and_membership(
A patch that only touches the temporary budget pair never copies permanent
limits into a new row, so the member keeps inheriting the live team default.
"""
if not budget_patch:
ignore_default_model_budget_clear: Final = (
existing_budget_id is not None
and existing_budget_id == team_default_budget_id
and "model_max_budget" in budget_patch
and (
budget_patch["model_max_budget"] is None
or (isinstance(budget_patch["model_max_budget"], Mapping) and not budget_patch["model_max_budget"])
)
)
effective_patch: Final = MappingProxyType(
{
field: value
for field, value in budget_patch.items()
if not (ignore_default_model_budget_clear and field == "model_max_budget")
}
)
if not effective_patch:
return
write_data: Final = dict(budget_patch)
write_data: Final = dict(effective_patch)
if "budget_duration" in write_data:
duration: Final = write_data["budget_duration"]
write_data["budget_reset_at"] = (
@ -675,7 +702,7 @@ async def upsert_budget_and_membership(
return
await tx.litellm_budgettable.update(
where={"budget_id": existing_budget_id},
data={"updated_by": user_api_key_dict.user_id or "", **write_data},
data={"updated_by": user_api_key_dict.user_id or "", **_prisma_payload(write_data)},
)
return
@ -689,14 +716,24 @@ async def upsert_budget_and_membership(
create_data: Final[dict[str, object]] = { # mutable-ok: Prisma create payloads are dict-shaped
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
"model_max_budget": {},
**MappingProxyType(
{f: source[f] for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS if _is_set_budget_value(source.get(f))}
{
f: source[f]
for f in _TEAM_MEMBER_BUDGET_LIMIT_FIELDS
if (
f != "model_max_budget"
or team_default_budget_id is None
or source.get("budget_id") != team_default_budget_id
)
and _is_set_budget_value(source.get(f))
}
),
**write_data,
}
# Restarting an inherited window on an unrelated edit hands the member a free period.
carried: Final = source.get("budget_reset_at") if "budget_duration" not in budget_patch else None
carried: Final = source.get("budget_reset_at") if "budget_duration" not in effective_patch else None
if carried is not None:
create_data["budget_reset_at"] = carried
if create_data.get("budget_reset_at") is None:
@ -708,7 +745,7 @@ async def upsert_budget_and_membership(
return
new_budget: Final = await tx.litellm_budgettable.create(
data=create_data,
data=_prisma_payload(create_data),
include={"team_membership": True},
)
await tx.litellm_teammembership.upsert(

View file

@ -37,6 +37,7 @@ from pydantic import BaseModel, Field, JsonValue, TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict, assert_never
import litellm
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.integrations.prometheus import PrometheusLogger
@ -228,7 +229,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
TeamUserSpendRow,
UpdateTeamMemberPermissionsRequest,
)
from litellm.types.utils import BudgetConfig
from litellm.types.utils import BudgetConfig, GenericBudgetConfigType
if TYPE_CHECKING:
from prisma import Prisma
@ -547,6 +548,7 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit: int | None = None,
team_member_tpm_limit: int | None = None,
team_member_budget_duration: str | None = None,
team_member_model_max_budget: GenericBudgetConfigType | None = None,
) -> bool:
"""Check if any team member limits are provided"""
return any(
@ -555,9 +557,22 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit is not None,
team_member_tpm_limit is not None,
team_member_budget_duration is not None,
team_member_model_max_budget is not None,
]
)
@staticmethod
@with_service_target("team_member_budgets")
async def _evict_default_budget_cache(
budget_id: str,
user_api_key_cache: UserApiKeyCache | None,
) -> None:
if user_api_key_cache is not None:
try:
await user_api_key_cache.async_delete_cache(key=f"team_member_default_budget:{budget_id}")
except Exception: # noqa: BLE001 # Cache failure must not fail a committed budget update.
verbose_proxy_logger.warning("Failed to evict the team member default budget cache", exc_info=True)
@staticmethod
async def create_team_member_budget_table(
data: NewTeamRequest | _TeamAliasBudgetRow,
@ -567,6 +582,8 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit: int | None = None,
team_member_tpm_limit: int | None = None,
team_member_budget_duration: str | None = None,
team_member_model_max_budget: GenericBudgetConfigType | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
explicitly_set_fields: AbstractSet[str] = frozenset(),
) -> dict:
"""Create team member budget table with provided limits.
@ -602,11 +619,17 @@ class TeamMemberBudgetHandler:
budget_request.tpm_limit = team_member_tpm_limit
if team_member_budget_duration is not None:
budget_request.budget_duration = team_member_budget_duration
if team_member_model_max_budget is not None:
budget_request.model_max_budget = team_member_model_max_budget
team_member_budget_table: Final = await _as_budget_write(new_budget)(
budget_obj=budget_request,
user_api_key_dict=user_api_key_dict,
)
await TeamMemberBudgetHandler._evict_default_budget_cache(
budget_id=team_member_budget_table.budget_id,
user_api_key_cache=user_api_key_cache,
)
# Add team_member_budget_id as metadata field to team table
if new_team_data_json.get("metadata") is None:
@ -627,6 +650,8 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit: int | None = None,
team_member_tpm_limit: int | None = None,
team_member_budget_duration: str | None = None,
team_member_model_max_budget: GenericBudgetConfigType | None = None,
user_api_key_cache: UserApiKeyCache | None = None,
explicitly_set_fields: AbstractSet[str] = frozenset(),
) -> dict:
"""Upsert team member budget table with provided limits.
@ -657,11 +682,19 @@ class TeamMemberBudgetHandler:
budget_request.budget_duration = team_member_budget_duration
if team_member_budget_duration is None:
budget_request.budget_reset_at = None
if team_member_model_max_budget is not None or "team_member_model_max_budget" in explicitly_set_fields:
budget_request.model_max_budget = (
{} if team_member_model_max_budget is None else team_member_model_max_budget
)
budget_row: Final = await _as_budget_write(update_budget)(
budget_obj=budget_request,
user_api_key_dict=user_api_key_dict,
)
await TeamMemberBudgetHandler._evict_default_budget_cache(
budget_id=team_member_budget_id,
user_api_key_cache=user_api_key_cache,
)
verbose_proxy_logger.info(
"Updated team member budget table: %s, with team_member_budget=%s, team_member_rpm_limit=%s, team_member_tpm_limit=%s",
budget_row.budget_id,
@ -682,6 +715,8 @@ class TeamMemberBudgetHandler:
team_member_rpm_limit=team_member_rpm_limit,
team_member_tpm_limit=team_member_tpm_limit,
team_member_budget_duration=team_member_budget_duration,
team_member_model_max_budget=team_member_model_max_budget,
user_api_key_cache=user_api_key_cache,
explicitly_set_fields=explicitly_set_fields,
)
@ -696,6 +731,9 @@ class TeamMemberBudgetHandler:
data_dict.pop("team_member_budget_duration", None)
data_dict.pop("team_member_rpm_limit", None)
data_dict.pop("team_member_tpm_limit", None)
data_dict.pop( # pyright: ignore[reportUnknownMemberType] # legacy raw dict lacks key typing
"team_member_model_max_budget", None
)
_clean_team_member_fields = clean_team_member_fields
@ -705,6 +743,7 @@ class TeamMemberBudgetHandler:
user_api_key_dict: "UserAPIKeyAuth",
updated_kv: dict,
explicitly_set_fields: set,
user_api_key_cache: UserApiKeyCache | None = None,
) -> dict:
"""Clear explicitly-nulled fields on the team member budget row."""
from litellm.proxy._types import BudgetNewRequest
@ -727,10 +766,16 @@ class TeamMemberBudgetHandler:
budget_request.rpm_limit = None
if "team_member_tpm_limit" in explicitly_set_fields:
budget_request.tpm_limit = None
if "team_member_model_max_budget" in explicitly_set_fields:
budget_request.model_max_budget = {}
await update_budget(
budget_obj=budget_request,
user_api_key_dict=user_api_key_dict,
)
await TeamMemberBudgetHandler._evict_default_budget_cache(
budget_id=team_member_budget_id,
user_api_key_cache=user_api_key_cache,
)
TeamMemberBudgetHandler.clean_team_member_fields(updated_kv)
return updated_kv
@ -1420,6 +1465,7 @@ async def new_team(
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
@ -1508,6 +1554,10 @@ async def new_team(
validate_budget_duration(data.budget_duration)
validate_budget_duration(data.team_member_budget_duration)
validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user)
validate_team_model_max_budget(
model_max_budget=data.team_member_model_max_budget,
premium_user=premium_user,
)
if data.soft_budget is not None:
if data.max_budget is not None:
@ -1702,6 +1752,7 @@ async def new_team(
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
team_member_model_max_budget=data.team_member_model_max_budget,
):
data_json = await TeamMemberBudgetHandler.create_team_member_budget_table(
data=data,
@ -1711,6 +1762,8 @@ async def new_team(
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
team_member_model_max_budget=data.team_member_model_max_budget,
user_api_key_cache=user_api_key_cache,
explicitly_set_fields=data.model_fields_set,
)
@ -2190,6 +2243,7 @@ async def update_team(
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
- team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
- team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
@ -2276,6 +2330,10 @@ async def update_team(
validate_budget_duration(data.budget_duration)
validate_budget_duration(data.team_member_budget_duration)
validate_team_model_max_budget(model_max_budget=data.model_max_budget, premium_user=premium_user)
validate_team_model_max_budget(
model_max_budget=data.team_member_model_max_budget,
premium_user=premium_user,
)
existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique(
where={"team_id": data.team_id}
@ -2495,6 +2553,7 @@ async def update_team(
"team_member_rpm_limit",
"team_member_tpm_limit",
"team_member_budget_duration",
"team_member_model_max_budget",
]
if field in updated_kv
}
@ -2524,6 +2583,7 @@ async def update_team(
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
team_member_model_max_budget=data.team_member_model_max_budget,
):
updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
team_table=existing_team_row,
@ -2533,6 +2593,8 @@ async def update_team(
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
team_member_budget_duration=data.team_member_budget_duration,
team_member_model_max_budget=data.team_member_model_max_budget,
user_api_key_cache=user_api_key_cache,
explicitly_set_fields=_team_member_fields_in_request,
)
# Backfill team_memberships for members who joined before the
@ -2556,6 +2618,7 @@ async def update_team(
user_api_key_dict=user_api_key_dict,
updated_kv=updated_kv,
explicitly_set_fields=_team_member_fields_in_request,
user_api_key_cache=user_api_key_cache,
)
else:
TeamMemberBudgetHandler.clean_team_member_fields(updated_kv)
@ -3866,6 +3929,7 @@ async def team_member_update(
)
validate_budget_duration(data.budget_duration)
validate_team_model_max_budget(data.model_max_budget, premium_user=premium_user)
_existing_team_row: Final = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
@ -3982,6 +4046,7 @@ async def team_member_update(
allowed_models=data.allowed_models,
temp_budget_increase=data.temp_budget_increase,
temp_budget_expiry=data.temp_budget_expiry,
model_max_budget=data.model_max_budget,
)

View file

@ -677,6 +677,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils):
_metadata["user_api_key_model_max_budget"] = user_api_key_dict.model_max_budget
_metadata["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget
_metadata["user_api_key_user_model_max_budget"] = user_api_key_dict.user_model_max_budget
_metadata["user_api_key_team_member_model_max_budget"] = user_api_key_dict.team_member_model_max_budget
_metadata["user_api_key_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget
else:
for field in (

View file

@ -359,7 +359,7 @@ from litellm.proxy.auth.auth_utils import (
log_once_if_budget_reservation_disabled,
warn_once_if_custom_auth_skips_common_checks,
)
from litellm.proxy.auth.fallback_budget import router_fallback_budget_check
from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck, router_fallback_budget_check
from litellm.proxy.auth.fallback_model_access import router_fallback_access_check
from litellm.proxy.auth.handle_jwt import JWTHandler
from litellm.proxy.auth.litellm_license import AUTO_ROUTER_LICENSE_REMEDY, LicenseCheck
@ -7196,7 +7196,10 @@ class ProxyConfig:
),
ignore_invalid_deployments=True, # don't raise an error if a deployment is invalid
fallback_access_check=router_fallback_access_check,
fallback_budget_check=router_fallback_budget_check,
fallback_budget_check=RouterFallbackBudgetCheck(
is_enforced=router_fallback_budget_check.is_enforced,
team_member_model_budget_limiter=model_max_budget_limiter,
),
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
@ -7691,7 +7694,10 @@ class ProxyConfig:
search_tools=search_tools,
ignore_invalid_deployments=True,
fallback_access_check=router_fallback_access_check,
fallback_budget_check=router_fallback_budget_check,
fallback_budget_check=RouterFallbackBudgetCheck(
is_enforced=router_fallback_budget_check.is_enforced,
team_member_model_budget_limiter=model_max_budget_limiter,
),
auto_router_capability_limit=_license_check.auto_router_capability_limit,
)
verbose_proxy_logger.debug("updated llm_router: %s", llm_router)

View file

@ -4,6 +4,8 @@ from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeVar
from pydantic import TypeAdapter
if TYPE_CHECKING:
from litellm.types.llms.openai import OpenAIFileObject
@ -19,6 +21,7 @@ from litellm.types.utils import LlmProviders
_V = TypeVar("_V")
ROUTER_ONLY_CALL_KWARGS: Final = frozenset({"silent_model", "include_fallback_errors"})
_OBJECT_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[object, object]]] = TypeAdapter(Mapping[object, object])
def without_router_only_kwargs(kwargs: Mapping[str, _V]) -> dict[str, _V]:
@ -72,6 +75,30 @@ def resolve_model_group_alias(model_group_alias: object, model: str) -> str | No
return target
def resolve_served_model(
model: str,
model_alias_map: object,
key_aliases: object,
model_group_alias: object,
) -> str:
"""Model the request preparation serves for `model`.
Apply the global alias map, then key aliases, then `router_settings.model_group_alias`.
"""
global_aliases: Final = (
_OBJECT_MAPPING_ADAPTER.validate_python(model_alias_map) if isinstance(model_alias_map, Mapping) else None
)
global_target: Final = global_aliases.get(model) if global_aliases is not None else None
after_global: Final = global_target if isinstance(global_target, str) else model
key_alias_mapping: Final = (
_OBJECT_MAPPING_ADAPTER.validate_python(key_aliases) if isinstance(key_aliases, Mapping) else None
)
key_target: Final = key_alias_mapping.get(after_global) if key_alias_mapping is not None else None
after_key: Final = key_target if isinstance(key_target, str) else after_global
router_target: Final = resolve_model_group_alias(model_group_alias, after_key)
return router_target if router_target is not None else after_key
def truncate_fallback_error_detail(detail: str) -> str:
"""
Bound a fallback failure detail before it is logged or appended to an exception message.

View file

@ -13,7 +13,7 @@ Coverage:
"""
import json
from typing import Any, Dict, List
from typing import Any, Dict, Final, List
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -1198,6 +1198,7 @@ def _fake_user_api_key_auth(
end_user_model_max_budget=None,
end_user_id=None,
user_model_max_budget=None,
team_member_model_max_budget=None,
user_id=None,
token=None,
):
@ -1219,6 +1220,7 @@ def _fake_user_api_key_auth(
auth.end_user_model_max_budget = end_user_model_max_budget
auth.end_user_id = end_user_id
auth.user_model_max_budget = user_model_max_budget
auth.team_member_model_max_budget = team_member_model_max_budget
auth.user_id = user_id
auth.token = token
return auth
@ -1659,6 +1661,120 @@ async def test_summary_model_denied_when_user_over_model_budget():
assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter no longer accepts"
async def test_summary_model_denied_when_team_member_over_model_budget(monkeypatch: pytest.MonkeyPatch):
from litellm.caching.dual_cache import DualCache
from litellm.llms.anthropic.pass_through.context_management.constants import COMPACT_SUMMARY_MODEL_SETTING_KEY
from litellm.models.team_membership import LiteLLM_TeamMembership
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_TeamTableCachedObj,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.auth.auth_checks import get_team_member_default_budget
from litellm.proxy.auth.user_api_key_auth import _load_team_member_default_model_budget
from litellm.proxy.common_utils.user_api_key_cache import (
UserApiKeyCache,
team_membership_reservation_cache_key,
)
from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter
from litellm.proxy.proxy_server import general_settings
import litellm.proxy.proxy_server as proxy_server
budget_id: Final = "summary-budget"
member_budget: Final = {"claude-haiku-4-5": {"max_budget": 0.0, "budget_duration": "1d"}}
user_id: Final = "member-over-budget"
team_id: Final = "team-1"
user_api_key_cache: Final = UserApiKeyCache()
prisma_client: Final = MagicMock()
team_object: Final = LiteLLM_TeamTableCachedObj(
team_id=team_id,
models=["all-proxy-models"],
metadata={"team_member_budget_id": budget_id},
)
membership: Final = LiteLLM_TeamMembership(
user_id=user_id,
team_id=team_id,
budget_id=budget_id,
litellm_budget_table=LiteLLM_BudgetTable(budget_id=budget_id, model_max_budget={}),
)
await user_api_key_cache.async_set_cache(
key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id),
value=membership,
model_type=LiteLLM_TeamMembership,
)
await user_api_key_cache.async_set_cache(
key=f"team_member_default_budget:{budget_id}",
value=LiteLLM_BudgetTable(budget_id=budget_id, model_max_budget=member_budget),
model_type=LiteLLM_BudgetTable,
)
await user_api_key_cache.async_set_cache(
key=f"team_id:{team_id}",
value=team_object,
model_type=LiteLLM_TeamTableCachedObj,
)
auth: Final = UserAPIKeyAuth(
token="hashed-token",
user_id=user_id,
team_id=team_id,
user_role=LitellmUserRoles.INTERNAL_USER,
models=["all-proxy-models"],
)
await _load_team_member_default_model_budget(
user_api_key_auth_obj=auth,
team_object=team_object,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
skip_budget_checks=False,
proxy_logging_obj=None,
)
assert auth.team_member_model_max_budget == member_budget
assert (
await get_team_member_default_budget(
budget_id=budget_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
raise_on_lookup_error=True,
)
).model_max_budget == member_budget
router: Final = litellm.Router(
model_list=[
{"model_name": MODEL, "litellm_params": {"model": MODEL}},
{
"model_name": "claude-haiku-4-5",
"litellm_params": {
"model": "claude-haiku-4-5",
"mock_response": "<summary>x</summary>",
},
},
]
)
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache)
monkeypatch.setattr(proxy_server, "llm_router", router)
monkeypatch.setattr(proxy_server, "model_max_budget_limiter", limiter)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock())
monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "claude-haiku-4-5")
messages: Final = [{"role": "user", "content": "hello " * 60_000}]
result: Final = await apply_compact_20260112(
model=MODEL,
messages=messages,
tools=None,
system=None,
edit_spec={"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 50_000}},
user_api_key_auth=auth,
llm_router=router,
)
assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded"
assert result.compaction_block is None
async def test_summary_model_denied_when_end_user_over_model_budget():
"""End-user per-model budget is enforced for the summary subrequest too."""
import litellm
@ -2022,9 +2138,7 @@ async def test_summary_model_denied_when_team_over_model_budget():
limiter = MagicMock()
limiter.is_key_within_model_budget = AsyncMock(return_value=True)
limiter.is_team_within_model_budget = AsyncMock(
side_effect=litellm.BudgetExceededError(
message="over budget", current_cost=10, max_budget=5
)
side_effect=litellm.BudgetExceededError(message="over budget", current_cost=10, max_budget=5)
)
with (
@ -2064,9 +2178,7 @@ async def test_summary_model_denied_when_team_over_model_budget():
PROXY_VirtualKeyModelMaxBudgetLimiter,
)
real_params = inspect.signature(
PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget
).parameters
real_params = inspect.signature(PROXY_VirtualKeyModelMaxBudgetLimiter.is_team_within_model_budget).parameters
for kwarg in ("team_id", "team_model_max_budget", "key_model_max_budget", "model"):
assert kwarg in real_params, f"compact.py passes {kwarg}=, which the limiter does not accept"
@ -2585,6 +2697,9 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata()
"user_api_key": "sk-parent",
"user_api_key_team_id": "team-abc",
"user_api_key_user_id": "user-xyz",
"user_api_key_team_member_model_max_budget": {
"claude-haiku-4-5": {"budget_limit": 4.0, "time_period": "1d"}
},
"litellm_call_id": "call-1",
},
additional_drop_params=None,
@ -2595,6 +2710,9 @@ async def test_prepare_context_managed_request_forwards_proxy_litellm_metadata()
assert captured_summary_metadata.get("user_api_key") == "sk-parent"
assert captured_summary_metadata.get("user_api_key_team_id") == "team-abc"
assert captured_summary_metadata.get("user_api_key_user_id") == "user-xyz"
assert captured_summary_metadata.get("user_api_key_team_member_model_max_budget") == {
"claude-haiku-4-5": {"budget_limit": 4.0, "time_period": "1d"}
}
assert captured_summary_metadata.get("litellm_call_id") == "call-1"
# Anthropic-shape ``metadata.user_id`` must not leak in as a propagated field.
assert "user_id" not in captured_summary_metadata

View file

@ -7703,6 +7703,47 @@ async def test_get_team_member_default_budget_caches_json_safe_payload():
mock_prisma_client.db.litellm_budgettable.find_unique.assert_awaited_once()
@pytest.mark.asyncio
async def test_get_team_member_default_budget_reraises_lookup_error_when_requested() -> None:
from litellm.proxy.auth.auth_checks import get_team_member_default_budget
lookup_error: Final = RuntimeError("budget database unavailable")
cache: Final = MagicMock()
cache.async_get_cache = AsyncMock(return_value=None)
prisma_client: Final = MagicMock()
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(side_effect=lookup_error)
with pytest.raises(RuntimeError) as exc_info:
await get_team_member_default_budget(
budget_id="default-budget",
prisma_client=prisma_client,
user_api_key_cache=cache,
raise_on_lookup_error=True,
)
assert exc_info.value is lookup_error
@pytest.mark.asyncio
async def test_get_team_member_default_budget_returns_none_on_lookup_error_by_default() -> None:
from litellm.proxy.auth.auth_checks import get_team_member_default_budget
prisma_client: Final = MagicMock()
prisma_client.db.litellm_budgettable.find_unique = AsyncMock(
side_effect=RuntimeError("budget database unavailable")
)
cache: Final = MagicMock()
cache.async_get_cache = AsyncMock(return_value=None)
budget: Final = await get_team_member_default_budget(
budget_id="default-budget",
prisma_client=prisma_client,
user_api_key_cache=cache,
)
assert budget is None
@pytest.mark.asyncio
async def test_get_end_user_object_db_fetch_returns_validated_end_user():
from litellm.proxy.auth.auth_checks import get_end_user_object

View file

@ -1,12 +1,23 @@
from collections.abc import Mapping, Sequence
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import litellm
import pytest
from litellm import Router
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
from litellm.proxy.auth.fallback_budget import (
RouterFallbackBudgetCheck,
is_token_within_budget_for_model,
router_fallback_budget_check,
)
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
team_member_budget_entity_id,
)
FREE_MODEL = {
"model_name": "free-model",
@ -48,6 +59,41 @@ def _token(**overrides) -> UserAPIKeyAuth:
return UserAPIKeyAuth(**fields)
async def _team_member_limiter(
spend_by_model: Mapping[str, float] | None = None,
) -> tuple[PROXY_VirtualKeyModelMaxBudgetLimiter, MagicMock]:
entity_id: Final = team_member_budget_entity_id(user_id="u1", team_id="t1")
spend_by_key: Final = {
model_budget_spend_cache_key(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=entity_id,
budget_model=model,
budget_duration="1d",
): spend
for model, spend in (spend_by_model or {}).items()
}
redis_cache: Final = MagicMock()
async def get_spend(*, key: str) -> object | None:
return spend_by_key.get(key)
redis_cache.async_get_cache = AsyncMock(side_effect=get_spend)
return PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=redis_cache)), redis_cache
def _member_budget_spend_keys(models: Sequence[str]) -> tuple[str, ...]:
entity_id: Final = team_member_budget_entity_id(user_id="u1", team_id="t1")
return tuple(
model_budget_spend_cache_key(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=entity_id,
budget_model=model,
budget_duration="1d",
)
for model in models
)
ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: True)
NOT_ENFORCED = RouterFallbackBudgetCheck(is_enforced=lambda: False)
@ -70,6 +116,201 @@ async def test_paid_target_refused_when_over_user_budget():
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
@pytest.mark.asyncio
async def test_team_member_budget_refuses_an_exhausted_paid_target() -> None:
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
token: Final = _token(
user_id="u1",
team_id="t1",
team_member_model_max_budget=member_caps,
)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 10.0})
within_budget: Final = await is_token_within_budget_for_model(
model="paid-model",
valid_token=token,
llm_router=_router(),
team_member_model_budget_limiter=limiter,
)
assert within_budget is False
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("paid-model",))
)
@pytest.mark.asyncio
async def test_router_fallback_budget_check_resolves_model_group_alias_before_member_cap_check() -> None:
expensive_model: Final = {
**PAID_MODEL,
"model_name": "expensive",
"model_info": {"id": "expensive-model-id"},
}
router: Final = Router(
model_list=[FREE_MODEL, expensive_model],
model_group_alias={"backup": "expensive"},
)
member_caps: Final = {
"backup": {"max_budget": 100.0, "budget_duration": "1d"},
"expensive": {"max_budget": 5.0, "budget_duration": "1d"},
}
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"backup": 1.0, "expensive": 10.0})
budget_check: Final = RouterFallbackBudgetCheck(
is_enforced=lambda: True,
team_member_model_budget_limiter=limiter,
)
within_budget: Final = await budget_check(
model="backup",
request_kwargs={"metadata": {"user_api_key_auth": token}},
llm_router=router,
)
assert within_budget is False
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("backup", "expensive"))
)
@pytest.mark.asyncio
async def test_router_fallback_budget_check_refuses_an_exhausted_alias_cap() -> None:
expensive_model: Final = {
**PAID_MODEL,
"model_name": "expensive",
"model_info": {"id": "expensive-model-id"},
}
router: Final = Router(
model_list=[FREE_MODEL, expensive_model],
model_group_alias={"backup": "expensive"},
)
member_caps: Final = {"backup": {"max_budget": 5.0, "budget_duration": "1d"}}
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"backup": 10.0})
budget_check: Final = RouterFallbackBudgetCheck(
is_enforced=lambda: True,
team_member_model_budget_limiter=limiter,
)
within_budget: Final = await budget_check(
model="backup",
request_kwargs={"metadata": {"user_api_key_auth": token}},
llm_router=router,
)
assert within_budget is False
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("backup",))
)
@pytest.mark.asyncio
async def test_router_fallback_budget_check_allows_alias_when_both_names_are_within_cap() -> None:
expensive_model: Final = {
**PAID_MODEL,
"model_name": "expensive",
"model_info": {"id": "expensive-model-id"},
}
router: Final = Router(
model_list=[FREE_MODEL, expensive_model],
model_group_alias={"backup": "expensive"},
)
member_caps: Final = {
"backup": {"max_budget": 5.0, "budget_duration": "1d"},
"expensive": {"max_budget": 5.0, "budget_duration": "1d"},
}
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"backup": 1.0, "expensive": 1.0})
budget_check: Final = RouterFallbackBudgetCheck(
is_enforced=lambda: True,
team_member_model_budget_limiter=limiter,
)
within_budget: Final = await budget_check(
model="backup",
request_kwargs={"metadata": {"user_api_key_auth": token}},
llm_router=router,
)
assert within_budget is True
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("backup", "expensive"))
)
@pytest.mark.asyncio
async def test_router_fallback_budget_check_deduplicates_non_alias_target() -> None:
router: Final = _router()
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
token: Final = _token(user_id="u1", team_id="t1", team_member_model_max_budget=member_caps)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 1.0})
budget_check: Final = RouterFallbackBudgetCheck(
is_enforced=lambda: True,
team_member_model_budget_limiter=limiter,
)
within_budget: Final = await budget_check(
model="paid-model",
request_kwargs={"metadata": {"user_api_key_auth": token}},
llm_router=router,
)
assert within_budget is True
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("paid-model",))
)
@pytest.mark.asyncio
async def test_team_member_budget_allows_a_paid_target_when_under_cap() -> None:
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
token: Final = _token(
user_id="u1",
team_id="t1",
team_member_model_max_budget=member_caps,
)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 1.0})
within_budget: Final = await is_token_within_budget_for_model(
model="paid-model",
valid_token=token,
llm_router=_router(),
team_member_model_budget_limiter=limiter,
)
assert within_budget is True
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("paid-model",))
)
@pytest.mark.asyncio
async def test_team_member_caps_are_not_evaluated_without_a_limiter() -> None:
token: Final = _token(
user_id="u1",
team_id="t1",
team_member_model_max_budget={"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}},
)
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True
@pytest.mark.asyncio
async def test_team_member_limiter_is_not_called_without_member_caps() -> None:
token: Final = _token(user_id="u1", team_id="t1")
limiter, redis_cache = await _team_member_limiter()
within_budget: Final = await is_token_within_budget_for_model(
model="paid-model",
valid_token=token,
llm_router=_router(),
team_member_model_budget_limiter=limiter,
)
assert within_budget is True
assert redis_cache.async_get_cache.call_args_list == []
@pytest.mark.asyncio
async def test_zero_cost_target_allowed_even_when_over_budget():
"""Refusing a free target would deny a request on spend some other model accrued."""
@ -123,6 +364,33 @@ async def test_enforced_check_reads_the_key_from_request_metadata(metadata_field
assert await ENFORCED(model="paid-model", request_kwargs=under, llm_router=_router()) is True
@pytest.mark.asyncio
async def test_router_check_passes_member_limiter_from_request_metadata() -> None:
member_caps: Final = {"paid-model": {"max_budget": 5.0, "budget_duration": "1d"}}
token: Final = _token(
user_id="u1",
team_id="t1",
team_member_model_max_budget=member_caps,
)
limiter, redis_cache = await _team_member_limiter(spend_by_model={"paid-model": 10.0})
budget_check: Final = RouterFallbackBudgetCheck(
is_enforced=lambda: True,
team_member_model_budget_limiter=limiter,
)
request_kwargs: Final = {"metadata": {"user_api_key_auth": token}}
within_budget: Final = await budget_check(
model="paid-model",
request_kwargs=request_kwargs,
llm_router=_router(),
)
assert within_budget is False
assert tuple(call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list) == (
_member_budget_spend_keys(("paid-model",))
)
@pytest.mark.asyncio
async def test_a_stale_low_counter_still_refuses_a_paid_target(monkeypatch):
"""

File diff suppressed because it is too large Load diff

View file

@ -1,6 +1,8 @@
# tests/litellm/proxy/common_utils/test_upsert_budget_membership.py
import json
import types
from datetime import datetime, timedelta, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import pytest
@ -230,7 +232,11 @@ async def test_create_from_temp_budget_pair_only(mock_tx, fake_user):
)
mock_tx.litellm_budgettable.create.assert_awaited_once()
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 5.0, "temp_budget_expiry": expiry}
assert stored_budget_row(mock_tx) == {
"model_max_budget": "{}",
"temp_budget_increase": 5.0,
"temp_budget_expiry": expiry,
}
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
mock_tx.litellm_teammembership.update.assert_not_called()
@ -252,7 +258,11 @@ async def test_create_from_temp_pair_never_snapshots_team_default(mock_tx, fake_
)
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}
assert stored_budget_row(mock_tx) == {
"model_max_budget": "{}",
"temp_budget_increase": 1.0,
"temp_budget_expiry": expiry,
}
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@ -274,7 +284,11 @@ async def test_temp_pair_on_shared_default_member_creates_bare_row(mock_tx, fake
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
mock_tx.litellm_budgettable.update.assert_not_called()
assert stored_budget_row(mock_tx) == {"temp_budget_increase": 1.0, "temp_budget_expiry": expiry}
assert stored_budget_row(mock_tx) == {
"model_max_budget": "{}",
"temp_budget_increase": 1.0,
"temp_budget_expiry": expiry,
}
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@ -296,6 +310,58 @@ async def test_clearing_temp_pair_on_shared_default_member_is_noop(mock_tx, fake
mock_tx.litellm_teammembership.upsert.assert_not_called()
@pytest.mark.asyncio
async def test_clearing_model_budget_to_none_on_shared_default_is_noop(mock_tx, fake_user):
await upsert_budget_and_membership(
mock_tx,
team_id="team-default",
user_id="user-on-default",
existing_budget_id="team-default-budget-1",
user_api_key_dict=fake_user,
budget_patch={"model_max_budget": None},
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_teammembership.update.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
@pytest.mark.asyncio
async def test_clearing_model_budget_to_none_with_tpm_clones_shared_default(mock_tx, fake_user):
shared_default_id: Final = "team-default-budget-1"
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(
budget_id=shared_default_id,
max_budget=100.0,
model_max_budget={"claude-sonnet-4-6": {"max_budget": 5.0}},
)
)
await upsert_budget_and_membership(
mock_tx,
team_id="team-default",
user_id="user-on-default",
existing_budget_id=shared_default_id,
user_api_key_dict=fake_user,
budget_patch={"model_max_budget": None, "tpm_limit": 5},
team_default_budget_id=shared_default_id,
)
create_data: Final = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert create_data == {
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"model_max_budget": "{}",
"max_budget": 100.0,
"tpm_limit": 5,
}
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@pytest.mark.asyncio
async def test_temp_pair_with_permanent_field_still_clones_shared_default(mock_tx, fake_user):
expiry = datetime(2100, 1, 1, tzinfo=timezone.utc)
@ -337,7 +403,7 @@ async def test_create_from_plain_patch_does_not_snapshot_team_default(mock_tx, f
)
mock_tx.litellm_budgettable.find_unique.assert_not_awaited()
assert stored_budget_row(mock_tx) == {"tpm_limit": 500}
assert stored_budget_row(mock_tx) == {"model_max_budget": "{}", "tpm_limit": 500}
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
@ -380,6 +446,7 @@ async def test_clone_on_write_from_shared_default(mock_tx, fake_user):
assert create_data == {
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"model_max_budget": "{}",
"max_budget": 50.0, # caller wins
"tpm_limit": 500, # cloned from default
"budget_duration": "1d", # cloned from default
@ -401,6 +468,69 @@ async def test_clone_on_write_from_shared_default(mock_tx, fake_user):
)
@pytest.mark.asyncio
async def test_clone_on_write_preserves_model_caps_from_shared_nondefault_row(mock_tx, fake_user):
shared_budget_id: Final = "shared-member-budget"
model_caps: Final = {"gpt-4o": {"max_budget": 4.0, "budget_duration": "1d"}}
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(
budget_id=shared_budget_id,
max_budget=100.0,
model_max_budget=model_caps,
)
)
await upsert_budget_and_membership(
mock_tx,
team_id="team-shared",
user_id="user-shared",
existing_budget_id=shared_budget_id,
user_api_key_dict=fake_user,
budget_patch={"tpm_limit": 1000},
team_default_budget_id="team-default-budget",
shared_budget_ids=frozenset({shared_budget_id}),
)
create_data: Final = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
persisted_model_caps: Final = create_data["model_max_budget"]
assert (
json.loads(persisted_model_caps) if isinstance(persisted_model_caps, str) else persisted_model_caps
) == model_caps
@pytest.mark.asyncio
async def test_clearing_model_budget_to_none_clones_shared_nondefault_row(mock_tx, fake_user):
shared_budget_id: Final = "shared-member-budget"
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(
budget_id=shared_budget_id,
max_budget=100.0,
model_max_budget={"gpt-4o": {"max_budget": 4.0}},
)
)
await upsert_budget_and_membership(
mock_tx,
team_id="team-shared",
user_id="user-shared",
existing_budget_id=shared_budget_id,
user_api_key_dict=fake_user,
budget_patch={"model_max_budget": None},
team_default_budget_id="team-default-budget",
shared_budget_ids=frozenset({shared_budget_id}),
)
create_data: Final = mock_tx.litellm_budgettable.create.await_args.kwargs["data"]
assert create_data == {
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"model_max_budget": "{}",
"max_budget": 100.0,
}
mock_tx.litellm_budgettable.update.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_awaited_once()
# TEST: forking the shared default while clearing its duration must drop the
# duration (and not carry a reset time) on the new private budget.
@pytest.mark.asyncio
@ -431,6 +561,7 @@ async def test_clone_on_write_clears_duration(mock_tx, fake_user):
assert create_data == {
"created_by": fake_user.user_id,
"updated_by": fake_user.user_id,
"model_max_budget": "{}",
"max_budget": 200.0,
"tpm_limit": 500,
"budget_duration": None,
@ -460,3 +591,31 @@ async def test_private_budget_updates_in_place(mock_tx, fake_user):
)
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()
@pytest.mark.asyncio
async def test_clearing_model_budget_to_none_updates_private_row(mock_tx, fake_user):
mock_tx.litellm_budgettable.find_unique = AsyncMock(
return_value=budget_row(
max_budget=10.0,
model_max_budget={"gpt-4o": {"max_budget": 4.0}},
)
)
await upsert_budget_and_membership(
mock_tx,
team_id="team-mixed",
user_id="user-private",
existing_budget_id="private-budget-xyz",
user_api_key_dict=fake_user,
budget_patch={"model_max_budget": None},
team_default_budget_id="team-default-budget-1",
)
mock_tx.litellm_budgettable.update.assert_awaited_once_with(
where={"budget_id": "private-budget-xyz"},
data={"updated_by": fake_user.user_id, "model_max_budget": "{}"},
)
mock_tx.litellm_budgettable.create.assert_not_called()
mock_tx.litellm_teammembership.update.assert_not_called()
mock_tx.litellm_teammembership.upsert.assert_not_called()

View file

@ -6,15 +6,20 @@ from typing import Final
import pytest
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
team_member_budget_entity_id,
)
from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import LiteLLMBatch, Usage
KEY_HASH: Final = "key-hash-batch"
USER_ID: Final = "user-batch"
TEAM_ID: Final = "team-batch"
MODEL_GROUP: Final = "batch-qa-primary"
BATCH_COST: Final = 2.925e-05
CHAT_COST: Final = 0.001
@ -35,7 +40,13 @@ def _batch(batch_id: str, status: str) -> LiteLLMBatch:
)
def _event(call_type: str, response_cost: float) -> dict[str, object]:
def _event(
call_type: str,
response_cost: float,
team_member_model_max_budget: Mapping[str, object] | None = None,
user_id: str = USER_ID,
team_id: str = TEAM_ID,
) -> dict[str, object]:
return {
"call_type": call_type,
"standard_logging_object": {
@ -43,12 +54,21 @@ def _event(call_type: str, response_cost: float) -> dict[str, object]:
"response_cost": response_cost,
"model": "openai/gpt-5.4-mini",
"model_group": MODEL_GROUP,
"metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_id": USER_ID},
"metadata": {
"user_api_key_hash": KEY_HASH,
"user_api_key_user_id": user_id,
"user_api_key_team_id": team_id,
},
},
"litellm_params": {
"metadata": {
"user_api_key_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}},
"user_api_key_user_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}},
**(
{"user_api_key_team_member_model_max_budget": team_member_model_max_budget}
if team_member_model_max_budget is not None
else {}
),
}
},
}
@ -60,9 +80,23 @@ async def _poll(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBa
)
async def _chat(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter) -> None:
async def _chat(
limiter: PROXY_VirtualKeyModelMaxBudgetLimiter,
team_member_model_max_budget: Mapping[str, object] | None = None,
user_id: str = USER_ID,
team_id: str = TEAM_ID,
) -> None:
await limiter.async_log_success_event(
_event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None
_event(
"acompletion",
CHAT_COST,
team_member_model_max_budget=team_member_model_max_budget,
user_id=user_id,
team_id=team_id,
),
response_obj=None,
start_time=None,
end_time=None,
)
@ -74,6 +108,87 @@ def _local_spend(limiter: PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str)
return limiter.dual_cache.in_memory_cache.get_cache(key=spend_key) or 0.0
@pytest.mark.asyncio
async def test_team_member_budget_is_charged_and_enforced_after_success():
budget: Final = {MODEL_GROUP: {"max_budget": 0.0005, "budget_duration": "1d"}}
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
await _chat(limiter, team_member_model_max_budget=budget)
entity_id: Final = team_member_budget_entity_id(user_id=USER_ID, team_id=TEAM_ID)
spend_key: Final = model_budget_spend_cache_key(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=entity_id,
budget_model=MODEL_GROUP,
budget_duration="1d",
)
assert await _spend(limiter, spend_key) == CHAT_COST
with pytest.raises(litellm.BudgetExceededError, match="LiteLLM Team Member"):
await limiter.is_team_member_within_model_budget(
user_id=USER_ID,
team_id=TEAM_ID,
team_member_model_max_budget=budget,
model=MODEL_GROUP,
)
def test_team_member_budget_entity_id_is_unambiguous():
first_entity_id: Final = team_member_budget_entity_id(user_id="a:b", team_id="c")
second_entity_id: Final = team_member_budget_entity_id(user_id="a", team_id="b:c")
assert first_entity_id == "3:a:b:c"
assert second_entity_id == "1:a:b:c"
first_spend_key: Final = model_budget_spend_cache_key(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=first_entity_id,
budget_model=MODEL_GROUP,
budget_duration="1d",
)
second_spend_key: Final = model_budget_spend_cache_key(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=second_entity_id,
budget_model=MODEL_GROUP,
budget_duration="1d",
)
assert first_spend_key != second_spend_key
@pytest.mark.asyncio
async def test_team_member_budget_spend_isolated_by_user_and_team():
budget: Final = {MODEL_GROUP: {"max_budget": CHAT_COST * 2, "budget_duration": "1d"}}
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
await _chat(limiter, team_member_model_max_budget=budget)
assert await limiter.is_team_member_within_model_budget(
user_id="different-user",
team_id=TEAM_ID,
team_member_model_max_budget=budget,
model=MODEL_GROUP,
)
assert await limiter.is_team_member_within_model_budget(
user_id=USER_ID,
team_id="different-team",
team_member_model_max_budget=budget,
model=MODEL_GROUP,
)
@pytest.mark.asyncio
async def test_zero_team_member_model_budget_rejects_first_request():
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
with pytest.raises(litellm.BudgetExceededError, match="LiteLLM Team Member"):
await limiter.is_team_member_within_model_budget(
user_id=USER_ID,
team_id=TEAM_ID,
team_member_model_max_budget={MODEL_GROUP: {"max_budget": 0, "budget_duration": "1d"}},
model=MODEL_GROUP,
)
class _Clock:
def __init__(self) -> None:
self.seconds = 0.0

View file

@ -85,6 +85,10 @@ def _matches(row: Mapping[str, object], where: Mapping[str, object]) -> bool:
return all((wanted := _wanted(where, field)) is not None and row.get(field) in wanted for field in where)
def _as_stored(data: Mapping[str, object]) -> dict[str, object]:
return {k: json.loads(v) if k == "model_max_budget" and isinstance(v, str) else v for k, v in data.items()}
class _BudgetTable:
def __init__(self, budgets: Sequence[_BudgetRow]) -> None:
self.rows: dict[str, _BudgetRow] = {b.budget_id: b for b in budgets}
@ -94,13 +98,13 @@ class _BudgetTable:
async def update(self, where: Mapping[str, str], data: Mapping[str, object]) -> _BudgetRow:
row: Final = self.rows[where["budget_id"]]
updated: Final = row.model_copy(update=dict(data))
updated: Final = row.model_copy(update=_as_stored(data))
self.rows[row.budget_id] = updated
return updated
async def create(self, data: Mapping[str, object], include: Mapping[str, bool] | None = None) -> _BudgetRow:
budget_id: Final = f"new-budget-{len(self.rows) + 1}"
row: Final = _BudgetRow.model_validate({**data, "budget_id": budget_id})
row: Final = _BudgetRow.model_validate({**_as_stored(data), "budget_id": budget_id})
self.rows[budget_id] = row
return row

View file

@ -11,6 +11,8 @@ from fastapi.testclient import TestClient
import litellm.proxy.proxy_server as ps
from litellm.proxy.proxy_server import app
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles, CommonProxyErrors
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -102,6 +104,25 @@ async def test_update_budget_rejects_a_duration_that_never_advances(client_and_m
mock_table.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_update_budget_evicts_local_cache_when_redis_delete_fails(
client_and_mocks, monkeypatch: pytest.MonkeyPatch
) -> None:
client, _, _ = client_and_mocks
cache_key = "team_member_default_budget:budget_456"
cache = UserApiKeyCache(in_memory_cache=InMemoryCache())
cache.set_cache(cache_key, {"budget_id": "budget_456"}, local_only=True)
redis_cache = MagicMock()
redis_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis unavailable"))
cache.redis_cache = redis_cache
monkeypatch.setattr(ps, "user_api_key_cache", cache)
response = client.post("/budget/update", json={"budget_id": "budget_456", "max_budget": 10.0})
assert response.status_code == 200, response.text
assert cache.get_cache(cache_key, local_only=True) is None
@pytest.mark.asyncio
async def test_new_budget_db_not_connected(client_and_mocks, monkeypatch):
client, mock_prisma, mock_table = client_and_mocks
@ -136,6 +157,26 @@ async def test_update_budget_success(client_and_mocks, monkeypatch):
assert body["updated_by"] == "test_user"
@pytest.mark.asyncio
async def test_update_budget_evicts_team_member_default_budget_cache(client_and_mocks, monkeypatch):
client, _, _ = client_and_mocks
cache: Final = UserApiKeyCache(in_memory_cache=InMemoryCache())
monkeypatch.setattr(ps, "user_api_key_cache", cache)
cache_key: Final = "team_member_default_budget:budget_456"
await cache.async_set_cache(key=cache_key, value={"budget_id": "budget_456"})
assert await cache.async_get_cache(key=cache_key) == {"budget_id": "budget_456"}
try:
response: Final = client.post(
"/budget/update",
json={"budget_id": "budget_456", "max_budget": 10.0},
)
assert response.status_code == 200, response.text
assert await cache.async_get_cache(key=cache_key) is None
finally:
await cache.async_delete_cache(key=cache_key)
@pytest.mark.asyncio
async def test_new_and_update_budget_persist_tpd_limit(client_and_mocks):
client, _, mock_table = client_and_mocks

View file

@ -18,6 +18,7 @@ from pydantic import BaseModel
from litellm import Router
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
@ -35,25 +36,132 @@ from litellm.proxy.management_endpoints.common_utils import (
update_metadata_fields,
user_api_key_has_admin_view,
user_has_admin_privileges,
upsert_budget_and_membership,
)
from litellm.types.utils import BudgetConfig
@pytest.mark.asyncio
@pytest.mark.parametrize(
"budget_patch",
[
{"tpm_limit": 200},
{"model_max_budget": {}, "tpm_limit": 5},
],
)
async def test_shared_default_clone_keeps_inheriting_updated_team_model_caps(
budget_patch: dict[str, object],
) -> None:
from litellm.proxy.auth.user_api_key_auth import _load_team_member_default_model_budget
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, team_membership_reservation_cache_key
from litellm.proxy._types import LiteLLM_TeamMembership
source_caps = {"claude-sonnet-4-6": {"max_budget": 3.0, "budget_duration": "1d"}}
default_caps = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}}
budget_table = SimpleNamespace(
find_unique=AsyncMock(
return_value=SimpleNamespace(
model_dump=MagicMock(
return_value={
"budget_id": "default-budget",
"model_max_budget": source_caps,
"tpm_limit": 100,
}
)
)
),
create=AsyncMock(return_value=SimpleNamespace(budget_id="private-budget")),
)
membership_table = SimpleNamespace(upsert=AsyncMock())
tx = SimpleNamespace(litellm_budgettable=budget_table, litellm_teammembership=membership_table)
token = UserAPIKeyAuth(user_id="member-1", team_id="team-1")
await upsert_budget_and_membership(
tx,
team_id="team-1",
user_id="member-1",
existing_budget_id="default-budget",
user_api_key_dict=token,
budget_patch=budget_patch,
team_default_budget_id="default-budget",
)
budget_table.create.assert_awaited_once()
membership_table.upsert.assert_awaited_once()
assert budget_table.create.await_args.kwargs["data"]["tpm_limit"] == budget_patch["tpm_limit"]
model_max_budget: Final = budget_table.create.await_args.kwargs["data"]["model_max_budget"]
assert isinstance(model_max_budget, str)
assert model_max_budget == "{}"
membership = LiteLLM_TeamMembership(
user_id="member-1",
team_id="team-1",
budget_id="private-budget",
litellm_budget_table=LiteLLM_BudgetTable(budget_id="private-budget", model_max_budget={}),
)
user_api_key_cache: Final = UserApiKeyCache()
await user_api_key_cache.async_set_cache(
key=team_membership_reservation_cache_key(user_id="member-1", team_id="team-1"),
value=membership,
model_type=LiteLLM_TeamMembership,
)
await user_api_key_cache.async_set_cache(
key="team_member_default_budget:default-budget",
value=LiteLLM_BudgetTable(budget_id="default-budget", model_max_budget=default_caps),
model_type=LiteLLM_BudgetTable,
)
team_object = SimpleNamespace(
team_id="team-1",
metadata={"team_member_budget_id": "default-budget"},
)
await _load_team_member_default_model_budget(
user_api_key_auth_obj=token,
team_object=team_object,
prisma_client=MagicMock(),
user_api_key_cache=user_api_key_cache,
skip_budget_checks=False,
proxy_logging_obj=None,
)
assert token.team_member_model_max_budget == default_caps
@pytest.mark.asyncio
async def test_empty_model_budget_patch_on_team_default_keeps_membership_attached() -> None:
tx: Final = MagicMock()
tx.litellm_budgettable.find_unique = AsyncMock(
return_value=SimpleNamespace(
model_dump=MagicMock(
return_value={
"budget_id": "default-budget",
"max_budget": 10.0,
"model_max_budget": {"claude-sonnet-4-6": {"max_budget": 1.0}},
}
)
)
)
tx.litellm_budgettable.create = AsyncMock()
tx.litellm_budgettable.update = AsyncMock()
tx.litellm_teammembership.upsert = AsyncMock()
tx.litellm_teammembership.update = AsyncMock()
await upsert_budget_and_membership(
tx,
team_id="team-1",
user_id="member-1",
existing_budget_id="default-budget",
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
budget_patch={"model_max_budget": {}},
team_default_budget_id="default-budget",
)
tx.litellm_budgettable.create.assert_not_awaited()
tx.litellm_budgettable.update.assert_not_awaited()
tx.litellm_teammembership.upsert.assert_not_awaited()
tx.litellm_teammembership.update.assert_not_awaited()
class TestUpdateMetadataFieldsEmptyCollections:
"""
Regression tests for issue #20304.
The UI sends empty arrays (`[]`) for enterprise-only fields like
guardrails, policies, and logging even when the user hasn't configured
these features. The backend must not treat empty collections as an
intent to use the feature, and therefore must not trigger the premium
license check.
However, empty collections must still be written into metadata so that
users can intentionally clear a previously-set field (e.g. removing all
guardrails by sending `guardrails: []`).
"""
@patch("litellm.proxy.management_endpoints.common_utils.premium_user_check")
def test_empty_list_does_not_trigger_premium_check(self, mock_premium_check):
"""Empty lists for premium fields must not trigger the premium check."""

View file

@ -2591,6 +2591,8 @@ async def test_update_team_team_member_budget_not_passed_to_db(
team_member_rpm_limit=None,
team_member_tpm_limit=None,
team_member_budget_duration=None,
team_member_model_max_budget=None,
user_api_key_cache=None,
explicitly_set_fields=frozenset(),
):
# Remove team_member_budget from updated_kv as the real function does
@ -3159,6 +3161,8 @@ async def test_update_team_with_team_member_budget_duration(
team_member_rpm_limit=None,
team_member_tpm_limit=None,
team_member_budget_duration=None,
team_member_model_max_budget=None,
user_api_key_cache=None,
explicitly_set_fields=frozenset(),
):
result_kv = updated_kv.copy()
@ -10891,6 +10895,7 @@ async def test_create_team_member_budget_table_with_duration():
team_alias="test-team",
team_member_budget=20.0,
team_member_budget_duration="30d",
team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}},
)
with patch(
@ -10904,15 +10909,181 @@ async def test_create_team_member_budget_table_with_duration():
user_api_key_dict=mock_admin,
team_member_budget=20.0,
team_member_budget_duration="30d",
team_member_model_max_budget=data.team_member_model_max_budget,
)
mock_new_budget.assert_awaited_once()
budget_request = mock_new_budget.call_args.kwargs["budget_obj"]
assert budget_request.budget_duration == "30d"
assert budget_request.max_budget == 20.0
assert budget_request.model_max_budget == data.team_member_model_max_budget
assert result["metadata"]["team_member_budget_id"] == "budget-abc"
@pytest.mark.asyncio
async def test_upsert_team_member_model_max_budget_updates_default_budget_and_evicts_cache():
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
model_budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
team_table: Final = LiteLLM_TeamTable(
team_id="test-team",
metadata={"team_member_budget_id": "member-budget-1"},
members_with_roles=[],
)
cache: Final = MagicMock()
cache.async_delete_cache = AsyncMock()
budget_cache: Final = MagicMock()
budget_cache.async_delete_cache = AsyncMock()
prisma: Final = MagicMock()
prisma.db.litellm_budgettable.update = AsyncMock(return_value=MagicMock(budget_id="member-budget-1"))
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.user_api_key_cache", budget_cache),
patch("litellm.proxy.proxy_server.premium_user", True),
):
result: Final = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
team_table=team_table,
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
updated_kv={"team_id": "test-team", "team_member_model_max_budget": model_budget},
team_member_model_max_budget=model_budget,
user_api_key_cache=cache,
explicitly_set_fields={"team_member_model_max_budget"},
)
budget_update: Final = prisma.db.litellm_budgettable.update.call_args.kwargs
assert budget_update["where"]["budget_id"] == "member-budget-1"
assert budget_update["data"]["budget_id"] == "member-budget-1"
assert json.loads(budget_update["data"]["model_max_budget"]) == model_budget
assert result["metadata"]["team_member_budget_id"] == "member-budget-1"
assert "team_member_model_max_budget" not in result
cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:member-budget-1")
budget_cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:member-budget-1")
@pytest.mark.asyncio
async def test_update_team_clears_explicitly_null_member_model_budget_and_evicts_cache(
disable_audit_logging_for_mocked_team,
):
from fastapi import Request
from litellm.proxy._types import UpdateTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import update_team
existing: Final = _existing_team_with_model_caps({})
existing.metadata = {"team_member_budget_id": "member-budget-1"}
existing.members_with_roles = []
existing.model_dump.return_value["metadata"] = existing.metadata
existing.model_dump.return_value["members_with_roles"] = []
cache: Final = MagicMock()
cache.async_get_cache = AsyncMock(return_value=None)
cache.async_set_cache = AsyncMock()
cache.async_delete_cache = AsyncMock()
updated: Final = _existing_team_with_model_caps({})
updated.metadata = {"team_member_budget_id": "member-budget-1"}
updated.model_dump.return_value.update({"metadata": updated.metadata, "members_with_roles": []})
updated.litellm_model_table = None
update_budget_response: Final = MagicMock(budget_id="member-budget-1")
prisma: Final = MagicMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch("litellm.proxy.proxy_server.premium_user", True),
):
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
prisma.jsonify_team_object = lambda db_data: db_data
prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated)
prisma.db.litellm_budgettable.update = AsyncMock(return_value=update_budget_response)
request_data: Final = UpdateTeamRequest(team_id="standalone-team-123", team_member_model_max_budget=None)
await update_team(
data=request_data,
http_request=MagicMock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
)
budget_update: Final = prisma.db.litellm_budgettable.update.call_args.kwargs
assert budget_update["where"]["budget_id"] == "member-budget-1"
assert budget_update["data"]["budget_id"] == "member-budget-1"
assert budget_update["data"]["model_max_budget"] == "{}"
assert "team_member_model_max_budget" in request_data.model_fields_set
assert "model_max_budget" in budget_update["data"]
assert cache.async_delete_cache.await_args_list == [
call(key="team_member_default_budget:member-budget-1"),
call(key="team_member_default_budget:member-budget-1"),
]
@pytest.mark.asyncio
async def test_update_team_persists_member_model_budget_to_default_budget_row(
disable_audit_logging_for_mocked_team,
):
from fastapi import Request
from litellm.proxy._types import UpdateTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import update_team
member_caps: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}
existing: Final = _existing_team_with_model_caps({})
existing.metadata = {"team_member_budget_id": "member-budget-1"}
existing.members_with_roles = []
existing.model_dump.return_value["metadata"] = existing.metadata
existing.model_dump.return_value["members_with_roles"] = []
cache: Final = MagicMock()
cache.async_get_cache = AsyncMock(return_value=None)
cache.async_set_cache = AsyncMock()
cache.async_delete_cache = AsyncMock()
updated: Final = _existing_team_with_model_caps({})
updated.metadata = {"team_member_budget_id": "member-budget-1"}
updated.model_dump.return_value.update({"metadata": updated.metadata, "members_with_roles": []})
updated.litellm_model_table = None
update_budget_response: Final = MagicMock(budget_id="member-budget-1")
prisma: Final = MagicMock()
with (
patch("litellm.proxy.proxy_server.prisma_client", prisma),
patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
patch("litellm.proxy.proxy_server.premium_user", True),
):
prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing)
prisma.jsonify_team_object = lambda db_data: db_data
prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated)
prisma.db.litellm_budgettable.update = AsyncMock(return_value=update_budget_response)
await update_team(
data=UpdateTeamRequest(team_id="standalone-team-123", team_member_model_max_budget=member_caps),
http_request=MagicMock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
)
budget_update: Final = prisma.db.litellm_budgettable.update.call_args.kwargs
assert budget_update["where"]["budget_id"] == "member-budget-1"
assert budget_update["data"]["budget_id"] == "member-budget-1"
stored_member_model_caps: Final = json.loads(budget_update["data"]["model_max_budget"])
assert {
model: {
**budget,
"tpm_limit": budget.get("tpm_limit"),
"rpm_limit": budget.get("rpm_limit"),
}
for model, budget in stored_member_model_caps.items()
} == {
"gpt-4o": {
"max_budget": 5.0,
"budget_duration": "1d",
"tpm_limit": None,
"rpm_limit": None,
}
}
assert cache.async_delete_cache.await_args_list == [
call(key="team_member_default_budget:member-budget-1"),
call(key="team_member_default_budget:member-budget-1"),
]
# ---------------------------------------------------------------------------
# Tests for _batch_resolve_access_group_resources
# ---------------------------------------------------------------------------
@ -12519,11 +12690,13 @@ async def _drive_team_write(
**(existing_kwargs or {}),
)
auth = user or UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="u")
user_api_key_cache: Final = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
with (
_patch("litellm.proxy.proxy_server.prisma_client") as pc,
_patch("litellm.proxy.proxy_server.llm_router", None),
_patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
_patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),
_patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
_patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
_patch(
@ -16245,13 +16418,24 @@ def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_a
@pytest.mark.asyncio
async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth):
@pytest.mark.parametrize(
("member_caps", "member_budget_duration"),
[
({"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, "30d"),
({}, None),
],
)
async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth, member_caps, member_budget_duration):
mock_db_client.jsonify_team_object = lambda db_data: db_data
mock_db_client.get_data = AsyncMock(return_value=None)
mock_db_client.update_data = AsyncMock(return_value=MagicMock())
mock_db_client.db = MagicMock()
mock_db_client.db.litellm_modeltable = MagicMock()
mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123"))
mock_db_client.db.litellm_budgettable = MagicMock()
mock_db_client.db.litellm_budgettable.create = AsyncMock(
side_effect=lambda *, data: MagicMock(budget_id=data["budget_id"])
)
team_create_result = MagicMock(team_id="team-model-caps")
team_create_result.model_dump.return_value = {"team_id": "team-model-caps"}
@ -16269,11 +16453,18 @@ async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_aut
from litellm.proxy._types import NewTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import new_team
with patch("litellm.proxy.proxy_server.premium_user", True): # test-quality-ok: proxy_server module global is the endpoint's only injection point
member_budget_cache: Final = MagicMock()
member_budget_cache.async_delete_cache = AsyncMock()
with (
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.proxy_server.user_api_key_cache", member_budget_cache),
):
await new_team(
data=NewTeamRequest(
team_alias="model-caps",
model_max_budget={"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}},
team_member_budget_duration=member_budget_duration,
team_member_model_max_budget=member_caps,
),
http_request=MagicMock(spec=Request),
user_api_key_dict=mock_admin_auth,
@ -16283,6 +16474,75 @@ async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_aut
assert team_data["model_max_budget"] == {
"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d", "tpm_limit": None, "rpm_limit": None}
}
member_budget_row: Final = mock_db_client.db.litellm_budgettable.create.call_args.kwargs["data"]
member_budget_id: Final = member_budget_row["budget_id"]
assert team_data["metadata"]["team_member_budget_id"] == member_budget_id
assert member_budget_row.get("budget_duration") == member_budget_duration
stored_member_budget_model_caps: Final = json.loads(member_budget_row["model_max_budget"])
expected_member_model_caps: Final = (
{
"gpt-4o": {
"max_budget": 5.0,
"budget_duration": "1d",
"tpm_limit": None,
"rpm_limit": None,
}
}
if member_caps
else {}
)
assert {
model: {
**budget,
"tpm_limit": budget.get("tpm_limit"),
"rpm_limit": budget.get("rpm_limit"),
}
for model, budget in stored_member_budget_model_caps.items()
} == expected_member_model_caps
member_budget_cache.async_delete_cache.assert_awaited_once_with(
key=f"team_member_default_budget:{member_budget_id}"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("endpoint_name", ["new", "update"])
@pytest.mark.parametrize(
("premium_user", "model_budget", "expected_status"),
[
(False, {"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}}, 403),
(True, {"gpt-4o": {"max_budget": 10.0}}, 400),
],
)
async def test_team_member_model_budget_validation(endpoint_name, premium_user, model_budget, expected_status):
from fastapi import Request
from litellm.proxy._types import NewTeamRequest, ProxyException, UpdateTeamRequest
from litellm.proxy.management_endpoints.team_endpoints import (
new_team,
)
from litellm.proxy.management_endpoints.team_endpoints import (
update_team as update_team_endpoint,
)
endpoint: Final = new_team if endpoint_name == "new" else update_team_endpoint
data: Final = (
NewTeamRequest(team_alias="member-model-validation", team_member_model_max_budget=model_budget)
if endpoint_name == "new"
else UpdateTeamRequest(team_id="member-model-validation", team_member_model_max_budget=model_budget)
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.premium_user", premium_user),
pytest.raises(ProxyException) as exc_info,
):
await endpoint(
data=data,
http_request=MagicMock(spec=Request),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
)
assert exc_info.value.code == str(expected_status)
@pytest.mark.asyncio
@ -17098,6 +17358,87 @@ def test_member_budget_patch_maps_temp_budget_fields() -> None:
}
def test_member_budget_patch_stores_model_max_budget_as_json_and_clears_on_empty() -> None:
from litellm.proxy.management_endpoints.common_utils import member_budget_patch
request: Final = TeamMemberUpdateRequest(
team_id="team-1",
user_id="user-1",
model_max_budget={"gpt-4o": {"budget_limit": 5, "time_period": "30d"}},
)
assert member_budget_patch(request) == {"model_max_budget": {"gpt-4o": {"max_budget": 5.0, "budget_duration": "30d"}}}
cleared: Final = TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", model_max_budget={})
assert member_budget_patch(cleared) == {"model_max_budget": {}}
@pytest.mark.asyncio
async def test_upsert_budget_and_membership_clears_own_model_max_budget_to_empty_json() -> None:
from litellm.proxy.management_endpoints.common_utils import upsert_budget_and_membership
own_row: Final = SimpleNamespace(
model_dump=lambda: {
"budget_id": "member-b",
"max_budget": 50.0,
"model_max_budget": {"gpt-4o": {"max_budget": 2.5, "budget_duration": "7d"}},
}
)
tx: Final = MagicMock()
tx.litellm_budgettable.find_unique = AsyncMock(return_value=own_row)
tx.litellm_budgettable.update = AsyncMock()
tx.litellm_teammembership.update = AsyncMock()
await upsert_budget_and_membership(
tx,
team_id="team-1",
user_id="user-1",
existing_budget_id="member-b",
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
budget_patch={"model_max_budget": None},
team_default_budget_id="default-b",
)
assert tx.litellm_budgettable.update.await_args.kwargs["data"]["model_max_budget"] == "{}"
tx.litellm_teammembership.update.assert_not_awaited()
@pytest.mark.asyncio
async def test_upsert_budget_and_membership_clone_keeps_shared_default_model_max_budget_as_json() -> None:
from litellm.proxy.management_endpoints.common_utils import upsert_budget_and_membership
default_model_budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "30d"}}
default_row: Final = SimpleNamespace(
model_dump=lambda: {
"budget_id": "default-b",
"max_budget": 10.0,
"model_max_budget": default_model_budget,
"budget_reset_at": None,
}
)
tx: Final = MagicMock()
tx.litellm_budgettable.find_unique = AsyncMock(return_value=default_row)
tx.litellm_budgettable.create = AsyncMock(return_value=SimpleNamespace(budget_id="member-b"))
tx.litellm_teammembership.upsert = AsyncMock()
await upsert_budget_and_membership(
tx,
team_id="team-1",
user_id="user-1",
existing_budget_id="default-b",
user_api_key_dict=UserAPIKeyAuth(user_id="admin"),
budget_patch={"tpm_limit": 7},
team_default_budget_id="default-b",
)
create_data: Final = tx.litellm_budgettable.create.await_args.kwargs["data"]
assert create_data["tpm_limit"] == 7
assert create_data["max_budget"] == 10.0
assert json.loads(create_data["model_max_budget"]) == {}
assert tx.litellm_teammembership.upsert.await_args.kwargs["data"]["update"] == {
"litellm_budget_table": {"connect": {"budget_id": "member-b"}}
}
def test_team_member_update_request_temp_budget_fields_must_be_set_together() -> None:
with pytest.raises(ValidationError, match="temp_budget_increase and temp_budget_expiry must be set together"):
TeamMemberUpdateRequest(team_id="team-1", user_id="user-1", temp_budget_increase=50.0)
@ -17211,3 +17552,18 @@ def test_aggregated_date_range_error_accepts_canonical_dates_and_keeps_range_che
assert aggregated_date_range_error("2026-09-24", "2026-09-26") is None
assert aggregated_date_range_error("2026-09-26", "2026-09-24") == "end_date must be on or after start_date"
assert aggregated_date_range_error("2020-01-01", "2026-12-31") == "Date range must be at most 400 days"
@pytest.mark.asyncio
async def test_default_budget_cache_eviction_does_not_fail_after_redis_error() -> None:
from litellm.proxy.management_endpoints.team_endpoints import TeamMemberBudgetHandler
user_api_key_cache = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock(side_effect=RuntimeError("redis unavailable"))
await TeamMemberBudgetHandler._evict_default_budget_cache(
budget_id="budget_456",
user_api_key_cache=user_api_key_cache,
)
user_api_key_cache.async_delete_cache.assert_awaited_once_with(key="team_member_default_budget:budget_456")

View file

@ -1,3 +1,4 @@
import json
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from types import SimpleNamespace
@ -817,6 +818,10 @@ async def test_add_new_member_with_user_email_links_default_budget():
mock_prisma_client.db.litellm_budgettable.create.assert_not_called()
def _as_stored(data: Mapping[str, object]) -> dict[str, object]:
return {k: json.loads(v) if k == "model_max_budget" and isinstance(v, str) else v for k, v in data.items()}
class _FakeBudgetTable:
def __init__(self) -> None:
self.rows: dict[str, dict[str, object]] = {}
@ -829,14 +834,14 @@ class _FakeBudgetTable:
self, *, data: Mapping[str, object], include: Mapping[str, bool] | None = None
) -> LiteLLM_BudgetTable:
budget_id: Final = str(data.get("budget_id") or uuid.uuid4())
self.rows[budget_id] = {**data, "budget_id": budget_id}
self.rows[budget_id] = {**_as_stored(data), "budget_id": budget_id}
return self._record(budget_id)
async def find_unique(self, *, where: Mapping[str, str]) -> LiteLLM_BudgetTable | None:
return self._record(where["budget_id"]) if where["budget_id"] in self.rows else None
async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> LiteLLM_BudgetTable:
self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **data}
self.rows[where["budget_id"]] = {**self.rows[where["budget_id"]], **_as_stored(data)}
return self._record(where["budget_id"])

View file

@ -7918,14 +7918,17 @@ def test_passthrough_carries_the_per_model_budgets():
key_budget = {"claude-opus-4-8": {"budget_limit": 1.0, "time_period": "18h"}}
user_budget = {"claude-opus-4-8": {"budget_limit": 2.0, "time_period": "1mo"}}
end_user_budget = {"claude-opus-4-8": {"budget_limit": 3.0, "time_period": "1d"}}
team_member_budget = {"claude-opus-4-8": {"budget_limit": 4.0, "time_period": "1d"}}
kwargs = _passthrough_kwargs_for_reservation(
UserAPIKeyAuth(
token="hash",
user_id="u-1",
team_id="team-1",
model_max_budget=key_budget,
user_model_max_budget=user_budget,
end_user_model_max_budget=end_user_budget,
team_member_model_max_budget=team_member_budget,
)
)
@ -7933,6 +7936,7 @@ def test_passthrough_carries_the_per_model_budgets():
assert metadata["user_api_key_model_max_budget"] == key_budget
assert metadata["user_api_key_user_model_max_budget"] == user_budget
assert metadata["user_api_key_end_user_model_max_budget"] == end_user_budget
assert metadata["user_api_key_team_member_model_max_budget"] == team_member_budget
def test_passthrough_budget_metadata_cannot_be_forged_by_the_request_body():
@ -7981,7 +7985,10 @@ async def test_user_defined_passthrough_is_neither_tracked_nor_enforced(metadata
budget: Final = {"managed-model": {"budget_limit": 0.1, "time_period": "1d"}}
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(DualCache())
auth: Final = UserAPIKeyAuth(
api_key="custom-key", token="custom-key", team_id="shared-team", team_model_max_budget=budget,
api_key="custom-key",
token="custom-key",
team_id="shared-team",
team_model_max_budget=budget,
)
endpoint: Final = create_pass_through_route(
endpoint="/custom-budget-test",

View file

@ -17,6 +17,7 @@ from fastapi.responses import JSONResponse, StreamingResponse
import litellm
from litellm._uuid import uuid
from litellm.anthropic_interface.exceptions import AnthropicErrorSseFrame, anthropic_error_sse_frame
from litellm.caching.caching import DualCache
from litellm.litellm_core_utils.bug_report import (
DISABLE_ENV_VAR,
ISSUE_URL_BASE,
@ -59,11 +60,18 @@ from litellm.proxy.common_request_processing import (
sse_error_payload,
)
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
from litellm.proxy.common_utils.sse_keepalive import ANTHROPIC_PING_SSE_CHUNK
from litellm.proxy.dd_span_tagger import DDSpanTagger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy._types import Litellm_EntityType, ProxyErrorTypes, ProxyException
from litellm.proxy._types import UserAPIKeyAuth as ProxyUserAPIKeyAuth
from litellm.proxy.auth.fallback_budget import RouterFallbackBudgetCheck
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
team_member_budget_entity_id,
)
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.router_utils.add_retry_fallback_headers import prepare_response_for_header_attachment
@ -7097,6 +7105,9 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
fallbacks: list[dict[str, list[str]]],
model_guardrails: dict[str, list[str]] | None = None,
untagged_rate_limited_models: frozenset[str] = frozenset(),
fallback_budget_check: RouterFallbackBudgetCheck | None = None,
free_models: frozenset[str] = frozenset(),
extra_models: frozenset[str] = frozenset(),
) -> tuple[ProxyLogging, litellm.Router, ProxyConfig, list[str]]:
"""Real v3 limiter (the default ``parallel_request_limiter``) wired in through the
``proxy_logging_obj`` seam, so ``common_processing_pre_call_logic`` runs for real:
@ -7145,11 +7156,17 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
"api_key": "fake",
**({"guardrails": guardrails_by_group[group]} if group in guardrails_by_group else {}),
},
**(
{"model_info": {"input_cost_per_token": 0, "output_cost_per_token": 0}}
if group in free_models
else {}
),
}
for chain in fallbacks
for group in (*chain.keys(), *(m for models in chain.values() for m in models))
for group in (*chain.keys(), *(m for models in chain.values() for m in models), *sorted(extra_models))
],
fallbacks=fallbacks,
fallback_budget_check=fallback_budget_check,
)
return proxy_logging_obj, router, proxy_server.ProxyConfig(), limiter_models
@ -7636,6 +7653,138 @@ class TestPreCallWithFallbacksOnLocalRateLimit:
assert data["metadata"]["guardrails"] == [structured_guardrail]
assert rig[3] == [primary_model, primary_model, fallback_model]
@staticmethod
async def _install_member_spend(monkeypatch: pytest.MonkeyPatch, model: str, spend: float) -> None:
from litellm.proxy import proxy_server
member_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
await member_limiter.dual_cache.async_set_cache(
key=model_budget_spend_cache_key(
entity_type=Litellm_EntityType.TEAM_MEMBER,
entity_id=team_member_budget_entity_id(user_id="user-1", team_id="team-1"),
budget_model=model,
budget_duration="1d",
),
value=spend,
)
monkeypatch.setattr(proxy_server, "model_max_budget_limiter", member_limiter)
def _member_key(self, primary_model: str, fallback_model: str) -> ProxyUserAPIKeyAuth:
return self._otel_key(model_rpm_limit={primary_model: 1}).model_copy(
update={
"user_id": "user-1",
"team_id": "team-1",
"team_member_model_max_budget": {fallback_model: {"max_budget": 5.0, "budget_duration": "1d"}},
}
)
@staticmethod
def _client_request(model: str) -> dict[str, object]:
return {"model": model, "messages": [{"role": "user", "content": "hi"}]}
@pytest.mark.asyncio
async def test_rate_limit_fallback_skips_backup_over_member_model_budget(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
primary_model: Final = "gpt-4.1"
fallback_model: Final = "gpt-4.1-mini"
await self._install_member_spend(monkeypatch, fallback_model, spend=5.0)
key: Final = self._member_key(primary_model, fallback_model)
rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
_, (first_data, _) = await self._pre_call(self._client_request(primary_model), key, rig)
processor: Final = ProxyBaseLLMRequestProcessing(data=self._client_request(primary_model))
with pytest.raises(ProxyRateLimitError):
await self._run(processor, key, rig)
assert first_data["model"] == primary_model
assert rig[3] == [primary_model, primary_model]
@pytest.mark.asyncio
async def test_rate_limit_fallback_skips_backup_aliased_to_model_over_member_budget(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
primary_model: Final = "gpt-4.1"
fallback_model: Final = "gpt-4.1-mini"
resolved_model: Final = "expensive"
await self._install_member_spend(monkeypatch, resolved_model, spend=5.0)
key: Final = self._member_key(primary_model, fallback_model).model_copy(
update={
"router_settings": {"model_group_alias": {fallback_model: resolved_model}},
"team_member_model_max_budget": {resolved_model: {"max_budget": 5.0, "budget_duration": "1d"}},
}
)
rig: Final = self._v3_limiter_rig(
monkeypatch,
key,
[{primary_model: [fallback_model]}],
extra_models=frozenset({resolved_model}),
)
await self._pre_call(self._client_request(primary_model), key, rig)
processor: Final = ProxyBaseLLMRequestProcessing(data=self._client_request(primary_model))
with pytest.raises(ProxyRateLimitError):
await self._run(processor, key, rig)
assert processor.data["model"] == primary_model
assert rig[3] == [primary_model, primary_model]
@pytest.mark.asyncio
async def test_rate_limit_fallback_uses_backup_within_member_model_budget(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
primary_model: Final = "gpt-4.1"
fallback_model: Final = "gpt-4.1-mini"
await self._install_member_spend(monkeypatch, fallback_model, spend=4.0)
key: Final = self._member_key(primary_model, fallback_model)
rig: Final = self._v3_limiter_rig(monkeypatch, key, [{primary_model: [fallback_model]}])
await self._pre_call(self._client_request(primary_model), key, rig)
_, (data, logging_obj) = await self._pre_call(self._client_request(primary_model), key, rig)
assert data["model"] == fallback_model
assert logging_obj.model == fallback_model
assert rig[3] == [primary_model, primary_model, fallback_model]
@pytest.mark.asyncio
async def test_rate_limit_fallback_uses_free_backup_with_spent_member_model_budget(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
primary_model: Final = "gpt-4.1"
fallback_model: Final = "gpt-4.1-mini"
await self._install_member_spend(monkeypatch, fallback_model, spend=5.0)
key: Final = self._member_key(primary_model, fallback_model)
rig: Final = self._v3_limiter_rig(
monkeypatch, key, [{primary_model: [fallback_model]}], free_models=frozenset({fallback_model})
)
await self._pre_call(self._client_request(primary_model), key, rig)
_, (data, _) = await self._pre_call(self._client_request(primary_model), key, rig)
assert data["model"] == fallback_model
assert rig[3] == [primary_model, primary_model, fallback_model]
@pytest.mark.asyncio
async def test_rate_limit_fallback_ignores_key_budget_already_admitted_at_auth(
self, monkeypatch: pytest.MonkeyPatch
) -> None:
primary_model: Final = "gpt-4.1"
fallback_model: Final = "gpt-4.1-mini"
key: Final = self._otel_key(model_rpm_limit={primary_model: 1}).model_copy(
update={"token": "hashed-key", "max_budget": 1.0, "spend": 2.0}
)
rig: Final = self._v3_limiter_rig(
monkeypatch,
key,
[{primary_model: [fallback_model]}],
fallback_budget_check=RouterFallbackBudgetCheck(is_enforced=lambda: True),
)
await self._pre_call(self._client_request(primary_model), key, rig)
_, (data, _) = await self._pre_call(self._client_request(primary_model), key, rig)
assert data["model"] == fallback_model
assert rig[3] == [primary_model, primary_model, fallback_model]
class _RecordingSuccessLogger(CustomLogger):
def __init__(self):

View file

@ -2016,6 +2016,8 @@ async def test_add_litellm_data_to_request_user_spend_and_budget():
api_key="hashed-key",
metadata={},
team_metadata={},
team_id="team-1",
team_member_model_max_budget={"gpt-3.5-turbo": {"max_budget": 15.0, "budget_duration": "1d"}},
user_spend=150.0,
user_max_budget=500.0,
)
@ -2032,6 +2034,9 @@ async def test_add_litellm_data_to_request_user_spend_and_budget():
metadata = updated_data.get("metadata", {})
assert metadata["user_api_key_user_spend"] == 150.0
assert metadata["user_api_key_user_max_budget"] == 500.0
assert metadata["user_api_key_team_member_model_max_budget"] == {
"gpt-3.5-turbo": {"max_budget": 15.0, "budget_duration": "1d"}
}
@pytest.mark.asyncio

View file

@ -14,6 +14,7 @@ from litellm.router_utils.common_utils import (
filter_web_search_deployments,
provider_for_generic_call,
resolve_model_group_alias,
resolve_served_model,
truncate_fallback_error_detail,
PROVIDER_SCOPED_CREDENTIAL_PARAMS,
warn_on_provider_credential_mismatch,
@ -566,6 +567,25 @@ class TestResolveModelGroupAlias:
assert router.get_model_from_alias("group-b") is None
@pytest.mark.parametrize(
"model, model_alias_map, key_aliases, model_group_alias, expected",
[
pytest.param("a", {"a": "b"}, {"b": "c"}, {"c": "d"}, "d", id="global-key-router-order"),
pytest.param("a", {}, ["a", "b"], {}, "a", id="non-mapping-key-aliases-ignored"),
pytest.param("a", {}, {}, {"a": {"model": "x"}}, "x", id="router-item-entry"),
pytest.param("a", {}, {}, {}, "a", id="no-alias"),
],
)
def test_resolves_served_model_aliases_in_order(
model: str,
model_alias_map: object,
key_aliases: object,
model_group_alias: object,
expected: str,
) -> None:
assert resolve_served_model(model, model_alias_map, key_aliases, model_group_alias) == expected
class TestTruncateFallbackErrorDetail:
def test_short_detail_is_returned_unchanged(self):
assert truncate_fallback_error_detail("boom") == "boom"

View file

@ -0,0 +1,83 @@
import { fireEvent, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders, testQueryClient } from "../../tests/test-utils";
import Teams from "./Teams";
const sessionCookie = () => {
const encode = (value: object) =>
btoa(JSON.stringify(value)).replaceAll("=", "").replaceAll("+", "-").replaceAll("/", "_");
const claims = {
key: "sk-session-test",
user_id: "test-admin",
user_role: "proxy_admin",
premium_user: true,
exp: Date.now() / 1000 + 3600,
};
document.cookie = `token=${encode({ alg: "none" })}.${encode(claims)}.test; Path=/`;
};
const responseFor = (url: string): object | object[] => {
if (url.includes("/v2/team/list")) return { teams: [], total: 0, page: 1, page_size: 50, total_pages: 1 };
if (url.includes("/models")) return { data: [{ id: "gpt-4" }] };
if (url.includes("/v1/access_group")) return [];
if (url.includes("/v1/mcp/access_groups")) return { access_groups: [] };
if (url.includes("/guardrails/list")) return { guardrails: [] };
if (url.includes("/policies/list")) return { policies: [] };
if (url.includes("/organization/list")) return { organizations: [] };
if (url.includes("/team/new")) return { team_id: "new-team-1" };
if (url.includes("/default")) return { values: {} };
if (url.includes("/litellm-ui-config")) return { server_root_path: "", proxy_base_url: "" };
return {};
};
beforeEach(() => {
testQueryClient.clear();
sessionCookie();
vi.stubGlobal(
"fetch",
vi.fn(async (input: RequestInfo | URL) => Response.json(responseFor(String(input)))),
);
});
afterEach(() => {
testQueryClient.clear();
vi.unstubAllGlobals();
document.cookie = "token=; Max-Age=0; Path=/";
});
describe("Teams creation integration", () => {
it("sends a configured default member per-model budget", async () => {
const user = userEvent.setup({ delay: null });
renderWithProviders(<Teams accessToken="sk-session-test" userID="test-admin" userRole="Admin" premiumUser />);
await user.click(screen.getAllByRole("button", { name: /create team/i })[0]);
await screen.findByLabelText(/team name/i);
await user.click(screen.getByText("Additional Settings"));
await screen.findByText(/Team Member Key Duration/);
const addBudgetButtons = screen.getAllByRole("button", { name: /Add Model Budget/i });
await user.click(addBudgetButtons[1]);
const modelSelectors = screen.getAllByPlaceholderText("Select model");
await user.click(modelSelectors[0]);
await user.click(await screen.findByRole("option", { name: "gpt-4" }));
fireEvent.change(screen.getAllByPlaceholderText("Max spend ($)")[0], { target: { value: "4" } });
await user.click(screen.getByText("Additional Settings"));
fireEvent.change(screen.getByLabelText(/team name/i), { target: { value: "Budget Team" } });
await user.click(screen.getAllByRole("button", { name: /create team/i }).at(-1)!);
await waitFor(() => {
expect(vi.mocked(fetch)).toHaveBeenCalled();
});
const createRequest = vi.mocked(fetch).mock.calls.find(([input]) => String(input).includes("/team/new"));
expect(createRequest).toBeDefined();
const payload = JSON.parse(String(createRequest?.[1]?.body)) as Record<string, unknown>;
expect(payload.team_member_model_max_budget).toStrictEqual({
"gpt-4": { budget_limit: 4, time_period: "30d" },
});
expect(JSON.parse(JSON.stringify(payload)).team_member_model_max_budget).toStrictEqual({
"gpt-4": { budget_limit: 4, time_period: "30d" },
});
});
});

View file

@ -51,6 +51,7 @@ import BudgetDurationDropdown, {
import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking";
import NumericalInput from "./shared/numerical_input";
import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor";
import { modelMaxBudgetUpdate } from "./key_team_helpers/modelMaxBudgetPayload";
import VectorStoreSelector from "./vector_store_management/VectorStoreSelector";
import SearchToolSelector from "./search_tools/SearchToolSelector";
import SkillSelector from "./skills/SkillSelector";
@ -280,6 +281,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
const [loggingSettings, setLoggingSettings] = useState<any[]>([]);
const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({});
const [modelMaxBudget, setModelMaxBudget] = useState<ModelMaxBudget>({});
const [teamMemberModelMaxBudget, setTeamMemberModelMaxBudget] = useState<ModelMaxBudget>({});
const [routerSettings, setRouterSettings] = useState<RouterSettingsAccordionValue | null>(null);
const [routerSettingsKey, setRouterSettingsKey] = useState<number>(0);
@ -358,6 +360,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
setLoggingSettings([]);
setModelAliases({});
setModelMaxBudget({});
setTeamMemberModelMaxBudget({});
setRouterSettings(null);
setRouterSettingsKey((prev) => prev + 1);
};
@ -551,9 +554,11 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
}
const { require_trace_id: requireTraceId, ...teamCreateValues } = formValues;
const teamMemberModelBudgets = modelMaxBudgetUpdate(teamMemberModelMaxBudget, {});
await teamCreateCall(accessToken, {
...teamCreateValues,
...(requireTraceId === undefined ? {} : { require_trace_id: requireTraceId }),
...(teamMemberModelBudgets !== undefined ? { team_member_model_max_budget: teamMemberModelBudgets } : {}),
models: normalizeTeamModelSelection(formValues.models),
});
toast.success("Team created");
@ -912,6 +917,15 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, isViewOnly
/>
)}
</FormField>
<ModelMaxBudgetField
key={`team-member-model-max-budget-${routerSettingsKey}`}
label="Default Per-Model Budget"
premiumUser={premiumUser}
value={teamMemberModelMaxBudget}
onChange={setTeamMemberModelMaxBudget}
availableModels={userModels}
hint="Set a default per-model spend cap for each team member, with an independent reset window."
/>
<FormField
control={form.control}
name="team_member_key_duration"

View file

@ -92,6 +92,13 @@ export function ModelMaxBudgetEditor({
}: ModelMaxBudgetEditorProps) {
const [entries, setEntries] = useState<ModelBudgetEntry[]>(() => modelMaxBudgetToEntries(value));
const incoming = JSON.stringify(value ?? {});
const [syncedValue, setSyncedValue] = useState(incoming);
if (incoming !== syncedValue) {
setSyncedValue(incoming);
if (incoming !== JSON.stringify(entriesToModelMaxBudget(entries))) setEntries(modelMaxBudgetToEntries(value));
}
const emitChange = (updated: ModelBudgetEntry[]) => {
setEntries(updated);
onChange(entriesToModelMaxBudget(updated));
@ -219,14 +226,15 @@ export function ModelMaxBudgetEditor({
interface ModelMaxBudgetFieldProps extends ModelMaxBudgetEditorProps {
hint: string;
label?: string;
}
/** The editor with its label, so every form that offers it presents it the same way. */
export function ModelMaxBudgetField({ hint, ...editorProps }: ModelMaxBudgetFieldProps) {
export function ModelMaxBudgetField({ hint, label = "Per-Model Budgets", ...editorProps }: ModelMaxBudgetFieldProps) {
return (
<Field>
<FieldLabel>
<span title={hint}>Per-Model Budgets</span>
<span title={hint}>{label}</span>
</FieldLabel>
<ModelMaxBudgetEditor {...editorProps} />
</Field>

View file

@ -94,6 +94,7 @@ import { Team } from "./key_team_helpers/key_list";
import { EmailEventSettingsResponse, EmailEventSettingsUpdateRequest } from "./email_events/types";
import type { ListPluginsResponse, SkillRegisterRequest } from "./claude_code_plugins/types";
import type { ModelBudgetUsage, ModelMaxBudget } from "./key_team_helpers/ModelMaxBudgetEditor";
import type { StoredModelMaxBudget } from "./key_team_helpers/modelMaxBudgetPayload";
import type { ObjectPermission } from "./object_permission_types";
import type { components } from "@/lib/http/schema";
import { fetchClient } from "@/lib/http/api";
@ -2831,6 +2832,7 @@ export interface Member {
allowed_models?: string[] | null;
temp_budget_increase?: number | null;
temp_budget_expiry?: string | null;
model_max_budget?: StoredModelMaxBudget | null;
}
export const teamMemberAddCall = async (accessToken: string, teamId: string, formValues: Member) => {
@ -2971,6 +2973,9 @@ export const teamMemberUpdateCall = async (
if ("temp_budget_expiry" in formValues) {
requestBody.temp_budget_expiry = orNull(formValues.temp_budget_expiry);
}
if (formValues.model_max_budget !== undefined) {
requestBody.model_max_budget = formValues.model_max_budget;
}
const response = await fetch(url, {
method: "POST",

View file

@ -1,4 +1,6 @@
import { useState } from "react";
import { fireEvent, screen, waitFor } from "@testing-library/react";
import userEvent from "@testing-library/user-event";
import { beforeEach, describe, expect, it, vi } from "vitest";
import { renderWithProviders } from "../../../tests/test-utils";
import EditMembership from "./EditMembership";
@ -104,6 +106,79 @@ describe("EditMembership submit payload", () => {
});
});
it("seeds the per-model budget editor from the member's stored budget and submits the edited cap", async () => {
const modelBudgetConfig = {
...teamMemberConfig,
additionalFields: [
...additionalFields,
{
name: "model_max_budget",
label: "Per-Model Budgets",
type: "model-max-budget" as const,
availableModels: ["gpt-4o"],
premiumUser: true,
},
],
};
const member = {
user_id: "u1",
user_email: "a@b.com",
role: "user",
model_max_budget: { "gpt-4o": { max_budget: 5, budget_duration: "30d" } },
};
renderEdit(modelBudgetConfig, member);
const cap = await screen.findByPlaceholderText("Max spend ($)");
expect(cap).toHaveValue(5);
fireEvent.change(cap, { target: { value: "9" } });
save();
await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce());
expect(submitted().model_max_budget).toStrictEqual({ "gpt-4o": { budget_limit: 9, time_period: "30d" } });
});
it("saves a role-only edit for a member with stored model caps", async () => {
const modelBudgetConfig = {
...teamMemberConfig,
additionalFields: [
...additionalFields,
{
name: "model_max_budget",
label: "Per-Model Budgets",
type: "model-max-budget" as const,
availableModels: ["gpt-4o"],
premiumUser: true,
},
],
};
const member = {
user_id: "u1",
user_email: "a@b.com",
role: "user",
model_max_budget: { "gpt-4o": { max_budget: 5, budget_duration: "30d" } },
};
renderEdit(modelBudgetConfig, member);
const user = userEvent.setup();
await user.click(screen.getByRole("combobox", { name: "Role(Current: User)" }));
await user.click(await screen.findByRole("option", { name: "Admin" }));
save();
await waitFor(() => expect(onSubmit).toHaveBeenCalledOnce());
const expectedPayload = {
user_email: "a@b.com",
user_id: "u1",
role: "admin",
max_budget_in_team: null,
budget_duration: null,
tpm_limit: null,
rpm_limit: null,
allowed_models: [],
model_max_budget: { "gpt-4o": { budget_limit: 5, time_period: "30d" } },
};
expect(submitted()).toStrictEqual(expectedPayload);
});
it("omits a field the config hides even when the member record carries it", async () => {
renderEdit({ ...orgMemberConfig, showUserId: false }, { user_id: "u1", user_email: "a@b.com", role: "admin" });
@ -342,3 +417,57 @@ describe("EditMembership submit payload", () => {
expect(screen.getByLabelText("Email")).toHaveValue("");
});
});
it("discards an unfinished per-model budget row when an edit session is reopened", async () => {
const user = userEvent.setup({ delay: null });
const modelBudgetConfig = {
...teamMemberConfig,
additionalFields: [
...additionalFields,
{
name: "model_max_budget",
label: "Per-Model Budgets",
type: "model-max-budget" as const,
availableModels: ["gpt-4o"],
premiumUser: true,
},
],
};
const member = { user_id: "u1", user_email: "a@b.com", role: "user", model_max_budget: {} };
const otherMember = { ...member, user_id: "u2", user_email: "c@d.com" };
const Harness = ({ initialData }: { initialData: typeof member }) => {
const [visible, setVisible] = useState(false);
return (
<>
<button type="button" onClick={() => setVisible(true)}>
Open Edit Member
</button>
<EditMembership
visible={visible}
onCancel={() => setVisible(false)}
onSubmit={vi.fn()}
mode="edit"
config={modelBudgetConfig as never}
initialData={initialData as never}
/>
</>
);
};
const { rerender } = renderWithProviders(<Harness initialData={member} />);
await user.click(screen.getByRole("button", { name: "Open Edit Member" }));
await user.click(screen.getByRole("button", { name: /Add Model Budget/i }));
expect(screen.getByPlaceholderText("Select model")).toBeInTheDocument();
await user.click(screen.getByRole("button", { name: "Cancel" }));
await user.click(screen.getByRole("button", { name: "Open Edit Member" }));
expect(screen.queryByPlaceholderText("Select model")).not.toBeInTheDocument();
await user.click(screen.getByRole("button", { name: /Add Model Budget/i }));
expect(screen.getByPlaceholderText("Select model")).toBeInTheDocument();
rerender(<Harness initialData={otherMember} />);
expect(screen.queryByPlaceholderText("Select model")).not.toBeInTheDocument();
});

View file

@ -12,11 +12,13 @@ import { Input } from "@/components/ui/input";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
import { UtcDateTimeInput } from "@/components/shared/form/UtcDateTimeInput";
import { ModelMaxBudgetEditor } from "../key_team_helpers/ModelMaxBudgetEditor";
import { useZodForm } from "@/lib/forms/useZodForm";
import {
buildMemberFormData,
buildMemberFormValues,
emptyMemberFormValues,
isModelMaxBudget,
TEMP_BUDGET_PAIR_MESSAGE,
tempBudgetPairError,
type MemberAdditionalField,
@ -51,7 +53,11 @@ const ROLE_REQUIRED_MESSAGE = "Please select a role!";
const isEmailish = (value: string): boolean =>
value === "" || z.email({ pattern: z.regexes.html5Email }).safeParse(value).success;
const memberFieldSchema = z.union([z.string(), z.number(), z.null(), z.array(z.string())]).optional();
const modelMaxBudgetSchema = z.record(z.string(), z.looseObject({ budget_limit: z.number(), time_period: z.string() }));
const memberFieldSchema = z
.union([z.string(), z.number(), z.null(), z.array(z.string()), modelMaxBudgetSchema])
.optional();
const buildMemberSchema = (config: ModalConfig): z.ZodType<MemberFormValues, MemberFormValues> => {
const shape = {
@ -171,6 +177,16 @@ const MemberModal = <T extends BaseMember>({
onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)}
/>
);
case "model-max-budget":
return (
<ModelMaxBudgetEditor
key={`${visible}-${initialData?.user_id ?? initialData?.user_email ?? ""}`}
value={isModelMaxBudget(value) ? value : {}}
onChange={onChange}
availableModels={field.availableModels ?? []}
premiumUser={field.premiumUser ?? false}
/>
);
case "utc-datetime":
return (
<UtcDateTimeInput

View file

@ -721,4 +721,98 @@ describe("TeamInfoView - member budget reset prompt", () => {
expect(screen.queryByText("Reset member budgets?")).not.toBeInTheDocument();
expect(bulkUpdatePOST).not.toHaveBeenCalled();
});
it("omits the default member model budget when it has not changed", async () => {
const user = userEvent.setup({ delay: null });
const storedBudget = { "gpt-4": { budget_limit: 3, time_period: "30d" } };
const data = createMockTeamData({
models: ["gpt-4"],
team_member_budget_table: {
max_budget: 10,
budget_duration: null,
tpm_limit: null,
rpm_limit: null,
model_max_budget: storedBudget,
},
});
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
renderWithProviders(<TeamInfoView {...props} premiumUser />);
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
await user.click(screen.getByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
await user.click(screen.getByText("Team Member Settings"));
await screen.findByPlaceholderText("Max spend ($)");
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1]).not.toHaveProperty("team_member_model_max_budget");
});
it("sends the default member model budget after an edit", async () => {
const user = userEvent.setup({ delay: null });
const storedBudget = { "gpt-4": { budget_limit: 3, time_period: "30d" } };
const data = createMockTeamData({
models: ["gpt-4"],
team_member_budget_table: {
max_budget: 10,
budget_duration: null,
tpm_limit: null,
rpm_limit: null,
model_max_budget: storedBudget,
},
});
vi.mocked(networking.teamInfoCall).mockResolvedValue(data);
vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: savedTeam, team_id: "123" });
renderWithProviders(<TeamInfoView {...props} premiumUser />);
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
await user.click(screen.getByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
await user.click(screen.getByText("Team Member Settings"));
fireEvent.change(await screen.findByPlaceholderText("Max spend ($)"), { target: { value: "7" } });
await user.click(screen.getByRole("button", { name: /save changes/i }));
await waitFor(() => expect(networking.teamUpdateCall).toHaveBeenCalled());
expect(vi.mocked(networking.teamUpdateCall).mock.calls[0][1].team_member_model_max_budget).toStrictEqual({
"gpt-4": { budget_limit: 7, time_period: "30d" },
});
});
it("offers access-group models in the default member budget picker", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(
createMockTeamData({
models: ["gpt-4o"],
access_group_models: ["claude-haiku-4-5"],
}) as TeamData,
);
renderWithProviders(<TeamInfoView {...props} userModels={["gpt-4o", "claude-haiku-4-5"]} premiumUser />);
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
await user.click(screen.getByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
await user.click(screen.getByText("Team Member Settings"));
await user.click(screen.getAllByRole("button", { name: /add model budget/i })[0]);
await user.click(screen.getByPlaceholderText("Select model"));
expect(await screen.findByRole("option", { name: "claude-haiku-4-5" })).toBeInTheDocument();
});
it("offers user models in the default member budget picker when the team model list is empty", async () => {
const user = userEvent.setup({ delay: null });
vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: [] }) as TeamData);
renderWithProviders(<TeamInfoView {...props} userModels={["gpt-4o", "claude-haiku-4-5"]} premiumUser />);
await waitFor(() => expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0));
await user.click(screen.getByRole("tab", { name: "Settings" }));
await user.click(await screen.findByRole("button", { name: /edit settings/i }));
await user.click(screen.getByText("Team Member Settings"));
await user.click(screen.getAllByRole("button", { name: /add model budget/i })[0]);
await user.click(screen.getByPlaceholderText("Select model"));
expect(await screen.findByRole("option", { name: "gpt-4o" })).toBeInTheDocument();
expect(screen.getByRole("option", { name: "claude-haiku-4-5" })).toBeInTheDocument();
});
});

View file

@ -286,7 +286,7 @@ export interface TeamMembership {
max_parallel_requests: number | null;
tpm_limit: number | null;
rpm_limit: number | null;
model_max_budget: Record<string, number> | null;
model_max_budget: StoredModelMaxBudget | null;
budget_duration: string | null;
budget_reset_at: string | null;
allowed_models?: string[] | null;
@ -339,6 +339,7 @@ export interface TeamData {
budget_duration: string | null;
tpm_limit: number | null;
rpm_limit: number | null;
model_max_budget?: StoredModelMaxBudget | null;
} | null;
};
keys: any[];
@ -631,6 +632,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const [isTeamSaving, setIsTeamSaving] = useState(false);
const [teamModelAliases, setTeamModelAliases] = useState<Record<string, string>>({});
const [teamModelMaxBudget, setTeamModelMaxBudget] = useState<ModelMaxBudget>({});
const [teamMemberModelMaxBudget, setTeamMemberModelMaxBudget] = useState<ModelMaxBudget>({});
const routerSettingsRef = React.useRef<RouterSettingsAccordionRef>(null);
const [organization, setOrganization] = useState<Organization | null>(null);
const { userRole, isViewOnly } = useAuthorized();
@ -661,12 +663,20 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
] as const
).find(([failed]) => failed)?.[1] ?? null;
const availableRateLimitModels = useMemo(() => {
const selected = watchedModels ?? teamData?.team_info?.models ?? [];
const selected = (isEditing ? watchedModels : undefined) ?? teamData?.team_info?.models ?? [];
if (selected.includes("all-proxy-models") || selected.includes("all-team-models")) {
return userModels;
}
return unfurlWildcardModelsInList(selected, userModels);
}, [watchedModels, teamData, userModels]);
}, [isEditing, watchedModels, teamData, userModels]);
const availableModelBudgetModels = useMemo(() => {
const selected = (isEditing ? watchedModels : undefined) ?? teamData?.team_info?.models ?? [];
if (selected.length === 0 || selected.includes("all-proxy-models") || selected.includes("all-team-models")) {
return userModels;
}
const granted = [...selected, ...(teamData?.team_info?.access_group_models ?? [])];
return unfurlWildcardModelsInList(Array.from(new Set(granted)), userModels);
}, [isEditing, watchedModels, teamData, userModels]);
const teamEditAccess = useMemo(() => parseTeamEditAccess(teamData?.team_info?.caller_edit_access), [teamData]);
const canEditTeam = is_team_admin || is_proxy_admin || teamEditAccess.kind !== "none";
@ -684,6 +694,9 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
const startEditing = () => {
form.reset(teamFormValues());
setTeamModelMaxBudget((teamData?.team_info?.model_max_budget ?? {}) as ModelMaxBudget);
setTeamMemberModelMaxBudget(
(teamData?.team_info?.team_member_budget_table?.model_max_budget ?? {}) as ModelMaxBudget,
);
setTeamMemberSettingsOpen(false);
setSearchToolSettingsOpen(false);
setIsEditing(true);
@ -866,6 +879,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
allowed_models: values.allowed_models,
temp_budget_increase: values.temp_budget_increase,
temp_budget_expiry: values.temp_budget_expiry,
model_max_budget: modelMaxBudgetUpdate(values.model_max_budget ?? {}, selectedEditMember?.model_max_budget),
};
toast.dismiss(); // Remove all existing toasts
@ -1223,6 +1237,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
updateData.model_max_budget = modelBudgets;
}
const teamMemberModelBudgets = modelMaxBudgetUpdate(
teamMemberModelMaxBudget,
info.team_member_budget_table?.model_max_budget,
);
if (teamMemberModelBudgets !== undefined) {
updateData.team_member_model_max_budget = teamMemberModelBudgets;
}
// Handle router_settings - read fresh values from DOM at save time.
const currentRouterSettings = routerSettingsRef.current?.getValue();
if (currentRouterSettings?.router_settings) {
@ -1658,6 +1680,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
/>
)}
</FormField>
<ModelMaxBudgetField
label="Default Per-Model Budget"
premiumUser={premiumUser}
value={teamMemberModelMaxBudget}
onChange={setTeamMemberModelMaxBudget}
availableModels={availableModelBudgetModels}
hint="Set a default per-model spend cap for each team member, with an independent reset window."
/>
<FormField
control={form.control}
name="team_member_key_duration"
@ -2347,6 +2377,14 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
</p>
<div>Max Budget: {info.team_member_budget_table?.max_budget ?? "No Limit"}</div>
<div>Budget Duration: {info.team_member_budget_table?.budget_duration || "No Limit"}</div>
{modelMaxBudgetToEntries(
info.team_member_budget_table?.model_max_budget as ModelMaxBudget | null | undefined,
).map(({ model, budgetLimit, timePeriod }) => (
<div key={model}>
Default Per-Model Budget ({model}):{" "}
{budgetLimit !== null ? `$${formatNumberWithCommas(budgetLimit, 4)}` : "No Limit"} per {timePeriod}
</div>
))}
<div>Key Duration: {info.metadata?.team_member_key_duration || "No Limit"}</div>
<div>TPM Limit: {info.team_member_budget_table?.tpm_limit ?? "No Limit"}</div>
<div>RPM Limit: {info.team_member_budget_table?.rpm_limit ?? "No Limit"}</div>
@ -2554,6 +2592,20 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
),
type: "utc-datetime" as const,
},
{
name: "model_max_budget",
label: (
<span>
Per-Model Budgets{" "}
<SimpleTooltip content="Spend cap per model for this member, each with its own reset period. Starts from the team's default per-model budget; saving a change stores a per-model budget for this member only.">
<Info className="ml-1 inline size-3.5 align-text-bottom" />
</SimpleTooltip>
</span>
),
type: "model-max-budget" as const,
availableModels: availableModelBudgetModels,
premiumUser,
},
{
name: "tpm_limit",
label: (

View file

@ -432,10 +432,31 @@ describe("TeamMembersComponent", () => {
allowed_models: [],
temp_budget_increase: 0,
temp_budget_expiry: "2030-01-02T03:04:00Z",
model_max_budget: null,
};
expect(seedMemberBudgetFields({ user_id: "user1@test.com", role: "member" }, budget)).toStrictEqual(seeded);
});
it("seeds the team's default per-model budget when the member has no per-model override", () => {
const teamDefault = { "gpt-4o": { max_budget: 10, budget_duration: "30d" } };
const ownRowWithoutModelCaps = {
...createMockTeamData().team_memberships[0].litellm_budget_table,
model_max_budget: {},
};
const memberOverride = { "gpt-4o": { max_budget: 2.5, budget_duration: "7d" } };
const ownRowWithModelCaps = { ...ownRowWithoutModelCaps, model_max_budget: memberOverride };
expect(
seedMemberBudgetFields({ user_id: "u", role: "user" }, ownRowWithoutModelCaps, teamDefault).model_max_budget,
).toBe(teamDefault);
expect(seedMemberBudgetFields({ user_id: "u", role: "user" }, undefined, teamDefault).model_max_budget).toBe(
teamDefault,
);
expect(
seedMemberBudgetFields({ user_id: "u", role: "user" }, ownRowWithModelCaps, teamDefault).model_max_budget,
).toBe(memberOverride);
});
it("seeds null temporary budget fields for a member without a budget row", () => {
expect(seedMemberBudgetFields({ user_id: "user2@test.com", role: "admin" }, undefined)).toMatchObject({
max_budget_in_team: null,

View file

@ -7,6 +7,7 @@ import { Button } from "@/components/ui/button";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
import { SimpleTooltip } from "@/components/ui/tooltip";
import MemberTable from "@/components/common_components/MemberTable";
import type { StoredModelMaxBudget } from "@/components/key_team_helpers/modelMaxBudgetPayload";
import { Member } from "@/components/networking";
import { parseErrorMessage } from "@/components/shared/errorUtils";
import { DateCell, MoneyCell } from "@/components/shared/table_cells";
@ -25,9 +26,13 @@ const BUDGET_SOURCE_LABELS: Record<Exclude<TeamMemberBudgetSource, "none">, stri
const formatBudget = (value: number | null): string =>
value === null ? "Unlimited" : `$${formatNumberWithCommas(value, 2)}`;
const hasModelBudgets = (budget: StoredModelMaxBudget | null | undefined): budget is StoredModelMaxBudget =>
budget != null && Object.keys(budget).length > 0;
export const seedMemberBudgetFields = (
record: Member,
budget: TeamMembership["litellm_budget_table"] | undefined,
teamDefaultModelMaxBudget?: StoredModelMaxBudget | null,
): Member => ({
...record,
max_budget_in_team: budget?.max_budget ?? null,
@ -37,6 +42,9 @@ export const seedMemberBudgetFields = (
allowed_models: budget?.allowed_models || [],
temp_budget_increase: budget?.temp_budget_increase ?? null,
temp_budget_expiry: budget?.temp_budget_expiry ?? null,
model_max_budget: hasModelBudgets(budget?.model_max_budget)
? budget.model_max_budget
: teamDefaultModelMaxBudget ?? null,
});
interface TeamMemberTabProps {
@ -305,7 +313,13 @@ export default function TeamMemberTab({
canEdit={canEditTeam}
onEdit={(record) => {
const membership = teamData.team_memberships.find((tm) => tm.user_id === record.user_id);
setSelectedEditMember(seedMemberBudgetFields(record, membership?.litellm_budget_table));
setSelectedEditMember(
seedMemberBudgetFields(
record,
membership?.litellm_budget_table,
teamData.team_info.team_member_budget_table?.model_max_budget,
),
);
setIsEditMemberModalVisible(true);
}}
onDelete={handleMemberDelete}

View file

@ -1,8 +1,21 @@
export type MemberFieldValue = string | number | null | undefined | string[];
import {
entriesToModelMaxBudget,
modelMaxBudgetToEntries,
type ModelMaxBudget,
} from "../key_team_helpers/ModelMaxBudgetEditor";
export type MemberFieldValue = string | number | null | undefined | string[] | ModelMaxBudget;
export type MemberFormValues = Record<string, MemberFieldValue>;
export type MemberFieldType = "input" | "select" | "numerical" | "multi-select" | "budget-duration" | "utc-datetime";
export type MemberFieldType =
| "input"
| "select"
| "numerical"
| "multi-select"
| "budget-duration"
| "utc-datetime"
| "model-max-budget";
export interface MemberAdditionalField {
name: string;
@ -12,6 +25,8 @@ export interface MemberAdditionalField {
step?: number;
min?: number;
placeholder?: string;
availableModels?: string[];
premiumUser?: boolean;
}
export interface MemberFieldsConfig {
@ -33,6 +48,9 @@ export const TEMP_BUDGET_PAIR_MESSAGE = "Set both a temporary budget increase an
const isUnset = (value: MemberFieldValue): boolean => value === null || value === undefined || value === "";
export const isModelMaxBudget = (value: MemberFieldValue): value is ModelMaxBudget =>
typeof value === "object" && value !== null && !Array.isArray(value);
export const tempBudgetPairError = (values: MemberFormValues): "temp_budget_increase" | "temp_budget_expiry" | null => {
const increaseUnset = isUnset(values.temp_budget_increase);
const expiryUnset = isUnset(values.temp_budget_expiry);
@ -66,6 +84,9 @@ export const buildMemberFormValues = (
allowed_models: initialData.allowed_models || [],
temp_budget_increase: initialData.temp_budget_increase ?? null,
temp_budget_expiry: initialData.temp_budget_expiry || null,
model_max_budget: isModelMaxBudget(initialData.model_max_budget)
? entriesToModelMaxBudget(modelMaxBudgetToEntries(initialData.model_max_budget))
: {},
};
return pickFieldNames(config, seeded);
@ -78,6 +99,8 @@ const emptyValueForType = (type: MemberFieldType | undefined): MemberFieldValue
switch (type) {
case "multi-select":
return [];
case "model-max-budget":
return {};
case "numerical":
case "budget-duration":
case "utc-datetime":

View file

@ -16882,6 +16882,7 @@ export interface paths {
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
* - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
* - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
* - team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
* - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
* - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
* - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
@ -17111,6 +17112,7 @@ export interface paths {
* - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
* - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
* - team_member_budget_duration: Optional[str] - The duration of the budget for the team member. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
* - team_member_model_max_budget: Optional[dict] - Per-model budgets shared by team members without an override.
* - team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
* - team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
* - team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
@ -38499,6 +38501,10 @@ export interface components {
team_member_budget_duration?: string | null;
/** Team Member Key Duration */
team_member_key_duration?: string | null;
/** Team Member Model Max Budget */
team_member_model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/** Team Member Permissions */
team_member_permissions?: string[] | null;
/** Team Member Rpm Limit */
@ -39584,6 +39590,10 @@ export interface components {
team_member_budget_duration?: string | null;
/** Team Member Key Duration */
team_member_key_duration?: string | null;
/** Team Member Model Max Budget */
team_member_model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/** Team Member Rpm Limit */
team_member_rpm_limit?: number | null;
/** Team Member Tpm Limit */
@ -45580,6 +45590,13 @@ export interface components {
budget_duration?: string | null;
/** Max Budget In Team */
max_budget_in_team?: number | null;
/**
* Model Max Budget
* @description Per-model spend caps for this team member, each with its own budget_duration. Overrides the team's default per-model member budget. Pass an empty dict to fall back to the team default.
*/
model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/** Role */
role?: ("admin" | "user") | null;
/**
@ -45617,6 +45634,10 @@ export interface components {
budget_duration?: string | null;
/** Max Budget In Team */
max_budget_in_team?: number | null;
/** Model Max Budget */
model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/** Rpm Limit */
rpm_limit?: number | null;
/** Team Id */
@ -47559,6 +47580,10 @@ export interface components {
team_member_budget_duration?: string | null;
/** Team Member Key Duration */
team_member_key_duration?: string | null;
/** Team Member Model Max Budget */
team_member_model_max_budget?: {
[key: string]: components["schemas"]["BudgetConfig"];
} | null;
/** Team Member Rpm Limit */
team_member_rpm_limit?: number | null;
/** Team Member Tpm Limit */