mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
feat(budgets): support annual calendar reset periods (#51508)
* feat(budgets): support annual calendar reset periods * fix(budgets): derive annual spend windows from calendar boundaries * fix(budgets): share timezone-aware calendar window boundaries * fix: preserve budget accounting across calendar boundaries * fix(budgets): isolate calendar changes to annual budgets
This commit is contained in:
tin-berri
GitHub
parent
b60246c352
commit
f4f94cc848
14 changed files with 322 additions and 26 deletions
|
|
@ -167,6 +167,15 @@ def _shift_months(moment: datetime, months: int) -> datetime:
|
|||
return moment.replace(year=year, month=month, day=min(moment.day, get_last_day_of_month(year, month)))
|
||||
|
||||
|
||||
def is_annual_budget_duration(duration: str) -> bool:
|
||||
return _parse_duration(normalize_duration(duration)) == (12, "mo")
|
||||
|
||||
|
||||
def get_annual_budget_window_start(reset_at: datetime, timezone_str: str) -> datetime:
|
||||
boundary: Final = _setup_timezone(reset_at, timezone_str)[0]
|
||||
return _shift_months(boundary, -12).astimezone(timezone.utc)
|
||||
|
||||
|
||||
def get_budget_window_start(duration: str, reset_at: datetime) -> datetime:
|
||||
"""Start of the budget period that ends at `reset_at`, under the same rules
|
||||
`get_next_standardized_reset_time` used to pick it: `30d` and `Nmo` span calendar
|
||||
|
|
@ -390,14 +399,22 @@ def _handle_month_reset(
|
|||
Args:
|
||||
current_time: Current datetime
|
||||
base_midnight: Midnight of current day
|
||||
value: Number of months (currently only supports 1 month resets)
|
||||
value: Number of months, 1 for monthly or 12 for annual resets
|
||||
reset_time_of_day: Wall-clock time the reset lands on
|
||||
|
||||
Returns:
|
||||
datetime: First day of the next reset month at `reset_time_of_day`
|
||||
"""
|
||||
if value == 12:
|
||||
first_of_year: Final = base_midnight.replace(month=1, day=1)
|
||||
annual_candidate: Final = _apply_time_of_day(first_of_year, reset_time_of_day)
|
||||
return (
|
||||
annual_candidate.replace(year=annual_candidate.year + 1)
|
||||
if annual_candidate <= current_time
|
||||
else annual_candidate
|
||||
)
|
||||
if value != 1:
|
||||
raise ValueError("Monthly resets currently only support 1 month intervals")
|
||||
raise ValueError("Calendar resets support 1 month or 12 month intervals")
|
||||
|
||||
first_of_this_month: Final = base_midnight.replace(day=1)
|
||||
candidate: Final = _apply_time_of_day(first_of_this_month, reset_time_of_day)
|
||||
|
|
|
|||
|
|
@ -30,7 +30,11 @@ from litellm.constants import (
|
|||
RESET_BUDGET_JOB_MAX_CHUNKS_PER_RUN,
|
||||
RESET_BUDGET_JOB_NAME,
|
||||
)
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.litellm_core_utils.duration_parser import (
|
||||
duration_in_seconds,
|
||||
get_annual_budget_window_start,
|
||||
is_annual_budget_duration,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
DB_RETRY_SAFE_ERROR_TYPES,
|
||||
LiteLLM_BudgetTableFull,
|
||||
|
|
@ -1514,16 +1518,21 @@ class ResetBudgetJob:
|
|||
reset_at_str: Final = window.get("reset_at")
|
||||
if not reset_at_str:
|
||||
return False
|
||||
reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00")).replace(tzinfo=None)
|
||||
if reset_at > now:
|
||||
budget_duration: Final = window["budget_duration"]
|
||||
annual: Final = isinstance(budget_duration, str) and is_annual_budget_duration(budget_duration)
|
||||
saved_reset_at: Final = datetime.fromisoformat(reset_at_str.replace("Z", "+00:00"))
|
||||
reset_at: Final = _as_utc(saved_reset_at) if annual else saved_reset_at.replace(tzinfo=None)
|
||||
cutoff: Final = _as_utc(now) if annual else now
|
||||
if reset_at > cutoff:
|
||||
return False
|
||||
new_value: Final = await ResetBudgetJob._window_carried_spend(window, counter_key, spend_counter_store)
|
||||
try:
|
||||
await spend_counter_store.set(counter_key, new_value)
|
||||
except Exception as redis_err:
|
||||
verbose_proxy_logger.warning("Failed to reset Redis counter %s: %s", counter_key, redis_err)
|
||||
budget_duration: Final = window["budget_duration"]
|
||||
next_reset_at: Final = compute_budget_reset_at(budget_duration=budget_duration, settings=reset_settings)
|
||||
next_reset_at: Final = compute_budget_reset_at(
|
||||
budget_duration=budget_duration, settings=reset_settings, now=now if annual else None
|
||||
)
|
||||
window["reset_at"] = next_reset_at.isoformat()
|
||||
await ResetBudgetJob._roll_window_spend_row(
|
||||
prisma_client=prisma_client,
|
||||
|
|
@ -1531,6 +1540,7 @@ class ResetBudgetJob:
|
|||
entity_id=entity_id,
|
||||
budget_duration=budget_duration,
|
||||
next_reset_at=next_reset_at,
|
||||
reset_settings=reset_settings,
|
||||
)
|
||||
return True
|
||||
|
||||
|
|
@ -1541,6 +1551,7 @@ class ResetBudgetJob:
|
|||
entity_id: str,
|
||||
budget_duration: str,
|
||||
next_reset_at: datetime,
|
||||
reset_settings: BudgetResetSettings | None = None,
|
||||
) -> None:
|
||||
"""Move this window's LiteLLM_BudgetWindowSpend row onto the window
|
||||
that just started, so the maintained total the read path uses starts
|
||||
|
|
@ -1551,7 +1562,11 @@ class ResetBudgetJob:
|
|||
windows from having their counters reset.
|
||||
"""
|
||||
try:
|
||||
window_start: Final = next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration))
|
||||
window_start: Final = (
|
||||
get_annual_budget_window_start(next_reset_at, (reset_settings or get_budget_reset_settings()).timezone)
|
||||
if is_annual_budget_duration(budget_duration)
|
||||
else next_reset_at - timedelta(seconds=duration_in_seconds(budget_duration))
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # duration_in_seconds raises bare exceptions on bad input
|
||||
verbose_proxy_logger.warning("Unparseable budget_duration %s: %s", budget_duration, e)
|
||||
return
|
||||
|
|
@ -1708,7 +1723,9 @@ class ResetBudgetJob:
|
|||
item.spend = _carried_spend(item.spend, _rollover_cap(item.max_budget)) if _rollover_enabled() else 0.0
|
||||
if hasattr(item, "budget_duration") and item.budget_duration is not None:
|
||||
item.budget_reset_at = compute_budget_reset_at(
|
||||
budget_duration=item.budget_duration, settings=reset_settings
|
||||
budget_duration=item.budget_duration,
|
||||
settings=reset_settings,
|
||||
now=current_time if is_annual_budget_duration(item.budget_duration) else None,
|
||||
)
|
||||
return item
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -63,11 +63,13 @@ def get_budget_reset_settings() -> BudgetResetSettings:
|
|||
)
|
||||
|
||||
|
||||
def compute_budget_reset_at(budget_duration: str, settings: BudgetResetSettings) -> datetime:
|
||||
def compute_budget_reset_at(
|
||||
budget_duration: str, settings: BudgetResetSettings, now: datetime | None = None
|
||||
) -> datetime:
|
||||
"""Compute the next reset time for a budget duration using injected settings."""
|
||||
return get_next_standardized_reset_time(
|
||||
duration=budget_duration,
|
||||
current_time=datetime.now(timezone.utc),
|
||||
current_time=now if now is not None else datetime.now(timezone.utc),
|
||||
timezone_str=settings.timezone,
|
||||
reset_time_of_day=settings.reset_time_of_day,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -18,7 +18,11 @@ import litellm
|
|||
from litellm._internal_context import with_service_target
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
|
||||
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
|
||||
from litellm.litellm_core_utils.duration_parser import (
|
||||
duration_in_seconds,
|
||||
get_annual_budget_window_start,
|
||||
is_annual_budget_duration,
|
||||
)
|
||||
from litellm.litellm_core_utils.llm_cost_calc.tiered_pricing import select_tier_for_input, tier_rate
|
||||
from litellm.proxy._types import (
|
||||
Litellm_EntityType,
|
||||
|
|
@ -30,6 +34,7 @@ from litellm.proxy._types import (
|
|||
from litellm.proxy.auth.auth_utils import get_model_from_request
|
||||
from litellm.proxy.auth.budget_throttle import should_throttle_budget_exceeded
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings, get_budget_reset_settings
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
AUTH_OBJECTS_TARGET,
|
||||
UserApiKeyCache,
|
||||
|
|
@ -1287,12 +1292,18 @@ def fixed_budget_window(window: object) -> tuple[str, datetime] | None:
|
|||
return duration, start
|
||||
|
||||
|
||||
def get_budget_window_start(window: object) -> datetime | None:
|
||||
def get_budget_window_start(window: object, reset_settings: BudgetResetSettings | None = None) -> datetime | None:
|
||||
window_dict: Final = _coerce_window(window)
|
||||
budget_duration: Final = window_dict.get("budget_duration")
|
||||
if budget_duration is None:
|
||||
return None
|
||||
try:
|
||||
if is_annual_budget_duration(str(budget_duration)):
|
||||
annual_reset_at: Final = _coerce_datetime(window_dict.get("reset_at"))
|
||||
if annual_reset_at is not None:
|
||||
return get_annual_budget_window_start(
|
||||
annual_reset_at, (reset_settings or get_budget_reset_settings()).timezone
|
||||
)
|
||||
duration_seconds: Final = duration_in_seconds(str(budget_duration))
|
||||
except Exception:
|
||||
return None
|
||||
|
|
|
|||
|
|
@ -175,7 +175,11 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
require_attributes,
|
||||
)
|
||||
from litellm.litellm_core_utils.deployment_pricing import pricing_models
|
||||
from litellm.litellm_core_utils.duration_parser import get_budget_window_start
|
||||
from litellm.litellm_core_utils.duration_parser import (
|
||||
get_annual_budget_window_start,
|
||||
get_budget_window_start,
|
||||
is_annual_budget_duration,
|
||||
)
|
||||
from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
BilledRoundsDispatch,
|
||||
|
|
@ -208,6 +212,7 @@ from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolv
|
|||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_settings
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.db import tool_registry_writer
|
||||
from litellm.proxy.db.create_views import (
|
||||
|
|
@ -9535,7 +9540,11 @@ def _as_aware(moment: datetime) -> datetime:
|
|||
|
||||
def _reset_window(budget_duration: str, budget_reset_at: datetime) -> tuple[datetime, datetime] | None:
|
||||
window_end: Final = _as_aware(budget_reset_at)
|
||||
window_start: Final = get_budget_window_start(budget_duration, window_end)
|
||||
window_start: Final = (
|
||||
get_annual_budget_window_start(window_end, get_budget_reset_settings().timezone)
|
||||
if is_annual_budget_duration(budget_duration)
|
||||
else get_budget_window_start(budget_duration, window_end)
|
||||
)
|
||||
return (window_start, window_end) if window_start < window_end else None
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -47,6 +47,7 @@ const RESET_OPTIONS: readonly { value: string; labelKey: ParseKeys }[] = [
|
|||
{ value: "24h", labelKey: "budgets.form.resetOptions.perDay" },
|
||||
{ value: "7d", labelKey: "budgets.form.resetOptions.perWeek" },
|
||||
{ value: "30d", labelKey: "budgets.form.resetOptions.perMonth" },
|
||||
{ value: "12mo", labelKey: "budgets.form.resetOptions.perYear" },
|
||||
];
|
||||
|
||||
const numberProps = <T extends { value?: number | null; onChange: (value: number | null) => void }>({
|
||||
|
|
|
|||
|
|
@ -25,7 +25,10 @@ describe("BudgetModal", () => {
|
|||
createMock.mockResolvedValue(undefined);
|
||||
});
|
||||
|
||||
it("submits the spend limit and reset period alongside the rate limits", async () => {
|
||||
it.each([
|
||||
["per month", "30d"],
|
||||
["per year", "12mo"],
|
||||
])("submits the spend limit %s alongside the rate limits", async (label, duration) => {
|
||||
const user = userEvent.setup();
|
||||
renderModal();
|
||||
|
||||
|
|
@ -33,7 +36,7 @@ describe("BudgetModal", () => {
|
|||
expect(resetPeriod()).toBeVisible();
|
||||
fireEvent.change(screen.getByLabelText("Budget ID"), { target: { value: "budget-alpha" } });
|
||||
fireEvent.change(screen.getByLabelText(/^Spend limit/), { target: { value: "42.567" } });
|
||||
await chooseSelectOption(user, resetPeriod(), "per month");
|
||||
await chooseSelectOption(user, resetPeriod(), label);
|
||||
fireEvent.change(screen.getByLabelText("Tokens per minute"), { target: { value: "500" } });
|
||||
fireEvent.change(screen.getByLabelText("Requests per minute"), { target: { value: "7" } });
|
||||
fireEvent.change(screen.getByLabelText("Batch tokens per day"), { target: { value: "90000" } });
|
||||
|
|
@ -43,7 +46,7 @@ describe("BudgetModal", () => {
|
|||
const expected = {
|
||||
budget_id: "budget-alpha",
|
||||
max_budget: 42.57,
|
||||
budget_duration: "30d",
|
||||
budget_duration: duration,
|
||||
tpm_limit: 500,
|
||||
rpm_limit: 7,
|
||||
tpd_limit: 90000,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import { render, screen } from "@testing-library/react";
|
|||
import userEvent from "@testing-library/user-event";
|
||||
import { describe, expect, it, vi } from "vitest";
|
||||
|
||||
import BudgetDurationDropdown from "./budget_duration_dropdown";
|
||||
import BudgetDurationDropdown, { getBudgetDurationLabel } from "./budget_duration_dropdown";
|
||||
|
||||
describe("BudgetDurationDropdown", () => {
|
||||
it("labels an unset period as Never resets in the trigger and option", async () => {
|
||||
|
|
@ -39,3 +39,13 @@ describe("BudgetDurationDropdown", () => {
|
|||
expect(await screen.findByRole("option", { name: "Not set" })).toBeInTheDocument();
|
||||
});
|
||||
});
|
||||
|
||||
it("saves and displays an annual budget period", async () => {
|
||||
const user = userEvent.setup();
|
||||
const onChange = vi.fn();
|
||||
render(<BudgetDurationDropdown value="30d" onChange={onChange} />);
|
||||
await user.click(screen.getByRole("combobox"));
|
||||
await user.click(await screen.findByRole("option", { name: "annually" }));
|
||||
expect(onChange).toHaveBeenCalledWith("12mo", expect.anything());
|
||||
expect(getBudgetDurationLabel("12mo")).toBe("annually");
|
||||
});
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ const DURATION_LABELS: Record<string, ParseKeys> = {
|
|||
"24h": "budgetDurationDropdown.periods.daily",
|
||||
"7d": "budgetDurationDropdown.periods.weekly",
|
||||
"30d": "budgetDurationDropdown.periods.monthly",
|
||||
"12mo": "budgetDurationDropdown.periods.annual",
|
||||
};
|
||||
|
||||
interface BudgetDurationDropdownProps {
|
||||
|
|
@ -51,6 +52,7 @@ const BudgetDurationDropdown: React.FC<BudgetDurationDropdownProps> = ({
|
|||
<SelectItem value="24h">{t("budgetDurationDropdown.periods.daily")}</SelectItem>
|
||||
<SelectItem value="7d">{t("budgetDurationDropdown.periods.weekly")}</SelectItem>
|
||||
<SelectItem value="30d">{t("budgetDurationDropdown.periods.monthly")}</SelectItem>
|
||||
<SelectItem value="12mo">{t("budgetDurationDropdown.periods.annual")}</SelectItem>
|
||||
</SelectContent>
|
||||
</Select>
|
||||
);
|
||||
|
|
@ -64,6 +66,7 @@ export const getBudgetDurationLabel = (value: string | null | undefined, t: TFun
|
|||
"24h": "budgetDurationDropdown.periods.daily",
|
||||
"7d": "budgetDurationDropdown.periods.weekly",
|
||||
"30d": "budgetDurationDropdown.periods.monthly",
|
||||
"12mo": "budgetDurationDropdown.periods.annual",
|
||||
};
|
||||
|
||||
return Object.hasOwn(budgetDurationMap, value) ? t(budgetDurationMap[value]) : value;
|
||||
|
|
|
|||
|
|
@ -3072,7 +3072,8 @@
|
|||
"hourly": "hourly",
|
||||
"daily": "daily",
|
||||
"weekly": "weekly",
|
||||
"monthly": "monthly"
|
||||
"monthly": "monthly",
|
||||
"annual": "annually"
|
||||
},
|
||||
"notSet": "Not set"
|
||||
},
|
||||
|
|
@ -3406,7 +3407,8 @@
|
|||
"inTotal": "in total",
|
||||
"perDay": "per day",
|
||||
"perWeek": "per week",
|
||||
"perMonth": "per month"
|
||||
"perMonth": "per month",
|
||||
"perYear": "per year"
|
||||
},
|
||||
"noLimit": "No limit",
|
||||
"optional": "Optional",
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ 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_annual_budget_window_start,
|
||||
get_budget_window_start,
|
||||
get_next_standardized_reset_time,
|
||||
)
|
||||
|
|
@ -195,6 +196,79 @@ class TestStandardizedResetTime(unittest.TestCase):
|
|||
self.assertEqual(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"now, zone, reset_time, expected",
|
||||
[
|
||||
(
|
||||
datetime(2026, 11, 12, tzinfo=timezone.utc),
|
||||
"UTC",
|
||||
time(0),
|
||||
datetime(2027, 1, 1, tzinfo=timezone.utc),
|
||||
),
|
||||
(
|
||||
datetime(2028, 2, 29, tzinfo=timezone.utc),
|
||||
"UTC",
|
||||
time(0),
|
||||
datetime(2029, 1, 1, tzinfo=timezone.utc),
|
||||
),
|
||||
(
|
||||
datetime(2027, 1, 1, tzinfo=timezone.utc),
|
||||
"UTC",
|
||||
time(0),
|
||||
datetime(2028, 1, 1, tzinfo=timezone.utc),
|
||||
),
|
||||
(
|
||||
datetime(2027, 1, 1, 8, tzinfo=timezone.utc),
|
||||
"UTC",
|
||||
time(12),
|
||||
datetime(2027, 1, 1, 12, tzinfo=timezone.utc),
|
||||
),
|
||||
(
|
||||
datetime(2027, 1, 1, 12, tzinfo=timezone.utc),
|
||||
"UTC",
|
||||
time(12),
|
||||
datetime(2028, 1, 1, 12, tzinfo=timezone.utc),
|
||||
),
|
||||
(
|
||||
datetime(2027, 1, 1, 2, tzinfo=timezone.utc),
|
||||
"America/New_York",
|
||||
time(0),
|
||||
datetime(2027, 1, 1, tzinfo=ZoneInfo("America/New_York")),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_annual_reset_and_window_start(
|
||||
now: datetime, zone: str, reset_time: time, expected: datetime
|
||||
) -> None:
|
||||
reset: Final = get_next_standardized_reset_time("12mo", now, zone, reset_time)
|
||||
assert reset == expected
|
||||
assert get_budget_window_start("12mo", reset) == expected.replace(year=expected.year - 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"zone,reset_at,expected",
|
||||
(
|
||||
("UTC", "2029-01-01T00:00:00+00:00", "2028-01-01T00:00:00+00:00"),
|
||||
("UTC", "2101-01-01T00:00:00+00:00", "2100-01-01T00:00:00+00:00"),
|
||||
("America/New_York", "2029-01-01T00:00:00-05:00", "2028-01-01T05:00:00+00:00"),
|
||||
("Asia/Kathmandu", "2029-01-01T12:30:00+05:45", "2028-01-01T06:45:00+00:00"),
|
||||
("Pacific/Apia", "2012-01-01T00:00:00+14:00", "2011-01-01T10:00:00+00:00"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("storage", ("offset", "utc", "naive_utc"))
|
||||
def test_annual_window_survives_timestamp_storage(
|
||||
zone: str, reset_at: str, expected: str, storage: str
|
||||
) -> None:
|
||||
reset: Final = datetime.fromisoformat(reset_at)
|
||||
stored: Final = {
|
||||
"offset": reset,
|
||||
"utc": reset.astimezone(timezone.utc),
|
||||
"naive_utc": reset.astimezone(timezone.utc).replace(tzinfo=None),
|
||||
}[storage]
|
||||
|
||||
assert get_annual_budget_window_start(stored, zone) == datetime.fromisoformat(expected)
|
||||
|
||||
|
||||
class TestResetTimeOfDay(unittest.TestCase):
|
||||
"""A configurable reset_time_of_day shifts day/week/month resets off midnight."""
|
||||
|
||||
|
|
|
|||
|
|
@ -28,11 +28,71 @@ from litellm.proxy.common_utils.reset_budget_job import (
|
|||
_write_team_windows,
|
||||
)
|
||||
from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings
|
||||
from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from tests.unit.proxy.db.fake_prisma_engine import engine_call
|
||||
from tests.unit.redis.stores.spend_counters_double import FakeSpendCounters
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"reset_at,now",
|
||||
(
|
||||
("2027-01-01T00:00:00-05:00", "2027-01-01T04:59:59+00:00"),
|
||||
("2027-01-01T00:00:00+05:45", "2026-12-31T18:14:59+00:00"),
|
||||
),
|
||||
)
|
||||
async def test_annual_window_reset_waits_for_the_saved_instant(reset_at: str, now: str) -> None:
|
||||
store: Final = FakeSpendCounters().store()
|
||||
counter: Final = "spend:key:offset-budget:window:12mo"
|
||||
await store.set(counter, 10.0)
|
||||
window: Final = {"budget_duration": "12mo", "reset_at": reset_at, "max_budget": 10.0}
|
||||
client: Final = MagicMock(spec=PrismaClient)
|
||||
changed: Final = await ResetBudgetJob._reset_expired_window(
|
||||
window,
|
||||
counter,
|
||||
store,
|
||||
datetime.fromisoformat(now),
|
||||
BudgetResetSettings(timezone="America/New_York"),
|
||||
client,
|
||||
reset_budget_job_module.Litellm_EntityType.KEY,
|
||||
"offset-budget",
|
||||
)
|
||||
|
||||
assert changed is False
|
||||
assert (await store.read(counter))[0] == 10.0
|
||||
assert window["reset_at"] == reset_at
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("today", ["2028-03-01", "2029-03-01"])
|
||||
async def test_annual_reset_persists_the_same_calendar_start_as_admission(today: str) -> None: # test-quality-ok: the SQL write is the reset job's observable persistence contract
|
||||
client: Final = MagicMock(spec=PrismaClient)
|
||||
client.db = MagicMock()
|
||||
client.db.execute_raw = AsyncMock(return_value=1)
|
||||
|
||||
with patch("time.time", return_value=datetime.fromisoformat(today).replace(tzinfo=timezone.utc).timestamp()):
|
||||
await ResetBudgetJob._roll_window_spend_row(
|
||||
prisma_client=client,
|
||||
entity_type=reset_budget_job_module.Litellm_EntityType.KEY,
|
||||
entity_id="annual-key",
|
||||
budget_duration="12mo",
|
||||
next_reset_at=datetime(2029, 1, 1, tzinfo=timezone.utc),
|
||||
reset_settings=BudgetResetSettings(),
|
||||
)
|
||||
|
||||
assert client.db.execute_raw.await_count == 1
|
||||
assert client.db.execute_raw.await_args.args[1:5] == (
|
||||
"key",
|
||||
"annual-key",
|
||||
"12mo",
|
||||
datetime(2028, 1, 1),
|
||||
)
|
||||
assert get_budget_window_start(
|
||||
{"budget_duration": "12mo", "reset_at": datetime(2029, 1, 1, tzinfo=timezone.utc)}
|
||||
) == client.db.execute_raw.await_args.args[4].replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
# Mock classes for testing
|
||||
class MockTable:
|
||||
"""A single prisma table: records reads/writes and replays canned rows."""
|
||||
|
|
|
|||
|
|
@ -2,7 +2,9 @@ from __future__ import annotations
|
|||
|
||||
import math
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -23,6 +25,7 @@ from litellm.proxy.common_utils.user_api_key_cache import (
|
|||
project_cache_key,
|
||||
team_membership_reservation_cache_key,
|
||||
)
|
||||
from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
_get_team_member_budget_counter,
|
||||
estimate_request_max_cost,
|
||||
|
|
@ -465,6 +468,88 @@ def test_budget_window_start_reads_json_encoded_windows(window: str, expected: d
|
|||
assert get_budget_window_start(window) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("today", ["2028-03-01", "2029-03-01"])
|
||||
@pytest.mark.parametrize(
|
||||
("duration", "reset_at", "expected"),
|
||||
[
|
||||
("12mo", "2029-01-01T00:00:00Z", datetime(2028, 1, 1, tzinfo=timezone.utc)),
|
||||
("12mo", "2028-01-01T00:00:00", datetime(2027, 1, 1, tzinfo=timezone.utc)),
|
||||
],
|
||||
)
|
||||
def test_annual_window_keeps_all_spend_from_its_reset_boundary(
|
||||
today: str, duration: str, reset_at: str, expected: datetime
|
||||
) -> None:
|
||||
with patch("time.time", return_value=datetime.fromisoformat(today).replace(tzinfo=timezone.utc).timestamp()):
|
||||
assert get_budget_window_start({"budget_duration": duration, "reset_at": reset_at}) == expected
|
||||
|
||||
|
||||
class _AnnualWindowDatabase:
|
||||
def __init__(self) -> None:
|
||||
self.db = self
|
||||
self.litellm_budgetwindowspend = self
|
||||
self.litellm_spendlogs = self
|
||||
self.start: datetime | None = None
|
||||
self.spend = 0.0
|
||||
self.scans = 0
|
||||
|
||||
async def execute_raw(self, sql: str, *parameters: object) -> int:
|
||||
start: Final = parameters[3]
|
||||
assert isinstance(start, datetime)
|
||||
self.start = start
|
||||
self.spend = 0.0
|
||||
return 1
|
||||
|
||||
async def find_unique(self, **query: object) -> SimpleNamespace:
|
||||
return SimpleNamespace(window_start=self.start, spend=self.spend)
|
||||
|
||||
async def group_by(self, **query: object) -> list[dict[str, object]]:
|
||||
self.scans += 1
|
||||
return [{"api_key": "calendar-key", "_sum": {"spend": 0.0}}]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"duration,zone,reset_at",
|
||||
(
|
||||
("12mo", "UTC", "2029-01-01T00:00:00Z"),
|
||||
("12mo", "America/New_York", "2029-01-01T00:00:00-05:00"),
|
||||
("12mo", "Pacific/Apia", "2012-01-01T00:00:00+14:00"),
|
||||
),
|
||||
)
|
||||
async def test_cold_annual_counter_preserves_reset_writer_spend(
|
||||
monkeypatch: pytest.MonkeyPatch, duration: str, zone: str, reset_at: str
|
||||
) -> None:
|
||||
from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.types.integrations.slack_alerting import Litellm_EntityType
|
||||
|
||||
database: Final = _AnnualWindowDatabase()
|
||||
store: Final = SpendCounterStore(connect(None).spend_counters())
|
||||
monkeypatch.setattr(litellm, "timezone", zone, raising=False)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", database)
|
||||
monkeypatch.setattr(proxy_server, "spend_counter_store", store)
|
||||
monkeypatch.setattr(proxy_server, "spend_counter_cache", DualCache())
|
||||
await ResetBudgetJob._roll_window_spend_row(
|
||||
prisma_client=database, # pyright: ignore[reportArgumentType] # in-memory external database boundary
|
||||
entity_type=Litellm_EntityType.KEY,
|
||||
entity_id="calendar-key",
|
||||
budget_duration=duration,
|
||||
next_reset_at=datetime.fromisoformat(reset_at.replace("Z", "+00:00")),
|
||||
reset_settings=BudgetResetSettings(timezone=zone),
|
||||
)
|
||||
database.spend = 10.0
|
||||
token: Final = UserAPIKeyAuth(
|
||||
token="calendar-key",
|
||||
budget_limits=[{"budget_duration": duration, "reset_at": reset_at, "max_budget": 10.0}],
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError) as failure:
|
||||
await _virtual_key_multi_budget_check(token)
|
||||
assert failure.value.current_cost == 10.0
|
||||
assert store.get_local(f"spend:key:calendar-key:window:{duration}") == 10.0
|
||||
assert database.scans == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reservation_counts_a_pdf_with_the_deployment_behind_a_gpt_named_group() -> None:
|
||||
from litellm.proxy.spend_tracking.input_tokens import count_input_tokens_for_model
|
||||
|
|
|
|||
|
|
@ -13933,12 +13933,13 @@ async def _drain(queue):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_window_spend_row_is_enqueued_with_the_actual_cost():
|
||||
@pytest.mark.parametrize("duration", ("30d", "12mo"))
|
||||
async def test_key_window_spend_row_is_enqueued_with_the_actual_cost(duration: str):
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
reset_at = datetime.now(timezone.utc) + timedelta(days=10)
|
||||
reset_at: Final = datetime(2029, 1, 1, tzinfo=timezone.utc)
|
||||
key_obj = MagicMock()
|
||||
key_obj.budget_limits = [{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}]
|
||||
key_obj.budget_limits = [{"budget_duration": duration, "max_budget": 100.0, "reset_at": reset_at.isoformat()}]
|
||||
|
||||
with _window_spend_enqueue_env({"hashed-token": key_obj}) as queue:
|
||||
await increment_spend_counters(token="hashed-token", team_id=None, user_id=None, response_cost=0.25)
|
||||
|
|
@ -13947,9 +13948,10 @@ async def test_key_window_spend_row_is_enqueued_with_the_actual_cost():
|
|||
assert len(enqueued) == 1
|
||||
assert enqueued[0]["entity_type"] == "key"
|
||||
assert enqueued[0]["entity_id"] == "hashed-token"
|
||||
assert enqueued[0]["window_duration"] == "30d"
|
||||
assert enqueued[0]["window_duration"] == duration
|
||||
assert enqueued[0]["spend"] == pytest.approx(0.25)
|
||||
assert enqueued[0]["window_start"] == (reset_at - timedelta(days=30)).astimezone(timezone.utc).replace(
|
||||
expected_start: Final = datetime(2028, 1, 1, tzinfo=timezone.utc) if duration == "12mo" else reset_at - timedelta(days=30)
|
||||
assert enqueued[0]["window_start"] == expected_start.astimezone(timezone.utc).replace(
|
||||
tzinfo=None
|
||||
).isoformat(timespec="microseconds")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue