Merge 488819ca25 into 230f653523

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(

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)