Feat: Filter the query variable drop-down box options by type #3221 (#8485)

### What problem does this PR solve?

Feat: Filter the query variable drop-down box options by type #3221

### Type of change


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

web/src/components/originui/select-with-search.tsx

@@ -168,3 +168,5 @@ export const SelectWithSearch = forwardRef<
     </Popover>
   );
 });
+
+SelectWithSearch.displayName = 'SelectWithSearch';

web/src/pages/agent/canvas/node/agent-node.tsx

@@ -1,8 +1,9 @@
 import { IAgentNode } from '@/interfaces/database/flow';
 import { Handle, NodeProps, Position } from '@xyflow/react';
 import { memo, useMemo } from 'react';
-import { NodeHandleId, Operator } from '../../constant';
+import { NodeHandleId } from '../../constant';
 import useGraphStore from '../../store';
+import { isBottomSubAgent } from '../../utils';
 import { CommonHandle } from './handle';
 import { LeftHandleStyle, RightHandleStyle } from './handle-icon';
 import styles from './index.less';
@@ -16,19 +17,16 @@ function InnerAgentNode({
   isConnectable = true,
   selected,
 }: NodeProps<IAgentNode>) {
-  const getNode = useGraphStore((state) => state.getNode);
   const edges = useGraphStore((state) => state.edges);
 
-  const isNotParentAgent = useMemo(() => {
-    const edge = edges.find((x) => x.target === id);
-    const label = getNode(edge?.source)?.data.label;
-    return label !== Operator.Agent;
-  }, [edges, getNode, id]);
+  const isHeadAgent = useMemo(() => {
+    return !isBottomSubAgent(edges, id);
+  }, [edges, id]);
 
   return (
     <ToolBar selected={selected} id={id} label={data.label}>
       <NodeWrapper>
-        {isNotParentAgent && (
+        {isHeadAgent && (
           <>
             <CommonHandle
               type="target"

web/src/pages/agent/constant.tsx

@@ -3058,3 +3058,8 @@ export enum NodeHandleId {
   AgentTop = 'agentTop',
   AgentBottom = 'agentBottom',
 }
+
+export enum VariableType {
+  String = 'string',
+  Array = 'array',
+}

web/src/pages/agent/form-sheet/use-form-config-map.tsx

@@ -1,5 +1,6 @@
 import { LlmSettingSchema } from '@/components/llm-setting-items/next';
 import { CodeTemplateStrMap, ProgrammingLanguage } from '@/constants/agent';
+import IterationStartForm from '@/pages/flow/form/iteration-start-from';
 import { useTranslation } from 'react-i18next';
 import { z } from 'zod';
 import { Operator } from '../constant';
@@ -367,7 +368,7 @@ export function useFormConfigMap() {
       schema: z.object({}),
     },
     [Operator.IterationStart]: {
-      component: () => <></>,
+      component: IterationStartForm,
       defaultValues: {},
       schema: z.object({}),
     },

web/src/pages/agent/form/begin-form/use-edit-query.ts

@@ -1,9 +1,12 @@
 import { useSetModalState } from '@/hooks/common-hooks';
 import { useSetSelectedRecord } from '@/hooks/logic-hooks';
 import { useCallback, useMemo, useState } from 'react';
+import { UseFormReturn } from 'react-hook-form';
 import { BeginQuery, INextOperatorForm } from '../../interface';
 
-export const useEditQueryRecord = ({ form, node }: INextOperatorForm) => {
+export const useEditQueryRecord = ({
+  form,
+}: INextOperatorForm & { form: UseFormReturn }) => {
   const { setRecord, currentRecord } = useSetSelectedRecord<BeginQuery>();
   const { visible, hideModal, showModal } = useSetModalState();
   const [index, setIndex] = useState(-1);
@@ -16,7 +19,6 @@ export const useEditQueryRecord = ({ form, node }: INextOperatorForm) => {
   const handleEditRecord = useCallback(
     (record: BeginQuery) => {
       const inputs: BeginQuery[] = form?.getValues('inputs') || [];
-      console.log('🚀 ~ useEditQueryRecord ~ inputs:', inputs);
 
       const nextQuery: BeginQuery[] =
         index > -1 ? inputs.toSpliced(index, 1, record) : [...inputs, record];

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

@@ -6,18 +6,29 @@ import {
   FormLabel,
   FormMessage,
 } from '@/components/ui/form';
+import { useMemo } from 'react';
 import { useFormContext } from 'react-hook-form';
 import { useTranslation } from 'react-i18next';
+import { VariableType } from '../../constant';
 import { useBuildQueryVariableOptions } from '../../hooks/use-get-begin-query';
 
-type QueryVariableProps = { name?: string };
+type QueryVariableProps = { name?: string; type?: VariableType };
 
-export function QueryVariable({ name = 'query' }: QueryVariableProps) {
+export function QueryVariable({
+  name = 'query',
+  type = VariableType.String,
+}: QueryVariableProps) {
   const { t } = useTranslation();
   const form = useFormContext();
 
   const nextOptions = useBuildQueryVariableOptions();
 
+  const finalOptions = useMemo(() => {
+    return nextOptions.map((x) => {
+      return { ...x, options: x.options.filter((y) => y.type === type) };
+    });
+  }, [nextOptions, type]);
+
   return (
     <FormField
       control={form.control}
@@ -27,7 +38,7 @@ export function QueryVariable({ name = 'query' }: QueryVariableProps) {
           <FormLabel tooltip={t('chat.modelTip')}>{t('flow.query')}</FormLabel>
           <FormControl>
             <SelectWithSearch
-              options={nextOptions}
+              options={finalOptions}
               {...field}
             ></SelectWithSearch>
           </FormControl>

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

@@ -4,7 +4,7 @@ import { zodResolver } from '@hookform/resolvers/zod';
 import { useMemo } from 'react';
 import { useForm } from 'react-hook-form';
 import { z } from 'zod';
-import { initialRetrievalValues } from '../../constant';
+import { initialRetrievalValues, VariableType } from '../../constant';
 import { useWatchFormChange } from '../../hooks/use-watch-form-change';
 import { INextOperatorForm } from '../../interface';
 import { Output } from '../components/output';
@@ -50,7 +50,10 @@ const IterationForm = ({ node }: INextOperatorForm) => {
         }}
       >
         <FormContainer>
-          <QueryVariable name="items_ref"></QueryVariable>
+          <QueryVariable
+            name="items_ref"
+            type={VariableType.Array}
+          ></QueryVariable>
         </FormContainer>
         <Output list={outputList}></Output>
       </form>

web/src/pages/agent/hooks/use-add-node.ts

@@ -232,8 +232,15 @@ function useAddToolNode() {
   return { addToolNode };
 }
 
+function isBottomSubAgent(type: string, position: Position) {
+  return (
+    (type === Operator.Agent && position === Position.Bottom) ||
+    type === Operator.Tool
+  );
+}
+
 export function useAddNode(reactFlowInstance?: ReactFlowInstance<any, any>) {
-  const { edges, nodes, addEdge, addNode, getNode } = useGraphStore(
+  const { edges, nodes, addEdge, addNode, getNode, updateNode } = useGraphStore(
     (state) => state,
   );
   const getNodeName = useGetNodeName();
@@ -291,6 +298,18 @@ export function useAddNode(reactFlowInstance?: ReactFlowInstance<any, any>) {
         if (node && node.parentId) {
           newNode.parentId = node.parentId;
           newNode.extent = 'parent';
+          const parentNode = getNode(node.parentId);
+          if (parentNode && !isBottomSubAgent(type, params.position)) {
+            const MoveRightDistance = 310;
+            updateNode({
+              ...parentNode,
+              width: (parentNode.width || 0) + MoveRightDistance,
+              position: {
+                x: parentNode.position.x + MoveRightDistance / 2,
+                y: parentNode.position.y,
+              },
+            });
+          }
         }
 
         if (type === Operator.Iteration) {
@@ -377,6 +396,7 @@ export function useAddNode(reactFlowInstance?: ReactFlowInstance<any, any>) {
       initializeOperatorParams,
       nodes,
       reactFlowInstance,
+      updateNode,
     ],
   );
 

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

@@ -5,7 +5,7 @@ import { DefaultOptionType } from 'antd/es/select';
 import { isEmpty } from 'lodash';
 import get from 'lodash/get';
 import { useCallback, useContext, useEffect, useMemo, useState } from 'react';
-import { BeginId, Operator } from '../constant';
+import { BeginId, BeginQueryType, Operator, VariableType } from '../constant';
 import { AgentFormContext } from '../context';
 import { buildBeginInputListFromObject } from '../form/begin-form/utils';
 import { BeginQuery } from '../interface';
@@ -65,6 +65,7 @@ function buildOutputOptions(
   return Object.keys(outputs).map((x) => ({
     label: x,
     value: `${nodeId}@${x}`,
+    type: outputs[x]?.type,
   }));
 }
 
@@ -104,6 +105,19 @@ const ExcludedNodes = [
   Operator.Note,
 ];
 
+const StringList = [
+  BeginQueryType.Line,
+  BeginQueryType.Paragraph,
+  BeginQueryType.Options,
+];
+
+function transferToVariableType(type: string) {
+  if (StringList.some((x) => x === type)) {
+    return VariableType.String;
+  }
+  return type;
+}
+
 export function useBuildBeginVariableOptions() {
   const getBeginNodeDataQuery = useGetBeginNodeDataQuery();
 
@@ -116,6 +130,7 @@ export function useBuildBeginVariableOptions() {
         options: query.map((x) => ({
           label: x.name,
           value: `begin@${x.key}`,
+          type: transferToVariableType(x.type),
         })),
       },
     ];
@@ -141,9 +156,11 @@ export function useBuildQueryVariableOptions() {
   const options = useBuildVariableOptions(node?.id);
 
   const nextOptions = useMemo(() => {
-    const globalOptions = Object.keys(data?.dsl?.globals ?? {}).map((x) => ({
-      label: x,
-      value: x,
+    const globals = data?.dsl?.globals ?? {};
+    const globalOptions = Object.entries(globals).map(([key, value]) => ({
+      label: key,
+      value: key,
+      type: Array.isArray(value) ? VariableType.Array : typeof value,
     }));
     return [
       { ...options[0], options: [...options[0]?.options, ...globalOptions] },

web/src/pages/agent/hooks/use-show-drawer.tsx

@@ -57,7 +57,7 @@ export const useShowSingleDebugDrawer = () => {
   };
 };
 
-const ExcludedNodes = [Operator.IterationStart, Operator.Note];
+const ExcludedNodes = [Operator.Note];
 
 export function useShowDrawer({
   drawerVisible,

web/src/pages/agent/store.ts

@@ -56,6 +56,7 @@ export type RFState = {
   onSelectionChange: OnSelectionChangeFunc;
   addNode: (nodes: RAGFlowNodeType) => void;
   getNode: (id?: string | null) => RAGFlowNodeType | undefined;
+  updateNode: (node: RAGFlowNodeType) => void;
   addEdge: (connection: Connection) => void;
   getEdge: (id: string) => Edge | undefined;
   updateFormDataOnConnect: (connection: Connection) => void;
@@ -192,6 +193,16 @@ const useGraphStore = create<RFState>()(
       addNode: (node: RAGFlowNodeType) => {
         set({ nodes: get().nodes.concat(node) });
       },
+      updateNode: (node) => {
+        const { nodes } = get();
+        const nextNodes = nodes.map((x) => {
+          if (x.id === node.id) {
+            return node;
+          }
+          return x;
+        });
+        set({ nodes: nextNodes });
+      },
       getNode: (id?: string | null) => {
         return get().nodes.find((x) => x.id === id);
       },

web/src/pages/flow/constant.tsx

@@ -639,7 +639,16 @@ export const initialEmailValues = {
 export const initialIterationValues = {
   delimiter: ',',
 };
-export const initialIterationStartValues = {};
+export const initialIterationStartValues = {
+  outputs: {
+    item: {
+      type: 'unkown',
+    },
+    index: {
+      type: 'integer',
+    },
+  },
+};
 
 export const initialCodeValues = {
   lang: 'python',

web/src/pages/flow/form/iteration-start-from/index.tsx

@@ -0,0 +1,22 @@
+import { Output, OutputType } from '@/pages/agent/form/components/output';
+import { initialIterationStartValues } from '../../constant';
+
+const outputs = initialIterationStartValues.outputs;
+
+const outputList = Object.entries(outputs).reduce<OutputType[]>(
+  (pre, [key, value]) => {
+    pre.push({ title: key, type: value.type });
+
+    return pre;
+  },
+  [],
+);
+const IterationStartForm = () => {
+  return (
+    <section className="space-y-6 p-4">
+      <Output list={outputList}></Output>
+    </section>
+  );
+};
+
+export default IterationStartForm;