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 e2b1f0a31c805c..7fb4ad68d8bbda 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 @@ -8,6 +8,7 @@ import { } from '@remixicon/react' import produce from 'immer' import { useStoreApi } from 'reactflow' +import useAvailableVarList from '../../hooks/use-available-var-list' import VarReferencePopup from './var-reference-popup' import { getNodeInfoById, isConversationVar, isENV, isSystemVar } from './utils' import ConstantField from './constant-field' @@ -26,7 +27,6 @@ import { } from '@/app/components/base/portal-to-follow-elem' import { useIsChatMode, - useWorkflow, useWorkflowVariables, } from '@/app/components/workflow/hooks' import { VarType as VarKindType } from '@/app/components/workflow/nodes/tool/types' @@ -67,7 +67,7 @@ const VarReferencePicker: FC = ({ onlyLeafNodeVar, filterVar = () => true, availableNodes: passedInAvailableNodes, - availableVars, + availableVars: passedInAvailableVars, isAddBtnTrigger, schema, valueTypePlaceHolder, @@ -79,11 +79,12 @@ const VarReferencePicker: FC = ({ } = store.getState() const isChatMode = useIsChatMode() - const { getTreeLeafNodes, getBeforeNodesInSameBranch } = useWorkflow() - const { getCurrentVariableType, getNodeAvailableVars } = useWorkflowVariables() - const availableNodes = useMemo(() => { - return passedInAvailableNodes || (onlyLeafNodeVar ? getTreeLeafNodes(nodeId) : getBeforeNodesInSameBranch(nodeId)) - }, [getBeforeNodesInSameBranch, getTreeLeafNodes, nodeId, onlyLeafNodeVar, passedInAvailableNodes]) + const { getCurrentVariableType } = useWorkflowVariables() + const { availableNodes, availableVars } = useAvailableVarList(nodeId, { + onlyLeafNodeVar, + passedInAvailableNodes, + filterVar, + }) const startNode = availableNodes.find((node: any) => { return node.data.type === BlockEnum.Start }) @@ -102,19 +103,8 @@ const VarReferencePicker: FC = ({ const [varKindType, setVarKindType] = useState(defaultVarKindType) const isConstant = isSupportConstantValue && varKindType === VarKindType.constant - const outputVars = useMemo(() => { - if (availableVars) - return availableVars - - const vars = getNodeAvailableVars({ - parentNode: iterationNode, - beforeNodes: availableNodes, - isChatMode, - filterVar, - }) - return vars - }, [iterationNode, availableNodes, isChatMode, filterVar, availableVars, getNodeAvailableVars]) + const outputVars = useMemo(() => (passedInAvailableVars || availableVars), [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 b81feab8053700..bd17bb1de859a7 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,12 +4,13 @@ import { useWorkflow, useWorkflowVariables, } from '@/app/components/workflow/hooks' -import type { ValueSelector, Var } from '@/app/components/workflow/types' +import type { Node, ValueSelector, Var } from '@/app/components/workflow/types' type Params = { onlyLeafNodeVar?: boolean hideEnv?: boolean hideChatVar?: boolean filterVar: (payload: Var, selector: ValueSelector) => boolean + passedInAvailableNodes?: Node[] } const useAvailableVarList = (nodeId: string, { @@ -17,6 +18,7 @@ const useAvailableVarList = (nodeId: string, { filterVar, hideEnv, hideChatVar, + passedInAvailableNodes, }: Params = { onlyLeafNodeVar: false, filterVar: () => true, @@ -25,7 +27,7 @@ const useAvailableVarList = (nodeId: string, { const { getNodeAvailableVars } = useWorkflowVariables() const isChatMode = useIsChatMode() - const availableNodes = onlyLeafNodeVar ? getTreeLeafNodes(nodeId) : getBeforeNodesInSameBranch(nodeId) + const availableNodes = passedInAvailableNodes || (onlyLeafNodeVar ? getTreeLeafNodes(nodeId) : getBeforeNodesInSameBranch(nodeId)) const { parentNode: iterationNode,