Skip to main content

abu_agent/toolbox/
mod.rs

1pub mod tools;
2pub mod mcp;
3use std::{ffi::OsStr, path::Path};
4use abu_base::chat::ToolDefinition;
5use abu_mcp::McpTool;
6use abu_tool::{Tool, ToolCallResult, ToolError, ToolRegister};
7use mcp::McpManager;
8use tracing::debug;
9
10use crate::AgentResult;
11
12pub struct ToolBox {
13    tools: ToolRegister,
14    mcp_manager: McpManager,
15    tool_definitions: Vec<ToolDefinition>,
16}
17
18impl ToolBox {
19    pub fn new() -> Self {
20        Self {
21            tools: ToolRegister::new(),
22            mcp_manager: McpManager::new(),
23            tool_definitions: vec![],
24        }
25    }
26
27    pub async fn load_mcpconfig(&mut self, path: impl AsRef<Path>) -> AgentResult<()> {
28        self.mcp_manager = McpManager::load_config(path).await?; 
29        Ok(())
30    }
31
32    pub fn add_tool<T: Tool + 'static>(&mut self, tool: T) {
33        debug!("add tool '{}'", tool.name());
34        self.tool_definitions.push(tool.to_function_define());
35        self.tools.add_tool(tool);
36    } 
37
38    pub fn add_tool_box(&mut self, tool: Box<dyn Tool>) {
39        debug!("add tool '{}'", tool.name());
40        self.tool_definitions.push(tool.to_function_define());
41        self.tools.add_tool_box(tool);
42    }
43
44    pub async fn add_mcp_server<I, S>(&mut self, cmd: S, args: I) -> AgentResult<()> 
45    where 
46        I: IntoIterator<Item = S>,
47        S: AsRef<OsStr>,
48    {
49        debug!("add mcp server");
50        let client = self.mcp_manager.add_stdio_server(cmd, args).await?;
51        for mcp_tool in client.server_tools.iter() {
52            self.tool_definitions.push(mcp_tool_to_tool_defintion(mcp_tool));
53        }
54        Ok(())
55    }
56
57    pub async fn execute_tool(&mut self, name: &str, arguments: &str) -> AgentResult<ToolCallResult> {
58        let arguments = match serde_json::from_str(arguments) {
59            Err(e) => return Ok(ToolCallResult::error(e.to_string())),
60            Ok(v) => v,
61        };
62
63        if self.tools.has_tool(&name) {
64            let result = self.tools.execute(name, arguments).await?;
65            Ok(result)
66        } else if self.mcp_manager.has_tool(&name) {
67            let result = self.mcp_manager.execute_toolcall(name, arguments).await?;
68            Ok(result)
69        } else {
70            Err(ToolError::ToolNotFound(name.to_string()))?
71        }
72    }
73
74    pub fn tool_definitions(&self) -> &[ToolDefinition] {
75        &self.tool_definitions
76    }
77}
78
79
80fn mcp_tool_to_tool_defintion(mcp_tool: &McpTool) -> ToolDefinition {
81    ToolDefinition {
82        name: mcp_tool.name.clone(),
83        description: mcp_tool.description.clone().unwrap_or_default(),
84        schema: serde_json::json!({
85            "type": "object",
86            "properties": mcp_tool.input_schema.properties.clone().unwrap_or(serde_json::json!({})),
87            "required": mcp_tool.input_schema.required.clone().unwrap_or(serde_json::json!([])),
88        })
89    }
90}