agent_base/engine/
builder.rs1use std::sync::Arc;
2
3use crate::llm::{LlmClient, ReasoningConfig};
4use crate::tool::{Tool, ToolPolicy, ToolRegistry};
5use crate::types::{AgentConfig, ResponseFormat, RetryConfig, SessionIdGenerator, AtomicU64SessionIdGenerator};
6
7use super::approval::ApprovalHandler;
8use super::context::ContextWindowManager;
9use super::middleware::{Middleware, MiddlewareRef};
10use super::recovery::{StopOnError, ToolErrorRecovery};
11use super::session_store::{InMemorySessionStore, SessionStore};
12use super::AgentRuntime;
13
14pub struct AgentBuilder {
15 client: Arc<dyn LlmClient>,
16 config: AgentConfig,
17 tools: ToolRegistry,
18 approval_handler: Option<Arc<dyn ApprovalHandler>>,
19 tool_policy: Option<Arc<dyn ToolPolicy>>,
20 middlewares: Vec<MiddlewareRef>,
21 context_manager: Option<ContextWindowManager>,
22 session_store: Option<Arc<dyn SessionStore>>,
23 error_recovery: Option<Arc<dyn ToolErrorRecovery>>,
24 event_bus_capacity: usize,
25 session_id_generator: Option<Arc<dyn SessionIdGenerator>>,
26}
27
28impl AgentBuilder {
29 pub fn new(client: Arc<dyn LlmClient>) -> Self {
30 Self {
31 client,
32 config: AgentConfig::default(),
33 tools: ToolRegistry::default(),
34 approval_handler: None,
35 tool_policy: None,
36 middlewares: Vec::new(),
37 context_manager: None,
38 session_store: None,
39 error_recovery: None,
40 event_bus_capacity: 2048,
41 session_id_generator: None,
42 }
43 }
44
45 pub fn event_bus_capacity(mut self, capacity: usize) -> Self {
46 self.event_bus_capacity = capacity;
47 self
48 }
49
50 pub fn session_id_generator(mut self, generator: Arc<dyn SessionIdGenerator>) -> Self {
51 self.session_id_generator = Some(generator);
52 self
53 }
54
55 pub fn system_prompt(mut self, system_prompt: impl Into<String>) -> Self {
56 self.config.system_prompt = Some(system_prompt.into());
57 self
58 }
59
60 pub fn enable_thought(mut self, enable: bool) -> Self {
61 self.config.enable_thought = enable;
62 self
63 }
64
65 pub fn reasoning(mut self, config: ReasoningConfig) -> Self {
66 self.config.reasoning = Some(config);
67 self
68 }
69
70 pub fn enable_thinking(mut self, enable: bool) -> Self {
71 let mut config = self.config.reasoning.take().unwrap_or_default();
72 config.enabled = Some(enable);
73 self.config.reasoning = Some(config);
74 self
75 }
76
77 pub fn thinking_budget(mut self, budget: u64) -> Self {
78 let mut config = self.config.reasoning.take().unwrap_or_default();
79 config.budget_tokens = Some(budget);
80 self.config.reasoning = Some(config);
81 self
82 }
83
84 pub fn tool_timeout(mut self, timeout_ms: u64) -> Self {
85 self.config.tool.tool_timeout_ms = Some(timeout_ms);
86 self
87 }
88
89 pub fn max_tool_output_chars(mut self, max_chars: usize) -> Self {
90 self.config.tool.max_tool_output_chars = Some(max_chars);
91 self
92 }
93
94 pub fn register_tool(mut self, tool: impl Tool + 'static) -> Self {
95 self.tools.register(tool);
96 self
97 }
98
99 pub fn register_tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
100 self.tools.register_arc(tool);
101 self
102 }
103
104 pub fn approval_handler(mut self, handler: Arc<dyn ApprovalHandler>) -> Self {
105 self.approval_handler = Some(handler);
106 self
107 }
108
109 pub fn tool_policy(mut self, policy: Arc<dyn ToolPolicy>) -> Self {
110 self.tool_policy = Some(policy);
111 self
112 }
113
114 pub fn middleware(mut self, mw: impl Middleware + 'static) -> Self {
115 self.middlewares.push(Arc::new(mw));
116 self
117 }
118
119 pub fn context_window(mut self, max_tokens: usize) -> Self {
120 self.context_manager = Some(ContextWindowManager::new(max_tokens));
121 self
122 }
123
124 pub fn context_window_manager(mut self, manager: ContextWindowManager) -> Self {
125 self.context_manager = Some(manager);
126 self
127 }
128
129 pub fn response_format(mut self, format: ResponseFormat) -> Self {
130 self.config.llm.response_format = Some(format);
131 self
132 }
133
134 pub fn llm_retry(mut self, retry: RetryConfig) -> Self {
135 self.config.llm.llm_retry = Some(retry);
136 self
137 }
138
139 pub fn session_store(mut self, store: Arc<dyn SessionStore>) -> Self {
140 self.session_store = Some(store);
141 self
142 }
143
144 pub fn error_recovery(mut self, recovery: Arc<dyn ToolErrorRecovery>) -> Self {
145 self.error_recovery = Some(recovery);
146 self
147 }
148
149 pub fn tool_error_retry_prompt(mut self, prompt: impl Into<String>) -> Self {
150 self.config.tool.tool_error_retry_prompt = Some(prompt.into());
151 self
152 }
153
154 pub fn language(mut self, language: crate::types::Language) -> Self {
155 self.config.language = language;
156 self
157 }
158
159 pub fn build(self) -> crate::types::AgentResult<AgentRuntime> {
160 let event_bus = super::runtime::EventBus::new(self.event_bus_capacity);
161 let session_store = self
162 .session_store
163 .unwrap_or_else(|| Arc::new(InMemorySessionStore::new()));
164 let error_recovery = self
165 .error_recovery
166 .unwrap_or_else(|| Arc::new(StopOnError));
167 let session_id_generator = self
168 .session_id_generator
169 .unwrap_or_else(|| Arc::new(AtomicU64SessionIdGenerator::default()));
170
171 let session_manager = super::runtime::SessionManager::new(
172 session_id_generator,
173 session_store,
174 );
175
176 let llm_engine = super::runtime::LlmEngine::new(
177 self.client.clone(),
178 event_bus.clone(),
179 );
180
181 let tool_engine = super::runtime::ToolEngine::new(
182 self.tools,
183 self.approval_handler,
184 self.tool_policy,
185 self.middlewares.clone(),
186 error_recovery,
187 event_bus.clone(),
188 );
189
190 Ok(AgentRuntime {
191 config: self.config,
192 llm_engine,
193 tool_engine,
194 session_manager,
195 event_bus,
196 context_manager: self.context_manager,
197 middlewares: self.middlewares,
198 })
199 }
200}