mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): send Retry-After on budget exceeded denials that have a reset time
This commit is contained in:
Denis Ćorić
parent
572832dea1
commit
1ff93ccdf7
10 changed files with 208 additions and 11 deletions
|
|
@ -1125,6 +1125,7 @@ class BudgetExceededError(Exception):
|
|||
llm_provider: str | None = None,
|
||||
entity_type: str | None = None,
|
||||
entity_id: str | None = None,
|
||||
budget_reset_at: datetime.datetime | None = None,
|
||||
):
|
||||
self.current_cost = current_cost
|
||||
self.max_budget = max_budget
|
||||
|
|
@ -1132,6 +1133,7 @@ class BudgetExceededError(Exception):
|
|||
self.llm_provider = llm_provider or ""
|
||||
self.entity_type = entity_type
|
||||
self.entity_id = entity_id
|
||||
self.budget_reset_at = budget_reset_at
|
||||
# Surface unified rate-limit fields without joining the RateLimitError
|
||||
# hierarchy so existing `except BudgetExceededError:` handlers keep
|
||||
# working; custom callbacks reading StandardLoggingPayload pick these
|
||||
|
|
|
|||
|
|
@ -70,3 +70,10 @@ class LiteLLM_TeamMemberTable(LiteLLM_BudgetTable):
|
|||
budget_id: str | None = None
|
||||
|
||||
model_config = ConfigDict(protected_namespaces=())
|
||||
|
||||
|
||||
def budget_table_reset_at(budget_table: LiteLLM_BudgetTable | None) -> datetime | None:
|
||||
"""The budget's next reset, known only when the row was loaded with it."""
|
||||
if isinstance(budget_table, LiteLLM_BudgetTableFull):
|
||||
return budget_table.budget_reset_at
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -51,6 +51,7 @@ from litellm.llms.base_llm.managed_resources.utils import (
|
|||
extract_provider_resource_id_from_unified_id,
|
||||
is_base64_encoded_unified_id,
|
||||
)
|
||||
from litellm.models.budget import budget_table_reset_at
|
||||
from litellm.models.project import LiteLLM_ProjectTable
|
||||
from litellm.proxy._types import (
|
||||
MODEL_LIST_ROUTES,
|
||||
|
|
@ -146,7 +147,7 @@ from litellm.proxy.guardrails.tool_name_extraction import (
|
|||
extract_request_tool_names,
|
||||
)
|
||||
from litellm.proxy.route_llm_request import route_request
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_reset_at, get_budget_window_start
|
||||
from litellm.proxy.spend_tracking.carried_budget_state import carry_organization_budget_state
|
||||
from litellm.proxy.spend_tracking.spend_counter_batch import SPEND_COUNTERS_TARGET
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics
|
||||
|
|
@ -1441,6 +1442,7 @@ async def common_checks(
|
|||
message=f"ExceededBudget: User={user_object.user_id} over budget. Spend={user_spend}, Budget={user_budget}",
|
||||
entity_type=Litellm_EntityType.USER.value,
|
||||
entity_id=user_object.user_id,
|
||||
budget_reset_at=user_object.budget_reset_at,
|
||||
)
|
||||
|
||||
# Each scope reads a distinct counter key with no cross-scope ordering
|
||||
|
|
@ -1977,6 +1979,7 @@ async def check_end_user_budget(
|
|||
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
|
||||
entity_type=Litellm_EntityType.END_USER.value,
|
||||
entity_id=end_user_obj.user_id,
|
||||
budget_reset_at=budget_table_reset_at(end_user_obj.litellm_budget_table),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -6837,6 +6840,7 @@ async def virtual_key_max_budget_check(
|
|||
message=f"Budget has been exceeded! Key={key_descriptor} Current cost: {spend}, Max budget: {valid_token.max_budget}",
|
||||
entity_type=Litellm_EntityType.KEY.value,
|
||||
entity_id=valid_token.token,
|
||||
budget_reset_at=valid_token.budget_reset_at,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -6884,6 +6888,7 @@ async def _virtual_key_multi_budget_check(
|
|||
),
|
||||
entity_type=Litellm_EntityType.KEY.value,
|
||||
entity_id=valid_token.token,
|
||||
budget_reset_at=get_budget_window_reset_at(window), # pyright: ignore[reportUnknownArgumentType] # key budget_limits rows are untyped dicts
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7198,6 +7203,11 @@ async def _check_team_member_budget(
|
|||
message=f"Budget has been exceeded! User={valid_token.user_id} in Team={team_object.team_id} Current cost: {team_member_spend}, Max budget: {team_member_budget}",
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
|
||||
entity_id=f"{valid_token.user_id}:{team_object.team_id}",
|
||||
budget_reset_at=budget_table_reset_at(
|
||||
member_budget_row
|
||||
if member_budget_row is not None and member_budget_row.max_budget is not None
|
||||
else team_default_budget
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7313,6 +7323,7 @@ async def team_max_budget_check(
|
|||
message=f"Budget has been exceeded! Team={team_object.team_id} Current cost: {spend}, Max budget: {team_object.max_budget}",
|
||||
entity_type=Litellm_EntityType.TEAM.value,
|
||||
entity_id=team_object.team_id,
|
||||
budget_reset_at=team_object.budget_reset_at,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7356,6 +7367,7 @@ async def _team_multi_budget_check(
|
|||
),
|
||||
entity_type=Litellm_EntityType.TEAM.value,
|
||||
entity_id=team_object.team_id,
|
||||
budget_reset_at=get_budget_window_reset_at(window),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7485,6 +7497,7 @@ async def _project_max_budget_check(
|
|||
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}",
|
||||
entity_type=Litellm_EntityType.PROJECT.value,
|
||||
entity_id=project_object.project_id,
|
||||
budget_reset_at=budget_table_reset_at(project_object.litellm_budget_table),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7693,6 +7706,7 @@ async def organization_max_budget_check(
|
|||
message=f"Budget has been exceeded! Organization={org_id} Current cost: {org_spend}, Max budget: {org_max_budget}",
|
||||
entity_type=Litellm_EntityType.ORGANIZATION.value,
|
||||
entity_id=org_id,
|
||||
budget_reset_at=budget_table_reset_at(org_table.litellm_budget_table),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -7765,6 +7779,7 @@ async def tag_max_budget_check_for_tags(
|
|||
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
|
||||
entity_type=Litellm_EntityType.TAG.value,
|
||||
entity_id=tag_name,
|
||||
budget_reset_at=budget_table_reset_at(tag_object.litellm_budget_table),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3,7 +3,9 @@ Handles Authentication Errors
|
|||
"""
|
||||
|
||||
import logging
|
||||
import math
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from fastapi import HTTPException, Request, status
|
||||
|
|
@ -48,6 +50,19 @@ else:
|
|||
Span = Any
|
||||
|
||||
|
||||
def _utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
def budget_retry_after_seconds(budget_reset_at: datetime, now: datetime) -> int:
|
||||
"""Whole seconds until the budget resets, never below 1 so a client does not retry at once."""
|
||||
reset_at: Final = (
|
||||
budget_reset_at if budget_reset_at.tzinfo is not None else budget_reset_at.replace(tzinfo=timezone.utc)
|
||||
)
|
||||
current: Final = now if now.tzinfo is not None else now.replace(tzinfo=timezone.utc)
|
||||
return max(1, math.ceil((reset_at - current).total_seconds()))
|
||||
|
||||
|
||||
def _as_proxy_exception(e: Exception) -> ProxyException:
|
||||
"""Convert an authentication failure into the ProxyException the client receives."""
|
||||
if isinstance(e, litellm.BudgetExceededError):
|
||||
|
|
@ -56,6 +71,11 @@ def _as_proxy_exception(e: Exception) -> ProxyException:
|
|||
type=ProxyErrorTypes.budget_exceeded,
|
||||
param=None,
|
||||
code=getattr(e, "status_code", status.HTTP_429_TOO_MANY_REQUESTS),
|
||||
headers=(
|
||||
{"retry-after": str(budget_retry_after_seconds(e.budget_reset_at, _utc_now()))}
|
||||
if e.budget_reset_at is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
if isinstance(e, ModelAccessDeniedHTTPException):
|
||||
return ModelAccessDeniedProxyException(
|
||||
|
|
|
|||
|
|
@ -41,6 +41,7 @@ from litellm.integrations.otel.runtime import phase_event, phase_span, seed_requ
|
|||
from litellm.litellm_core_utils.dd_tracing import tracer
|
||||
from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value
|
||||
from litellm.litellm_core_utils.routing_common_utils import resolve_model_group_alias, resolve_served_model
|
||||
from litellm.models.budget import budget_table_reset_at
|
||||
from litellm.models.end_user import LiteLLM_EndUserTable
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers
|
||||
|
|
@ -3085,6 +3086,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
|
|||
),
|
||||
entity_type=Litellm_EntityType.TEAM_MEMBER.value,
|
||||
entity_id=_entity_id,
|
||||
budget_reset_at=budget_table_reset_at(team_member_info.litellm_budget_table),
|
||||
)
|
||||
|
||||
# Check 3. If token is expired
|
||||
|
|
|
|||
|
|
@ -67,6 +67,7 @@ class _BudgetCounter:
|
|||
spend_log_entity_id: str | None = None
|
||||
window_duration: str | None = None
|
||||
window_start: datetime | None = None
|
||||
budget_reset_at: datetime | None = None
|
||||
|
||||
|
||||
_COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = {
|
||||
|
|
@ -234,6 +235,7 @@ def _raise_counter_budget_exceeded(
|
|||
),
|
||||
entity_type=_COUNTER_ENTITY_TYPES.get(counter.entity_type),
|
||||
entity_id=counter.spend_log_entity_id or counter.entity_id,
|
||||
budget_reset_at=counter.budget_reset_at,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -495,6 +497,7 @@ async def _get_budget_counters(
|
|||
fallback_spend=float(valid_token.spend or 0.0),
|
||||
entity_type="Key",
|
||||
entity_id=valid_token.token,
|
||||
budget_reset_at=valid_token.budget_reset_at,
|
||||
)
|
||||
)
|
||||
counters.extend(
|
||||
|
|
@ -518,6 +521,7 @@ async def _get_budget_counters(
|
|||
fallback_spend=float(team_object.spend or 0.0),
|
||||
entity_type="Team",
|
||||
entity_id=team_id,
|
||||
budget_reset_at=team_object.budget_reset_at,
|
||||
)
|
||||
)
|
||||
counters.extend(
|
||||
|
|
@ -546,6 +550,7 @@ async def _get_budget_counters(
|
|||
fallback_spend=float(user_object.spend or 0.0),
|
||||
entity_type="User",
|
||||
entity_id=user_object.user_id,
|
||||
budget_reset_at=user_object.budget_reset_at,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -629,6 +634,7 @@ async def _get_end_user_budget_counter(
|
|||
fallback_spend=fallback_spend,
|
||||
entity_type="EndUser",
|
||||
entity_id=end_user_id,
|
||||
budget_reset_at=_reset_at_of(_get_value(end_user_object, "litellm_budget_table")),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -669,6 +675,7 @@ async def _get_tag_budget_counters(
|
|||
fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0,
|
||||
entity_type="Tag",
|
||||
entity_id=tag_name,
|
||||
budget_reset_at=_reset_at_of(budget_table),
|
||||
)
|
||||
)
|
||||
return counters
|
||||
|
|
@ -725,6 +732,17 @@ def _dedupe_tags(tags: list[str]) -> list[str]:
|
|||
return deduped_tags
|
||||
|
||||
|
||||
async def _cached_team_member_default_budget(
|
||||
team_object: LiteLLM_TeamTable, user_api_key_cache: UserApiKeyCache
|
||||
) -> object:
|
||||
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
|
||||
if not isinstance(default_budget_id, str):
|
||||
return None
|
||||
return _CACHED_VALUE.validate_python(
|
||||
await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}")
|
||||
)
|
||||
|
||||
|
||||
@with_service_target(AUTH_OBJECTS_TARGET)
|
||||
async def _get_team_member_budget_counter(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
|
|
@ -749,20 +767,20 @@ async def _get_team_member_budget_counter(
|
|||
|
||||
member_budget_row: Final = team_membership.litellm_budget_table if team_membership is not None else None
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
member_row_sets_cap: Final = member_budget_row is not None and member_budget_row.max_budget is not None
|
||||
default_budget: Final = (
|
||||
None if member_row_sets_cap else await _cached_team_member_default_budget(team_object, user_api_key_cache)
|
||||
)
|
||||
team_member_reset_at: Final = _reset_at_of(member_budget_row if member_row_sets_cap else default_budget)
|
||||
team_member_budget: float | None = None
|
||||
if member_budget_row is not None and member_budget_row.max_budget is not None:
|
||||
team_member_budget = member_budget_row.effective_max_budget(now=now)
|
||||
else:
|
||||
default_budget_id: Final = (team_object.metadata or {}).get("team_member_budget_id")
|
||||
if isinstance(default_budget_id, str):
|
||||
default_budget: Final = _CACHED_VALUE.validate_python(
|
||||
await user_api_key_cache.async_get_cache(key=f"team_member_default_budget:{default_budget_id}")
|
||||
default_cap: Final = _to_float(_get_value(default_budget, "max_budget"))
|
||||
if default_cap is not None and default_cap > 0:
|
||||
team_member_budget = default_cap + (
|
||||
member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
|
||||
)
|
||||
default_cap: Final = _to_float(_get_value(default_budget, "max_budget"))
|
||||
if default_cap is not None and default_cap > 0:
|
||||
team_member_budget = default_cap + (
|
||||
member_budget_row.active_temp_budget_increase(now=now) if member_budget_row is not None else 0.0
|
||||
)
|
||||
|
||||
if team_member_budget is None or team_member_budget <= 0:
|
||||
return None
|
||||
|
|
@ -775,6 +793,7 @@ async def _get_team_member_budget_counter(
|
|||
fallback_spend=float(team_member_spend or 0.0),
|
||||
entity_type="TeamMember",
|
||||
entity_id=f"{valid_token.user_id}:{team_object.team_id}",
|
||||
budget_reset_at=team_member_reset_at,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -814,6 +833,7 @@ async def _get_org_budget_counter(
|
|||
fallback_spend=org_spend,
|
||||
entity_type="Organization",
|
||||
entity_id=org_id,
|
||||
budget_reset_at=_reset_at_of(org_budget_table),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -847,6 +867,7 @@ async def _get_project_budget_counter(
|
|||
fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0,
|
||||
entity_type="Project",
|
||||
entity_id=valid_token.project_id,
|
||||
budget_reset_at=_reset_at_of(project_budget_table),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -886,6 +907,7 @@ def _get_budget_limit_counters(
|
|||
spend_log_entity_id=entity_id,
|
||||
window_duration=str(budget_duration),
|
||||
window_start=window_start,
|
||||
budget_reset_at=get_budget_window_reset_at(window_dict),
|
||||
)
|
||||
)
|
||||
return counters
|
||||
|
|
@ -1282,6 +1304,17 @@ def get_budget_window_start(window: object) -> datetime | None:
|
|||
return reset_at - timedelta(seconds=duration_seconds)
|
||||
|
||||
|
||||
def get_budget_window_reset_at(window: object) -> datetime | None:
|
||||
reset_at: Final = _coerce_datetime(_coerce_window(window).get("reset_at"))
|
||||
if reset_at is None:
|
||||
return None
|
||||
return reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _reset_at_of(budget_table: object) -> datetime | None:
|
||||
return _coerce_datetime(_get_value(budget_table, "budget_reset_at"))
|
||||
|
||||
|
||||
def _coerce_datetime(value: object) -> datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -14,6 +14,7 @@ from litellm.models.budget import (
|
|||
LiteLLM_BudgetTable,
|
||||
LiteLLM_BudgetTableFull,
|
||||
LiteLLM_TeamMemberTable,
|
||||
budget_table_reset_at,
|
||||
)
|
||||
from litellm.models.config import LiteLLM_Config
|
||||
from litellm.models.credentials import CreateCredentialItem, CredentialItem
|
||||
|
|
@ -690,3 +691,12 @@ class TestAutoRouterSession:
|
|||
|
||||
def test_a_partial_comparison_against_one_baseline_keeps_its_label(self) -> None:
|
||||
assert self._row({"anthropic/claude-opus-5": 3}, estimated_turns=1).baseline_model == "anthropic/claude-opus-5"
|
||||
|
||||
|
||||
def test_budget_table_reset_at_is_known_only_on_the_full_row():
|
||||
reset_at = datetime(2026, 11, 3, tzinfo=timezone.utc)
|
||||
full_row = LiteLLM_BudgetTableFull(max_budget=1.0, budget_reset_at=reset_at, created_at=reset_at)
|
||||
|
||||
assert budget_table_reset_at(full_row) == reset_at
|
||||
assert budget_table_reset_at(LiteLLM_BudgetTable(max_budget=1.0)) is None
|
||||
assert budget_table_reset_at(None) is None
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
|
||||
import sys, os, asyncio, time, random, uuid
|
||||
import hashlib
|
||||
from datetime import datetime, timezone
|
||||
import json
|
||||
import traceback
|
||||
from dotenv import load_dotenv
|
||||
|
|
@ -991,6 +992,25 @@ async def test_virtual_key_max_budget_check(
|
|||
assert alert_called, "Budget alert should be triggered"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_virtual_key_max_budget_check_carries_the_budget_reset_at():
|
||||
from litellm.proxy.auth.auth_checks import virtual_key_max_budget_check
|
||||
|
||||
reset_at: Final = datetime(2026, 11, 3, tzinfo=timezone.utc)
|
||||
valid_token = UserAPIKeyAuth(token="test-token", spend=10.0, max_budget=10.0, budget_reset_at=reset_at)
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=None)
|
||||
|
||||
async def mock_budget_alert(*args, **kwargs):
|
||||
return None
|
||||
|
||||
proxy_logging_obj.budget_alerts = mock_budget_alert
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await virtual_key_max_budget_check(valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, user_obj=None)
|
||||
|
||||
assert exc_info.value.budget_reset_at == reset_at
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model, team_models, expect_to_work",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
|
@ -38,7 +39,12 @@ from litellm.proxy._types import (
|
|||
ProxyException,
|
||||
UserAPIKeyAuth,
|
||||
)
|
||||
from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler, _as_proxy_exception
|
||||
from litellm.proxy.auth import auth_exception_handler
|
||||
from litellm.proxy.auth.auth_exception_handler import (
|
||||
UserAPIKeyAuthExceptionHandler,
|
||||
_as_proxy_exception,
|
||||
budget_retry_after_seconds,
|
||||
)
|
||||
from litellm.proxy.auth.model_access_denied import ModelAccessDeniedHTTPException
|
||||
|
||||
|
||||
|
|
@ -1312,3 +1318,47 @@ def test_auth_failure_error_body_carries_a_json_null_param(auth_failure: Excepti
|
|||
|
||||
assert body["error"]["type"] == ProxyErrorTypes.auth_error.value
|
||||
assert body["error"]["param"] is None
|
||||
|
||||
|
||||
FIXED_NOW: Final = datetime(2026, 11, 2, 12, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"budget_reset_at, expected",
|
||||
[
|
||||
(FIXED_NOW + timedelta(seconds=90.4), 91),
|
||||
(FIXED_NOW + timedelta(seconds=60), 60),
|
||||
(FIXED_NOW - timedelta(seconds=5), 1),
|
||||
(datetime(2026, 11, 2, 12, 0, 30), 30),
|
||||
],
|
||||
)
|
||||
def test_budget_retry_after_seconds(budget_reset_at: datetime, expected: int):
|
||||
assert budget_retry_after_seconds(budget_reset_at, FIXED_NOW) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"budget_reset_at, expected_header",
|
||||
[
|
||||
(FIXED_NOW + timedelta(days=1), "86400"),
|
||||
(FIXED_NOW - timedelta(minutes=1), "1"),
|
||||
(None, None),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_authentication_error_budget_exceeded_retry_after(
|
||||
monkeypatch: pytest.MonkeyPatch, budget_reset_at: datetime | None, expected_header: str | None
|
||||
):
|
||||
monkeypatch.setattr(auth_exception_handler, "_utc_now", lambda: FIXED_NOW)
|
||||
handler = UserAPIKeyAuthExceptionHandler()
|
||||
budget_error = litellm.BudgetExceededError(
|
||||
message="Budget exceeded",
|
||||
current_cost=100,
|
||||
max_budget=100,
|
||||
budget_reset_at=budget_reset_at,
|
||||
)
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await handler.handle_authentication_error(budget_error, MagicMock(), {}, "/test", None, "test-key")
|
||||
|
||||
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
|
||||
assert exc_info.value.headers.get("retry-after") == expected_header
|
||||
|
|
|
|||
|
|
@ -24,7 +24,10 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
_BudgetCounter,
|
||||
_get_team_member_budget_counter,
|
||||
_raise_counter_budget_exceeded,
|
||||
get_budget_window_reset_at,
|
||||
estimate_request_max_cost,
|
||||
get_budget_window_start,
|
||||
release_unbound_budget_reservation,
|
||||
|
|
@ -588,3 +591,38 @@ async def test_group_reserves_on_its_current_deployments_after_another_vendors_d
|
|||
by_own_name: Final = litellm.token_counter(model="prod", messages=_MULTILINGUAL_MESSAGES)
|
||||
assert openai_only != by_own_name
|
||||
assert (await _input_tokens(router, "prod"))["prod"] == by_own_name
|
||||
|
||||
|
||||
def test_counter_budget_exceeded_carries_the_counter_reset_at():
|
||||
reset_at: Final = datetime(2026, 11, 3, tzinfo=timezone.utc)
|
||||
counter = _BudgetCounter(
|
||||
counter_key="spend:key:k",
|
||||
max_budget=1.0,
|
||||
fallback_spend=1.0,
|
||||
entity_type="Key",
|
||||
entity_id="k",
|
||||
budget_reset_at=reset_at,
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
_raise_counter_budget_exceeded(counter=counter, current_cost=1.5)
|
||||
|
||||
assert exc_info.value.budget_reset_at == reset_at
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"window, expected",
|
||||
[
|
||||
(
|
||||
{"budget_duration": "1d", "max_budget": 1.0, "reset_at": "2026-11-03T00:00:00Z"},
|
||||
datetime(2026, 11, 3, tzinfo=timezone.utc),
|
||||
),
|
||||
(
|
||||
{"budget_duration": "1d", "max_budget": 1.0, "reset_at": datetime(2026, 11, 3)},
|
||||
datetime(2026, 11, 3, tzinfo=timezone.utc),
|
||||
),
|
||||
({"budget_duration": "1d", "max_budget": 1.0}, None),
|
||||
],
|
||||
)
|
||||
def test_get_budget_window_reset_at(window: dict[str, object], expected: datetime | None):
|
||||
assert get_budget_window_reset_at(window) == expected
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue