Skip to main content

abu_agent/kit/
mod.rs

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