Incorporate Greptile reviews

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):

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:

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: