From 300c6e74061e44d5eb82e8d9b385b26b2e707a7a Mon Sep 17 00:00:00 2001 From: Jin Yi Date: Thu, 31 Jul 2025 14:15:21 +0900 Subject: [PATCH] feat: Add Pydantic validation to import_fail_info_bulk endpoint - Regenerated Pydantic models from updated OpenAPI specification - Updated import_fail_info_bulk route handler to use ImportFailInfoBulkRequest/Response models - Replaced manual JSON validation with Pydantic model validation - Added proper error handling with ValidationError - Updated data_models/__init__.py to export new models Following the process outlined in data_models/README.md for type safety and consistency. --- comfyui_manager/data_models/__init__.py | 12 ++++++ .../data_models/generated_models.py | 24 +++++++++++- comfyui_manager/glob/manager_server.py | 39 +++++++++---------- 3 files changed, 53 insertions(+), 22 deletions(-) diff --git a/comfyui_manager/data_models/__init__.py b/comfyui_manager/data_models/__init__.py index 5d115e6a..d0d5c182 100644 --- a/comfyui_manager/data_models/__init__.py +++ b/comfyui_manager/data_models/__init__.py @@ -30,6 +30,12 @@ from .generated_models import ( InstalledModelInfo, ComfyUIVersionInfo, + # Import Fail Info Models + ImportFailInfoBulkRequest, + ImportFailInfoBulkResponse, + ImportFailInfoItem, + ImportFailInfoItem1, + # Other models OperationType, OperationResult, @@ -88,6 +94,12 @@ __all__ = [ "InstalledModelInfo", "ComfyUIVersionInfo", + # Import Fail Info Models + "ImportFailInfoBulkRequest", + "ImportFailInfoBulkResponse", + "ImportFailInfoItem", + "ImportFailInfoItem1", + # Other models "OperationType", "OperationResult", diff --git a/comfyui_manager/data_models/generated_models.py b/comfyui_manager/data_models/generated_models.py index 69771009..0f3b9cfb 100644 --- a/comfyui_manager/data_models/generated_models.py +++ b/comfyui_manager/data_models/generated_models.py @@ -1,6 +1,6 @@ # generated by datamodel-codegen: # filename: openapi.yaml -# timestamp: 2025-06-27T04:01:45+00:00 +# timestamp: 2025-07-31T04:52:26+00:00 from __future__ import annotations @@ -454,6 +454,24 @@ class BatchExecutionRecord(BaseModel): ) +class ImportFailInfoBulkRequest(BaseModel): + cnr_ids: Optional[List[str]] = Field( + None, description="A list of CNR IDs to check." + ) + urls: Optional[List[str]] = Field( + None, description="A list of repository URLs to check." + ) + + +class ImportFailInfoItem1(BaseModel): + error: Optional[str] = None + traceback: Optional[str] = None + + +class ImportFailInfoItem(RootModel[Optional[ImportFailInfoItem1]]): + root: Optional[ImportFailInfoItem1] + + class QueueTaskItem(BaseModel): ui_id: str = Field(..., description="Unique identifier for the task") client_id: str = Field(..., description="Client identifier that initiated the task") @@ -537,3 +555,7 @@ class HistoryResponse(BaseModel): history: Optional[Dict[str, TaskHistoryItem]] = Field( None, description="Map of task IDs to their history items" ) + + +class ImportFailInfoBulkResponse(RootModel[Optional[Dict[str, ImportFailInfoItem]]]): + root: Optional[Dict[str, ImportFailInfoItem]] = None diff --git a/comfyui_manager/glob/manager_server.py b/comfyui_manager/glob/manager_server.py index e2699389..cc06a58e 100644 --- a/comfyui_manager/glob/manager_server.py +++ b/comfyui_manager/glob/manager_server.py @@ -61,6 +61,8 @@ from ..data_models import ( ManagerMessageName, BatchExecutionRecord, ComfyUISystemState, + ImportFailInfoBulkRequest, + ImportFailInfoBulkResponse, BatchOperation, InstalledNodeInfo, ComfyUIVersionInfo, @@ -1660,12 +1662,12 @@ async def import_fail_info(request): async def import_fail_info_bulk(request): try: json_data = await request.json() - - # Basic validation - ensure we have either cnr_ids or urls - if not isinstance(json_data, dict): - return web.Response(status=400, text="Request body must be a JSON object") - - if "cnr_ids" not in json_data and "urls" not in json_data: + + # Validate input using Pydantic model + request_data = ImportFailInfoBulkRequest.model_validate(json_data) + + # Ensure we have either cnr_ids or urls + if not request_data.cnr_ids and not request_data.urls: return web.Response( status=400, text="Either 'cnr_ids' or 'urls' field is required" ) @@ -1675,13 +1677,8 @@ async def import_fail_info_bulk(request): results = {} - if "cnr_ids" in json_data: - if not isinstance(json_data["cnr_ids"], list): - return web.Response(status=400, text="'cnr_ids' must be an array") - for cnr_id in json_data["cnr_ids"]: - if not isinstance(cnr_id, str): - results[cnr_id] = {"error": "cnr_id must be a string"} - continue + if request_data.cnr_ids: + for cnr_id in request_data.cnr_ids: module_name = core.unified_manager.get_module_name(cnr_id) if module_name is not None: info = cm_global.error_dict.get(module_name) @@ -1692,13 +1689,8 @@ async def import_fail_info_bulk(request): else: results[cnr_id] = None - if "urls" in json_data: - if not isinstance(json_data["urls"], list): - return web.Response(status=400, text="'urls' must be an array") - for url in json_data["urls"]: - if not isinstance(url, str): - results[url] = {"error": "url must be a string"} - continue + if request_data.urls: + for url in request_data.urls: module_name = core.unified_manager.get_module_name(url) if module_name is not None: info = cm_global.error_dict.get(module_name) @@ -1709,7 +1701,12 @@ async def import_fail_info_bulk(request): else: results[url] = None - return web.json_response(results) + # Create response using Pydantic model + response_data = ImportFailInfoBulkResponse(root=results) + return web.json_response(response_data.root) + except ValidationError as e: + logging.error(f"[ComfyUI-Manager] Invalid request data: {e}") + return web.Response(status=400, text=f"Invalid request data: {e}") except Exception as e: logging.error(f"[ComfyUI-Manager] Error processing bulk import fail info: {e}") return web.Response(status=500, text="Internal server error")