From f65752c18b6914b17e128bfede0e53e06551f0ae Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 13 Jun 2024 16:52:17 -0700
Subject: [PATCH] feat(proxy/utils.py): allow budget duration in months

Closes https://github.com/BerriAI/litellm/issues/4042
---
 litellm/proxy/utils.py      | 40 ++++++++++++++++++++++++++++++++++---
 litellm/tests/test_utils.py | 38 +++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index 54782c0887e..ba5df81b7c6 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -1,4 +1,4 @@
-from typing import Optional, List, Any, Literal, Union, TYPE_CHECKING
+from typing import Optional, List, Any, Literal, Union, TYPE_CHECKING, Tuple
 import os
 import subprocess
 import hashlib
@@ -2093,14 +2093,32 @@ def get_logging_payload(
         raise e
 
 
-def _duration_in_seconds(duration: str):
-    match = re.match(r"(\d+)([smhd]?)", duration)
+def _extract_from_regex(duration: str) -> Tuple[int, str]:
+    match = re.match(r"(\d+)(mo|[smhd]?)", duration)
+
     if not match:
         raise ValueError("Invalid duration format")
 
     value, unit = match.groups()
     value = int(value)
 
+    return value, unit
+
+
+def _duration_in_seconds(duration: str) -> int:
+    """
+    Parameters:
+    - duration:
+        - "<number>s" - seconds
+        - "<number>m" - minutes
+        - "<number>h" - hours
+        - "<number>d" - days
+        - "<number>mo" - months
+
+    Returns time in seconds till when budget needs to be reset
+    """
+    value, unit = _extract_from_regex(duration=duration)
+
     if unit == "s":
         return value
     elif unit == "m":
@@ -2109,6 +2127,22 @@ def _duration_in_seconds(duration: str):
         return value * 3600
     elif unit == "d":
         return value * 86400
+    elif unit == "mo":
+        now = time.time()
+        current_time = datetime.fromtimestamp(now)
+
+        # Calculate the first day of the next month
+        if current_time.month == 12:
+            next_month = datetime(year=current_time.year + 1, month=1, day=1)
+        else:
+            next_month = datetime(
+                year=current_time.year, month=current_time.month + value, day=1
+            )
+
+        # Calculate the duration until the first day of the next month
+        duration_until_next_month = next_month - current_time
+        return int(duration_until_next_month.total_seconds())
+
     else:
         raise ValueError("Unsupported duration unit")
 
diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py
index 2e32e32df7c..742199c7f9c 100644
--- a/litellm/tests/test_utils.py
+++ b/litellm/tests/test_utils.py
@@ -26,6 +26,7 @@ from litellm.utils import (
     get_max_tokens,
     get_supported_openai_params,
 )
+from litellm.proxy.utils import _duration_in_seconds, _extract_from_regex
 
 # Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils'
 
@@ -445,3 +446,40 @@ def test_redact_msgs_from_logs():
 
     litellm.turn_off_message_logging = False
     print("Test passed")
+
+
+@pytest.mark.parametrize(
+    "duration, unit",
+    [("7s", "s"), ("7m", "m"), ("7h", "h"), ("7d", "d"), ("7mo", "mo")],
+)
+def test_extract_from_regex(duration, unit):
+    value, _unit = _extract_from_regex(duration=duration)
+
+    assert value == 7
+    assert _unit == unit
+
+
+def test_duration_in_seconds():
+    """
+    Test if duration int is correctly calculated for different str
+    """
+    import time
+
+    now = time.time()
+    current_time = datetime.fromtimestamp(now)
+    print("current_time={}".format(current_time))
+    # Calculate the first day of the next month
+    if current_time.month == 12:
+        next_month = datetime(year=current_time.year + 1, month=1, day=1)
+    else:
+        next_month = datetime(
+            year=current_time.year, month=current_time.month + 1, day=1
+        )
+    print("next_month={}".format(next_month))
+    # Calculate the duration until the first day of the next month
+    duration_until_next_month = next_month - current_time
+    expected_duration = int(duration_until_next_month.total_seconds())
+
+    value = _duration_in_seconds(duration="1mo")
+
+    assert value - expected_duration < 2