Merge bfd1e96add into 0c98afa780

litellm/proxy/proxy_server.py

@@ -687,7 +687,6 @@ from litellm.proxy.utils import (
     _get_openapi_url,
     _get_projected_spend_over_limit,
     _get_redoc_url,
-    _is_projected_spend_over_limit,
     _is_valid_team_configs,
     evict_config_param,
     get_config_param,
@@ -3515,21 +3514,18 @@ async def update_cache(
         new_spend: Final = existing_spend + response_cost
 
         ## CHECK IF USER PROJECTED SPEND > SOFT LIMIT
-        if (
-            existing_spend_obj.soft_budget_cooldown is False
-            and existing_spend_obj.soft_budget is not None
-            and (
-                _is_projected_spend_over_limit(
-                    current_spend=new_spend,
-                    soft_budget_limit=existing_spend_obj.soft_budget,
-                )
-                is True
-            )
-        ):
-            projected_spend, projected_exceeded_date = _get_projected_spend_over_limit(
+        projection: Final = (
+            _get_projected_spend_over_limit(
                 current_spend=new_spend,
                 soft_budget_limit=existing_spend_obj.soft_budget,
+                budget_duration=existing_spend_obj.budget_duration,
+                budget_reset_at=existing_spend_obj.budget_reset_at,
             )
+            if existing_spend_obj.soft_budget_cooldown is False and existing_spend_obj.soft_budget is not None
+            else None
+        )
+        if projection is not None:
+            projected_spend, projected_exceeded_date = projection
             soft_limit: Final = existing_spend_obj.soft_budget
             call_info: Final = CallInfo(
                 token=existing_spend_obj.token or "",

litellm/proxy/utils.py

@@ -99,6 +99,7 @@ from litellm.litellm_core_utils.core_helpers import (
     independent_snapshot,
     is_expected_client_error,
 )
+from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.litellm_core_utils.litellm_logging import Logging
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
@@ -7327,63 +7328,65 @@ def _get_month_end_date(today: date) -> date:
     return date(today.year, today.month + 1, 1) - timedelta(days=1)
 
 
-def _is_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None):
-    if soft_budget_limit is None:
-        # If there's no limit, we can't exceed it.
-        return False
-
-    today: Final = date.today()
-
-    # Finding the first day of the next month, then subtracting one day to get the end of the current month.
-    end_month: Final = _get_month_end_date(today)
-
-    remaining_days: Final = (end_month - today).days
-
-    # Check for the start of the month to avoid division by zero
-    if today.day == 1:
-        daily_spend_estimate = current_spend
-    else:
-        daily_spend_estimate = current_spend / (today.day - 1)
-
-    # Total projected spend for the month
-    projected_spend: Final = current_spend + (daily_spend_estimate * remaining_days)
-
-    if projected_spend > soft_budget_limit:
-        print_verbose("Projected spend exceeds soft budget limit!")
-        return True
-    return False
+def _get_budget_window(
+    today: date,
+    budget_duration: str | None,
+    budget_reset_at: datetime | None,
+) -> tuple[date, int]:
+    if budget_duration is None or budget_reset_at is None:
+        return _get_month_end_date(today), max(today.day - 1, 1)
+    try:
+        window_days: Final = max(duration_in_seconds(duration=budget_duration) // 86400, 1)
+    except ValueError:
+        return _get_month_end_date(today), max(today.day - 1, 1)
+    window_end: Final = budget_reset_at.date()
+    window_start: Final = window_end - timedelta(days=window_days)
+    return window_end, max((today - window_start).days, 1)
 
 
-def _get_projected_spend_over_limit(current_spend: float, soft_budget_limit: float | None) -> tuple | None:
+def _is_projected_spend_over_limit(
+    current_spend: float,
+    soft_budget_limit: float | None,
+    budget_duration: str | None = None,
+    budget_reset_at: datetime | None = None,
+) -> bool:
+    return (
+        _get_projected_spend_over_limit(
+            current_spend=current_spend,
+            soft_budget_limit=soft_budget_limit,
+            budget_duration=budget_duration,
+            budget_reset_at=budget_reset_at,
+        )
+        is not None
+    )
+
+
+def _get_projected_spend_over_limit(
+    current_spend: float,
+    soft_budget_limit: float | None,
+    budget_duration: str | None = None,
+    budget_reset_at: datetime | None = None,
+) -> tuple[float, date] | None:
     if soft_budget_limit is None:
         return None
 
-    today: Final = date.today()
-    end_month: Final = _get_month_end_date(today)
-    remaining_days: Final = (end_month - today).days
-
-    # assuming the current spend till today (not including today)
-    if today.day == 1:
-        daily_spend = current_spend
-    else:
-        daily_spend = current_spend / (today.day - 1)
+    today: Final = (
+        datetime.now(budget_reset_at.tzinfo).date()
+        if budget_reset_at is not None and budget_reset_at.tzinfo is not None
+        else date.today()
+    )
+    window_end, elapsed_days = _get_budget_window(today, budget_duration, budget_reset_at)
+    remaining_days: Final = max((window_end - today).days, 0)
+    daily_spend: Final = current_spend / elapsed_days
     projected_spend: Final = current_spend + (daily_spend * remaining_days)
 
-    if projected_spend > soft_budget_limit:
-        if daily_spend <= 0:
-            limit_exceed_date = today
-        else:
-            remaining_budget: Final = soft_budget_limit - current_spend
-            if remaining_budget <= 0:
-                limit_exceed_date = today
-            else:
-                approx_days: Final = remaining_budget / daily_spend
-                limit_exceed_date = today + timedelta(days=approx_days)
+    if projected_spend <= soft_budget_limit:
+        return None
 
-        # return the projected spend and the date it will exceeded
-        return projected_spend, limit_exceed_date
-
-    return None
+    remaining_budget: Final = soft_budget_limit - current_spend
+    if daily_spend <= 0 or remaining_budget <= 0:
+        return projected_spend, today
+    return projected_spend, min(today + timedelta(days=remaining_budget / daily_spend), window_end)
 
 
 def _is_valid_team_configs(team_id=None, team_config=None, request_data=None):

tests/test_litellm/proxy/utils/helpers/test_month_end_projection.py

@@ -1,4 +1,4 @@
-from datetime import date, timedelta
+from datetime import date, datetime, timedelta, timezone
 
 import pytest
 
@@ -230,3 +230,100 @@ def test_get_projected_spend_over_limit_raises_when_today_missing(monkeypatch):
     monkeypatch.setattr("litellm.proxy.utils.date", _Broken)
     with pytest.raises(RuntimeError):
         _get_projected_spend_over_limit(current_spend=1.0, soft_budget_limit=1.0)
+
+
+def test_get_projected_spend_over_limit_daily_budget_projects_to_reset_not_month_end(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 15))
+    result = _get_projected_spend_over_limit(
+        current_spend=8.0,
+        soft_budget_limit=10.0,
+        budget_duration="1d",
+        budget_reset_at=datetime(2024, 1, 16, 0, 0, 0),
+    )
+    assert result is not None
+    projected, exceed_date = result
+    assert projected == 16.0
+    assert exceed_date <= date(2024, 1, 16)
+
+
+def test_get_projected_spend_over_limit_weekly_budget_uses_window_daily_rate(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 15))
+    result = _get_projected_spend_over_limit(
+        current_spend=8.0,
+        soft_budget_limit=10.0,
+        budget_duration="7d",
+        budget_reset_at=datetime(2024, 1, 18, 0, 0, 0),
+    )
+    assert result is not None
+    projected, exceed_date = result
+    assert projected == 14.0
+    assert exceed_date == date(2024, 1, 16)
+
+
+def test_get_projected_spend_over_limit_daily_budget_under_limit_no_false_alert(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 15))
+    assert (
+        _get_projected_spend_over_limit(
+            current_spend=4.0,
+            soft_budget_limit=10.0,
+            budget_duration="1d",
+            budget_reset_at=datetime(2024, 1, 16, 0, 0, 0),
+        )
+        is None
+    )
+
+
+def test_is_projected_spend_over_limit_daily_budget_under_limit(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 15))
+    assert (
+        _is_projected_spend_over_limit(
+            current_spend=4.0,
+            soft_budget_limit=10.0,
+            budget_duration="1d",
+            budget_reset_at=datetime(2024, 1, 16, 0, 0, 0),
+        )
+        is False
+    )
+
+
+def test_get_projected_spend_over_limit_without_duration_keeps_month_end_behavior(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 15))
+    result = _get_projected_spend_over_limit(current_spend=8.0, soft_budget_limit=10.0)
+    assert result is not None
+    projected, _ = result
+    assert projected == pytest.approx(8.0 + (8.0 / 14) * 16)
+
+
+def test_get_projected_spend_over_limit_uses_reset_timezone_not_host_timezone(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 16))
+
+    class _FrozenDatetime(datetime):
+        @classmethod
+        def now(cls, tz=None):
+            return datetime(2024, 1, 15, 23, 30, tzinfo=tz)
+
+    monkeypatch.setattr("litellm.proxy.utils.datetime", _FrozenDatetime)
+    result = _get_projected_spend_over_limit(
+        current_spend=8.0,
+        soft_budget_limit=10.0,
+        budget_duration="1d",
+        budget_reset_at=datetime(2024, 1, 16, 0, 0, 0, tzinfo=timezone.utc),
+    )
+    assert result is not None
+    projected, exceed_date = result
+    assert projected == 16.0
+    assert exceed_date <= date(2024, 1, 16)
+
+
+def test_get_projected_spend_over_limit_exceed_date_never_past_reset(monkeypatch):
+    _freeze_today(monkeypatch, date(2024, 1, 15))
+    for spend, soft in [(8.0, 9.99), (8.0, 13.9), (8.0, 13.99)]:
+        result = _get_projected_spend_over_limit(
+            current_spend=spend,
+            soft_budget_limit=soft,
+            budget_duration="7d",
+            budget_reset_at=datetime(2024, 1, 18, 0, 0, 0),
+        )
+        assert result is not None
+        _, exceed_date = result
+        assert exceed_date <= date(2024, 1, 18)