diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index d52716951df..16caddc8343 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -21,21 +21,64 @@ _BUDGET_DURATION_WORD_ALIASES: Final[dict[str, str]] = { "monthly": "30d", } +_DURATION_UNIT_ALIASES: Final[dict[str, str]] = { # mutable-ok: constant lookup table + "s": "s", + "sec": "s", + "secs": "s", + "second": "s", + "seconds": "s", + "m": "m", + "min": "m", + "mins": "m", + "minute": "m", + "minutes": "m", + "h": "h", + "hr": "h", + "hrs": "h", + "hour": "h", + "hours": "h", + "d": "d", + "day": "d", + "days": "d", + "w": "w", + "wk": "w", + "wks": "w", + "week": "w", + "weeks": "w", + "mo": "mo", + "mos": "mo", + "mon": "mo", + "mons": "mo", + "month": "mo", + "months": "mo", +} + def _normalize_duration(duration: str) -> str: return _BUDGET_DURATION_WORD_ALIASES.get(duration.strip().lower(), duration) def _extract_from_regex(duration: str) -> tuple[int, str]: - match: Final = re.match(r"(\d+)(mo|[smhdw]?)", duration) + # Use fullmatch so trailing characters after a valid unit are rejected + # rather than silently dropped (e.g. "10mb" must not parse as 10 minutes). + match: Final = re.fullmatch(r"(\d+)([a-z]+)", duration) if not match: raise ValueError("Invalid duration format") - value, unit = match.groups() - value = int(value) + value, raw_unit = match.groups() + unit = _DURATION_UNIT_ALIASES.get(raw_unit) + if unit is None: + raise ValueError(f"Unsupported duration unit, passed duration: {duration}") - return value, unit + return int(value), unit + + +def _parse_canonical_duration(duration: str) -> tuple[int | None, str | None]: + match: Final = re.fullmatch(r"(\d+)(mo|[smhdw])", _normalize_duration(duration)) + if not match: + return None, None + return int(match.group(1)), match.group(2) def get_last_day_of_month(year, month): @@ -129,7 +172,7 @@ def get_next_standardized_reset_time( current_time, _ = _setup_timezone(current_time, timezone_str) # Parse duration - value, unit = _parse_duration(_normalize_duration(duration)) + value, unit = _parse_canonical_duration(duration) if value is None: verbose_logger.warning( "Unrecognized budget_duration %r; falling back to a next-midnight reset. " @@ -171,7 +214,7 @@ 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 months and a duration it does not recognize resets at the next midnight.""" - value, unit = _parse_duration(_normalize_duration(duration)) + value, unit = _parse_canonical_duration(duration) if value is None: return reset_at - timedelta(days=1) match unit: @@ -214,16 +257,6 @@ def _setup_timezone(current_time: datetime, timezone_str: str = "UTC") -> tuple[ return current_time, tz -def _parse_duration(duration: str) -> tuple[int | None, str | None]: - """Parse the duration string into value and unit.""" - match: Final = re.match(r"(\d+)([a-z]+)", duration) - if not match: - return None, None - - value, unit = match.groups() - return int(value), unit - - def _apply_time_of_day(dt: datetime, reset_time_of_day: time) -> datetime: """Set the wall-clock time of `dt` to `reset_time_of_day`, keeping its date and tzinfo.""" return dt.replace( diff --git a/tests/unit/litellm_core_utils/test_duration_parser.py b/tests/unit/litellm_core_utils/test_duration_parser.py index f14015effe6..7d1fd64be1a 100644 --- a/tests/unit/litellm_core_utils/test_duration_parser.py +++ b/tests/unit/litellm_core_utils/test_duration_parser.py @@ -6,7 +6,7 @@ from zoneinfo import ZoneInfo import pytest -import litellm.litellm_core_utils.duration_parser as duration_parser +from litellm.litellm_core_utils import duration_parser from litellm.litellm_core_utils.duration_parser import ( duration_in_seconds, duration_in_seconds_from, @@ -144,6 +144,14 @@ class TestStandardizedResetTime(unittest.TestCase): invalid_tz_result = get_next_standardized_reset_time("1d", on_hour, "NonExistentTimeZone") self.assertEqual(invalid_tz_result, invalid_tz_expected) + def test_malformed_duration_with_trailing_text_falls_back(self): + base_time = datetime(2023, 5, 15, 15, 20, 30, tzinfo=timezone.utc) + expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc) + + result = get_next_standardized_reset_time("10m!", base_time, "UTC") + + self.assertEqual(result, expected) + def test_iana_timezones_previously_unsupported(self): """Test IANA timezones that were previously unsupported by the hardcoded map.""" # Base time: 2023-05-15 15:00:00 UTC @@ -346,6 +354,38 @@ class TestWordFormBudgetDurations(unittest.TestCase): self.assertIn("garbage", mock_warning.call_args.args) +class TestDurationInSeconds(unittest.TestCase): + def test_valid_durations(self): + self.assertEqual(duration_in_seconds("30s"), 30) + self.assertEqual(duration_in_seconds("10m"), 600) + self.assertEqual(duration_in_seconds("2h"), 7200) + self.assertEqual(duration_in_seconds("3d"), 259200) + self.assertEqual(duration_in_seconds("1w"), 604800) + self.assertGreater(duration_in_seconds("1mo"), 0) + + def test_malformed_durations_raise(self): + # Trailing characters after a valid unit must be rejected, not silently + # dropped: "10mb" previously parsed as 10 minutes. + for bad in ["10mb", "1h30m", "1ms", "10dogs", "30x", "abc", "10 m"]: + with self.assertRaises(ValueError): + duration_in_seconds(bad) + + def test_existing_long_unit_spellings_remain_valid(self): + aliases = { + "30seconds": 30, + "15minutes": 900, + "6hours": 21600, + "3days": 259200, + "2weeks": 1209600, + } + for duration, expected in aliases.items(): + with self.subTest(duration=duration): + self.assertEqual(duration_in_seconds(duration), expected) + + self.assertEqual(duration_in_seconds("1mos"), duration_in_seconds("1mo")) + self.assertEqual(duration_in_seconds("2mons"), duration_in_seconds("2mo")) + + class TestGetBudgetWindowStart(unittest.TestCase): def test_window_start_is_the_previous_reset_on_the_same_schedule(self): created_at: Final = datetime(2024, 10, 1, 0, 30, tzinfo=timezone.utc)