fix(proxy/utils.py): fix reset monthly budget

fix to reset at the same time each month (not at start of month)
This commit is contained in:
Krrish Dholakia 2024-06-14 14:39:36 -07:00
parent 056913fd70
commit d8acda9f39
3 changed files with 60 additions and 14 deletions

View file

@ -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())