diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index 7ac2299b7d8..76a25254557 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -21,21 +21,57 @@ _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 get_last_day_of_month(year, month): @@ -191,11 +227,14 @@ def _setup_timezone(current_time: datetime, timezone_str: str = "UTC") -> tuple[ 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) + match: Final = re.fullmatch(r"(\d+)([a-z]+)", duration) if not match: return None, None - value, unit = match.groups() + value, raw_unit = match.groups() + unit = _DURATION_UNIT_ALIASES.get(raw_unit) + if unit is None: + return None, None return int(value), unit diff --git a/tests/test_litellm/litellm_core_utils/test_duration_parser.py b/tests/test_litellm/litellm_core_utils/test_duration_parser.py index cb9f273a0a7..df0fc9938c6 100644 --- a/tests/test_litellm/litellm_core_utils/test_duration_parser.py +++ b/tests/test_litellm/litellm_core_utils/test_duration_parser.py @@ -149,6 +149,14 @@ class TestStandardizedResetTime(unittest.TestCase): ) 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 @@ -385,5 +393,44 @@ 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")) + + def test_long_unit_spelling_sets_reset_interval(self): + now = datetime(2023, 5, 15, 15, 0, 0, tzinfo=timezone.utc) + self.assertEqual( + get_next_standardized_reset_time("3days", now, "UTC"), + datetime(2023, 5, 18, 0, 0, 0, tzinfo=timezone.utc), + ) + + if __name__ == "__main__": unittest.main()