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, EphemeralConfig};
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        // Build ephemeral config from tools
71        let ephemeral_config = self
72            .tools
73            .iter()
74            .filter_map(|t| {
75                let cfg = t.ephemeral();
76                if cfg != crate::tools::EphemeralConfig::None {
77                    let keep_count = match cfg {
78                        crate::tools::EphemeralConfig::Single => 1,
79                        crate::tools::EphemeralConfig::Count(n) => n,
80                        crate::tools::EphemeralConfig::None => 0,
81                    };
82                    Some((t.name().to_string(), EphemeralConfig { keep_count }))
83                } else {
84                    None
85                }
86            })
87            .collect();
88
89        Ok(Agent::new_with_config(
90            llm,
91            self.tools,
92            self.config.unwrap_or_default(),
93            ephemeral_config,
94            self.memory,
95        ))
96    }
97}
agent_io/agent/builder.rs

agent_io/agent/
builder.rs
