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

fix: more bug fixes

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-11 17:41:59 -07:00
parent c27ae34a39
commit a1ca70a8a7
1 changed files with 10 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

16
litellm/proxy/proxy_server.py
View file

@ -1075,8 +1075,9 @@ async def user_api_key_auth(
_user_id = _user.get("user_id", None) _user_id = _user.get("user_id", None)
if user_current_spend > user_max_budget: if user_current_spend > user_max_budget:
raise Exception( raise litellm.BudgetExceededError(
f"ExceededBudget: User {_user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}" current_cost=user_current_spend,
max_budget=user_max_budget,
) )
else: else:
# Token exists, not expired now check if its in budget for the user # Token exists, not expired now check if its in budget for the user
@ -1107,9 +1108,11 @@ async def user_api_key_auth(
) )
if user_current_spend > user_max_budget: if user_current_spend > user_max_budget:
raise Exception( raise litellm.BudgetExceededError(
f"ExceededBudget: User {valid_token.user_id} has exceeded their budget. Current spend: {user_current_spend}; Max Budget: {user_max_budget}" current_cost=user_current_spend,
max_budget=user_max_budget,
) )
# Check 3. Check if user is in their team budget # Check 3. Check if user is in their team budget
if valid_token.team_member_spend is not None: if valid_token.team_member_spend is not None:
if prisma_client is not None: if prisma_client is not None:
@ -1143,8 +1146,9 @@ async def user_api_key_auth(
) )
if team_member_budget is not None and team_member_budget > 0: if team_member_budget is not None and team_member_budget > 0:
if valid_token.team_member_spend > team_member_budget: if valid_token.team_member_spend > team_member_budget:
raise Exception( raise litellm.BudgetExceededError(
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}" current_cost=valid_token.team_member_spend,
max_budget=team_member_budget,
) )
# Check 3. If token is expired # Check 3. If token is expired