Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/timezone_utils.py: 63%

1from datetime import datetime, time, timezone 

2 

3from pydantic import BaseModel, ConfigDict 

4 

5import litellm 

6from litellm.litellm_core_utils.duration_parser import get_next_standardized_reset_time 

7 

8 

9class BudgetResetSettings(BaseModel): 

10 """Immutable, validated settings that govern when budgets reset. 

11 

12 Parsed once from `litellm_settings` and injected into consumers (the reset 

13 job, management endpoints) so reset times never depend on reaching into 

14 module-level globals at call time. 

15 """ 

16 

17 model_config = ConfigDict(frozen=True) 

18 

19 timezone: str = "UTC" 

20 reset_time_of_day: time = time(0, 0) 

21 

22 

23def parse_budget_reset_time(raw: object) -> time: 

24 """Parse a `budget_reset_time` config value (e.g. "12:00") into a `time`. 

25 

26 Falls back to midnight when unset; raises a clear error on a malformed value 

27 so a bad config fails loudly at startup instead of silently resetting at midnight. 

28 """ 

29 if raw is None or raw == "": 29 ↛ 31line 29 didn't jump to line 31 because the condition on line 29 was always true

30 return time(0, 0) 

31 if not isinstance(raw, str): 

32 raise ValueError(f"Invalid budget_reset_time {raw!r}; must be a quoted 24-hour 'HH:MM' string, e.g. \"12:00\"") 

33 for fmt in ("%H:%M", "%H:%M:%S"): 

34 try: 

35 parsed = datetime.strptime(raw, fmt) 

36 return time(hour=parsed.hour, minute=parsed.minute, second=parsed.second) 

37 except ValueError: 

38 continue 

39 raise ValueError( 

40 f"Invalid budget_reset_time {raw!r}; expected a 24-hour 'HH:MM' or 'HH:MM:SS' string, e.g. \"12:00\"" 

41 ) 

42 

43 

44def get_budget_reset_timezone() -> str: 

45 """ 

46 Get the budget reset timezone from litellm_settings. 

47 Falls back to UTC if not specified. 

48 

49 litellm_settings values are set as attributes on the litellm module 

50 by proxy_server.py at startup (via setattr(litellm, key, value)). 

51 """ 

52 return getattr(litellm, "timezone", None) or "UTC" 

53 

54 

55def get_budget_reset_settings() -> BudgetResetSettings: 

56 """Build validated reset settings from litellm_settings. Raises on a malformed 

57 `budget_reset_time`, which lets the proxy fail fast at startup.""" 

58 return BudgetResetSettings( 

59 timezone=get_budget_reset_timezone(), 

60 reset_time_of_day=parse_budget_reset_time(getattr(litellm, "budget_reset_time", None)), 

61 ) 

62 

63 

64def compute_budget_reset_at(budget_duration: str, settings: BudgetResetSettings) -> datetime: 

65 """Compute the next reset time for a budget duration using injected settings.""" 

66 return get_next_standardized_reset_time( 

67 duration=budget_duration, 

68 current_time=datetime.now(timezone.utc), 

69 timezone_str=settings.timezone, 

70 reset_time_of_day=settings.reset_time_of_day, 

71 ) 

72 

73 

74def get_budget_reset_time(budget_duration: str) -> datetime: 

75 """Get the budget reset time using the globally-configured timezone and reset time. 

76 

77 Thin wrapper over `compute_budget_reset_at` for callers that don't yet receive 

78 `BudgetResetSettings` by injection (creation/update endpoints, startup backfill). 

79 """ 

80 return compute_budget_reset_at(budget_duration, get_budget_reset_settings()) 

81 

82 

83def _is_persistable_budget_duration(budget_duration: str) -> bool: 

84 from litellm.litellm_core_utils.duration_parser import duration_in_seconds 

85 

86 try: 

87 if duration_in_seconds(budget_duration) <= 0: 87 ↛ anywhereline 87 didn't jump anywhere: it always raised an exception.

88 return False 

89 get_budget_reset_time(budget_duration=budget_duration) 

90 except (ValueError, OverflowError): 

91 return False 

92 return True 

93 

94 

95def budget_duration_error(budget_duration: str | None) -> str | None: 

96 """Why `budget_duration` cannot be persisted, or None when it is usable. 

97 

98 A non-positive duration resolves to a reset time of "now", which leaves the row 

99 permanently due: the reset job re-reads it every tick and, once enough of them 

100 exist, they fill each batch and starve every other tenant's reset. 

101 """ 

102 if budget_duration is None or _is_persistable_budget_duration(budget_duration): 

103 return None 

104 return f"Invalid budget_duration '{budget_duration}'. Use a format like '1h', '24h', '7d', or '30d'."