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 c715878ad..801e28c48 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 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")