Allow selecting several tools

Author: brichet

src/chat-handler.ts

@@ -120,12 +120,12 @@ export class ChatHandler extends AbstractChatModel {
   /**
    * Get/set a tool, which will build an agent.
    */
-  get tool(): Tool | null {
-    return this._tool;
+  get tools(): Tool[] {
+    return this._tools;
   }
-  set tool(value: Tool | null) {
-    this._tool = value;
-    this._providerRegistry.buildAgent(this._tool);
+  set tools(value: Tool[]) {
+    this._tools = value;
+    this._providerRegistry.buildAgent(this._tools);
   }
 
   /**
@@ -338,7 +338,7 @@ export class ChatHandler extends AbstractChatModel {
   private _defaultErrorMessage = 'AI provider not configured';
   private _controller: AbortController | null = null;
   private _useTool: boolean = false;
-  private _tool: Tool | null = null;
+  private _tools: Tool[] = [];
   private _toolRegistry?: IToolRegistry;
   private _useToolChanged = new Signal<ChatHandler, boolean>(this);
 }
@@ -382,11 +382,11 @@ export namespace ChatHandler {
     /**
      * Getter/setter of the tool to use.
      */
-    get tool(): Tool | null {
-      return (this._model as ChatHandler).tool;
+    get tools(): Tool[] {
+      return (this._model as ChatHandler).tools;
     }
-    set tool(value: Tool | null) {
-      (this._model as ChatHandler).tool = value;
+    set tools(value: Tool[]) {
+      (this._model as ChatHandler).tools = value;
     }
   }
 

src/components/tool-select.tsx

@@ -24,7 +24,7 @@ export function toolSelect(
   const toolRegistry = chatContext.toolsRegistry;
 
   const [useTool, setUseTool] = useState<boolean>(chatContext.useTool);
-  const [selectedTool, setSelectedTool] = useState<Tool | null>(null);
+  const [selectedTools, setSelectedTools] = useState<Tool[]>([]);
   const [tools, setTools] = useState<Tool[]>(toolRegistry?.tools || []);
   const [menuAnchorEl, setMenuAnchorEl] = useState<HTMLElement | null>(null);
   const [menuOpen, setMenuOpen] = useState(false);
@@ -38,13 +38,17 @@ export function toolSelect(
     setMenuOpen(false);
   }, []);
 
-  const onClick = useCallback(
-    (tool: Tool | null) => {
-      setSelectedTool(tool);
-      chatContext.tool = tool;
-    },
-    [props.model]
-  );
+  const onClick = (tool: Tool) => {
+    const currentTools = [...selectedTools];
+    const index = currentTools.indexOf(tool);
+    if (index !== -1) {
+      currentTools.splice(index, 1);
+    } else {
+      currentTools.push(tool);
+    }
+    setSelectedTools(currentTools);
+    chatContext.tools = currentTools;
+  };
 
   useEffect(() => {
     const updateTools = () => setTools(toolRegistry?.tools || []);
@@ -83,7 +87,7 @@ export function toolSelect(
           }
         }}
         sx={
-          selectedTool === null
+          selectedTools.length === 0
             ? { backgroundColor: 'var(--jp-layout-color3)' }
             : {}
         }
@@ -109,23 +113,6 @@ export function toolSelect(
           }
         }}
       >
-        <Tooltip title={'Do not use a tool'}>
-          <MenuItem
-            className={SELECT_ITEM_CLASS}
-            onClick={e => {
-              onClick(null);
-              // prevent sending second message with no selection
-              e.stopPropagation();
-            }}
-          >
-            {selectedTool === null ? (
-              <checkIcon.react className={'lm-Menu-itemIcon'} />
-            ) : (
-              <div className={'lm-Menu-itemIcon'} />
-            )}
-            <Typography display="block">No tool</Typography>
-          </MenuItem>
-        </Tooltip>
         {tools.map(tool => (
           <Tooltip title={tool.description}>
             <MenuItem
@@ -136,7 +123,7 @@ export function toolSelect(
                 e.stopPropagation();
               }}
             >
-              {selectedTool === tool ? (
+              {selectedTools.includes(tool) ? (
                 <checkIcon.react className={'lm-Menu-itemIcon'} />
               ) : (
                 <div className={'lm-Menu-itemIcon'} />

src/provider.ts

@@ -393,19 +393,19 @@ export class AIProviderRegistry implements IAIProviderRegistry {
   /**
    * Build an agent with a given tool.
    */
-  buildAgent(tool: Tool | null) {
-    if (tool !== null) {
+  buildAgent(tools: Tool[]) {
+    if (tools.length) {
       const chatModel = Private.getChatModel();
       if (chatModel === null) {
         Private.setAgent(null);
         return;
       }
-      chatModel.bindTools?.([tool], { tool_choice: tool.name });
+      chatModel.bindTools?.(tools);
       Private.setChatModel(chatModel);
       Private.setAgent(
         createReactAgent({
           llm: chatModel,
-          tools: [tool]
+          tools
         })
       );
     } else {

src/tokens.ts

@@ -144,7 +144,7 @@ export interface IAIProviderRegistry {
   /**
    * Build an agent with a given tool.
    */
-  buildAgent(tool: Tool | null): void;
+  buildAgent(tools: Tool[]): void;
   /**
    * A signal emitting when the provider or its settings has changed.
    */