forked from phoenix/litellm-mirror
92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
"""
|
|
Helper utilities for parsing durations - 1s, 1d, 10d, 30d, 1mo, 2mo
|
|
|
|
duration_in_seconds is used in diff parts of the code base, example
|
|
- Router - Provider budget routing
|
|
- Proxy - Key, Team Generation
|
|
"""
|
|
|
|
import re
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from typing import Tuple
|
|
|
|
|
|
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 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:
|
|
- 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":
|
|
return value * 60
|
|
elif unit == "h":
|
|
return value * 3600
|
|
elif unit == "d":
|
|
return value * 86400
|
|
elif unit == "mo":
|
|
now = time.time()
|
|
current_time = datetime.fromtimestamp(now)
|
|
|
|
if current_time.month == 12:
|
|
target_year = current_time.year + 1
|
|
target_month = 1
|
|
else:
|
|
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
|
|
return int(duration_until_next_month.total_seconds())
|
|
|
|
else:
|
|
raise ValueError(f"Unsupported duration unit, passed duration: {duration}")
|