Feat: The query variables of the subsequent operators can reference the structured variables defined in the agent operator. #10866 (#10902)

### What problem does this PR solve?

Feat: The query variables of the subsequent operators can reference the
structured variables defined in the agent operator. #10866

### Type of change


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

web/src/components/prompt-editor/variable-node.tsx

@@ -1,5 +1,5 @@
 import i18n from '@/locales/config';
-import { BeginId } from '@/pages/flow/constant';
+import { BeginId } from '@/pages/agent/constant';
 import { DecoratorNode, LexicalNode, NodeKey } from 'lexical';
 import { ReactNode } from 'react';
 const prefix = BeginId + '@';

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

@@ -114,7 +114,7 @@ export default function VariablePickerMenuPlugin({
     minLength: 0,
   });
 
-  const [queryString, setQueryString] = React.useState<string| null>('');
+  const [queryString, setQueryString] = React.useState<string | null>('');
 
   const options = useBuildComponentIdSelectOptions(node?.id, node?.parentId);
 

web/src/locales/en.ts

@@ -1623,7 +1623,7 @@ This delimiter is used to split the input text into several text pieces echo of
       extractorDescription:
         'Use an LLM to extract structured insights from document chunks—such as summaries, classifications, etc.',
       outputFormat: 'Output format',
-      fileFormats: 'File format',
+      fileFormats: 'File type',
       fileFormatOptions: {
         pdf: 'PDF',
         spreadsheet: 'Spreadsheet',
@@ -1644,7 +1644,7 @@ This delimiter is used to split the input text into several text pieces echo of
       searchMethodTip: `Defines how the content can be searched — by full-text, embedding, or both.
 The Indexer will store the content in the corresponding data structures for the selected methods.`,
       // file: 'File',
-      parserMethod: 'Parsing method',
+      parserMethod: 'PDF parser',
       // systemPrompt: 'System Prompt',
       systemPromptPlaceholder:
         'Enter system prompt for image analysis, if empty the system default value will be used',

web/src/locales/zh.ts

@@ -1529,7 +1529,7 @@ General：实体和关系提取提示来自 GitHub - microsoft/graphrag：基于
       extractorDescription:
         '使用 LLM 从文档块（例如摘要、分类等）中提取结构化见解。',
       outputFormat: '输出格式',
-      fileFormats: '文件格式',
+      fileFormats: '文件类型',
       fields: '字段',
       addParser: '增加解析器',
       hierarchy: '层次结构',

web/src/pages/agent/constant/index.tsx

@@ -617,7 +617,10 @@ export const initialAgentValues = {
       type: 'string',
       value: '',
     },
-    [AgentStructuredOutputField]: {},
+    [AgentStructuredOutputField]: {
+      type: 'Object Array String Number Boolean',
+      value: '',
+    },
   },
 };
 

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

@@ -31,20 +31,15 @@ import * as ReactDOM from 'react-dom';
 import { $createVariableNode } from './variable-node';
 
 import {
-  HoverCard,
+  useFilterStructuredOutputByValue,
-  HoverCardContent,
+  useFindAgentStructuredOutputLabel,
-  HoverCardTrigger,
+  useShowSecondaryMenu,
-} from '@/components/ui/hover-card';
+} from '@/pages/agent/hooks/use-build-structured-output';
-import { Operator } from '@/constants/agent';
-import { cn } from '@/lib/utils';
-import { AgentStructuredOutputField } from '@/pages/agent/constant';
 import { useBuildQueryVariableOptions } from '@/pages/agent/hooks/use-get-begin-query';
-import useGraphStore from '@/pages/agent/store';
-import { get, isPlainObject } from 'lodash';
 import { PromptIdentity } from '../../agent-form/use-build-prompt-options';
+import { StructuredOutputSecondaryMenu } from '../structured-output-secondary-menu';
 import { ProgrammaticTag } from './constant';
 import './index.css';
