diff --git a/web/app/components/workflow/nodes/_base/components/variable/utils.ts b/web/app/components/workflow/nodes/_base/components/variable/utils.ts index c59ebfdc68..b68a506c18 100644 --- a/web/app/components/workflow/nodes/_base/components/variable/utils.ts +++ b/web/app/components/workflow/nodes/_base/components/variable/utils.ts @@ -62,7 +62,7 @@ export const isSpecialVar = (prefix: string): boolean => { return ['sys', 'env', 'conversation', 'rag'].includes(prefix) } -const inputVarTypeToVarType = (type: InputVarType): VarType => { +export const inputVarTypeToVarType = (type: InputVarType): VarType => { return ({ [InputVarType.number]: VarType.number, [InputVarType.singleFile]: VarType.file, 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 18ee82d865..eb59beebd4 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 @@ -9,14 +9,18 @@ import { RiMoreLine, } from '@remixicon/react' import produce from 'immer' -import { useReactFlow, useStoreApi } from 'reactflow' +import { + useNodes, + useReactFlow, + useStoreApi, +} from 'reactflow' import RemoveButton from '../remove-button' import useAvailableVarList from '../../hooks/use-available-var-list' import VarReferencePopup from './var-reference-popup' -import { getNodeInfoById, isConversationVar, isENV, isRagVariableVar, isSystemVar, varTypeToStructType } from './utils' +import { getNodeInfoById, inputVarTypeToVarType, isConversationVar, isENV, isRagVariableVar, isSystemVar, varTypeToStructType } from './utils' import ConstantField from './constant-field' import cn from '@/utils/classnames' -import type { Node, NodeOutPutVar, ValueSelector, Var } from '@/app/components/workflow/types' +import type { CommonNodeType, NodeOutPutVar, ValueSelector, Var } from '@/app/components/workflow/types' import type { CredentialFormSchema } from '@/app/components/header/account-setting/model-provider-page/declarations' import { BlockEnum } from '@/app/components/workflow/types' import { VarBlockIcon } from '@/app/components/workflow/block-icon' @@ -41,7 +45,7 @@ import { isExceptionVariable } from '@/app/components/workflow/utils' import VarFullPathPanel from './var-full-path-panel' import { noop } from 'lodash-es' import { InputField } from '@/app/components/base/icons/src/vender/pipeline' - +import { useStore as useWorkflowStore } from '@/app/components/workflow/store' const TRIGGER_DEFAULT_WIDTH = 227 type Props = { @@ -101,11 +105,8 @@ const VarReferencePicker: FC = ({ }) => { const { t } = useTranslation() const store = useStoreApi() - const { - getNodes, - } = store.getState() + const nodes = useNodes() const isChatMode = useIsChatMode() - const { getCurrentVariableType } = useWorkflowVariables() const { availableVars, availableNodesWithParent: availableNodes } = useAvailableVarList(nodeId, { onlyLeafNodeVar, @@ -119,12 +120,13 @@ const VarReferencePicker: FC = ({ return node.data.type === BlockEnum.Start }) - const node = getNodes().find(n => n.id === nodeId) - const isInIteration = !!node?.data.isInIteration - const iterationNode = isInIteration ? getNodes().find(n => n.id === node.parentId) : null + const node = nodes.find(n => n.id === nodeId) + const ragPipelineVariables = useWorkflowStore(s => s.ragPipelineVariables) + const isInIteration = !!(node?.data as any).isInIteration + const iterationNode = isInIteration ? nodes.find(n => n.id === node?.parentId) : null - const isInLoop = !!node?.data.isInLoop - const loopNode = isInLoop ? getNodes().find(n => n.id === node.parentId) : null + const isInLoop = !!(node?.data as any).isInLoop + const loopNode = isInLoop ? nodes.find(n => n.id === node?.parentId) : null const triggerRef = useRef(null) const [triggerWidth, setTriggerWidth] = useState(TRIGGER_DEFAULT_WIDTH) @@ -137,7 +139,27 @@ const VarReferencePicker: FC = ({ const [varKindType, setVarKindType] = useState(defaultVarKindType) const isConstant = isSupportConstantValue && varKindType === VarKindType.constant - const outputVars = useMemo(() => (passedInAvailableVars || availableVars), [passedInAvailableVars, availableVars]) + 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]) const [open, setOpen] = useState(false) useEffect(() => { @@ -268,7 +290,7 @@ const VarReferencePicker: FC = ({ }, [availableNodes, reactflow, store]) const type = getCurrentVariableType({ - parentNode: isInIteration ? iterationNode : loopNode, + parentNode: (isInIteration ? iterationNode : loopNode) as any, valueSelector: value as ValueSelector, availableNodes, isChatMode,