diff --git a/litellm/exceptions.py b/litellm/exceptions.py index 0abe3e33af..dd9953a321 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -730,10 +730,15 @@ LITELLM_EXCEPTION_TYPES = [ class BudgetExceededError(Exception): - def __init__(self, current_cost, max_budget): + def __init__( + self, current_cost: float, max_budget: float, message: Optional[str] = None + ): self.current_cost = current_cost self.max_budget = max_budget - message = f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" + message = ( + message + or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" + ) self.message = message super().__init__(message) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 950e457f71..6ab63d8cb2 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -88,8 +88,10 @@ def common_checks( and team_object.spend is not None and team_object.spend > team_object.max_budget ): - raise Exception( - f"Team={team_object.team_id} over budget. Spend={team_object.spend}, Budget={team_object.max_budget}" + raise litellm.BudgetExceededError( + current_cost=team_object.spend, + max_budget=team_object.max_budget, + message=f"Team={team_object.team_id} over budget. Spend={team_object.spend}, Budget={team_object.max_budget}", ) # 4. If user is in budget ## 4.1 check personal budget, if personal key @@ -100,16 +102,20 @@ def common_checks( ): user_budget = user_object.max_budget if user_budget < user_object.spend: - raise Exception( - f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}" + raise litellm.BudgetExceededError( + current_cost=user_object.spend, + max_budget=user_budget, + message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}", ) ## 4.2 check team member budget, if team key # 5. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget if end_user_object is not None and end_user_object.litellm_budget_table is not None: end_user_budget = end_user_object.litellm_budget_table.max_budget if end_user_budget is not None and end_user_object.spend > end_user_budget: - raise Exception( - f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}" + raise litellm.BudgetExceededError( + current_cost=end_user_object.spend, + max_budget=end_user_budget, + message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}", ) # 6. [OPTIONAL] If 'enforce_user_param' enabled - did developer pass in 'user' param for openai endpoints if ( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 3e03888526..115c90bb6a 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -552,7 +552,6 @@ async def user_api_key_auth( key=api_key ) if valid_token is None: - user_obj: Optional[LiteLLM_UserTable] = None ## check db verbose_proxy_logger.debug("api key: %s", api_key) if prisma_client is not None: @@ -584,6 +583,7 @@ async def user_api_key_auth( user_id_information: Optional[List] = None if valid_token is not None: + user_obj: Optional[LiteLLM_UserTable] = None # Got Valid Token from Cache, DB # Run checks for # 1. If token can call model diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index d210b1d23e..12204ec066 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -504,7 +504,7 @@ def test_call_with_user_over_budget(prisma_client): asyncio.run(test()) except Exception as e: error_detail = e.message - assert "Budget has been exceeded" in error_detail + assert "ExceededBudget:" in error_detail assert isinstance(e, ProxyException) assert e.type == ProxyErrorTypes.budget_exceeded print(vars(e)) @@ -607,7 +607,7 @@ def test_call_with_end_user_over_budget(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print("result from user auth with new key", result) - pytest.fail(f"This should have failed!. They key crossed it's budget") + pytest.fail("This should have failed!. They key crossed it's budget") asyncio.run(test()) except Exception as e: @@ -779,12 +779,12 @@ def test_call_with_user_over_budget_stream(prisma_client): # use generated key to auth in result = await user_api_key_auth(request=request, api_key=bearer_token) print("result from user auth with new key", result) - pytest.fail(f"This should have failed!. They key crossed it's budget") + pytest.fail("This should have failed!. They key crossed it's budget") asyncio.run(test()) except Exception as e: error_detail = e.message - assert "Budget has been exceeded" in error_detail + assert "ExceededBudget:" in error_detail assert isinstance(e, ProxyException) assert e.type == ProxyErrorTypes.budget_exceeded print(vars(e)) @@ -2511,7 +2511,6 @@ async def test_update_user_role(prisma_client): Tests if we update user role, incorrect values are not stored in cache -> create a user with role == INTERNAL_USER -> access an Admin only route -> expect to fail - -> update user role to == PROXY_ADMIN -> access an Admin only route -> expect to succeed """