-import { filterAgentStructuredOutput } from './utils';
 class VariableInnerOption extends MenuOption {
   label: string;
   value: string;
@@ -82,10 +77,6 @@ class VariableOption extends MenuOption {
   }
 }
 
-function getNodeId(value: string) {
-  return value.split('@').at(0);
-}
-
 function VariablePickerMenuItem({
   index,
   option,
@@ -97,58 +88,9 @@ function VariablePickerMenuItem({
     option: VariableOption | VariableInnerOption,
   ) => void;
 }) {
-  const { getOperatorTypeFromId, getNode, clickedNodeId } = useGraphStore(
+  const filterStructuredOutput = useFilterStructuredOutputByValue();
-    (state) => state,
-  );
 
-  const showSecondaryMenu = useCallback(
+  const showSecondaryMenu = useShowSecondaryMenu();
-    (value: string, outputLabel: string) => {
-      const nodeId = getNodeId(value);
-      return (
-        getOperatorTypeFromId(nodeId) === Operator.Agent &&
-        outputLabel === AgentStructuredOutputField
-      );
-    },
-    [getOperatorTypeFromId],
-  );
-
-  const renderAgentStructuredOutput = useCallback(
-    (values: any, option: VariableInnerOption) => {
-      if (isPlainObject(values) && 'properties' in values) {
-        return (
-          <ul className="border-l">
-            {Object.entries(values.properties).map(([key, value]) => {
-              const nextOption = new VariableInnerOption(
-                option.label + `.${key}`,
-                option.value + `.${key}`,
-                option.parentLabel,
-                option.icon,
-              );
-
-              const dataType = get(value, 'type');
-
-              return (
-                <li key={key} className="pl-1">
-                  <div
-                    onClick={() => selectOptionAndCleanUp(nextOption)}
-                    className="hover:bg-bg-card p-1 text-text-primary rounded-sm flex justify-between"
-                  >
-                    {key}
-                    <span className="text-text-secondary">{dataType}</span>
-                  </div>
-                  {dataType === 'object' &&
-                    renderAgentStructuredOutput(value, nextOption)}
-                </li>
-              );
-            })}
-          </ul>
-        );
-      }
-
-      return <div></div>;
-    },
-    [selectOptionAndCleanUp],
-  );
 
   return (
     <li
@@ -165,39 +107,20 @@ function VariablePickerMenuItem({
             const shouldShowSecondary = showSecondaryMenu(x.value, x.label);
 
             if (shouldShowSecondary) {
-              const node = getNode(getNodeId(x.value));
+              const filteredStructuredOutput = filterStructuredOutput(x.value);
-              const structuredOutput = get(
-                node,
-                `data.form.outputs.${AgentStructuredOutputField}`,
-              );
-
-              const filteredStructuredOutput = filterAgentStructuredOutput(
-                structuredOutput,
-                getOperatorTypeFromId(clickedNodeId),
-              );
 
               return (
-                <HoverCard key={x.value} openDelay={100} closeDelay={100}>
+                <StructuredOutputSecondaryMenu
-                  <HoverCardTrigger asChild>
+                  key={x.value}
-                    <li className="hover:bg-bg-card p-1 text-text-primary rounded-sm">
+                  data={x}
-                      {x.label}
+                  click={(y) =>
-                    </li>
+                    selectOptionAndCleanUp({
-                  </HoverCardTrigger>
+                      ...x,
-                  <HoverCardContent
+                      ...y,
-                    side="left"
+                    } as VariableInnerOption)
-                    align="start"
+                  }
-                    className={cn(
+                  filteredStructuredOutput={filteredStructuredOutput}
-                      'min-w-[140px]  border border-border rounded-md shadow-lg p-0',
+                ></StructuredOutputSecondaryMenu>
-                    )}
-                  >
-                    <section className="p-2">
-                      <div className="p-1">
-                        {x.parentLabel} structured output:
-                      </div>
-                      {renderAgentStructuredOutput(filteredStructuredOutput, x)}
-                    </section>
-                  </HoverCardContent>
-                </HoverCard>
               );
             }
 
@@ -239,9 +162,8 @@ export default function VariablePickerMenuPlugin({
   baseOptions,
 }: VariablePickerMenuPluginProps): JSX.Element {
   const [editor] = useLexicalComposerContext();
-  const getOperatorTypeFromId = useGraphStore(
+
-    (state) => state.getOperatorTypeFromId,
+  const findAgentStructuredOutputLabel = useFindAgentStructuredOutputLabel();
-  );
 
   // const checkForTriggerMatch = useBasicTypeaheadTriggerMatch('/', {
   //   minLength: 0,
@@ -313,27 +235,17 @@ export default function VariablePickerMenuPlugin({
       }, []);
 
       // agent structured output
-      const fields = value.split('@');
+      const agentStructuredOutput = findAgentStructuredOutputLabel(
-      if (
+        value,
-        getOperatorTypeFromId(fields.at(0)) === Operator.Agent &&
+        children,
-        fields.at(1)?.startsWith(AgentStructuredOutputField)
+      );
-      ) {
+      if (agentStructuredOutput) {
-        // is agent structured output
+        return agentStructuredOutput;
-        const agentOption = children.find((x) => value.includes(x.value));
-        const jsonSchemaFields = fields
-          .at(1)
-          ?.slice(AgentStructuredOutputField.length);
-
-        return {
-          ...agentOption,
-          label: (agentOption?.label ?? '') + jsonSchemaFields,
-          value: value,
-        };
       }
 
       return children.find((x) => x.value === value);
     },
-    [getOperatorTypeFromId, options],
+    [findAgentStructuredOutputLabel, options],
   );
 
   const onSelectOption = useCallback(

web/src/pages/agent/form/components/query-variable.tsx

@@ -1,4 +1,3 @@
-import { SelectWithSearch } from '@/components/originui/select-with-search';
 import {
   FormControl,
   FormField,
@@ -12,6 +11,7 @@ import { useFormContext } from 'react-hook-form';
 import { useTranslation } from 'react-i18next';
 import { VariableType } from '../../constant';
 import { useBuildQueryVariableOptions } from '../../hooks/use-get-begin-query';
+import { GroupedSelectWithSecondaryMenu } from './select-with-secondary-menu';
 
 type QueryVariableProps = {
   name?: string;
@@ -52,11 +52,11 @@ export function QueryVariable({
             </FormLabel>
           )}
           <FormControl>
-            <SelectWithSearch
+            <GroupedSelectWithSecondaryMenu
               options={finalOptions}
               {...field}
-              allowClear
+              // allowClear
-            ></SelectWithSearch>
+            ></GroupedSelectWithSecondaryMenu>
           </FormControl>
           <FormMessage />
         </FormItem>

web/src/pages/agent/form/components/select-with-secondary-menu.tsx

@@ -0,0 +1,206 @@
+import { Button } from '@/components/ui/button';
+import {
+  Command,
+  CommandGroup,
+  CommandInput,
+  CommandItem,
+  CommandList,
+} from '@/components/ui/command';
+import {
+  HoverCard,
+  HoverCardContent,
+  HoverCardTrigger,
+} from '@/components/ui/hover-card';
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from '@/components/ui/popover';
+import { cn } from '@/lib/utils';
+import { ChevronDown, X } from 'lucide-react';
+import * as React from 'react';
+import { useCallback } from 'react';
+import {
+  useFilterStructuredOutputByValue,
+  useFindAgentStructuredOutputLabel,
+  useShowSecondaryMenu,
+} from '../../hooks/use-build-structured-output';
+import { StructuredOutputSecondaryMenu } from './structured-output-secondary-menu';
+
+type Item = {
+  label: string;
+  value: string;
+};
+
+type Option = {
+  label: string;
+  value: string;
+  children?: Item[];
+};
+
+type Group = {
+  label: string | React.ReactNode;
+  options: Option[];
+};
+
+interface GroupedSelectWithSecondaryMenuProps {
+  options: Group[];
+  value?: string;
+  onChange?: (value: string) => void;
+  placeholder?: string;
+}
+
+export function GroupedSelectWithSecondaryMenu({
+  options,
+  value,
+  onChange,
+  placeholder = 'Select an option...',
+}: GroupedSelectWithSecondaryMenuProps) {
+  const [open, setOpen] = React.useState(false);
+
+  const showSecondaryMenu = useShowSecondaryMenu();
+  const filterStructuredOutput = useFilterStructuredOutputByValue();
+  const findAgentStructuredOutputLabel = useFindAgentStructuredOutputLabel();
+
+  // Find the label of the selected item
+  const flattenedOptions = options.flatMap((g) => g.options);
+  let selectedLabel =
+    flattenedOptions
+      .flatMap((o) => [o, ...(o.children || [])])
+      .find((o) => o.value === value)?.label || '';
+
+  if (!selectedLabel && value) {
+    selectedLabel =
+      findAgentStructuredOutputLabel(value, flattenedOptions)?.label ?? '';
+  }
+
+  // Handle clear click
+  const handleClear = (e: React.MouseEvent) => {
+    e.stopPropagation();
+    onChange?.('');
+    setOpen(false);
+  };
+
+  const handleSecondaryMenuClick = useCallback(
+    (record: Item) => {
+      onChange?.(record.value);
+      setOpen(false);
+    },
+    [onChange],
+  );
+
+  return (
+    <Popover open={open} onOpenChange={setOpen}>
+      <PopoverTrigger asChild>
+        <Button
+          variant="outline"
+          role="combobox"
+          aria-expanded={open}
+          className={cn(
+            'w-full justify-between text-sm font-normal',
+            !value && 'text-muted-foreground',
+          )}
+        >
+          <span className="truncate">{selectedLabel || placeholder}</span>
+          <div className="flex items-center gap-1">
+            {value && (
+              <X
+                className="h-4 w-4 text-muted-foreground hover:text-foreground cursor-pointer"
+                onClick={handleClear}
+              />
+            )}
+            <ChevronDown className="h-4 w-4 opacity-50" />
+          </div>
+        </Button>
+      </PopoverTrigger>
+
+      <PopoverContent className="p-0" align="start">
+        <Command>
+          <CommandInput placeholder="Search..." />
+          <CommandList className="overflow-visible">
+            {options.map((group, idx) => (
+              <CommandGroup key={idx} heading={group.label}>
+                {group.options.map((option) => {
+                  const shouldShowSecondary = showSecondaryMenu(
+                    option.value,
+                    option.label,
+                  );
+
+                  if (shouldShowSecondary) {
+                    const filteredStructuredOutput = filterStructuredOutput(
+                      option.value,
+                    );
+                    return (
+                      <StructuredOutputSecondaryMenu
+                        key={option.value}
+                        data={option}
+                        click={handleSecondaryMenuClick}
+                        filteredStructuredOutput={filteredStructuredOutput}
+                      ></StructuredOutputSecondaryMenu>
+                    );
+                  }
+
+                  return option.children ? (
+                    <HoverCard
+                      key={option.value}
+                      openDelay={100}
+                      closeDelay={150}
+                    >
+                      <HoverCardTrigger asChild>
+                        <CommandItem
+                          onSelect={() => {}}
+                          className="flex items-center justify-between cursor-default"
+                        >
+                          {option.label}
+                          <span className="ml-auto text-muted-foreground">
+                            ›
+                          </span>
+                        </CommandItem>
+                      </HoverCardTrigger>
+                      <HoverCardContent
+                        side="right"
+                        align="start"
+                        className="w-[180px] p-1"
+                      >
+                        {option.children.map((child) => (
+                          <div
+                            key={child.value}
+                            className={cn(
+                              'cursor-pointer rounded-sm px-2 py-1.5 text-sm hover:bg-accent hover:text-accent-foreground',
+                              value === child.value &&
+                                'bg-accent text-accent-foreground',
+                            )}
+                            onClick={() => {
+                              onChange?.(child.value);
+                              setOpen(false);
+                            }}
+                          >
+                            {child.label}
+                          </div>
+                        ))}
+                      </HoverCardContent>
+                    </HoverCard>
+                  ) : (
+                    <CommandItem
+                      key={option.value}
+                      onSelect={() => {
+                        onChange?.(option.value);
+                        setOpen(false);
+                      }}
+                      className={cn(
+                        value === option.value &&
+                          'bg-accent text-accent-foreground',
+                      )}
+                    >
+                      {option.label}
+                    </CommandItem>
+                  );
+                })}
+              </CommandGroup>
+            ))}
+          </CommandList>
+        </Command>
+      </PopoverContent>
+    </Popover>
+  );
+}

web/src/pages/agent/form/components/structured-output-secondary-menu.tsx

@@ -0,0 +1,80 @@
+import { JSONSchema } from '@/components/jsonjoy-builder';
+import {
+  HoverCard,
+  HoverCardContent,
+  HoverCardTrigger,
+} from '@/components/ui/hover-card';
+import { cn } from '@/lib/utils';
+import { get, isPlainObject } from 'lodash';
+import { PropsWithChildren, ReactNode, useCallback } from 'react';
+
+type DataItem = { label: ReactNode; value: string; parentLabel?: ReactNode };
+
+type StructuredOutputSecondaryMenuProps = {
+  data: DataItem;
+  click(option: { label: ReactNode; value: string }): void;
+  filteredStructuredOutput: JSONSchema;
+} & PropsWithChildren;
+export function StructuredOutputSecondaryMenu({
+  data,
+  click,
+  filteredStructuredOutput,
+}: StructuredOutputSecondaryMenuProps) {
+  const renderAgentStructuredOutput = useCallback(
+    (values: any, option: { label: ReactNode; value: string }) => {
+      if (isPlainObject(values) && 'properties' in values) {
+        return (
+          <ul className="border-l">
+            {Object.entries(values.properties).map(([key, value]) => {
+              const nextOption = {
+                label: option.label + `.${key}`,
+                value: option.value + `.${key}`,
+              };
+
+              const dataType = get(value, 'type');
+
+              return (
+                <li key={key} className="pl-1">
+                  <div
+                    onClick={() => click(nextOption)}
+                    className="hover:bg-bg-card p-1 text-text-primary rounded-sm flex justify-between"
+                  >
+                    {key}
+                    <span className="text-text-secondary">{dataType}</span>
+                  </div>
+                  {dataType === 'object' &&
+                    renderAgentStructuredOutput(value, nextOption)}
+                </li>
+              );
+            })}
+          </ul>
+        );
+      }
+
+      return <div></div>;
+    },
+    [click],
+  );
+
+  return (
+    <HoverCard key={data.value} openDelay={100} closeDelay={100}>
+      <HoverCardTrigger asChild>
+        <li className="hover:bg-bg-card p-1 text-text-primary rounded-sm">
+          {data.label}
+        </li>
+      </HoverCardTrigger>
+      <HoverCardContent
+        side="left"
+        align="start"
+        className={cn(
+          'min-w-[140px]  border border-border rounded-md shadow-lg p-0',
+        )}
+      >
+        <section className="p-2">
+          <div className="p-1">{data?.parentLabel} structured output:</div>
+          {renderAgentStructuredOutput(filteredStructuredOutput, data)}
+        </section>
+      </HoverCardContent>
+    </HoverCard>
+  );
+}

web/src/pages/agent/form/invoke-form/index.tsx

@@ -2,6 +2,7 @@ import { Collapse } from '@/components/collapse';
 import { FormContainer } from '@/components/form-container';
 import NumberInput from '@/components/originui/number-input';
 import { SelectWithSearch } from '@/components/originui/select-with-search';
+import { useIsDarkTheme } from '@/components/theme-provider';
 import { Button } from '@/components/ui/button';
 import {
   Form,
@@ -86,6 +87,8 @@ function InvokeForm({ node }: INextOperatorForm) {
 
   const variables = useWatch({ control: form.control, name: 'variables' });
 
+  const isDarkTheme = useIsDarkTheme();
+
   useWatchFormChange(node?.id, form);
 
   return (
@@ -147,7 +150,7 @@ function InvokeForm({ node }: INextOperatorForm) {
                   <Editor
                     height={200}
                     defaultLanguage="json"
-                    theme="vs-dark"
+                    theme={isDarkTheme ? 'vs-dark' : undefined}
                     {...field}
                   />
                 </FormControl>

web/src/pages/agent/hooks/use-build-structured-output.ts

@@ -0,0 +1,92 @@
+import { get } from 'lodash';
+import { ReactNode, useCallback } from 'react';
+import { AgentStructuredOutputField, Operator } from '../constant';
+import useGraphStore from '../store';
+import { filterAgentStructuredOutput } from '../utils/filter-agent-structured-output';
+
+function getNodeId(value: string) {
+  return value.split('@').at(0);
+}
+
+export function useShowSecondaryMenu() {
+  const { getOperatorTypeFromId } = useGraphStore((state) => state);
+
+  const showSecondaryMenu = useCallback(
+    (value: string, outputLabel: string) => {
+      const nodeId = getNodeId(value);
+      return (
+        getOperatorTypeFromId(nodeId) === Operator.Agent &&
+        outputLabel === AgentStructuredOutputField
+      );
+    },
+    [getOperatorTypeFromId],
+  );
+
+  return showSecondaryMenu;
+}
+
+export function useFilterStructuredOutputByValue() {
+  const { getOperatorTypeFromId, getNode, clickedNodeId } = useGraphStore(
+    (state) => state,
+  );
+
+  const filterStructuredOutput = useCallback(
+    (value: string) => {
+      const node = getNode(getNodeId(value));
+      const structuredOutput = get(
+        node,
+        `data.form.outputs.${AgentStructuredOutputField}`,
+      );
+
+      const filteredStructuredOutput = filterAgentStructuredOutput(
+        structuredOutput,
+        getOperatorTypeFromId(clickedNodeId),
+      );
+
+      return filteredStructuredOutput;
+    },
+    [clickedNodeId, getNode, getOperatorTypeFromId],
+  );
+
+  return filterStructuredOutput;
+}
+
+export function useFindAgentStructuredOutputLabel() {
+  const getOperatorTypeFromId = useGraphStore(
+    (state) => state.getOperatorTypeFromId,
+  );
+
+  const findAgentStructuredOutputLabel = useCallback(
+    (
+      value: string,
+      options: Array<{
+        label: string;
+        value: string;
+        parentLabel?: string | ReactNode;
+        icon?: ReactNode;
+      }>,
+    ) => {
+      // agent structured output
+      const fields = value.split('@');
+      if (
+        getOperatorTypeFromId(fields.at(0)) === Operator.Agent &&
+        fields.at(1)?.startsWith(AgentStructuredOutputField)
+      ) {
+        // is agent structured output
+        const agentOption = options.find((x) => value.includes(x.value));
+        const jsonSchemaFields = fields
+          .at(1)
+          ?.slice(AgentStructuredOutputField.length);
+
+        return {
+          ...agentOption,
+          label: (agentOption?.label ?? '') + jsonSchemaFields,
+          value: value,
+        };
+      }
+    },
+    [getOperatorTypeFromId],
+  );
+
+  return findAgentStructuredOutputLabel;
+}