insert node

langgenius · Mar 17, 2024 · 722ff77 · 722ff77
1 parent 843db3d
commit 722ff77
Show file tree

Hide file tree

Showing 7 changed files with 50 additions and 13 deletions.
diff --git a/web/app/components/workflow/custom-edge.tsx b/web/app/components/workflow/custom-edge.tsx
@@ -3,16 +3,23 @@ import {
   useCallback,
   useState,
 } from 'react'
+import { union } from 'lodash-es'
 import type { EdgeProps } from 'reactflow'
 import {
   BaseEdge,
   EdgeLabelRenderer,
   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)}
           />
         </div>
       </EdgeLabelRenderer>

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -40,9 +40,11 @@ export type CommonNodeType<T = {}> = {
 } & T & Partial<Pick<ToolDefaultValue, 'provider_id' | 'provider_type' | 'provider_name' | 'tool_name'>>
 
 export type CommonEdgeType = {
-  _hovering: boolean
-  _connectedNodeIsHovering: boolean
+  _hovering?: boolean
+  _connectedNodeIsHovering?: boolean
   _runned?: boolean
+  sourceType: BlockEnum
+  targetType: BlockEnum
 }
 
 export type Node<T = {}> = ReactFlowNode<CommonNodeType<T>>