feat: can support choose current node var

web/app/components/workflow/hooks/use-workflow.ts

@@ -63,6 +63,15 @@ export const useWorkflow = () => {
     workflowStore.setState({ panelWidth: width })
   }, [workflowStore])
 
+  const getNodeById = useCallback((nodeId: string) => {
+    const {
+      getNodes,
+    } = store.getState()
+    const nodes = getNodes()
+    const currentNode = nodes.find(node => node.id === nodeId)
+    return currentNode
+  }, [store])
+
   const getTreeLeafNodes = useCallback((nodeId: string) => {
     const {
       getNodes,
@@ -445,6 +454,7 @@ export const useWorkflow = () => {
 
   return {
     setPanelWidth,
+    getNodeById,
     getTreeLeafNodes,
     getBeforeNodesInSameBranch,
     getBeforeNodesInSameBranchIncludeParent,

web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx

@@ -17,7 +17,7 @@ import {
 import RemoveButton from '../remove-button'
 import useAvailableVarList from '../../hooks/use-available-var-list'
 import VarReferencePopup from './var-reference-popup'
-import { getNodeInfoById, inputVarTypeToVarType, isConversationVar, isENV, isRagVariableVar, isSystemVar, varTypeToStructType } from './utils'
+import { getNodeInfoById, isConversationVar, isENV, isRagVariableVar, isSystemVar, varTypeToStructType } from './utils'
 import ConstantField from './constant-field'
 import cn from '@/utils/classnames'
 import type { CommonNodeType, NodeOutPutVar, ValueSelector, Var } from '@/app/components/workflow/types'
@@ -141,25 +141,8 @@ const VarReferencePicker: FC<Props> = ({
 
   const outputVars = useMemo(() => {
     const results = passedInAvailableVars || availableVars
-
-    if (node?.data.type === BlockEnum.DataSource) {
-      const ragVariablesInDataSource = ragPipelineVariables?.find(ragVariable => ragVariable.belong_to_node_id === node.id)
-
-      if (ragVariablesInDataSource) {
-        results.unshift({
-          nodeId: node.id,
-          title: node.data?.title,
-          vars: [{
-            variable: `rag.${node.id}.${ragVariablesInDataSource.variable}`,
-            type: inputVarTypeToVarType(ragVariablesInDataSource.type as any),
-            description: ragVariablesInDataSource.label,
-            isRagVariable: true,
-          } as Var],
-        })
-      }
-    }
     return results
-  }, [passedInAvailableVars, availableVars, node, ragPipelineVariables])
+  }, [passedInAvailableVars, availableVars])
 
   const [open, setOpen] = useState(false)
   useEffect(() => {

web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts

@@ -4,7 +4,11 @@ import {
   useWorkflow,
   useWorkflowVariables,
 } from '@/app/components/workflow/hooks'
-import type { Node, ValueSelector, Var } from '@/app/components/workflow/types'
+import type { NodeOutPutVar } from '@/app/components/workflow/types'
+import { BlockEnum, type Node, type ValueSelector, type Var } from '@/app/components/workflow/types'
+import { useStore as useWorkflowStore } from '@/app/components/workflow/store'
+import { inputVarTypeToVarType } from '../../data-source/utils'
+
 type Params = {
   onlyLeafNodeVar?: boolean
   hideEnv?: boolean
@@ -24,24 +28,49 @@ const useAvailableVarList = (nodeId: string, {
   onlyLeafNodeVar: false,
   filterVar: () => true,
 }) => {
-  const { getTreeLeafNodes, getBeforeNodesInSameBranchIncludeParent } = useWorkflow()
+  const { getTreeLeafNodes, getNodeById, getBeforeNodesInSameBranchIncludeParent } = useWorkflow()
   const { getNodeAvailableVars } = useWorkflowVariables()
   const isChatMode = useIsChatMode()
-
   const availableNodes = passedInAvailableNodes || (onlyLeafNodeVar ? getTreeLeafNodes(nodeId) : getBeforeNodesInSameBranchIncludeParent(nodeId))
-
   const {
     parentNode: iterationNode,
   } = useNodeInfo(nodeId)
 
-  const availableVars = getNodeAvailableVars({
+  const currNode = getNodeById(nodeId)
+  const ragPipelineVariables = useWorkflowStore(s => s.ragPipelineVariables)
+  const isDataSourceNode = currNode?.data?.type === BlockEnum.DataSource
+  const dataSourceRagVars: NodeOutPutVar[] = []
+  if(isDataSourceNode) {
+    const ragVariablesInDataSource = ragPipelineVariables?.filter(ragVariable => ragVariable.belong_to_node_id === nodeId)
+    const filterVars = ragVariablesInDataSource?.filter(v => filterVar({
+      variable: v.variable,
+      type: inputVarTypeToVarType(v.type),
+      nodeId,
+      isRagVariable: true,
+    }, ['rag', nodeId, v.variable]))
+    if(filterVars?.length) {
+      dataSourceRagVars.push({
+          nodeId,
+          title: currNode.data?.title,
+          vars: filterVars.map((v) => {
+            return {
+              variable: `rag.${nodeId}.${v.variable}`,
+              type: inputVarTypeToVarType(v.type),
+              description: v.label,
+              isRagVariable: true,
+          } as Var
+          }),
+      })
+    }
+  }
+  const availableVars = [...getNodeAvailableVars({
     parentNode: iterationNode,
     beforeNodes: availableNodes,
     isChatMode,
     filterVar,
     hideEnv,
     hideChatVar,
-  })
+  }), ...dataSourceRagVars]
 
   return {
     availableVars,