diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e0855333e26..5dcb0339739 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -7,7 +7,7 @@ import threading import time import traceback -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText from typing import ( @@ -3891,11 +3891,15 @@ def _raise_failed_update_spend_exception( raise e +def _get_month_end_date(today: date) -> date: + if today.month == 12: + return date(today.year + 1, 1, 1) - timedelta(days=1) + return date(today.year, today.month + 1, 1) - timedelta(days=1) + + def _is_projected_spend_over_limit( current_spend: float, soft_budget_limit: Optional[float] ): - from datetime import date - if soft_budget_limit is None: # If there's no limit, we can't exceed it. return False @@ -3903,10 +3907,7 @@ def _is_projected_spend_over_limit( today = date.today() # Finding the first day of the next month, then subtracting one day to get the end of the current month. - if today.month == 12: # December edge case - end_month = date(today.year + 1, 1, 1) - timedelta(days=1) - else: - end_month = date(today.year, today.month + 1, 1) - timedelta(days=1) + end_month = _get_month_end_date(today) remaining_days = (end_month - today).days @@ -3928,25 +3929,30 @@ def _is_projected_spend_over_limit( def _get_projected_spend_over_limit( current_spend: float, soft_budget_limit: Optional[float] ) -> Optional[tuple]: - import datetime - if soft_budget_limit is None: return None - today = datetime.date.today() - end_month = datetime.date(today.year, today.month + 1, 1) - datetime.timedelta( - days=1 - ) + today = date.today() + end_month = _get_month_end_date(today) remaining_days = (end_month - today).days - daily_spend = current_spend / ( - today.day - 1 - ) # assuming the current spend till today (not including today) - projected_spend = daily_spend * remaining_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) + projected_spend = current_spend + (daily_spend * remaining_days) if projected_spend > soft_budget_limit: - approx_days = soft_budget_limit / daily_spend - limit_exceed_date = today + datetime.timedelta(days=approx_days) + if daily_spend <= 0: + limit_exceed_date = today + else: + remaining_budget = soft_budget_limit - current_spend + if remaining_budget <= 0: + limit_exceed_date = today + else: + approx_days = remaining_budget / daily_spend + limit_exceed_date = today + timedelta(days=approx_days) # return the projected spend and the date it will exceeded return projected_spend, limit_exceed_date diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index 9d0d5e6c0f3..7deda21c215 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1,3 +1,4 @@ +import datetime as real_datetime import json import os import sys @@ -132,3 +133,48 @@ def test_join_paths_nested_path(): """Test path joining with nested paths""" result = join_paths(base_path="http://0.0.0.0:4000/v1", route="chat/completions") assert result == "http://0.0.0.0:4000/v1/chat/completions" + + +def _patch_today(monkeypatch, year, month, day): + class PatchedDate(real_datetime.date): + @classmethod + def today(cls): + return real_datetime.date(year, month, day) + + monkeypatch.setattr("litellm.proxy.utils.date", PatchedDate) + + +def test_get_projected_spend_over_limit_day_one(monkeypatch): + from litellm.proxy.utils import _get_projected_spend_over_limit + + _patch_today(monkeypatch, 2026, 1, 1) + result = _get_projected_spend_over_limit(100.0, 1.0) + + assert result is not None + projected_spend, projected_exceeded_date = result + assert projected_spend == 3100.0 + assert projected_exceeded_date == real_datetime.date(2026, 1, 1) + + +def test_get_projected_spend_over_limit_december(monkeypatch): + from litellm.proxy.utils import _get_projected_spend_over_limit + + _patch_today(monkeypatch, 2026, 12, 15) + result = _get_projected_spend_over_limit(100.0, 1.0) + + assert result is not None + projected_spend, projected_exceeded_date = result + assert projected_spend == pytest.approx(214.28571428571428) + assert projected_exceeded_date == real_datetime.date(2026, 12, 15) + + +def test_get_projected_spend_over_limit_includes_current_spend(monkeypatch): + from litellm.proxy.utils import _get_projected_spend_over_limit + + _patch_today(monkeypatch, 2026, 4, 11) + result = _get_projected_spend_over_limit(100.0, 200.0) + + assert result is not None + projected_spend, projected_exceeded_date = result + assert projected_spend == 290.0 + assert projected_exceeded_date == real_datetime.date(2026, 4, 21)