phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

Merge pull request #4139 from BerriAI/litellm_fix_budget_exceeded_error_code

Browse source 
fix(proxy_server.py): use consistent 400-status code error code for exceeded budget errors
This commit is contained in:
Krish Dholakia 2024-06-11 18:36:58 -07:00 • committed by GitHub
commit 75b0d7821b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 30 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

31
litellm/proxy/proxy_server.py
View file

@ -1078,8 +1078,9 @@ async def user_api_key_auth(
_user_id = _user.get("user_id", None)
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {_user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
raise litellm.BudgetExceededError(
current_cost=user_current_spend,
max_budget=user_max_budget,
)
else:
# Token exists, not expired now check if its in budget for the user
@ -1110,9 +1111,11 @@ async def user_api_key_auth(
)
if user_current_spend > user_max_budget:
raise Exception(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}"
raise litellm.BudgetExceededError(
current_cost=user_current_spend,
max_budget=user_max_budget,
)
# Check 3. Check if user is in their team budget
if valid_token.team_member_spend is not None:
if prisma_client is not None:
@ -1146,8 +1149,9 @@ async def user_api_key_auth(
)
if team_member_budget is not None and team_member_budget > 0:
if valid_token.team_member_spend > team_member_budget:
raise Exception(
f"ExceededBudget: Crossed spend within team. UserID: {valid_token.user_id}, in team {valid_token.team_id} has exceeded their budget. Current spend: {valid_token.team_member_spend}; Max Budget: {team_member_budget}"
raise litellm.BudgetExceededError(
current_cost=valid_token.team_member_spend,
max_budget=team_member_budget,
)
# Check 3. If token is expired
@ -1205,8 +1209,9 @@ async def user_api_key_auth(
####################################
if valid_token.spend >= valid_token.max_budget:
raise Exception(
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Token: {valid_token.max_budget}"
raise litellm.BudgetExceededError(
current_cost=valid_token.spend,
max_budget=valid_token.max_budget,
)
# Check 5. Token Model Spend is under Model budget
@ -1242,8 +1247,9 @@ async def user_api_key_auth(
):
current_model_spend = model_spend[0]["_sum"]["spend"]
current_model_budget = max_budget_per_model[current_model]
raise Exception(
f"ExceededModelBudget: Current spend for model: {current_model_spend}; Max Budget for Model: {current_model_budget}"
raise litellm.BudgetExceededError(
current_cost=current_model_spend,
max_budget=current_model_budget,
)
# Check 6. Team spend is under Team budget
@ -1267,8 +1273,9 @@ async def user_api_key_auth(
)
if valid_token.team_spend >= valid_token.team_max_budget:
raise Exception(
f"ExceededTokenBudget: Current Team Spend: {valid_token.team_spend}; Max Budget for Team: {valid_token.team_max_budget}"
raise litellm.BudgetExceededError(
current_cost=valid_token.team_spend,
max_budget=valid_token.team_max_budget,
)
# Check 8: Additional Common Checks across jwt + key auth