From 7d893e78f59589170fc18c796e9b74f0b24fc4f8 Mon Sep 17 00:00:00 2001 From: Aayush Malviya Date: Thu, 5 Mar 2026 17:16:35 +0530 Subject: [PATCH] Incorporate Greptile reviews --- litellm/proxy/_types.py | 46 +++++++++++-------- .../internal_user_endpoints.py | 5 +- .../key_management_endpoints.py | 8 ++-- 3 files changed, 36 insertions(+), 23 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 920bb25453c..a4e11398558 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -786,7 +786,6 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None - initial_budget_reset_at: Optional[datetime] = None allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} @@ -803,23 +802,6 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): aliases: Optional[dict] = {} object_permission: Optional[LiteLLM_ObjectPermissionBase] = None - @model_validator(mode="after") - def validate_initial_budget_reset_at(self) -> "GenerateRequestBase": - """Validate that initial_budget_reset_at is not in the past""" - if self.initial_budget_reset_at is not None: - # Ensure timezone-aware for comparison - reset_at = self.initial_budget_reset_at - if reset_at.tzinfo is None: - reset_at = reset_at.replace(tzinfo=timezone.utc) - - current_time = datetime.now(timezone.utc) - if reset_at < current_time: - raise ValueError( - f"initial_budget_reset_at cannot be in the past. " - f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}" - ) - return self - class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): index_name: str @@ -1128,10 +1110,26 @@ class NewUserRequestTeam(LiteLLMPydanticObjectBase): user_role: Literal["user", "admin"] = "user" +def _validate_initial_budget_reset_at(value: Optional[datetime]) -> Optional[datetime]: + """Shared validator: ensure initial_budget_reset_at is not in the past.""" + if value is not None: + reset_at = value + if reset_at.tzinfo is None: + reset_at = reset_at.replace(tzinfo=timezone.utc) + current_time = datetime.now(timezone.utc) + if reset_at < current_time: + raise ValueError( + f"initial_budget_reset_at cannot be in the past. " + f"Provided: {reset_at.isoformat()}, Current time: {current_time.isoformat()}" + ) + return value + + class NewUserRequest(GenerateRequestBase): max_budget: Optional[float] = None user_email: Optional[str] = None user_alias: Optional[str] = None + initial_budget_reset_at: Optional[datetime] = None user_role: Optional[ Literal[ LitellmUserRoles.PROXY_ADMIN, @@ -1148,6 +1146,11 @@ class NewUserRequest(GenerateRequestBase): sso_user_id: Optional[str] = None organizations: Optional[List[str]] = None + @model_validator(mode="after") + def validate_initial_budget_reset_at(self) -> "NewUserRequest": + _validate_initial_budget_reset_at(self.initial_budget_reset_at) + return self + class NewUserResponse(GenerateKeyResponse): max_budget: Optional[float] = None @@ -1165,6 +1168,7 @@ class NewUserResponse(GenerateKeyResponse): model_max_budget: Optional[dict] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None + initial_budget_reset_at: Optional[datetime] = None class UpdateUserRequestNoUserIDorEmail( @@ -1183,6 +1187,12 @@ class UpdateUserRequestNoUserIDorEmail( ] ] = None max_budget: Optional[float] = None + initial_budget_reset_at: Optional[datetime] = None + + @model_validator(mode="after") + def validate_initial_budget_reset_at(self) -> "UpdateUserRequestNoUserIDorEmail": + _validate_initial_budget_reset_at(self.initial_budget_reset_at) + return self class UpdateUserRequest(UpdateUserRequestNoUserIDorEmail): diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 1026b967a8f..196605689fc 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -780,7 +780,10 @@ def _update_internal_user_params( # ensure timezone aware if _reset_at.tzinfo is None: _reset_at = _reset_at.replace(tzinfo=timezone.utc) - except Exception: + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to parse initial_budget_reset_at value: {_reset_at}. Error: {e}. Falling back to None." + ) _reset_at = None non_default_values["budget_reset_at"] = _reset_at elif "initial_budget_reset_at" in non_default_values: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index d0715eae959..a0aabfb3121 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -582,9 +582,6 @@ async def _common_key_generation_helper( # noqa: PLR0915 if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) - if "initial_budget_reset_at" in data_json: - data_json["budget_reset_at"] = data_json.pop("initial_budget_reset_at", None) - if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id data_json["updated_by"] = user_api_key_dict.user_id @@ -2091,7 +2088,10 @@ async def generate_key_helper_fn( # noqa: PLR0915 if _reset_candidate.tzinfo is None: _reset_candidate = _reset_candidate.replace(tzinfo=timezone.utc) reset_at = _reset_candidate - except Exception: + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to parse budget_reset_at value: {budget_reset_at}. Error: {e}. Falling back to budget_duration-based computation." + ) if budget_duration is None: # fallback to one-time budget reset_at = None else: