290e2528cd
* Schedule budget resets at expectable times (#10331) * Enhance budget reset functionality with timezone support and standardized reset times - Added `get_next_standardized_reset_time` function to calculate budget reset times based on specified durations and timezones. - Introduced `timezone_utils.py` to manage timezone retrieval and budget reset time calculations. - Updated budget reset logic in `reset_budget_job.py`, `internal_user_endpoints.py`, `key_management_endpoints.py`, and `team_endpoints.py` to utilize the new timezone-aware reset time calculations. - Added unit tests for the new reset time functionality in `test_duration_parser.py`. - Updated `.gitignore` to include `test.py` and made minor formatting adjustments in `docker-compose.yml` for consistency. * Fixed linting * Fix for mypy * Fixed testcase for reset * fix(duration_parser.py): move off zoneinfo - doesn't work with python 3.8 * test: update test * refactor: improve budget reset time calculation and update related tests for accuracy * clean up imports in team_endpoints.py * test: update budget remaining hours assertions to reflect new reset time logic * build(model_prices_and_context_window.json): update model --------- Co-authored-by: Prathamesh Saraf <pratamesh1867@gmail.com>
29 lines
990 B
Python
29 lines
990 B
Python
from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time
|
|
|
|
from datetime import datetime, timezone
|
|
def get_budget_reset_timezone():
|
|
"""
|
|
Get the budget reset timezone from general_settings.
|
|
Falls back to UTC if not specified.
|
|
"""
|
|
# Import at function level to avoid circular imports
|
|
from litellm.proxy.proxy_server import general_settings
|
|
if general_settings:
|
|
litellm_settings = general_settings.get("litellm_settings", {})
|
|
if litellm_settings and "timezone" in litellm_settings:
|
|
return litellm_settings["timezone"]
|
|
|
|
return "UTC"
|
|
|
|
|
|
def get_budget_reset_time(budget_duration: str):
|
|
"""
|
|
Get the budget reset time from general_settings.
|
|
Falls back to UTC if not specified.
|
|
"""
|
|
reset_at = get_next_standardized_reset_time(
|
|
duration=budget_duration,
|
|
current_time=datetime.now(timezone.utc),
|
|
timezone_str=get_budget_reset_timezone()
|
|
)
|
|
return reset_at |