diff --git a/litellm/main.py b/litellm/main.py index bf4931168..8133e3517 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -15,7 +15,6 @@ from functools import partial import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import httpx - import litellm from ._logging import verbose_logger from litellm import ( # type: ignore diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 17f0822e6..e32d56706 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -151,8 +151,8 @@ def common_checks( and route != "/models" ): if global_proxy_spend > litellm.max_budget: - raise Exception( - f"ExceededBudget: LiteLLM Proxy has exceeded its budget. Current spend: {global_proxy_spend}; Max Budget: {litellm.max_budget}" + raise litellm.BudgetExceededError( + current_cost=global_proxy_spend, max_budget=litellm.max_budget ) return True diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f76bf2257..a03b02d95 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 277d9718f..8a02c4c7a 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -471,7 +471,7 @@ def test_call_with_user_over_budget(prisma_client): asyncio.run(test()) except Exception as e: error_detail = e.message - assert "Authentication Error, ExceededBudget:" in error_detail + assert "Budget has been exceeded" in error_detail print(vars(e)) @@ -652,7 +652,7 @@ def test_call_with_proxy_over_budget(prisma_client): error_detail = e.message else: error_detail = traceback.format_exc() - assert "Authentication Error, ExceededBudget:" in error_detail + assert "Budget has been exceeded" in error_detail print(vars(e)) @@ -730,7 +730,7 @@ def test_call_with_user_over_budget_stream(prisma_client): asyncio.run(test()) except Exception as e: error_detail = e.message - assert "Authentication Error, ExceededBudget:" in error_detail + assert "Budget has been exceeded" in error_detail print(vars(e)) @@ -827,7 +827,7 @@ def test_call_with_proxy_over_budget_stream(prisma_client): asyncio.run(test()) except Exception as e: error_detail = e.message - assert "Authentication Error, ExceededBudget:" in error_detail + assert "Budget has been exceeded" in error_detail print(vars(e)) @@ -1362,7 +1362,7 @@ def test_call_with_key_over_budget(prisma_client): error_detail = e.message else: error_detail = str(e) - assert "Authentication Error, ExceededTokenBudget:" in error_detail + assert "Budget has been exceeded" in error_detail print(vars(e)) @@ -1476,7 +1476,7 @@ def test_call_with_key_over_model_budget(prisma_client): # print(f"Error - {str(e)}") traceback.print_exc() error_detail = e.message - assert "Authentication Error, ExceededModelBudget:" in error_detail + assert "Budget has been exceeded!" in error_detail print(vars(e)) @@ -1637,7 +1637,7 @@ async def test_call_with_key_over_budget_stream(prisma_client): except Exception as e: print("Got Exception", e) error_detail = e.message - assert "Authentication Error, ExceededTokenBudget:" in error_detail + assert "Budget has been exceeded" in error_detail print(vars(e)) diff --git a/tests/test_keys.py b/tests/test_keys.py index bbe4cc35e..877c239bc 100644 --- a/tests/test_keys.py +++ b/tests/test_keys.py @@ -664,7 +664,7 @@ async def test_key_crossing_budget(): response = await chat_completion(session=session, key=key) pytest.fail("Should have failed - Key crossed it's budget") except Exception as e: - assert "ExceededTokenBudget: Current spend for token:" in str(e) + assert "Budget has been exceeded!" in str(e) @pytest.mark.skip(reason="AWS Suspended Account") diff --git a/tests/test_team.py b/tests/test_team.py index 467767be0..544273c2e 100644 --- a/tests/test_team.py +++ b/tests/test_team.py @@ -578,4 +578,4 @@ async def test_users_in_team_budget(): except Exception as e: print("got exception, this is expected") print(e) - assert "Crossed spend within team" in str(e) + assert "Budget has been exceeded" in str(e)