Feat: Allows the extractor operator's prompt to reference the output of an upstream operator #9869 (#10279)

### What problem does this PR solve?

Feat: Allows the extractor operator's prompt to reference the output of
an upstream operator #9869

### Type of change


- [x] New Feature (non-breaking change which adds functionality)

web/src/locales/en.ts

@@ -1705,9 +1705,9 @@ This delimiter is used to split the input text into several text pieces echo of
       exportJson: 'Export JSON',
       viewResult: 'View Result',
       running: 'Running',
-      context: 'Context Generator',
-      contextDescription: 'Context Generator',
-      summary: 'Summary',
+      extractor: 'Extractor',
+      extractorDescription: 'Extractor',
+      summary: 'Augmented Context',
       keywords: 'Keywords',
       questions: 'Questions',
       metadata: 'Metadata',

web/src/locales/zh.ts

@@ -1623,9 +1623,9 @@ General：实体和关系提取提示来自 GitHub - microsoft/graphrag：基于
       exportJson: '导出 JSON',
       viewResult: '查看结果',
       running: '运行中',
-      context: '上下文生成器',
-      contextDescription: '上下文生成器',
-      summary: '摘要',
+      extractor: '提取器',
+      extractorDescription: '提取器',
+      summary: '增强上下文',
       keywords: '关键词',
       questions: '问题',
       metadata: '元数据',

web/src/pages/agent/form/components/prompt-editor/index.tsx

@@ -55,7 +55,7 @@ type IProps = {
   onChange?: (value?: string) => void;
   placeholder?: ReactNode;
 } & PromptContentProps &
