Add base model and specific implementations for Vae, Lora, and Checkpoint; introduce exception handling and event types

lib/src/comfyui_api.dart

@@ -8,15 +8,13 @@ import 'package:web_socket_channel/web_socket_channel.dart';
 import 'models/websocket_event.dart';
 import 'models/progress_event.dart';
 import 'models/execution_event.dart';
-
+import 'exceptions/comfyui_api_exception.dart';
-/// Callback function type for prompt events
+import 'types/callback_types.dart';
-typedef PromptEventCallback = void Function(String promptId);
+import 'utils/websocket_event_handler.dart';
-
+import 'models/history_response.dart';
-/// Callback function type for typed WebSocket events
+import 'models/checkpoint.dart';
-typedef WebSocketEventCallback = void Function(WebSocketEvent event);
+import 'models/vae.dart';
-
+import 'models/lora.dart';
-/// Callback function type for progress events
-typedef ProgressEventCallback = void Function(ProgressEvent event);
 
 /// A Dart SDK for interacting with the ComfyUI API
 class ComfyUiApi {
@@ -149,44 +147,19 @@ class ComfyUiApi {
 
   /// Attempts to create a ProgressEvent from a WebSocketEvent
   void _tryCreateProgressEvent(WebSocketEvent event) {
-    if (event.data.containsKey('value') &&
+    WebSocketEventHandler.tryCreateProgressEvent(
-        event.data.containsKey('max') &&
+      event,
-        event.data.containsKey('prompt_id')) {
+      _progressEventController,
-      try {
+      _progressEventCallbacks,
-        final progressEvent = ProgressEvent(
+    );
-          value: event.data['value'] as int,
-          max: event.data['max'] as int,
-          promptId: event.data['prompt_id'] as String,
-          node: event.data['node']?.toString(),
-        );
-        _progressEventController.add(progressEvent);
-
-        // Trigger all registered progress callbacks
-        for (final callback in _progressEventCallbacks) {
-          callback(progressEvent);
-        }
-      } catch (e) {
-        print('Error creating ProgressEvent: $e');
-      }
-    }
   }
 
   /// Attempts to create an ExecutionEvent from a WebSocketEvent
   void _tryCreateExecutionEvent(WebSocketEvent event) {
-    if (event.data.containsKey('prompt_id')) {
+    WebSocketEventHandler.tryCreateExecutionEvent(
-      try {
+      event,
-        final executionEvent = ExecutionEvent(
+      _executionEventController,
-          promptId: event.data['prompt_id'] as String,
+    );
-          timestamp: event.data['timestamp'] as int? ??
-              DateTime.now().millisecondsSinceEpoch,
-          node: event.data['node']?.toString(),
-          extra: event.data['extra'] as Map<String, dynamic>?,
-        );
-        _executionEventController.add(executionEvent);
-      } catch (e) {
-        print('Error creating ExecutionEvent: $e');
-      }
-    }
   }
 
   /// Closes the WebSocket connection and cleans up resources
@@ -207,11 +180,11 @@ class ComfyUiApi {
   }
 
   /// Gets the history of the queue
-  Future<Map<String, dynamic>> getHistory({int maxItems = 64}) async {
+  Future<HistoryResponse> getHistory({int maxItems = 64}) async {
     final response = await _httpClient
         .get(Uri.parse('$host/api/history?max_items=$maxItems'));
     _validateResponse(response);
-    return jsonDecode(response.body);
+    return HistoryResponse.fromJson(jsonDecode(response.body));
   }
 
   /// Gets image data by filename
@@ -231,11 +204,12 @@ class ComfyUiApi {
   }
 
   /// Gets a list of checkpoints
-  Future<List<dynamic>> getCheckpoints() async {
+  Future<List<Checkpoint>> getCheckpoints() async {
     final response = await _httpClient
         .get(Uri.parse('$host/api/experiment/models/checkpoints'));
     _validateResponse(response);
-    return jsonDecode(response.body);
+    final List<dynamic> jsonData = jsonDecode(response.body);
+    return jsonData.map((item) => Checkpoint.fromJson(item)).toList();
   }
 
   /// Gets details for a specific checkpoint
@@ -248,11 +222,12 @@ class ComfyUiApi {
   }
 
   /// Gets a list of LoRAs
-  Future<List<dynamic>> getLoras() async {
+  Future<List<Lora>> getLoras() async {
     final response =
         await _httpClient.get(Uri.parse('$host/api/experiment/models/loras'));
     _validateResponse(response);
-    return jsonDecode(response.body);
+    final List<dynamic> jsonData = jsonDecode(response.body);
+    return jsonData.map((item) => Lora.fromJson(item)).toList();
   }
 
   /// Gets details for a specific LoRA
@@ -264,11 +239,12 @@ class ComfyUiApi {
   }
 
   /// Gets a list of VAEs
-  Future<List<dynamic>> getVaes() async {
+  Future<List<Vae>> getVaes() async {
     final response =
         await _httpClient.get(Uri.parse('$host/api/experiment/models/vae'));
     _validateResponse(response);
-    return jsonDecode(response.body);
+    final List<dynamic> jsonData = jsonDecode(response.body);
+    return jsonData.map((item) => Vae.fromJson(item)).toList();
   }
 
   /// Gets details for a specific VAE
@@ -350,14 +326,3 @@ class ComfyUiApi {
     }
   }
 }
-
-/// Exception thrown when the ComfyUI API returns an error
-class ComfyUiApiException implements Exception {
-  final int statusCode;
-  final String message;
-
-  ComfyUiApiException({required this.statusCode, required this.message});
-
-  @override
-  String toString() => 'ComfyUiApiException: $statusCode - $message';
-}

lib/src/exceptions/comfyui_api_exception.dart

@@ -0,0 +1,10 @@
+/// Exception thrown when the ComfyUI API returns an error
+class ComfyUiApiException implements Exception {
+  final int statusCode;
+  final String message;
+
+  ComfyUiApiException({required this.statusCode, required this.message});
+
+  @override
+  String toString() => 'ComfyUiApiException: $statusCode - $message';
+}

lib/src/models/base_model.dart

@@ -0,0 +1,6 @@
+abstract class BaseModel {
+  final String name;
+  final int pathIndex;
+
+  BaseModel({required this.name, required this.pathIndex});
+}

lib/src/models/checkpoint.dart

@@ -0,0 +1,13 @@
+import 'base_model.dart';
+
+class Checkpoint extends BaseModel {
+  Checkpoint({required String name, required int pathIndex})
+      : super(name: name, pathIndex: pathIndex);
+
+  factory Checkpoint.fromJson(Map<String, dynamic> json) {
+    return Checkpoint(
+      name: json['name'] as String,
+      pathIndex: json['pathIndex'] as int,
+    );
+  }
+}

lib/src/models/history_response.dart

@@ -0,0 +1,228 @@
+class HistoryResponse {
+  final Map<String, HistoryItem> items;
+
+  HistoryResponse({required this.items});
+
+  factory HistoryResponse.fromJson(Map<String, dynamic> json) {
+    return HistoryResponse(
+      items:
+          json.map((key, value) => MapEntry(key, HistoryItem.fromJson(value))),
+    );
+  }
+}
+
+class HistoryItem {
+  final Prompt prompt;
+  final Outputs outputs;
+  final Status status;
+  final Map<String, Meta> meta;
+
+  HistoryItem({
+    required this.prompt,
+    required this.outputs,
+    required this.status,
+    required this.meta,
+  });
+
+  factory HistoryItem.fromJson(Map<String, dynamic> json) {
+    return HistoryItem(
+      prompt: Prompt.fromJson(json['prompt']),
+      outputs: Outputs.fromJson(json['outputs']),
+      status: Status.fromJson(json['status']),
+      meta: (json['meta'] as Map<String, dynamic>).map(
+        (key, value) => MapEntry(key, Meta.fromJson(value)),
+      ),
+    );
+  }
+}
+
+class Prompt {
+  final int id;
+  final String promptId;
+  final Map<String, Node> nodes;
+  final ExtraPngInfo extraPngInfo;
+
+  Prompt({
+    required this.id,
+    required this.promptId,
+    required this.nodes,
+    required this.extraPngInfo,
+  });
+
+  factory Prompt.fromJson(List<dynamic> json) {
+    return Prompt(
+      id: json[0] as int,
+      promptId: json[1] as String,
+      nodes: (json[2] as Map<String, dynamic>).map(
+        (key, value) => MapEntry(key, Node.fromJson(value)),
+      ),
+      extraPngInfo: ExtraPngInfo.fromJson(json[3]['extra_pnginfo']),
+    );
+  }
+}
+
+class Node {
+  final Map<String, dynamic> inputs;
+  final String classType;
+  final Meta meta;
+
+  Node({
+    required this.inputs,
+    required this.classType,
+    required this.meta,
+  });
+
+  factory Node.fromJson(Map<String, dynamic> json) {
+    return Node(
+      inputs: json['inputs'] as Map<String, dynamic>,
+      classType: json['class_type'] as String,
+      meta: Meta.fromJson(json['_meta']),
+    );
+  }
+}
+
+class ExtraPngInfo {
+  final Workflow workflow;
+
+  ExtraPngInfo({required this.workflow});
+
+  factory ExtraPngInfo.fromJson(Map<String, dynamic> json) {
+    return ExtraPngInfo(
+      workflow: Workflow.fromJson(json['workflow']),
+    );
+  }
+}
+
+class Workflow {
+  final List<NodeInfo> nodes;
+  final List<Link> links;
+
+  Workflow({
+    required this.nodes,
+    required this.links,
+  });
+
+  factory Workflow.fromJson(Map<String, dynamic> json) {
+    return Workflow(
+      nodes: (json['nodes'] as List<dynamic>)
+          .map((e) => NodeInfo.fromJson(e))
+          .toList(),
+      links: (json['links'] as List<dynamic>)
+          .map((e) => Link.fromJson(e))
+          .toList(),
+    );
+  }
+}
+
+class NodeInfo {
+  final int id;
+  final String type;
+
+  NodeInfo({
+    required this.id,
+    required this.type,
+  });
+
+  factory NodeInfo.fromJson(Map<String, dynamic> json) {
+    return NodeInfo(
+      id: json['id'] as int,
+      type: json['type'] as String,
+    );
+  }
+}
+
+class Link {
+  final int id;
+  final int sourceNodeId;
+  final int targetNodeId;
+
+  Link({
+    required this.id,
+    required this.sourceNodeId,
+    required this.targetNodeId,
+  });
+
+  factory Link.fromJson(List<dynamic> json) {
+    return Link(
+      id: json[0] as int,
+      sourceNodeId: json[1] as int,
+      targetNodeId: json[2] as int,
+    );
+  }
+}
+
+class Outputs {
+  final Map<String, OutputNode> nodes;
+
+  Outputs({required this.nodes});
+
+  factory Outputs.fromJson(Map<String, dynamic> json) {
+    return Outputs(
+      nodes:
+          json.map((key, value) => MapEntry(key, OutputNode.fromJson(value))),
+    );
+  }
+}
+
+class OutputNode {
+  final List<Image> images;
+
+  OutputNode({required this.images});
+
+  factory OutputNode.fromJson(Map<String, dynamic> json) {
+    return OutputNode(
+      images: (json['images'] as List<dynamic>)
+          .map((e) => Image.fromJson(e))
+          .toList(),
+    );
+  }
+}
+
+class Image {
+  final String filename;
+  final String subfolder;
+  final String type;
+
+  Image({
+    required this.filename,
+    required this.subfolder,
+    required this.type,
+  });
+
+  factory Image.fromJson(Map<String, dynamic> json) {
+    return Image(
+      filename: json['filename'] as String,
+      subfolder: json['subfolder'] as String,
+      type: json['type'] as String,
+    );
+  }
+}
+
+class Status {
+  final String statusStr;
+  final bool completed;
+
+  Status({
+    required this.statusStr,
+    required this.completed,
+  });
+
+  factory Status.fromJson(Map<String, dynamic> json) {
+    return Status(
+      statusStr: json['status_str'] as String,
+      completed: json['completed'] as bool,
+    );
+  }
+}
+
+class Meta {
+  final String? nodeId;
+
+  Meta({this.nodeId});
+
+  factory Meta.fromJson(Map<String, dynamic> json) {
+    return Meta(
+      nodeId: json['node_id'] as String?,
+    );
+  }
+}

lib/src/models/lora.dart

@@ -0,0 +1,13 @@
+import 'base_model.dart';
+
+class Lora extends BaseModel {
+  Lora({required String name, required int pathIndex})
+      : super(name: name, pathIndex: pathIndex);
+
+  factory Lora.fromJson(Map<String, dynamic> json) {
+    return Lora(
+      name: json['name'] as String,
+      pathIndex: json['pathIndex'] as int,
+    );
+  }
+}

lib/src/models/models.dart

@@ -1,87 +0,0 @@
-/// Models that represent ComfyUI API responses
-
-/// Represents queue information from ComfyUI
-class QueueInfo {
-  final int queueRunning;
-  final List<Map<String, dynamic>> queue;
-  final Map<String, dynamic> queuePending;
-
-  QueueInfo({
-    required this.queueRunning,
-    required this.queue,
-    required this.queuePending,
-  });
-
-  factory QueueInfo.fromJson(Map<String, dynamic> json) {
-    return QueueInfo(
-      queueRunning: json['queue_running'] as int,
-      queue: List<Map<String, dynamic>>.from(json['queue'] ?? []),
-      queuePending: Map<String, dynamic>.from(json['queue_pending'] ?? {}),
-    );
-  }
-}
-
-/// Represents a prompt execution status
-class PromptExecutionStatus {
-  final String? promptId;
-  final int? number;
-  final String? status;
-  final dynamic error;
-
-  PromptExecutionStatus({
-    this.promptId,
-    this.number,
-    this.status,
-    this.error,
-  });
-
-  factory PromptExecutionStatus.fromJson(Map<String, dynamic> json) {
-    return PromptExecutionStatus(
-      promptId: json['prompt_id'] as String?,
-      number: json['number'] as int?,
-      status: json['status'] as String?,
-      error: json['error'],
-    );
-  }
-}
-
-/// Represents history data
-class HistoryItem {
-  final String promptId;
-  final Map<String, dynamic> prompt;
-  final Map<String, dynamic>? outputs;
-
-  HistoryItem({
-    required this.promptId,
-    required this.prompt,
-    this.outputs,
-  });
-
-  factory HistoryItem.fromJson(Map<String, dynamic> json) {
-    return HistoryItem(
-      promptId: json['prompt_id'] as String,
-      prompt: Map<String, dynamic>.from(json['prompt'] ?? {}),
-      outputs: json['outputs'] != null
-          ? Map<String, dynamic>.from(json['outputs'])
-          : null,
-    );
-  }
-}
-
-/// Represents a progress update received via WebSocket
-class ProgressUpdate {
-  final String type;
-  final Map<String, dynamic> data;
-
-  ProgressUpdate({
-    required this.type,
-    required this.data,
-  });
-
-  factory ProgressUpdate.fromJson(Map<String, dynamic> json) {
-    return ProgressUpdate(
-      type: json['type'] as String,
-      data: Map<String, dynamic>.from(json['data'] ?? {}),
-    );
-  }
-}

lib/src/models/vae.dart

@@ -0,0 +1,13 @@
+import 'base_model.dart';
+
+class Vae extends BaseModel {
+  Vae({required String name, required int pathIndex})
+      : super(name: name, pathIndex: pathIndex);
+
+  factory Vae.fromJson(Map<String, dynamic> json) {
+    return Vae(
+      name: json['name'] as String,
+      pathIndex: json['pathIndex'] as int,
+    );
+  }
+}

lib/src/types/callback_types.dart

@@ -0,0 +1,10 @@
+import 'package:comfyui_api_sdk/comfyui_api_sdk.dart';
+
+/// Callback function type for prompt events
+typedef PromptEventCallback = void Function(String promptId);
+
+/// Callback function type for typed WebSocket events
+typedef WebSocketEventCallback = void Function(WebSocketEvent event);
+
+/// Callback function type for progress events
+typedef ProgressEventCallback = void Function(ProgressEvent event);

lib/src/utils/websocket_event_handler.dart

@@ -0,0 +1,54 @@
+import 'dart:async';
+
+import '../models/websocket_event.dart';
+import '../models/progress_event.dart';
+import '../models/execution_event.dart';
+
+class WebSocketEventHandler {
+  static void tryCreateProgressEvent(
+    WebSocketEvent event,
+    StreamController<ProgressEvent> progressEventController,
+    List<Function(ProgressEvent)> progressEventCallbacks,
+  ) {
+    if (event.data.containsKey('value') &&
+        event.data.containsKey('max') &&
+        event.data.containsKey('prompt_id')) {
+      try {
+        final progressEvent = ProgressEvent(
+          value: event.data['value'] as int,
+          max: event.data['max'] as int,
+          promptId: event.data['prompt_id'] as String,
+          node: event.data['node']?.toString(),
+        );
+        progressEventController.add(progressEvent);
+
+        // Trigger all registered progress callbacks
+        for (final callback in progressEventCallbacks) {
+          callback(progressEvent);
+        }
+      } catch (e) {
+        print('Error creating ProgressEvent: $e');
+      }
+    }
+  }
+
+  static void tryCreateExecutionEvent(
+    WebSocketEvent event,
+    StreamController<ExecutionEvent> executionEventController,
+  ) {
+    if (event.data.containsKey('prompt_id')) {
+      try {
+        final executionEvent = ExecutionEvent(
+          promptId: event.data['prompt_id'] as String,
+          timestamp: event.data['timestamp'] as int? ??
+              DateTime.now().millisecondsSinceEpoch,
+          node: event.data['node']?.toString(),
+          extra: event.data['extra'] as Map<String, dynamic>?,
+        );
+        executionEventController.add(executionEvent);
+      } catch (e) {
+        print('Error creating ExecutionEvent: $e');
+      }
+    }
+  }
+}