diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py index f743320b702..85a20ea4435 100644 --- a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py @@ -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( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 26d9e0e1b44..7ac8b1f8538 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 7c62e4f4147..d687b2ad7a0 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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: diff --git a/litellm/proxy/auth/fallback_budget.py b/litellm/proxy/auth/fallback_budget.py index 5aa28a08bb8..a126a53821b 100644 --- a/litellm/proxy/auth/fallback_budget.py +++ b/litellm/proxy/auth/fallback_budget.py @@ -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) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d6a51bc2fc9..2231706b282 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, ) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 7f0e0b1fc87..ed9ac258e16 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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: diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index 8bca51784f2..eaa185917e5 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -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) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 913310952cd..0d693b85ce8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index d001dc666af..b3f372327dc 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -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 diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 6a70a0137a4..adf89c29279 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c78434eda83..544bc471ce5 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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, ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 5d73a04f9e0..d6358c638e7 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -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 ( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index dac46855fd6..089647c6297 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 4d6af710bb9..6976a64ef39 100644 --- a/litellm/router_utils/common_utils.py +++ b/litellm/router_utils/common_utils.py @@ -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. diff --git a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py index ced53794cc2..494e742e426 100644 --- a/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py +++ b/tests/unit/llms/anthropic/pass_through/context_management/test_compact.py @@ -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": "x", + }, + }, + ] + ) + 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 diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index e26845e5b2a..8b712772230 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -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 diff --git a/tests/unit/proxy/auth/test_fallback_budget.py b/tests/unit/proxy/auth/test_fallback_budget.py index 00c1a7cdefc..3d31f1e0c78 100644 --- a/tests/unit/proxy/auth/test_fallback_budget.py +++ b/tests/unit/proxy/auth/test_fallback_budget.py @@ -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): """ diff --git a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py index 1a636e28fb9..492ab408fc8 100644 --- a/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py +++ b/tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py @@ -5,10 +5,11 @@ import os import subprocess import sys import time -from collections.abc import Mapping -from contextlib import contextmanager +from collections.abc import Mapping, Sequence +from contextlib import contextmanager, nullcontext from datetime import datetime, timedelta, timezone from functools import partial +from io import BytesIO from pathlib import Path from textwrap import dedent from types import SimpleNamespace @@ -16,13 +17,24 @@ from typing import Final from unittest.mock import ANY, AsyncMock, MagicMock, patch +import orjson import pytest -from fastapi import HTTPException, status +from fastapi import HTTPException, Request, status +from starlette.datastructures import UploadFile import litellm import litellm.proxy.proxy_server from litellm.caching.dual_cache import DualCache -from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.common_utils.user_api_key_cache import ( + NO_TEAM_MEMBERSHIP_SENTINEL, + UserApiKeyCache, + team_membership_reservation_cache_key, +) +from litellm.proxy.hooks.model_max_budget_limiter import ( + PROXY_VirtualKeyModelMaxBudgetLimiter, + model_budget_spend_cache_key, + team_member_budget_entity_id, +) from litellm.proxy._types import ( LiteLLMRoutes, LiteLLM_JWTAuth, @@ -30,6 +42,9 @@ from litellm.proxy._types import ( LiteLLM_EndUserTable, LiteLLM_ObjectPermissionTable, LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TeamMembership, + LiteLLM_TeamTable, LiteLLM_TeamTableCachedObj, LiteLLM_UserTable, Litellm_EntityType, @@ -52,8 +67,10 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.http_parsing_utils import get_client_requested_model from litellm.proxy.auth.user_api_key_auth import ( _check_key_model_budget_with_fallback, + _check_team_member_model_budget_with_fallback, _ensure_litellm_received_at_on_request_state, _ensure_parent_otel_span_on_request_state, + _load_team_member_default_model_budget, _PendingAutoRegister, _matches_routing_override, _reserve_budget_after_common_checks, @@ -78,6 +95,82 @@ class _RoutingRequest: self.state = SimpleNamespace() +async def _real_model_budget_limiter( + *, + member_spend: Mapping[str, float] | None = None, + key_spend: Mapping[str, float] | None = None, + member_errors: Mapping[str, Exception] | None = None, +) -> tuple[PROXY_VirtualKeyModelMaxBudgetLimiter, MagicMock]: + member_entity_id: Final = team_member_budget_entity_id(user_id="user-1", team_id="team-1") + scope_spend: Final = ( + *( + (Litellm_EntityType.TEAM_MEMBER, member_entity_id, model, spend) + for model, spend in (member_spend or {}).items() + ), + *((Litellm_EntityType.KEY, "test-key", model, spend) for model, spend in (key_spend or {}).items()), + ) + scope_errors: Final = tuple( + (Litellm_EntityType.TEAM_MEMBER, member_entity_id, model, error) + for model, error in (member_errors or {}).items() + ) + spend_by_key: Final = { + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=model, + budget_duration="1d", + ): spend + for entity_type, entity_id, model, spend in scope_spend + } + errors_by_key: Final = { + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=model, + budget_duration="1d", + ): error + for entity_type, entity_id, model, error in scope_errors + } + redis_cache: Final = MagicMock() + + async def get_spend(*, key: str) -> object | None: + if key in errors_by_key: + raise errors_by_key[key] + 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 _model_budget_spend_keys( + entity_type: Litellm_EntityType, + entity_id: str, + models: Sequence[str], +) -> list[str]: + return [ + model_budget_spend_cache_key( + entity_type=entity_type, + entity_id=entity_id, + budget_model=model, + budget_duration="1d", + ) + for model in models + ] + + +def _team_member_model_budget_spend_keys(models: Sequence[str]) -> list[str]: + return _model_budget_spend_keys( + Litellm_EntityType.TEAM_MEMBER, + team_member_budget_entity_id(user_id="user-1", team_id="team-1"), + models, + ) + + +def _assert_budget_exceeded_values(error: litellm.BudgetExceededError) -> None: + assert error.current_cost == 10 + assert error.max_budget == 5 + + def test_get_api_key(): bearer_token = "Bearer sk-98765678" api_key = "sk-98765678" @@ -4730,9 +4823,11 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + database = MagicMock() + database.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) attrs = { **_proxy_attrs_for_centralized_checks(user_custom_auth=None), - "prisma_client": MagicMock(), + "prisma_client": database, "user_api_key_cache": user_api_key_cache, "proxy_logging_obj": proxy_logging_obj, } @@ -5011,6 +5106,503 @@ async def test_centralized_common_checks_enforces_team_model_max_budget_from_the assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")] +@pytest.mark.asyncio +async def test_centralized_common_checks_loads_member_budget_and_exempts_discovery_routes(): + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter + import litellm.proxy.proxy_server as _proxy_server_mod + + budget_id: Final = "member-budget-1" + member_caps: Final = {"gpt-4o": {"max_budget": 0.0, "budget_duration": "1d"}} + membership: Final = LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + budget_id=budget_id, + litellm_budget_table=LiteLLM_BudgetTable( + budget_id=budget_id, + model_max_budget={"gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"}}, + ), + ) + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="u1", team_id="t1"), + value=membership, + model_type=LiteLLM_TeamMembership, + ) + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj( + team_id="t1", + models=["gpt-4o"], + metadata={"team_member_budget_id": budget_id}, + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + 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_caps), + model_type=LiteLLM_BudgetTable, + ) + await user_api_key_cache.async_set_cache( + key="u1", + value=LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER), + model_type=LiteLLM_UserTable, + ) + router: Final = litellm.Router( + model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], + ) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + attrs: Final = { + "general_settings": {"disable_budget_reservation": True}, + "litellm_proxy_admin_name": "admin", + "llm_model_list": router.model_list, + "llm_router": router, + "master_key": "sk-master-key", + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "proxy_logging_obj": MagicMock(), + "model_max_budget_limiter": limiter, + "user_custom_auth": None, + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + token: Final = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + user_id="u1", + team_id="t1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + request: Final = Request( + scope={ + "type": "http", + "method": "POST", + "path": "/v1/models", + "raw_path": b"/v1/models", + "query_string": b"", + "headers": [], + "server": ("test", 80), + "client": ("test", 1), + "scheme": "http", + } + ) + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/v1/models", + ) + + assert token.team_member_model_max_budget == member_caps + + non_exempt_token: Final = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + user_id="u1", + team_id="t1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + non_exempt_request: Final = Request( + scope={ + "type": "http", + "method": "POST", + "path": "/v1/chat/completions", + "raw_path": b"/v1/chat/completions", + "query_string": b"", + "headers": [], + "server": ("test", 80), + "client": ("test", 1), + "scheme": "http", + } + ) + with pytest.raises(litellm.BudgetExceededError): + await run_centralized_common_checks( + user_api_key_auth_obj=non_exempt_token, + request=non_exempt_request, + request_data={"model": "gpt-4o"}, + route="/v1/chat/completions", + ) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + +@pytest.mark.parametrize( + ("route", "path", "must_raise"), + [ + pytest.param("/v1/models", "/v1/models", False, id="budget-exempt-discovery-route"), + pytest.param("/chat/completions", "/v1/chat/completions", True, id="paid-model-route"), + ], +) +@pytest.mark.asyncio +async def test_centralized_common_checks_skips_default_member_budget_lookup_errors_on_exempt_routes( + route: str, + path: str, + must_raise: bool, +): + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.hooks.model_max_budget_limiter import PROXY_VirtualKeyModelMaxBudgetLimiter + import litellm.proxy.proxy_server as _proxy_server_mod + + budget_id: Final = "member-budget-lookup-fails" + membership: Final = LiteLLM_TeamMembership( + user_id="u1", + team_id="t1", + budget_id=budget_id, + litellm_budget_table=LiteLLM_BudgetTable( + budget_id=budget_id, + model_max_budget={"gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"}}, + ), + ) + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="u1", team_id="t1"), + value=membership, + model_type=LiteLLM_TeamMembership, + ) + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj( + team_id="t1", + models=["gpt-4o"], + metadata={"team_member_budget_id": budget_id}, + ), + model_type=LiteLLM_TeamTableCachedObj, + ) + await user_api_key_cache.async_set_cache( + key="u1", + value=LiteLLM_UserTable(user_id="u1", user_role=LitellmUserRoles.INTERNAL_USER), + model_type=LiteLLM_UserTable, + ) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + side_effect=RuntimeError("budget database unavailable") + ) + router: Final = litellm.Router( + model_list=[{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], + ) + limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + attrs: Final = { + "general_settings": {"disable_budget_reservation": True}, + "litellm_proxy_admin_name": "admin", + "llm_model_list": router.model_list, + "llm_router": router, + "master_key": "sk-master-key", + "prisma_client": prisma_client, + "user_api_key_cache": user_api_key_cache, + "proxy_logging_obj": MagicMock(), + "model_max_budget_limiter": limiter, + "user_custom_auth": None, + } + originals: Final = {attribute: getattr(_proxy_server_mod, attribute, None) for attribute in attrs} + + try: + for attribute, value in attrs.items(): + setattr(_proxy_server_mod, attribute, value) + token: Final = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + user_id="u1", + team_id="t1", + user_role=LitellmUserRoles.PROXY_ADMIN, + ) + request: Final = Request( + scope={ + "type": "http", + "method": "POST", + "path": path, + "raw_path": path.encode(), + "query_string": b"", + "headers": [], + "server": ("test", 80), + "client": ("test", 1), + "scheme": "http", + } + ) + with pytest.raises(RuntimeError, match="budget database unavailable") if must_raise else nullcontext(): + await run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route=route, + ) + + assert token.team_member_model_max_budget is None + prisma_client.db.litellm_budgettable.find_unique.assert_awaited_once_with(where={"budget_id": budget_id}) + finally: + for attribute, value in originals.items(): + setattr(_proxy_server_mod, attribute, value) + + +def _team_member_model_budget_load_inputs() -> tuple[ + UserAPIKeyAuth, LiteLLM_TeamTableCachedObj, MagicMock, UserApiKeyCache +]: + token: Final = UserAPIKeyAuth(user_id="member-1", team_id="team-1") + team_object: Final = LiteLLM_TeamTableCachedObj( + team_id="team-1", + metadata={"team_member_budget_id": "default-budget"}, + ) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_budgettable.find_unique = AsyncMock() + user_api_key_cache: Final = UserApiKeyCache() + return token, team_object, prisma_client, user_api_key_cache + + +@pytest.mark.asyncio +async def test_member_model_budget_load_prefers_override_row() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + member_caps: Final = {"claude-sonnet-4-6": {"max_budget": 2.0, "budget_duration": "1d"}} + membership: Final = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="member-budget", model_max_budget=member_caps), + ) + 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 _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget == member_caps + prisma_client.db.litellm_budgettable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_member_model_budget_load_uses_override_without_team_default_id() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + team_object = LiteLLM_TeamTableCachedObj(team_id=team_object.team_id, metadata={}) + member_caps: Final = {"claude-sonnet-4-6": {"max_budget": 2.0, "budget_duration": "1d"}} + membership: Final = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="member-budget", model_max_budget=member_caps), + ) + 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 _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget == member_caps + prisma_client.db.litellm_budgettable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_member_model_budget_load_reads_current_default_for_plain_team_object() -> None: + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + token: Final = UserAPIKeyAuth(user_id="member-1", team_id="team-1") + team_object: Final = LiteLLM_TeamTable( + team_id="team-1", + metadata={"team_member_budget_id": "default-budget"}, + ) + prisma_client: Final = MagicMock() + user_api_key_cache: Final = UserApiKeyCache() + stale_caps: Final = {"claude-sonnet-4-6": {"max_budget": 3.0, "budget_duration": "1d"}} + current_caps: Final = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}} + membership: Final = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="default-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="default-budget", model_max_budget=stale_caps), + ) + 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=current_caps), + model_type=LiteLLM_BudgetTable, + ) + + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget == current_caps + + +@pytest.mark.asyncio +async def test_member_model_budget_load_overall_only_override_keeps_default_caps() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + default_caps: Final = {"claude-sonnet-4-6": {"max_budget": 3.0, "budget_duration": "1d"}} + membership: Final = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="member-budget", max_budget=5.0, model_max_budget=None), + ) + 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, + ) + + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget == default_caps + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cleared_member_caps", [None, {}], ids=["null", "empty"]) +async def test_member_model_budget_load_falls_back_when_override_has_no_model_caps( + cleared_member_caps: dict[str, object] | None, +) -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + default_caps: Final = {"claude-sonnet-4-6": {"max_budget": 3.0, "budget_duration": "1d"}} + membership: Final = LiteLLM_TeamMembership( + user_id="member-1", + team_id="team-1", + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="member-budget", model_max_budget=cleared_member_caps), + ) + default_budget: Final = LiteLLM_BudgetTable(budget_id="default-budget", model_max_budget=default_caps) + 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=default_budget, + model_type=LiteLLM_BudgetTable, + ) + + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget == default_caps + + +@pytest.mark.asyncio +async def test_member_model_budget_load_uses_default_without_membership() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + default_caps: Final = {"claude-sonnet-4-6": {"max_budget": 3.0, "budget_duration": "1d"}} + default_budget: Final = LiteLLM_BudgetTable(budget_id="default-budget", model_max_budget=default_caps) + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + await user_api_key_cache.async_set_cache( + key="team_member_default_budget:default-budget", + value=default_budget, + model_type=LiteLLM_BudgetTable, + ) + + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget == default_caps + + +@pytest.mark.asyncio +async def test_member_model_budget_load_fails_closed_on_default_lookup_error() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + side_effect=RuntimeError("budget database unavailable") + ) + + with pytest.raises(RuntimeError, match="budget database unavailable"): + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + skip_budget_checks=False, + proxy_logging_obj=None, + ) + + +@pytest.mark.asyncio +async def test_member_model_budget_load_skips_default_lookup_error_when_checks_are_skipped() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock( + side_effect=RuntimeError("budget database unavailable") + ) + + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + skip_budget_checks=True, + proxy_logging_obj=None, + ) + + assert token.team_member_model_max_budget is None + + +@pytest.mark.asyncio +async def test_member_model_budget_load_allows_missing_default_row() -> None: + token, team_object, prisma_client, user_api_key_cache = _team_member_model_budget_load_inputs() + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_budgettable.find_unique = AsyncMock(return_value=None) + + await _load_team_member_default_model_budget( + user_api_key_auth_obj=token, + 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 token.team_member_model_max_budget is None + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without @@ -5331,11 +5923,13 @@ async def test_team_key_vector_store_access_when_team_cannot_be_resolved( request._url = URL(url="/v1/rag/query") database: Final = MagicMock() database.db.litellm_objectpermissiontable.find_unique = AsyncMock( - side_effect=lambda where: LiteLLM_ObjectPermissionTable( - object_permission_id=where["object_permission_id"], vector_stores=["KBSTOREA"] + side_effect=lambda where: ( + LiteLLM_ObjectPermissionTable( + object_permission_id=where["object_permission_id"], vector_stores=["KBSTOREA"] + ) + if where["object_permission_id"] == "key-permission" + else None ) - if where["object_permission_id"] == "key-permission" - else None ) attrs: Final = { **_proxy_attrs_for_centralized_checks(), @@ -5387,6 +5981,7 @@ async def test_keyless_team_member_vector_store_access_uses_only_the_resolved_te ): from fastapi import Request from starlette.datastructures import URL + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache token: Final = UserAPIKeyAuth( user_id="user-1", @@ -5408,9 +6003,27 @@ async def test_keyless_team_member_vector_store_access_uses_only_the_resolved_te object_permission_id=where["object_permission_id"], vector_stores=grants[where["object_permission_id"]] ) ) + database.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + if isinstance(team_lookup, Exception): + database.db.litellm_teamtable.find_unique = AsyncMock(side_effect=team_lookup) + else: + database.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=lambda where, **_: LiteLLM_TeamTable( + team_id=where["team_id"], + models=["gpt-4o-mini"], + object_permission_id=team_lookup[where["team_id"]], + ) + ) + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key="user-1", + value=LiteLLM_UserTable(user_id="user-1", teams=["team-1", "team-2"], object_permission_id="user-grants-a"), + model_type=LiteLLM_UserTable, + ) attrs: Final = { **_proxy_attrs_for_centralized_checks(), "prisma_client": database, + "user_api_key_cache": user_api_key_cache, "proxy_logging_obj": MagicMock(service_logging_obj=MagicMock(async_service_success_hook=AsyncMock())), "general_settings": {"vector_store_deny_by_default": True}, } @@ -5418,22 +6031,6 @@ async def test_keyless_team_member_vector_store_access_uses_only_the_resolved_te monkeypatch.setattr(litellm.proxy.proxy_server, name, value) monkeypatch.setattr(litellm, "vector_store_registry", None) - async def get_team(team_id: str, **_: object) -> LiteLLM_TeamTableCachedObj: - if isinstance(team_lookup, Exception): - raise team_lookup - return LiteLLM_TeamTableCachedObj(team_id=team_id, models=["gpt-4o-mini"], object_permission_id=team_lookup[team_id]) - - monkeypatch.setattr("litellm.proxy.auth.user_api_key_auth.get_team_object", get_team) - monkeypatch.setattr("litellm.proxy.auth.auth_checks.get_team_membership", AsyncMock(return_value=None)) - monkeypatch.setattr( - "litellm.proxy.auth.user_api_key_auth.get_user_object", - AsyncMock( - return_value=LiteLLM_UserTable( - user_id="user-1", teams=["team-1", "team-2"], object_permission_id="user-grants-a" - ) - ), - ) - checks: Final = run_centralized_common_checks( user_api_key_auth_obj=token, request=request, @@ -6683,6 +7280,7 @@ async def test_centralized_common_checks_inherits_org_identity( attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None) attrs["prisma_client"] = MagicMock() + attrs["prisma_client"].db.litellm_teammembership.find_unique = AsyncMock(return_value=None) attrs["general_settings"] = {"allow_requests_on_db_unavailable": allow_db_unavailable} originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} try: @@ -7690,7 +8288,8 @@ async def _authenticate_session_token_against_db( from starlette.datastructures import URL attrs = _proxy_attrs_for_db_lookup() - attrs["prisma_client"].db.litellm_teammembership.find_first = AsyncMock(return_value=None) + attrs["prisma_client"].db.litellm_teammembership.find_unique = AsyncMock(return_value=membership_row) + attrs["prisma_client"].db.litellm_teammembership.find_first = AsyncMock(return_value=membership_row) originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} assemble = AsyncMock(return_value=UserAPIKeyAuth(user_id=SESSION_USER_ID, is_session_token=True)) try: @@ -7710,10 +8309,6 @@ async def _authenticate_session_token_against_db( "litellm.proxy.auth.user_api_key_auth.get_team_object", AsyncMock(return_value=team_row, side_effect=team_error), ), - patch( # test-quality-ok: the builder reads its DB row loaders off module globals - "litellm.proxy.auth.user_api_key_auth.get_team_membership", - AsyncMock(return_value=membership_row), - ), ): await user_api_key_auth_builder( request=request, @@ -8226,25 +8821,16 @@ class TestCheckKeyModelBudgetWithFallback: request_data = {"model": "gpt-4o"} request = self._make_request() - with patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - side_effect=ProxyException( - message="model not allowed", - type=ProxyErrorTypes.budget_exceeded, - param="model", - code=status.HTTP_403_FORBIDDEN, - ), - ): - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_key_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + model_name="gpt-4o", + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=None, + ) assert exc_info.value is original_error assert request_data["model"] == "gpt-4o" @@ -8264,19 +8850,15 @@ class TestCheckKeyModelBudgetWithFallback: request_data = {"model": "gpt-4o"} request = self._make_request() - with patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - return_value=True, - ): - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + await _check_key_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + model_name="gpt-4o", + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=None, + ) assert request_data["model"] == "gpt-4o-mini" @@ -8298,20 +8880,16 @@ class TestCheckKeyModelBudgetWithFallback: request_data = {"model": "gpt-4o"} request = self._make_request() - with patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - return_value=True, - ): - with pytest.raises(litellm.BudgetExceededError) as exc_info: - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_key_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + model_name="gpt-4o", + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=None, + ) assert exc_info.value is original_error assert request_data["model"] == "gpt-4o" @@ -8333,19 +8911,15 @@ class TestCheckKeyModelBudgetWithFallback: request = self._make_request() request.scope["path_params"] = {"model": "gpt-4o"} - with patch( - "litellm.proxy.auth.user_api_key_auth.can_key_call_model", - return_value=True, - ): - await _check_key_model_budget_with_fallback( - valid_token=valid_token, - model_max_budget_limiter=limiter, - model_name="gpt-4o", - request_data=request_data, - request=request, - llm_model_list=None, - llm_router=None, - ) + await _check_key_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + model_name="gpt-4o", + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=None, + ) assert request_data["model"] == "gpt-4o-mini" assert request.scope["path_params"]["model"] == "gpt-4o-mini" @@ -8379,6 +8953,1287 @@ class TestCheckKeyModelBudgetWithFallback: assert "model" not in request_data +class TestCheckTeamMemberModelBudgetWithFallback: + @pytest.fixture(autouse=True) + async def _seed_auth_cache(self, monkeypatch: pytest.MonkeyPatch) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, team_membership_reservation_cache_key + + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="user-1", team_id="team-1"), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) + await user_api_key_cache.async_set_cache( + key="team_id:team-1", + value=LiteLLM_TeamTableCachedObj(team_id="team-1"), + model_type=LiteLLM_TeamTableCachedObj, + ) + prisma_client: Final = MagicMock() + prisma_client.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + prisma_client.db.litellm_teammembership.find_first = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "user_api_key_cache", user_api_key_cache) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock()) + + def _router( + self, + fallbacks: list[dict[str, object]] | None, + model_group_alias: dict[str, str] | None = None, + free_models: frozenset[str] = frozenset(), + ) -> litellm.Router: + model_list: Final = [ + { + "model_name": model, + "litellm_params": {"model": "gpt-4o"}, + **( + {"model_info": {"input_cost_per_token": 0, "output_cost_per_token": 0}} + if model in free_models + else {} + ), + } + for model in ( + "gpt-4o", + "gpt-4o-mini", + "restricted-model", + "backup", + "backup-unavailable", + "backup-valid", + "claude-sonnet-4-6", + "claude-haiku-4-5", + "claude-haiku", + "expensive", + *sorted(free_models), + ) + ] + return litellm.Router(model_list=model_list, fallbacks=fallbacks, model_group_alias=model_group_alias) + + def _request(self, model: str) -> Request: + return Request(scope={"type": "http", "path_params": {"model": model}, "headers": []}) + + @pytest.mark.asyncio + async def test_within_budget_does_not_rewrite_request(self): + budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget=budget, + ) + limiter, _ = await _real_model_budget_limiter() + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + assert request_data["model"] == "gpt-4o" + assert request.scope["path_params"]["model"] == "gpt-4o" + assert "parsed_body" not in request.scope + + @pytest.mark.asyncio + async def test_generic_fallback_rewrites_every_request_view(self): + budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini"], + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={**budget, "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o", "auth_injected_metadata": {"internal": "value"}} + request: Final = self._request("gpt-4o") + client_body: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "hello"}]} + request.scope["parsed_body"] = (tuple(client_body), client_body) + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"*": ["gpt-4o-mini"]}]), + ) + + assert request_data["model"] == "gpt-4o-mini" + expected_body: Final = {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hello"}]} + assert request._json == expected_body + assert orjson.loads(request._body) == expected_body + assert request.scope["parsed_body"][1] == expected_body + assert "auth_injected_metadata" not in request._json + assert request.scope["path_params"]["model"] == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_specific_fallback_rewrites_requested_model(self): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + assert request_data["model"] == "gpt-4o-mini" + assert request.scope["path_params"]["model"] == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_multipart_request_keeps_member_budget_error_without_rewriting(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["claude-sonnet-4-6", "claude-haiku-4-5"], + user_id="user-1", + team_id="team-1", + router_settings={"fallbacks": [{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]}, + team_member_model_max_budget={"claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, redis_cache = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="user-1", team_id="team-1"), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) + uploaded_file: Final = UploadFile(filename="audio.wav", file=BytesIO(b"audio")) + request_data: Final = {"model": "claude-sonnet-4-6", "file": uploaded_file} + request: Final = Request( + scope={ + "type": "http", + "path_params": {"model": "claude-sonnet-4-6"}, + "headers": [(b"content-type", b"Multipart/Form-Data; boundary=test")], + } + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=litellm.Router( + model_list=[], + fallbacks=[{"claude-sonnet-4-6": ["claude-haiku-4-5"]}], + ), + user_api_key_cache=user_api_key_cache, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + assert request_data["file"] is uploaded_file + assert request.scope["path_params"]["model"] == "claude-sonnet-4-6" + assert getattr(request, "_body", None) is None + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6"]) + + @pytest.mark.asyncio + async def test_provider_scoped_resource_keeps_member_budget_error_without_fallback(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["claude-sonnet-4-6", "claude-haiku-4-5"], + user_id="user-1", + team_id="team-1", + router_settings={"fallbacks": [{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]}, + team_member_model_max_budget={"claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, redis_cache = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final[dict[str, object]] = { + "model": "claude-sonnet-4-6", + "input_file_id": "file-123", + } + request: Final = self._request("claude-sonnet-4-6") + user_api_key_cache: Final = UserApiKeyCache() + await user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="user-1", team_id="team-1"), + value=NO_TEAM_MEMBERSHIP_SENTINEL, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=litellm.Router( + model_list=[], + fallbacks=[{"claude-sonnet-4-6": ["claude-haiku-4-5"]}], + ), + user_api_key_cache=user_api_key_cache, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data == { + "model": "claude-sonnet-4-6", + "input_file_id": "file-123", + } + assert request.scope["path_params"]["model"] == "claude-sonnet-4-6" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6"]) + + @pytest.mark.asyncio + async def test_fallback_prefers_hierarchical_router_settings_fallbacks(self) -> None: + import litellm.proxy.proxy_server as proxy_server + + router_settings: Final = {"fallbacks": [{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]} + await proxy_server.user_api_key_cache.async_set_cache( + key="team_id:team-1", + value=LiteLLM_TeamTableCachedObj(team_id="team-1", router_settings=router_settings), + model_type=LiteLLM_TeamTableCachedObj, + ) + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "claude-haiku-4-5": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + router: Final = self._router([{"claude-sonnet-4-6": ["global-backup"]}]) + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=router, + ) + + assert request_data["model"] == "claude-haiku-4-5" + assert request.scope["path_params"]["model"] == "claude-haiku-4-5" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6", "claude-haiku-4-5"]) + + @pytest.mark.asyncio + async def test_fallback_resolves_router_alias_before_budget_check(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + router_settings={"model_group_alias": {"backup": "expensive"}}, + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "backup": {"max_budget": 100.0, "budget_duration": "1d"}, + "expensive": {"max_budget": 5.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={"claude-sonnet-4-6": 10.0, "expensive": 10.0} + ) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["backup"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6", "backup", "expensive"]) + + @pytest.mark.asyncio + async def test_fallback_checks_global_router_alias_budget_before_rerouting(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "backup": {"max_budget": 100.0, "budget_duration": "1d"}, + "expensive": {"max_budget": 5.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={"claude-sonnet-4-6": 10.0, "expensive": 10.0} + ) + request_data: Final[dict[str, object]] = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + router: Final = self._router( + [{"claude-sonnet-4-6": ["backup"]}], + model_group_alias={"backup": "expensive"}, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=router, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6", "backup", "expensive"]) + + @pytest.mark.asyncio + async def test_fallback_checks_key_team_alias_budget_before_rerouting(self) -> None: + model_group_alias: Final = {"backup": "expensive"} + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + router_settings={"model_group_alias": model_group_alias}, + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "backup": {"max_budget": 5.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={"claude-sonnet-4-6": 10.0, "backup": 10.0} + ) + request_data: Final[dict[str, object]] = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + router: Final = self._router([{"claude-sonnet-4-6": ["backup"]}]) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=router, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6", "backup"]) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "key_aliases, global_alias_map", + [ + pytest.param({"backup": "expensive"}, {}, id="key_aliases"), + pytest.param({}, {"backup": "expensive"}, id="global_alias_map"), + ], + ) + async def test_fallback_aliased_to_spent_member_model_budget_is_skipped( + self, + monkeypatch: pytest.MonkeyPatch, + key_aliases: dict[str, str], + global_alias_map: dict[str, str], + ) -> None: + primary_model: Final = "claude-sonnet-4-6" + fallback_model: Final = "backup" + expensive_model: Final = "expensive" + monkeypatch.setattr(litellm, "model_alias_map", global_alias_map) + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + aliases=key_aliases, + team_member_model_max_budget={ + primary_model: {"max_budget": 5.0, "budget_duration": "1d"}, + expensive_model: {"max_budget": 5.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={primary_model: 10.0, expensive_model: 10.0} + ) + request_data: Final[dict[str, object]] = {"model": primary_model} + request: Final = self._request(primary_model) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=[primary_model], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{primary_model: [fallback_model]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == primary_model + assert request.scope["path_params"]["model"] == primary_model + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys([primary_model, expensive_model]) + + @pytest.mark.asyncio + async def test_free_fallback_aliased_to_spent_paid_model_is_skipped(self) -> None: + primary_model: Final = "claude-sonnet-4-6" + fallback_model: Final = "free-backup" + expensive_model: Final = "expensive" + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + aliases={fallback_model: expensive_model}, + team_member_model_max_budget={ + primary_model: {"max_budget": 5.0, "budget_duration": "1d"}, + expensive_model: {"max_budget": 5.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={primary_model: 10.0, expensive_model: 10.0} + ) + request_data: Final[dict[str, object]] = {"model": primary_model} + request: Final = self._request(primary_model) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=[primary_model], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router( + [{primary_model: [fallback_model]}], + free_models=frozenset({fallback_model}), + ), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == primary_model + assert request.scope["path_params"]["model"] == primary_model + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys([primary_model, expensive_model]) + + @pytest.mark.asyncio + async def test_fallback_rewrites_router_alias_to_budget_checked_target(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + router_settings={"model_group_alias": {"backup": "expensive"}}, + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "backup": {"max_budget": 100.0, "budget_duration": "1d"}, + "expensive": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final[dict[str, object]] = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["backup"]}]), + ) + + assert request_data["model"] == "expensive" + assert request.scope["path_params"]["model"] == "expensive" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6", "backup", "expensive"]) + + @pytest.mark.asyncio + async def test_zero_cost_fallback_skips_spent_member_model_budget(self) -> None: + primary_model: Final = "claude-sonnet-4-6" + fallback_model: Final = "free-backup" + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + primary_model: {"max_budget": 5.0, "budget_duration": "1d"}, + fallback_model: {"max_budget": 5.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={primary_model: 10.0, fallback_model: 10.0} + ) + request_data: Final[dict[str, object]] = {"model": primary_model} + request: Final = self._request(primary_model) + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=[primary_model], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router( + [{primary_model: [fallback_model]}], + free_models=frozenset({fallback_model}), + ), + ) + + assert request_data["model"] == fallback_model + assert request.scope["path_params"]["model"] == fallback_model + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys([primary_model]) + + @pytest.mark.asyncio + async def test_fallback_settings_lookup_error_keeps_original_budget_error(self) -> None: + import litellm.proxy.proxy_server as proxy_server + + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={"claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, redis_cache = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + await proxy_server.user_api_key_cache.async_delete_cache(key="team_id:team-1") + proxy_server.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( + side_effect=RuntimeError("settings unavailable") + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router(None), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6"]) + proxy_server.prisma_client.db.litellm_teamtable.find_unique.assert_awaited_once() + + @pytest.mark.asyncio + async def test_fallback_skips_candidate_after_budget_check_error(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + router_settings={"model_group_alias": {}}, + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "backup-unavailable": {"max_budget": 100.0, "budget_duration": "1d"}, + "backup-valid": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={"claude-sonnet-4-6": 10.0}, + member_errors={"backup-unavailable": RuntimeError("budget counter unavailable")}, + ) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["backup-unavailable", "backup-valid"]}]), + ) + + assert request_data["model"] == "backup-valid" + assert request.scope["path_params"]["model"] == "backup-valid" + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["claude-sonnet-4-6", "backup-unavailable", "backup-valid"]) + + @pytest.mark.asyncio + async def test_fallback_candidate_check_errors_keep_original_budget_error(self) -> None: + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + router_settings={"model_group_alias": {}}, + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "backup-unavailable": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter( + member_spend={"claude-sonnet-4-6": 10.0}, + member_errors={"backup-unavailable": RuntimeError("budget counter unavailable")}, + ) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["backup-unavailable"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + + @pytest.mark.asyncio + async def test_fallback_outside_project_models_raises_original_member_budget_error( + self, + ) -> None: + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "claude-haiku-4-5": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + project_object: Final = LiteLLM_ProjectTableCachedObj( + project_id="project-1", + models=["claude-sonnet-4-6"], + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]), + project_object=project_object, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + + @pytest.mark.asyncio + async def test_fallback_within_project_models_rewrites_request( + self, + ) -> None: + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "claude-haiku-4-5": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + project_object: Final = LiteLLM_ProjectTableCachedObj( + project_id="project-1", + models=["claude-sonnet-4-6", "claude-haiku-4-5"], + ) + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]), + project_object=project_object, + ) + + assert request_data["model"] == "claude-haiku-4-5" + assert request.scope["path_params"]["model"] == "claude-haiku-4-5" + + @pytest.mark.asyncio + @pytest.mark.parametrize("grant_source", ["team", "key"]) + async def test_fallback_uses_team_and_key_access_group_grants( + self, grant_source: str, monkeypatch: pytest.MonkeyPatch + ) -> None: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy._types import LiteLLM_AccessGroupTable + + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_models=["claude-sonnet-4-6"], + access_group_ids=["key-group"] if grant_source == "key" else [], + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "claude-haiku-4-5": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + team_object: Final = LiteLLM_TeamTable( + team_id="team-1", + models=["claude-sonnet-4-6"], + access_group_ids=["team-group"] if grant_source == "team" else [], + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + prisma_client: Final = MagicMock() + access_group_id: Final = "team-group" if grant_source == "team" else "key-group" + group_row: Final = LiteLLM_AccessGroupTable( + access_group_id=access_group_id, + access_group_name=access_group_id, + access_model_names=["claude-haiku-4-5"], + assigned_team_ids=["team-1"], + ) + prisma_client.db.litellm_accessgrouptable.find_unique = AsyncMock(return_value=group_row) + monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]), + team_object=team_object, + prisma_client=prisma_client, + user_api_key_cache=proxy_server.user_api_key_cache, + proxy_logging_obj=proxy_server.proxy_logging_obj, + ) + + assert request_data["model"] == "claude-haiku-4-5" + prisma_client.db.litellm_accessgrouptable.find_unique.assert_awaited_once_with( + where={"access_group_id": access_group_id} + ) + + @pytest.mark.asyncio + async def test_fallback_without_team_or_key_access_group_grants_keeps_budget_error(self) -> None: + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_models=["claude-sonnet-4-6"], + team_member_model_max_budget={ + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "1d"}, + "claude-haiku-4-5": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + team_object: Final = LiteLLM_TeamTable(team_id="team-1", models=["claude-sonnet-4-6"]) + limiter, _ = await _real_model_budget_limiter(member_spend={"claude-sonnet-4-6": 10.0}) + request_data: Final = {"model": "claude-sonnet-4-6"} + request: Final = self._request("claude-sonnet-4-6") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["claude-sonnet-4-6"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"claude-sonnet-4-6": ["claude-haiku-4-5"]}]), + team_object=team_object, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "claude-sonnet-4-6" + + @pytest.mark.asyncio + async def test_exhausted_fallback_raises_original_member_budget_error(self): + budget: Final = { + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 5.0, "budget_duration": "1d"}, + } + valid_token: Final = UserAPIKeyAuth( + token="test-key", + user_id="user-1", + team_id="team-1", + team_member_model_max_budget=budget, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0, "gpt-4o-mini": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + + @pytest.mark.asyncio + async def test_no_router_fallback_raises_original_member_budget_error(self): + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=None, + ) + + _assert_budget_exceeded_values(exc_info.value) + + @pytest.mark.asyncio + async def test_router_without_fallbacks_raises_original_member_budget_error(self): + budget: Final = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_member_model_max_budget=budget, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router(None), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + + @pytest.mark.asyncio + async def test_fallback_exceeding_key_budget_is_rejected(self): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + model_max_budget={"gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, redis_cache = await _real_model_budget_limiter( + member_spend={"gpt-4o": 10.0}, key_spend={"gpt-4o-mini": 2.0} + ) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert [ + call.kwargs["key"] + for call in redis_cache.async_get_cache.call_args_list + if call.kwargs["key"].startswith("virtual_key_spend:") + ] == _model_budget_spend_keys(Litellm_EntityType.KEY, "test-key", ["gpt-4o-mini"]) + + @pytest.mark.asyncio + async def test_fallback_outside_key_models_is_rejected(self): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o"], + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + denied: Final = ProxyException( + message="model not allowed", + type=ProxyErrorTypes.budget_exceeded, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + + @pytest.mark.asyncio + async def test_agent_key_does_not_reroute_to_fallback(self): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini"], + user_id="user-1", + team_id="team-1", + agent_id="agent-1", + team_member_model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + assert getattr(request, "_json", None) is None + + @pytest.mark.asyncio + async def test_agent_key_can_reroute_on_key_model_budget_fallback(self): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini"], + agent_id="agent-1", + budget_fallbacks={"gpt-4o": ["gpt-4o-mini"]}, + model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(key_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + await _check_key_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + model_name="gpt-4o", + request_data=request_data, + request=request, + ) + + assert request_data["model"] == "gpt-4o-mini" + + @pytest.mark.asyncio + async def test_fallback_outside_member_allowed_models_is_rejected(self): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini"], + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + import litellm.proxy.proxy_server as proxy_server + + membership: Final = LiteLLM_TeamMembership( + user_id="user-1", + team_id="team-1", + budget_id="member-budget", + litellm_budget_table=LiteLLM_BudgetTable(budget_id="member-budget", allowed_models=["gpt-4o"]), + ) + await proxy_server.user_api_key_cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id="user-1", team_id="team-1"), + value=membership, + model_type=LiteLLM_TeamMembership, + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + ) + + _assert_budget_exceeded_values(exc_info.value) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("targets", "expected_model"), + [ + ([{"model": "gpt-4o-mini"}], "gpt-4o-mini"), + ([{"model": "gpt-4o-mini", "temperature": 0.2}, {"model": "claude-haiku"}], "claude-haiku"), + ([{"model": "gpt-4o-mini", "temperature": 0.2}], None), + ], + ) + async def test_dict_fallback_targets_only_reroute_model_only_entries(self, targets, expected_model): + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini", "claude-haiku"], + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + "claude-haiku": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + + if expected_model is not None: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": targets}]), + ) + assert request_data["model"] == expected_model + else: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": targets}]), + ) + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + + @pytest.mark.asyncio + async def test_fallback_skips_candidate_denied_to_end_user(self): + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=self._router([{"gpt-4o": ["gpt-4o-mini"]}]), + end_user_object=SimpleNamespace(models=["gpt-4o"]), + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + + @pytest.mark.asyncio + async def test_fallback_skips_candidates_without_pricing(self, monkeypatch: pytest.MonkeyPatch): + valid_token: Final = UserAPIKeyAuth( + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, + ) + limiter, redis_cache = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + {"model_name": "unpriced-model", "litellm_params": {"model": "openai/unpriced-model"}}, + ], + fallbacks=[{"gpt-4o": ["unpriced-model"]}], + ) + monkeypatch.setattr(litellm, "block_requests_for_models_without_pricing", True) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=router, + route="/v1/chat/completions", + ) + + _assert_budget_exceeded_values(exc_info.value) + assert [ + call.kwargs["key"] for call in redis_cache.async_get_cache.call_args_list + ] == _team_member_model_budget_spend_keys(["gpt-4o"]) + + @pytest.mark.asyncio + async def test_fallback_checks_candidate_access_group_budget_and_stamps_candidate_groups(self): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.user_api_key_cache import model_access_group_cache_key + from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini", "fallback-group"], + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + await proxy_server.user_api_key_cache.async_set_cache( + key=model_access_group_cache_key("fallback-group"), + value=ModelAccessGroupBudget(access_group_name="fallback-group", spend=0.0, max_budget=10.0), + model_type=ModelAccessGroupBudget, + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "gpt-4o-mini"}, + "model_info": {"access_groups": ["fallback-group"]}, + }, + ], + fallbacks=[{"gpt-4o": ["gpt-4o-mini"]}], + ) + + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=router, + ) + + assert request_data["model"] == "gpt-4o-mini" + assert request.scope["path_params"]["model"] == "gpt-4o-mini" + assert valid_token.matched_model_access_groups == ["fallback-group"] + + @pytest.mark.asyncio + async def test_access_group_budget_exhaustion_skips_fallback_candidate(self): + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.common_utils.user_api_key_cache import model_access_group_cache_key + from litellm.types.proxy.model_access_group_budget import ModelAccessGroupBudget + + valid_token: Final = UserAPIKeyAuth( + token="test-key", + models=["gpt-4o", "gpt-4o-mini", "fallback-group"], + user_id="user-1", + team_id="team-1", + team_member_model_max_budget={ + "gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}, + "gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ) + limiter, _ = await _real_model_budget_limiter(member_spend={"gpt-4o": 10.0}) + request_data: Final = {"model": "gpt-4o"} + request: Final = self._request("gpt-4o") + await proxy_server.user_api_key_cache.async_set_cache( + key=model_access_group_cache_key("fallback-group"), + value=ModelAccessGroupBudget(access_group_name="fallback-group", spend=5.0, max_budget=5.0), + model_type=ModelAccessGroupBudget, + ) + router: Final = litellm.Router( + model_list=[ + {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "gpt-4o-mini"}, + "model_info": {"access_groups": ["fallback-group"]}, + }, + ], + fallbacks=[{"gpt-4o": ["gpt-4o-mini"]}], + ) + + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _check_team_member_model_budget_with_fallback( + valid_token=valid_token, + model_max_budget_limiter=limiter, + models=["gpt-4o"], + request_data=request_data, + request=request, + llm_model_list=None, + llm_router=router, + ) + + _assert_budget_exceeded_values(exc_info.value) + assert request_data["model"] == "gpt-4o" + + @pytest.mark.asyncio async def test_global_proxy_spend_reads_resettable_proxy_budget_row(): """Regression for the global proxy budget ignoring budget_duration @@ -10185,9 +12040,7 @@ async def test_managed_agent_cannot_bypass_grants_with_server_default( monkeypatch.setattr(proxy_server, name, value) data = {"messages": [{"role": "user", "content": "hi"}], **({"model": requested} if requested else {})} auth = UserAPIKeyAuth(agent_id="managed") - auth.managed_agent_context = ManagedAgentContext( - agent_id="managed", binding_revision="revision", mode="autonomous" - ) + auth.managed_agent_context = ManagedAgentContext(agent_id="managed", binding_revision="revision", mode="autonomous") if not grant_default: with pytest.raises(ProxyException) as denied: await _authorize_authenticated_request(auth, _alias_request(route, data), data, route, "persisted-key") @@ -10217,12 +12070,21 @@ async def test_managed_jwt_cannot_be_downgraded_into_virtual_key_mapping(monkeyp from litellm.types.proxy.agent_identity import AgentIdentityBinding binding: Final = AgentIdentityBinding( - agent_id="managed", provider="microsoft_entra", issuer="issuer", tenant_id="tenant", - client_id="client", service_principal_id="principal", revision="current", + agent_id="managed", + provider="microsoft_entra", + issuer="issuer", + tenant_id="tenant", + client_id="client", + service_principal_id="principal", + revision="current", ) agent: Final = AgentResponse( - agent_id="managed", agent_name="Managed", agent_card_params={}, - identity_managed=True, identity=binding, execution_mode="autonomous", + agent_id="managed", + agent_name="Managed", + agent_card_params={}, + identity_managed=True, + identity=binding, + execution_mode="autonomous", ) client: Final = MagicMock() client.writer_db.litellm_agentidentity.find_unique = AsyncMock(return_value=binding) @@ -10230,22 +12092,35 @@ async def test_managed_jwt_cannot_be_downgraded_into_virtual_key_mapping(monkeyp handler: Final = MagicMock() handler.is_jwt.return_value = True handler.litellm_jwtauth = LiteLLM_JWTAuth(virtual_key_claim_field="sub") - handler.auth_jwt = AsyncMock(return_value={ - "iss": "issuer", "tid": "tenant", "azp": "client", "oid": "principal", "sub": "mapped-key", - }) + handler.auth_jwt = AsyncMock( + return_value={ + "iss": "issuer", + "tid": "tenant", + "azp": "client", + "oid": "principal", + "sub": "mapped-key", + } + ) for name, value in { **_proxy_attrs_for_centralized_checks(), - "general_settings": {"enable_jwt_auth": True}, "premium_user": True, - "prisma_client": client, "jwt_handler": handler, "user_api_key_cache": UserApiKeyCache(), + "general_settings": {"enable_jwt_auth": True}, + "premium_user": True, + "prisma_client": client, + "jwt_handler": handler, + "user_api_key_cache": UserApiKeyCache(), "proxy_logging_obj": MagicMock(post_call_failure_hook=AsyncMock(return_value=None)), }.items(): monkeypatch.setattr(proxy_server, name, value) for _ in range(2): with pytest.raises(ProxyException) as failure: await user_api_key_auth_builder( - request=_alias_request("/v1/chat/completions", {}), api_key="Bearer verified.jwt.token", - azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, - azure_apim_header=None, request_data={}, + request=_alias_request("/v1/chat/completions", {}), + api_key="Bearer verified.jwt.token", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, ) assert failure.value.code == "403" assert "without virtual-key mapping" in failure.value.message @@ -10262,9 +12137,17 @@ async def test_virtual_key_cannot_enter_checks_as_an_identity_managed_actor(monk from litellm.types.proxy.agent_identity import AgentIdentityBinding target: Final = AgentResponse( - agent_id="bound", agent_name="Bound", agent_card_params={}, identity_managed=True, + agent_id="bound", + agent_name="Bound", + agent_card_params={}, + identity_managed=True, identity=AgentIdentityBinding( - agent_id="bound", provider="microsoft_entra", tenant_id="tenant", client_id="client", issuer="issuer", revision="current" + agent_id="bound", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="current", ), ) client: Final = MagicMock() @@ -10272,7 +12155,9 @@ async def test_virtual_key_cannot_enter_checks_as_an_identity_managed_actor(monk monkeypatch.setattr(proxy_server, "prisma_client", client) checks: Final = AsyncMock() monkeypatch.setattr(auth_module, "run_centralized_common_checks", checks) - monkeypatch.setattr(proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None))) + monkeypatch.setattr( + proxy_server, "proxy_logging_obj", MagicMock(post_call_failure_hook=AsyncMock(return_value=None)) + ) data: Final = {"model": "allowed", "messages": [{"role": "user", "content": "hello"}]} request: Final = _alias_request("/v1/chat/completions", data) with pytest.raises(ProxyException): @@ -10300,16 +12185,24 @@ async def test_custom_auth_grants_reach_managed_targets_without_a_virtual_key_ro from litellm.types.proxy.agent_identity import AgentIdentityBinding target: Final = AgentResponse( - agent_id="target", agent_name="Target", agent_card_params={}, identity_managed=True, + agent_id="target", + agent_name="Target", + agent_card_params={}, + identity_managed=True, identity=AgentIdentityBinding( - agent_id="target", provider="microsoft_entra", tenant_id="tenant", client_id="client", - issuer="issuer", revision="current", + agent_id="target", + provider="microsoft_entra", + tenant_id="tenant", + client_id="client", + issuer="issuer", + revision="current", ), ) registry: Final = AgentRegistry() registry.register_agent(target) trusted: Final = UserAPIKeyAuth( - api_key=credential, object_permission={"object_permission_id": "custom", "agents": ["target"] if granted else ["other"]} + api_key=credential, + object_permission={"object_permission_id": "custom", "agents": ["target"] if granted else ["other"]}, ) custom: Final = AsyncMock(return_value=trusted) database: Final = MagicMock() @@ -10325,9 +12218,13 @@ async def test_custom_auth_grants_reach_managed_targets_without_a_virtual_key_ro monkeypatch.setattr(agent_registry, "global_agent_registry", registry) monkeypatch.setattr(litellm, "enable_post_custom_auth_checks", False, raising=False) admitted: Final = await user_api_key_auth_builder( - request=_alias_request("/a2a/target/message/send", {}), api_key=f"Bearer {credential}", - azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, - azure_apim_header=None, request_data={}, + request=_alias_request("/a2a/target/message/send", {}), + api_key=f"Bearer {credential}", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, ) assert await AgentRequestHandler.is_agent_allowed("target", admitted) is granted custom.assert_awaited_once() @@ -10347,9 +12244,13 @@ async def test_enterprise_custom_auth_key_return_stays_a_proxy_validated_key(mon module: Final = importlib.import_module("litellm.proxy.auth.user_api_key_auth") monkeypatch.setattr(module, "enterprise_custom_auth", custom) admitted: Final = await user_api_key_auth_builder( - request=_alias_request("/v1/chat/completions", {}), api_key="Bearer external-credential", - azure_api_key_header="", anthropic_api_key_header=None, google_ai_studio_api_key_header=None, - azure_apim_header=None, request_data={}, + request=_alias_request("/v1/chat/completions", {}), + api_key="Bearer external-credential", + azure_api_key_header="", + anthropic_api_key_header=None, + google_ai_studio_api_key_header=None, + azure_apim_header=None, + request_data={}, ) assert admitted.authenticated_by_custom_auth is False assert admitted.via_virtual_key is True diff --git a/tests/unit/proxy/common_utils/test_upsert_budget_membership.py b/tests/unit/proxy/common_utils/test_upsert_budget_membership.py index 74d1748d4de..1f7154f0a74 100644 --- a/tests/unit/proxy/common_utils/test_upsert_budget_membership.py +++ b/tests/unit/proxy/common_utils/test_upsert_budget_membership.py @@ -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() diff --git a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py index 8de34e6391e..2678f4205d5 100644 --- a/tests/unit/proxy/hooks/test_model_max_budget_limiter.py +++ b/tests/unit/proxy/hooks/test_model_max_budget_limiter.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_teams.py b/tests/unit/proxy/management_endpoints/management_v1/test_teams.py index 3f60384cfa7..5d1c658b143 100644 --- a/tests/unit/proxy/management_endpoints/management_v1/test_teams.py +++ b/tests/unit/proxy/management_endpoints/management_v1/test_teams.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py index 2f3be61d00f..24129488f21 100644 --- a/tests/unit/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_budget_endpoints.py @@ -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 diff --git a/tests/unit/proxy/management_endpoints/test_common_utils.py b/tests/unit/proxy/management_endpoints/test_common_utils.py index d783f4c64f6..21d017addca 100644 --- a/tests/unit/proxy/management_endpoints/test_common_utils.py +++ b/tests/unit/proxy/management_endpoints/test_common_utils.py @@ -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.""" diff --git a/tests/unit/proxy/management_endpoints/test_team_endpoints.py b/tests/unit/proxy/management_endpoints/test_team_endpoints.py index b81c1c9f3d5..e921e4df5c1 100644 --- a/tests/unit/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_team_endpoints.py @@ -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") diff --git a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py index b450a4dfd02..644f0656c7a 100644 --- a/tests/unit/proxy/management_helpers/test_management_helpers_utils.py +++ b/tests/unit/proxy/management_helpers/test_management_helpers_utils.py @@ -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"]) diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index fb8d98ef61e..bd66bc24837 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -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", diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index b040cc861bf..0f66674daf9 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -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): diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index cf04b9ff91e..3a6aac9be05 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -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 diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index e0829e42557..3694ecea4ed 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -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" diff --git a/ui/litellm-dashboard/src/components/Teams.integration.test.tsx b/ui/litellm-dashboard/src/components/Teams.integration.test.tsx new file mode 100644 index 00000000000..a6801e3c996 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Teams.integration.test.tsx @@ -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(); + + 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; + 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" }, + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 0b773fd0f04..9971010c3ef 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -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 = ({ accessToken, userID, userRole, isViewOnly const [loggingSettings, setLoggingSettings] = useState([]); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); const [modelMaxBudget, setModelMaxBudget] = useState({}); + const [teamMemberModelMaxBudget, setTeamMemberModelMaxBudget] = useState({}); const [routerSettings, setRouterSettings] = useState(null); const [routerSettingsKey, setRouterSettingsKey] = useState(0); @@ -358,6 +360,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, isViewOnly setLoggingSettings([]); setModelAliases({}); setModelMaxBudget({}); + setTeamMemberModelMaxBudget({}); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); }; @@ -551,9 +554,11 @@ const Teams: React.FC = ({ 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 = ({ accessToken, userID, userRole, isViewOnly /> )} + (() => 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 ( - Per-Model Budgets + {label} diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 54b2066c9c5..0fbf3d36355 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -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", diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx index 45800a18ff8..2c98d17f961 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.integration.test.tsx @@ -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 ( + <> + + setVisible(false)} + onSubmit={vi.fn()} + mode="edit" + config={modelBudgetConfig as never} + initialData={initialData as never} + /> + + ); + }; + + const { rerender } = renderWithProviders(); + + 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(); + expect(screen.queryByPlaceholderText("Select model")).not.toBeInTheDocument(); +}); diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index 4aaeb42da98..4e1fa4cfbcf 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -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 => { const shape = { @@ -171,6 +177,16 @@ const MemberModal = ({ onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)} /> ); + case "model-max-budget": + return ( + + ); case "utc-datetime": return ( { 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(); + 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(); + 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(); + 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(); + 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(); + }); }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index a782112edb2..7663ef8c588 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -286,7 +286,7 @@ export interface TeamMembership { max_parallel_requests: number | null; tpm_limit: number | null; rpm_limit: number | null; - model_max_budget: Record | 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 = ({ const [isTeamSaving, setIsTeamSaving] = useState(false); const [teamModelAliases, setTeamModelAliases] = useState>({}); const [teamModelMaxBudget, setTeamModelMaxBudget] = useState({}); + const [teamMemberModelMaxBudget, setTeamMemberModelMaxBudget] = useState({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, isViewOnly } = useAuthorized(); @@ -661,12 +663,20 @@ const TeamInfoView: React.FC = ({ ] 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 = ({ 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 = ({ 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 = ({ 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 = ({ /> )} + = ({

Max Budget: {info.team_member_budget_table?.max_budget ?? "No Limit"}
Budget Duration: {info.team_member_budget_table?.budget_duration || "No Limit"}
+ {modelMaxBudgetToEntries( + info.team_member_budget_table?.model_max_budget as ModelMaxBudget | null | undefined, + ).map(({ model, budgetLimit, timePeriod }) => ( +
+ Default Per-Model Budget ({model}):{" "} + {budgetLimit !== null ? `$${formatNumberWithCommas(budgetLimit, 4)}` : "No Limit"} per {timePeriod} +
+ ))}
Key Duration: {info.metadata?.team_member_key_duration || "No Limit"}
TPM Limit: {info.team_member_budget_table?.tpm_limit ?? "No Limit"}
RPM Limit: {info.team_member_budget_table?.rpm_limit ?? "No Limit"}
@@ -2554,6 +2592,20 @@ const TeamInfoView: React.FC = ({ ), type: "utc-datetime" as const, }, + { + name: "model_max_budget", + label: ( + + Per-Model Budgets{" "} + + + + + ), + type: "model-max-budget" as const, + availableModels: availableModelBudgetModels, + premiumUser, + }, { name: "tpm_limit", label: ( diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx index 8953a06796f..b8f292222ed 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.test.tsx @@ -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, diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 6778b40e93c..f50c47dd12f 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -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, 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} diff --git a/ui/litellm-dashboard/src/components/team/memberFormValues.ts b/ui/litellm-dashboard/src/components/team/memberFormValues.ts index 77cf2dd3e43..df80e2dd9ac 100644 --- a/ui/litellm-dashboard/src/components/team/memberFormValues.ts +++ b/ui/litellm-dashboard/src/components/team/memberFormValues.ts @@ -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; -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": diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index f36456628e5..731a92d8b98 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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 */