mirror of https://github.com/langgenius/dify.git
70 lines
1.7 KiB
TypeScript
70 lines
1.7 KiB
TypeScript
import type {
|
|
Dispatch,
|
|
SetStateAction,
|
|
} from 'react'
|
|
import {
|
|
useCallback,
|
|
useMemo,
|
|
useState,
|
|
} from 'react'
|
|
import produce from 'immer'
|
|
import type { Edge } from 'reactflow'
|
|
import type {
|
|
BlockEnum,
|
|
Node,
|
|
} from './types'
|
|
import { NodeInitialData } from './constants'
|
|
|
|
export const useWorkflow = (
|
|
nodes: Node[],
|
|
edges: Edge[],
|
|
setNodes: Dispatch<SetStateAction<Node[]>>,
|
|
setEdges: Dispatch<SetStateAction<Edge[]>>,
|
|
initialSelectedNodeId?: string,
|
|
) => {
|
|
const [selectedNodeId, setSelectedNodeId] = useState(initialSelectedNodeId)
|
|
|
|
const handleSelectedNodeIdChange = useCallback((nodeId: string) => setSelectedNodeId(nodeId), [])
|
|
|
|
const selectedNode = useMemo(() => {
|
|
return nodes.find(node => node.id === selectedNodeId)
|
|
}, [nodes, selectedNodeId])
|
|
|
|
const handleAddNextNode = useCallback((prevNode: Node, nextNodeType: BlockEnum) => {
|
|
const prevNodeDom = document.querySelector(`.react-flow__node-custom[data-id="${prevNode.id}"]`)
|
|
const prevNodeDomHeight = prevNodeDom?.getBoundingClientRect().height || 0
|
|
|
|
const nextNode = {
|
|
id: `node-${Date.now()}`,
|
|
type: 'custom',
|
|
position: {
|
|
x: prevNode.position.x,
|
|
y: prevNode.position.y + prevNodeDomHeight + 64,
|
|
},
|
|
data: NodeInitialData[nextNodeType],
|
|
}
|
|
const newEdge = {
|
|
id: `edge-${Date.now()}`,
|
|
source: prevNode.id,
|
|
target: nextNode.id,
|
|
}
|
|
setNodes((oldNodes) => {
|
|
return produce(oldNodes, (draft) => {
|
|
draft.push(nextNode)
|
|
})
|
|
})
|
|
setEdges((oldEdges) => {
|
|
return produce(oldEdges, (draft) => {
|
|
draft.push(newEdge)
|
|
})
|
|
})
|
|
}, [setNodes, setEdges])
|
|
|
|
return {
|
|
selectedNodeId,
|
|
selectedNode,
|
|
handleSelectedNodeIdChange,
|
|
handleAddNextNode,
|
|
}
|
|
}
|