diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index fb6a1c40253..ecaf8f2e7e1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -78,6 +78,7 @@ _PROPAGATED_METADATA_KEYS: Final = ( "user_api_key_end_user_id", "user_api_end_user_max_budget", "user_api_key_model_max_budget", + "user_api_key_team_model_max_budget", "user_api_key_user_model_max_budget", "user_api_key_end_user_model_max_budget", "litellm_call_id", @@ -395,9 +396,9 @@ async def _check_summary_model_budget( ``user_api_key_auth`` runs for the client-requested model. Returns True outside the proxy or when no per-model budget is configured. - All three scopes are checked because the summary's spend is charged to all - three: this file propagates the key, user and end-user budgets into the - subrequest's metadata, so enforcing only two of them would let compaction + 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. """ if user_api_key_auth is None: @@ -444,6 +445,26 @@ async def _check_summary_model_budget( ) return False + team_model_max_budget: Final = user_api_key_auth.team_model_max_budget + team_id: Final = user_api_key_auth.team_id + if isinstance(team_model_max_budget, dict) and team_model_max_budget and team_id is not None: + try: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=team_id, + team_model_max_budget=team_model_max_budget, + key_model_max_budget=model_max_budget if isinstance(model_max_budget, dict) else None, + model=summary_model, + ) + except litellm.BudgetExceededError: + return False + except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the other scopes do + verbose_logger.warning( + "compact_20260112: unexpected error during team model-budget check for summary_model=%s; denying: %s", + summary_model, + e, + ) + return False + end_user_model_max_budget: Final[dict[str, object] | None] = getattr( user_api_key_auth, "end_user_model_max_budget", None ) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ad55fa5d2be..72852c7c20c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2002,6 +2002,13 @@ RouterSettingsDict = Annotated[ class NewTeamRequest(TeamBase): router_settings: RouterSettingsDict | None = None model_aliases: dict | None = None + model_max_budget: GenericBudgetConfigType | None = Field( + default=None, + description=( + "Max budget per model for every key on the team, overridable per key " + "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})" + ), + ) tags: list | None = None guardrails: list[str] | None = None policies: list[str] | None = None @@ -2103,6 +2110,13 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): access_group_ids: list[str] | None = None budget_limits: list[BudgetLimitEntry] | None = None # multiple concurrent budget windows default_team_member_models: list[str] | None = None # default allowed_models seeded onto new team members + model_max_budget: GenericBudgetConfigType | None = Field( + default=None, + description=( + "Max budget per model for every key on the team, overridable per key " + "(e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}})" + ), + ) class PatchTeamRequest(UpdateTeamRequest): @@ -3030,6 +3044,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_tpd_limit: int | None = None team_max_budget: float | None = None team_soft_budget: float | None = None + team_model_max_budget: dict[str, object] | None = None team_models: list = [] team_blocked: bool = False soft_budget: float | None = None @@ -4444,6 +4459,7 @@ class TeamInfoResponseObjectTeamTable(LiteLLM_TeamTable): # Parent org's model ceiling, reported only to callers who can manage the team. # None = no org or not a manager; [] or ["all-proxy-models"] = no ceiling. organization_models: list[str] | None = None + model_max_budget_usage: Mapping[str, Mapping[str, object]] | None = None class TeamInfoResponseObject(TypedDict): diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 0421659c331..2029ee342ae 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -59,6 +59,7 @@ class TeamGrants(TypedDict, total=False): team_tpd_limit: ReadOnly[int | None] team_max_budget: ReadOnly[float | None] team_soft_budget: ReadOnly[float | None] + team_model_max_budget: ReadOnly[dict[str, object] | None] team_spend: ReadOnly[float | None] team_models: ReadOnly[Sequence[str]] team_blocked: ReadOnly[bool] @@ -101,6 +102,7 @@ def team_grants( team_tpd_limit=team_object.tpd_limit, team_max_budget=team_object.max_budget, team_soft_budget=team_object.soft_budget, + team_model_max_budget=team_object.model_max_budget, team_spend=team_object.spend, team_models=tuple(team_object.models), team_blocked=team_object.blocked, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 5958a68f975..0a5aa9ea793 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -264,6 +264,16 @@ class _UserModelBudgetLimiter(Protocol): ) -> bool: ... +class _TeamModelBudgetLimiter(Protocol): + 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: ... + + class _TokenTeamModels(Protocol): @property def team_models(self) -> list[str]: ... @@ -334,6 +344,25 @@ async def _check_user_model_budget( ) +async def _check_team_model_budget( + valid_token: UserAPIKeyAuth, + model_max_budget_limiter: _TeamModelBudgetLimiter, + models: list[str], +) -> None: + """Enforce the team's `model_max_budget` for every requested model the key does not override.""" + team_model_max_budget: Final = valid_token.team_model_max_budget + if valid_token.team_id is None or not team_model_max_budget: + return + key_model_max_budget: Final[Mapping[str, object] | None] = valid_token.model_max_budget + for model_name in models: + 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_name, + ) + + async def _check_key_model_budget_with_fallback( valid_token: UserAPIKeyAuth, model_max_budget_limiter: _KeyModelBudgetLimiter, @@ -2369,6 +2398,7 @@ async def _user_api_key_auth_builder( team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, @@ -2523,6 +2553,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached team_id=valid_token.team_id, max_budget=valid_token.team_max_budget, soft_budget=valid_token.team_soft_budget, + model_max_budget=valid_token.team_model_max_budget, spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, @@ -2592,6 +2623,7 @@ async def _run_centralized_common_checks( litellm_proxy_admin_name, llm_router, master_key, + model_max_budget_limiter, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2864,6 +2896,21 @@ async def _run_centralized_common_checks( finally: release_spend_counter_batch() + if not skip_budget_checks: + 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, + ) + ), + ) + await _reserve_budget_after_common_checks( user_api_key_auth_obj=user_api_key_auth_obj, request=request, diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index d3f3de730ab..f7131091c0b 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -78,6 +78,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None: v.*, t.spend AS team_spend, t.max_budget AS team_max_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit, diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index efaaab277a9..bbfc7325f40 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -19,12 +19,14 @@ from litellm.types.utils import BudgetConfig, StandardLoggingPayload 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" _SPEND_CACHE_KEY_PREFIXES: Final = MappingProxyType( { Litellm_EntityType.KEY: VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, 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, } ) @@ -37,6 +39,7 @@ _BUDGET_START_TIME_KEY_PREFIXES: Final = MappingProxyType( Litellm_EntityType.KEY: "virtual_key_budget_start_time", 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", } ) @@ -139,6 +142,18 @@ def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> return None +def team_model_budget_applies(model: str, key_model_max_budget: Mapping[str, object] | None) -> bool: + """A key entry that spend-gates `model` overrides the team cap: it is then gated on and billed to the key alone.""" + if not key_model_max_budget: + return True + resolved: Final = resolve_model_budget(model=model, model_max_budget=key_model_max_budget) + return resolved is None or not _spend_gated(resolved.budget_config) + + +def _spend_gated(budget_config: BudgetConfig) -> bool: + return budget_config.max_budget is not None and budget_config.max_budget >= 0 + + def _budget_model_candidates(model: str) -> tuple[str, ...]: """Names a budget may be configured under for a request on `model`, most specific first. @@ -346,6 +361,30 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): exceeded_message=f"LiteLLM End User: {end_user_id}, exceeded budget for model={model}", ) + 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: + """ + Check if the team is within the model budget, unless the key's own + `model_max_budget` overrides it for `model` + + Raises: + BudgetExceededError: If the team has exceeded the model budget + """ + if not team_model_budget_applies(model=model, key_model_max_budget=key_model_max_budget): + return True + return await self._is_entity_within_model_budget( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + model_max_budget=team_model_max_budget, + model=model, + exceeded_message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", + ) + async def _is_entity_within_model_budget( self, entity_type: Litellm_EntityType, @@ -456,11 +495,26 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return response_cost: Final[float] = standard_logging_payload.get("response_cost", 0) + key_model_max_budget: Final = _metadata.get("user_api_key_model_max_budget") entity_budgets: Final = ( ( Litellm_EntityType.KEY, payload_metadata.get("user_api_key_hash"), - _metadata.get("user_api_key_model_max_budget"), + key_model_max_budget, + ), + ( + Litellm_EntityType.TEAM, + payload_metadata.get("user_api_key_team_id"), + ( + _metadata.get("user_api_key_team_model_max_budget") + if team_model_budget_applies( + model=model, + key_model_max_budget=( + key_model_max_budget if isinstance(key_model_max_budget, Mapping) else None + ), + ) + else None + ), ), ( Litellm_EntityType.USER, @@ -478,7 +532,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): if not resolved_budgets: verbose_proxy_logger.debug( "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event: " - "no key, user or end-user model_max_budget covers model=%s", + "no key, team, user or end-user model_max_budget covers model=%s", model, ) return diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 563db811edc..93374eb099b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2287,6 +2287,7 @@ async def add_litellm_data_to_request( # Team spend, budget - used by prometheus.py data[_metadata_variable_name]["user_api_key_team_max_budget"] = user_api_key_dict.team_max_budget data[_metadata_variable_name]["user_api_key_team_spend"] = user_api_key_dict.team_spend + data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget data[_metadata_variable_name]["user_api_key_request_route"] = user_api_key_dict.request_route # API Key spend, budget - used by prometheus.py diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 98155ad6839..973311608ed 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -55,6 +55,7 @@ def validate_budget_duration(budget_duration: str | None, status_code: int = 400 from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import ( + CommonProxyErrors, KeyRequestBase, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, @@ -73,12 +74,62 @@ from litellm.proxy._types import ( # noqa: F401 re-exported from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.utils import _premium_user_check from litellm.repositories.team_repository import TeamRepository +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from litellm.proxy._types import NewProjectRequest, UpdateProjectRequest from litellm.proxy.utils import PrismaClient, ProxyLogging +def validate_team_model_max_budget( + model_max_budget: Mapping[str, BudgetConfig] | None, + premium_user: bool, +) -> None: + """Reject a team `model_max_budget` the limiter could not enforce (no duration, bad cap, tpm/rpm limits).""" + if not model_max_budget: + return + if premium_user is not True: + raise HTTPException( + status_code=403, + detail={ + "error": f"Setting model_max_budget on a team is an enterprise feature. {CommonProxyErrors.not_premium_user.value}" + }, + ) + for model_name, budget_config in model_max_budget.items(): + if not model_name.strip(): + raise HTTPException( + status_code=400, + detail={"error": "model_max_budget keys must be non-empty model names"}, + ) + max_budget = budget_config.max_budget + if max_budget is None or not math.isfinite(max_budget) or max_budget < 0: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"model_max_budget[{model_name!r}].max_budget must be a non-negative finite number. " + f"Received: {max_budget}" + ) + }, + ) + if budget_config.budget_duration is None: + raise HTTPException( + status_code=400, + detail={"error": f"model_max_budget[{model_name!r}] requires a budget_duration, e.g. '1d' or '30d'"}, + ) + validate_budget_duration(budget_config.budget_duration) + if budget_config.tpm_limit is not None or budget_config.rpm_limit is not None: + raise HTTPException( + status_code=400, + detail={ + "error": ( + f"model_max_budget[{model_name!r}] tpm_limit/rpm_limit are not enforced on a team; " + "set per-model rate limits on the key instead" + ) + }, + ) + + def require_caller_user_id_for_non_admin( user_api_key_dict: UserAPIKeyAuth, ) -> str: diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e719d6d761a..6fb1ef5ec93 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -22,7 +22,7 @@ from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protoc import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status -from pydantic import BaseModel, JsonValue +from pydantic import BaseModel, JsonValue, ValidationError from typing_extensions import ReadOnly, TypedDict import litellm @@ -38,6 +38,7 @@ from litellm.proxy._types import ( DeleteTeamRequest, LiteLLM_AuditLogs, LiteLLM_DeletedTeamTable, + Litellm_EntityType, LiteLLM_ManagementEndpoint_MetadataFields, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_ModelTable, @@ -95,6 +96,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.hooks.model_max_budget_limiter import build_model_max_budget_usage from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity_aggregated, ) @@ -108,6 +110,7 @@ from litellm.proxy.management_endpoints.common_utils import ( _upsert_budget_and_membership, _user_has_admin_view, validate_budget_duration, + validate_team_model_max_budget, ) from litellm.proxy.management_endpoints.organization_endpoints import ( add_member_to_organization, @@ -177,6 +180,7 @@ from litellm.types.proxy.management_endpoints.team_endpoints import ( TeamUserSpendRow, UpdateTeamMemberPermissionsRequest, ) +from litellm.types.utils import BudgetConfig if TYPE_CHECKING: from prisma import Prisma @@ -1170,6 +1174,56 @@ def _check_team_budget_update_authority( ) +def _existing_model_cap(raw_budget_config: object) -> BudgetConfig | None: + try: + return BudgetConfig.model_validate(raw_budget_config) + except ValidationError: + return None + + +def _check_team_model_budget_update_authority( + data: UpdateTeamRequest, + user_api_key_dict: UserAPIKeyAuth, + existing_model_max_budget: Mapping[str, object] | None, +) -> None: + """Like `_check_team_budget_update_authority`: only a proxy admin may raise, re-window or drop a per-model cap.""" + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + return + if "model_max_budget" not in data.model_fields_set or not existing_model_max_budget: + return + requested: Final[Mapping[str, BudgetConfig]] = data.model_max_budget or {} + for model_name, raw_existing in existing_model_max_budget.items(): + existing = _existing_model_cap(raw_existing) + if existing is None or existing.max_budget is None: + continue + proposed = requested.get(model_name) + if proposed is None: + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can remove a team's model_max_budget for {model_name!r}. " + f"Current max_budget={existing.max_budget}." + ) + }, + ) + if ( + proposed.max_budget is None + or proposed.max_budget > existing.max_budget + or proposed.budget_duration != existing.budget_duration + ): + raise HTTPException( + status_code=403, + detail={ + "error": ( + f"Only a proxy admin can raise a team's model_max_budget for {model_name!r} or change its " + f"budget_duration. Current max_budget={existing.max_budget} per {existing.budget_duration}, " + f"requested={proposed.max_budget} per {proposed.budget_duration}." + ) + }, + ) + + def _should_auto_add_team_creator( user_api_key_dict: UserAPIKeyAuth, general_settings: Mapping[str, object], @@ -1230,6 +1284,7 @@ async def new_team( - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -1291,6 +1346,7 @@ async def new_team( general_settings, litellm_proxy_admin_name, llm_router, + premium_user, prisma_client, user_api_key_cache, ) @@ -1321,6 +1377,7 @@ 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) if data.soft_budget is not None: if data.max_budget is not None: @@ -1980,6 +2037,7 @@ async def update_team( - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -2031,6 +2089,7 @@ async def update_team( from litellm.proxy.proxy_server import ( litellm_proxy_admin_name, llm_router, + premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, @@ -2069,6 +2128,7 @@ 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) existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique( where={"team_id": data.team_id} @@ -2204,8 +2264,15 @@ async def update_team( user_api_key_dict=user_api_key_dict, existing_team_max_budget=existing_team_row.max_budget, ) + _check_team_model_budget_update_authority( + data=data, + user_api_key_dict=user_api_key_dict, + existing_model_max_budget=existing_team_row.model_max_budget, + ) updated_kv = data.json(exclude_unset=True) + if "model_max_budget" in updated_kv and updated_kv["model_max_budget"] is None: + updated_kv["model_max_budget"] = {} # Drop server-owned metadata keys from caller input so they can only # be written by the same code path that creates the underlying rows. @@ -4473,7 +4540,7 @@ async def team_info( ``` """ from litellm.proxy._types import TeamInfoResponseObjectTeamTable - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import model_max_budget_limiter, prisma_client try: if prisma_client is None: @@ -4573,6 +4640,12 @@ async def team_info( update={ # mutable-ok: pydantic update payload "members_with_roles": hydrated_members, "organization_models": organization_models, + "model_max_budget_usage": await build_model_max_budget_usage( + entity_type=Litellm_EntityType.TEAM, + entity_id=team_id, + model_max_budget=resolved_team_info.model_max_budget, + cache=model_max_budget_limiter.dual_cache, + ), } ) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 686544d352c..de6f9cb7647 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -609,6 +609,7 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): # merely shares the name. if not request_dispatched_to_pass_through_endpoint(request): _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_end_user_model_max_budget"] = user_api_key_dict.end_user_model_max_budget _metadata.update( diff --git a/litellm/proxy/spend_tracking/carried_budget_state.py b/litellm/proxy/spend_tracking/carried_budget_state.py index efd3a78d211..da8bf60ebda 100644 --- a/litellm/proxy/spend_tracking/carried_budget_state.py +++ b/litellm/proxy/spend_tracking/carried_budget_state.py @@ -25,6 +25,7 @@ def carry_team_and_user_budget_state( budget_reset_at=team_object.budget_reset_at, max_budget=team_object.max_budget, ) + valid_token.team_model_max_budget = team_object.model_max_budget # rebind-ok: caller keeps this object if user_object is not None: valid_token.user_budget_snapshot = UserBudgetSnapshot( # rebind-ok: same object the caller keeps using budget_reset_at=user_object.budget_reset_at, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 479bd0a55af..b77f25389c6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4319,6 +4319,7 @@ class PrismaClient: v.*, t.spend AS team_spend, t.max_budget AS team_max_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit @@ -4758,6 +4759,7 @@ class PrismaClient: t.spend AS team_spend, t.max_budget AS team_max_budget, t.soft_budget AS team_soft_budget, + t.model_max_budget AS team_model_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.tpd_limit AS team_tpd_limit, diff --git a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py index 096efc33aaf..efe41e1da9a 100644 --- a/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py +++ b/tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py @@ -587,6 +587,8 @@ def _success_kwargs( response_cost=0.5, key_hash=None, key_model_max_budget=None, + team_id=None, + team_model_max_budget=None, user_id=None, user_model_max_budget=None, end_user_id=None, @@ -600,6 +602,7 @@ def _success_kwargs( "end_user": end_user_id, "metadata": { "user_api_key_hash": key_hash, + "user_api_key_team_id": team_id, "user_api_key_user_id": user_id, "user_api_key_end_user_id": end_user_id, }, @@ -607,6 +610,7 @@ def _success_kwargs( "litellm_params": { "metadata": { "user_api_key_model_max_budget": key_model_max_budget, + "user_api_key_team_model_max_budget": team_model_max_budget, "user_api_key_user_model_max_budget": user_model_max_budget, "user_api_key_end_user_model_max_budget": end_user_model_max_budget, }, @@ -1417,3 +1421,266 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another() replica_c = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis)) with pytest.raises(litellm.BudgetExceededError): await replica_c.is_key_within_model_budget(user_api_key, "gpt-4") + + +def _log_success(limiter, **kwargs): + return limiter.async_log_success_event( + _success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_model", + ["gpt-4", "openai/gpt-4"], + ids=["bare_model", "provider_prefixed_model"], +) +async def test_team_model_budget_is_shared_by_every_key_without_an_override(request_model): + """ + Two keys on the same team, neither carrying a matching key-level entry, + charge one team counter and are both refused once it is spent. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + check = lambda: limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model=request_model, + ) + + assert await check() is True + await _log_success( + limiter, + model_group=request_model, + response_cost=0.6, + key_hash="vk-a", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + assert await check() is True + await _log_success( + limiter, + model_group=request_model, + response_cost=0.6, + key_hash="vk-b", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == pytest.approx(1.2) + with pytest.raises(litellm.BudgetExceededError) as exc: + await check() + assert exc.value.entity_type == Litellm_EntityType.TEAM.value + assert await build_model_max_budget_usage( + entity_type=Litellm_EntityType.TEAM, + entity_id="team-1", + model_max_budget=team_model_max_budget, + cache=dual_cache, + ) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}} + + +@pytest.mark.asyncio +async def test_key_override_replaces_the_team_cap_for_that_model(): + """ + A key with its own entry for the model is gated on the key counter alone: + the exhausted team counter does not block it, and its spend never lands on + the team counter. + """ + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": {"budget_limit": 5.0, "time_period": "1d"}} + await dual_cache.async_set_cache(key="team_model_spend:team-1:gpt-4:1d", value=9.0) + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="openai/gpt-4", + ) + is True + ) + + await _log_success( + limiter, + model_group="openai/gpt-4", + response_cost=2.0, + key_hash="vk-override", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 9.0 + assert await dual_cache.async_get_cache(key="virtual_key_spend:vk-override:gpt-4:1d") == 2.0 + + +@pytest.mark.asyncio +async def test_key_entry_for_another_model_does_not_lift_the_team_cap(): + """A key override only covers the model it names; other models stay on the team counter.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"claude-3": {"budget_limit": 5.0, "time_period": "1d"}} + + await _log_success( + limiter, + model_group="gpt-4", + response_cost=1.5, + key_hash="vk-other", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="gpt-4", + ) + + +@pytest.mark.asyncio +async def test_team_budget_leaves_unconfigured_models_alone(): + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 0.0, "time_period": "1d"}} + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="claude-3", + ) + is True + ) + with patch.object(limiter, "_increment_spend_for_key", new_callable=AsyncMock) as mock_increment: + await _log_success( + limiter, + model_group="claude-3", + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + mock_increment.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_counters_are_isolated_by_team_model_and_window(): + """Same model on two teams, and two models with different windows on one team, never share a counter.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = { + "gpt-4": {"budget_limit": 10.0, "time_period": "1d"}, + "claude-3": {"budget_limit": 10.0, "time_period": "30d"}, + } + + for team_id, model in (("team-1", "gpt-4"), ("team-2", "gpt-4"), ("team-1", "claude-3")): + await _log_success( + limiter, + model_group=model, + response_cost=1.0, + team_id=team_id, + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_spend:team-2:gpt-4:1d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:claude-3:30d") == 1.0 + assert await dual_cache.async_get_cache(key="team_model_budget_start_time:team-1:claude-3:30d") is not None + + +@pytest.mark.asyncio +async def test_malformed_team_entry_is_skipped_and_its_sibling_still_enforced(): + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache()) + team_model_max_budget = { + "gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}, + "claude-3": {"budget_limit": 0.0, "time_period": "1d"}, + } + + assert ( + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="gpt-4", + ) + is True + ) + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=None, + model="claude-3", + ) + + +@pytest.mark.asyncio +async def test_malformed_key_entry_does_not_count_as_an_override(): + """A key entry the limiter cannot enforce must not also switch the team cap off.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": {"budget_limit": "not-a-number", "time_period": "1d"}} + + await _log_success( + limiter, + model_group="gpt-4", + response_cost=1.5, + key_hash="vk-bad", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="gpt-4", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_entry", + [ + {"time_period": "1d", "tpm_limit": 100}, + {"time_period": "1d", "rpm_limit": 10}, + {"budget_limit": -1.0, "time_period": "1d"}, + ], +) +async def test_key_entry_without_a_spend_cap_does_not_lift_the_team_cap(key_entry): + """A key row that only rate-limits the model, or has no enforceable cap, leaves the team cap in force.""" + dual_cache = DualCache() + limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=dual_cache) + team_model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}} + key_model_max_budget = {"gpt-4": key_entry} + + await _log_success( + limiter, + model_group="openai/gpt-4", + response_cost=1.5, + key_hash="vk-rate-limited", + key_model_max_budget=key_model_max_budget, + team_id="team-1", + team_model_max_budget=team_model_max_budget, + ) + + assert await dual_cache.async_get_cache(key="team_model_spend:team-1:gpt-4:1d") == 1.5 + with pytest.raises(litellm.BudgetExceededError): + await limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget=team_model_max_budget, + key_model_max_budget=key_model_max_budget, + model="openai/gpt-4", + ) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 7660a8649b5..fc5d807bc23 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -1200,6 +1200,7 @@ def _fake_user_api_key_auth( team_models=None, team_id=None, model_max_budget=None, + team_model_max_budget=None, end_user_model_max_budget=None, end_user_id=None, user_model_max_budget=None, @@ -1220,6 +1221,7 @@ def _fake_user_api_key_auth( auth.team_id = team_id auth.team_model_aliases = None auth.model_max_budget = model_max_budget + auth.team_model_max_budget = team_model_max_budget 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 @@ -1860,6 +1862,78 @@ async def test_summary_model_rate_limit_skipped_for_legacy_limiter(): assert not result.applied_edits[0].get("error") +async def test_summary_model_denied_when_team_over_model_budget(): + """The team per-model budget gates the summary subrequest, whose spend is + charged to the team counter via the propagated `user_api_key_team_model_max_budget`. + The key's own `model_max_budget` is handed to the limiter so a key-level + override keeps taking precedence over the team cap here as it does in auth.""" + import litellm + + messages = _simple_messages() + mock_call = AsyncMock(return_value=_make_mock_response("x")) + key_budget = {"claude-opus-4-8": {"budget_limit": 1}} + team_budget = {"claude-haiku-4-5": {"budget_limit": 5, "time_period": "1d"}} + + auth = _fake_user_api_key_auth( + key_models=["all-proxy-models"], + model_max_budget=key_budget, + team_model_max_budget=team_budget, + team_id="team-over-budget", + token="hashed-token", + ) + + 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 + ) + ) + + with ( + patch( # test-quality-ok: apply_compact_20260112 reads the summary model setting as a module global, no seam + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._read_summary_model_setting", + return_value="claude-haiku-4-5", + ), + patch("litellm.token_counter", return_value=200_000), # test-quality-ok: forces the over-threshold branch + patch( # test-quality-ok: the summary call is the observable that must NOT happen when the team is over budget + "litellm.llms.anthropic.experimental_pass_through.context_management.editors.compact._call_summary_model", + mock_call, + ), + patch( # test-quality-ok: the limiter is a proxy_server module global the editor imports, no injection seam + "litellm.proxy.proxy_server.model_max_budget_limiter", limiter + ), + ): + result = await apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec=_EDIT_SPEC_DEFAULT, + user_api_key_auth=auth, + ) + + mock_call.assert_not_awaited() + assert result.applied_edits[0].get("error") == "summary_model_budget_exceeded" + limiter.is_team_within_model_budget.assert_awaited_once_with( + team_id="team-over-budget", + team_model_max_budget=team_budget, + key_model_max_budget=key_budget, + model="claude-haiku-4-5", + ) + import inspect + + from litellm.proxy.hooks.model_max_budget_limiter import ( + _PROXY_VirtualKeyModelMaxBudgetLimiter, + ) + + 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" + + async def test_scoped_budget_metadata_propagated_to_summary_call(): """The end-user/project scope identifiers and the end-user budget the post-call spend and rate-limit hooks key on are forwarded to the summary subrequest, and diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py index 7b6717f804f..f74531beaf0 100644 --- a/tests/test_litellm/proxy/auth/test_team_grants.py +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -31,6 +31,7 @@ def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: max_budget=50.0, soft_budget=25.0, spend=12.5, + model_max_budget={"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}}, models=["gpt-4o", "gpt-4o-mini"], blocked=True, metadata={"tier": "gold"}, @@ -72,6 +73,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets(): assert token.team_max_budget == 50.0 assert token.team_soft_budget == 25.0 assert token.team_spend == 12.5 + assert token.team_model_max_budget == {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} assert token.team_models == ["gpt-4o", "gpt-4o-mini"] assert token.team_blocked is True assert token.team_metadata == {"tier": "gold"} diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bd7ff62ac8b..c55dd966b2b 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -4372,6 +4372,74 @@ async def test_centralized_common_checks_carries_team_and_user_budget_state_on_t } +class _RecordingTeamModelBudgetLimiter: + def __init__(self): + self.calls = [] + + async def is_team_within_model_budget(self, team_id, team_model_max_budget, key_model_max_budget, model): + self.calls.append((team_id, dict(team_model_max_budget), key_model_max_budget, model)) + return True + + +@pytest.mark.asyncio +async def test_centralized_common_checks_enforces_team_model_max_budget_from_the_resolved_team(): + """The team's model_max_budget is enforced at the single authz gate, off the + team object auth resolved (not the possibly stale token copy), and the key's + own model_max_budget is handed to the limiter so a matching key entry can + override the team cap.""" + from fastapi import Request + from starlette.datastructures import URL + + import litellm.proxy.proxy_server as _proxy_server_mod + + team_caps = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + key_caps = {"claude-sonnet-4-6": {"max_budget": 1.0, "budget_duration": "1d"}} + token = UserAPIKeyAuth( + api_key="sk-test", + token="hashed", + team_id="t1", + team_model_max_budget={"gpt-4o": {"max_budget": 999.0, "budget_duration": "30d"}}, + model_max_budget=key_caps, + ) + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + user_api_key_cache = DualCache() + await user_api_key_cache.async_set_cache( + key="team_id:t1", + value=LiteLLM_TeamTableCachedObj(team_id="t1", model_max_budget=team_caps), + ) + limiter = _RecordingTeamModelBudgetLimiter() + attrs = { + **_proxy_attrs_for_centralized_checks(user_custom_auth=None), + "prisma_client": MagicMock(), + "user_api_key_cache": user_api_key_cache, + "model_max_budget_limiter": limiter, + } + originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs} + try: + for k, v in attrs.items(): + setattr(_proxy_server_mod, k, v) + with ( + patch("litellm.proxy.auth.user_api_key_auth.common_checks", new_callable=AsyncMock), # test-quality-ok: stubs the sibling check so only the team model-budget gate is under test + patch( # test-quality-ok: stubs the budget reservation so only the team model-budget gate is under test + "litellm.proxy.auth.user_api_key_auth._reserve_budget_after_common_checks", + new_callable=AsyncMock, + ), + ): + await _run_centralized_common_checks( + user_api_key_auth_obj=token, + request=request, + request_data={"model": "gpt-4o"}, + route="/chat/completions", + ) + finally: + for k, v in originals.items(): + setattr(_proxy_server_mod, k, v) + + assert limiter.calls == [("t1", team_caps, key_caps, "gpt-4o")] + + @pytest.mark.asyncio async def test_centralized_common_checks_skipped_for_custom_auth_without_flag(): """Existing RPS guarantee: custom-auth deployments without diff --git a/tests/test_litellm/proxy/db/test_create_views.py b/tests/test_litellm/proxy/db/test_create_views.py index ecc6d70123e..54418e10bdf 100644 --- a/tests/test_litellm/proxy/db/test_create_views.py +++ b/tests/test_litellm/proxy/db/test_create_views.py @@ -71,6 +71,7 @@ async def test_create_views_creates_view_on_does_not_exist(): mock_db.execute_raw.assert_called_once() created_sql = mock_db.execute_raw.call_args[0][0] assert 'CREATE VIEW "LiteLLM_VerificationTokenView"' in created_sql + assert "t.model_max_budget AS team_model_max_budget" in created_sql @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py index 7352ca0e9ee..2b614632346 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_utils.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_utils.py @@ -35,6 +35,7 @@ from litellm.proxy.management_endpoints.common_utils import ( admin_can_invite_user, ) from litellm.proxy.management_endpoints.common_utils import _has_non_empty_value +from litellm.types.utils import BudgetConfig class TestUpdateMetadataFieldsEmptyCollections: @@ -1162,3 +1163,54 @@ async def test_router_weights_validate_current_deployment_scope( assert exc.value.detail == error else: await validation + + +@pytest.mark.parametrize( + "model_max_budget, error", + [ + ({"gpt-4o": BudgetConfig(max_budget=-1.0, budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=float("inf"), budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=float("nan"), budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(budget_duration="1d")}, "non-negative finite"), + ({"gpt-4o": BudgetConfig(max_budget=5.0)}, "requires a budget_duration"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="fortnight")}, "budget_duration"), + ({" ": BudgetConfig(max_budget=5.0, budget_duration="1d")}, "non-empty model names"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", tpm_limit=1000)}, "not enforced on a team"), + ({"gpt-4o": BudgetConfig(max_budget=5.0, budget_duration="1d", rpm_limit=10)}, "not enforced on a team"), + ], + ids=["negative", "inf", "nan", "no_cap", "no_duration", "bad_duration", "blank_model", "tpm_limit", "rpm_limit"], +) +def test_validate_team_model_max_budget_rejects_unenforceable_entries(model_max_budget, error) -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + with pytest.raises(HTTPException) as exc: + validate_team_model_max_budget(model_max_budget=model_max_budget, premium_user=True) + assert exc.value.status_code == 400 + assert error in exc.value.detail["error"] + + +def test_validate_team_model_max_budget_accepts_a_zero_cap_and_prefixed_models() -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + assert ( + validate_team_model_max_budget( + model_max_budget={ + "gpt-4o": BudgetConfig(max_budget=0.0, budget_duration="1d"), + "openai/gpt-4o-mini": BudgetConfig(max_budget=2.5, budget_duration="30d"), + }, + premium_user=True, + ) + is None + ) + + +def test_validate_team_model_max_budget_is_license_gated_only_when_set() -> None: + from litellm.proxy.management_endpoints.common_utils import validate_team_model_max_budget + + validate_team_model_max_budget(model_max_budget=None, premium_user=False) + validate_team_model_max_budget(model_max_budget={}, premium_user=False) + with pytest.raises(HTTPException) as exc: + validate_team_model_max_budget( + model_max_budget={"gpt-4o": BudgetConfig(max_budget=1.0, budget_duration="1d")}, premium_user=False + ) + assert exc.value.status_code == 403 diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index ebbedc6541e..dda5bb344b4 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -14651,3 +14651,232 @@ async def test_team_info_reports_parent_organization_models_only_to_team_manager ) assert response["team_info"].organization_models == expected_models + + +_EXISTING_TEAM_MODEL_CAPS: Final = { + "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}, + "claude-sonnet-4-6": {"max_budget": 5.0, "budget_duration": "7d"}, +} + + +@pytest.mark.parametrize( + "requested", + [ + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 20.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 10.0, "budget_duration": "30d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"budget_duration": "1d"}}, + {"claude-sonnet-4-6": _EXISTING_TEAM_MODEL_CAPS["claude-sonnet-4-6"]}, + {}, + None, + ], + ids=["raise", "change_duration", "drop_cap_value", "remove_model", "clear_all", "clear_with_null"], +) +def test_team_admin_cannot_loosen_team_model_caps(requested) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + with pytest.raises(HTTPException) as exc: + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=requested), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + assert exc.value.status_code == 403 + assert "proxy admin" in exc.value.detail["error"] + + +@pytest.mark.parametrize( + "requested", + [ + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o": {"max_budget": 2.0, "budget_duration": "1d"}}, + {**_EXISTING_TEAM_MODEL_CAPS, "gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}}, + dict(_EXISTING_TEAM_MODEL_CAPS), + ], + ids=["lower", "add_model", "unchanged"], +) +def test_team_admin_can_tighten_or_keep_team_model_caps(requested) -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + assert ( + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=requested), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin"), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ) + is None + ) + + +def test_team_model_cap_authority_skips_omitted_field_malformed_rows_and_proxy_admins() -> None: + from litellm.proxy.management_endpoints.team_endpoints import _check_team_model_budget_update_authority + + team_admin = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin") + outcomes = ( + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", max_budget=1.0), + user_api_key_dict=team_admin, + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ), + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget={}), + user_api_key_dict=team_admin, + existing_model_max_budget={"gpt-4o": "not-a-budget", "gpt-4o-mini": {"budget_duration": "1d"}}, + ), + _check_team_model_budget_update_authority( + data=UpdateTeamRequest(team_id="t1", model_max_budget=None), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + existing_model_max_budget=_EXISTING_TEAM_MODEL_CAPS, + ), + ) + assert outcomes == (None, None, None) + + +@pytest.mark.asyncio +async def test_new_team_persists_model_max_budget(mock_db_client, mock_admin_auth): + 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")) + + team_create_result = MagicMock(team_id="team-model-caps") + team_create_result.model_dump.return_value = {"team_id": "team-model-caps"} + mock_team_create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + from fastapi import Request + + 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 + await new_team( + data=NewTeamRequest( + team_alias="model-caps", + model_max_budget={"gpt-4o": {"max_budget": 10.0, "budget_duration": "1d"}}, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data["model_max_budget"] == { + "gpt-4o": {"max_budget": 10.0, "budget_duration": "1d", "tpm_limit": None, "rpm_limit": None} + } + + +@pytest.mark.asyncio +async def test_new_team_rejects_unenforceable_model_max_budget(mock_db_client, mock_admin_auth): + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest, ProxyException + from litellm.proxy.management_endpoints.team_endpoints import new_team + + mock_db_client.db.litellm_teamtable.create = AsyncMock() + + with patch("litellm.proxy.proxy_server.premium_user", True), pytest.raises(ProxyException) as exc: # test-quality-ok: proxy_server module global is the endpoint's only injection point + await new_team( + data=NewTeamRequest(team_alias="model-caps", model_max_budget={"gpt-4o": {"max_budget": 10.0}}), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + assert exc.value.code == "400" + assert "budget_duration" in str(exc.value.message) + mock_db_client.db.litellm_teamtable.create.assert_not_awaited() + + +def _existing_team_with_model_caps(caps): + existing = MagicMock() + existing.team_id = "standalone-team-123" + existing.organization_id = None + existing.max_budget = None + existing.model_id = None + existing.model_max_budget = caps + existing.model_dump.return_value = { + "team_id": "standalone-team-123", + "organization_id": None, + "model_max_budget": caps, + "members_with_roles": [{"user_id": "team-admin-model-caps", "role": "admin"}], + } + return existing + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cleared_with", [{}, None], ids=["empty_mapping", "null"]) +async def test_update_team_clearing_model_max_budget_writes_an_empty_mapping( + disable_audit_logging_for_mocked_team, cleared_with +): + from fastapi import Request + + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS) + ) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + updated = _existing_team_with_model_caps({}) + updated.litellm_model_table = None + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated) + + await update_team( + data=UpdateTeamRequest(team_id="standalone-team-123", model_max_budget=cleared_with), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + assert mock_prisma.db.litellm_teamtable.update.call_args.kwargs["data"]["model_max_budget"] == {} + + +@pytest.mark.asyncio +async def test_update_team_model_max_budget_raise_blocked_for_team_admin(): + from fastapi import Request + + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache, # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch("litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock()), # test-quality-ok: stubs the audit write so the test observes only the team update result + ): + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock( + return_value=_existing_team_with_model_caps(_EXISTING_TEAM_MODEL_CAPS) + ) + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_prisma.db.litellm_teamtable.update = AsyncMock() + + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest( + team_id="standalone-team-123", + model_max_budget={ + **_EXISTING_TEAM_MODEL_CAPS, + "gpt-4o": {"max_budget": 100.0, "budget_duration": "1d"}, + }, + ), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="team-admin-model-caps", models=[] + ), + ) + + assert exc.value.code == "403" + assert "proxy admin" in str(exc.value.message).lower() + mock_prisma.db.litellm_teamtable.update.assert_not_awaited() diff --git a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py index fe852be775c..0bdf43b396c 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py +++ b/tests/test_litellm/proxy/spend_tracking/test_carried_budget_state.py @@ -47,6 +47,19 @@ def test_team_and_user_state_round_trips_through_metadata(): ) +def test_team_model_max_budget_rides_on_the_token(): + """The team's per-model caps must reach the token, or the auth check and the spend hook never see them.""" + token = UserAPIKeyAuth(token="hashed", team_id="t1") + team_model_max_budget = {"gpt-4o": {"max_budget": 5.0, "budget_duration": "1d"}} + carry_team_and_user_budget_state( + valid_token=token, + team_object=LiteLLM_TeamTable(team_id="t1", model_max_budget=team_model_max_budget), + user_object=None, + ) + + assert token.team_model_max_budget == team_model_max_budget + + def test_missing_objects_leave_no_metadata_and_no_snapshot(): token = UserAPIKeyAuth(token="hashed", team_id="t1", user_id="u1") carry_team_and_user_budget_state(valid_token=token, team_object=None, user_object=None) diff --git a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py index ce6ecc2ea65..672dd1eb674 100644 --- a/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py +++ b/tests/test_litellm/proxy/utils/prisma_and_spend/test_prisma_client_get_data.py @@ -401,19 +401,20 @@ async def test_check_view_exists_creates_token_view_when_missing( prisma_client.db.execute_raw = AsyncMock() prisma_client.health_check = AsyncMock(return_value=[{"?column?": 1}]) result = await prisma_client.check_view_exists() + created_sql = prisma_client.db.execute_raw.await_args.args[0] actual = { "result": result, "create_called": prisma_client.db.execute_raw.await_count, - "create_sql_starts_with_create_view": prisma_client.db.execute_raw.await_args.args[ - 0 - ] - .strip() - .startswith('CREATE VIEW "LiteLLM_VerificationTokenView"'), + "create_sql_starts_with_create_view": created_sql.strip().startswith( + 'CREATE VIEW "LiteLLM_VerificationTokenView"' + ), + "projects_team_model_max_budget": "t.model_max_budget AS team_model_max_budget" in created_sql, } assert actual == { "result": None, "create_called": 1, "create_sql_starts_with_create_view": True, + "projects_team_model_max_budget": True, } diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 851c9e6d487..f2d7cff2ec4 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -7,6 +7,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import { toast } from "@/lib/toast"; import { fetchAvailableModelsForTeamOrKey } from "./key_team_helpers/fetch_available_models_team_key"; +import { MODEL_MAX_BUDGET_PREMIUM_HINT } from "./key_team_helpers/ModelMaxBudgetEditor"; import { fetchMCPAccessGroups, getDefaultTeamSettings, @@ -1547,6 +1548,36 @@ describe("Teams - the exact bytes the create call sends", () => { expect(await screen.findByText("Please input a team name")).toBeInTheDocument(); expect(teamCreateCall).not.toHaveBeenCalled(); }); + + it("locks the per-model budget editor and says why when the proxy has no enterprise license", async () => { + await openCreateModal({ premiumUser: false }); + + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + expect(screen.getByText(MODEL_MAX_BUDGET_PREMIUM_HINT)).toBeInTheDocument(); + }); + + it("sends the per-model budget a licensed operator fills in, keyed by model", async () => { + const user = userEvent.setup({ delay: null }); + await openCreateModal({ premiumUser: true }); + + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await chooseSelectOption(user, screen.getByPlaceholderText("Select model"), "gpt-4"); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "3" } }); + + const payload = await submit(); + + expect(payload.model_max_budget).toStrictEqual({ "gpt-4": { budget_limit: 3, time_period: "30d" } }); + }); + + it("leaves model_max_budget out when a started row is removed again", async () => { + const user = userEvent.setup({ delay: null }); + await openCreateModal({ premiumUser: true }); + + await user.click(screen.getByRole("button", { name: /Add Model Budget/i })); + await user.click(screen.getByRole("button", { name: "Remove model budget" })); + + expect(wireBody(await submit())).not.toHaveProperty("model_max_budget"); + }); }); describe("Teams - the create form keeps the organization and models picks while it is open", () => { diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index 4f3367d8b98..7214d16f665 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -48,6 +48,7 @@ import BudgetDurationDropdown, { } from "./common_components/budget_duration_dropdown"; import { Organization, getDefaultTeamSettings, getGuardrailsList, getPoliciesList, teamDeleteCall } from "./networking"; import NumericalInput from "./shared/numerical_input"; +import { ModelMaxBudget, ModelMaxBudgetField } from "./key_team_helpers/ModelMaxBudgetEditor"; import VectorStoreSelector from "./vector_store_management/VectorStoreSelector"; import SearchToolSelector from "./search_tools/SearchToolSelector"; import SkillSelector from "./skills/SkillSelector"; @@ -271,6 +272,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser const [policiesList, setPoliciesList] = useState([]); const [loggingSettings, setLoggingSettings] = useState([]); const [modelAliases, setModelAliases] = useState<{ [key: string]: string }>({}); + const [modelMaxBudget, setModelMaxBudget] = useState({}); const [routerSettings, setRouterSettings] = useState(null); const [routerSettingsKey, setRouterSettingsKey] = useState(0); @@ -348,6 +350,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser setSearchToolSettingsOpen(false); setLoggingSettings([]); setModelAliases({}); + setModelMaxBudget({}); setRouterSettings(null); setRouterSettingsKey((prev) => prev + 1); }; @@ -525,6 +528,10 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser formValues.model_aliases = modelAliases; } + if (Object.keys(modelMaxBudget).length > 0) { + formValues.model_max_budget = modelMaxBudget; + } + // Add router_settings if any are defined if (routerSettings?.router_settings) { // Only include router_settings if it has at least one non-null value @@ -813,6 +820,14 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser /> )} + {({ ref, value, ...field }) => ( diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx index 0fab5555343..4d2a0e88f84 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/ModelMaxBudgetEditor.tsx @@ -144,6 +144,7 @@ export function ModelMaxBudgetEditor({ onClick={() => removeEntry(entry.id)} disabled={!premiumUser} title={hintWhenLocked} + aria-label="Remove model budget" className="absolute top-2 right-2 text-muted-foreground hover:text-destructive transition-colors p-1" > diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index eb912ffa3cc..16eb70fdf5e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1609,6 +1609,99 @@ describe("TeamInfoView", () => { }); }); + describe("per-model budgets", () => { + const teamWithModelBudget = () => + createMockTeamData({ + models: ["gpt-4"], + model_max_budget: { "gpt-4": { max_budget: 5, budget_duration: "1d" } }, + model_max_budget_usage: { "gpt-4": { current_spend: 1.25, budget_limit: 5, time_period: "1d" } }, + }); + + const openSettingsEditor = async (user: ReturnType) => { + 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 screen.findByLabelText("Team Name"); + }; + + const savedPayload = async () => { + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + return vi.mocked(networking.teamUpdateCall).mock.calls[0][1] as Record; + }; + + it("shows the stored per-model budget and its current spend in the read-only settings view", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + + renderWithProviders(); + + await waitFor(() => { + expect(screen.queryAllByText("Test Team").length).toBeGreaterThan(0); + }); + await user.click(screen.getByRole("tab", { name: "Settings" })); + + expect(await screen.findByText("Per-Model Budget (gpt-4): $5 per 1d, spent $1.25")).toBeInTheDocument(); + }); + + it("seeds the editor from the stored budget and keeps it read-only without an enterprise license", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + + renderWithProviders(); + + await openSettingsEditor(user); + + expect(screen.getByPlaceholderText("Max spend ($)")).toHaveValue(5); + expect(screen.getByPlaceholderText("Max spend ($)")).toBeDisabled(); + expect(screen.getByRole("button", { name: /Add Model Budget/i })).toBeDisabled(); + }); + + it("leaves model_max_budget out of a save that did not touch it", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect(await savedPayload()).not.toHaveProperty("model_max_budget"); + }); + + it("sends the edited cap for the model", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + fireEvent.change(screen.getByPlaceholderText("Max spend ($)"), { target: { value: "2.5" } }); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect((await savedPayload()).model_max_budget).toEqual({ "gpt-4": { budget_limit: 2.5, time_period: "1d" } }); + }); + + it("sends an empty model_max_budget when the last row is removed, so the stored cap is cleared", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(teamWithModelBudget()); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + + await openSettingsEditor(user); + await user.click(screen.getByRole("button", { name: "Remove model budget" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + expect((await savedPayload()).model_max_budget).toEqual({}); + }); + }); + describe("team member settings", () => { it("should populate Default Key Duration from the team's stored metadata", async () => { const user = userEvent.setup({ delay: null }); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ce705008678..c476f3492a3 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -51,6 +51,13 @@ import GuardrailsSelect from "./GuardrailsSelect"; import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown, { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; +import { + ModelBudgetUsage, + ModelMaxBudget, + ModelMaxBudgetField, + modelMaxBudgetToEntries, +} from "../key_team_helpers/ModelMaxBudgetEditor"; +import { modelMaxBudgetUpdate, StoredModelMaxBudget } from "../key_team_helpers/modelMaxBudgetPayload"; import { computeTeamModelBadges, normalizeTeamModelSelection, @@ -268,6 +275,8 @@ export interface TeamData { max_budget: number | null; soft_budget?: number | null; budget_duration: string | null; + model_max_budget?: StoredModelMaxBudget | null; + model_max_budget_usage?: Record | null; models: string[]; blocked: boolean; spend: number; @@ -563,6 +572,7 @@ const TeamInfoView: React.FC = ({ const [isDeleting, setIsDeleting] = useState(false); const [isTeamSaving, setIsTeamSaving] = useState(false); const [teamModelAliases, setTeamModelAliases] = useState>({}); + const [teamModelMaxBudget, setTeamModelMaxBudget] = useState({}); const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); @@ -628,6 +638,7 @@ const TeamInfoView: React.FC = ({ const startEditing = () => { form.reset(teamFormValues()); + setTeamModelMaxBudget((teamData?.team_info?.model_max_budget ?? {}) as ModelMaxBudget); setTeamMemberSettingsOpen(false); setSearchToolSettingsOpen(false); setIsEditing(true); @@ -1078,6 +1089,11 @@ const TeamInfoView: React.FC = ({ updateData.model_aliases = teamModelAliases; } + const modelBudgets = modelMaxBudgetUpdate(teamModelMaxBudget, info.model_max_budget); + if (modelBudgets !== undefined) { + updateData.model_max_budget = modelBudgets; + } + // Handle router_settings - read fresh values from DOM at save time. const currentRouterSettings = routerSettingsRef.current?.getValue(); if (currentRouterSettings?.router_settings) { @@ -1536,6 +1552,15 @@ const TeamInfoView: React.FC = ({ )} + + {({ ref, value, ...field }) => } @@ -2051,6 +2076,17 @@ const TeamInfoView: React.FC = ({ : "No Limit"}
Budget Reset: {info.budget_duration || "Never"}
+ {modelMaxBudgetToEntries(info.model_max_budget as ModelMaxBudget | null | undefined).map( + ({ model, budgetLimit, timePeriod }) => { + const spent = model === null ? undefined : info.model_max_budget_usage?.[model]?.current_spend; + return ( +
+ Per-Model Budget ({model}): ${budgetLimit ?? "?"} per {timePeriod} + {spent !== undefined && `, spent $${spent}`} +
+ ); + }, + )} {info.metadata?.soft_budget_alerting_emails && Array.isArray(info.metadata.soft_budget_alerting_emails) && info.metadata.soft_budget_alerting_emails.length > 0 && ( diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 7ca30d5c4f0..f7a26a26370 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15626,6 +15626,7 @@ export interface paths { * - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -15852,6 +15853,7 @@ export interface paths { * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`. * - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias) + * - model_max_budget: Optional[dict] - Per-model max budget every key on the team inherits unless the key sets its own for that model. Example: {"gpt-4o": {"max_budget": 10, "budget_duration": "1d"}} * - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails) * - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies) * - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. @@ -33328,6 +33330,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -34086,6 +34095,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -39225,6 +39241,13 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** + * Model Max Budget + * @description Max budget per model for every key on the team, overridable per key (e.g. {'gpt-4o': {'max_budget': 10, 'budget_duration': '1d'}}) + */ + model_max_budget?: { + [key: string]: components["schemas"]["BudgetConfig"]; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -39936,6 +39959,10 @@ export interface components { team_model_aliases?: { [key: string]: unknown; } | null; + /** Team Model Max Budget */ + team_model_max_budget?: { + [key: string]: unknown; + } | null; /** * Team Models * @default []