[Refactor/TS] Simplify node filter logic (#3275)

src/components/searchbox/NodeSearchBox.vue

@@ -60,12 +60,17 @@
       <!-- FilterAndValue -->
       <template v-slot:chip="{ value }">
         <SearchFilterChip
-          v-if="Array.isArray(value) && value.length === 2"
-          :key="`${value[0].id}-${value[1]}`"
-          @remove="onRemoveFilter($event, value as FilterAndValue)"
-          :text="value[1]"
-          :badge="value[0].invokeSequence.toUpperCase()"
-          :badge-class="value[0].invokeSequence + '-badge'"
+          v-if="value.filterDef && value.value"
+          :key="`${value.filterDef.id}-${value.value}`"
+          @remove="
+            onRemoveFilter(
+              $event,
+              value as FuseFilterWithValue<ComfyNodeDefImpl, string>
+            )
+          "
+          :text="value.value"
+          :badge="value.filterDef.invokeSequence.toUpperCase()"
+          :badge-class="value.filterDef.invokeSequence + '-badge'"
         />
       </template>
     </AutoCompletePlus>
@@ -82,13 +87,13 @@ import NodePreview from '@/components/node/NodePreview.vue'
 import AutoCompletePlus from '@/components/primevueOverride/AutoCompletePlus.vue'
 import NodeSearchFilter from '@/components/searchbox/NodeSearchFilter.vue'
 import NodeSearchItem from '@/components/searchbox/NodeSearchItem.vue'
-import { type FilterAndValue } from '@/services/nodeSearchService'
 import {
   ComfyNodeDefImpl,
   useNodeDefStore,
   useNodeFrequencyStore
 } from '@/stores/nodeDefStore'
 import { useSettingStore } from '@/stores/settingStore'
+import type { FuseFilterWithValue } from '@/utils/fuseUtil'
 
 import SearchFilterChip from '../common/SearchFilterChip.vue'
 
@@ -100,7 +105,7 @@ const enableNodePreview = computed(() =>
 )
 
 const { filters, searchLimit = 64 } = defineProps<{
-  filters: FilterAndValue[]
+  filters: FuseFilterWithValue<ComfyNodeDefImpl, string>[]
   searchLimit?: number
 }>()
 
@@ -139,11 +144,16 @@ const reFocusInput = () => {
 }
 
 onMounted(reFocusInput)
-const onAddFilter = (filterAndValue: FilterAndValue) => {
+const onAddFilter = (
+  filterAndValue: FuseFilterWithValue<ComfyNodeDefImpl, string>
+) => {
   nodeSearchFilterVisible.value = false
   emit('addFilter', filterAndValue)
 }
-const onRemoveFilter = (event: Event, filterAndValue: FilterAndValue) => {
+const onRemoveFilter = (
+  event: Event,
+  filterAndValue: FuseFilterWithValue<ComfyNodeDefImpl, string>
+) => {
   event.stopPropagation()
   event.preventDefault()
   emit('removeFilter', filterAndValue)

src/components/searchbox/NodeSearchBoxPopover.vue

@@ -46,13 +46,13 @@ import Dialog from 'primevue/dialog'
 import { computed, ref, toRaw, watchEffect } from 'vue'
 
 import { useLitegraphService } from '@/services/litegraphService'
-import { FilterAndValue } from '@/services/nodeSearchService'
 import { useCanvasStore } from '@/stores/graphStore'
 import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
 import { useSettingStore } from '@/stores/settingStore'
 import { useSearchBoxStore } from '@/stores/workspace/searchBoxStore'
 import { ConnectingLinkImpl } from '@/types/litegraphTypes'
 import { LinkReleaseTriggerAction } from '@/types/searchBoxTypes'
+import { FuseFilterWithValue } from '@/utils/fuseUtil'
 
 import NodeSearchBox from './NodeSearchBox.vue'
 
@@ -71,11 +71,13 @@ const getNewNodeLocation = (): Vector2 => {
     .originalEvent
   return [originalEvent.canvasX, originalEvent.canvasY]
 }
-const nodeFilters = ref<FilterAndValue[]>([])
-const addFilter = (filter: FilterAndValue) => {
+const nodeFilters = ref<FuseFilterWithValue<ComfyNodeDefImpl, string>[]>([])
+const addFilter = (filter: FuseFilterWithValue<ComfyNodeDefImpl, string>) => {
   nodeFilters.value.push(filter)
 }
-const removeFilter = (filter: FilterAndValue) => {
+const removeFilter = (
+  filter: FuseFilterWithValue<ComfyNodeDefImpl, string>
+) => {
   nodeFilters.value = nodeFilters.value.filter(
     (f) => toRaw(f) !== toRaw(filter)
   )
@@ -136,13 +138,16 @@ const showNewSearchBox = (e: LiteGraphCanvasEvent) => {
       return
     }
     const firstLink = ConnectingLinkImpl.createFromPlainObject(links[0])
-    const filter = nodeDefStore.nodeSearchService.getFilterById(
-      firstLink.releaseSlotType
-    )
-    // @ts-expect-error fixme ts strict error
-    const dataType = firstLink.type.toString()
-    // @ts-expect-error fixme ts strict error
-    addFilter([filter, dataType])
+    const filter =
+      firstLink.releaseSlotType === 'input'
+        ? nodeDefStore.nodeSearchService.inputTypeFilter
+        : nodeDefStore.nodeSearchService.outputTypeFilter
+
+    const dataType = firstLink.type?.toString() ?? ''
+    addFilter({
+      filterDef: filter,
+      value: dataType
+    })
   }
 
   visible.value = true

src/components/searchbox/NodeSearchFilter.vue

@@ -27,11 +27,11 @@ import Select from 'primevue/select'
 import SelectButton from 'primevue/selectbutton'
 import { computed, onMounted, ref } from 'vue'
 
-import { type FilterAndValue, NodeFilter } from '@/services/nodeSearchService'
-import { useNodeDefStore } from '@/stores/nodeDefStore'
+import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import { FuseFilter, FuseFilterWithValue } from '@/utils/fuseUtil'
 
 const filters = computed(() => nodeDefStore.nodeSearchService.nodeFilters)
-const selectedFilter = ref<NodeFilter>()
+const selectedFilter = ref<FuseFilter<ComfyNodeDefImpl, string>>()
 const filterValues = computed(() => selectedFilter.value?.fuseSearch.data ?? [])
 const selectedFilterValue = ref<string>('')
 
@@ -43,7 +43,10 @@ onMounted(() => {
 })
 
 const emit = defineEmits<{
-  (event: 'addFilter', filterAndValue: FilterAndValue): void
+  (
+    event: 'addFilter',
+    filterAndValue: FuseFilterWithValue<ComfyNodeDefImpl, string>
+  ): void
 }>()
 
 const updateSelectedFilterValue = () => {
@@ -54,10 +57,13 @@ const updateSelectedFilterValue = () => {
 }
 
 const submit = () => {
-  emit('addFilter', [
-    selectedFilter.value,
-    selectedFilterValue.value
-  ] as FilterAndValue)
+  if (!selectedFilter.value) {
+    return
+  }
+  emit('addFilter', {
+    filterDef: selectedFilter.value,
+    value: selectedFilterValue.value
+  })
 }
 </script>
 

src/components/sidebar/tabs/NodeLibrarySidebarTab.vue

@@ -76,7 +76,6 @@ import SidebarTabTemplate from '@/components/sidebar/tabs/SidebarTabTemplate.vue
 import NodeTreeLeaf from '@/components/sidebar/tabs/nodeLibrary/NodeTreeLeaf.vue'
 import { useTreeExpansion } from '@/composables/useTreeExpansion'
 import { useLitegraphService } from '@/services/litegraphService'
-import { FilterAndValue } from '@/services/nodeSearchService'
 import { useNodeBookmarkStore } from '@/stores/nodeBookmarkStore'
 import {
   ComfyNodeDefImpl,
@@ -85,6 +84,7 @@ import {
 } from '@/stores/nodeDefStore'
 import type { TreeNode } from '@/types/treeExplorerTypes'
 import type { TreeExplorerNode } from '@/types/treeExplorerTypes'
+import { FuseFilterWithValue } from '@/utils/fuseUtil'
 import { sortedTree } from '@/utils/treeUtil'
 
 import NodeBookmarkTreeExplorer from './nodeLibrary/NodeBookmarkTreeExplorer.vue'
@@ -150,8 +150,9 @@ const filteredRoot = computed<TreeNode | null>(() => {
   }
   return buildNodeDefTree(filteredNodeDefs.value)
 })
-const filters: Ref<Array<SearchFilter & { filter: FilterAndValue<string> }>> =
-  ref([])
+const filters: Ref<
+  (SearchFilter & { filter: FuseFilterWithValue<ComfyNodeDefImpl, string> })[]
+> = ref([])
 const handleSearch = (query: string) => {
   // Don't apply a min length filter because it does not make sense in
   // multi-byte languages like Chinese, Japanese, Korean, etc.
@@ -161,7 +162,7 @@ const handleSearch = (query: string) => {
     return
   }
 
-  const f = filters.value.map((f) => f.filter as FilterAndValue<string>)
+  const f = filters.value.map((f) => f.filter)
   filteredNodeDefs.value = nodeDefStore.nodeSearchService.searchNode(
     query,
     f,
@@ -179,12 +180,14 @@ const handleSearch = (query: string) => {
   })
 }
 
-const onAddFilter = (filterAndValue: FilterAndValue) => {
+const onAddFilter = (
+  filterAndValue: FuseFilterWithValue<ComfyNodeDefImpl, string>
+) => {
   filters.value.push({
     filter: filterAndValue,
-    badge: filterAndValue[0].invokeSequence.toUpperCase(),
-    badgeClass: filterAndValue[0].invokeSequence + '-badge',
-    text: filterAndValue[1],
+    badge: filterAndValue.filterDef.invokeSequence.toUpperCase(),
+    badgeClass: filterAndValue.filterDef.invokeSequence + '-badge',
+    text: filterAndValue.value,
     id: +new Date()
   })
 

src/services/nodeSearchService.ts

@@ -1,74 +1,14 @@
-import { FuseSearchOptions, IFuseOptions } from 'fuse.js'
-import _ from 'lodash'
+import { FuseSearchOptions } from 'fuse.js'
 
 import { ComfyNodeDefImpl } from '@/stores/nodeDefStore'
-import { FuseSearch } from '@/utils/fuseUtil'
-
-export type SearchAuxScore = number[]
-
-interface ExtraSearchOptions {
-  matchWildcards?: boolean
-}
-
-export type FilterAndValue<T = string> = [NodeFilter<T>, T]
-
-export class NodeFilter<FilterOptionT = string> {
-  public readonly fuseSearch: FuseSearch<FilterOptionT>
-
-  constructor(
-    public readonly id: string,
-    public readonly name: string,
-    public readonly invokeSequence: string,
-    public readonly longInvokeSequence: string,
-    public readonly nodeOptions:
-      | FilterOptionT[]
-      | ((node: ComfyNodeDefImpl) => FilterOptionT[]),
-    nodeDefs: ComfyNodeDefImpl[],
-    options?: IFuseOptions<FilterOptionT>
-  ) {
-    this.fuseSearch = new FuseSearch(this.getAllNodeOptions(nodeDefs), {
-      fuseOptions: options
-    })
-  }
-
-  public getNodeOptions(node: ComfyNodeDefImpl): FilterOptionT[] {
-    return this.nodeOptions instanceof Function
-      ? this.nodeOptions(node)
-      : this.nodeOptions
-  }
-
-  public getAllNodeOptions(nodeDefs: ComfyNodeDefImpl[]): FilterOptionT[] {
-    // @ts-expect-error fixme ts strict error
-    return [
-      ...new Set(
-        // @ts-expect-error fixme ts strict error
-        nodeDefs.reduce((acc, nodeDef) => {
-          return [...acc, ...this.getNodeOptions(nodeDef)]
-        }, [])
-      )
-    ]
-  }
-
-  public matches(
-    node: ComfyNodeDefImpl,
-    value: FilterOptionT,
-    extraOptions?: ExtraSearchOptions
-  ): boolean {
-    const matchWildcards = extraOptions?.matchWildcards !== false
-    if (matchWildcards && value === '*') {
-      return true
-    }
-    const options = this.getNodeOptions(node)
-    return (
-      options.includes(value) ||
-      (matchWildcards && _.some(options, (option) => option === '*'))
-    )
-  }
-}
+import { FuseFilter, FuseFilterWithValue, FuseSearch } from '@/utils/fuseUtil'
 
 export class NodeSearchService {
   public readonly nodeFuseSearch: FuseSearch<ComfyNodeDefImpl>
-  public readonly nodeFilters: NodeFilter<string>[]
+  public readonly inputTypeFilter: FuseFilter<ComfyNodeDefImpl, string>
+  public readonly outputTypeFilter: FuseFilter<ComfyNodeDefImpl, string>
+  public readonly nodeCategoryFilter: FuseFilter<ComfyNodeDefImpl, string>
+  public readonly nodeSourceFilter: FuseFilter<ComfyNodeDefImpl, string>
 
   constructor(data: ComfyNodeDefImpl[]) {
     this.nodeFuseSearch = new FuseSearch(data, {
@@ -83,83 +23,74 @@ export class NodeSearchService {
       advancedScoring: true
     })
 
-    const filterSearchOptions = {
+    const fuseOptions = {
       includeScore: true,
       threshold: 0.3,
       shouldSort: true
     }
 
-    const inputTypeFilter = new NodeFilter<string>(
-      /* id */ 'input',
-      /* name */ 'Input Type',
-      /* invokeSequence */ 'i',
-      /* longInvokeSequence */ 'input',
-      (node) => Object.values(node.inputs).map((input) => input.type),
-      data,
-      filterSearchOptions
-    )
+    this.inputTypeFilter = new FuseFilter<ComfyNodeDefImpl, string>(data, {
+      id: 'input',
+      name: 'Input Type',
+      invokeSequence: 'i',
+      getItemOptions: (node) =>
+        Object.values(node.inputs).map((input) => input.type),
+      fuseOptions
+    })
 
-    const outputTypeFilter = new NodeFilter<string>(
-      /* id */ 'output',
-      /* name */ 'Output Type',
-      /* invokeSequence */ 'o',
-      /* longInvokeSequence */ 'output',
-      (node) => node.outputs.map((output) => output.type),
-      data,
-      filterSearchOptions
-    )
+    this.outputTypeFilter = new FuseFilter<ComfyNodeDefImpl, string>(data, {
+      id: 'output',
+      name: 'Output Type',
+      invokeSequence: 'o',
+      getItemOptions: (node) => node.outputs.map((output) => output.type),
+      fuseOptions
+    })
 
-    const nodeCategoryFilter = new NodeFilter<string>(
-      /* id */ 'category',
-      /* name */ 'Category',
-      /* invokeSequence */ 'c',
-      /* longInvokeSequence */ 'category',
-      (node) => [node.category],
-      data,
-      filterSearchOptions
-    )
+    this.nodeCategoryFilter = new FuseFilter<ComfyNodeDefImpl, string>(data, {
+      id: 'category',
+      name: 'Category',
+      invokeSequence: 'c',
+      getItemOptions: (node) => [node.category],
+      fuseOptions
+    })
 
-    const nodeSourceFilter = new NodeFilter<string>(
-      /* id */ 'source',
-      /* name */ 'Source',
-      /* invokeSequence */ 's',
-      /* longInvokeSequence */ 'source',
-      (node) => [node.nodeSource.displayText],
-      data,
-      filterSearchOptions
-    )
-
-    this.nodeFilters = [
-      inputTypeFilter,
-      outputTypeFilter,
-      nodeCategoryFilter,
-      nodeSourceFilter
-    ]
-  }
-
-  public endsWithFilterStartSequence(query: string): boolean {
-    return query.endsWith(':')
+    this.nodeSourceFilter = new FuseFilter<ComfyNodeDefImpl, string>(data, {
+      id: 'source',
+      name: 'Source',
+      invokeSequence: 's',
+      getItemOptions: (node) => [node.nodeSource.displayText],
+      fuseOptions
+    })
   }
 
   public searchNode(
     query: string,
-    filters: FilterAndValue<string>[] = [],
+    filters: FuseFilterWithValue<ComfyNodeDefImpl, string>[] = [],
     options?: FuseSearchOptions,
-    extraOptions?: ExtraSearchOptions
+    extraOptions: {
+      matchWildcards?: boolean
+    } = {}
   ): ComfyNodeDefImpl[] {
+    const { matchWildcards = true } = extraOptions
+    const wildcard = matchWildcards ? '*' : undefined
     const matchedNodes = this.nodeFuseSearch.search(query)
 
     const results = matchedNodes.filter((node) => {
-      return _.every(filters, (filterAndValue) => {
-        const [filter, value] = filterAndValue
-        return filter.matches(node, value, extraOptions)
+      return filters.every((filterAndValue) => {
+        const { filterDef, value } = filterAndValue
+        return filterDef.matches(node, value, { wildcard })
       })
     })
 
     return options?.limit ? results.slice(0, options.limit) : results
   }
 
-  public getFilterById(id: string): NodeFilter<string> | undefined {
-    return this.nodeFilters.find((filter) => filter.id === id)
+  get nodeFilters(): FuseFilter<ComfyNodeDefImpl, string>[] {
+    return [
+      this.inputTypeFilter,
+      this.outputTypeFilter,
+      this.nodeCategoryFilter,
+      this.nodeSourceFilter
+    ]
   }
 }

src/stores/nodeDefStore.ts

@@ -14,19 +14,19 @@ import type {
   ComfyNodeDef as ComfyNodeDefV1,
   ComfyOutputTypesSpec as ComfyOutputSpecV1
 } from '@/schemas/nodeDefSchema'
-import {
-  NodeSearchService,
-  type SearchAuxScore
-} from '@/services/nodeSearchService'
+import { NodeSearchService } from '@/services/nodeSearchService'
 import {
   type NodeSource,
   NodeSourceType,
   getNodeSource
 } from '@/types/nodeSource'
 import type { TreeNode } from '@/types/treeExplorerTypes'
+import type { FuseSearchable, SearchAuxScore } from '@/utils/fuseUtil'
 import { buildTree } from '@/utils/treeUtil'
 
-export class ComfyNodeDefImpl implements ComfyNodeDefV1, ComfyNodeDefV2 {
+export class ComfyNodeDefImpl
+  implements ComfyNodeDefV1, ComfyNodeDefV2, FuseSearchable
+{
   // ComfyNodeDef fields (V1)
   readonly name: string
   readonly display_name: string

src/utils/fuseUtil.ts

@@ -1,6 +1,80 @@
 import Fuse, { FuseOptionKey, FuseSearchOptions, IFuseOptions } from 'fuse.js'
 
-type SearchAuxScore = number[]
+export type SearchAuxScore = number[]
+
+export interface FuseFilterWithValue<T, O = string> {
+  filterDef: FuseFilter<T, O>
+  value: O
+}
+
+export class FuseFilter<T, O = string> {
+  public readonly fuseSearch: FuseSearch<O>
+  /** The unique identifier for the filter. */
+  public readonly id: string
+  /** The name of the filter for display purposes. */
+  public readonly name: string
+  /** The sequence of characters to invoke the filter. */
+  public readonly invokeSequence: string
+  /** A function that returns the options for the filter. */
+  public readonly getItemOptions: (item: T) => O[]
+
+  constructor(
+    data: T[],
+    options: {
+      id: string
+      name: string
+      invokeSequence: string
+      getItemOptions: (item: T) => O[]
+      fuseOptions?: IFuseOptions<O>
+    }
+  ) {
+    this.id = options.id
+    this.name = options.name
+    this.invokeSequence = options.invokeSequence
+    this.getItemOptions = options.getItemOptions
+
+    this.fuseSearch = new FuseSearch(this.getAllNodeOptions(data), {
+      fuseOptions: options.fuseOptions
+    })
+  }
+
+  public getAllNodeOptions(data: T[]): O[] {
+    const options = new Set<O>()
+    for (const item of data) {
+      for (const option of this.getItemOptions(item)) {
+        options.add(option)
+      }
+    }
+    return Array.from(options)
+  }
+
+  public matches(
+    item: T,
+    value: O,
+    extraOptions: {
+      wildcard?: O
+    } = {}
+  ): boolean {
+    const { wildcard } = extraOptions
+
+    if (wildcard && value === wildcard) {
+      return true
+    }
+    const options = this.getItemOptions(item)
+    return (
+      options.includes(value) ||
+      (!!wildcard && options.some((option) => option === wildcard))
+    )
+  }
+}
+
+export interface FuseSearchable {
+  postProcessSearchScores: (scores: SearchAuxScore) => SearchAuxScore
+}
+
+function isFuseSearchable(item: any): item is FuseSearchable {
+  return 'postProcessSearchScores' in item
+}
 
 /**
  * A wrapper around Fuse.js that provides a more type-safe API.
@@ -56,20 +130,22 @@ export class FuseSearch<T> {
 
   public calcAuxScores(query: string, entry: T, score: number): SearchAuxScore {
     let values: string[] = []
-    if (!this.keys.length) values = [entry as string]
-    // @ts-expect-error fixme ts strict error
-    else values = this.keys.map((x) => entry[x])
+    if (typeof entry === 'string') {
+      values = [entry]
+    } else if (typeof entry === 'object' && entry !== null) {
+      values = this.keys
+        .map((x) => entry[x as keyof T])
+        .filter((x) => typeof x === 'string') as string[]
+    }
     const scores = values.map((x) => this.calcAuxSingle(query, x, score))
     let result = scores.sort(this.compareAux)[0]
 
     const deprecated = values.some((x) =>
       x.toLocaleLowerCase().includes('deprecated')
     )
-    result[0] += deprecated && result[0] != 0 ? 5 : 0
-    // @ts-expect-error fixme ts strict error
-    if (entry['postProcessSearchScores']) {
-      // @ts-expect-error fixme ts strict error
-      result = entry['postProcessSearchScores'](result) as SearchAuxScore
+    result[0] += deprecated && result[0] !== 0 ? 5 : 0
+    if (isFuseSearchable(entry)) {
+      result = entry.postProcessSearchScores(result)
     }
     return result
   }

src/views/GraphView.vue

@@ -232,7 +232,7 @@ const onGraphReady = () => {
       // Node defs now available after comfyApp.setup.
       // Explicitly initialize nodeSearchService to avoid indexing delay when
       // node search is triggered
-      useNodeDefStore().nodeSearchService.endsWithFilterStartSequence('')
+      useNodeDefStore().nodeSearchService.searchNode('')
     },
     { timeout: 1000 }
   )

tests-ui/tests/nodeSearchService.test.ts

@@ -64,12 +64,14 @@ const EXAMPLE_NODE_DEFS: ComfyNodeDefImpl[] = (
 describe('nodeSearchService', () => {
   it('searches with input filter', () => {
     const service = new NodeSearchService(EXAMPLE_NODE_DEFS)
-    const inputFilter = service.getFilterById('input')
-    // @ts-expect-error fixme ts strict error
-    expect(service.searchNode('L', [[inputFilter, 'LATENT']])).toHaveLength(1)
+    const inputFilter = service.inputTypeFilter
+    expect(
+      service.searchNode('L', [{ filterDef: inputFilter, value: 'LATENT' }])
+    ).toHaveLength(1)
     // Wildcard should match all.
-    // @ts-expect-error fixme ts strict error
-    expect(service.searchNode('L', [[inputFilter, '*']])).toHaveLength(2)
+    expect(
+      service.searchNode('L', [{ filterDef: inputFilter, value: '*' }])
+    ).toHaveLength(2)
     expect(service.searchNode('L')).toHaveLength(2)
   })
 })