'use client'
import { useCallback, useEffect, useMemo } from 'react'
import { useNodes } from 'reactflow'
import { useNodesInteractions } from '@/app/components/workflow/hooks/use-nodes-interactions'
import type { CommonNodeType } from '@/app/components/workflow/types'
import { ragPipelineNodesAction } from '@/app/components/goto-anything/actions/rag-pipeline-nodes'
import BlockIcon from '@/app/components/workflow/block-icon'
import { setupNodeSelectionListener } from '@/app/components/workflow/utils/node-navigation'
import { BlockEnum } from '@/app/components/workflow/types'
import type { LLMNodeType } from '@/app/components/workflow/nodes/llm/types'
import type { ToolNodeType } from '@/app/components/workflow/nodes/tool/types'
import type { KnowledgeRetrievalNodeType } from '@/app/components/workflow/nodes/knowledge-retrieval/types'
import { useGetToolIcon } from '@/app/components/workflow/hooks/use-tool-icon'
/**
* Hook to register RAG pipeline nodes search functionality
*/
export const useRagPipelineSearch = () => {
const nodes = useNodes()
const { handleNodeSelect } = useNodesInteractions()
const getToolIcon = useGetToolIcon()
// Process nodes to create searchable data structure
const searchableNodes = useMemo(() => {
return nodes.map((node) => {
const nodeData = node.data as CommonNodeType
const title = nodeData.title || nodeData.type || 'Untitled Node'
let desc = nodeData.desc || ''
// Keep the original node title for consistency with workflow display
// Only enhance description for better search context
if (nodeData.type === BlockEnum.Tool) {
const toolData = nodeData as ToolNodeType
desc = toolData.tool_description || toolData.tool_label || desc
}
if (nodeData.type === BlockEnum.LLM) {
const llmData = nodeData as LLMNodeType
if (llmData.model?.provider && llmData.model?.name)
desc = `${llmData.model.name} (${llmData.model.provider}) - ${llmData.model.mode || desc}`
}
if (nodeData.type === BlockEnum.KnowledgeRetrieval) {
const knowledgeData = nodeData as KnowledgeRetrievalNodeType
if (knowledgeData.dataset_ids?.length)
desc = `Knowledge Retrieval with ${knowledgeData.dataset_ids.length} datasets - ${desc}`
}
return {
id: node.id,
title,
desc,
type: nodeData.type,
blockType: nodeData.type,
nodeData,
toolIcon: getToolIcon(nodeData),
modelInfo: nodeData.type === BlockEnum.LLM ? {
provider: (nodeData as LLMNodeType).model?.provider,
name: (nodeData as LLMNodeType).model?.name,
mode: (nodeData as LLMNodeType).model?.mode,
} : {
provider: undefined,
name: undefined,
mode: undefined,
},
}
})
}, [nodes, getToolIcon])
// Calculate relevance score for search results
const calculateScore = useCallback((node: {
title: string;
type: string;
desc: string;
modelInfo: { provider?: string; name?: string; mode?: string }
}, searchTerm: string): number => {
if (!searchTerm) return 1
let score = 0
const term = searchTerm.toLowerCase()
// Title match (highest priority)
if (node.title.toLowerCase().includes(term))
score += 10
// Type match
if (node.type.toLowerCase().includes(term))
score += 8
// Description match
if (node.desc.toLowerCase().includes(term))
score += 5
// Model info matches (for LLM nodes)
if (node.modelInfo.provider?.toLowerCase().includes(term))
score += 6
if (node.modelInfo.name?.toLowerCase().includes(term))
score += 6
if (node.modelInfo.mode?.toLowerCase().includes(term))
score += 4
return score
}, [])
// Create search function for RAG pipeline nodes
const searchRagPipelineNodes = useCallback((query: string) => {
if (!searchableNodes.length) return []
const searchTerm = query.toLowerCase().trim()
const results = searchableNodes
.map((node) => {
const score = calculateScore(node, searchTerm)
return score > 0 ? {
id: node.id,
title: node.title,
description: node.desc || node.type,
type: 'workflow-node' as const,
path: `#${node.id}`,
icon: (
),
metadata: {
nodeId: node.id,
nodeData: node.nodeData,
},
data: node.nodeData,
score,
} : null
})
.filter((node): node is NonNullable => node !== null)
.sort((a, b) => {
// If no search term, sort alphabetically
if (!searchTerm) return a.title.localeCompare(b.title)
// Sort by relevance score (higher score first)
return (b.score || 0) - (a.score || 0)
})
return results
}, [searchableNodes, calculateScore])
// Directly set the search function on the action object
useEffect(() => {
if (searchableNodes.length > 0) {
// Set the search function directly on the action
ragPipelineNodesAction.searchFn = searchRagPipelineNodes
}
return () => {
// Clean up when component unmounts
ragPipelineNodesAction.searchFn = undefined
}
}, [searchableNodes, searchRagPipelineNodes])
// Set up node selection event listener using the utility function
useEffect(() => {
return setupNodeSelectionListener(handleNodeSelect)
}, [handleNodeSelect])
return null
}