From 6021db5bbaef9bf9dda6400fa0190b428328af43 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 13 Nov 2026 11:08:55 -0800 Subject: [PATCH] fix(budgets): reset 12months and other long unit spellings on their short form's schedule (#51672) * fix(budgets): reset long unit spellings like 12months on the schedule they are charged for * test(budgets): match the 3mo error message in the long month spelling test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/duration_parser.py | 8 +++-- .../test_listing_and_health_contracts.py | 17 ++++++++++ .../test_duration_parser.py | 34 ++++++++++++++++++- 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 759ffbffcb1..2abb1c7e4da 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -229,8 +229,12 @@ def _parse_duration(duration: str) -> tuple[int | None, str | None]: if not match: return None, None - value, unit = match.groups() - return int(value), unit + value: Final = int(match.group(1)) + spelled_unit: Final = match.group(2) + unit: Final = _extract_from_regex(duration)[1] or spelled_unit + if unit == "mo" and spelled_unit != "mo" and value not in (1, 12): + return value, spelled_unit + return value, unit def _apply_time_of_day(dt: datetime, reset_time_of_day: time) -> datetime: diff --git a/tests/integration/management/test_listing_and_health_contracts.py b/tests/integration/management/test_listing_and_health_contracts.py index 8492dc50adb..9b93b85c5de 100644 --- a/tests/integration/management/test_listing_and_health_contracts.py +++ b/tests/integration/management/test_listing_and_health_contracts.py @@ -2,6 +2,7 @@ import uuid from datetime import datetime, timezone from typing import Final +import pytest from pydantic import JsonValue from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time @@ -57,6 +58,22 @@ def test_budget_duration_schedules_reset_at_the_next_standardized_boundary(gatew assert abs((_utc(rows[0]["reset_at"]) - expected).total_seconds()) <= 3, (rows, expected) +@pytest.mark.parametrize("budget_duration", ["12mo", "12months", "12mos"]) +def test_twelve_month_budget_resets_on_january_first_however_the_unit_is_spelled( + gateway: Gateway, budget_duration: str +) -> None: + with gateway.scenario() as scenario: + budget_id: Final = scenario.budget(max_budget=1200.0, budget_duration=budget_duration) + rows: Final = read_rows( + 'SELECT created_at::text AS created_at, budget_reset_at::text AS reset_at FROM "LiteLLM_BudgetTable" ' + "WHERE budget_id = %s", + (budget_id,), + ) + assert len(rows) == 1, rows + created_at: Final = _utc(rows[0]["created_at"]) + assert _utc(rows[0]["reset_at"]) == datetime(created_at.year + 1, 1, 1, tzinfo=timezone.utc), rows + + def test_admin_health_counts_every_deployment_and_reports_a_live_one_healthy(gateway: Gateway) -> None: with gateway.scenario() as scenario: model: Final = scenario.model() diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index 2f74762edaa..c5e662e0045 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -13,6 +13,7 @@ from litellm.litellm_core_utils.duration_parser import ( get_annual_budget_window_start, get_budget_window_start, get_next_standardized_reset_time, + is_annual_budget_duration, ) @@ -245,6 +246,37 @@ def test_annual_reset_and_window_start( assert get_budget_window_start("12mo", reset) == expected.replace(year=expected.year - 1) +@pytest.mark.parametrize( + "long_spelling,short_form", + [ + ("12months", "12mo"), + ("12mos", "12mo"), + ("1month", "1mo"), + ("7days", "7d"), + ("2weeks", "2w"), + ("1hour", "1h"), + ("5min", "5m"), + ("30sec", "30s"), + ], +) +def test_long_unit_spelling_resets_on_the_schedule_it_is_charged_for(long_spelling: str, short_form: str) -> None: + now: Final = datetime(2026, 11, 13, 18, 0, tzinfo=timezone.utc) + assert duration_in_seconds_from(now, long_spelling) == duration_in_seconds_from(now, short_form) + reset: Final = get_next_standardized_reset_time(long_spelling, now, "UTC") + assert reset == get_next_standardized_reset_time(short_form, now, "UTC") + assert get_budget_window_start(long_spelling, reset) == get_budget_window_start(short_form, reset) + assert is_annual_budget_duration(long_spelling) == is_annual_budget_duration(short_form) + + +def test_long_month_spelling_without_a_calendar_schedule_keeps_its_nightly_reset() -> None: + now: Final = datetime(2026, 11, 13, 18, 0, tzinfo=timezone.utc) + with pytest.raises(ValueError, match="1 month or 12 month"): + get_next_standardized_reset_time("3mo", now, "UTC") + reset: Final = get_next_standardized_reset_time("3months", now, "UTC") + assert reset == datetime(2026, 11, 14, tzinfo=timezone.utc) + assert get_budget_window_start("3months", reset) == datetime(2026, 11, 13, tzinfo=timezone.utc) + + @pytest.mark.parametrize( "zone,reset_at,expected", ( @@ -445,7 +477,7 @@ class TestGetBudgetWindowStart(unittest.TestCase): def test_unrecognized_duration_is_the_one_day_window_the_scheduler_falls_back_to(self): self.assertEqual( - get_budget_window_start("1hr", datetime(2024, 1, 16, tzinfo=timezone.utc)), + get_budget_window_start("1y", datetime(2024, 1, 16, tzinfo=timezone.utc)), datetime(2024, 1, 15, tzinfo=timezone.utc), )