feat: implement usePipeline hook for managing pipeline variables and refactor input field handling

web/app/components/rag-pipeline/components/input-field/field-list/hooks.ts

@@ -8,9 +8,11 @@ import type { InputVar } from '@/models/pipeline'
 import type { SortableItem } from './types'
 import type { MoreInfo, ValueSelector } from '@/app/components/workflow/types'
 import { ChangeType } from '@/app/components/workflow/types'
-import { useWorkflow } from '@/app/components/workflow/hooks'
 import { useBoolean } from 'ahooks'
 import Toast from '@/app/components/base/toast'
+import { usePipeline } from '../../../hooks/use-pipeline'
+
+const VARIABLE_PREFIX = 'rag'
 
 export const useFieldList = (
   initialInputFields: InputVar[],
@@ -22,7 +24,7 @@ export const useFieldList = (
   const [removedVar, setRemovedVar] = useState<ValueSelector>([])
   const [removedIndex, setRemoveIndex] = useState(0)
 
-  const { handleOutVarRenameChange, isVarUsedInNodes, removeUsedVarInNodes } = useWorkflow()
+  const { handleInputVarRename, isVarUsedInNodes, removeUsedVarInNodes } = usePipeline()
 
   const [isShowRemoveVarConfirm, {
     setTrue: showRemoveVarConfirm,
@@ -61,9 +63,9 @@ export const useFieldList = (
   const handleRemoveField = useCallback((index: number) => {
     const itemToRemove = inputFieldsRef.current[index]
     // Check if the variable is used in other nodes
-    if (isVarUsedInNodes([nodeId, itemToRemove.variable || ''])) {
+    if (isVarUsedInNodes([VARIABLE_PREFIX, nodeId, itemToRemove.variable || ''])) {
       showRemoveVarConfirm()
-      setRemovedVar([nodeId, itemToRemove.variable || ''])
+      setRemovedVar([VARIABLE_PREFIX, nodeId, itemToRemove.variable || ''])
       setRemoveIndex(index as number)
       return
     }
@@ -99,9 +101,9 @@ export const useFieldList = (
     handleInputFieldsChange(newInputFields)
     // Update variable name in nodes if it has changed
     if (moreInfo?.type === ChangeType.changeVarName)
-      handleOutVarRenameChange(nodeId, [nodeId, moreInfo.payload?.beforeKey || ''], [nodeId, moreInfo.payload?.afterKey || ''])
+      handleInputVarRename(nodeId, [VARIABLE_PREFIX, nodeId, moreInfo.payload?.beforeKey || ''], [VARIABLE_PREFIX, nodeId, moreInfo.payload?.afterKey || ''])
     handleCloseInputFieldEditor()
-  }, [editingField?.variable, handleCloseInputFieldEditor, handleInputFieldsChange, handleOutVarRenameChange, nodeId])
+  }, [editingField?.variable, handleCloseInputFieldEditor, handleInputFieldsChange, handleInputVarRename, nodeId])
 
   return {
     inputFields,

web/app/components/rag-pipeline/hooks/use-pipeline.tsx

@@ -0,0 +1,115 @@
+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,
+  }
+}

web/app/components/workflow/nodes/_base/components/variable/utils.ts

@@ -838,9 +838,9 @@ export const getVarType = ({
   })
 
   const targetVarNodeId = (() => {
-    if(isSystem)
+    if (isSystem)
       return startNode?.id
-    if(isInNodeRagVariable)
+    if (isInNodeRagVariable)
       return valueSelector[1]
     return valueSelector[0]
   })()
@@ -857,14 +857,14 @@ export const getVarType = ({
   }
   else {
     const targetVar = curr.find((v: any) => {
-      if(isInNodeRagVariable)
+      if (isInNodeRagVariable)
         return v.variable === valueSelector.join('.')
       return v.variable === valueSelector[1]
-  })
+    })
     if (!targetVar)
       return VarType.string
 
-    if(isInNodeRagVariable)
+    if (isInNodeRagVariable)
       return targetVar.type
 
     const isStructuredOutputVar = !!targetVar.children?.schema?.properties
@@ -1084,6 +1084,13 @@ export const getNodeUsedVars = (node: Node): ValueSelector[] => {
       res = [...(mixVars as ValueSelector[]), ...(vars as any)]
       break
     }
+    case BlockEnum.DataSource: {
+      const payload = data as DataSourceNodeType
+      const mixVars = matchNotSystemVars(Object.keys(payload.datasource_parameters)?.filter(key => payload.datasource_parameters[key].type === ToolVarType.mixed).map(key => payload.datasource_parameters[key].value) as string[])
+      const vars = Object.keys(payload.datasource_parameters).filter(key => payload.datasource_parameters[key].type === ToolVarType.variable).map(key => payload.datasource_parameters[key].value as string) || []
+      res = [...(mixVars as ValueSelector[]), ...(vars as any)]
+      break
+    }
 
     case BlockEnum.VariableAssigner: {
       res = (data as VariableAssignerNodeType)?.variables
@@ -1357,6 +1364,30 @@ export const updateNodeVars = (oldNode: Node, oldVarSelector: ValueSelector, new
         }
         break
       }
+      case BlockEnum.DataSource: {
+        const payload = data as DataSourceNodeType
+        const hasShouldRenameVar = Object.keys(payload.datasource_parameters)?.filter(key => payload.datasource_parameters[key].type !== ToolVarType.constant)
+        if (hasShouldRenameVar) {
+          Object.keys(payload.datasource_parameters).forEach((key) => {
+            const value = payload.datasource_parameters[key]
+            const { type } = value
+            if (type === ToolVarType.variable) {
+              payload.datasource_parameters[key] = {
+                ...value,
+                value: newVarSelector,
+              }
+            }
+
+            if (type === ToolVarType.mixed) {
+              payload.datasource_parameters[key] = {
+                ...value,
+                value: replaceOldVarInText(payload.datasource_parameters[key].value as string, oldVarSelector, newVarSelector),
+              }
+            }
+          })
+        }
+        break
+      }
       case BlockEnum.VariableAssigner: {
         const payload = data as VariableAssignerNodeType
         if (payload.variables) {