diff --git a/web/app/components/workflow/hooks/use-workflow.ts b/web/app/components/workflow/hooks/use-workflow.ts index 03334ae58b..0bae109c27 100644 --- a/web/app/components/workflow/hooks/use-workflow.ts +++ b/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, diff --git a/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx b/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx index eb59beebd4..3e4506b832 100644 --- a/web/app/components/workflow/nodes/_base/components/variable/var-reference-picker.tsx +++ b/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 = ({ 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(() => { diff --git a/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts b/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts index e1a6a8bd8d..38e104c90f 100644 --- a/web/app/components/workflow/nodes/_base/hooks/use-available-var-list.ts +++ b/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,