From b4d773585513cacd2b06c723e3369ef174221bb5 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Fri, 26 Jul 2024 13:41:24 -0400
Subject: [PATCH] Use ComfyNodeDefImpl on nodeSearchService  (#233)

* Make nodeSearchService use ComfyNodeDefImpl

* Fix test
---
 src/components/NodeSearchBox.vue         | 10 ++-
 src/services/nodeSearchService.ts        | 84 +++++-------------------
 src/stores/nodeDefStore.ts               |  6 +-
 tests-ui/tests/nodeSearchService.test.ts |  7 +-
 4 files changed, 30 insertions(+), 77 deletions(-)

diff --git a/src/components/NodeSearchBox.vue b/src/components/NodeSearchBox.vue
index 1979a56a4..cfa041492 100644
--- a/src/components/NodeSearchBox.vue
+++ b/src/components/NodeSearchBox.vue
@@ -2,7 +2,7 @@
   <div class="comfy-vue-node-search-container">
     <div class="comfy-vue-node-preview-container">
       <NodePreview
-        :nodeDef="plainToClass(ComfyNodeDefImpl, hoveredSuggestion)"
+        :nodeDef="hoveredSuggestion"
         :key="hoveredSuggestion?.name || ''"
         v-if="hoveredSuggestion"
       />
@@ -53,17 +53,15 @@
 </template>
 
 <script setup lang="ts">
-import { computed, inject, onMounted, Ref, ref } from 'vue'
+import { computed, onMounted, ref } from 'vue'
 import AutoCompletePlus from './primevueOverride/AutoCompletePlus.vue'
 import Chip from 'primevue/chip'
 import Badge from 'primevue/badge'
 import NodeSearchFilter from '@/components/NodeSearchFilter.vue'
 import NodeSourceChip from '@/components/NodeSourceChip.vue'
-import { ComfyNodeDef } from '@/types/apiTypes'
 import { type FilterAndValue } from '@/services/nodeSearchService'
 import NodePreview from './NodePreview.vue'
 import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
-import { plainToClass } from 'class-transformer'
 
 const props = defineProps({
   filters: {
@@ -76,8 +74,8 @@ const props = defineProps({
 })
 
 const inputId = `comfy-vue-node-search-box-input-${Math.random()}`
-const suggestions = ref<ComfyNodeDef[]>([])
-const hoveredSuggestion = ref<ComfyNodeDef | null>(null)
+const suggestions = ref<ComfyNodeDefImpl[]>([])
+const hoveredSuggestion = ref<ComfyNodeDefImpl | null>(null)
 const placeholder = computed(() => {
   return props.filters.length === 0 ? 'Search for nodes' : ''
 })
diff --git a/src/services/nodeSearchService.ts b/src/services/nodeSearchService.ts
index b13822d08..7492c4cb9 100644
--- a/src/services/nodeSearchService.ts
+++ b/src/services/nodeSearchService.ts
@@ -1,44 +1,8 @@
-import { ComfyNodeDef } from '@/types/apiTypes'
+import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
 import { getNodeSource } from '@/types/nodeSource'
 import Fuse, { IFuseOptions, FuseSearchOptions } from 'fuse.js'
 import _ from 'lodash'
 
-export const SYSTEM_NODE_DEFS: ComfyNodeDef[] = [
-  {
-    name: 'PrimitiveNode',
-    display_name: 'Primitive',
-    category: 'utils',
-    input: { required: {}, optional: {} },
-    output: ['*'],
-    output_name: ['connect to widget input'],
-    output_is_list: [false],
-    python_module: 'nodes',
-    description: 'Primitive values like numbers, strings, and booleans.'
-  },
-  {
-    name: 'Reroute',
-    display_name: 'Reroute',
-    category: 'utils',
-    input: { required: { '': ['*'] }, optional: {} },
-    output: ['*'],
-    output_name: [''],
-    output_is_list: [false],
-    python_module: 'nodes',
-    description: 'Reroute the connection to another node.'
-  },
-  {
-    name: 'Note',
-    display_name: 'Note',
-    category: 'utils',
-    input: { required: {}, optional: {} },
-    output: [],
-    output_name: [],
-    output_is_list: [],
-    python_module: 'nodes',
-    description: 'Node that add notes to your project'
-  }
-]
-
 export class FuseSearch<T> {
   private fuse: Fuse<T>
   public readonly data: T[]
@@ -73,11 +37,14 @@ export abstract class NodeFilter<FilterOptionT = string> {
   public abstract readonly longInvokeSequence: string
   public readonly fuseSearch: FuseSearch<FilterOptionT>
 
-  constructor(nodeDefs: ComfyNodeDef[], options?: IFuseOptions<FilterOptionT>) {
+  constructor(
+    nodeDefs: ComfyNodeDefImpl[],
+    options?: IFuseOptions<FilterOptionT>
+  ) {
     this.fuseSearch = new FuseSearch(this.getAllNodeOptions(nodeDefs), options)
   }
 
-  private getAllNodeOptions(nodeDefs: ComfyNodeDef[]): FilterOptionT[] {
+  private getAllNodeOptions(nodeDefs: ComfyNodeDefImpl[]): FilterOptionT[] {
     return [
       ...new Set(
         nodeDefs.reduce((acc, nodeDef) => {
@@ -87,9 +54,9 @@ export abstract class NodeFilter<FilterOptionT = string> {
     ]
   }
 
-  public abstract getNodeOptions(node: ComfyNodeDef): FilterOptionT[]
+  public abstract getNodeOptions(node: ComfyNodeDefImpl): FilterOptionT[]
 
-  public matches(node: ComfyNodeDef, value: FilterOptionT): boolean {
+  public matches(node: ComfyNodeDefImpl, value: FilterOptionT): boolean {
     return this.getNodeOptions(node).includes(value)
   }
 }
@@ -100,15 +67,8 @@ export class InputTypeFilter extends NodeFilter<string> {
   public readonly invokeSequence = 'i'
   public readonly longInvokeSequence = 'input'
 
-  public override getNodeOptions(node: ComfyNodeDef): string[] {
-    const inputs = {
-      ...(node.input.required || {}),
-      ...(node.input.optional || {})
-    }
-    return Object.values(inputs).map((input) => {
-      const [inputType, inputSpec] = input
-      return typeof inputType === 'string' ? inputType : 'COMBO'
-    })
+  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
+    return node.input.all.map((input) => input.type)
   }
 }
 
@@ -118,18 +78,8 @@ export class OutputTypeFilter extends NodeFilter<string> {
   public readonly invokeSequence = 'o'
   public readonly longInvokeSequence = 'output'
 
-  public override getNodeOptions(node: ComfyNodeDef): string[] {
-    const outputs = node.output || []
-    // "custom_nodes.was-node-suite-comfyui"
-    // has a custom node with an output that is not an array.
-    // https://github.com/WASasquatch/was-node-suite-comfyui/pull/440
-    if (!(outputs instanceof Array)) {
-      console.error('Invalid output type', node)
-      return []
-    }
-    return outputs.map((output) => {
-      return typeof output === 'string' ? output : output[0]
-    })
+  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
+    return node.output.all.map((output) => output.type)
   }
 }
 
@@ -139,7 +89,7 @@ export class NodeSourceFilter extends NodeFilter<string> {
   public readonly invokeSequence = 's'
   public readonly longInvokeSequence = 'source'
 
-  public override getNodeOptions(node: ComfyNodeDef): string[] {
+  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
     return [getNodeSource(node.python_module).displayText]
   }
 }
@@ -150,16 +100,16 @@ export class NodeCategoryFilter extends NodeFilter<string> {
   public readonly invokeSequence = 'c'
   public readonly longInvokeSequence = 'category'
 
-  public override getNodeOptions(node: ComfyNodeDef): string[] {
+  public override getNodeOptions(node: ComfyNodeDefImpl): string[] {
     return [node.category]
   }
 }
 
 export class NodeSearchService {
-  public readonly nodeFuseSearch: FuseSearch<ComfyNodeDef>
+  public readonly nodeFuseSearch: FuseSearch<ComfyNodeDefImpl>
   public readonly nodeFilters: NodeFilter<string>[]
 
-  constructor(data: ComfyNodeDef[]) {
+  constructor(data: ComfyNodeDefImpl[]) {
     this.nodeFuseSearch = new FuseSearch(data, {
       keys: ['name', 'display_name', 'description'],
       includeScore: true,
@@ -192,7 +142,7 @@ export class NodeSearchService {
     query: string,
     filters: FilterAndValue<string>[] = [],
     options?: FuseSearchOptions
-  ): ComfyNodeDef[] {
+  ): ComfyNodeDefImpl[] {
     const matchedNodes = this.nodeFuseSearch.search(query)
 
     const results = matchedNodes.filter((node) => {
diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts
index f95995072..5cc502b2c 100644
--- a/src/stores/nodeDefStore.ts
+++ b/src/stores/nodeDefStore.ts
@@ -236,7 +236,11 @@ export const useNodeDefStore = defineStore('nodeDef', {
       return Object.values(state.nodeDefsByName)
     },
     nodeSearchService(state) {
-      return new NodeSearchService(Object.values(state.nodeDefsByName))
+      return new NodeSearchService(
+        Object.values(state.nodeDefsByName).map((nodeDef) =>
+          plainToClass(ComfyNodeDefImpl, nodeDef)
+        )
+      )
     }
   },
   actions: {
diff --git a/tests-ui/tests/nodeSearchService.test.ts b/tests-ui/tests/nodeSearchService.test.ts
index 8901828a1..b3cdf9ad5 100644
--- a/tests-ui/tests/nodeSearchService.test.ts
+++ b/tests-ui/tests/nodeSearchService.test.ts
@@ -1,7 +1,8 @@
 import { NodeSearchService } from '@/services/nodeSearchService'
-import { ComfyNodeDef } from '@/types/apiTypes'
+import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
+import { plainToClass } from 'class-transformer'
 
-const EXAMPLE_NODE_DEFS: ComfyNodeDef[] = [
+const EXAMPLE_NODE_DEFS: ComfyNodeDefImpl[] = [
   {
     input: {
       required: {
@@ -50,7 +51,7 @@ const EXAMPLE_NODE_DEFS: ComfyNodeDef[] = [
     category: 'latent/batch',
     output_node: false
   }
-]
+].map((nodeDef) => plainToClass(ComfyNodeDefImpl, nodeDef))
 
 describe('nodeSearchService', () => {
   it('searches with input filter', () => {