import type { Connection } from 'reactflow' import type { IterationNodeType } from '../nodes/iteration/types' import type { LoopNodeType } from '../nodes/loop/types' import type { BlockEnum, Edge, Node, ValueSelector } from '../types' import { uniqBy } from 'es-toolkit/compat' import { useCallback } from 'react' import { getIncomers, getOutgoers } from 'reactflow' import { useStore as useAppStore } from '@/app/components/app/store' import { useCollaborativeWorkflow } from '@/app/components/workflow/hooks/use-collaborative-workflow' import { CUSTOM_ITERATION_START_NODE } from '@/app/components/workflow/nodes/iteration-start/constants' import { CUSTOM_LOOP_START_NODE } from '@/app/components/workflow/nodes/loop-start/constants' import { AppModeEnum } from '@/types/app' import { useNodesMetaData } from '.' import { SUPPORT_OUTPUT_VARS_NODE } from '../constants' import { useHooksStore } from '../hooks-store' import { findUsedVarNodes, getNodeOutputVars, updateNodeVars, } from '../nodes/_base/components/variable/utils' import { CUSTOM_NOTE_NODE } from '../note-node/constants' import { useStore, useWorkflowStore } from '../store' import { WorkflowRunningStatus } from '../types' import { getNodeCatalogType } from '../utils' import { getWorkflowEntryNode, isWorkflowEntryNode } from '../utils/workflow-entry' import { useAvailableBlocks } from './use-available-blocks' export const useIsChatMode = () => { const appDetail = useAppStore((s) => s.appDetail) return appDetail?.mode === AppModeEnum.ADVANCED_CHAT } export const useWorkflow = () => { const collaborativeWorkflow = useCollaborativeWorkflow() const { getAvailableBlocks } = useAvailableBlocks() const { nodesMap } = useNodesMetaData() const getNodeById = useCallback( (nodeId: string) => { const { nodes } = collaborativeWorkflow.getState() const currentNode = nodes.find((node) => node.id === nodeId) return currentNode }, [collaborativeWorkflow], ) const getTreeLeafNodes = useCallback( (nodeId: string) => { const { nodes, edges } = collaborativeWorkflow.getState() const currentNode = nodes.find((node) => node.id === nodeId) let startNodes = nodes.filter((node) => nodesMap?.[node.data.type as BlockEnum]?.metaData.isStart) || [] if (currentNode?.parentId) { const startNode = nodes.find( (node) => node.parentId === currentNode.parentId && (node.type === CUSTOM_ITERATION_START_NODE || node.type === CUSTOM_LOOP_START_NODE), ) if (startNode) startNodes = [startNode] } if (!startNodes.length) return [] const list: Node[] = [] const preOrder = (root: Node, callback: (node: Node) => void) => { if (root.id === nodeId) return const outgoers = getOutgoers(root, nodes, edges) if (outgoers.length) { outgoers.forEach((outgoer) => { preOrder(outgoer, callback) }) } else { if (root.id !== nodeId) callback(root) } } startNodes.forEach((startNode) => { preOrder(startNode, (node) => { list.push(node) }) }) const incomers = getIncomers({ id: nodeId } as Node, nodes, edges) list.push(...incomers) return uniqBy(list, 'id').filter((item: Node) => { return SUPPORT_OUTPUT_VARS_NODE.includes(item.data.type) }) }, [collaborativeWorkflow, nodesMap], ) const getBeforeNodesInSameBranch = useCallback( (nodeId: string, newNodes?: Node[], newEdges?: Edge[]) => { const { nodes: oldNodes, edges } = collaborativeWorkflow.getState() const nodes = newNodes || oldNodes const currentNode = nodes.find((node) => node.id === nodeId) const list: Node[] = [] if (!currentNode) return list if (currentNode.parentId) { const parentNode = nodes.find((node) => node.id === currentNode.parentId) if (parentNode) { const parentList = getBeforeNodesInSameBranch(parentNode.id) list.push(...parentList) } } const traverse = (root: Node, callback: (node: Node) => void) => { if (root) { const incomers = getIncomers(root, nodes, newEdges || edges) if (incomers.length) { incomers.forEach((node) => { if (!list.find((n) => node.id === n.id)) { callback(node) traverse(node, callback) } }) } } } traverse(currentNode, (node) => { list.push(node) }) const length = list.length if (length) { return uniqBy(list, 'id') .reverse() .filter((item: Node) => { return SUPPORT_OUTPUT_VARS_NODE.includes(item.data.type) }) } return [] }, [collaborativeWorkflow], ) const getBeforeNodesInSameBranchIncludeParent = useCallback( (nodeId: string, newNodes?: Node[], newEdges?: Edge[]) => { const nodes = getBeforeNodesInSameBranch(nodeId, newNodes, newEdges) const { nodes: allNodes } = collaborativeWorkflow.getState() const node = allNodes.find((n) => n.id === nodeId) const parentNodeId = node?.parentId const parentNode = allNodes.find((n) => n.id === parentNodeId) if (parentNode) nodes.push(parentNode) return nodes }, [getBeforeNodesInSameBranch, collaborativeWorkflow], ) const getAfterNodesInSameBranch = useCallback( (nodeId: string) => { const { nodes, edges } = collaborativeWorkflow.getState() const currentNode = nodes.find((node) => node.id === nodeId)! if (!currentNode) return [] const list: Node[] = [currentNode] 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) }) } } } traverse(currentNode, (node) => { list.push(node) }) return uniqBy(list, 'id') }, [collaborativeWorkflow], ) const getBeforeNodeById = useCallback( (nodeId: string) => { const { nodes, edges } = collaborativeWorkflow.getState() const node = nodes.find((node) => node.id === nodeId)! return getIncomers(node, nodes, edges) }, [collaborativeWorkflow], ) const getIterationNodeChildren = useCallback( (nodeId: string) => { const { nodes } = collaborativeWorkflow.getState() return nodes.filter((node) => node.parentId === nodeId) }, [collaborativeWorkflow], ) const getLoopNodeChildren = useCallback( (nodeId: string) => { const { nodes } = collaborativeWorkflow.getState() return nodes.filter((node) => node.parentId === nodeId) }, [collaborativeWorkflow], ) const isFromStartNode = useCallback( (nodeId: string) => { const { nodes } = collaborativeWorkflow.getState() const currentNode = nodes.find((node) => node.id === nodeId) if (!currentNode) return false if (isWorkflowEntryNode(currentNode.data.type)) return true const checkPreviousNodes = (node: Node) => { const previousNodes = getBeforeNodeById(node.id) for (const prevNode of previousNodes) { if (isWorkflowEntryNode(prevNode.data.type)) return true if (checkPreviousNodes(prevNode)) return true } return false } return checkPreviousNodes(currentNode) }, [collaborativeWorkflow, getBeforeNodeById], ) const handleOutVarRenameChange = useCallback( (nodeId: string, oldValeSelector: ValueSelector, newVarSelector: ValueSelector) => { const { nodes: allNodes, setNodes } = collaborativeWorkflow.getState() const affectedNodes = findUsedVarNodes(oldValeSelector, allNodes) if (affectedNodes.length > 0) { const newNodes = allNodes.map((node) => { if (affectedNodes.find((n) => n.id === node.id)) return updateNodeVars(node, oldValeSelector, newVarSelector) return node }) setNodes(newNodes) } }, [collaborativeWorkflow], ) const isVarUsedInNodes = useCallback( (varSelector: ValueSelector) => { const nodeId = varSelector[0] const afterNodes = getAfterNodesInSameBranch(nodeId!) const effectNodes = findUsedVarNodes(varSelector, afterNodes) return effectNodes.length > 0 }, [getAfterNodesInSameBranch], ) const removeUsedVarInNodes = useCallback( (varSelector: ValueSelector) => { const nodeId = varSelector[0] const { nodes, setNodes } = collaborativeWorkflow.getState() const afterNodes = getAfterNodesInSameBranch(nodeId!) const effectNodes = findUsedVarNodes(varSelector, afterNodes) if (effectNodes.length > 0) { const newNodes = nodes.map((node) => { if (effectNodes.find((n) => n.id === node.id)) return updateNodeVars(node, varSelector, []) return node }) setNodes(newNodes) } }, [getAfterNodesInSameBranch, collaborativeWorkflow], ) const isNodeVarsUsedInNodes = useCallback( (node: Node, isChatMode: boolean) => { const outputVars = getNodeOutputVars(node, isChatMode) const isUsed = outputVars.some((varSelector) => { return isVarUsedInNodes(varSelector) }) return isUsed }, [isVarUsedInNodes], ) const getRootNodesById = useCallback( (nodeId: string) => { const { nodes, edges } = collaborativeWorkflow.getState() const currentNode = nodes.find((node) => node.id === nodeId) const rootNodes: Node[] = [] if (!currentNode) return rootNodes if (currentNode.parentId) { const parentNode = nodes.find((node) => node.id === currentNode.parentId) if (parentNode) { const parentList = getRootNodesById(parentNode.id) rootNodes.push(...parentList) } } const traverse = (root: Node, callback: (node: Node) => void) => { if (root) { const incomers = getIncomers(root, nodes, edges) if (incomers.length) { incomers.forEach((node) => { traverse(node, callback) }) } else { callback(root) } } } traverse(currentNode, (node) => { rootNodes.push(node) }) const length = rootNodes.length if (length) return uniqBy(rootNodes, 'id') return [] }, [collaborativeWorkflow], ) const getStartNodes = useCallback( (nodes: Node[], currentNode?: Node) => { const { id, parentId } = currentNode || {} let startNodes: Node[] = [] if (parentId) { const parentNode = nodes.find((node) => node.id === parentId) if (!parentNode) throw new Error('Parent node not found') const startNode = nodes.find( (node) => node.id === (parentNode.data as IterationNodeType | LoopNodeType).start_node_id, ) if (startNode) startNodes = [startNode] } else { startNodes = nodes.filter((node) => nodesMap?.[node.data.type as BlockEnum]?.metaData.isStart) || [] } if (!startNodes.length) startNodes = getRootNodesById(id || '') return startNodes }, [nodesMap, getRootNodesById], ) const isValidConnection = useCallback( ({ source, sourceHandle: _sourceHandle, target }: Connection) => { const { nodes, edges } = collaborativeWorkflow.getState() const sourceNode: Node = nodes.find((node) => node.id === source)! const targetNode: Node = nodes.find((node) => node.id === target)! if (sourceNode.type === CUSTOM_NOTE_NODE || targetNode.type === CUSTOM_NOTE_NODE) return false if (sourceNode.parentId !== targetNode.parentId) return false if (sourceNode && targetNode) { const sourceNodeCatalogType = getNodeCatalogType(sourceNode.data) const targetNodeCatalogType = getNodeCatalogType(targetNode.data) const sourceNodeAvailableNextNodes = getAvailableBlocks( sourceNodeCatalogType, !!sourceNode.parentId, ).availableNextBlocks const targetNodeAvailablePrevNodes = getAvailableBlocks( targetNodeCatalogType, !!targetNode.parentId, ).availablePrevBlocks if (!sourceNodeAvailableNextNodes.includes(targetNodeCatalogType)) return false if (!targetNodeAvailablePrevNodes.includes(sourceNodeCatalogType)) return false } const hasCycle = (node: Node, visited = new Set()) => { if (visited.has(node.id)) return false visited.add(node.id) for (const outgoer of getOutgoers(node, nodes, edges)) { if (outgoer.id === source) return true if (hasCycle(outgoer, visited)) return true } } return !hasCycle(targetNode) }, [collaborativeWorkflow, getAvailableBlocks], ) const getNode = useCallback( (nodeId?: string) => { const { nodes } = collaborativeWorkflow.getState() return nodes.find((node) => node.id === nodeId) || getWorkflowEntryNode(nodes) }, [collaborativeWorkflow], ) return { getNodeById, getTreeLeafNodes, getBeforeNodesInSameBranch, getBeforeNodesInSameBranchIncludeParent, getAfterNodesInSameBranch, handleOutVarRenameChange, isVarUsedInNodes, removeUsedVarInNodes, isNodeVarsUsedInNodes, isValidConnection, getBeforeNodeById, getIterationNodeChildren, getLoopNodeChildren, getRootNodesById, getStartNodes, isFromStartNode, getNode, } } export const useWorkflowReadOnly = () => { const workflowStore = useWorkflowStore() const workflowRunningData = useStore((s) => s.workflowRunningData) const canvasReadOnly = useStore((s) => s.canvasReadOnly) const getWorkflowReadOnly = useCallback(() => { const { canvasReadOnly, workflowRunningData } = workflowStore.getState() return canvasReadOnly || workflowRunningData?.result.status === WorkflowRunningStatus.Running }, [workflowStore]) return { workflowReadOnly: canvasReadOnly || workflowRunningData?.result.status === WorkflowRunningStatus.Running, getWorkflowReadOnly, } } const useNodesReadOnlyBase = (canEdit: boolean) => { const workflowStore = useWorkflowStore() const canvasReadOnly = useStore((s) => s.canvasReadOnly) const workflowRunningData = useStore((s) => s.workflowRunningData) const historyWorkflowData = useStore((s) => s.historyWorkflowData) const isRestoring = useStore((s) => s.isRestoring) const getNodesReadOnly = useCallback((): boolean => { const { workflowRunningData, historyWorkflowData, isRestoring, canvasReadOnly } = workflowStore.getState() return !!( canvasReadOnly || !canEdit || workflowRunningData?.result.status === WorkflowRunningStatus.Running || workflowRunningData?.result.status === WorkflowRunningStatus.Paused || historyWorkflowData || isRestoring ) }, [workflowStore, canEdit]) return { nodesReadOnly: !!( canvasReadOnly || !canEdit || workflowRunningData?.result.status === WorkflowRunningStatus.Running || workflowRunningData?.result.status === WorkflowRunningStatus.Paused || historyWorkflowData || isRestoring ), getNodesReadOnly, } } export const useNodesReadOnlyByCanEdit = (canEdit: boolean) => { return useNodesReadOnlyBase(canEdit) } export const useNodesReadOnly = () => { const canEdit = useHooksStore((s) => s.accessControl.canEdit) return useNodesReadOnlyBase(canEdit) } export const useIsNodeInIteration = (iterationId: string) => { const collaborativeWorkflow = useCollaborativeWorkflow() const isNodeInIteration = useCallback( (nodeId: string) => { const { nodes } = collaborativeWorkflow.getState() const node = nodes.find((node) => node.id === nodeId) if (!node) return false if (node.parentId === iterationId) return true return false }, [iterationId, collaborativeWorkflow], ) return { isNodeInIteration, } } export const useIsNodeInLoop = (loopId: string) => { const collaborativeWorkflow = useCollaborativeWorkflow() const isNodeInLoop = useCallback( (nodeId: string) => { const { nodes } = collaborativeWorkflow.getState() const node = nodes.find((node) => node.id === nodeId) if (!node) return false if (node.parentId === loopId) return true return false }, [loopId, collaborativeWorkflow], ) return { isNodeInLoop, } }