mirror of
https://github.com/langgenius/dify.git
synced 2025-07-14 20:49:56 +00:00
116 lines
3.5 KiB
TypeScript
116 lines
3.5 KiB
TypeScript
import { useCallback } from 'react'
|
|
import { getOutgoers, useStoreApi } from 'reactflow'
|
|
import { BlockEnum, type Node, type ValueSelector } from '../../workflow/types'
|
|
import { uniqBy } from 'lodash-es'
|
|
import { findUsedVarNodes, updateNodeVars } from '../../workflow/nodes/_base/components/variable/utils'
|
|
import type { DataSourceNodeType } from '../../workflow/nodes/data-source/types'
|
|
|
|
export const usePipeline = () => {
|
|
const store = useStoreApi()
|
|
|
|
const getAllDatasourceNodes = useCallback(() => {
|
|
const {
|
|
getNodes,
|
|
} = store.getState()
|
|
const nodes = getNodes() as Node<DataSourceNodeType>[]
|
|
const datasourceNodes = nodes.filter(node => node.data.type === BlockEnum.DataSource)
|
|
|
|
return datasourceNodes
|
|
}, [store])
|
|
|
|
const getAllNodesInSameBranch = useCallback((nodeId: string) => {
|
|
const {
|
|
getNodes,
|
|
edges,
|
|
} = store.getState()
|
|
const nodes = getNodes()
|
|
const list: Node[] = []
|
|
|
|
const traverse = (root: Node, callback: (node: Node) => void) => {
|
|
if (root) {
|
|
const outgoers = getOutgoers(root, nodes, edges)
|
|
|
|
if (outgoers.length) {
|
|
outgoers.forEach((node) => {
|
|
callback(node)
|
|
traverse(node, callback)
|
|
})
|
|
}
|
|
}
|
|
}
|
|
|
|
if (nodeId === 'shared') {
|
|
const allDatasourceNodes = getAllDatasourceNodes()
|
|
|
|
if (allDatasourceNodes.length === 0)
|
|
return []
|
|
|
|
list.push(...allDatasourceNodes)
|
|
|
|
allDatasourceNodes.forEach((node) => {
|
|
traverse(node, (childNode) => {
|
|
list.push(childNode)
|
|
})
|
|
})
|
|
}
|
|
else {
|
|
const currentNode = nodes.find(node => node.id === nodeId)!
|
|
|
|
if (!currentNode)
|
|
return []
|
|
|
|
list.push(currentNode)
|
|
|
|
traverse(currentNode, (node) => {
|
|
list.push(node)
|
|
})
|
|
}
|
|
|
|
return uniqBy(list, 'id')
|
|
}, [getAllDatasourceNodes, store])
|
|
|
|
const isVarUsedInNodes = useCallback((varSelector: ValueSelector) => {
|
|
const nodeId = varSelector[1] // Assuming the first element is always 'VARIABLE_PREFIX'(rag)
|
|
const afterNodes = getAllNodesInSameBranch(nodeId)
|
|
const effectNodes = findUsedVarNodes(varSelector, afterNodes)
|
|
return effectNodes.length > 0
|
|
}, [getAllNodesInSameBranch])
|
|
|
|
const handleInputVarRename = useCallback((nodeId: string, oldValeSelector: ValueSelector, newVarSelector: ValueSelector) => {
|
|
const { getNodes, setNodes } = store.getState()
|
|
const afterNodes = getAllNodesInSameBranch(nodeId)
|
|
const effectNodes = findUsedVarNodes(oldValeSelector, afterNodes)
|
|
if (effectNodes.length > 0) {
|
|
const newNodes = getNodes().map((node) => {
|
|
if (effectNodes.find(n => n.id === node.id))
|
|
return updateNodeVars(node, oldValeSelector, newVarSelector)
|
|
|
|
return node
|
|
})
|
|
setNodes(newNodes)
|
|
}
|
|
}, [getAllNodesInSameBranch, store])
|
|
|
|
const removeUsedVarInNodes = useCallback((varSelector: ValueSelector) => {
|
|
const nodeId = varSelector[1] // Assuming the first element is always 'VARIABLE_PREFIX'(rag)
|
|
const { getNodes, setNodes } = store.getState()
|
|
const afterNodes = getAllNodesInSameBranch(nodeId)
|
|
const effectNodes = findUsedVarNodes(varSelector, afterNodes)
|
|
if (effectNodes.length > 0) {
|
|
const newNodes = getNodes().map((node) => {
|
|
if (effectNodes.find(n => n.id === node.id))
|
|
return updateNodeVars(node, varSelector, [])
|
|
|
|
return node
|
|
})
|
|
setNodes(newNodes)
|
|
}
|
|
}, [getAllNodesInSameBranch, store])
|
|
|
|
return {
|
|
handleInputVarRename,
|
|
isVarUsedInNodes,
|
|
removeUsedVarInNodes,
|
|
}
|
|
}
|