agent_io/agent/
builder.rs

1//! Agent builder
2
3use std::sync::Arc;
4
5use crate::Result;
6use crate::llm::BaseChatModel;
7use crate::memory::MemoryManager;
8use crate::tools::Tool;
9
10use super::config::{AgentConfig, build_ephemeral_config};
11use super::service::Agent;
12
13/// Agent builder
14#[derive(Default)]
15pub struct AgentBuilder {
16    llm: Option<Arc<dyn BaseChatModel>>,
17    tools: Vec<Arc<dyn Tool>>,
18    config: Option<AgentConfig>,
19    memory: Option<Arc<RwLock<MemoryManager>>>,
20}
21
22use tokio::sync::RwLock;
23
24impl AgentBuilder {
25    pub fn with_llm(mut self, llm: Arc<dyn BaseChatModel>) -> Self {
26        self.llm = Some(llm);
27        self
28    }
29
30    pub fn tool(mut self, tool: Arc<dyn Tool>) -> Self {
31        self.tools.push(tool);
32        self
33    }
34
35    pub fn tools(mut self, tools: Vec<Arc<dyn Tool>>) -> Self {
36        self.tools = tools;
37        self
38    }
39
40    pub fn config(mut self, config: AgentConfig) -> Self {
41        self.config = Some(config);
42        self
43    }
44
45    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
46        let mut config = self.config.unwrap_or_default();
47        config.system_prompt = Some(prompt.into());
48        self.config = Some(config);
49        self
50    }
51
52    pub fn max_iterations(mut self, max: usize) -> Self {
53        let mut config = self.config.unwrap_or_default();
54        config.max_iterations = max;
55        self.config = Some(config);
56        self
57    }
58
59    /// Add memory manager to the agent
60    pub fn with_memory(mut self, memory: Arc<RwLock<MemoryManager>>) -> Self {
61        self.memory = Some(memory);
62        self
63    }
64
65    pub fn build(self) -> Result<Agent> {
66        let llm = self
67            .llm
68            .ok_or_else(|| crate::Error::Config("LLM is required".into()))?;
69
70        let ephemeral_config = build_ephemeral_config(&self.tools);
71
72        Ok(Agent::new_with_config(
73            llm,
74            self.tools,
75            self.config.unwrap_or_default(),
76            ephemeral_config,
77            self.memory,
78        ))
79    }
80}
agent_io/agent/builder.rs

agent_io/agent/
builder.rs
