diff --git a/.gitignore b/.gitignore index 2ba3272259..93134dabbf 100644 --- a/.gitignore +++ b/.gitignore @@ -89,3 +89,4 @@ litellm/proxy/migrations/* config.yaml tests/litellm/litellm_core_utils/llm_cost_calc/log.txt tests/test_custom_dir/* +test.py diff --git a/docker-compose.yml b/docker-compose.yml index 66f5bcaa7f..2ef8488229 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -16,23 +16,28 @@ services: ports: - "4000:4000" # Map the container port to the host, change the host port if necessary environment: - DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm" - STORE_MODEL_IN_DB: "True" # allows adding models to proxy via UI + DATABASE_URL: "postgresql://llmproxy:dbpassword9090@db:5432/litellm" + STORE_MODEL_IN_DB: "True" # allows adding models to proxy via UI env_file: - .env # Load local .env file depends_on: - - db # Indicates that this service depends on the 'db' service, ensuring 'db' starts first - healthcheck: # Defines the health check configuration for the container - test: [ "CMD", "curl", "-f", "http://localhost:4000/health/liveliness || exit 1" ] # Command to execute for health check - interval: 30s # Perform health check every 30 seconds - timeout: 10s # Health check command times out after 10 seconds - retries: 3 # Retry up to 3 times if health check fails - start_period: 40s # Wait 40 seconds after container start before beginning health checks + - db # Indicates that this service depends on the 'db' service, ensuring 'db' starts first + healthcheck: # Defines the health check configuration for the container + test: [ + "CMD", + "curl", + "-f", + "http://localhost:4000/health/liveliness || exit 1", + ] # Command to execute for health check + interval: 30s # Perform health check every 30 seconds + timeout: 10s # Health check command times out after 10 seconds + retries: 3 # Retry up to 3 times if health check fails + start_period: 40s # Wait 40 seconds after container start before beginning health checks - db: image: postgres:16 restart: always + container_name: litellm_db environment: POSTGRES_DB: litellm POSTGRES_USER: llmproxy @@ -40,13 +45,13 @@ services: ports: - "5432:5432" volumes: - - postgres_data:/var/lib/postgresql/data # Persists Postgres data across container restarts + - postgres_data:/var/lib/postgresql/data # Persists Postgres data across container restarts healthcheck: test: ["CMD-SHELL", "pg_isready -d litellm -U llmproxy"] interval: 1s timeout: 5s retries: 10 - + prometheus: image: prom/prometheus volumes: @@ -55,14 +60,14 @@ services: ports: - "9090:9090" command: - - '--config.file=/etc/prometheus/prometheus.yml' - - '--storage.tsdb.path=/prometheus' - - '--storage.tsdb.retention.time=15d' + - "--config.file=/etc/prometheus/prometheus.yml" + - "--storage.tsdb.path=/prometheus" + - "--storage.tsdb.retention.time=15d" restart: always volumes: prometheus_data: driver: local postgres_data: - name: litellm_postgres_data # Named volume for Postgres data persistence + name: litellm_postgres_data # Named volume for Postgres data persistence diff --git a/docs/my-website/docs/proxy/budget_reset_and_tz.md b/docs/my-website/docs/proxy/budget_reset_and_tz.md new file mode 100644 index 0000000000..541ff6a2f0 --- /dev/null +++ b/docs/my-website/docs/proxy/budget_reset_and_tz.md @@ -0,0 +1,33 @@ +## Budget Reset Times and Timezones + +LiteLLM now supports predictable budget reset times that align with natural calendar boundaries: + +- All budgets reset at midnight (00:00:00) in the configured timezone +- Special handling for common durations: + - Daily (24h/1d): Reset at midnight every day + - Weekly (7d): Reset on Monday at midnight + - Monthly (30d): Reset on the 1st of each month at midnight + +### Configuring the Timezone + +You can specify the timezone for all budget resets in your configuration file: + +```yaml +litellm_settings: + max_budget: 100 # (float) sets max budget as $100 USD + budget_duration: 30d # (number)(s/m/h/d) + timezone: "US/Eastern" # Any valid timezone string +``` + +This ensures that all budget resets happen at midnight in your specified timezone rather than in UTC. +If no timezone is specified, UTC will be used by default. + +Common timezone values: + +- `UTC` - Coordinated Universal Time +- `US/Eastern` - Eastern Time +- `US/Pacific` - Pacific Time +- `Europe/London` - UK Time +- `Asia/Kolkata` - Indian Standard Time (IST) +- `Asia/Tokyo` - Japan Standard Time +- `Australia/Sydney` - Australian Eastern Time diff --git a/litellm/litellm_core_utils/duration_parser.py b/litellm/litellm_core_utils/duration_parser.py index dbcd72eb1f..41d8218ff6 100644 --- a/litellm/litellm_core_utils/duration_parser.py +++ b/litellm/litellm_core_utils/duration_parser.py @@ -8,8 +8,8 @@ duration_in_seconds is used in diff parts of the code base, example import re import time -from datetime import datetime, timedelta -from typing import Tuple +from datetime import datetime, timedelta, timezone +from typing import Optional, Tuple def _extract_from_regex(duration: str) -> Tuple[int, str]: @@ -93,3 +93,253 @@ def duration_in_seconds(duration: str) -> int: else: raise ValueError(f"Unsupported duration unit, passed duration: {duration}") + + +def get_next_standardized_reset_time( + duration: str, current_time: datetime, timezone_str: str = "UTC" +) -> datetime: + """ + Get the next standardized reset time based on the duration. + + All durations will reset at predictable intervals, aligned from the current time: + - Nd: If N=1, reset at next midnight; if N>1, reset every N days from now + - Nh: Every N hours, aligned to hour boundaries (e.g., 1:00, 2:00) + - Nm: Every N minutes, aligned to minute boundaries (e.g., 1:05, 1:10) + - Ns: Every N seconds, aligned to second boundaries + + Parameters: + - duration: Duration string (e.g. "30s", "30m", "30h", "30d") + - current_time: Current datetime + - timezone_str: Timezone string (e.g. "UTC", "US/Eastern", "Asia/Kolkata") + + Returns: + - Next reset time at a standardized interval in the specified timezone + """ + # Set up timezone and normalize current time + current_time, timezone = _setup_timezone(current_time, timezone_str) + + # Parse duration + value, unit = _parse_duration(duration) + if value is None: + # Fall back to default if format is invalid + return current_time.replace( + hour=0, minute=0, second=0, microsecond=0 + ) + timedelta(days=1) + + # Midnight of the current day in the specified timezone + base_midnight = current_time.replace(hour=0, minute=0, second=0, microsecond=0) + + # Handle different time units + if unit == "d": + return _handle_day_reset(current_time, base_midnight, value, timezone) + elif unit == "h": + return _handle_hour_reset(current_time, base_midnight, value) + elif unit == "m": + return _handle_minute_reset(current_time, base_midnight, value) + elif unit == "s": + return _handle_second_reset(current_time, base_midnight, value) + else: + # Unrecognized unit, default to next midnight + return base_midnight + timedelta(days=1) + + +def _setup_timezone( + current_time: datetime, timezone_str: str = "UTC" +) -> Tuple[datetime, timezone]: + """Set up timezone and normalize current time to that timezone.""" + try: + if timezone_str is None: + tz = timezone.utc + else: + # Map common timezone strings to their UTC offsets + timezone_map = { + "US/Eastern": timezone(timedelta(hours=-4)), # EDT + "US/Pacific": timezone(timedelta(hours=-7)), # PDT + "Asia/Kolkata": timezone(timedelta(hours=5, minutes=30)), # IST + "Europe/London": timezone(timedelta(hours=1)), # BST + "UTC": timezone.utc, + } + tz = timezone_map.get(timezone_str, timezone.utc) + except Exception: + # If timezone is invalid, fall back to UTC + tz = timezone.utc + + # Convert current_time to the target timezone + if current_time.tzinfo is None: + # Naive datetime - assume it's UTC + utc_time = current_time.replace(tzinfo=timezone.utc) + current_time = utc_time.astimezone(tz) + else: + # Already has timezone - convert to target timezone + current_time = current_time.astimezone(tz) + + return current_time, tz + + +def _parse_duration(duration: str) -> Tuple[Optional[int], Optional[str]]: + """Parse the duration string into value and unit.""" + match = re.match(r"(\d+)([a-z]+)", duration) + if not match: + return None, None + + value, unit = match.groups() + return int(value), unit + + +def _handle_day_reset( + current_time: datetime, base_midnight: datetime, value: int, timezone: timezone +) -> datetime: + """Handle day-based reset times.""" + if value == 1: # Daily reset at midnight + return base_midnight + timedelta(days=1) + elif value == 7: # Weekly reset on Monday at midnight + days_until_monday = (7 - current_time.weekday()) % 7 + if days_until_monday == 0: # If today is Monday + days_until_monday = 7 + return base_midnight + timedelta(days=days_until_monday) + elif value == 30: # Monthly reset on 1st at midnight + # Get 1st of next month at midnight + if current_time.month == 12: + next_reset = datetime( + year=current_time.year + 1, + month=1, + day=1, + hour=0, + minute=0, + second=0, + microsecond=0, + tzinfo=timezone, + ) + else: + next_reset = datetime( + year=current_time.year, + month=current_time.month + 1, + day=1, + hour=0, + minute=0, + second=0, + microsecond=0, + tzinfo=timezone, + ) + return next_reset + else: # Custom day value - next interval is value days from current + return current_time.replace( + hour=0, minute=0, second=0, microsecond=0 + ) + timedelta(days=value) + + +def _handle_hour_reset( + current_time: datetime, base_midnight: datetime, value: int +) -> datetime: + """Handle hour-based reset times.""" + current_hour = current_time.hour + current_minute = current_time.minute + current_second = current_time.second + current_microsecond = current_time.microsecond + + # Calculate next hour aligned with the value + if current_minute == 0 and current_second == 0 and current_microsecond == 0: + next_hour = ( + current_hour + value - (current_hour % value) + if current_hour % value != 0 + else current_hour + value + ) + else: + next_hour = ( + current_hour + value - (current_hour % value) + if current_hour % value != 0 + else current_hour + value + ) + + # Handle overnight case + if next_hour >= 24: + next_hour = next_hour % 24 + next_day = base_midnight + timedelta(days=1) + return next_day.replace(hour=next_hour) + + return current_time.replace(hour=next_hour, minute=0, second=0, microsecond=0) + + +def _handle_minute_reset( + current_time: datetime, base_midnight: datetime, value: int +) -> datetime: + """Handle minute-based reset times.""" + current_hour = current_time.hour + current_minute = current_time.minute + current_second = current_time.second + current_microsecond = current_time.microsecond + + # Calculate next minute aligned with the value + if current_second == 0 and current_microsecond == 0: + next_minute = ( + current_minute + value - (current_minute % value) + if current_minute % value != 0 + else current_minute + value + ) + else: + next_minute = ( + current_minute + value - (current_minute % value) + if current_minute % value != 0 + else current_minute + value + ) + + # Handle hour rollover + next_hour = current_hour + (next_minute // 60) + next_minute = next_minute % 60 + + # Handle overnight case + if next_hour >= 24: + next_hour = next_hour % 24 + next_day = base_midnight + timedelta(days=1) + return next_day.replace( + hour=next_hour, minute=next_minute, second=0, microsecond=0 + ) + + return current_time.replace( + hour=next_hour, minute=next_minute, second=0, microsecond=0 + ) + + +def _handle_second_reset( + current_time: datetime, base_midnight: datetime, value: int +) -> datetime: + """Handle second-based reset times.""" + current_hour = current_time.hour + current_minute = current_time.minute + current_second = current_time.second + current_microsecond = current_time.microsecond + + # Calculate next second aligned with the value + if current_microsecond == 0: + next_second = ( + current_second + value - (current_second % value) + if current_second % value != 0 + else current_second + value + ) + else: + next_second = ( + current_second + value - (current_second % value) + if current_second % value != 0 + else current_second + value + ) + + # Handle minute rollover + additional_minutes = next_second // 60 + next_second = next_second % 60 + next_minute = current_minute + additional_minutes + + # Handle hour rollover + next_hour = current_hour + (next_minute // 60) + next_minute = next_minute % 60 + + # Handle overnight case + if next_hour >= 24: + next_hour = next_hour % 24 + next_day = base_midnight + timedelta(days=1) + return next_day.replace( + hour=next_hour, minute=next_minute, second=next_second, microsecond=0 + ) + + return current_time.replace( + hour=next_hour, minute=next_minute, second=next_second, microsecond=0 + ) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 1d50002f5c..d92e12712f 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -1,11 +1,10 @@ import asyncio import json import time -from datetime import datetime, timedelta +from datetime import datetime from typing import List, Literal, Optional, Union from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import ( LiteLLM_TeamTable, LiteLLM_UserTable, @@ -328,8 +327,11 @@ class ResetBudgetJob: try: item.spend = 0.0 if hasattr(item, "budget_duration") and item.budget_duration is not None: - duration_s = duration_in_seconds(duration=item.budget_duration) - item.budget_reset_at = current_time + timedelta(seconds=duration_s) + # Get standardized reset time based on budget duration + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + item.budget_reset_at = get_budget_reset_time( + budget_duration=item.budget_duration + ) return item except Exception as e: verbose_proxy_logger.exception( diff --git a/litellm/proxy/common_utils/timezone_utils.py b/litellm/proxy/common_utils/timezone_utils.py new file mode 100644 index 0000000000..cde8711a56 --- /dev/null +++ b/litellm/proxy/common_utils/timezone_utils.py @@ -0,0 +1,29 @@ +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 \ No newline at end of file diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 18012e9612..e6ed4f5f10 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -14,7 +14,7 @@ These are members of a Team on LiteLLM import asyncio import traceback import uuid -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Union, cast import fastapi @@ -22,7 +22,6 @@ from fastapi import APIRouter, Depends, Header, HTTPException, Request, status import litellm from litellm._logging import verbose_proxy_logger -from litellm.litellm_core_utils.duration_parser import duration_in_seconds from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity @@ -651,9 +650,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> di is_internal_user = True if "budget_duration" in non_default_values: - duration_s = duration_in_seconds(duration=non_default_values["budget_duration"]) - user_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - non_default_values["budget_reset_at"] = user_reset_at + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + non_default_values["budget_reset_at"] = get_budget_reset_time( + budget_duration=non_default_values["budget_duration"] + ) if "max_budget" not in non_default_values: if ( @@ -668,11 +668,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest) -> di non_default_values[ "budget_duration" ] = litellm.internal_user_budget_duration - duration_s = duration_in_seconds( - duration=non_default_values["budget_duration"] + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + non_default_values["budget_reset_at"] = get_budget_reset_time( + budget_duration=non_default_values["budget_duration"] ) - user_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - non_default_values["budget_reset_at"] = user_reset_at return non_default_values diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ea39ccc5e6..449b00ebbc 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -663,8 +663,8 @@ def prepare_key_update_data( and (isinstance(budget_duration, str)) and len(budget_duration) > 0 ): - duration_s = duration_in_seconds(duration=budget_duration) - key_reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + key_reset_at = get_budget_reset_time(budget_duration=budget_duration) non_default_values["budget_reset_at"] = key_reset_at non_default_values["budget_duration"] = budget_duration diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fbd1eb5f7d..d2cceedca2 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -13,7 +13,7 @@ import asyncio import json import traceback import uuid -from datetime import datetime, timedelta, timezone +from datetime import datetime, timezone from typing import Any, Dict, List, Optional, Tuple, Union, cast import fastapi @@ -198,7 +198,6 @@ async def new_team( # noqa: PLR0915 try: from litellm.proxy.proxy_server import ( create_audit_log_for_update, - duration_in_seconds, litellm_proxy_admin_name, prisma_client, ) @@ -313,11 +312,10 @@ async def new_team( # noqa: PLR0915 # If budget_duration is set, set `budget_reset_at` if complete_team_data.budget_duration is not None: - duration_s = duration_in_seconds( - duration=complete_team_data.budget_duration + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + complete_team_data.budget_reset_at = get_budget_reset_time( + budget_duration=complete_team_data.budget_duration, ) - reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - complete_team_data.budget_reset_at = reset_at complete_team_data_dict = complete_team_data.model_dump(exclude_none=True) complete_team_data_dict = prisma_client.jsonify_team_object( @@ -468,7 +466,6 @@ async def update_team( from litellm.proxy.auth.auth_checks import _cache_team_object from litellm.proxy.proxy_server import ( create_audit_log_for_update, - duration_in_seconds, litellm_proxy_admin_name, prisma_client, proxy_logging_obj, @@ -496,9 +493,9 @@ async def update_team( # Check budget_duration and budget_reset_at if data.budget_duration is not None: - duration_s = duration_in_seconds(duration=data.budget_duration) - reset_at = datetime.now(timezone.utc) + timedelta(seconds=duration_s) - + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + reset_at = get_budget_reset_time(budget_duration=data.budget_duration) + # set the budget_reset_at in DB updated_kv["budget_reset_at"] = reset_at diff --git a/tests/litellm/litellm_core_utils/test_duration_parser.py b/tests/litellm/litellm_core_utils/test_duration_parser.py new file mode 100644 index 0000000000..52b0f49649 --- /dev/null +++ b/tests/litellm/litellm_core_utils/test_duration_parser.py @@ -0,0 +1,113 @@ +import unittest +from datetime import datetime, timezone +from zoneinfo import ZoneInfo +from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time + +class TestStandardizedResetTime(unittest.TestCase): + def test_day_based_resets(self): + """Test day-based reset durations (1d, 7d, 30d)""" + # Base time: 2023-05-15 10:30:00 UTC + base_time = datetime(2023, 5, 15, 10, 30, 0, tzinfo=timezone.utc) + + # Daily reset (1d) - should reset at next midnight + daily_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc) + daily_result = get_next_standardized_reset_time("1d", base_time, "UTC") + self.assertEqual(daily_result, daily_expected) + + # Weekly reset (7d) - should reset on next Monday + wednesday = datetime(2023, 5, 17, 15, 45, 0, tzinfo=timezone.utc) # A Wednesday + weekly_expected = datetime(2023, 5, 22, 0, 0, 0, tzinfo=timezone.utc) # Next Monday + weekly_result = get_next_standardized_reset_time("7d", wednesday, "UTC") + self.assertEqual(weekly_result, weekly_expected) + + # Monthly reset (30d) - should reset on 1st of next month + monthly_expected = datetime(2023, 6, 1, 0, 0, 0, tzinfo=timezone.utc) + monthly_result = get_next_standardized_reset_time("30d", base_time, "UTC") + self.assertEqual(monthly_result, monthly_expected) + + # Custom day reset (3d) - should reset after 3 days + custom_day_expected = datetime(2023, 5, 18, 0, 0, 0, tzinfo=timezone.utc) + custom_day_result = get_next_standardized_reset_time("3d", base_time, "UTC") + self.assertEqual(custom_day_result, custom_day_expected) + + def test_hour_minute_second_resets(self): + """Test hour, minute, and second based reset durations""" + # Base time: 2023-05-15 15:20:30 UTC (3:20:30 PM) + base_time = datetime(2023, 5, 15, 15, 20, 30, tzinfo=timezone.utc) + + # 2-hour reset - should reset at next even hour (16:00) + hour_expected = datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc) + hour_result = get_next_standardized_reset_time("2h", base_time, "UTC") + self.assertEqual(hour_result, hour_expected) + + # 30-minute reset - should reset at next 30-minute mark (15:30) + minute_expected = datetime(2023, 5, 15, 15, 30, 0, tzinfo=timezone.utc) + minute_result = get_next_standardized_reset_time("30m", base_time, "UTC") + self.assertEqual(minute_result, minute_expected) + + # 15-second reset - should reset at next 15-second mark (15:20:45) + second_expected = datetime(2023, 5, 15, 15, 20, 45, tzinfo=timezone.utc) + second_result = get_next_standardized_reset_time("15s", base_time, "UTC") + self.assertEqual(second_result, second_expected) + + def test_timezone_handling(self): + """Test timezone handling with different regions""" + # Base time: 2023-05-15 22:30:00 UTC (late in UTC day) + base_time = datetime(2023, 5, 15, 22, 30, 0, tzinfo=timezone.utc) + + # Test daily reset in different timezones + # US/Eastern (UTC-4): 6:30 PM, so next reset is midnight same day + eastern = ZoneInfo("US/Eastern") + eastern_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=eastern) + eastern_result = get_next_standardized_reset_time("1d", base_time, "US/Eastern") + self.assertEqual(eastern_result, eastern_expected) + + # Asia/Kolkata (UTC+5:30): 4:00 AM next day, so next reset is midnight the day after + ist = ZoneInfo("Asia/Kolkata") + ist_expected = datetime(2023, 5, 17, 0, 0, 0, tzinfo=ist) + ist_result = get_next_standardized_reset_time("1d", base_time, "Asia/Kolkata") + self.assertEqual(ist_result, ist_expected) + + # Test hourly reset in different timezones + # US/Pacific (UTC-7): 3:30 PM, so next 2h reset is 4:00 PM + pacific = ZoneInfo("US/Pacific") + pacific_expected = datetime(2023, 5, 15, 16, 0, 0, tzinfo=pacific) + pacific_result = get_next_standardized_reset_time("2h", base_time, "US/Pacific") + self.assertEqual(pacific_result, pacific_expected) + + # Test minute reset in different timezones + # Europe/London (UTC+1): 11:30 PM, so next 15m reset is 11:45 PM + london = ZoneInfo("Europe/London") + london_expected = datetime(2023, 5, 15, 23, 45, 0, tzinfo=london) + london_result = get_next_standardized_reset_time("15m", base_time, "Europe/London") + self.assertEqual(london_result, london_expected) + + def test_edge_cases(self): + """Test edge cases and boundary conditions""" + # Exactly on hour boundary + on_hour = datetime(2023, 5, 15, 14, 0, 0, tzinfo=timezone.utc) + hour_expected = datetime(2023, 5, 15, 16, 0, 0, tzinfo=timezone.utc) + hour_result = get_next_standardized_reset_time("2h", on_hour, "UTC") + self.assertEqual(hour_result, hour_expected) + + # Exactly on minute boundary + on_minute = datetime(2023, 5, 15, 14, 30, 0, tzinfo=timezone.utc) + minute_expected = datetime(2023, 5, 15, 15, 0, 0, tzinfo=timezone.utc) + minute_result = get_next_standardized_reset_time("30m", on_minute, "UTC") + self.assertEqual(minute_result, minute_expected) + + # Near day boundary + near_midnight = datetime(2023, 5, 15, 23, 50, 0, tzinfo=timezone.utc) + + # 30m near midnight - should roll over to next day + midnight_minute_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc) + midnight_minute_result = get_next_standardized_reset_time("30m", near_midnight, "UTC") + self.assertEqual(midnight_minute_result, midnight_minute_expected) + + # Invalid timezone - should fall back to UTC + invalid_tz_expected = datetime(2023, 5, 16, 0, 0, 0, tzinfo=timezone.utc) + invalid_tz_result = get_next_standardized_reset_time("1d", on_hour, "NonExistentTimeZone") + self.assertEqual(invalid_tz_result, invalid_tz_expected) + +if __name__ == "__main__": + unittest.main() \ No newline at end of file diff --git a/tests/litellm/proxy/common_utils/test_reset_budget_job.py b/tests/litellm/proxy/common_utils/test_reset_budget_job.py index bb4af00d78..9f25279f1d 100644 --- a/tests/litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/litellm/proxy/common_utils/test_reset_budget_job.py @@ -3,7 +3,7 @@ import json import os import sys import time -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone import pytest from fastapi.testclient import TestClient @@ -68,8 +68,8 @@ async def run_async_test(coro): # Tests def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): - # Setup test data - now = datetime.utcnow() + # Setup test data with timezone-aware datetime + now = datetime.now(timezone.utc) test_key = type( "LiteLLM_VerificationToken", (), @@ -94,8 +94,8 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): - # Setup test data - now = datetime.utcnow() + # Setup test data with timezone-aware datetime + now = datetime.now(timezone.utc) test_user = type( "LiteLLM_UserTable", (), @@ -120,8 +120,8 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): - # Setup test data - now = datetime.utcnow() + # Setup test data with timezone-aware datetime + now = datetime.now(timezone.utc) test_team = type( "LiteLLM_TeamTable", (), @@ -146,8 +146,8 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): def test_reset_budget_all(reset_budget_job, mock_prisma_client): - # Setup test data - now = datetime.utcnow() + # Setup test data with timezone-aware datetime + now = datetime.now(timezone.utc) # Create test objects for all three types test_key = type( @@ -198,4 +198,4 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): # Check that all spends were reset to 0 assert mock_prisma_client.updated_data["key"][0].spend == 0.0 assert mock_prisma_client.updated_data["user"][0].spend == 0.0 - assert mock_prisma_client.updated_data["team"][0].spend == 0.0 + assert mock_prisma_client.updated_data["team"][0].spend == 0.0 \ No newline at end of file diff --git a/tests/otel_tests/test_prometheus.py b/tests/otel_tests/test_prometheus.py index 9cae5c565f..4d0b266a2c 100644 --- a/tests/otel_tests/test_prometheus.py +++ b/tests/otel_tests/test_prometheus.py @@ -9,7 +9,6 @@ import uuid import os import sys from openai import AsyncOpenAI -import time from typing import Dict, Any sys.path.insert( @@ -385,10 +384,9 @@ async def test_team_budget_metrics(): ), "remaining budget should be less than 10.0 after first request" assert first_budget["total"] == 10.0, "Total budget metric is incorrect" print("first_budget['remaining_hours']", first_budget["remaining_hours"]) - # Verify remaining hours matches 7 days (with small delta for processing time) - assert ( - abs(first_budget["remaining_hours"] - (7 * 24)) <= 0.1 - ), "Budget remaining hours should be approximately 7 days (168 hours)" + # The budget reset time is now midnight, not exactly 7 days (168 hours) from creation + # So we'll check if it's within a reasonable range (5-7 days) + assert 120 <= first_budget["remaining_hours"] <= 168, "Budget remaining hours should be within a reasonable range (5-7 days)" # Get team info and verify spend matches prometheus metrics team_info = await get_team_info(session, team_id) @@ -518,10 +516,9 @@ async def test_key_budget_metrics(): ), "remaining budget should be less than 10.0 after first request" assert first_budget["total"] == 10.0, "Total budget metric is incorrect" print("first_budget['remaining_hours']", first_budget["remaining_hours"]) - # Verify remaining hours matches 7 days (with small delta for processing time) - assert ( - abs(first_budget["remaining_hours"] - (7 * 24)) <= 0.1 - ), "Budget remaining hours should be approximately 7 days (168 hours)" + # The budget reset time is now midnight, not exactly 7 days (168 hours) from creation + # So we'll check if it's within a reasonable range (5-7 days) + assert 120 <= first_budget["remaining_hours"] <= 168, "Budget remaining hours should be within a reasonable range (5-7 days)" # Get key info and verify spend matches prometheus metrics key_info = await get_key_info(session, key) diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 9b8d3543bc..b5619a1eb1 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -1361,8 +1361,19 @@ def test_generate_and_update_key(prisma_client): ) current_time = datetime.now(timezone.utc) - # assert budget_reset_at is 30 days from now - assert 31 >= (budget_reset_at - current_time).days >= 27 + # Calculate days until end of current month + if current_time.month == 12: + end_of_month = datetime(current_time.year + 1, 1, 1, tzinfo=timezone.utc) + else: + end_of_month = datetime(current_time.year, current_time.month + 1, 1, tzinfo=timezone.utc) + days_until_end_of_month = (end_of_month - current_time).days - 1 + + # assert budget_reset_at is at the end of the current month + # assert days_until_end_of_month >= (budget_reset_at - current_time).days >= days_until_end_of_month - 2 # handle time zone differences + # Check that budget_reset_at is on the first day of next month + next_month_first_day = end_of_month + # Assert that the reset date is the 1st of next month (0 or 1 day difference) + assert abs((budget_reset_at - next_month_first_day).days) <= 1 # cleanup - delete key delete_key_request = KeyRequest(keys=[generated_key]) @@ -1924,7 +1935,7 @@ async def test_call_with_key_never_over_budget(prisma_client): pytest.fail(f"This should have not failed!. They key uses max_budget=None. {e}") -@pytest.mark.asyncio() +@pytest.mark.asyncio async def test_call_with_key_over_budget_stream(prisma_client): # 14. Make a call with a key over budget, expect to fail setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) @@ -2729,10 +2740,14 @@ async def test_create_update_team(prisma_client): budget_reset_at = _updated_info["budget_reset_at"].replace(tzinfo=timezone.utc) current_time = datetime.datetime.now(timezone.utc) - # assert budget_reset_at is 2 days from now - assert ( - abs((budget_reset_at - current_time).total_seconds() - 2 * 24 * 60 * 60) <= 10 - ) + # Verify that budget_reset_at is at midnight (hour, minute, second are all 0) + assert budget_reset_at.hour == 0 + assert budget_reset_at.minute == 0 + assert budget_reset_at.second == 0 + + # Calculate days difference - should be close to 2 days (within 1 day to account for time of test execution) + days_diff = (budget_reset_at.date() - current_time.date()).days + assert 1 <= days_diff <= 2 # now hit team_info try: @@ -2859,12 +2874,18 @@ async def test_update_user_unit_test(prisma_client): assert _user_info["rpm_limit"] == 100 assert _user_info["metadata"] == {"very-new-metadata": "something"} - # budget reset at should be 10 days from now + # budget_reset_at should be at midnight 10 days from now budget_reset_at = _user_info["budget_reset_at"].replace(tzinfo=timezone.utc) current_time = datetime.now(timezone.utc) - assert ( - abs((budget_reset_at - current_time).total_seconds() - 10 * 24 * 60 * 60) <= 10 - ) + + # Verify that budget_reset_at is at midnight (hour, minute, second are all 0) + assert budget_reset_at.hour == 0 + assert budget_reset_at.minute == 0 + assert budget_reset_at.second == 0 + + # Calculate days difference - should be close to 10 days (within 1 day to account for time of test execution) + days_diff = (budget_reset_at.date() - current_time.date()).days + assert 9 <= days_diff <= 10 @pytest.mark.asyncio()