From 06463135ef40504d9d598190872e286150ee786b Mon Sep 17 00:00:00 2001 From: balibabu Date: Fri, 6 Jun 2025 19:27:29 +0800 Subject: [PATCH] Feat: Reference the output variable of the upstream operator #3221 (#8111) ### What problem does this PR solve? Feat: Reference the output variable of the upstream operator #3221 ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- .../prompt-editor/variable-picker-plugin.tsx | 11 +-- .../form/invoke-form/dynamic-variables.tsx | 3 +- .../pages/agent/hooks/use-get-begin-query.tsx | 71 +++++++++++++++++-- 3 files changed, 70 insertions(+), 15 deletions(-) diff --git a/web/src/pages/agent/form/components/prompt-editor/variable-picker-plugin.tsx b/web/src/pages/agent/form/components/prompt-editor/variable-picker-plugin.tsx index 4592e64a0..c059c2acb 100644 --- a/web/src/pages/agent/form/components/prompt-editor/variable-picker-plugin.tsx +++ b/web/src/pages/agent/form/components/prompt-editor/variable-picker-plugin.tsx @@ -117,13 +117,9 @@ export default function VariablePickerMenuPlugin({ const [queryString, setQueryString] = React.useState(''); - const buildGroupedOptions = useBuildComponentIdSelectOptions( - node?.id, - node?.parentId, - ); + const options = useBuildComponentIdSelectOptions(node?.id, node?.parentId); const buildNextOptions = useCallback(() => { - const options = buildGroupedOptions(); let filteredOptions = options; if (queryString) { @@ -150,11 +146,10 @@ export default function VariablePickerMenuPlugin({ ); return nextOptions; - }, [buildGroupedOptions, queryString]); + }, [options, queryString]); const findLabelByValue = useCallback( (value: string) => { - const options = buildGroupedOptions(); const children = options.reduce>( (pre, cur) => { return pre.concat(cur.options); @@ -164,7 +159,7 @@ export default function VariablePickerMenuPlugin({ return children.find((x) => x.value === value)?.label; }, - [buildGroupedOptions], + [options], ); const onSelectOption = useCallback( diff --git a/web/src/pages/agent/form/invoke-form/dynamic-variables.tsx b/web/src/pages/agent/form/invoke-form/dynamic-variables.tsx index 3538b8b72..9a67c1f78 100644 --- a/web/src/pages/agent/form/invoke-form/dynamic-variables.tsx +++ b/web/src/pages/agent/form/invoke-form/dynamic-variables.tsx @@ -4,9 +4,10 @@ import { DeleteOutlined } from '@ant-design/icons'; import { Button, Collapse, Flex, Input, Select, Table, TableProps } from 'antd'; import { trim } from 'lodash'; import { useBuildComponentIdSelectOptions } from '../../hooks/use-get-begin-query'; -import { IInvokeVariable, RAGFlowNodeType } from '../../interface'; +import { IInvokeVariable } from '../../interface'; import { useHandleOperateParameters } from './hooks'; +import { RAGFlowNodeType } from '@/interfaces/database/flow'; import styles from './index.less'; interface IProps { diff --git a/web/src/pages/agent/hooks/use-get-begin-query.tsx b/web/src/pages/agent/hooks/use-get-begin-query.tsx index 052b51af0..0c167f2e4 100644 --- a/web/src/pages/agent/hooks/use-get-begin-query.tsx +++ b/web/src/pages/agent/hooks/use-get-begin-query.tsx @@ -1,5 +1,7 @@ import { RAGFlowNodeType } from '@/interfaces/database/flow'; +import { Edge } from '@xyflow/react'; import { DefaultOptionType } from 'antd/es/select'; +import { isEmpty } from 'lodash'; import get from 'lodash/get'; import { useCallback, useEffect, useMemo, useState } from 'react'; import { BeginId, Operator } from '../constant'; @@ -34,6 +36,61 @@ export const useGetBeginNodeDataQueryIsSafe = () => { return isBeginNodeDataQuerySafe; }; +function filterAllUpstreamNodeIds(edges: Edge[], nodeIds: string[]) { + return nodeIds.reduce((pre, nodeId) => { + const currentEdges = edges.filter((x) => x.target === nodeId); + + const upstreamNodeIds: string[] = currentEdges.map((x) => x.source); + + const ids = upstreamNodeIds.concat( + filterAllUpstreamNodeIds(edges, upstreamNodeIds), + ); + + ids.forEach((x) => { + if (pre.every((y) => y !== x)) { + pre.push(x); + } + }); + + return pre; + }, []); +} + +function buildOutputOptions(outputs: Record = {}) { + return Object.keys(outputs).map((x) => ({ + label: x, + value: x, + })); +} + +export function useBuildNodeOutputOptions(nodeId?: string) { + const nodes = useGraphStore((state) => state.nodes); + const edges = useGraphStore((state) => state.edges); + + const nodeOutputOptions = useMemo(() => { + if (!nodeId) { + return []; + } + const upstreamIds = filterAllUpstreamNodeIds(edges, [nodeId]); + + const nodeWithOutputList = nodes.filter( + (x) => + upstreamIds.some((y) => y === x.id) && !isEmpty(x.data?.form?.outputs), + ); + + return nodeWithOutputList + .filter((x) => x.id !== nodeId) + .map((x) => ({ + label: x.data.name, + value: x.id, + title: x.data.name, + options: buildOutputOptions(x.data.form.outputs), + })); + }, [edges, nodeId, nodes]); + + return nodeOutputOptions; +} + // exclude nodes with branches const ExcludedNodes = [ Operator.Categorize, @@ -49,6 +106,8 @@ export const useBuildComponentIdSelectOptions = ( const nodes = useGraphStore((state) => state.nodes); const getBeginNodeDataQuery = useGetBeginNodeDataQuery(); + const nodeOutputOptions = useBuildNodeOutputOptions(nodeId); + // Limit the nodes inside iteration to only reference peer nodes with the same parentId and other external nodes other than their parent nodes const filterChildNodesToSameParentOrExternal = useCallback( (node: RAGFlowNodeType) => { @@ -76,7 +135,7 @@ export const useBuildComponentIdSelectOptions = ( .map((x) => ({ label: x.data.name, value: x.id })); }, [nodes, nodeId, filterChildNodesToSameParentOrExternal]); - const buildGroupedOptions = useCallback(() => { + const options = useMemo(() => { const query: BeginQuery[] = getBeginNodeDataQuery(); return [ { @@ -92,21 +151,21 @@ export const useBuildComponentIdSelectOptions = ( value: `begin@${x.key}`, })), }, + ...nodeOutputOptions, ]; - }, [componentIdOptions, getBeginNodeDataQuery]); + }, [componentIdOptions, getBeginNodeDataQuery, nodeOutputOptions]); - return buildGroupedOptions; + return options; }; export const useGetComponentLabelByValue = (nodeId: string) => { - const buildGroupedOptions = useBuildComponentIdSelectOptions(nodeId); + const options = useBuildComponentIdSelectOptions(nodeId); const flattenOptions = useMemo(() => { - const options = buildGroupedOptions(); return options.reduce((pre, cur) => { return [...pre, ...cur.options]; }, []); - }, [buildGroupedOptions]); + }, [options]); const getLabel = useCallback( (val?: string) => {