Skip to main content

abu_agent/
build.rs

1use std::path::PathBuf;
2use abu_provider::{deepseek::DeepSeek, ChatProvide};
3use abu_tool::Tool;
4use crate::{context::ContextBuilder, kit::tools::{bash::Bash, calculate::Calculator, fs::{FileCreator, FileReader, FileWritor}}, memory::{Memory, SequentialMemory}, AgentResult};
5use super::{Agent, AgentConfig, AgentKit};
6
7const DEFAULT_SYSTEM_PROMPT: &str = "You are an agent.";
8
9pub struct AgentBuilder<C: ChatProvide = DeepSeek, M: Memory = SequentialMemory> {
10    pub llm: C,
11    pub model: String,
12    pub config: AgentConfig,
13    pub memory: M,
14    pub system_prompt: String,
15    pub with_skills: Option<PathBuf>,
16    pub with_builin_tools: bool,
17    pub with_subagent: bool,
18    pub tools: Vec<Box<dyn Tool>>,
19    pub mcpservers: Vec<(String, Vec<String>)>,
20    pub mcpconfig_path: Option<PathBuf>,
21}
22
23impl Default for AgentConfig {
24    fn default() -> Self {
25        Self {
26            max_iteration: 10,
27            temperature: 0.7,
28        }
29    }
30}
31
32impl<C: ChatProvide, M: Memory> AgentBuilder<C, M> {
33    pub async fn build(mut self) -> AgentResult<Agent<C, M>> {
34        // let mut system_prompt = format!("{}\nOnce you consider the work complete or do task to do, call the terminate method.", self.system_prompt);
35        let mut kit = AgentKit::new();
36        // kit.add_tool(Terminator::new());
37
38        // tool
39        if self.with_builin_tools {
40            kit.add_tool(Bash::new());
41            kit.add_tool(Calculator::new());
42            kit.add_tool(FileCreator::new());
43            kit.add_tool(FileWritor::new());
44            kit.add_tool(FileReader::new());
45        }
46
47        for tool in self.tools {
48            kit.add_tool_box(tool);
49        }
50
51        // mcp
52        if let Some(path) = self.mcpconfig_path {
53            kit.load_mcpconfig(&path).await?;
54        }
55
56        for (cmd, args) in self.mcpservers {
57            kit.add_mcp_server(&cmd, &args).await?;
58        }
59
60        // skill
61        if let Some(skill_path) = self.with_skills {
62            kit.load_skill(skill_path)?;
63            self.system_prompt = kit.attach_system_prompt(&self.system_prompt);   
64        }
65
66        // context builder
67        let context_builder = ContextBuilder::new(self.system_prompt);
68
69        Ok(Agent {
70            config: self.config,
71            llm: self.llm,
72            model: self.model,
73            memory: self.memory,
74            kit,
75            context_builder
76        })
77
78    }
79}
80
81impl<C: ChatProvide> AgentBuilder<C> {
82    pub fn new(llm: C, model: impl Into<String>) -> Self {
83        Self {
84            llm,
85            model: model.into(),
86            config: AgentConfig::default(),
87            memory: SequentialMemory::default(),
88            system_prompt: DEFAULT_SYSTEM_PROMPT.to_string(),
89            with_skills: None,
90            with_builin_tools: true,
91            with_subagent: false,
92            tools: vec![],
93            mcpservers: vec![],
94            mcpconfig_path: None,
95        }
96    }
97}
98
99impl<C: ChatProvide, M: Memory> AgentBuilder<C, M> {
100    pub fn temperature(mut self, temperature: f64) -> Self {
101        self.config.temperature = temperature;
102        self
103    }
104
105    pub fn max_iteration(mut self, max_iteration: usize) -> Self {
106        self.config.max_iteration = max_iteration;
107        self
108    }
109
110    pub fn memory<NM: Memory>(self, memory: NM) -> AgentBuilder<C, NM> {
111        AgentBuilder {
112            memory,
113            llm: self.llm,
114            model: self.model,
115            config: self.config,
116            system_prompt: self.system_prompt,
117            with_skills: self.with_skills,
118            with_builin_tools: self.with_builin_tools,
119            with_subagent: self.with_subagent,
120            tools: self.tools,
121            mcpservers: self.mcpservers,
122            mcpconfig_path: self.mcpconfig_path
123        }
124    }
125
126    pub fn llm<NC: ChatProvide>(self, llm: NC) -> AgentBuilder<NC, M> {
127        AgentBuilder {
128            memory: self.memory,
129            llm,
130            model: self.model,
131            config: self.config,
132            system_prompt: self.system_prompt,
133            with_skills: self.with_skills,
134            with_builin_tools: self.with_builin_tools,
135            with_subagent: self.with_subagent,
136            tools: self.tools,
137            mcpservers: self.mcpservers,
138            mcpconfig_path: self.mcpconfig_path
139        }
140    }
141
142    pub fn model(mut self, model: impl Into<String>) -> Self {
143        self.model = model.into();
144        self
145    }
146
147    pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
148        self.system_prompt = system_prompt.into();
149        self
150    }
151
152    pub fn with_skills(mut self, skill_path: impl Into<PathBuf>) -> Self {
153        self.with_skills = Some(skill_path.into());
154        self
155    }
156
157    pub fn with_builin_tools(mut self, enabled: bool) -> Self {
158        self.with_builin_tools = enabled;
159        self
160    }
161
162    pub fn with_tool(mut self, tool: impl Tool + 'static) -> Self {
163        self.tools.push(Box::new(tool));
164        self
165    }
166
167    pub fn with_tools(mut self, tools: impl IntoIterator<Item = Box<dyn Tool>>) -> Self {
168        for tool in tools.into_iter() {
169            self.tools.push(tool);
170        }
171        self
172    }
173
174    pub fn with_mcpconfig(mut self, path: impl Into<PathBuf>) -> Self {
175        self.mcpconfig_path = Some(path.into());
176        self
177    }
178
179    pub fn with_mcpserver<'a>(mut self, cmd: &str, args: impl IntoIterator<Item = &'a str>) -> Self {
180        let args = args.into_iter().collect::<Vec<_>>();
181        let cmd = cmd.to_string();
182        let args = args.into_iter()
183            .map(|arg| arg.to_string())
184            .collect();
185        self.mcpservers.push((cmd, args));
186        self
187    }
188}
189
190#[cfg(test)]
191mod test {
192    use abu_provider::deepseek::DeepSeek;
193
194    use super::AgentBuilder;
195
196    #[tokio::test]
197    async fn test_build() {
198        dotenv::from_filename(".env").unwrap();
199        let deepseek = DeepSeek::from_env().expect("new deepseek");
200        let model = std::env::var("MODEL_ID").unwrap();
201        AgentBuilder::new(deepseek, model)
202            .system_prompt("hihi")
203            .with_builin_tools(true)
204            .build()
205            .await
206            .expect("build llm");
207        
208    }
209    
210}