fix(proxy_server.py): use consistent 400-status code error code for exceeded budget errors

standardizes error code for budget exceeded errors to status code 400
This commit is contained in:
Krrish Dholakia 2024-06-11 16:10:58 -07:00
parent 8e6905a5cf
commit 2a5a30fa7a
4 changed files with 19 additions and 16 deletions

View file

@ -1202,8 +1202,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
@ -1239,8 +1240,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
@ -1264,8 +1266,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