diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e8355f4941a..f5727a89189 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,6 +1,7 @@ import enum import json import os +from collections.abc import Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union @@ -1785,6 +1786,7 @@ from litellm.models.team import TeamBase as TeamBase # noqa: E402 class NewTeamRequest(TeamBase): model_aliases: Optional[dict] = None + model_max_budget: Mapping[str, Mapping[str, str | float]] | None = None tags: Optional[list] = None guardrails: Optional[List[str]] = None policies: Optional[List[str]] = None @@ -1847,6 +1849,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None max_budget: Optional[float] = None soft_budget: Optional[float] = None + model_max_budget: Mapping[str, Mapping[str, str | float]] | None = None models: Optional[list] = None blocked: Optional[bool] = None budget_duration: Optional[str] = None @@ -2542,6 +2545,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_blocked: bool = False soft_budget: Optional[float] = None team_model_aliases: Optional[Dict] = None + team_model_max_budget: Mapping[str, Mapping[str, str | float]] | None = None team_member: Optional[Member] = None team_metadata: Optional[Dict] = None team_object_permission_id: Optional[str] = None diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 263fec77d12..6d9aaf05805 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,6 +13,7 @@ import asyncio import math import re import time +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Type, Union, cast from fastapi import HTTPException, Request, status @@ -680,6 +681,11 @@ async def common_checks( valid_token=valid_token, ), _team_multi_budget_check(team_object=team_object), + _team_model_max_budget_check( + team_object=team_object, + valid_token=valid_token, + model=_model, + ), _virtual_key_multi_budget_check(valid_token=valid_token) if valid_token is not None else None, _team_soft_budget_check( team_object=team_object, @@ -4011,6 +4017,39 @@ async def _team_max_budget_check( ) +async def _team_model_max_budget_check( + team_object: LiteLLM_TeamTable | None, + valid_token: UserAPIKeyAuth | None, + model: str | Sequence[str] | None, +) -> None: + """ + Enforce team-level per-model max budgets (team_object.model_max_budget). + + Every key on the team shares one spend counter per (model, duration) window. + A key whose own model_max_budget covers the request model is exempt — the + key-level check in user_api_key_auth already enforced its private cap. + + Raises BudgetExceededError when the shared team counter is over the cap. + """ + if team_object is None or team_object.team_id is None or model is None: + return + team_model_max_budget = team_object.model_max_budget + if not isinstance(team_model_max_budget, dict) or len(team_model_max_budget) == 0: + return + + from litellm.proxy.proxy_server import model_max_budget_limiter + + models = (model,) if isinstance(model, str) else tuple(model) + key_model_max_budget = valid_token.model_max_budget if valid_token is not None else None + for model_name in models: + await model_max_budget_limiter.is_team_within_model_budget( + team_id=team_object.team_id, + team_model_max_budget=team_model_max_budget, + model=model_name, + key_model_max_budget=key_model_max_budget, + ) + + async def _team_multi_budget_check( team_object: Optional[LiteLLM_TeamTable], ): diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index f4d07c1a674..f67ae9bc808 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -2319,6 +2319,8 @@ async def _run_centralized_common_checks( if user_api_key_auth_obj.org_id is None and team_object is not None and team_object.organization_id is not None: user_api_key_auth_obj.org_id = team_object.organization_id + user_api_key_auth_obj.team_model_max_budget = team_object.model_max_budget if team_object is not None else None + # common_checks identifies admin via user_object, not the token # (non_proxy_admin_allowed_routes_check). JWT admin shortcut and # master_key tokens get admin from the token; the DB row for the diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py index e8cf5fbc718..7371c053f9f 100644 --- a/litellm/proxy/hooks/model_max_budget_limiter.py +++ b/litellm/proxy/hooks/model_max_budget_limiter.py @@ -1,4 +1,6 @@ import json +from collections.abc import Mapping +from types import MappingProxyType from typing import List, Optional import litellm @@ -16,6 +18,7 @@ from litellm.types.utils import ( VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX = "virtual_key_spend" END_USER_SPEND_CACHE_KEY_PREFIX = "end_user_model_spend" +TEAM_MODEL_SPEND_CACHE_KEY_PREFIX = "team_model_spend" class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): @@ -148,6 +151,144 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return True + def _coerce_budget_configs( + self, model_max_budget: Mapping[str, Mapping[str, str | float]] + ) -> Mapping[str, BudgetConfig]: + return MappingProxyType( + {_model: BudgetConfig(**_budget_info) for _model, _budget_info in model_max_budget.items()} + ) + + def _key_already_covers_model( + self, key_model_max_budget: Mapping[str, Mapping[str, str | float]] | None, model: str + ) -> bool: + """ + True iff the requesting key declares its own model_max_budget entry for `model` + (exact or `{provider}/{model}`-normalized match). Key entries take precedence + over team-level defaults, so the team budget is neither checked nor incremented + for such (key, model) pairs. + """ + if not key_model_max_budget: + return False + return ( + self._get_request_model_budget_config( + model=model, internal_model_max_budget=self._coerce_budget_configs(key_model_max_budget) + ) + is not None + ) + + async def is_team_within_model_budget( + self, + team_id: str, + team_model_max_budget: Mapping[str, Mapping[str, str | float]], + model: str, + key_model_max_budget: Mapping[str, Mapping[str, str | float]] | None = None, + ) -> bool: + """ + Check if the team is within its team-level model budget. + + The team budget is the shared default for every key on the team; a key-level + model_max_budget entry covering `model` takes precedence and short-circuits + this check. + + Raises: + BudgetExceededError: If the team has exceeded the model budget + """ + if self._key_already_covers_model(key_model_max_budget, model): + verbose_proxy_logger.debug(f"Team model budget check skipped for model={model}: key has its own entry") + return True + + internal_model_max_budget = self._coerce_budget_configs(team_model_max_budget) + + verbose_proxy_logger.debug("team internal_model_max_budget %s", internal_model_max_budget) + + budget_config = self._get_request_model_budget_config( + model=model, internal_model_max_budget=internal_model_max_budget + ) + if budget_config is None: + verbose_proxy_logger.debug(f"Model {model} not found in team_model_max_budget") + return True + + if not budget_config.max_budget or budget_config.max_budget <= 0: + return True + + current_spend = await self._get_team_spend_for_model( + team_id=team_id, + model=model, + budget_config=budget_config, + ) + verbose_proxy_logger.debug( + f"team model budget check: team={team_id} model={model} spend={current_spend} max_budget={budget_config.max_budget}" + ) + if current_spend is not None and current_spend > budget_config.max_budget: + raise litellm.BudgetExceededError( + message=f"LiteLLM Team: {team_id}, exceeded budget for model={model}", + current_cost=current_spend, + max_budget=budget_config.max_budget, + entity_type=Litellm_EntityType.TEAM.value, + entity_id=team_id, + ) + + return True + + async def _get_team_spend_for_model( + self, + team_id: str, + model: str, + budget_config: BudgetConfig, + ) -> float | None: + """ + Get the current team spend for a model. + + Lookup model in this order: + 1. model: directly look up `model` + 2. If 1, does not exist, check if passed as {custom_llm_provider}/model + """ + team_model_spend_cache_key = ( + f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{model}:{budget_config.budget_duration}" + ) + _current_spend = await self.dual_cache.async_get_cache( + key=team_model_spend_cache_key, + ) + + if _current_spend is None: + team_model_spend_cache_key = f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{self._get_model_without_custom_llm_provider(model)}:{budget_config.budget_duration}" + _current_spend = await self.dual_cache.async_get_cache( + key=team_model_spend_cache_key, + ) + return _current_spend + + async def _track_team_spend_for_model( + self, + team_id: str, + team_model_max_budget: Mapping[str, Mapping[str, str | float]], + key_model_max_budget: Mapping[str, Mapping[str, str | float]] | None, + model: str, + response_cost: float, + ) -> None: + """ + Increment the shared team counter for `model`, unless the requesting key + declares its own model_max_budget entry for it — mirror of the precedence + rule in is_team_within_model_budget, so keys with private caps never drive + the team counter and block sibling keys. + """ + if self._key_already_covers_model(key_model_max_budget, model): + return + + budget_config = self._get_request_model_budget_config( + model=model, internal_model_max_budget=self._coerce_budget_configs(team_model_max_budget) + ) + if budget_config is None or not budget_config.budget_duration: + return + + team_spend_key = f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:{team_id}:{model}:{budget_config.budget_duration}" + team_start_time_key = f"team_model_budget_start_time:{team_id}" + await self._increment_spend_for_key( + budget_config=budget_config, + spend_key=team_spend_key, + start_time_key=team_start_time_key, + response_cost=response_cost, + ) + async def _get_end_user_spend_for_model( self, end_user_id: str, @@ -202,7 +343,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): return _current_spend def _get_request_model_budget_config( - self, model: str, internal_model_max_budget: GenericBudgetConfigType + self, model: str, internal_model_max_budget: Mapping[str, BudgetConfig] ) -> Optional[BudgetConfig]: """ Get the budget config for the request model @@ -249,11 +390,16 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): user_api_key_end_user_model_max_budget: Optional[dict] = _metadata.get( "user_api_key_end_user_model_max_budget", None ) - if (user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0) and ( - user_api_key_end_user_model_max_budget is None or len(user_api_key_end_user_model_max_budget) == 0 + user_api_key_team_model_max_budget: Mapping[str, Mapping[str, str | float]] | None = _metadata.get( + "user_api_key_team_model_max_budget", None + ) + if ( + (user_api_key_model_max_budget is None or len(user_api_key_model_max_budget) == 0) + and (user_api_key_end_user_model_max_budget is None or len(user_api_key_end_user_model_max_budget) == 0) + and (user_api_key_team_model_max_budget is None or len(user_api_key_team_model_max_budget) == 0) ): verbose_proxy_logger.debug( - "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget and user_api_key_end_user_model_max_budget are None or empty." + "Not running _PROXY_VirtualKeyModelMaxBudgetLimiter.async_log_success_event because user_api_key_model_max_budget, user_api_key_end_user_model_max_budget, and user_api_key_team_model_max_budget are None or empty." ) return @@ -320,6 +466,20 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting): response_cost=response_cost, ) + user_api_key_team_id: str | None = _metadata.get("user_api_key_team_id", None) + if ( + user_api_key_team_id is not None + and user_api_key_team_model_max_budget is not None + and len(user_api_key_team_model_max_budget) > 0 + ): + await self._track_team_spend_for_model( + team_id=user_api_key_team_id, + team_model_max_budget=user_api_key_team_model_max_budget, + key_model_max_budget=user_api_key_model_max_budget, + model=model, + response_cost=response_cost, + ) + if self.dual_cache.redis_cache is not None: await self._push_in_memory_increments_to_redis() diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 1fad1954dc4..47215b92e9d 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1672,6 +1672,7 @@ async def add_litellm_data_to_request( data[_metadata_variable_name]["user_api_key_end_user_model_max_budget"] = ( user_api_key_dict.end_user_model_max_budget ) + data[_metadata_variable_name]["user_api_key_team_model_max_budget"] = user_api_key_dict.team_model_max_budget # User spend, budget - used by prometheus.py # Follow same pattern as team and API key budgets diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index c336e9321e0..dfff4eec182 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -6551,7 +6551,7 @@ async def _enforce_unique_key_alias( ) -def validate_model_max_budget(model_max_budget: Optional[Dict]) -> None: +def validate_model_max_budget(model_max_budget: Mapping[str, Mapping[str, str | float]] | None) -> None: """ Validate the model_max_budget is GenericBudgetConfigType + enforce user has an enterprise license diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 1efd0a747df..19eda04ece8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1082,6 +1082,7 @@ async def new_team( - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. + - model_max_budget: Optional[dict] - Per-model max budgets shared by every key on the team, e.g. {"gpt-4o": {"budget_limit": 100.0, "time_period": "1d"}}. A key's own model_max_budget entry for a model takes precedence; the team entry is the default cap for keys without one. - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. - blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id. @@ -1188,6 +1189,21 @@ async def new_team( }, ) + if data.model_max_budget is not None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_model_max_budget, + ) + + try: + validate_model_max_budget(data.model_max_budget) + except ValueError as e: + raise ProxyException( + message=str(e), + type=ProxyErrorTypes.bad_request_error, + param="model_max_budget", + code="400", + ) + # Check if license is over limit total_teams = await _team_db(prisma_client).count() if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): @@ -1751,6 +1767,7 @@ async def update_team( - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. + - model_max_budget: Optional[dict] - Per-model max budgets shared by every key on the team, e.g. {"gpt-4o": {"budget_limit": 100.0, "time_period": "1d"}}. A key's own model_max_budget entry for a model takes precedence; the team entry is the default cap for keys without one. - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. @@ -1954,6 +1971,21 @@ async def update_team( existing_team_max_budget=existing_team_row.max_budget, ) + if data.model_max_budget is not None: + from litellm.proxy.management_endpoints.key_management_endpoints import ( + validate_model_max_budget, + ) + + try: + validate_model_max_budget(data.model_max_budget) + except ValueError as e: + raise ProxyException( + message=str(e), + type=ProxyErrorTypes.bad_request_error, + param="model_max_budget", + code="400", + ) + updated_kv = data.json(exclude_unset=True) # Drop server-owned metadata keys from caller input so they can only 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 55459721906..7eceadbcc2d 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 @@ -558,3 +558,281 @@ async def test_async_log_success_event_skips_redis_push_without_redis(budget_lim kwargs, response_obj=None, start_time=None, end_time=None ) mock_push.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_is_team_within_model_budget_under_budget(budget_limiter): + """Team-level cap admits requests while the shared counter is below the limit.""" + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=10.0) + ): + assert ( + await budget_limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget={ + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + model="gpt-4", + ) + is True + ) + + +@pytest.mark.asyncio +async def test_is_team_within_model_budget_exceeded_raises(budget_limiter): + """Exceeding the team cap raises BudgetExceededError attributed to the TEAM entity.""" + from litellm.proxy._types import Litellm_EntityType + + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=150.0) + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await budget_limiter.is_team_within_model_budget( + team_id="team-finance", + team_model_max_budget={ + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + model="gpt-4", + ) + assert "team-finance" in str(exc_info.value) + assert "gpt-4" in str(exc_info.value) + assert exc_info.value.entity_type == Litellm_EntityType.TEAM.value + assert exc_info.value.entity_id == "team-finance" + + +@pytest.mark.asyncio +async def test_is_team_within_model_budget_model_not_configured(budget_limiter): + """Models absent from team_model_max_budget stay uncapped: no spend lookup, no error.""" + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=999.0) + ) as mock_get_spend: + assert ( + await budget_limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget={ + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + model="claude-3", + ) + is True + ) + mock_get_spend.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_check_skipped_when_key_covers_model(budget_limiter): + """A key with its own entry for the model is exempt from the team cap, so an + over-budget team counter must not reject it (the key check already ran).""" + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=10_000.0) + ) as mock_team_spend: + assert ( + await budget_limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget={ + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + model="gpt-4", + key_model_max_budget={ + "gpt-4": {"budget_limit": 50.0, "time_period": "1d"} + }, + ) + is True + ) + mock_team_spend.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_check_runs_when_key_covers_different_model(budget_limiter): + """A key entry for another model does not exempt the requested model from the team cap.""" + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=10.0) + ) as mock_team_spend: + assert ( + await budget_limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget={ + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + model="gpt-4", + key_model_max_budget={ + "claude-3": {"budget_limit": 50.0, "time_period": "1d"} + }, + ) + is True + ) + mock_team_spend.assert_called_once() + + +@pytest.mark.asyncio +async def test_team_check_skipped_for_provider_prefixed_key_match(budget_limiter): + """Key precedence must apply through {provider}/{model} normalization, matching + how is_key_within_model_budget resolves the key's own entries.""" + with patch.object( + budget_limiter, "_get_team_spend_for_model", AsyncMock(return_value=10_000.0) + ) as mock_team_spend: + assert ( + await budget_limiter.is_team_within_model_budget( + team_id="team-1", + team_model_max_budget={ + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + model="openai/gpt-4", + key_model_max_budget={ + "gpt-4": {"budget_limit": 50.0, "time_period": "1d"} + }, + ) + is True + ) + mock_team_spend.assert_not_called() + + +@pytest.mark.asyncio +async def test_team_spend_cache_key_format(budget_limiter): + """The shared counter must be keyed by team_id (not key hash) so every key on + the team reads and writes the same window.""" + from litellm.proxy.hooks.model_max_budget_limiter import ( + TEAM_MODEL_SPEND_CACHE_KEY_PREFIX, + ) + + budget_cfg = GenericBudgetInfo(budget_limit=10.0, time_period="1d") + captured = {} + + async def fake_get(key): + captured["key"] = key + return 0.0 + + with patch.object( + budget_limiter.dual_cache, "async_get_cache", side_effect=fake_get + ): + await budget_limiter._get_team_spend_for_model( + team_id="team-xyz", model="gpt-4", budget_config=budget_cfg + ) + + assert captured["key"] == f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:team-xyz:gpt-4:1d" + + +@pytest.mark.asyncio +async def test_async_log_success_event_increments_team_spend(budget_limiter): + """Successful calls under a team cap must drive the team-scoped counter.""" + from litellm.proxy.hooks.model_max_budget_limiter import ( + TEAM_MODEL_SPEND_CACHE_KEY_PREFIX, + ) + + incremented = {} + + async def fake_increment(budget_config, spend_key, start_time_key, response_cost): + incremented[spend_key] = response_cost + + kwargs = { + "standard_logging_object": { + "response_cost": 7.5, + "model_group": "gpt-4", + "model": "openai/gpt-4", + "metadata": {"user_api_key_hash": "hash-1"}, + }, + "litellm_params": { + "metadata": { + "user_api_key_team_id": "team-tracking", + "user_api_key_team_model_max_budget": { + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + } + }, + } + + with patch.object( + budget_limiter, "_increment_spend_for_key", side_effect=fake_increment + ): + await budget_limiter.async_log_success_event(kwargs, None, 0, 0) + + assert ( + incremented.get(f"{TEAM_MODEL_SPEND_CACHE_KEY_PREFIX}:team-tracking:gpt-4:1d") + == 7.5 + ) + + +@pytest.mark.asyncio +async def test_async_log_success_event_skips_team_increment_when_key_covers_model( + budget_limiter, +): + """Keys with their own entry must not drive the shared team counter (mirror of + the enforcement-side precedence), while their key counter still increments.""" + from litellm.proxy.hooks.model_max_budget_limiter import ( + TEAM_MODEL_SPEND_CACHE_KEY_PREFIX, + VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX, + ) + + incremented = {} + + async def fake_increment(budget_config, spend_key, start_time_key, response_cost): + incremented[spend_key] = response_cost + + kwargs = { + "standard_logging_object": { + "response_cost": 5.0, + "model_group": "gpt-4", + "model": "openai/gpt-4", + "metadata": {"user_api_key_hash": "hash-1"}, + }, + "litellm_params": { + "metadata": { + "user_api_key_team_id": "team-1", + "user_api_key_team_model_max_budget": { + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + "user_api_key_model_max_budget": { + "gpt-4": {"budget_limit": 50.0, "time_period": "1d"} + }, + } + }, + } + + with patch.object( + budget_limiter, "_increment_spend_for_key", side_effect=fake_increment + ): + await budget_limiter.async_log_success_event(kwargs, None, 0, 0) + + assert not any( + k.startswith(TEAM_MODEL_SPEND_CACHE_KEY_PREFIX) for k in incremented + ) + assert any(k.startswith(VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX) for k in incremented) + + +@pytest.mark.asyncio +async def test_async_log_success_event_skips_team_without_team_id(budget_limiter): + """No team_id in request metadata means no team counter write, even when a + team budget dict is present.""" + from litellm.proxy.hooks.model_max_budget_limiter import ( + TEAM_MODEL_SPEND_CACHE_KEY_PREFIX, + ) + + incremented = {} + + async def fake_increment(budget_config, spend_key, start_time_key, response_cost): + incremented[spend_key] = response_cost + + kwargs = { + "standard_logging_object": { + "response_cost": 3.0, + "model_group": "gpt-4", + "model": "openai/gpt-4", + "metadata": {"user_api_key_hash": "hash-1"}, + }, + "litellm_params": { + "metadata": { + "user_api_key_team_model_max_budget": { + "gpt-4": {"budget_limit": 100.0, "time_period": "1d"} + }, + } + }, + } + + with patch.object( + budget_limiter, "_increment_spend_for_key", side_effect=fake_increment + ): + await budget_limiter.async_log_success_event(kwargs, None, 0, 0) + + assert not any( + k.startswith(TEAM_MODEL_SPEND_CACHE_KEY_PREFIX) for k in incremented + ) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 5f3b0f36b95..c31d337609d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -42,6 +42,7 @@ from litellm.proxy.auth.auth_checks import ( _log_budget_lookup_failure, _tag_max_budget_check, _team_max_budget_check, + _team_model_max_budget_check, _virtual_key_max_budget_alert_check, _virtual_key_max_budget_check, _virtual_key_soft_budget_check, @@ -3093,6 +3094,81 @@ async def test_team_budget_check_reads_from_spend_counter(): assert exc_info.value.entity_id == "test-team" +@pytest.mark.asyncio +async def test_team_model_max_budget_check_delegates_to_limiter(): + """The auth bridge must pass the team's per-model budgets and the key's own + model_max_budget (for precedence) through to the limiter, once per request model.""" + team_object = LiteLLM_TeamTable( + team_id="team-mmb", + model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + valid_token = UserAPIKeyAuth( + token="test-token", + team_id="team-mmb", + model_max_budget={"claude-3": {"budget_limit": 5.0, "time_period": "1d"}}, + ) + mock_limiter = AsyncMock() + + with patch("litellm.proxy.proxy_server.model_max_budget_limiter", mock_limiter): + await _team_model_max_budget_check( + team_object=team_object, + valid_token=valid_token, + model="gpt-4o", + ) + + mock_limiter.is_team_within_model_budget.assert_awaited_once_with( + team_id="team-mmb", + team_model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}}, + model="gpt-4o", + key_model_max_budget={"claude-3": {"budget_limit": 5.0, "time_period": "1d"}}, + ) + + +@pytest.mark.asyncio +async def test_team_model_max_budget_check_noop_without_budgets(): + """Teams without model_max_budget (the default {}) must not touch the limiter.""" + team_object = LiteLLM_TeamTable(team_id="team-plain") + mock_limiter = AsyncMock() + + with patch("litellm.proxy.proxy_server.model_max_budget_limiter", mock_limiter): + await _team_model_max_budget_check( + team_object=team_object, + valid_token=UserAPIKeyAuth(token="t", team_id="team-plain"), + model="gpt-4o", + ) + await _team_model_max_budget_check( + team_object=None, + valid_token=None, + model="gpt-4o", + ) + + mock_limiter.is_team_within_model_budget.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_team_model_max_budget_check_propagates_budget_error(): + """BudgetExceededError from the limiter must escape the bridge so common_checks + rejects the request.""" + team_object = LiteLLM_TeamTable( + team_id="team-mmb", + model_max_budget={"gpt-4o": {"budget_limit": 10.0, "time_period": "1d"}}, + ) + mock_limiter = AsyncMock() + mock_limiter.is_team_within_model_budget.side_effect = litellm.BudgetExceededError( + current_cost=11.0, + max_budget=10.0, + message="over", + ) + + with patch("litellm.proxy.proxy_server.model_max_budget_limiter", mock_limiter): + with pytest.raises(litellm.BudgetExceededError): + await _team_model_max_budget_check( + team_object=team_object, + valid_token=None, + model=["gpt-4o"], + ) + + @pytest.mark.asyncio async def test_end_user_budget_check_reads_from_spend_counter(): """End-user budget check should use get_current_spend when counter exists.""" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 109d638fb9c..756ea8bcc1d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -13735,6 +13735,7 @@ export interface paths { * - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget * - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. + * - model_max_budget: Optional[dict] - Per-model max budgets shared by every key on the team, e.g. {"gpt-4o": {"budget_limit": 100.0, "time_period": "1d"}}. A key's own model_max_budget entry for a model takes precedence; the team entry is the default cap for keys without one. * - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) * - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. * - blocked: bool - Flag indicating if the team is blocked or not - will stop all calls from keys with this team_id. @@ -13913,6 +13914,7 @@ export interface paths { * - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget * - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. + * - model_max_budget: Optional[dict] - Per-model max budgets shared by every key on the team, e.g. {"gpt-4o": {"budget_limit": 100.0, "time_period": "1d"}}. A key's own model_max_budget entry for a model takes precedence; the team entry is the default cap for keys without one. * - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) * - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed. * - prompts: Optional[List[str]] - List of prompts that the team is allowed to use. @@ -23578,7 +23580,7 @@ export interface components { * @description Default role assigned to new users created * @default internal_user_viewer */ - user_role: ("internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer") | null; + user_role: ("proxy_admin" | "proxy_admin_viewer" | "internal_user" | "internal_user_viewer") | null; }; /** * DefaultTeamSSOParams @@ -26172,7 +26174,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: number | string | null; + stream_timeout?: string | number | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */ @@ -28488,6 +28490,12 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** Model Max Budget */ + model_max_budget?: { + [key: string]: { + [key: string]: string | number; + }; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -29230,6 +29238,12 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** Model Max Budget */ + model_max_budget?: { + [key: string]: { + [key: string]: string | number; + }; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -32977,6 +32991,12 @@ export interface components { model_aliases?: { [key: string]: unknown; } | null; + /** Model Max Budget */ + model_max_budget?: { + [key: string]: { + [key: string]: string | number; + }; + } | null; /** Model Rpm Limit */ model_rpm_limit?: { [key: string]: number; @@ -33597,6 +33617,12 @@ export interface components { team_model_aliases?: { [key: string]: unknown; } | null; + /** Team Model Max Budget */ + team_model_max_budget?: { + [key: string]: { + [key: string]: string | number; + }; + } | null; /** * Team Models * @default [] @@ -34304,7 +34330,7 @@ export interface components { [key: string]: unknown; } | null; /** Stream Timeout */ - stream_timeout?: number | string | null; + stream_timeout?: string | number | null; /** Tag Regex */ tag_regex?: string[] | null; /** Tags */