-  Pick<VariablePickerMenuPluginProps, 'extraOptions'>;
+  Pick<VariablePickerMenuPluginProps, 'extraOptions' | 'baseOptions'>;
 
 function PromptContent({
   showToolbar = true,
@@ -126,6 +126,7 @@ export function PromptEditor({
   showToolbar,
   multiLine = true,
   extraOptions,
+  baseOptions,
 }: IProps) {
   const { t } = useTranslation();
   const initialConfig: InitialConfigType = {
@@ -177,6 +178,7 @@ export function PromptEditor({
         <VariablePickerMenuPlugin
           value={value}
           extraOptions={extraOptions}
+          baseOptions={baseOptions}
         ></VariablePickerMenuPlugin>
         <PasteHandlerPlugin />
         <VariableOnChangePlugin

web/src/pages/agent/form/components/prompt-editor/variable-picker-plugin.tsx

@@ -109,17 +109,26 @@ function VariablePickerMenuItem({
   );
 }
 
+export type VariablePickerMenuOptionType = {
+  label: string;
+  title: string;
+  value?: string;
+  options: Array<{
+    label: string;
+    value: string;
+    icon: ReactNode;
+  }>;
+};
+
 export type VariablePickerMenuPluginProps = {
   value?: string;
-  extraOptions?: Array<{
-    label: string;
-    title: string;
-    options: Array<{ label: string; value: string; icon?: ReactNode }>;
-  }>;
+  extraOptions?: VariablePickerMenuOptionType[];
+  baseOptions?: VariablePickerMenuOptionType[];
 };
 export default function VariablePickerMenuPlugin({
   value,
   extraOptions,
+  baseOptions,
 }: VariablePickerMenuPluginProps): JSX.Element {
   const [editor] = useLexicalComposerContext();
   const isFirstRender = useRef(true);
@@ -132,6 +141,10 @@ export default function VariablePickerMenuPlugin({
 
   let options = useBuildQueryVariableOptions();
 
+  if (baseOptions) {
+    options = baseOptions as typeof options;
+  }
+
   const buildNextOptions = useCallback(() => {
     let filteredOptions = [...options, ...(extraOptions ?? [])];
     if (queryString) {

web/src/pages/agent/form/iteration-form/use-build-options.ts

@@ -1,7 +1,7 @@
+import { buildOutputOptions } from '@/utils/canvas-util';
 import { isEmpty } from 'lodash';
 import { useMemo } from 'react';
 import { Operator } from '../../constant';
-import { buildOutputOptions } from '../../hooks/use-get-begin-query';
 import useGraphStore from '../../store';
 
 export function useBuildSubNodeOutputOptions(nodeId?: string) {

web/src/pages/agent/hooks/use-get-begin-query.tsx

@@ -1,19 +1,11 @@
 import { AgentGlobals } from '@/constants/agent';
 import { useFetchAgent } from '@/hooks/use-agent-request';
 import { RAGFlowNodeType } from '@/interfaces/database/flow';
-import { Edge } from '@xyflow/react';
+import { buildNodeOutputOptions } from '@/utils/canvas-util';
 import { DefaultOptionType } from 'antd/es/select';
 import { t } from 'i18next';
-import { isEmpty } from 'lodash';
 import get from 'lodash/get';
-import {
-  ReactNode,
-  useCallback,
-  useContext,
-  useEffect,
-  useMemo,
-  useState,
-} from 'react';
+import { useCallback, useContext, useEffect, useMemo, useState } from 'react';
 import {
   AgentDialogueMode,
   BeginId,
@@ -83,72 +75,18 @@ export const useGetBeginNodeDataQueryIsSafe = () => {
   return isBeginNodeDataQuerySafe;
 };
 
-function filterAllUpstreamNodeIds(edges: Edge[], nodeIds: string[]) {
-  return nodeIds.reduce<string[]>((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;
-  }, []);
-}
-
-export function buildOutputOptions(
-  outputs: Record<string, any> = {},
-  nodeId?: string,
-  parentLabel?: string | ReactNode,
-  icon?: ReactNode,
-) {
-  return Object.keys(outputs).map((x) => ({
-    label: x,
-    value: `${nodeId}@${x}`,
-    parentLabel,
-    icon,
-    type: outputs[x]?.type,
-  }));
-}
-
 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,
-          x.id,
-          x.data.name,
-          <OperatorIcon name={x.data.label as Operator} />,
-        ),
-      }));
+  return useMemo(() => {
+    return buildNodeOutputOptions({
+      nodes,
+      edges,
+      nodeId,
+      Icon: ({ name }) => <OperatorIcon name={name as Operator}></OperatorIcon>,
+    });
   }, [edges, nodeId, nodes]);
-
-  return nodeOutputOptions;
 }
 
 // exclude nodes with branches

web/src/pages/data-flow/canvas/node/dropdown/next-step-dropdown.tsx

@@ -124,7 +124,7 @@ function AccordionOperators({
         Operator.Tokenizer,
         Operator.Splitter,
         Operator.HierarchicalMerger,
-        Operator.Context,
+        Operator.Extractor,
       ]}
       isCustomDropdown={isCustomDropdown}
       mousePosition={mousePosition}

web/src/pages/data-flow/constant.tsx

@@ -119,7 +119,7 @@ export enum Operator {
   Tokenizer = 'Tokenizer',
   Splitter = 'Splitter',
   HierarchicalMerger = 'HierarchicalMerger',
-  Context = 'Context',
+  Extractor = 'Extractor',
 }
 
 export const SwitchLogicOperatorOptions = ['and', 'or'];
@@ -291,7 +291,7 @@ export const initialHierarchicalMergerValues = {
 
 export const initialContextValues = {
   ...initialLlmBaseValues,
-  field_name: [ContextGeneratorFieldName.Summary],
+  field_name: ContextGeneratorFieldName.Summary,
   outputs: {},
 };
 
@@ -327,7 +327,7 @@ export const NodeMap = {
   [Operator.Tokenizer]: 'tokenizerNode',
   [Operator.Splitter]: 'splitterNode',
   [Operator.HierarchicalMerger]: 'hierarchicalMergerNode',
-  [Operator.Context]: 'contextNode',
+  [Operator.Extractor]: 'contextNode',
 };
 
 export enum BeginQueryType {

web/src/pages/data-flow/form-sheet/form-config-map.tsx

@@ -1,5 +1,5 @@
 import { Operator } from '../constant';
-import ContextForm from '../form/context-form';
+import ExtractorForm from '../form/extractor-form';
 import HierarchicalMergerForm from '../form/hierarchical-merger-form';
 import ParserForm from '../form/parser-form';
 import SplitterForm from '../form/splitter-form';
@@ -24,7 +24,7 @@ export const FormConfigMap = {
   [Operator.HierarchicalMerger]: {
     component: HierarchicalMergerForm,
   },
-  [Operator.Context]: {
-    component: ContextForm,
+  [Operator.Extractor]: {
+    component: ExtractorForm,
   },
 };

web/src/pages/data-flow/form/extractor-form/index.tsx

@@ -1,9 +1,8 @@
 import { LargeModelFormField } from '@/components/large-model-form-field';
 import { LlmSettingSchema } from '@/components/llm-setting-items/next';
+import { SelectWithSearch } from '@/components/originui/select-with-search';
 import { RAGFlowFormItem } from '@/components/ragflow-form';
 import { Form } from '@/components/ui/form';
-import { MultiSelect } from '@/components/ui/multi-select';
-import { useBuildPromptExtraPromptOptions } from '@/pages/agent/form/agent-form/use-build-prompt-options';
 import { PromptEditor } from '@/pages/agent/form/components/prompt-editor';
 import { buildOptions } from '@/utils/form';
 import { zodResolver } from '@hookform/resolvers/zod';
@@ -15,15 +14,11 @@ import {
   ContextGeneratorFieldName,
   initialContextValues,
 } from '../../constant';
+import { useBuildNodeOutputOptions } from '../../hooks/use-build-options';
 import { useFormValues } from '../../hooks/use-form-values';
 import { useWatchFormChange } from '../../hooks/use-watch-form-change';
 import { INextOperatorForm } from '../../interface';
-import useGraphStore from '../../store';
-import { buildOutputList } from '../../utils/build-output-list';
 import { FormWrapper } from '../components/form-wrapper';
-import { Output } from '../components/output';
-
-const outputList = buildOutputList(initialContextValues.outputs);
 
 export const FormSchema = z.object({
   sys_prompt: z.string(),
@@ -32,20 +27,18 @@ export const FormSchema = z.object({
   field_name: z.array(z.string()),
 });
 
-export type ContextFormSchemaType = z.infer<typeof FormSchema>;
+export type ExtractorFormSchemaType = z.infer<typeof FormSchema>;
 
-const ContextForm = ({ node }: INextOperatorForm) => {
+const ExtractorForm = ({ node }: INextOperatorForm) => {
   const defaultValues = useFormValues(initialContextValues, node);
   const { t } = useTranslation();
 
-  const form = useForm<ContextFormSchemaType>({
+  const form = useForm<ExtractorFormSchemaType>({
     defaultValues,
     resolver: zodResolver(FormSchema),
   });
 
-  const { edges } = useGraphStore((state) => state);
-
-  const { extraOptions } = useBuildPromptExtraPromptOptions(edges, node?.id);
+  const promptOptions = useBuildNodeOutputOptions(node?.id);
 
   const options = buildOptions(ContextGeneratorFieldName, t, 'dataflow');
 
@@ -55,32 +48,31 @@ const ContextForm = ({ node }: INextOperatorForm) => {
     <Form {...form}>
       <FormWrapper>
         <LargeModelFormField></LargeModelFormField>
+        <RAGFlowFormItem label={t('dataflow.fieldName')} name="field_name">
+          {(field) => (
+            <SelectWithSearch
+              {...field}
+              placeholder={t('dataFlowPlaceholder')}
+              options={options}
+            ></SelectWithSearch>
+          )}
+        </RAGFlowFormItem>
         <RAGFlowFormItem label={t('flow.systemPrompt')} name="sys_prompt">
           <PromptEditor
             placeholder={t('flow.messagePlaceholder')}
             showToolbar={true}
-            extraOptions={extraOptions}
+            baseOptions={promptOptions}
           ></PromptEditor>
         </RAGFlowFormItem>
         <RAGFlowFormItem label={t('flow.userPrompt')} name="prompts">
-          <PromptEditor showToolbar={true}></PromptEditor>
-        </RAGFlowFormItem>
-        <RAGFlowFormItem label={t('dataflow.fieldName')} name="field_name">
-          {(field) => (
-            <MultiSelect
-              onValueChange={field.onChange}
-              placeholder={t('dataFlowPlaceholder')}
-              defaultValue={field.value}
-              options={options}
-            ></MultiSelect>
-          )}
+          <PromptEditor
+            showToolbar={true}
+            baseOptions={promptOptions}
+          ></PromptEditor>
         </RAGFlowFormItem>
       </FormWrapper>
-      <div className="p-5">
-        <Output list={outputList}></Output>
-      </div>
     </Form>
   );
 };
 
-export default memo(ContextForm);
+export default memo(ExtractorForm);

web/src/pages/data-flow/hooks/use-add-node.ts

@@ -33,7 +33,7 @@ export const useInitializeOperatorParams = () => {
       [Operator.Tokenizer]: initialTokenizerValues,
       [Operator.Splitter]: initialSplitterValues,
       [Operator.HierarchicalMerger]: initialHierarchicalMergerValues,
-      [Operator.Context]: { ...initialContextValues, llm_id: llmId },
+      [Operator.Extractor]: { ...initialContextValues, llm_id: llmId },
     };
   }, [llmId]);
 

web/src/pages/data-flow/hooks/use-build-options.tsx

@@ -0,0 +1,19 @@
+import { buildNodeOutputOptions } from '@/utils/canvas-util';
+import { useMemo } from 'react';
+import { Operator } from '../constant';
+import OperatorIcon from '../operator-icon';
+import useGraphStore from '../store';
+
+export function useBuildNodeOutputOptions(nodeId?: string) {
+  const nodes = useGraphStore((state) => state.nodes);
+  const edges = useGraphStore((state) => state.edges);
+
+  return useMemo(() => {
+    return buildNodeOutputOptions({
+      nodes,
+      edges,
+      nodeId,
+      Icon: ({ name }) => <OperatorIcon name={name as Operator}></OperatorIcon>,
+    });
+  }, [edges, nodeId, nodes]);
+}

web/src/pages/data-flow/operator-icon.tsx

@@ -25,7 +25,7 @@ export const SVGIconMap = {
   [Operator.Tokenizer]: ListMinus,
   [Operator.Splitter]: Blocks,
   [Operator.HierarchicalMerger]: Heading,
-  [Operator.Context]: FileStack,
+  [Operator.Extractor]: FileStack,
 };
 
 const Empty = () => {

web/src/utils/canvas-util.tsx

@@ -0,0 +1,75 @@
+import { BaseNode } from '@/interfaces/database/agent';
+import { Edge } from '@xyflow/react';
+import { isEmpty } from 'lodash';
+import { ComponentType, ReactNode } from 'react';
+
+export function filterAllUpstreamNodeIds(edges: Edge[], nodeIds: string[]) {
+  return nodeIds.reduce<string[]>((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;
+  }, []);
+}
+
+export function buildOutputOptions(
+  outputs: Record<string, any> = {},
+  nodeId?: string,
+  parentLabel?: string | ReactNode,
+  icon?: ReactNode,
+) {
+  return Object.keys(outputs).map((x) => ({
+    label: x,
+    value: `${nodeId}@${x}`,
+    parentLabel,
+    icon,
+    type: outputs[x]?.type,
+  }));
+}
+
+export function buildNodeOutputOptions({
+  nodes,
+  edges,
+  nodeId,
+  Icon,
+}: {
+  nodes: BaseNode[];
+  edges: Edge[];
+  nodeId?: string;
+  Icon: ComponentType<{ name: string }>;
+}) {
+  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,
+        x.id,
+        x.data.name,
+        <Icon name={x.data.name} />,
+      ),
+    }));
+}