mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +00:00
fix: show budget reset date in BudgetExceededError message
This commit is contained in:
parent
aff0d1a18c
commit
731263b32e
10 changed files with 33 additions and 8 deletions
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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 {}
|
||||
|
@ -440,7 +443,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
|
||||
|
@ -1365,7 +1370,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}.",
|
||||
)
|
||||
|
||||
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -1035,6 +1035,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
|
||||
|
@ -1283,6 +1284,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
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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))
|
||||
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue