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

fix: check all models in list against per-model budget, not just model[0]

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-24 18:37:28 -07:00
parent 138e182459
commit 1a9edf3a6c
No known key found for this signature in database
1 changed files with 29 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3367,40 +3367,41 @@ async def _check_team_member_model_budget(
if not team_member_model_max_budget:
return
# resolve a single model string (skip list checks — model was already validated above)
model_str = model if isinstance(model, str) else (model[0] if model else None)
if model_str is None:
return
model_budget_config = team_member_model_max_budget.get(model_str)
if model_budget_config is None:
return
max_budget = (
model_budget_config.get("max_budget")
if isinstance(model_budget_config, dict)
else None
)
if max_budget is None:
# Build the list of candidate models to check
models_to_check: List[str] = [model] if isinstance(model, str) else list(model)
if not models_to_check:
return
from litellm.proxy.proxy_server import get_current_spend
model_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}:model:{model_str}",
fallback_spend=0.0,
)
for model_str in models_to_check:
model_budget_config = team_member_model_max_budget.get(model_str)
if model_budget_config is None:
continue
if model_spend >= max_budget:
raise litellm.BudgetExceededError(
current_cost=model_spend,
max_budget=max_budget,
message=(
f"ExceededBudget: Team member model budget exceeded. "
f"User={valid_token.user_id}, Team={team_object.team_id}, "
f"Model={model_str}. Spend=${model_spend:.6f}, Budget=${max_budget:.6f}"
),
max_budget = (
model_budget_config.get("max_budget")
if isinstance(model_budget_config, dict)
else None
)
if max_budget is None:
continue
model_spend = await get_current_spend(
counter_key=f"spend:team_member:{valid_token.user_id}:{team_object.team_id}:model:{model_str}",
fallback_spend=0.0,
)
if model_spend >= max_budget:
raise litellm.BudgetExceededError(
current_cost=model_spend,
max_budget=max_budget,
message=(
f"ExceededBudget: Team member model budget exceeded. "
f"User={valid_token.user_id}, Team={team_object.team_id}, "
f"Model={model_str}. Spend=${model_spend:.6f}, Budget=${max_budget:.6f}"
),
)
async def _check_team_member_model_access(