diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d6b2d664f9f..b4f64a76fd1 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -786,6 +786,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None + budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 66085b69b3d..681087451bc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -306,6 +306,7 @@ async def new_user( - user_role: Optional[str] - Specify a user role - "proxy_admin", "proxy_admin_viewer", "internal_user", "internal_user_viewer", "team", "customer". Info about each role here: `https://github.com/BerriAI/litellm/litellm/proxy/_types.py#L20` - max_budget: Optional[float] - Specify max budget for a given user. - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"), months ("1mo"). + - budget_reset_at: Optional[datetime] - Specify the exact datetime when the budget should first reset (ISO 8601 format, e.g., "2025-11-15T00:00:00Z"). If provided along with budget_duration, this value takes precedence for the initial reset; subsequent resets will follow the budget_duration interval. - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models). Set to ['no-default-models'] to block all model access. Restricting user to only team-based model access. - tpm_limit: Optional[int] - Specify tpm limit for a given user (Tokens per minute) - rpm_limit: Optional[int] - Specify rpm limit for a given user (Requests per minute) @@ -763,12 +764,26 @@ def _update_internal_user_params( if data.user_role == LitellmUserRoles.INTERNAL_USER: is_internal_user = True + if "budget_reset_at" in non_default_values and non_default_values["budget_reset_at"] is not None: + _reset_at = non_default_values["budget_reset_at"] + try: + # convert to datetime if not already + if not isinstance(_reset_at, datetime): + _reset_at = datetime.fromisoformat(str(_reset_at).replace("Z", "+00:00")) + # ensure timezone aware + if _reset_at.tzinfo is None: + _reset_at = _reset_at.replace(tzinfo=timezone.utc) + except Exception: + _reset_at = None + non_default_values["budget_reset_at"] = _reset_at + if "budget_duration" in non_default_values: 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"] - ) + # only auto-compute budget_reset_at if not explicitly provided + if "budget_reset_at" not in non_default_values or non_default_values["budget_reset_at"] is None: + 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 ( @@ -785,9 +800,11 @@ def _update_internal_user_params( ) 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"] - ) + # only auto-compute budget_reset_at if not explicitly provided + if "budget_reset_at" not in non_default_values or non_default_values["budget_reset_at"] is None: + non_default_values["budget_reset_at"] = get_budget_reset_time( + budget_duration=non_default_values["budget_duration"] + ) 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 a8def13869f..74ed92ce1f8 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -2015,6 +2015,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 max_budget: Optional[float] = None, # max_budget is used to Budget Per user blocked: Optional[bool] = None, budget_duration: Optional[str] = None, # max_budget is used to Budget Per user + budget_reset_at: Optional[datetime] = None, # manual override for initial user/key budget reset time token: Optional[str] = None, key: Optional[ str @@ -2076,10 +2077,27 @@ async def generate_key_helper_fn( # noqa: PLR0915 else: key_reset_at = get_budget_reset_time(budget_duration=key_budget_duration) - if budget_duration is None: # one-time budget - reset_at = None + if budget_reset_at is not None: + try: + # try to parse budget_reset_at + if isinstance(budget_reset_at, datetime): + _reset_candidate = budget_reset_at + else: + _reset_candidate = datetime.fromisoformat(str(budget_reset_at).replace("Z", "+00:00")) + # ensure tz-aware datetime + if _reset_candidate.tzinfo is None: + _reset_candidate = _reset_candidate.replace(tzinfo=timezone.utc) + reset_at = _reset_candidate + except Exception: + if budget_duration is None: # fallback to one-time budget + reset_at = None + else: + reset_at = get_budget_reset_time(budget_duration=budget_duration) # calculate based on budget duration else: - reset_at = get_budget_reset_time(budget_duration=budget_duration) + if budget_duration is None: # one-time budget + reset_at = None + else: + reset_at = get_budget_reset_time(budget_duration=budget_duration) # calculate based on budget duration aliases_json = json.dumps(aliases) config_json = json.dumps(config)