diff --git a/comfyui_manager/data_models/generated_models.py b/comfyui_manager/data_models/generated_models.py index 6e3349f1..d26207c0 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-17T21:37:15+00:00 +# timestamp: 2025-06-17T22:03:51+00:00 from __future__ import annotations @@ -8,158 +8,114 @@ from datetime import datetime from enum import Enum from typing import Any, Dict, List, Optional, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, RootModel -class OperationType(Enum): - """ - Type of operation or task being performed - """ - - install = 'install' - uninstall = 'uninstall' - update = 'update' - update_comfyui = 'update-comfyui' - fix = 'fix' - disable = 'disable' - enable = 'enable' - install_model = 'install-model' +class OperationType(str, Enum): + install = "install" + uninstall = "uninstall" + update = "update" + update_comfyui = "update-comfyui" + fix = "fix" + disable = "disable" + enable = "enable" + install_model = "install-model" -class OperationResult(Enum): - """ - Result status of an operation (failed/error and skipped/skip are aliases) - """ - - success = 'success' - failed = 'failed' - skipped = 'skipped' - error = 'error' - skip = 'skip' +class OperationResult(str, Enum): + success = "success" + failed = "failed" + skipped = "skipped" + error = "error" + skip = "skip" class TaskExecutionStatus(BaseModel): status_str: OperationResult - completed: bool = Field(..., description='Whether the task completed') - messages: List[str] = Field(..., description='Additional status messages') + completed: bool = Field(..., description="Whether the task completed") + messages: List[str] = Field(..., description="Additional status messages") -class ManagerMessageName(Enum): - """ - WebSocket message type constants for manager events - """ - - cm_task_completed = 'cm-task-completed' - cm_task_started = 'cm-task-started' - cm_queue_status = 'cm-queue-status' +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', + 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') + 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)', + 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' + 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)', + description="The name of the pack if installed from github (author/repo-name format)", ) - enabled: bool = Field(..., description='Whether the pack is enabled') + enabled: bool = Field(..., description="Whether the pack is enabled") -class SelectedVersion(Enum): - """ - Version selection for pack installation - """ - - latest = 'latest' - nightly = 'nightly' +class SelectedVersion(str, Enum): + latest = "latest" + nightly = "nightly" -class ManagerChannel(Enum): - """ - Channel for pack sources - """ - - default = 'default' - recent = 'recent' - legacy = 'legacy' - forked = 'forked' - dev = 'dev' - tutorial = 'tutorial' +class ManagerChannel(str, Enum): + default = "default" + recent = "recent" + legacy = "legacy" + forked = "forked" + dev = "dev" + tutorial = "tutorial" -class ManagerDatabaseSource(Enum): - """ - Source for pack information - """ - - remote = 'remote' - local = 'local' - cache = 'cache' +class ManagerDatabaseSource(str, Enum): + remote = "remote" + local = "local" + cache = "cache" -class ManagerPackState(Enum): - """ - Current state of a pack - """ - - installed = 'installed' - disabled = 'disabled' - not_installed = 'not_installed' - import_failed = 'import_failed' - needs_update = 'needs_update' +class ManagerPackState(str, Enum): + installed = "installed" + disabled = "disabled" + not_installed = "not_installed" + import_failed = "import_failed" + needs_update = "needs_update" -class ManagerPackInstallType(Enum): - """ - Type of installation used for the pack - """ - - git_clone = 'git-clone' - copy = 'copy' - cnr = 'cnr' +class ManagerPackInstallType(str, Enum): + git_clone = "git-clone" + copy = "copy" + cnr = "cnr" -class SecurityLevel(Enum): - """ - Security level configuration (from most to least restrictive) - """ - - strong = 'strong' - normal = 'normal' - normal_ = 'normal-' - weak = 'weak' +class SecurityLevel(str, Enum): + strong = "strong" + normal = "normal" + normal_ = "normal-" + weak = "weak" -class RiskLevel(Enum): - """ - Risk classification for operations - """ - - block = 'block' - high = 'high' - middle = 'middle' +class RiskLevel(str, Enum): + block = "block" + high = "high" + middle = "middle" class UpdateState(Enum): - """ - Update availability status - """ - - false = 'false' - true = 'true' + false = "false" + true = "true" class ManagerPack(ManagerPackInfo): @@ -168,339 +124,333 @@ class ManagerPack(ManagerPackInfo): ) files: Optional[List[str]] = Field( None, - description='Repository URLs for installation (typically contains one GitHub URL)', + description="Repository URLs for installation (typically contains one GitHub URL)", ) reference: Optional[str] = Field( - None, description='The type of installation reference' + None, description="The type of installation reference" ) - title: Optional[str] = Field(None, description='The display name of the pack') + 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') + 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' + 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') + 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' + ..., description="Semantic version, Git commit hash, latest, or nightly" ) repository: Optional[str] = Field( None, - description='GitHub repository URL (required if selected_version is nightly)', + description="GitHub repository URL (required if selected_version is nightly)", ) - pip: Optional[List[str]] = Field(None, description='PyPi dependency names') + 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' + 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') + ui_id: Optional[str] = Field(None, description="Task ID - generated internally") class UpdatePackParams(BaseModel): - node_name: str = Field(..., description='Name of the node package to update') + node_name: str = Field(..., description="Name of the node package to update") node_ver: Optional[str] = Field( - None, description='Current version of the node package' + None, description="Current version of the node package" ) class UpdateComfyUIParams(BaseModel): is_stable: Optional[bool] = Field( True, - description='Whether to update to stable version (true) or nightly (false)', + description="Whether to update to stable version (true) or nightly (false)", ) target_version: Optional[str] = Field( None, - description='Specific version to switch to (for version switching operations)', + description="Specific version to switch to (for version switching operations)", ) class FixPackParams(BaseModel): - node_name: str = Field(..., description='Name of the node package to fix') - node_ver: str = Field(..., description='Version of the node package') + node_name: str = Field(..., description="Name of the node package to fix") + node_ver: str = Field(..., description="Version of the node package") class UninstallPackParams(BaseModel): - node_name: str = Field(..., description='Name of the node package to uninstall') + node_name: str = Field(..., description="Name of the node package to uninstall") is_unknown: Optional[bool] = Field( - False, description='Whether this is an unknown/unregistered package' + False, description="Whether this is an unknown/unregistered package" ) class DisablePackParams(BaseModel): - node_name: str = Field(..., description='Name of the node package to disable') + node_name: str = Field(..., description="Name of the node package to disable") is_unknown: Optional[bool] = Field( - False, description='Whether this is an unknown/unregistered package' + False, description="Whether this is an unknown/unregistered package" ) class EnablePackParams(BaseModel): cnr_id: str = Field( - ..., description='ComfyUI Node Registry ID of the package to enable' + ..., description="ComfyUI Node Registry ID of the package to enable" ) class UpdateAllQueryParams(BaseModel): client_id: str = Field( - ..., description='Client identifier that initiated the request' + ..., description="Client identifier that initiated the request" ) - ui_id: str = Field(..., description='Base UI identifier for task tracking') + ui_id: str = Field(..., description="Base UI identifier for task tracking") mode: Optional[ManagerDatabaseSource] = None class UpdateComfyUIQueryParams(BaseModel): client_id: str = Field( - ..., description='Client identifier that initiated the request' + ..., description="Client identifier that initiated the request" ) - ui_id: str = Field(..., description='UI identifier for task tracking') + ui_id: str = Field(..., description="UI identifier for task tracking") stable: Optional[bool] = Field( True, - description='Whether to update to stable version (true) or nightly (false)', + description="Whether to update to stable version (true) or nightly (false)", ) class ComfyUISwitchVersionQueryParams(BaseModel): - ver: str = Field(..., description='Version to switch to') + ver: str = Field(..., description="Version to switch to") client_id: str = Field( - ..., description='Client identifier that initiated the request' + ..., description="Client identifier that initiated the request" ) - ui_id: str = Field(..., description='UI identifier for task tracking') + ui_id: str = Field(..., description="UI identifier for task tracking") class QueueStatus(BaseModel): total_count: int = Field( - ..., description='Total number of tasks (pending + running)' + ..., 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') + 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' + None, description="Number of tasks waiting to be executed" ) - is_processing: bool = Field(..., description='Whether the task worker is active') + is_processing: bool = Field(..., description="Whether the task worker is active") client_id: Optional[str] = Field( - None, description='Client ID (when filtered by client)' + None, description="Client ID (when filtered by client)" ) class ManagerMappings1(BaseModel): - title_aux: Optional[str] = Field(None, description='The display name of the pack') + title_aux: Optional[str] = Field(None, description="The display name of the pack") -class ManagerMappings(BaseModel): - __root__: Optional[Dict[str, List[Union[List[str], ManagerMappings1]]]] = Field( - None, description='Tuple of [node_names, metadata]' +class ManagerMappings( + RootModel[Optional[Dict[str, List[Union[List[str], ManagerMappings1]]]]] +): + root: Optional[Dict[str, List[Union[List[str], ManagerMappings1]]]] = Field( + None, description="Tuple of [node_names, metadata]" ) 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') + 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(Enum): - """ - Installation method - """ - - git = 'git' - copy = 'copy' - pip = 'pip' +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') + 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' + 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') + 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' + 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') + 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(BaseModel): - __root__: str = Field(..., description='Name of the snapshot') +class SnapshotItem(RootModel[str]): + root: str = Field(..., description="Name of the snapshot") class Error(BaseModel): - error: str = Field(..., description='Error message') + error: str = Field(..., description="Error message") -class InstalledPacksResponse(BaseModel): - """ - Map of pack names to their installation info - """ - - __root__: Optional[Dict[str, ManagerPackInstalled]] = None +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' + 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') + 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.)' + ..., description="Installation method (cnr, git, pip, etc.)" ) enabled: Optional[bool] = Field( - True, description='Whether the node is currently enabled' + True, description="Whether the node is currently enabled" ) install_date: Optional[datetime] = Field( - None, description='ISO timestamp of installation' + 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') + 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' + 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') + 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' + False, description="Whether this is a stable release" ) last_updated: Optional[datetime] = Field( - None, description='ISO timestamp of last update' + None, description="ISO timestamp of last update" ) class BatchOperation(BaseModel): - operation_id: str = Field(..., description='Unique operation identifier') + operation_id: str = Field(..., description="Unique operation identifier") operation_type: OperationType target: str = Field( - ..., description='Target of the operation (node name, model name, etc.)' + ..., description="Target of the operation (node name, model name, etc.)" ) target_version: Optional[str] = Field( - None, description='Target version for the operation' + None, description="Target version for the operation" ) result: OperationResult error_message: Optional[str] = Field( - None, description='Error message if operation failed' + None, description="Error message if operation failed" ) start_time: datetime = Field( - ..., description='ISO timestamp when operation started' + ..., description="ISO timestamp when operation started" ) end_time: Optional[datetime] = Field( - None, description='ISO timestamp when operation completed' + None, description="ISO timestamp when operation completed" ) client_id: Optional[str] = Field( - None, description='Client that initiated the operation' + None, description="Client that initiated the operation" ) class ComfyUISystemState(BaseModel): snapshot_time: datetime = Field( - ..., description='ISO timestamp when snapshot was taken' + ..., description="ISO timestamp when snapshot was taken" ) comfyui_version: ComfyUIVersionInfo frontend_version: Optional[str] = Field( - None, description='ComfyUI frontend version if available' + None, description="ComfyUI frontend version if available" ) - python_version: str = Field(..., description='Python interpreter version') + python_version: str = Field(..., description="Python interpreter version") platform_info: str = Field( - ..., description='Operating system and platform information' + ..., description="Operating system and platform information" ) installed_nodes: Optional[Dict[str, InstalledNodeInfo]] = Field( - None, description='Map of installed node packages by name' + None, description="Map of installed node packages by name" ) installed_models: Optional[Dict[str, InstalledModelInfo]] = Field( - None, description='Map of installed models by name' + None, description="Map of installed models by name" ) manager_config: Optional[Dict[str, Any]] = Field( - None, description='ComfyUI Manager configuration settings' + None, description="ComfyUI Manager configuration settings" ) comfyui_root_path: Optional[str] = Field( - None, description='ComfyUI root installation directory' + None, description="ComfyUI root installation directory" ) model_paths: Optional[Dict[str, List[str]]] = Field( - None, description='Map of model types to their configured paths' + None, description="Map of model types to their configured paths" ) - manager_version: Optional[str] = Field(None, description='ComfyUI Manager version') + manager_version: Optional[str] = Field(None, description="ComfyUI Manager version") security_level: Optional[SecurityLevel] = None network_mode: Optional[str] = Field( - None, description='Network mode (online, offline, private)' + None, description="Network mode (online, offline, private)" ) cli_args: Optional[Dict[str, Any]] = Field( - None, description='Selected ComfyUI CLI arguments' + None, description="Selected ComfyUI CLI arguments" ) custom_nodes_count: Optional[int] = Field( - None, description='Total number of custom node packages', ge=0 + None, description="Total number of custom node packages", ge=0 ) failed_imports: Optional[List[str]] = Field( - None, description='List of custom nodes that failed to import' + None, description="List of custom nodes that failed to import" ) pip_packages: Optional[Dict[str, str]] = Field( - None, description='Map of installed pip packages to their versions' + None, description="Map of installed pip packages to their versions" ) class BatchExecutionRecord(BaseModel): - batch_id: str = Field(..., description='Unique batch identifier') - start_time: datetime = Field(..., description='ISO timestamp when batch started') + 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' + None, description="ISO timestamp when batch completed" ) state_before: ComfyUISystemState state_after: Optional[ComfyUISystemState] = Field( - None, description='System state after batch execution' + None, description="System state after batch execution" ) operations: Optional[List[BatchOperation]] = Field( - None, description='List of operations performed in this batch' + None, description="List of operations performed in this batch" ) total_operations: Optional[int] = Field( - 0, description='Total number of operations in batch', ge=0 + 0, description="Total number of operations in batch", ge=0 ) successful_operations: Optional[int] = Field( - 0, description='Number of successful operations', ge=0 + 0, description="Number of successful operations", ge=0 ) failed_operations: Optional[int] = Field( - 0, description='Number of failed operations', ge=0 + 0, description="Number of failed operations", ge=0 ) skipped_operations: Optional[int] = Field( - 0, description='Number of skipped operations', ge=0 + 0, description="Number of skipped operations", ge=0 ) class QueueTaskItem(BaseModel): - ui_id: str = Field(..., description='Unique identifier for the task') - client_id: str = Field(..., description='Client identifier that initiated the task') + ui_id: str = Field(..., description="Unique identifier for the task") + client_id: str = Field(..., description="Client identifier that initiated the task") kind: OperationType params: Union[ InstallPackParams, @@ -516,60 +466,62 @@ class QueueTaskItem(BaseModel): 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') + 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' + ..., description="Map of task IDs to their history items" ) running_queue: List[QueueTaskItem] = Field( - ..., description='Currently executing tasks' + ..., description="Currently executing tasks" ) pending_queue: List[QueueTaskItem] = Field( - ..., description='Tasks waiting to be executed' + ..., description="Tasks waiting to be executed" ) installed_packs: Dict[str, ManagerPackInstalled] = Field( - ..., description='Map of currently installed node packages by name' + ..., description="Map of currently installed node packages by name" ) class MessageTaskDone(BaseModel): - ui_id: str = Field(..., description='Task identifier') - result: str = Field(..., description='Task result message') - kind: str = Field(..., description='Type of task') + 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') + 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') + 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') + 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(BaseModel): - __root__: Union[MessageTaskDone, MessageTaskStarted, MessageTaskFailed] = Field( - ..., description='Union type for all possible WebSocket message updates' +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' + None, description="Map of task IDs to their history items" )