mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 03:04:13 +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
|
```shell
|
||||||
{
|
{
|
||||||
"error": {
|
"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",
|
"type": "auth_error",
|
||||||
"param": null,
|
"param": null,
|
||||||
"code": 400
|
"code": 400
|
||||||
|
|
|
@ -462,7 +462,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
|
||||||
|
|
||||||
Error
|
Error
|
||||||
```shell
|
```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
|
## Reset Budgets
|
||||||
|
|
|
@ -9,6 +9,7 @@
|
||||||
|
|
||||||
## LiteLLM versions of the OpenAI Exception Types
|
## LiteLLM versions of the OpenAI Exception Types
|
||||||
|
|
||||||
|
from datetime import datetime
|
||||||
from typing import Optional
|
from typing import Optional
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
@ -740,13 +741,18 @@ LITELLM_EXCEPTION_TYPES = [
|
||||||
|
|
||||||
class BudgetExceededError(Exception):
|
class BudgetExceededError(Exception):
|
||||||
def __init__(
|
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.current_cost = current_cost
|
||||||
self.max_budget = max_budget
|
self.max_budget = max_budget
|
||||||
|
self.budget_reset_at = budget_reset_at
|
||||||
message = (
|
message = (
|
||||||
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
|
self.message = message
|
||||||
super().__init__(message)
|
super().__init__(message)
|
||||||
|
|
|
@ -138,6 +138,7 @@ async def common_checks(
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=user_object.spend,
|
current_cost=user_object.spend,
|
||||||
max_budget=user_budget,
|
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}",
|
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:
|
if global_proxy_spend > litellm.max_budget:
|
||||||
raise litellm.BudgetExceededError(
|
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 {}
|
_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
|
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:
|
if end_user_budget is not None and end_user_obj.spend > end_user_budget:
|
||||||
raise litellm.BudgetExceededError(
|
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
|
# check if in cache
|
||||||
|
@ -1365,7 +1370,8 @@ async def _team_max_budget_check(
|
||||||
raise litellm.BudgetExceededError(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=team_object.spend,
|
current_cost=team_object.spend,
|
||||||
max_budget=team_object.max_budget,
|
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(
|
raise litellm.BudgetExceededError(
|
||||||
current_cost=valid_token.team_member_spend,
|
current_cost=valid_token.team_member_spend,
|
||||||
max_budget=team_member_budget,
|
max_budget=team_member_budget,
|
||||||
|
budget_reset_at=valid_token.budget_reset_at,
|
||||||
)
|
)
|
||||||
|
|
||||||
# Check 3. If token is expired
|
# 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}",
|
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,
|
current_cost=_current_spend,
|
||||||
max_budget=_current_model_budget_info.max_budget,
|
max_budget=_current_model_budget_info.max_budget,
|
||||||
|
budget_reset_at=_current_model_budget_info.budget_reset_at,
|
||||||
)
|
)
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
|
@ -1035,6 +1035,7 @@ def client(original_function): # noqa: PLR0915
|
||||||
raise BudgetExceededError(
|
raise BudgetExceededError(
|
||||||
current_cost=litellm._current_cost,
|
current_cost=litellm._current_cost,
|
||||||
max_budget=litellm.max_budget,
|
max_budget=litellm.max_budget,
|
||||||
|
budget_reset_at=litellm.budget_reset_at,
|
||||||
)
|
)
|
||||||
|
|
||||||
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
|
# [OPTIONAL] CHECK MAX RETRIES / REQUEST
|
||||||
|
@ -1283,6 +1284,7 @@ def client(original_function): # noqa: PLR0915
|
||||||
raise BudgetExceededError(
|
raise BudgetExceededError(
|
||||||
current_cost=litellm._current_cost,
|
current_cost=litellm._current_cost,
|
||||||
max_budget=litellm.max_budget,
|
max_budget=litellm.max_budget,
|
||||||
|
budget_reset_at=litellm.budget_reset_at,
|
||||||
)
|
)
|
||||||
|
|
||||||
# [OPTIONAL] CHECK CACHE
|
# [OPTIONAL] CHECK CACHE
|
||||||
|
|
|
@ -2,6 +2,7 @@ import asyncio
|
||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
|
from datetime import datetime
|
||||||
from unittest.mock import AsyncMock, MagicMock, patch
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
@ -98,7 +99,7 @@ async def test_handle_authentication_error_budget_exceeded():
|
||||||
from litellm.exceptions import BudgetExceededError
|
from litellm.exceptions import BudgetExceededError
|
||||||
|
|
||||||
budget_error = 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(
|
await handler._handle_authentication_error(
|
||||||
budget_error,
|
budget_error,
|
||||||
|
@ -110,6 +111,7 @@ async def test_handle_authentication_error_budget_exceeded():
|
||||||
)
|
)
|
||||||
|
|
||||||
assert exc_info.value.type == ProxyErrorTypes.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
|
@pytest.mark.asyncio
|
||||||
|
|
|
@ -143,6 +143,7 @@ def test_handle_db_exception_with_non_db_error():
|
||||||
regular_error = litellm.BudgetExceededError(
|
regular_error = litellm.BudgetExceededError(
|
||||||
current_cost=10,
|
current_cost=10,
|
||||||
max_budget=10,
|
max_budget=10,
|
||||||
|
budget_reset_at="2025-01-01T00:00:00Z",
|
||||||
)
|
)
|
||||||
with pytest.raises(litellm.BudgetExceededError):
|
with pytest.raises(litellm.BudgetExceededError):
|
||||||
PrismaDBExceptionHandler.handle_db_exception(regular_error)
|
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()}")
|
print(f"raised error: {e}, traceback: {traceback.format_exc()}")
|
||||||
error_detail = e.message
|
error_detail = e.message
|
||||||
assert "Budget has been exceeded! Current" in error_detail
|
assert "Budget has been exceeded! Current" in error_detail
|
||||||
|
assert "Budget resets at:" in error_detail
|
||||||
assert isinstance(e, ProxyException)
|
assert isinstance(e, ProxyException)
|
||||||
assert e.type == ProxyErrorTypes.budget_exceeded
|
assert e.type == ProxyErrorTypes.budget_exceeded
|
||||||
print(vars(e))
|
print(vars(e))
|
||||||
|
@ -759,6 +760,7 @@ def test_call_with_proxy_over_budget(prisma_client):
|
||||||
else:
|
else:
|
||||||
error_detail = traceback.format_exc()
|
error_detail = traceback.format_exc()
|
||||||
assert "Budget has been exceeded" in error_detail
|
assert "Budget has been exceeded" in error_detail
|
||||||
|
assert "Budget resets at:" in error_detail
|
||||||
assert isinstance(e, ProxyException)
|
assert isinstance(e, ProxyException)
|
||||||
assert e.type == ProxyErrorTypes.budget_exceeded
|
assert e.type == ProxyErrorTypes.budget_exceeded
|
||||||
print(vars(e))
|
print(vars(e))
|
||||||
|
@ -953,6 +955,7 @@ def test_call_with_proxy_over_budget_stream(prisma_client):
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
error_detail = e.message
|
error_detail = e.message
|
||||||
assert "Budget has been exceeded" in error_detail
|
assert "Budget has been exceeded" in error_detail
|
||||||
|
assert "Budget resets at:" in error_detail
|
||||||
print(vars(e))
|
print(vars(e))
|
||||||
|
|
||||||
|
|
||||||
|
@ -1598,6 +1601,7 @@ def test_call_with_key_over_budget(prisma_client):
|
||||||
else:
|
else:
|
||||||
error_detail = str(e)
|
error_detail = str(e)
|
||||||
assert "Budget has been exceeded" in error_detail
|
assert "Budget has been exceeded" in error_detail
|
||||||
|
assert "Budget resets at:" in error_detail
|
||||||
assert isinstance(e, ProxyException)
|
assert isinstance(e, ProxyException)
|
||||||
assert e.type == ProxyErrorTypes.budget_exceeded
|
assert e.type == ProxyErrorTypes.budget_exceeded
|
||||||
print(vars(e))
|
print(vars(e))
|
||||||
|
@ -1722,6 +1726,7 @@ def test_call_with_key_over_budget_no_cache(prisma_client):
|
||||||
else:
|
else:
|
||||||
error_detail = str(e)
|
error_detail = str(e)
|
||||||
assert "Budget has been exceeded" in error_detail
|
assert "Budget has been exceeded" in error_detail
|
||||||
|
assert "Budget resets at:" in error_detail
|
||||||
assert isinstance(e, ProxyException)
|
assert isinstance(e, ProxyException)
|
||||||
assert e.type == ProxyErrorTypes.budget_exceeded
|
assert e.type == ProxyErrorTypes.budget_exceeded
|
||||||
print(vars(e))
|
print(vars(e))
|
||||||
|
@ -2016,6 +2021,7 @@ async def test_call_with_key_over_budget_stream(prisma_client):
|
||||||
print("Got Exception", e)
|
print("Got Exception", e)
|
||||||
error_detail = e.message
|
error_detail = e.message
|
||||||
assert "Budget has been exceeded" in error_detail
|
assert "Budget has been exceeded" in error_detail
|
||||||
|
assert "Budget resets at:" in error_detail
|
||||||
|
|
||||||
print(vars(e))
|
print(vars(e))
|
||||||
|
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue