scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-14 04:09:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 97

Merge fcdd92c2bc into c0ca4a51f7

Browse source
This commit is contained in:
Thijmen Stavenuiter 2026-10-11 15:54:11 -04:00 • committed by GitHub
commit fcb0b3a61d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 399 additions and 61 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
litellm/litellm_core_utils/duration_parser.py
View file

@ -74,38 +74,31 @@ def duration_in_seconds(duration: str) -> int:
elif unit == "w":
return value * 604800
elif unit == "mo":
now: Final = time_module.time()
current_time: Final = datetime.fromtimestamp(now)
# Calculate target month and year, handling overflow past December
total_months: Final = current_time.month - 1 + value # 0-indexed months
target_year: Final = current_time.year + total_months // 12
target_month: Final = total_months % 12 + 1 # back to 1-indexed
# Determine the day to set for next month
target_day = current_time.day
last_day_of_target_month: Final = get_last_day_of_month(target_year, target_month)
target_day = min(target_day, last_day_of_target_month)
next_month: Final = datetime(
year=target_year,
month=target_month,
day=target_day,
hour=current_time.hour,
minute=current_time.minute,
second=current_time.second,
microsecond=current_time.microsecond,
return _months_in_seconds(
start=datetime.fromtimestamp(time_module.time(), tz=timezone.utc).astimezone(), months=value
)
# Calculate the duration until the first day of the next month
duration_until_next_month: Final = next_month - current_time
return int(duration_until_next_month.total_seconds())
else:
raise ValueError(f"Unsupported duration unit, passed duration: {duration}")
def duration_in_seconds_from(start: datetime, duration: str) -> int:
"""How long a period of `duration` lasts when it begins at `start`.
`duration_in_seconds` measures `Nmo` from now, so a monthly window that opened in a
month of a different length needs its own start to know when it ends. Every other
unit is a fixed length and ignores `start`.
"""
value, unit = _extract_from_regex(duration=_normalize_duration(duration))
if unit == "mo":
return _months_in_seconds(start=start, months=value)
return duration_in_seconds(duration)
def _months_in_seconds(start: datetime, months: int) -> int:
return int((_shift_months(start, months) - start).total_seconds())
def get_next_standardized_reset_time(
duration: str,
current_time: datetime,
@ -166,8 +159,9 @@ def get_next_standardized_reset_time(
return base_midnight + timedelta(days=1)
def _subtract_months(moment: datetime, months: int) -> datetime:
total_months: Final = moment.year * 12 + moment.month - 1 - months
def _shift_months(moment: datetime, months: int) -> datetime:
"""`moment` moved by `months` calendar months, back when negative, clamped to the target month's last day."""
total_months: Final = moment.year * 12 + moment.month - 1 + months
year, month_index = divmod(total_months, 12)
month: Final = month_index + 1
return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month)))
@ -182,9 +176,9 @@ def get_budget_window_start(duration: str, reset_at: datetime) -> datetime:
return reset_at - timedelta(days=1)
match unit:
case "mo":
return _subtract_months(reset_at, value)
return _shift_months(reset_at, -value)
case "d":
return _subtract_months(reset_at, 1) if value == 30 else reset_at - timedelta(days=value)
return _shift_months(reset_at, -1) if value == 30 else reset_at - timedelta(days=value)
case "w":
return reset_at - timedelta(weeks=value)
case "h":

98
litellm/proxy/hooks/model_max_budget_limiter.py
View file

@ -3,6 +3,7 @@ import json
import time
from collections.abc import Iterable, Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timezone
from types import MappingProxyType
from typing import Final, cast # noqa: TID251 # untyped callback payload requires explicit narrowing
@ -13,7 +14,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import Span
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.litellm_core_utils.duration_parser import duration_in_seconds, duration_in_seconds_from
from litellm.llms.bedrock.common_utils import get_bedrock_base_model
from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
@ -205,17 +206,28 @@ def _bedrock_candidates(model: str) -> tuple[str, ...]:
return (base_model, without_vendor) if without_vendor else (base_model,)
@dataclass(frozen=True, slots=True)
class _WindowReading:
current_spend: float
started_at: float | None
async def build_model_max_budget_usage(
entity_type: Litellm_EntityType,
entity_id: str | None,
model_max_budget: Mapping[str, object] | None,
cache: DualCache | None,
now: datetime | None = None,
) -> dict[str, dict[str, object]]:
"""Current-window spend per configured budget model, as `/key/info` and `/user/info` report it.
`cache` must be the DualCache the limiter writes the counters to; callers
read it off the limiter rather than re-deriving it, so a scope that is being
blocked can never report zero usage.
`budget_reset_at` is when the open window ends and its spend counter
expires, which is the window's start plus its period. It is None when no
window is open, because nothing has been charged since the last one ended.
"""
if cache is None or entity_id is None or not model_max_budget:
return {}
@ -228,41 +240,82 @@ async def build_model_max_budget_usage(
)
if not budgets:
return {}
spend_keys: Final = tuple(
model_budget_spend_cache_key(
entity_type=entity_type,
entity_id=entity_id,
budget_model=budget_model,
budget_duration=budget_config.budget_duration,
)
for budget_model, budget_config in budgets
windows: Final = await _read_windows(
cache=cache,
spend_keys=tuple(
model_budget_spend_cache_key(
entity_type=entity_type,
entity_id=entity_id,
budget_model=budget_model,
budget_duration=budget_config.budget_duration,
)
for budget_model, budget_config in budgets
),
start_time_keys=tuple(
model_budget_start_time_cache_key(
entity_type=entity_type,
entity_id=entity_id,
budget_model=budget_model,
budget_duration=budget_config.budget_duration,
)
for budget_model, budget_config in budgets
),
)
current_spends: Final = await _current_window_spends(cache=cache, spend_keys=spend_keys)
current_time: Final = (now or datetime.now(timezone.utc)).timestamp()
return {
budget_model: {
"current_spend": round(current_spend, 4),
"current_spend": round(window.current_spend, 4),
"budget_limit": budget_config.max_budget,
"time_period": budget_config.budget_duration,
"budget_reset_at": _window_reset_at(
started_at=window.started_at,
budget_duration=budget_config.budget_duration,
now=current_time,
),
}
for (budget_model, budget_config), current_spend in zip(budgets, current_spends, strict=True)
for (budget_model, budget_config), window in zip(budgets, windows, strict=True)
}
@with_service_target("model_budgets")
async def _current_window_spends(cache: DualCache, spend_keys: Sequence[str]) -> tuple[float, ...]:
"""Redis holds the window total across replicas; the in-memory copy is one replica's share."""
keys: Final = list(spend_keys)
async def _read_windows(
cache: DualCache, spend_keys: Sequence[str], start_time_keys: Sequence[str]
) -> tuple[_WindowReading, ...]:
"""Every window's spend and start time in one batched read, however many models are budgeted.
Redis holds the window total across replicas; the in-memory copy is one replica's share.
"""
keys: Final = [*spend_keys, *start_time_keys]
redis_cache: Final = cache.redis_cache
if redis_cache is not None:
shared: Final = await redis_cache.async_batch_get_cache(key_list=keys)
return tuple(_as_spend(shared.get(key)) for key in keys)
return _readings(tuple(_as_float(shared.get(key)) for key in keys), window_count=len(spend_keys))
# async_batch_get_cache returns None if it fails internally, and its result is
# index-aligned with `keys` otherwise. An unusable result reads as a miss,
# which is what a never-written counter already reads as.
batched: Final = await cache.async_batch_get_cache(keys=keys)
if not isinstance(batched, list) or len(batched) != len(keys):
return (0.0,) * len(keys)
return tuple(_as_spend(current_spend) for current_spend in batched)
return tuple(_WindowReading(current_spend=0.0, started_at=None) for _ in spend_keys)
return _readings(tuple(_as_float(value) for value in batched), window_count=len(spend_keys))
def _readings(values: Sequence[float | None], window_count: int) -> tuple[_WindowReading, ...]:
"""`values` holds every window's spend followed by every window's start, in the same order."""
return tuple(
_WindowReading(current_spend=spend or 0.0, started_at=started_at)
for spend, started_at in zip(values[:window_count], values[window_count:], strict=True)
)
def _window_reset_at(started_at: float | None, budget_duration: str | None, now: float) -> str | None:
if started_at is None or budget_duration is None:
return None
ends_at: Final = started_at + duration_in_seconds_from(
start=datetime.fromtimestamp(started_at, tz=timezone.utc).astimezone(), duration=budget_duration
)
if ends_at <= now:
return None
return datetime.fromtimestamp(ends_at, tz=timezone.utc).isoformat()
def _usable_budget_config(raw_budget_config: object) -> BudgetConfig | None:
@ -283,6 +336,15 @@ def _as_spend(current_spend: object) -> float:
return 0.0
def _as_float(value: object) -> float | None:
if value is None:
return None
try:
return float(value) # pyright: ignore[reportArgumentType] # non-numeric falls to the except
except (TypeError, ValueError):
return None
def _resolve_entity_model_budgets(
model: str,
entity_budgets: Iterable[tuple[Litellm_EntityType, str | None, object]],

3
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -4612,7 +4612,8 @@ async def info_key_fn(
hour), so subtracting budget_duration from it does not give the window's start
- model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
- model_max_budget_usage: dict | None - Current-window spend per model, present only when
the key has per-model budgets
the key has per-model budgets. Each entry's budget_reset_at is when that model's open
window ends and its spend resets, or None when no window is open
- budget_limits: list | None - Concurrent budget windows, exactly as stored
- budget_limits_usage: dict | None - Current-window spend per budget window, e.g.
{"1h": {"current_spend": 0.0009}}, present only when the key has budget windows

32
tests/unit/litellm_core_utils/test_duration_parser.py
View file

@ -4,9 +4,12 @@ from typing import Final
from unittest.mock import patch
from zoneinfo import ZoneInfo
import pytest
import litellm.litellm_core_utils.duration_parser as duration_parser
from litellm.litellm_core_utils.duration_parser import (
duration_in_seconds,
duration_in_seconds_from,
get_budget_window_start,
get_next_standardized_reset_time,
)
@ -373,5 +376,34 @@ class TestGetBudgetWindowStart(unittest.TestCase):
)
@pytest.mark.parametrize(
"start,duration,days",
[
(datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "1mo", 31),
(datetime(2026, 2, 10, 12, tzinfo=timezone.utc), "1mo", 28),
(datetime(2026, 1, 31, 12, tzinfo=timezone.utc), "1mo", 28),
(datetime(2028, 1, 31, 12, tzinfo=timezone.utc), "1mo", 29),
(datetime(2026, 12, 15, 12, tzinfo=timezone.utc), "1mo", 31),
(datetime(2026, 1, 15, 12, tzinfo=timezone.utc), "2mo", 59),
],
ids=["january", "february", "clamped_to_february", "leap_february", "across_new_year", "two_months"],
)
def test_a_month_lasts_as_long_as_the_month_it_starts_in(start, duration, days):
assert duration_in_seconds_from(start, duration) == days * 86400
@pytest.mark.parametrize("duration", ["45s", "30m", "1h", "1d", "7d", "2w", "30d", "daily", "monthly"])
def test_fixed_length_units_do_not_depend_on_the_start(duration):
assert duration_in_seconds_from(datetime(2026, 2, 10, 12, tzinfo=timezone.utc), duration) == duration_in_seconds(
duration
)
def test_a_month_starting_now_matches_duration_in_seconds():
now: Final = datetime(2026, 1, 31, 9, 30, 0, 250000, tzinfo=timezone.utc).astimezone()
with patch.object(duration_parser.time_module, "time", return_value=now.timestamp()):
assert duration_in_seconds("1mo") == duration_in_seconds_from(now, "1mo")
if __name__ == "__main__":
unittest.main()

268
tests/unit/proxy/hooks/test_unit_test_max_model_budget_limiter.py
View file

@ -8,7 +8,8 @@ import pytest
import litellm
from litellm.caching.caching import DualCache
from litellm.caching.redis_cache import RedisCache
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from typing import Final
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy._types import Litellm_EntityType
@ -16,6 +17,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import (
_budget_model_candidates,
PROXY_VirtualKeyModelMaxBudgetLimiter,
build_model_max_budget_usage,
model_budget_spend_cache_key,
model_budget_start_time_cache_key,
resolve_model_budget,
)
from litellm.proxy._types import UserAPIKeyAuth
@ -23,6 +26,17 @@ from litellm.types.caching import RedisPipelineIncrementOperation
from litellm.types.utils import BudgetConfig as GenericBudgetInfo
async def _reset_at_of_open_window(
cache: DualCache, entity_type: Litellm_EntityType, entity_id: str, budget_model: str, period: str
) -> str:
started_at = await cache.async_get_cache(
key=model_budget_start_time_cache_key(
entity_type=entity_type, entity_id=entity_id, budget_model=budget_model, budget_duration=period
)
)
return datetime.fromtimestamp(started_at + duration_in_seconds(period), tz=timezone.utc).isoformat()
# Test class setup
@pytest.fixture
def budget_limiter():
@ -680,6 +694,9 @@ async def test_logged_spend_is_visible_to_key_info_usage_and_enforcement(request
"current_spend": 0.75,
"budget_limit": 1.0,
"time_period": "1d",
"budget_reset_at": await _reset_at_of_open_window(
dual_cache, Litellm_EntityType.KEY, key_hash, "gpt-4", "1d"
),
}
}
@ -747,7 +764,16 @@ async def test_user_model_budget_is_tracked_and_enforced():
entity_id=user_id,
model_max_budget=user_model_max_budget,
cache=dual_cache,
) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}}
) == {
"gpt-4": {
"current_spend": 1.5,
"budget_limit": 1.0,
"time_period": "1mo",
"budget_reset_at": await _reset_at_of_open_window(
dual_cache, Litellm_EntityType.USER, user_id, "gpt-4", "1mo"
),
}
}
with pytest.raises(litellm.BudgetExceededError) as exc:
await limiter.is_user_within_model_budget(
@ -857,7 +883,9 @@ async def test_build_model_max_budget_usage_skips_unusable_entries():
},
cache=dual_cache,
)
assert usage == {"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d"}}
assert usage == {
"gpt-4": {"current_spend": 3.0, "budget_limit": 10.0, "time_period": "1d", "budget_reset_at": None}
}
@pytest.mark.asyncio
@ -896,6 +924,9 @@ async def test_bedrock_traffic_charges_the_bare_family_name_budget():
"current_spend": 1.5,
"budget_limit": 1.0,
"time_period": "18h",
"budget_reset_at": await _reset_at_of_open_window(
dual_cache, Litellm_EntityType.KEY, key_hash, "claude-opus-4-8", "18h"
),
}
}
@ -963,7 +994,16 @@ async def test_user_model_budget_window_resets_when_the_period_elapses():
entity_id=user_id,
model_max_budget=user_model_max_budget,
cache=dual_cache,
) == {"gpt-4": {"current_spend": 1.5, "budget_limit": 1.0, "time_period": "1mo"}}
) == {
"gpt-4": {
"current_spend": 1.5,
"budget_limit": 1.0,
"time_period": "1mo",
"budget_reset_at": await _reset_at_of_open_window(
dual_cache, Litellm_EntityType.USER, user_id, "gpt-4", "1mo"
),
}
}
@pytest.mark.asyncio
@ -995,7 +1035,7 @@ async def test_a_zero_dollar_cap_is_reported_as_a_cap_not_as_absent():
entity_id="hash-zero",
model_max_budget={"gpt-4": {"budget_limit": 0, "time_period": "1d"}},
cache=DualCache(),
) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d"}}
) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 0.0, "time_period": "1d", "budget_reset_at": None}}
@pytest.mark.asyncio
@ -1033,7 +1073,7 @@ async def test_usage_report_reads_every_counter_in_one_batched_lookup():
budget = {f"model-{i}": {"budget_limit": 1.0, "time_period": "1d"} for i in range(50)}
with (
patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 50)) as batched,
patch.object(dual_cache, "async_batch_get_cache", new=AsyncMock(return_value=[None] * 100)) as batched,
patch.object(dual_cache, "async_get_cache", new=AsyncMock()) as single,
):
usage = await build_model_max_budget_usage(
@ -1044,7 +1084,8 @@ async def test_usage_report_reads_every_counter_in_one_batched_lookup():
)
assert batched.await_count == 1
assert len(batched.await_args.kwargs["keys"]) == 50
# One spend counter and one window start per budgeted model, all in the same round trip.
assert len(batched.await_args.kwargs["keys"]) == 100
assert single.await_count == 0
assert len(usage) == 50
@ -1063,7 +1104,161 @@ async def test_usage_report_survives_a_batch_lookup_that_returns_nothing():
entity_id="hash-none",
model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1d"}},
cache=dual_cache,
) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d"}}
) == {"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None}}
_WINDOW_OPENED_AT: Final = datetime(2026, 10, 10, 12, 0, tzinfo=timezone.utc)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"seconds_since_window_opened,expected_reset_at",
[
(0, "2026-10-10T13:00:00+00:00"),
(600, "2026-10-10T13:00:00+00:00"),
(3599, "2026-10-10T13:00:00+00:00"),
(3600, None),
(7200, None),
],
ids=["just_opened", "mid_window", "last_second", "ended_exactly_now", "long_ended"],
)
async def test_usage_reports_when_the_open_window_resets(seconds_since_window_opened, expected_reset_at):
"""
The counter expires at the window's start plus its period, and that is when the
model is admitted again. A window that has already run its period is not open,
even if its start time is still sitting in the cache, so no reset is reported.
"""
dual_cache = DualCache()
key = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="gpt-4", budget_duration="1h")
await dual_cache.async_set_cache(key=model_budget_spend_cache_key(**key), value=0.4)
await dual_cache.async_set_cache(key=model_budget_start_time_cache_key(**key), value=_WINDOW_OPENED_AT.timestamp())
usage = await build_model_max_budget_usage(
entity_type=Litellm_EntityType.KEY,
entity_id="vk-hash",
model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1h"}},
cache=dual_cache,
now=_WINDOW_OPENED_AT + timedelta(seconds=seconds_since_window_opened),
)
assert usage == {
"gpt-4": {"current_spend": 0.4, "budget_limit": 1.0, "time_period": "1h", "budget_reset_at": expected_reset_at}
}
_DAY: Final = 86400
@pytest.mark.asyncio
@pytest.mark.parametrize(
"opened_at,read_after_days,window_days",
[
(datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 26, 31),
(datetime(2026, 1, 15, 12, tzinfo=timezone.utc), 29, 31),
(datetime(2026, 1, 31, 12, tzinfo=timezone.utc), 20, 28),
(datetime(2026, 12, 15, 12, tzinfo=timezone.utc), 30, 31),
],
ids=["read_in_a_shorter_month", "read_past_the_shorter_month_length", "clamped_to_february", "across_new_year"],
)
async def test_a_monthly_window_resets_a_calendar_month_after_it_opened(opened_at, read_after_days, window_days):
"""
A month is only as long as the month the window opened in, which is the lifetime the
counter was given when it opened. Measuring the month from the read instead made a
window opened on January 15 and read on February 10 report February 12, and then null
while the counter still refused requests until February 15.
"""
dual_cache = DualCache()
started_at: Final = opened_at.timestamp()
await dual_cache.async_set_cache(
key=model_budget_start_time_cache_key(
entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo"
),
value=started_at,
)
usage = await build_model_max_budget_usage(
entity_type=Litellm_EntityType.USER,
entity_id="user-1",
model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}},
cache=dual_cache,
now=opened_at + timedelta(days=read_after_days),
)
assert usage == {
"gpt-4": {
"current_spend": 0.0,
"budget_limit": 1.0,
"time_period": "1mo",
"budget_reset_at": (opened_at + timedelta(days=window_days)).isoformat(),
}
}
@pytest.mark.asyncio
async def test_a_monthly_window_reports_no_reset_once_its_month_has_passed():
dual_cache = DualCache()
opened_at: Final = datetime(2026, 1, 15, 12, tzinfo=timezone.utc)
await dual_cache.async_set_cache(
key=model_budget_start_time_cache_key(
entity_type=Litellm_EntityType.USER, entity_id="user-1", budget_model="gpt-4", budget_duration="1mo"
),
value=opened_at.timestamp(),
)
usage = await build_model_max_budget_usage(
entity_type=Litellm_EntityType.USER,
entity_id="user-1",
model_max_budget={"gpt-4": {"budget_limit": 1.0, "time_period": "1mo"}},
cache=dual_cache,
now=opened_at + timedelta(days=31),
)
assert usage == {
"gpt-4": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1mo", "budget_reset_at": None}
}
@pytest.mark.asyncio
async def test_each_budgeted_model_reports_its_own_window_reset():
"""Two models on one key own separate windows, so one model's reset must never be reported for the other."""
dual_cache = DualCache()
hourly = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="gpt-4", budget_duration="1h")
daily = dict(entity_type=Litellm_EntityType.KEY, entity_id="vk-hash", budget_model="claude-3", budget_duration="1d")
await dual_cache.async_set_cache(
key=model_budget_start_time_cache_key(**hourly), value=_WINDOW_OPENED_AT.timestamp()
)
await dual_cache.async_set_cache(
key=model_budget_start_time_cache_key(**daily),
value=(_WINDOW_OPENED_AT - timedelta(hours=5)).timestamp(),
)
usage = await build_model_max_budget_usage(
entity_type=Litellm_EntityType.KEY,
entity_id="vk-hash",
model_max_budget={
"gpt-4": {"budget_limit": 1.0, "time_period": "1h"},
"claude-3": {"budget_limit": 1.0, "time_period": "1d"},
"never-used": {"budget_limit": 1.0, "time_period": "1d"},
},
cache=dual_cache,
now=_WINDOW_OPENED_AT,
)
assert usage == {
"gpt-4": {
"current_spend": 0.0,
"budget_limit": 1.0,
"time_period": "1h",
"budget_reset_at": "2026-10-10T13:00:00+00:00",
},
"claude-3": {
"current_spend": 0.0,
"budget_limit": 1.0,
"time_period": "1d",
"budget_reset_at": "2026-10-11T07:00:00+00:00",
},
"never-used": {"current_spend": 0.0, "budget_limit": 1.0, "time_period": "1d", "budget_reset_at": None},
}
@pytest.mark.asyncio
@ -1097,7 +1292,16 @@ async def test_one_malformed_scope_does_not_abort_the_other_scopes():
entity_id="hash-mixed",
model_max_budget=key_budget,
cache=dual_cache,
) == {"gpt-4": {"current_spend": 0.25, "budget_limit": 10.0, "time_period": "1d"}}
) == {
"gpt-4": {
"current_spend": 0.25,
"budget_limit": 10.0,
"time_period": "1d",
"budget_reset_at": await _reset_at_of_open_window(
dual_cache, Litellm_EntityType.KEY, "hash-mixed", "gpt-4", "1d"
),
}
}
@pytest.mark.asyncio
@ -1441,6 +1645,41 @@ async def test_spend_logged_on_one_replica_is_enforced_and_reported_on_another()
await replica_c.is_key_within_model_budget(user_api_key, "gpt-4")
@pytest.mark.asyncio
async def test_a_replica_reports_the_reset_of_a_window_another_replica_opened():
"""
The window start is written by whichever replica served the first charge, so a
replica that has never seen this key must read it from the shared Redis, or it
would report no reset for a model it is actively refusing.
"""
shared_redis = _SharedFakeRedis()
opening_replica = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis))
other_replica = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache(redis_cache=shared_redis))
model_max_budget = {"gpt-4": {"budget_limit": 1.0, "time_period": "1h"}}
await _log_spend(opening_replica, key_hash="vk-shared", model_max_budget=model_max_budget, response_cost=0.5)
usage_on_other: Final = await build_model_max_budget_usage(
entity_type=Litellm_EntityType.KEY,
entity_id="vk-shared",
model_max_budget=model_max_budget,
cache=other_replica.dual_cache,
)
started_at: Final = shared_redis._store[
model_budget_start_time_cache_key(
entity_type=Litellm_EntityType.KEY, entity_id="vk-shared", budget_model="gpt-4", budget_duration="1h"
)
]
assert usage_on_other == {
"gpt-4": {
"current_spend": 0.5,
"budget_limit": 1.0,
"time_period": "1h",
"budget_reset_at": (datetime.fromtimestamp(started_at, tz=timezone.utc) + timedelta(hours=1)).isoformat(),
}
}
def _log_success(limiter, **kwargs):
return limiter.async_log_success_event(
_success_kwargs(**kwargs), response_obj=None, start_time=None, end_time=None
@ -1496,7 +1735,16 @@ async def test_team_model_budget_is_shared_by_every_key_without_an_override(requ
entity_id="team-1",
model_max_budget=team_model_max_budget,
cache=dual_cache,
) == {"gpt-4": {"current_spend": pytest.approx(1.2), "budget_limit": 1.0, "time_period": "1d"}}
) == {
"gpt-4": {
"current_spend": pytest.approx(1.2),
"budget_limit": 1.0,
"time_period": "1d",
"budget_reset_at": await _reset_at_of_open_window(
dual_cache, Litellm_EntityType.TEAM, "team-1", "gpt-4", "1d"
),
}
}
@pytest.mark.asyncio

2
tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py
View file

@ -3499,7 +3499,7 @@ async def test_user_info_v2_response_shape(mocker):
# these two, so dropping them makes a save replace the user's budgets.
assert response_dict["model_max_budget"] == {"gpt-3.5-turbo": {"budget_limit": 5.0, "time_period": "30d"}}
assert response_dict["model_max_budget_usage"] == {
"gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d"}
"gpt-3.5-turbo": {"current_spend": 0.0, "budget_limit": 5.0, "time_period": "30d", "budget_reset_at": None}
}
# Verify teams is a list of strings (team IDs), not team objects

3
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -8514,7 +8514,8 @@ export interface paths {
* hour), so subtracting budget_duration from it does not give the window's start
* - model_max_budget: dict - Per-model budgets, e.g. {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}
* - model_max_budget_usage: dict | None - Current-window spend per model, present only when
* the key has per-model budgets
* the key has per-model budgets. Each entry's budget_reset_at is when that model's open
* window ends and its spend resets, or None when no window is open
* - budget_limits: list | None - Concurrent budget windows, exactly as stored
* - budget_limits_usage: dict | None - Current-window spend per budget window, e.g.
* {"1h": {"current_spend": 0.0009}}, present only when the key has budget windows