diff --git a/docs/my-website/docs/proxy/team_budgets.md b/docs/my-website/docs/proxy/team_budgets.md index 9e897d507e..54f5efddff 100644 --- a/docs/my-website/docs/proxy/team_budgets.md +++ b/docs/my-website/docs/proxy/team_budgets.md @@ -46,7 +46,7 @@ Possible values for `budget_duration` | `budget_duration="1m"` | every 1 min | | `budget_duration="1h"` | every 1 hour | | `budget_duration="1d"` | every 1 day | -| `budget_duration="1mo"` | start of every month | +| `budget_duration="1mo"` | every 1 month | ### 2. Create a key for the `team` diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e40440ee35..ebe30789d7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2115,6 +2115,16 @@ def _extract_from_regex(duration: str) -> Tuple[int, str]: return value, unit +def get_last_day_of_month(year, month): + # Handle December case + if month == 12: + return 31 + # Next month is January, so subtract a day from March 1st + next_month = datetime(year=year, month=month + 1, day=1) + last_day_of_month = (next_month - timedelta(days=1)).day + return last_day_of_month + + def _duration_in_seconds(duration: str) -> int: """ Parameters: @@ -2141,13 +2151,29 @@ def _duration_in_seconds(duration: str) -> int: 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) + target_year = current_time.year + 1 + target_month = 1 else: - next_month = datetime( - year=current_time.year, month=current_time.month + value, day=1 - ) + target_year = current_time.year + target_month = current_time.month + value + + # Determine the day to set for next month + target_day = current_time.day + last_day_of_target_month = get_last_day_of_month(target_year, target_month) + + if target_day > last_day_of_target_month: + target_day = last_day_of_target_month + + next_month = datetime( + year=target_year, + month=target_month, + day=target_day, + hour=current_time.hour, + minute=current_time.minute, + second=current_time.second, + microsecond=current_time.microsecond, + ) # Calculate the duration until the first day of the next month duration_until_next_month = next_month - current_time diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py index 742199c7f9..7d581a0fb0 100644 --- a/litellm/tests/test_utils.py +++ b/litellm/tests/test_utils.py @@ -26,7 +26,11 @@ from litellm.utils import ( get_max_tokens, get_supported_openai_params, ) -from litellm.proxy.utils import _duration_in_seconds, _extract_from_regex +from litellm.proxy.utils import ( + _duration_in_seconds, + _extract_from_regex, + get_last_day_of_month, +) # Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils' @@ -467,15 +471,31 @@ def test_duration_in_seconds(): 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) + target_year = current_time.year + 1 + target_month = 1 else: - next_month = datetime( - year=current_time.year, month=current_time.month + 1, day=1 - ) - print("next_month={}".format(next_month)) + target_year = current_time.year + target_month = current_time.month + 1 + + # Determine the day to set for next month + target_day = current_time.day + last_day_of_target_month = get_last_day_of_month(target_year, target_month) + + if target_day > last_day_of_target_month: + target_day = last_day_of_target_month + + next_month = datetime( + year=target_year, + month=target_month, + day=target_day, + hour=current_time.hour, + minute=current_time.minute, + second=current_time.second, + microsecond=current_time.microsecond, + ) + # 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())