From 722ff7795d8d09aecba5788f98af0f9a76b69929 Mon Sep 17 00:00:00 2001 From: StyleZhang Date: Sun, 17 Mar 2024 20:19:58 +0800 Subject: [PATCH] insert node --- web/app/components/workflow/custom-edge.tsx | 15 +++++++++-- .../workflow/hooks/use-edges-interactions.ts | 4 +-- .../workflow/hooks/use-nodes-interactions.ts | 27 ++++++++++++++++--- .../workflow/hooks/use-nodes-sync-draft.ts | 5 +++- .../workflow/hooks/use-workflow-run.ts | 4 +-- .../components/var-list/use-var-list.ts | 2 +- web/app/components/workflow/types.ts | 6 +++-- 7 files changed, 50 insertions(+), 13 deletions(-) diff --git a/web/app/components/workflow/custom-edge.tsx b/web/app/components/workflow/custom-edge.tsx index c14c8e52c3..919c69d5da 100644 --- a/web/app/components/workflow/custom-edge.tsx +++ b/web/app/components/workflow/custom-edge.tsx @@ -3,6 +3,7 @@ import { useCallback, useState, } from 'react' +import { union } from 'lodash-es' import type { EdgeProps } from 'reactflow' import { BaseEdge, @@ -10,9 +11,15 @@ import { Position, getSimpleBezierPath, } from 'reactflow' -import { useNodesInteractions } from './hooks' +import { + useNodesExtraData, + useNodesInteractions, +} from './hooks' import BlockSelector from './block-selector' -import type { OnSelectBlock } from './types' +import type { + Edge, + OnSelectBlock, +} from './types' const CustomEdge = ({ id, @@ -41,6 +48,9 @@ const CustomEdge = ({ }) const [open, setOpen] = useState(false) const { handleNodeAdd } = useNodesInteractions() + const nodesExtraData = useNodesExtraData() + const availablePrevNodes = nodesExtraData[(data as Edge['data'])!.targetType].availablePrevNodes + const availableNextNodes = nodesExtraData[(data as Edge['data'])!.sourceType].availableNextNodes const handleOpenChange = useCallback((v: boolean) => { setOpen(v) }, []) @@ -88,6 +98,7 @@ const CustomEdge = ({ onOpenChange={handleOpenChange} asChild onSelect={handleInsert} + availableBlocksTypes={union(availablePrevNodes, availableNextNodes)} /> diff --git a/web/app/components/workflow/hooks/use-edges-interactions.ts b/web/app/components/workflow/hooks/use-edges-interactions.ts index 01b96a3c42..56d67c336e 100644 --- a/web/app/components/workflow/hooks/use-edges-interactions.ts +++ b/web/app/components/workflow/hooks/use-edges-interactions.ts @@ -34,7 +34,7 @@ export const useEdgesInteractions = () => { const newEdges = produce(edges, (draft) => { const currentEdge = draft.find(e => e.id === edge.id)! - currentEdge.data = { ...currentEdge.data, _hovering: true } + currentEdge.data._hovering = true }) setEdges(newEdges) }, [store, workflowStore]) @@ -52,7 +52,7 @@ export const useEdgesInteractions = () => { const newEdges = produce(edges, (draft) => { const currentEdge = draft.find(e => e.id === edge.id)! - currentEdge.data = { ...currentEdge.data, _hovering: false } + currentEdge.data._hovering = false }) setEdges(newEdges) }, [store, workflowStore]) diff --git a/web/app/components/workflow/hooks/use-nodes-interactions.ts b/web/app/components/workflow/hooks/use-nodes-interactions.ts index 589ed0dae6..9f182b404b 100644 --- a/web/app/components/workflow/hooks/use-nodes-interactions.ts +++ b/web/app/components/workflow/hooks/use-nodes-interactions.ts @@ -212,7 +212,7 @@ export const useNodesInteractions = () => { connectedEdges.forEach((edge) => { const currentEdge = draft.find(e => e.id === edge.id) if (currentEdge) - currentEdge.data = { ...currentEdge.data, _connectedNodeIsHovering: true } + currentEdge.data._connectedNodeIsHovering = true }) }) setEdges(newEdges) @@ -238,7 +238,7 @@ export const useNodesInteractions = () => { setNodes(newNodes) const newEdges = produce(edges, (draft) => { draft.forEach((edge) => { - edge.data = { ...edge.data, _connectedNodeIsHovering: false } + edge.data._connectedNodeIsHovering = false }) }) setEdges(newEdges) @@ -301,6 +301,7 @@ export const useNodesInteractions = () => { edges, setEdges, } = store.getState() + const nodes = getNodes() const needDeleteEdges = edges.filter(edge => (edge.source === source && edge.sourceHandle === sourceHandle) || (edge.target === target && edge.targetHandle === targetHandle)) const needDeleteEdgesIds = needDeleteEdges.map(edge => edge.id) const newEdge = { @@ -310,8 +311,11 @@ export const useNodesInteractions = () => { target: target!, sourceHandle, targetHandle, + data: { + sourceType: nodes.find(node => node.id === source)!.data.type, + targetType: nodes.find(node => node.id === target)!.data.type, + }, } - const nodes = getNodes() const nodesConnectedSourceOrTargetHandleIdsMap = getNodesConnectedSourceOrTargetHandleIdsMap( [ ...needDeleteEdges.map(edge => ({ type: 'remove', edge })), @@ -451,6 +455,10 @@ export const useNodesInteractions = () => { sourceHandle: prevNodeSourceHandle, target: newNode.id, targetHandle, + data: { + sourceType: prevNode.data.type, + targetType: newNode.data.type, + }, } const newNodes = produce(nodes, (draft: Node[]) => { draft.forEach((node) => { @@ -484,6 +492,10 @@ export const useNodesInteractions = () => { sourceHandle, target: nextNodeId, targetHandle: nextNodeTargetHandle, + data: { + sourceType: newNode.data.type, + targetType: nextNode.data.type, + }, } const afterNodesInSameBranch = getAfterNodesInSameBranch(nextNodeId!) const afterNodesInSameBranchIds = afterNodesInSameBranch.map(node => node.id) @@ -506,6 +518,7 @@ export const useNodesInteractions = () => { setEdges(newEdges) } if (prevNodeId && nextNodeId) { + const prevNode = nodes.find(node => node.id === prevNodeId)! const nextNode = nodes.find(node => node.id === nextNodeId)! newNode.data._connectedTargetHandleIds = [targetHandle] newNode.data._connectedSourceHandleIds = [sourceHandle] @@ -522,6 +535,10 @@ export const useNodesInteractions = () => { sourceHandle: prevNodeSourceHandle, target: newNode.id, targetHandle, + data: { + sourceType: prevNode.data.type, + targetType: newNode.data.type, + }, } const newNextEdge = { id: `${newNode.id}-${nextNodeId}`, @@ -530,6 +547,10 @@ export const useNodesInteractions = () => { sourceHandle, target: nextNodeId, targetHandle: nextNodeTargetHandle, + data: { + sourceType: newNode.data.type, + targetType: nextNode.data.type, + }, } const nodesConnectedSourceOrTargetHandleIdsMap = getNodesConnectedSourceOrTargetHandleIdsMap( [ diff --git a/web/app/components/workflow/hooks/use-nodes-sync-draft.ts b/web/app/components/workflow/hooks/use-nodes-sync-draft.ts index a695a4598f..717942a029 100644 --- a/web/app/components/workflow/hooks/use-nodes-sync-draft.ts +++ b/web/app/components/workflow/hooks/use-nodes-sync-draft.ts @@ -36,7 +36,10 @@ export const useNodesSyncDraft = () => { }) const producedEdges = produce(edges, (draft) => { draft.forEach((edge) => { - delete edge.data + Object.keys(edge.data).forEach((key) => { + if (key.startsWith('_')) + delete edge.data[key] + }) }) }) syncWorkflowDraft({ diff --git a/web/app/components/workflow/hooks/use-workflow-run.ts b/web/app/components/workflow/hooks/use-workflow-run.ts index e5503c68b5..9cf32696ce 100644 --- a/web/app/components/workflow/hooks/use-workflow-run.ts +++ b/web/app/components/workflow/hooks/use-workflow-run.ts @@ -81,7 +81,7 @@ export const useWorkflowRun = () => { setNodes(newNodes) const newEdges = produce(edges, (draft) => { draft.forEach((edge) => { - edge.data = { ...edge.data, _runned: false } + edge.data._runned = false }) }) setEdges(newEdges) @@ -154,7 +154,7 @@ export const useWorkflowRun = () => { const edge = draft.find(edge => edge.target === data.node_id) if (edge) - edge.data = { ...edge.data, _runned: true } + edge.data._runned = true }) setEdges(newEdges) }, diff --git a/web/app/components/workflow/nodes/variable-assigner/components/var-list/use-var-list.ts b/web/app/components/workflow/nodes/variable-assigner/components/var-list/use-var-list.ts index 1d8cd5ace9..fcfb4b2e49 100644 --- a/web/app/components/workflow/nodes/variable-assigner/components/var-list/use-var-list.ts +++ b/web/app/components/workflow/nodes/variable-assigner/components/var-list/use-var-list.ts @@ -19,8 +19,8 @@ function useVarList({ const newInputs = produce(inputs, (draft) => { draft.variables = newList }) - handleVariableAssignerEdgesChange(id, newList) setInputs(newInputs) + handleVariableAssignerEdgesChange(id, newList) }, [inputs, setInputs, id, handleVariableAssignerEdgesChange]) const handleAddVariable = useCallback(() => { diff --git a/web/app/components/workflow/types.ts b/web/app/components/workflow/types.ts index 024f16a498..20cd9946d0 100644 --- a/web/app/components/workflow/types.ts +++ b/web/app/components/workflow/types.ts @@ -40,9 +40,11 @@ export type CommonNodeType = { } & T & Partial> export type CommonEdgeType = { - _hovering: boolean - _connectedNodeIsHovering: boolean + _hovering?: boolean + _connectedNodeIsHovering?: boolean _runned?: boolean + sourceType: BlockEnum + targetType: BlockEnum } export type Node = ReactFlowNode>