Merge 731263b32e into b82af5b826

docs/my-website/docs/proxy/team_budgets.md

@@ -116,7 +116,7 @@ On the 2nd response - expect to see the following exception
 ```shell
 {
  "error": {
-   "message": "Budget has been exceeded! Current cost: 3.5e-06, Max budget: 1e-09",
+   "message": "Budget has been exceeded! Current cost: 3.5e-06, Max budget: 1e-09, Budget resets at: 2025-01-01T00:00:00Z",
    "type": "auth_error",
    "param": null,
    "code": 400

docs/my-website/docs/proxy/users.md

@@ -462,7 +462,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
 
 Error
 ```shell
-{"error":{"message":"Budget has been exceeded: User ishaan3 has exceeded their budget. Current spend: 0.0008869999999999999; Max Budget: 0.0001","type":"auth_error","param":"None","code":401}}%                
+{"error":{"message":"Budget has been exceeded: User ishaan3 has exceeded their budget. Current spend: 0.0008869999999999999, Max Budget: 0.0001, Budget resets at: 2025-01-01T00:00:00Z","type":"auth_error","param":"None","code":400}}
 ```
 
 ## Reset Budgets 

litellm/exceptions.py

@@ -9,6 +9,7 @@
 
 ## LiteLLM versions of the OpenAI Exception Types
 
+from datetime import datetime
 from typing import Optional
 
 import httpx
@@ -740,13 +741,18 @@ LITELLM_EXCEPTION_TYPES = [
 
 class BudgetExceededError(Exception):
     def __init__(
-        self, current_cost: float, max_budget: float, message: Optional[str] = None
+        self,
+        current_cost: float,
+        max_budget: float,
+        budget_reset_at: datetime | None,
+        message: Optional[str] = None,
     ):
         self.current_cost = current_cost
         self.max_budget = max_budget
+        self.budget_reset_at = budget_reset_at
         message = (
             message
-            or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}"
+            or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}, Budget resets at: {str(budget_reset_at)}."
         )
         self.message = message
         super().__init__(message)

litellm/proxy/auth/auth_checks.py

@@ -138,6 +138,7 @@ async def common_checks(
             raise litellm.BudgetExceededError(
                 current_cost=user_object.spend,
                 max_budget=user_budget,
+                budget_reset_at=user_object.budget_reset_at,
                 message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_object.spend}, Budget={user_budget}",
             )
 
@@ -173,7 +174,9 @@ async def common_checks(
     ):
         if global_proxy_spend > litellm.max_budget:
             raise litellm.BudgetExceededError(
-                current_cost=global_proxy_spend, max_budget=litellm.max_budget
+                current_cost=global_proxy_spend,
+                max_budget=litellm.max_budget,
+                budget_reset_at=getattr(litellm, "budget_reset_at", None),
             )
 
     _request_metadata: dict = request_body.get("metadata", {}) or {}
@@ -445,7 +448,9 @@ async def get_end_user_object(
         end_user_budget = end_user_obj.litellm_budget_table.max_budget
         if end_user_budget is not None and end_user_obj.spend > end_user_budget:
             raise litellm.BudgetExceededError(
-                current_cost=end_user_obj.spend, max_budget=end_user_budget
+                current_cost=end_user_obj.spend,
+                max_budget=end_user_budget,
+                budget_reset_at=end_user_obj.litellm_budget_table.budget_reset_at,
             )
 
     # check if in cache
@@ -1372,7 +1377,8 @@ async def _team_max_budget_check(
         raise litellm.BudgetExceededError(
             current_cost=team_object.spend,
             max_budget=team_object.max_budget,
-            message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}",
+            budget_reset_at=team_object.budget_reset_at,
+            message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {team_object.spend}, Max budget: {team_object.max_budget}, Budget resets at: {team_object.budget_reset_at}.",
         )
 
 

litellm/proxy/auth/user_api_key_auth.py

@@ -832,6 +832,7 @@ async def _user_api_key_auth_builder(  # noqa: PLR0915
                                 raise litellm.BudgetExceededError(
                                     current_cost=valid_token.team_member_spend,
                                     max_budget=team_member_budget,
+                                    budget_reset_at=valid_token.budget_reset_at,
                                 )
 
             # Check 3. If token is expired

litellm/proxy/hooks/model_max_budget_limiter.py

@@ -79,6 +79,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
                     message=f"LiteLLM Virtual Key: {user_api_key_dict.token}, key_alias: {user_api_key_dict.key_alias}, exceeded budget for model={model}",
                     current_cost=_current_spend,
                     max_budget=_current_model_budget_info.max_budget,
+                    budget_reset_at=_current_model_budget_info.budget_reset_at,
                 )
 
         return True

litellm/utils.py

@@ -1043,6 +1043,7 @@ def client(original_function):  # noqa: PLR0915
                     raise BudgetExceededError(
                         current_cost=litellm._current_cost,
                         max_budget=litellm.max_budget,
+                        budget_reset_at=litellm.budget_reset_at,
                     )
 
             # [OPTIONAL] CHECK MAX RETRIES / REQUEST
@@ -1291,6 +1292,7 @@ def client(original_function):  # noqa: PLR0915
                     raise BudgetExceededError(
                         current_cost=litellm._current_cost,
                         max_budget=litellm.max_budget,
+                        budget_reset_at=litellm.budget_reset_at,
                     )
 
             # [OPTIONAL] CHECK CACHE

tests/litellm/proxy/auth/test_auth_exception_handler.py

@@ -2,6 +2,7 @@ import asyncio
 import json
 import os
 import sys
+from datetime import datetime
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -98,7 +99,7 @@ async def test_handle_authentication_error_budget_exceeded():
         from litellm.exceptions import BudgetExceededError
 
         budget_error = BudgetExceededError(
-            message="Budget exceeded", current_cost=100, max_budget=100
+            current_cost=100, max_budget=100, budget_reset_at=datetime.fromisoformat("2025-01-01T00:00:00Z")
         )
         await handler._handle_authentication_error(
             budget_error,
@@ -110,6 +111,7 @@ async def test_handle_authentication_error_budget_exceeded():
         )
 
     assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
+    assert "Budget has been exceeded! Current cost: 100, Max budget: 100, Budget resets at: 2025-01-01 00:00:00+00:00." == str(exc_info.value.message)
 
 
 @pytest.mark.asyncio

tests/litellm/proxy/db/test_exception_handler.py

@@ -143,6 +143,7 @@ def test_handle_db_exception_with_non_db_error():
     regular_error = litellm.BudgetExceededError(
         current_cost=10,
         max_budget=10,
+        budget_reset_at="2025-01-01T00:00:00Z",
     )
     with pytest.raises(litellm.BudgetExceededError):
         PrismaDBExceptionHandler.handle_db_exception(regular_error)

tests/proxy_unit_tests/test_key_generate_prisma.py

@@ -661,6 +661,7 @@ def test_call_with_end_user_over_budget(prisma_client):
         print(f"raised error: {e}, traceback: {traceback.format_exc()}")
         error_detail = e.message
         assert "Budget has been exceeded! Current" in error_detail
+        assert "Budget resets at:" in error_detail
         assert isinstance(e, ProxyException)
         assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
@@ -759,6 +760,7 @@ def test_call_with_proxy_over_budget(prisma_client):
         else:
             error_detail = traceback.format_exc()
         assert "Budget has been exceeded" in error_detail
+        assert "Budget resets at:" in error_detail
         assert isinstance(e, ProxyException)
         assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
@@ -953,6 +955,7 @@ def test_call_with_proxy_over_budget_stream(prisma_client):
     except Exception as e:
         error_detail = e.message
         assert "Budget has been exceeded" in error_detail
+        assert "Budget resets at:" in error_detail
         print(vars(e))
 
 
@@ -1598,6 +1601,7 @@ def test_call_with_key_over_budget(prisma_client):
         else:
             error_detail = str(e)
         assert "Budget has been exceeded" in error_detail
+        assert "Budget resets at:" in error_detail
         assert isinstance(e, ProxyException)
         assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
@@ -1722,6 +1726,7 @@ def test_call_with_key_over_budget_no_cache(prisma_client):
         else:
             error_detail = str(e)
         assert "Budget has been exceeded" in error_detail
+        assert "Budget resets at:" in error_detail
         assert isinstance(e, ProxyException)
         assert e.type == ProxyErrorTypes.budget_exceeded
         print(vars(e))
@@ -2016,6 +2021,7 @@ async def test_call_with_key_over_budget_stream(prisma_client):
         print("Got Exception", e)
         error_detail = e.message
         assert "Budget has been exceeded" in error_detail
+        assert "Budget resets at:" in error_detail
 
         print(vars(e))