[feat] Implement comprehensive batch tracking and OpenAPI-driven data models
Enhances ComfyUI Manager with robust batch execution tracking and unified data model architecture: - Implemented automatic batch history serialization with before/after system state snapshots - Added comprehensive state management capturing installed nodes, models, and ComfyUI version info - Enhanced task queue with proper client ID handling and WebSocket notifications - Migrated all data models to OpenAPI-generated Pydantic models for consistency - Added documentation for new TaskQueue methods (done_count, total_count, finalize) - Fixed 64 linting errors with proper imports and code cleanup Technical improvements: - All models now auto-generated from openapi.yaml ensuring API/implementation consistency - Batch tracking captures complete system state at operation start and completion - Enhanced REST endpoints with comprehensive documentation - Removed manual model files in favor of single source of truth - Added helper methods for system state capture and batch lifecycle management
This commit is contained in:
bymyself
417
comfyui_manager/data_models/generated_models.py
Normal file
417
comfyui_manager/data_models/generated_models.py
Normal file
@@ -0,0 +1,417 @@
|
||||
# generated by datamodel-codegen:
|
||||
# filename: openapi.yaml
|
||||
# timestamp: 2025-06-08T08:07:38+00:00
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, List, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, Field, RootModel
|
||||
|
||||
|
||||
class Kind(str, Enum):
|
||||
install = 'install'
|
||||
uninstall = 'uninstall'
|
||||
update = 'update'
|
||||
update_all = 'update-all'
|
||||
update_comfyui = 'update-comfyui'
|
||||
fix = 'fix'
|
||||
disable = 'disable'
|
||||
enable = 'enable'
|
||||
install_model = 'install-model'
|
||||
|
||||
|
||||
class QueueTaskItem(BaseModel):
|
||||
ui_id: str = Field(..., description='Unique identifier for the task')
|
||||
client_id: str = Field(..., description='Client identifier that initiated the task')
|
||||
kind: Kind = Field(..., description='Type of task being performed')
|
||||
|
||||
|
||||
class StatusStr(str, Enum):
|
||||
success = 'success'
|
||||
error = 'error'
|
||||
skip = 'skip'
|
||||
|
||||
|
||||
class TaskExecutionStatus(BaseModel):
|
||||
status_str: StatusStr = Field(..., description='Overall task execution status')
|
||||
completed: bool = Field(..., description='Whether the task completed')
|
||||
messages: List[str] = Field(..., description='Additional status messages')
|
||||
|
||||
|
||||
class ManagerMessageName(str, Enum):
|
||||
cm_task_completed = 'cm-task-completed'
|
||||
cm_task_started = 'cm-task-started'
|
||||
cm_queue_status = 'cm-queue-status'
|
||||
|
||||
|
||||
class ManagerPackInfo(BaseModel):
|
||||
id: str = Field(
|
||||
...,
|
||||
description='Either github-author/github-repo or name of pack from the registry',
|
||||
)
|
||||
version: str = Field(..., description='Semantic version or Git commit hash')
|
||||
ui_id: Optional[str] = Field(None, description='Task ID - generated internally')
|
||||
|
||||
|
||||
class ManagerPackInstalled(BaseModel):
|
||||
ver: str = Field(
|
||||
...,
|
||||
description='The version of the pack that is installed (Git commit hash or semantic version)',
|
||||
)
|
||||
cnr_id: Optional[str] = Field(
|
||||
None, description='The name of the pack if installed from the registry'
|
||||
)
|
||||
aux_id: Optional[str] = Field(
|
||||
None,
|
||||
description='The name of the pack if installed from github (author/repo-name format)',
|
||||
)
|
||||
enabled: bool = Field(..., description='Whether the pack is enabled')
|
||||
|
||||
|
||||
class SelectedVersion(str, Enum):
|
||||
latest = 'latest'
|
||||
nightly = 'nightly'
|
||||
|
||||
|
||||
class ManagerChannel(str, Enum):
|
||||
default = 'default'
|
||||
recent = 'recent'
|
||||
legacy = 'legacy'
|
||||
forked = 'forked'
|
||||
dev = 'dev'
|
||||
tutorial = 'tutorial'
|
||||
|
||||
|
||||
class ManagerDatabaseSource(str, Enum):
|
||||
remote = 'remote'
|
||||
local = 'local'
|
||||
cache = 'cache'
|
||||
|
||||
|
||||
class ManagerPackState(str, Enum):
|
||||
installed = 'installed'
|
||||
disabled = 'disabled'
|
||||
not_installed = 'not_installed'
|
||||
import_failed = 'import_failed'
|
||||
needs_update = 'needs_update'
|
||||
|
||||
|
||||
class ManagerPackInstallType(str, Enum):
|
||||
git_clone = 'git-clone'
|
||||
copy = 'copy'
|
||||
cnr = 'cnr'
|
||||
|
||||
|
||||
class UpdateState(str, Enum):
|
||||
false = 'false'
|
||||
true = 'true'
|
||||
|
||||
|
||||
class ManagerPack(ManagerPackInfo):
|
||||
author: Optional[str] = Field(
|
||||
None, description="Pack author name or 'Unclaimed' if added via GitHub crawl"
|
||||
)
|
||||
files: Optional[List[str]] = Field(None, description='Files included in the pack')
|
||||
reference: Optional[str] = Field(
|
||||
None, description='The type of installation reference'
|
||||
)
|
||||
title: Optional[str] = Field(None, description='The display name of the pack')
|
||||
cnr_latest: Optional[SelectedVersion] = None
|
||||
repository: Optional[str] = Field(None, description='GitHub repository URL')
|
||||
state: Optional[ManagerPackState] = None
|
||||
update_state: Optional[UpdateState] = Field(
|
||||
None, alias='update-state', description='Update availability status'
|
||||
)
|
||||
stars: Optional[int] = Field(None, description='GitHub stars count')
|
||||
last_update: Optional[datetime] = Field(None, description='Last update timestamp')
|
||||
health: Optional[str] = Field(None, description='Health status of the pack')
|
||||
description: Optional[str] = Field(None, description='Pack description')
|
||||
trust: Optional[bool] = Field(None, description='Whether the pack is trusted')
|
||||
install_type: Optional[ManagerPackInstallType] = None
|
||||
|
||||
|
||||
class InstallPackParams(ManagerPackInfo):
|
||||
selected_version: Union[str, SelectedVersion] = Field(
|
||||
..., description='Semantic version, Git commit hash, latest, or nightly'
|
||||
)
|
||||
repository: Optional[str] = Field(
|
||||
None,
|
||||
description='GitHub repository URL (required if selected_version is nightly)',
|
||||
)
|
||||
pip: Optional[List[str]] = Field(None, description='PyPi dependency names')
|
||||
mode: ManagerDatabaseSource
|
||||
channel: ManagerChannel
|
||||
skip_post_install: Optional[bool] = Field(
|
||||
None, description='Whether to skip post-installation steps'
|
||||
)
|
||||
|
||||
|
||||
class UpdateAllPacksParams(BaseModel):
|
||||
mode: Optional[ManagerDatabaseSource] = None
|
||||
ui_id: Optional[str] = Field(None, description='Task ID - generated internally')
|
||||
|
||||
|
||||
class QueueStatus(BaseModel):
|
||||
total_count: int = Field(
|
||||
..., description='Total number of tasks (pending + running)'
|
||||
)
|
||||
done_count: int = Field(..., description='Number of completed tasks')
|
||||
in_progress_count: int = Field(..., description='Number of tasks currently running')
|
||||
pending_count: Optional[int] = Field(
|
||||
None, description='Number of tasks waiting to be executed'
|
||||
)
|
||||
is_processing: bool = Field(..., description='Whether the task worker is active')
|
||||
client_id: Optional[str] = Field(
|
||||
None, description='Client ID (when filtered by client)'
|
||||
)
|
||||
|
||||
|
||||
class ManagerMapping(BaseModel):
|
||||
title_aux: Optional[str] = Field(None, description='The display name of the pack')
|
||||
|
||||
|
||||
class ManagerMappings(
|
||||
RootModel[Optional[Dict[str, List[Union[List[str], ManagerMapping]]]]]
|
||||
):
|
||||
root: Optional[Dict[str, List[Union[List[str], ManagerMapping]]]] = None
|
||||
|
||||
|
||||
class ModelMetadata(BaseModel):
|
||||
name: str = Field(..., description='Name of the model')
|
||||
type: str = Field(..., description='Type of model')
|
||||
base: Optional[str] = Field(None, description='Base model type')
|
||||
save_path: Optional[str] = Field(None, description='Path for saving the model')
|
||||
url: str = Field(..., description='Download URL')
|
||||
filename: str = Field(..., description='Target filename')
|
||||
ui_id: Optional[str] = Field(None, description='ID for UI reference')
|
||||
|
||||
|
||||
class InstallType(str, Enum):
|
||||
git = 'git'
|
||||
copy = 'copy'
|
||||
pip = 'pip'
|
||||
|
||||
|
||||
class NodePackageMetadata(BaseModel):
|
||||
title: Optional[str] = Field(None, description='Display name of the node package')
|
||||
name: Optional[str] = Field(None, description='Repository/package name')
|
||||
files: Optional[List[str]] = Field(None, description='Source URLs for the package')
|
||||
description: Optional[str] = Field(
|
||||
None, description='Description of the node package functionality'
|
||||
)
|
||||
install_type: Optional[InstallType] = Field(None, description='Installation method')
|
||||
version: Optional[str] = Field(None, description='Version identifier')
|
||||
id: Optional[str] = Field(
|
||||
None, description='Unique identifier for the node package'
|
||||
)
|
||||
ui_id: Optional[str] = Field(None, description='ID for UI reference')
|
||||
channel: Optional[str] = Field(None, description='Source channel')
|
||||
mode: Optional[str] = Field(None, description='Source mode')
|
||||
|
||||
|
||||
class SnapshotItem(RootModel[str]):
|
||||
root: str = Field(..., description='Name of the snapshot')
|
||||
|
||||
|
||||
class Error(BaseModel):
|
||||
error: str = Field(..., description='Error message')
|
||||
|
||||
|
||||
class InstalledPacksResponse(RootModel[Optional[Dict[str, ManagerPackInstalled]]]):
|
||||
root: Optional[Dict[str, ManagerPackInstalled]] = None
|
||||
|
||||
|
||||
class HistoryListResponse(BaseModel):
|
||||
ids: Optional[List[str]] = Field(
|
||||
None, description='List of available batch history IDs'
|
||||
)
|
||||
|
||||
|
||||
class InstalledNodeInfo(BaseModel):
|
||||
name: str = Field(..., description='Node package name')
|
||||
version: str = Field(..., description='Installed version')
|
||||
repository_url: Optional[str] = Field(None, description='Git repository URL')
|
||||
install_method: str = Field(
|
||||
..., description='Installation method (cnr, git, pip, etc.)'
|
||||
)
|
||||
enabled: Optional[bool] = Field(
|
||||
True, description='Whether the node is currently enabled'
|
||||
)
|
||||
install_date: Optional[datetime] = Field(
|
||||
None, description='ISO timestamp of installation'
|
||||
)
|
||||
|
||||
|
||||
class InstalledModelInfo(BaseModel):
|
||||
name: str = Field(..., description='Model filename')
|
||||
path: str = Field(..., description='Full path to model file')
|
||||
type: str = Field(..., description='Model type (checkpoint, lora, vae, etc.)')
|
||||
size_bytes: Optional[int] = Field(None, description='File size in bytes', ge=0)
|
||||
hash: Optional[str] = Field(None, description='Model file hash for verification')
|
||||
install_date: Optional[datetime] = Field(
|
||||
None, description='ISO timestamp when added'
|
||||
)
|
||||
|
||||
|
||||
class ComfyUIVersionInfo(BaseModel):
|
||||
version: str = Field(..., description='ComfyUI version string')
|
||||
commit_hash: Optional[str] = Field(None, description='Git commit hash')
|
||||
branch: Optional[str] = Field(None, description='Git branch name')
|
||||
is_stable: Optional[bool] = Field(
|
||||
False, description='Whether this is a stable release'
|
||||
)
|
||||
last_updated: Optional[datetime] = Field(
|
||||
None, description='ISO timestamp of last update'
|
||||
)
|
||||
|
||||
|
||||
class OperationType(str, Enum):
|
||||
install = 'install'
|
||||
update = 'update'
|
||||
uninstall = 'uninstall'
|
||||
fix = 'fix'
|
||||
disable = 'disable'
|
||||
enable = 'enable'
|
||||
install_model = 'install-model'
|
||||
|
||||
|
||||
class Result(str, Enum):
|
||||
success = 'success'
|
||||
failed = 'failed'
|
||||
skipped = 'skipped'
|
||||
|
||||
|
||||
class BatchOperation(BaseModel):
|
||||
operation_id: str = Field(..., description='Unique operation identifier')
|
||||
operation_type: OperationType = Field(..., description='Type of operation')
|
||||
target: str = Field(
|
||||
..., description='Target of the operation (node name, model name, etc.)'
|
||||
)
|
||||
target_version: Optional[str] = Field(
|
||||
None, description='Target version for the operation'
|
||||
)
|
||||
result: Result = Field(..., description='Operation result')
|
||||
error_message: Optional[str] = Field(
|
||||
None, description='Error message if operation failed'
|
||||
)
|
||||
start_time: datetime = Field(
|
||||
..., description='ISO timestamp when operation started'
|
||||
)
|
||||
end_time: Optional[datetime] = Field(
|
||||
None, description='ISO timestamp when operation completed'
|
||||
)
|
||||
client_id: Optional[str] = Field(
|
||||
None, description='Client that initiated the operation'
|
||||
)
|
||||
|
||||
|
||||
class ComfyUISystemState(BaseModel):
|
||||
snapshot_time: datetime = Field(
|
||||
..., description='ISO timestamp when snapshot was taken'
|
||||
)
|
||||
comfyui_version: ComfyUIVersionInfo
|
||||
frontend_version: Optional[str] = Field(
|
||||
None, description='ComfyUI frontend version if available'
|
||||
)
|
||||
python_version: str = Field(..., description='Python interpreter version')
|
||||
platform_info: str = Field(
|
||||
..., description='Operating system and platform information'
|
||||
)
|
||||
installed_nodes: Optional[Dict[str, InstalledNodeInfo]] = Field(
|
||||
None, description='Map of installed node packages by name'
|
||||
)
|
||||
installed_models: Optional[Dict[str, InstalledModelInfo]] = Field(
|
||||
None, description='Map of installed models by name'
|
||||
)
|
||||
manager_config: Optional[Dict[str, Any]] = Field(
|
||||
None, description='ComfyUI Manager configuration settings'
|
||||
)
|
||||
|
||||
|
||||
class BatchExecutionRecord(BaseModel):
|
||||
batch_id: str = Field(..., description='Unique batch identifier')
|
||||
start_time: datetime = Field(..., description='ISO timestamp when batch started')
|
||||
end_time: Optional[datetime] = Field(
|
||||
None, description='ISO timestamp when batch completed'
|
||||
)
|
||||
state_before: ComfyUISystemState
|
||||
state_after: Optional[ComfyUISystemState] = Field(
|
||||
None, description='System state after batch execution'
|
||||
)
|
||||
operations: Optional[List[BatchOperation]] = Field(
|
||||
None, description='List of operations performed in this batch'
|
||||
)
|
||||
total_operations: Optional[int] = Field(
|
||||
0, description='Total number of operations in batch', ge=0
|
||||
)
|
||||
successful_operations: Optional[int] = Field(
|
||||
0, description='Number of successful operations', ge=0
|
||||
)
|
||||
failed_operations: Optional[int] = Field(
|
||||
0, description='Number of failed operations', ge=0
|
||||
)
|
||||
skipped_operations: Optional[int] = Field(
|
||||
0, description='Number of skipped operations', ge=0
|
||||
)
|
||||
|
||||
|
||||
class TaskHistoryItem(BaseModel):
|
||||
ui_id: str = Field(..., description='Unique identifier for the task')
|
||||
client_id: str = Field(..., description='Client identifier that initiated the task')
|
||||
kind: str = Field(..., description='Type of task that was performed')
|
||||
timestamp: datetime = Field(..., description='ISO timestamp when task completed')
|
||||
result: str = Field(..., description='Task result message or details')
|
||||
status: Optional[TaskExecutionStatus] = None
|
||||
|
||||
|
||||
class TaskStateMessage(BaseModel):
|
||||
history: Dict[str, TaskHistoryItem] = Field(
|
||||
..., description='Map of task IDs to their history items'
|
||||
)
|
||||
running_queue: List[QueueTaskItem] = Field(
|
||||
..., description='Currently executing tasks'
|
||||
)
|
||||
pending_queue: List[QueueTaskItem] = Field(
|
||||
..., description='Tasks waiting to be executed'
|
||||
)
|
||||
|
||||
|
||||
class MessageTaskDone(BaseModel):
|
||||
ui_id: str = Field(..., description='Task identifier')
|
||||
result: str = Field(..., description='Task result message')
|
||||
kind: str = Field(..., description='Type of task')
|
||||
status: Optional[TaskExecutionStatus] = None
|
||||
timestamp: datetime = Field(..., description='ISO timestamp when task completed')
|
||||
state: TaskStateMessage
|
||||
|
||||
|
||||
class MessageTaskStarted(BaseModel):
|
||||
ui_id: str = Field(..., description='Task identifier')
|
||||
kind: str = Field(..., description='Type of task')
|
||||
timestamp: datetime = Field(..., description='ISO timestamp when task started')
|
||||
state: TaskStateMessage
|
||||
|
||||
|
||||
class MessageTaskFailed(BaseModel):
|
||||
ui_id: str = Field(..., description='Task identifier')
|
||||
error: str = Field(..., description='Error message')
|
||||
kind: str = Field(..., description='Type of task')
|
||||
timestamp: datetime = Field(..., description='ISO timestamp when task failed')
|
||||
state: TaskStateMessage
|
||||
|
||||
|
||||
class MessageUpdate(
|
||||
RootModel[Union[MessageTaskDone, MessageTaskStarted, MessageTaskFailed]]
|
||||
):
|
||||
root: Union[MessageTaskDone, MessageTaskStarted, MessageTaskFailed] = Field(
|
||||
..., description='Union type for all possible WebSocket message updates'
|
||||
)
|
||||
|
||||
|
||||
class HistoryResponse(BaseModel):
|
||||
history: Optional[Dict[str, TaskHistoryItem]] = Field(
|
||||
None, description='Map of task IDs to their history items'
|
||||
)
|
||||
Reference in New Issue
Block a user