abu_agent/agent/
build.rs

1use std::{path::PathBuf, sync::Arc};
2use abu_provider::{deepseek::DeepSeek, ChatProvide};
3use abu_skill::SkillLoader;
4use abu_tool::Tool;
5use crate::{
6    context::ContextBuilder, hook::{Hook, HookManager}, memory::{Memory, SequentialMemory}, middleware::{LlmOutMiddleware, Middleware, MiddlewareManager, ToolCallMiddleware, ToolResultMiddleware}, model::{ChatConfig, ChatModel}, toolbox::tools::{bash::Bash, calculate::Calculator, fs::{FileCreator, FileReader, FileWriter}, skill::SkillTool}, AgentResult
7};
8use super::{Agent, AgentConfig, ToolBox};
9
10const DEFAULT_SYSTEM_PROMPT: &str = "You are an agent.";
11
12pub struct AgentBuilder<P: ChatProvide = DeepSeek, M: Memory = SequentialMemory> {
13    pub llm: ChatModel<P>,
14    pub config: AgentConfig,
15    pub memory: M,
16    pub system_prompt: String,
17    pub with_skills: Option<PathBuf>,
18    pub with_builtin_tools: bool,
19    pub with_subagent: bool,
20    pub tools: Vec<Box<dyn Tool>>,
21    pub mcpservers: Vec<(String, Vec<String>)>,
22    pub mcpconfig_path: Option<PathBuf>,
23    pub hooks: HookManager,
24    pub middlewares: MiddlewareManager,
25}
26
27impl Default for AgentConfig {
28    fn default() -> Self {
29        Self {
30            max_iteration: 10,
31            temperature: 0.7,
32        }
33    }
34}
35
36impl<C: ChatProvide, M: Memory> AgentBuilder<C, M> {
37    pub async fn build(mut self) -> AgentResult<Agent<C, M>> {
38        let mut toolbox = ToolBox::new();
39        let mut context_builder = ContextBuilder::new(self.system_prompt);
40
41        // tool
42        if self.with_builtin_tools {
43            toolbox.add_tool(Bash::new());
44            toolbox.add_tool(Calculator::new());
45            toolbox.add_tool(FileCreator::new());
46            toolbox.add_tool(FileWriter::new());
47            toolbox.add_tool(FileReader::new());
48        }
49        for tool in self.tools {
50            toolbox.add_tool_box(tool);
51        }
52
53        // mcp
54        if let Some(path) = self.mcpconfig_path {
55            toolbox.load_mcpconfig(&path).await?;
56        }
57        for (cmd, args) in self.mcpservers {
58            toolbox.add_mcp_server(&cmd, &args).await?;
59        }
60
61        // skill
62        if let Some(skill_dir) = self.with_skills {
63            let skill_loader = Arc::new(SkillLoader::load(skill_dir)?);
64            context_builder.with_skill(skill_loader.clone());
65            toolbox.add_tool(SkillTool::new(skill_loader));
66            
67        }
68
69        // llm init
70        self.llm.bind_tool_defines(toolbox.tool_definitions());
71        self.llm.set_config(ChatConfig { temperature: Some(self.config.temperature) });
72
73        Ok(Agent {
74            config: self.config,
75            llm: self.llm,
76            memory: self.memory,
77            toolbox,
78            context_builder,
79            hooks: self.hooks,
80            middlewares: self.middlewares,
81        })
82
83    }
84}
85
86impl<P: ChatProvide> AgentBuilder<P> {
87    pub fn new(llm: ChatModel<P>) -> Self {
88        Self {
89            llm,
90            config: AgentConfig::default(),
91            memory: SequentialMemory::default(),
92            system_prompt: DEFAULT_SYSTEM_PROMPT.to_string(),
93            with_skills: None,
94            with_builtin_tools: true,
95            with_subagent: false,
96            tools: vec![],
97            mcpservers: vec![],
98            mcpconfig_path: None,
99            hooks: HookManager::new(),
100            middlewares: MiddlewareManager::new(),
101        }
102    }
103}
104
105impl<C: ChatProvide, M: Memory> AgentBuilder<C, M> {
106    pub fn temperature(mut self, temperature: f64) -> Self {
107        self.config.temperature = temperature;
108        self
109    }
110
111    pub fn max_iteration(mut self, max_iteration: usize) -> Self {
112        self.config.max_iteration = max_iteration;
113        self
114    }
115
116    pub fn memory<NM: Memory>(self, memory: NM) -> AgentBuilder<C, NM> {
117        AgentBuilder {
118            memory,
119            llm: self.llm,
120            config: self.config,
121            system_prompt: self.system_prompt,
122            with_skills: self.with_skills,
123            with_builtin_tools: self.with_builtin_tools,
124            with_subagent: self.with_subagent,
125            tools: self.tools,
126            mcpservers: self.mcpservers,
127            mcpconfig_path: self.mcpconfig_path,
128            hooks: self.hooks,
129            middlewares: self.middlewares,
130        }
131    }
132
133    pub fn llm<NC: ChatProvide>(self, llm: ChatModel<NC>) -> AgentBuilder<NC, M> {
134        AgentBuilder {
135            memory: self.memory,
136            llm,
137            config: self.config,
138            system_prompt: self.system_prompt,
139            with_skills: self.with_skills,
140            with_builtin_tools: self.with_builtin_tools,
141            with_subagent: self.with_subagent,
142            tools: self.tools,
143            mcpservers: self.mcpservers,
144            mcpconfig_path: self.mcpconfig_path,
145            hooks: self.hooks,
146            middlewares: self.middlewares,
147        }
148    }
149
150    pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
151        self.system_prompt = system_prompt.into();
152        self
153    }
154
155    pub fn with_skills(mut self, skill_path: impl Into<PathBuf>) -> Self {
156        self.with_skills = Some(skill_path.into());
157        self
158    }
159
160    pub fn with_builtin_tools(mut self, enabled: bool) -> Self {
161        self.with_builtin_tools = enabled;
162        self
163    }
164
165    pub fn with_tool(mut self, tool: impl Tool + 'static) -> Self {
166        self.tools.push(Box::new(tool));
167        self
168    }
169
170    pub fn with_hook(mut self, hook: impl Hook + 'static) -> Self {
171        self.hooks.add_hook(hook);
172        self
173    }
174
175    pub fn with_middleware(mut self, middleware: impl Into<Middleware>) -> Self {
176        self.middlewares.add_middleware(middleware);
177        self
178    }
179
180    pub fn with_llm_out_middleware<LM: LlmOutMiddleware + 'static>(mut self, middleware: LM) -> Self {
181        self.middlewares.add_llm_out(middleware);
182        self
183    }
184
185    pub fn with_tool_call_middleware<TM: ToolCallMiddleware + 'static>(mut self, middleware: TM) -> Self {
186        self.middlewares.add_tool_call(middleware);
187        self
188    }
189
190    pub fn with_tool_result_middleware<TM: ToolResultMiddleware + 'static>(mut self, middleware: TM) -> Self {
191        self.middlewares.add_tool_result(middleware);
192        self
193    }
194
195    pub fn with_tools(mut self, tools: impl IntoIterator<Item = Box<dyn Tool>>) -> Self {
196        for tool in tools.into_iter() {
197            self.tools.push(tool);
198        }
199        self
200    }
201
202    pub fn with_mcpconfig(mut self, path: impl Into<PathBuf>) -> Self {
203        self.mcpconfig_path = Some(path.into());
204        self
205    }
206
207    pub fn with_mcpserver<S1: Into<String>, S2: Into<String>, I: IntoIterator<Item = S2>>(mut self, cmd: S1, args: I) -> Self {
208        let args = args.into_iter().collect::<Vec<_>>();
209        let cmd = cmd.into();
210        let args = args.into_iter()
211            .map(|arg| arg.into())
212            .collect();
213        self.mcpservers.push((cmd, args));
214        self
215    }
216}
217
218#[cfg(test)]
219mod test {
220    use crate::model::ChatModel;
221    use super::AgentBuilder;
222
223    #[tokio::test]
224    async fn test_build() {
225        dotenv::from_filename(".env").unwrap();
226        let model = ChatModel::deepseek("deepseek-chat").unwrap();
227        AgentBuilder::new(model)
228            .system_prompt("hihi")
229            .with_builtin_tools(true)
230            .build()
231            .await
232            .expect("build llm");
233        
234    }
235    
236}
abu_agent/agent/build.rs

abu_agent/agent/
build.rs
