From fcc22f06ac5b80800d37af839e0406408fa552cb Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sat, 29 Mar 2025 13:00:18 -0400 Subject: [PATCH] [Refactor/TS] Simplify node filter logic (#3275) --- src/components/searchbox/NodeSearchBox.vue | 30 ++- .../searchbox/NodeSearchBoxPopover.vue | 27 +-- src/components/searchbox/NodeSearchFilter.vue | 22 ++- .../sidebar/tabs/NodeLibrarySidebarTab.vue | 19 +- src/services/nodeSearchService.ts | 173 ++++++------------ src/stores/nodeDefStore.ts | 10 +- src/utils/fuseUtil.ts | 94 +++++++++- src/views/GraphView.vue | 2 +- tests-ui/tests/nodeSearchService.test.ts | 12 +- 9 files changed, 211 insertions(+), 178 deletions(-) diff --git a/src/components/searchbox/NodeSearchBox.vue b/src/components/searchbox/NodeSearchBox.vue index 4409f9384..d9dddeb1a 100644 --- a/src/components/searchbox/NodeSearchBox.vue +++ b/src/components/searchbox/NodeSearchBox.vue @@ -60,12 +60,17 @@ @@ -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[] searchLimit?: number }>() @@ -139,11 +144,16 @@ const reFocusInput = () => { } onMounted(reFocusInput) -const onAddFilter = (filterAndValue: FilterAndValue) => { +const onAddFilter = ( + filterAndValue: FuseFilterWithValue +) => { nodeSearchFilterVisible.value = false emit('addFilter', filterAndValue) } -const onRemoveFilter = (event: Event, filterAndValue: FilterAndValue) => { +const onRemoveFilter = ( + event: Event, + filterAndValue: FuseFilterWithValue +) => { event.stopPropagation() event.preventDefault() emit('removeFilter', filterAndValue) diff --git a/src/components/searchbox/NodeSearchBoxPopover.vue b/src/components/searchbox/NodeSearchBoxPopover.vue index 987a15913..6d48e55e4 100644 --- a/src/components/searchbox/NodeSearchBoxPopover.vue +++ b/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([]) -const addFilter = (filter: FilterAndValue) => { +const nodeFilters = ref[]>([]) +const addFilter = (filter: FuseFilterWithValue) => { nodeFilters.value.push(filter) } -const removeFilter = (filter: FilterAndValue) => { +const removeFilter = ( + filter: FuseFilterWithValue +) => { 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 diff --git a/src/components/searchbox/NodeSearchFilter.vue b/src/components/searchbox/NodeSearchFilter.vue index 8ec4057a7..07ada2173 100644 --- a/src/components/searchbox/NodeSearchFilter.vue +++ b/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() +const selectedFilter = ref>() const filterValues = computed(() => selectedFilter.value?.fuseSearch.data ?? []) const selectedFilterValue = ref('') @@ -43,7 +43,10 @@ onMounted(() => { }) const emit = defineEmits<{ - (event: 'addFilter', filterAndValue: FilterAndValue): void + ( + event: 'addFilter', + filterAndValue: FuseFilterWithValue + ): 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 + }) } diff --git a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue b/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue index b988c4734..6fe36f7ff 100644 --- a/src/components/sidebar/tabs/NodeLibrarySidebarTab.vue +++ b/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(() => { } return buildNodeDefTree(filteredNodeDefs.value) }) -const filters: Ref }>> = - ref([]) +const filters: Ref< + (SearchFilter & { filter: FuseFilterWithValue })[] +> = 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) + 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 +) => { 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() }) diff --git a/src/services/nodeSearchService.ts b/src/services/nodeSearchService.ts index 03e2fd38c..aa242666b 100644 --- a/src/services/nodeSearchService.ts +++ b/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 = [NodeFilter, T] - -export class NodeFilter { - public readonly fuseSearch: FuseSearch - - 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 - ) { - 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 - public readonly nodeFilters: NodeFilter[] + public readonly inputTypeFilter: FuseFilter + public readonly outputTypeFilter: FuseFilter + public readonly nodeCategoryFilter: FuseFilter + public readonly nodeSourceFilter: FuseFilter 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( - /* id */ 'input', - /* name */ 'Input Type', - /* invokeSequence */ 'i', - /* longInvokeSequence */ 'input', - (node) => Object.values(node.inputs).map((input) => input.type), - data, - filterSearchOptions - ) + this.inputTypeFilter = new FuseFilter(data, { + id: 'input', + name: 'Input Type', + invokeSequence: 'i', + getItemOptions: (node) => + Object.values(node.inputs).map((input) => input.type), + fuseOptions + }) - const outputTypeFilter = new NodeFilter( - /* id */ 'output', - /* name */ 'Output Type', - /* invokeSequence */ 'o', - /* longInvokeSequence */ 'output', - (node) => node.outputs.map((output) => output.type), - data, - filterSearchOptions - ) + this.outputTypeFilter = new FuseFilter(data, { + id: 'output', + name: 'Output Type', + invokeSequence: 'o', + getItemOptions: (node) => node.outputs.map((output) => output.type), + fuseOptions + }) - const nodeCategoryFilter = new NodeFilter( - /* id */ 'category', - /* name */ 'Category', - /* invokeSequence */ 'c', - /* longInvokeSequence */ 'category', - (node) => [node.category], - data, - filterSearchOptions - ) + this.nodeCategoryFilter = new FuseFilter(data, { + id: 'category', + name: 'Category', + invokeSequence: 'c', + getItemOptions: (node) => [node.category], + fuseOptions + }) - const nodeSourceFilter = new NodeFilter( - /* 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(data, { + id: 'source', + name: 'Source', + invokeSequence: 's', + getItemOptions: (node) => [node.nodeSource.displayText], + fuseOptions + }) } public searchNode( query: string, - filters: FilterAndValue[] = [], + filters: FuseFilterWithValue[] = [], 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 | undefined { - return this.nodeFilters.find((filter) => filter.id === id) + get nodeFilters(): FuseFilter[] { + return [ + this.inputTypeFilter, + this.outputTypeFilter, + this.nodeCategoryFilter, + this.nodeSourceFilter + ] } } diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index 4db86b576..fd8c6c559 100644 --- a/src/stores/nodeDefStore.ts +++ b/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 diff --git a/src/utils/fuseUtil.ts b/src/utils/fuseUtil.ts index 94d5888d9..313382d0e 100644 --- a/src/utils/fuseUtil.ts +++ b/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 { + filterDef: FuseFilter + value: O +} + +export class FuseFilter { + public readonly fuseSearch: FuseSearch + /** 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 + } + ) { + 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() + 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 { 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 } diff --git a/src/views/GraphView.vue b/src/views/GraphView.vue index cf08c2bee..fd52e79dd 100644 --- a/src/views/GraphView.vue +++ b/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 } ) diff --git a/tests-ui/tests/nodeSearchService.test.ts b/tests-ui/tests/nodeSearchService.test.ts index d0c5d7cc2..7f6c1f377 100644 --- a/tests-ui/tests/nodeSearchService.test.ts +++ b/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) }) })