diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 606a590c24b..e5b60452ecd 100644 --- a/litellm/proxy/proxy_server.py +++ b/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 "", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 89625021e37..077393d2a3c 100644 --- a/litellm/proxy/utils.py +++ b/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): diff --git a/tests/test_litellm/proxy/utils/helpers/test_month_end_projection.py b/tests/test_litellm/proxy/utils/helpers/test_month_end_projection.py index 5afe1f4faf8..bbdc0124a63 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_month_end_projection.py +++ b/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)