feat: supports multiple retrieval tool under an agent (#12046)

### What problem does this PR solve?

Add support for multiple Retrieval tools under an agent

### Type of change

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

Author: JimmyBenKlieve

web/src/components/large-model-form-field.tsx

@@ -68,7 +68,7 @@ export function LargeModelFormField({
                   <FormItem>
                     <FormControl>
                       <DropdownMenu>
-                        <DropdownMenuTrigger>
+                        <DropdownMenuTrigger asChild>
                           <Button variant={'ghost'}>
                             <Funnel className="text-text-disabled" />
                           </Button>

web/src/interfaces/database/agent.ts

@@ -170,6 +170,7 @@ export interface IAgentForm {
   tools: Array<{
     name: string;
     component_name: string;
+    id: string;
     params: Record<string, any>;
   }>;
   mcp: Array<{

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

@@ -2,7 +2,7 @@ import { NodeCollapsible } from '@/components/collapse';
 import { IAgentForm, IToolNode } from '@/interfaces/database/agent';
 import { Handle, NodeProps, Position } from '@xyflow/react';
 import { get } from 'lodash';
-import { MouseEventHandler, memo, useCallback } from 'react';
+import { memo } from 'react';
 import { NodeHandleId, Operator } from '../../constant';
 import { ToolCard } from '../../form/agent-form/agent-tools';
 import { useFindMcpById } from '../../hooks/use-find-mcp-by-id';
@@ -15,22 +15,11 @@ function InnerToolNode({
   isConnectable = true,
   selected,
 }: NodeProps<IToolNode>) {
-  const { edges, getNode } = useGraphStore((state) => state);
+  const { edges, getNode, setClickedToolId } = useGraphStore();
   const upstreamAgentNodeId = edges.find((x) => x.target === id)?.source;
   const upstreamAgentNode = getNode(upstreamAgentNodeId);
   const { findMcpById } = useFindMcpById();
 
-  const handleClick = useCallback(
-    (operator: string): MouseEventHandler<HTMLLIElement> =>
-      (e) => {
-        if (operator === Operator.Code) {
-          e.preventDefault();
-          e.stopPropagation();
-        }
-      },
-    [],
-  );
-
   const tools: IAgentForm['tools'] = get(
     upstreamAgentNode,
     'data.form.tools',
@@ -51,36 +40,53 @@ function InnerToolNode({
         position={Position.Top}
         isConnectable={isConnectable}
         className="!bg-accent-primary !size-2"
-      ></Handle>
+      />
+
       <NodeCollapsible items={[tools, mcpList]}>
         {(x) => {
-          if ('mcp_id' in x) {
+          if (Reflect.has(x, 'mcp_id')) {
             const mcp = x as unknown as IAgentForm['mcp'][number];
+
             return (
               <ToolCard
                 key={mcp.mcp_id}
-                onClick={handleClick(mcp.mcp_id)}
+                onClick={(e) => {
+                  if (mcp.mcp_id === Operator.Code) {
+                    e.preventDefault();
+                    e.stopPropagation();
+                  }
+                }}
                 className="cursor-pointer"
-                data-tool={x.mcp_id}
+                data-tool={mcp.mcp_id}
               >
                 {findMcpById(mcp.mcp_id)?.name}
               </ToolCard>
             );
           }
 
           const tool = x as unknown as IAgentForm['tools'][number];
+
           return (
             <ToolCard
-              key={tool.component_name}
-              onClick={handleClick(tool.component_name)}
+              key={tool.id}
+              onClick={(e) => {
+                if (tool.component_name === Operator.Code) {
+                  e.preventDefault();
+                  e.stopPropagation();
+                }
+
+                setClickedToolId(tool.id || tool.component_name);
+              }}
               className="cursor-pointer"
               data-tool={tool.component_name}
+              data-tool-id={tool.id}
             >
               <div className="flex gap-1 items-center pointer-events-none">
-                <OperatorIcon
-                  name={tool.component_name as Operator}
-                ></OperatorIcon>
-                {tool.component_name}
+                <OperatorIcon name={tool.component_name as Operator} />
+
+                {tool.component_name === Operator.Retrieval
+                  ? tool.name
+                  : tool.component_name}
               </div>
             </ToolCard>
           );

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

@@ -1,3 +1,4 @@
+import { Button, ButtonProps } from '@/components/ui/button';
 import {
   TooltipContent,
   TooltipNode,
@@ -6,31 +7,24 @@ import {
 import { cn } from '@/lib/utils';
 import { Position } from '@xyflow/react';
 import { Copy, Play, Trash2 } from 'lucide-react';
-import {
-  HTMLAttributes,
-  MouseEventHandler,
-  PropsWithChildren,
-  useCallback,
-} from 'react';
+import { MouseEventHandler, PropsWithChildren, useCallback } from 'react';
 import { Operator } from '../../constant';
 import { useDuplicateNode } from '../../hooks';
 import useGraphStore from '../../store';
 
-function IconWrapper({
-  children,
-  className,
-  ...props
-}: HTMLAttributes<HTMLDivElement>) {
+function IconWrapper({ children, className, ...props }: ButtonProps) {
   return (
-    <div
+    <Button
+      variant="secondary"
+      size="icon"
       className={cn(
-        'p-1.5 bg-bg-component border border-border-button rounded-sm cursor-pointer hover:text-text-primary',
+        'size-7 p-0 bg-bg-component text-current hover:text-text-primary focus-visible:text-text-primary',
         className,
       )}
       {...props}
     >
       {children}
-    </div>
+    </Button>
   );
 }
 
@@ -55,7 +49,7 @@ export function ToolBar({
     (store) => store.deleteIterationNodeById,
   );
 
-  const deleteNode: MouseEventHandler<HTMLDivElement> = useCallback(
+  const deleteNode: MouseEventHandler<HTMLButtonElement> = useCallback(
     (e) => {
       e.stopPropagation();
       if ([Operator.Iteration, Operator.Loop].includes(label as Operator)) {
@@ -69,7 +63,7 @@ export function ToolBar({
 
   const duplicateNode = useDuplicateNode();
 
-  const handleDuplicate: MouseEventHandler<HTMLDivElement> = useCallback(
+  const handleDuplicate: MouseEventHandler<HTMLButtonElement> = useCallback(
     (e) => {
       e.stopPropagation();
       duplicateNode(id, label);
@@ -82,7 +76,7 @@ export function ToolBar({
       <TooltipTrigger className="h-full">{children}</TooltipTrigger>
 
       <TooltipContent position={Position.Top}>
-        <section className="flex gap-2 items-center text-text-secondary">
+        <section className="flex gap-2 items-center text-text-secondary pb-2">
           {showRun && (
             <IconWrapper>
               <Play className="size-3.5" data-play />
@@ -94,8 +88,8 @@ export function ToolBar({
             </IconWrapper>
           )}
           <IconWrapper
-            onClick={deleteNode}
             className="hover:text-state-error hover:border-state-error"
+            onClick={deleteNode}
           >
             <Trash2 className="size-3.5" />
           </IconWrapper>

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

@@ -778,6 +778,7 @@ export const NoDebugOperatorsList = [
   Operator.Splitter,
   Operator.HierarchicalMerger,
   Operator.Extractor,
+  Operator.Tool,
 ];
 
 export const NoCopyOperatorsList = [

web/src/pages/agent/form-sheet/next.tsx

@@ -1,3 +1,4 @@
+import { Button } from '@/components/ui/button';
 import {
   Sheet,
   SheetContent,
@@ -41,49 +42,69 @@ const FormSheet = ({
   showSingleDebugDrawer,
 }: IModalProps<any> & IProps) => {
   const operatorName: Operator = node?.data.label as Operator;
-  const clickedToolId = useGraphStore((state) => state.clickedToolId);
+  const { clickedToolId, getAgentToolById } = useGraphStore();
 
   const currentFormMap = FormConfigMap[operatorName];
-
   const OperatorForm = currentFormMap?.component ?? EmptyContent;
-
   const isMcp = useIsMcp(operatorName);
-
   const { t } = useTranslate('flow');
+  const { component_name: toolComponentName } = (getAgentToolById(
+    clickedToolId,
+  ) ?? {}) as {
+    component_name: Operator;
+    name: string;
+    id: string;
+  };
 
   return (
     <Sheet open={visible} modal={false}>
       <SheetContent
-        className={cn('top-20 p-0 flex flex-col pb-20', {
+        className={cn('top-20 p-0 flex flex-col pb-20 gap-0', {
           'right-[clamp(0px,34%,620px)]': chatVisible,
         })}
         closeIcon={false}
       >
         <SheetHeader>
           <SheetTitle className="hidden"></SheetTitle>
-          <section className="flex-col border-b py-2 px-5">
+          <section className="flex-col border-b pt-2 pb-4 px-5">
             <div className="flex items-center gap-2 pb-3">
-              <OperatorIcon name={operatorName}></OperatorIcon>
+              <OperatorIcon
+                name={toolComponentName || operatorName}
+              ></OperatorIcon>
               <TitleInput node={node}></TitleInput>
               {needsSingleStepDebugging(operatorName) && (
                 <RunTooltip>
-                  <CirclePlay
-                    className="size-3.5 cursor-pointer"
+                  <Button
+                    variant="ghost"
+                    size="icon"
+                    className="size-6 !p-0 bg-transparent"
                     onClick={showSingleDebugDrawer}
-                  />
+                  >
+                    <CirclePlay className="size-3.5 cursor-pointer" />
+                  </Button>
                 </RunTooltip>
               )}
-              <X onClick={hideModal} className="size-3.5 cursor-pointer" />
+
+              <Button
+                variant="ghost"
+                size="icon"
+                className="size-6 !p-0 bg-transparent"
+                onClick={hideModal}
+              >
+                <X className="size-3.5 cursor-pointer" />
+              </Button>
             </div>
-            {isMcp || (
-              <span className="text-text-secondary">
+
+            {!isMcp && (
+              <p className="text-text-secondary">
                 {t(
-                  `${lowerFirst(operatorName === Operator.Tool ? clickedToolId : operatorName)}Description`,
+                  `${lowerFirst(operatorName === Operator.Tool ? toolComponentName : operatorName)}Description`,
                 )}
-              </span>
+              </p>
             )}
           </section>
         </SheetHeader>
+
         <section className="pt-4 overflow-auto flex-1">
           {visible && (
             <AgentFormContext.Provider value={node}>

web/src/pages/agent/form-sheet/title-input.tsx

@@ -1,7 +1,8 @@
+import { Button } from '@/components/ui/button';
 import { Input } from '@/components/ui/input';
 import { RAGFlowNodeType } from '@/interfaces/database/agent';
 import { PenLine } from 'lucide-react';
-import { useCallback, useState } from 'react';
+import { useCallback, useLayoutEffect, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { BeginId, Operator } from '../constant';
 import { useHandleNodeNameChange } from '../hooks/use-change-node-name';
@@ -13,47 +14,75 @@ type TitleInputProps = {
 
 export function TitleInput({ node }: TitleInputProps) {
   const { t } = useTranslation();
+  const inputRef = useRef<HTMLInputElement>(null);
   const { name, handleNameBlur, handleNameChange } = useHandleNodeNameChange({
     id: node?.id,
     data: node?.data,
   });
 
   const operatorName: Operator = node?.data.label as Operator;
-
   const isMcp = useIsMcp(operatorName);
-
   const [isEditingMode, setIsEditingMode] = useState(false);
 
   const switchIsEditingMode = useCallback(() => {
     setIsEditingMode((prev) => !prev);
   }, []);
 
-  const handleBlur = useCallback(() => {
-    handleNameBlur();
-    setIsEditingMode(false);
-  }, [handleNameBlur]);
+  const handleBlur = useCallback(
+    (e: React.FocusEvent<HTMLInputElement>) => {
+      if (handleNameBlur()) {
+        setIsEditingMode(false);
+      } else {
+        // Re-focus the input if name doesn't change successfully
+        e.target.focus();
+        e.target.select();
+      }
+    },
+    [handleNameBlur],
+  );
+
+  useLayoutEffect(() => {
+    if (isEditingMode && inputRef.current) {
+      inputRef.current.focus();
+      inputRef.current.select();
+    }
+  }, [isEditingMode]);
 
   if (isMcp) {
     return <div className="flex-1 text-base">MCP Config</div>;
   }
 
   return (
-    <div className="flex items-center gap-1 flex-1">
+    // Give a fixed height to prevent layout shift when switching between edit and view modes
+    <div className="flex items-center gap-1 flex-1 h-8 mr-2">
       {node?.id === BeginId ? (
-        <span>{t(BeginId)}</span>
+        // Begin node is not editable
+        <span>{t(`flow.${BeginId}`)}</span>
       ) : isEditingMode ? (
         <Input
+          ref={inputRef}
           value={name}
           onBlur={handleBlur}
+          onKeyDown={(e) => {
+            // Support committing the value changes by pressing Enter
+            if (e.key === 'Enter') {
+              handleBlur(e as unknown as React.FocusEvent<HTMLInputElement>);
+            }
+          }}
           onChange={handleNameChange}
-        ></Input>
+        />
       ) : (
         <div className="flex items-center gap-2.5 text-base">
           {name}
-          <PenLine
+
+          <Button
+            variant="transparent"
+            size="icon"
+            className="size-6 !p-0 border-0 bg-transparent"
             onClick={switchIsEditingMode}
-            className="size-3.5 text-text-secondary cursor-pointer"
-          />
+          >
+            <PenLine className="size-3.5 text-text-secondary cursor-pointer" />
+          </Button>
         </div>
       )}
     </div>