scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

feat(team): enforce team-level model_max_budget with key-level overrides

Browse source 
Teams can now carry per-model max budgets that every key on the team
inherits by default. The LiteLLM_TeamTable.model_max_budget column already
existed but was inert; /team/new and /team/update now accept and validate
it, auth hydrates it onto UserAPIKeyAuth.team_model_max_budget from the
cached team object, and a new _team_model_max_budget_check in common_checks
enforces it via is_team_within_model_budget on the shared limiter. Spend is
tracked in a team-scoped counter (team_model_spend:{team_id}:{model}:{duration})
shared by all keys on the team. A key with its own model_max_budget entry
covering the request model takes precedence: the team counter is neither
consulted nor incremented for that key and model.
This commit is contained in:
ryan-crabbe-berri 2026-08-01 16:56:00 -07:00
parent 4a3819a602
commit 7a9a1a0d45
10 changed files with 626 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

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

39
litellm/proxy/auth/auth_checks.py
View file

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

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

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

168
litellm/proxy/hooks/model_max_budget_limiter.py
View file

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

1
litellm/proxy/litellm_pre_call_utils.py
View file

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

2
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

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

32
litellm/proxy/management_endpoints/team_endpoints.py
View file

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

278
tests/proxy_unit_tests/test_unit_test_max_model_budget_limiter.py
View file

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

76
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

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

32
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

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