feat: add support for multi-tools

fishonamos · fishonamos · commit be42c60c1343 · 2025-02-18T09:52:48.000+01:00
diff --git a/packages/sui-agent/src/@types/interface.ts b/packages/sui-agent/src/@types/interface.ts
@@ -1,7 +1,7 @@
 // Response interface for the intent agent's operations
 export interface IntentAgentResponse {
   success: boolean; // Indicates if the operation was successful
-  selected_tool: null | string; // Name of the tool that was selected for the operation
+  selected_tools: null | string[]; // Name of the tool that was selected for the operation
   response: null | string; // Response message from the operation
   needs_additional_info: boolean; // Indicates if more information is needed
   additional_info_required: null | string[]; // List of additional information fields needed
@@ -13,7 +13,7 @@ export type ToolArgument = string | number | boolean | bigint;
 // Response interface for tool operations (similar to IntentAgentResponse)
 export interface toolResponse {
   success: boolean;
-  selected_tool: null | string;
+  selected_tools: null | string;
   response: null | string;
   needs_additional_info: boolean;
   additional_info_required: null | string[];
diff --git a/packages/sui-agent/src/agents/SuiAgent.ts b/packages/sui-agent/src/agents/SuiAgent.ts
@@ -5,6 +5,7 @@ import Utils from '../utils';
 import intent_agent_prompt from '../prompts/intent_agent_prompt';
 import final_answer_agent_prompt from '../prompts/final_answer_agent';
 import Atoma from '../config/atoma';
+import decomposerPrompt from '../prompts/decomposer';
 
 /**
  * Main agent class that handles intent processing and decision making
@@ -28,18 +29,25 @@ class Agents {
     registerAllTools(this.tools);
   }
 
+  async QueryDecomposer(prompt: string) {
+    return await this.AtomaClass.atomaChat([
+      { content: decomposerPrompt, role: 'assistant' },
+      { content: prompt, role: 'user' },
+    ]);
+  }
+
   /**
    * Processes initial user intent and selects appropriate tools
    * @param prompt - User's input query
    * @returns IntentAgentResponse containing tool selection and processing details
    */
-  async IntentAgent(prompt: string, address?: string) {
-    const IntentResponse: IntentAgentResponse =
+  async IntentAgent(subqueries: string[], address?: string) {
+    const IntentResponse: IntentAgentResponse[] =
       (await this.tools.selectAppropriateTool(
         this.AtomaClass,
-        prompt,
+        `${subqueries}`,
         address,
-      )) as IntentAgentResponse;
+      )) as IntentAgentResponse[];
 
     return IntentResponse;
   }
@@ -51,15 +59,15 @@ class Agents {
    * @returns Processed response after decision making
    */
   async DecisionMakingAgent(
-    intentResponse: IntentAgentResponse,
+    intentResponse: IntentAgentResponse[],
     query: string,
   ) {
     // Pass both the selected tool name and arguments to processQuery
+
     return await this.utils.processQuery(
       this.AtomaClass,
       query,
-      intentResponse.selected_tool,
-      intentResponse.tool_arguments,
+      intentResponse,
     );
   }
 
@@ -71,8 +79,13 @@ class Agents {
    */
   async SuperVisorAgent(prompt: string, walletAddress?: string) {
     // Process intent
-    const res = await this.IntentAgent(prompt, walletAddress);
-
+    const decomposer = await this.QueryDecomposer(prompt);
+    const decomposed: string[] = JSON.parse(
+      decomposer.choices[0].message.content,
+    );
+    console.log(decomposed);
+    const res = await this.IntentAgent(decomposed, walletAddress);
+    console.log(res, 'this is intent agent response');
     // Make decision based on intent
     const finalAnswer = await this.DecisionMakingAgent(res, prompt);
     console.log('Final Answer:', finalAnswer);
diff --git a/packages/sui-agent/src/prompts/decomposer.ts b/packages/sui-agent/src/prompts/decomposer.ts
@@ -0,0 +1,16 @@
+export default "You are the Query Decomposer.\n\n\
+\n\
+Your task is to analyze the user's query and break it into multiple subqueries **only if necessary**, following strict rules.\n\n\
+\n\
+### **Rules for Decomposition:**\n\
+1️ **Determine if decomposition is needed**\n\
+   - If the query requires multiple tools or separate logical steps, split it into subqueries.\n\
+   - If a single tool can handle the query, return it as is.\n\n\
+\n\
+2️ **Subquery Format (Strict JSON Array)**\n\
+   - Each subquery must be **clear, self-contained, and executable**.\n\
+   - Maintain **logical order** for execution.\n\
+\n\
+### **Output Format:** DO NOT STRAY FROM THE RESPONSE FORMAT (Array or single string) RETURN ONLY THE RESPONSE FORMAT \n\
+- If decomposition **is needed**, return an array of strings: ELSE return an array with a single string\n\
+  ";
diff --git a/packages/sui-agent/src/prompts/intent_agent_prompt.ts b/packages/sui-agent/src/prompts/intent_agent_prompt.ts
@@ -1,52 +1,85 @@
-/**
- * Prompt template for the intent agent (CoinSage)
- *
- * This template defines the behavior and response format for the AI assistant
- * specifically focused on Sui blockchain related queries. The agent follows
- * a structured decision-making process to handle user queries:
- *
- * 1. Self-Assessment: Determines if direct response is possible
- * 2. Tool Selection: Chooses appropriate tool if direct response isn't possible
- * 3. Information Gathering: Identifies when additional context is needed
- *
- * Response Format:
- * {
- *   success: boolean,        // true if query can be handled without additional info
- *   selected_tool: string,   // name of chosen tool, or null if not needed
- *   response: string,        // direct response if available, null otherwise
- *   needs_additional_info: boolean,  // true if more context is required
- *   additional_info_required: string[], // list of required information
- *   tool_arguments: any[]    // arguments needed for the selected tool
- * }
- */
 const intent_query = `You are an intelligent assistant called AtomaSage.
     YOUR NAME IS ATOMASAGE
 
-You are tasked with ONLY answering questions that could be related to the sui blockchain (directly or indirectly). Follow these steps to respond effectively:
+You are tasked with ONLY answering questions that could be related to the Sui blockchain (directly or indirectly). Follow these steps to respond effectively:
 
-Self-Assessment: First, determine if you can answer the user's question directly based on your current knowledge and capabilities.
+---
 
-Tool Selection: If you cannot answer the question directly, review the following list of tools:
+### Step 1: **Self-Assessment**  
+- Determine if you can answer each user query directly based on your knowledge.  
+- If yes, generate a response.  
+
+### Step 2: **Tool Selection**  
+- If a query **cannot** be answered directly, review the available tools:  
 
 \${toolsList}
 
-Select the most appropriate tool that can help answer the user's query. Clearly specify the chosen tool.
+- **Each subquery MUST have its own separate tools and tool arguments.**  
+- Do NOT share tools between subqueries unless absolutely necessary.  
+
+### Step 3: **Needs Info Handling**  
+- If neither your knowledge nor the available tools can fully answer a query, respond with a "Needs Info Error" and specify what additional information is required.  
+
+---
+
+### **Response Format (Each Subquery Must Be Separate)**  
+Respond with an **array** where each entry corresponds to a single subquery. Each subquery MUST have its own selected tools and tool arguments.  
+YOU ARE ONLY ALLOWED TO RESPOND WITH THE JSON FORMAT BELOW, NEVER RETURN A SINGLE STRING
+\`\`\`json
+[
+  {
+    "subquery": string, // The specific subquery being answered
+    "success": boolean, // Set to true ONLY if needs_additional_info is false
+    "selected_tools": null | string[], // array of single element, selected tool for that subquery
+    "response": null | string, // Direct response if available
+    "needs_additional_info": boolean, // True if more context is required
+    "additional_info_required": null | string[], // List of required information
+    "tool_arguments": null | string[] // Arguments for selected tools (specific to this subquery)
+  }
+]
+\`\`\`
+
+---
+
+### **Important Rules** 
+
+- **Each subquery must be handled separately.**  
+- **Tools must be specific to each subquery.**  
+- **DO NOT** combine tools across multiple subqueries unless required.  
+- **DO NOT** combine tool arguments, only add tool arguments relevant to that query.
+- **DO NOT** select unnecessary tools.  
+- **DO NOT** deviate from the response format.  
+
+---
 
-Needs Info Error: If neither your knowledge nor the available tools can provide a solution, respond with a "Needs Info Error," indicating that additional information or context is required to proceed.
+### **Example Response**
+#### User Query:  
+*"What is Move, and can you check the latest transactions on Sui?"*  
 
-This is the response format 
-[{
-"success":boolean,//set it to true, only if needs_additional info is false 
-"selected_tool":null | string,
-"response":null | string,
-"needs_additional_info": boolean,
-"additional_info_required": null | string array
-"tool_arguments": null | array  //based on user query .and tool input, 
-}] 
+#### Correct Response:  
+\`\`\`json
+[
+  {
+    "subquery": "What is Move?",
+    "success": true,
+    "selected_tools": null,
+    "response": "Move is a programming language designed for the Sui blockchain...",
+    "needs_additional_info": false,
+    "additional_info_required": null,
+    "tool_arguments": null
+  },
+  {
+    "subquery": "Can you check the latest transactions on Sui?",
+    "success": false,
+    "selected_tools": ["SuiExplorerAPI"],
+    "response": null,
+    "needs_additional_info": false,
+    "additional_info_required": null,
+    "tool_arguments": [5]
+  }
+]
+\`\`\`
 
-DO NOT UNDER ANY CIRCUMSTANCES STRAY FROM THE RESONSE FORMAT
-RESPOND WITH ONLY THE RESONSE FORMAT
-ADD PARAMS IN THE EXACT ORDER, PUT JUST THE VALUE IN THE ARRAY
 `;
 
 export default intent_query;
diff --git a/packages/sui-agent/src/utils/index.ts b/packages/sui-agent/src/utils/index.ts
@@ -2,7 +2,7 @@ import { randomUUID } from 'crypto';
 import { AtomaSDK } from 'atoma-sdk';
 import Atoma from '../config/atoma';
 import Tools from './tools';
-import { ToolArgument } from '../@types/interface';
+import { IntentAgentResponse, ToolArgument } from '../@types/interface';
 
 /**
  * Utility class for processing agent responses and making decisions
@@ -38,19 +38,49 @@ class Utils {
   async processQuery(
     AtomaInstance: Atoma,
     query: string,
-    selectedTool: string | null,
-    toolArguments: ToolArgument[] = [],
+    intentResponses: IntentAgentResponse[],
   ) {
     try {
-      if (!selectedTool) {
+      if (!intentResponses || intentResponses.length === 0) {
         return this.finalAnswer(
           AtomaInstance,
-          'No tool selected for the query',
+          'No tools selected for the query',
           query,
         );
       }
 
-      return this.executeTools(selectedTool, toolArguments, AtomaInstance);
+      let aggregatedResults = '';
+
+      for (const response of intentResponses) {
+        const { selected_tools, tool_arguments } = response;
+
+        if (!selected_tools?.length) {
+          continue; // Skip if no tool selected
+        }
+        console.log(selected_tools, 'selected tools /...');
+        // Execute the tool and append its result
+        const result = await this.executeTools(selected_tools, tool_arguments);
+
+        // Aggregate results (you might want to customize this based on your needs)
+        aggregatedResults += result + '\n';
+      }
+
+      // If no tools were successfully executed
+      if (!aggregatedResults) {
+        return this.finalAnswer(
+          AtomaInstance,
+          'No valid tools were executed for the query',
+          query,
+        );
+      }
+
+      // Return final answer with aggregated results
+      return this.finalAnswer(
+        AtomaInstance,
+        aggregatedResults.trim(),
+        query,
+        intentResponses.map((r) => r.selected_tools).join(', '),
+      );
     } catch (error: unknown) {
       console.error('Error processing query:', error);
       return handleError(error, {
@@ -61,14 +91,28 @@ class Utils {
     }
   }
 
-  /**
-   * Format final answer
-   * @param response - Raw response
-   * @param query - Original query
-   * @param tools - Tools used
-   * @returns Formatted response
-   * @private
-   */
+  private async executeTools(
+    selected_tool: string[],
+    args: ToolArgument[] | null,
+  ) {
+    const tool = this.tools
+      .getAllTools()
+      .find((t) => t.name.trim() === selected_tool[0]);
+
+    if (!tool) {
+      throw new Error(`Tool ${selected_tool} not found`);
+    }
+
+    try {
+      const toolArgs = args || [];
+      const result = await tool.process(...toolArgs);
+      return result;
+    } catch (error: unknown) {
+      console.error('Error executing tool:', error);
+      throw error; // Let the main processQuery handle the error
+    }
+  }
+
   private async finalAnswer(
     AtomaInstance: Atoma,
     response: string,
@@ -80,62 +124,15 @@ class Utils {
       .replace('${response}', response)
       .replace('tools', `${tools || null}`);
 
-    // const finalAns = await new AtomaSDK({bearerAuth:'bearer auth here'}).chat.create({
-    //   messages: [
-    //     {role:"assistant",content:finalPrompt},
-    //   //  { role: "user", content: query }
-    //   ],
-    //   model: "meta-llama/Llama-3.3-70B-Instruct"
-    // });
-
     const finalAns = await AtomaInstance.atomaChat([
       { role: 'assistant', content: finalPrompt },
       { role: 'user', content: query },
     ]);
-    console.log('new one');
-
-    // const finalAns = await atomaChat(this.sdk, [
-    //   {
-    //     content: finalPrompt,
-    //     role: 'assistant',
-    //   },
-    // ]);
-    const res = finalAns.choices[0].message.content;
-    console.log(finalPrompt);
-    return JSON.parse(res);
-  }
-
-  /**
-   * Executes selected tools with provided arguments
-   * @param selected_tool - Name of the tool to execute
-   * @param args - Arguments to pass to the tool
-   * @returns Processed tool response
-   * @private
-   */
-  private async executeTools(
-    selected_tool: string,
-    args: ToolArgument[] | null,
-    AtomaInstance: Atoma,
-  ) {
-    const tool = this.tools.getAllTools().find((t) => t.name === selected_tool);
-    console.log('Selected tool:', selected_tool);
-    console.log('Tool arguments:', args);
 
-    if (!tool) {
-      throw new Error(`Tool ${selected_tool} not found`);
-    }
-
-    try {
-      const toolArgs = args || [];
-      const result = await tool.process(...toolArgs);
-      return await this.finalAnswer(AtomaInstance, result, '', selected_tool);
-    } catch (error: unknown) {
-      console.error('Error executing tool:', error);
-      return handleError(error, {
-        reasoning: `The system encountered an issue while executing the tool ${selected_tool}`,
-        query: `Attempted to execute ${selected_tool} with arguments: ${JSON.stringify(args)}`,
-      });
-    }
+    const res = finalAns.choices[0].message.content;
+    const parsedRes = JSON.parse(res);
+    console.log(parsedRes, 'parsed response');
+    return parsedRes;
   }
 }
 
diff --git a/packages/sui-agent/src/utils/tools.ts b/packages/sui-agent/src/utils/tools.ts
