1use std::sync::Arc;
6
7use crate::agent::Agent;
8use crate::config::AgentConfig;
9use crate::context_managers::RuleBasedCompressor;
10use crate::default_strategy::DefaultStrategy;
11use crate::memory::in_memory::InMemoryMemory;
12use crate::traits::context_manager::ContextManager;
13use crate::traits::execution_strategy::{ExecutionStrategy, SequentialStrategy};
14use crate::traits::guard::{Guard, NoopGuard};
15use crate::traits::hint::{Hint, NoopHint};
16use crate::traits::hook::AgentHook;
17use crate::traits::memory::Memory;
18use crate::traits::output_transformer::OutputTransformer;
19use crate::traits::provider::Provider;
20use crate::traits::strategy::AgentStrategy;
21use crate::traits::tool::ErasedTool;
22use crate::traits::tool_registry::{SimpleRegistry, ToolRegistry};
23use crate::traits::tracker::{NoopTracker, Tracker};
24use crate::transformers::BudgetAwareTruncator;
25use crate::Result;
26
27pub struct AgentBuilder {
45 provider: Option<Arc<dyn Provider>>,
46 tools: Vec<Arc<dyn ErasedTool>>,
47 memory: Option<Arc<dyn Memory>>,
48 guards: Vec<Arc<dyn Guard>>,
49 hints: Vec<Arc<dyn Hint>>,
50 tracker: Option<Arc<dyn Tracker>>,
51 context_manager: Option<Arc<dyn ContextManager>>,
52 execution_strategy: Option<Arc<dyn ExecutionStrategy>>,
53 output_transformer: Option<Arc<dyn OutputTransformer>>,
54 tool_registry: Option<Arc<dyn ToolRegistry>>,
55 strategy: Option<Box<dyn AgentStrategy>>,
56 hooks: Vec<Arc<dyn AgentHook>>,
57 config: AgentConfig,
58}
59
60impl AgentBuilder {
61 #[must_use]
63 pub fn new() -> Self {
64 Self {
65 provider: None,
66 tools: Vec::new(),
67 memory: None,
68 guards: Vec::new(),
69 hints: Vec::new(),
70 tracker: None,
71 context_manager: None,
72 execution_strategy: None,
73 output_transformer: None,
74 tool_registry: None,
75 strategy: None,
76 hooks: Vec::new(),
77 config: AgentConfig::default(),
78 }
79 }
80
81 #[must_use]
87 pub fn provider(mut self, provider: impl Provider) -> Self {
88 self.provider = Some(Arc::new(provider));
89 self
90 }
91
92 #[must_use]
97 pub fn provider_arc(mut self, provider: Arc<dyn Provider>) -> Self {
98 self.provider = Some(provider);
99 self
100 }
101
102 #[must_use]
106 pub fn model(self, provider: impl Provider) -> Self {
107 self.provider(provider)
108 }
109
110 #[must_use]
115 pub fn with_retry(mut self, config: crate::retry::RetryConfig) -> Self {
116 if let Some(inner) = self.provider.take() {
117 self.provider = Some(Arc::new(crate::retry::RetryProvider::new(inner, config)));
118 } else {
119 tracing::warn!("with_retry() called before provider() — retry config will be ignored");
120 }
121 self
122 }
123
124 #[must_use]
126 pub fn system(mut self, prompt: impl Into<String>) -> Self {
127 self.config.system_prompt = Some(prompt.into());
128 self
129 }
130
131 #[must_use]
133 pub fn tool(mut self, tool: impl ErasedTool) -> Self {
134 self.tools.push(Arc::new(tool));
135 self
136 }
137
138 #[must_use]
143 pub fn tool_arc(mut self, tool: Arc<dyn ErasedTool>) -> Self {
144 self.tools.push(tool);
145 self
146 }
147
148 #[must_use]
150 pub fn tools<I, T>(mut self, tools: I) -> Self
151 where
152 I: IntoIterator<Item = T>,
153 T: ErasedTool,
154 {
155 for tool in tools {
156 self.tools.push(Arc::new(tool));
157 }
158 self
159 }
160
161 #[must_use]
163 pub fn tools_arc<I>(mut self, tools: I) -> Self
164 where
165 I: IntoIterator<Item = Arc<dyn ErasedTool>>,
166 {
167 self.tools.extend(tools);
168 self
169 }
170
171 #[must_use]
173 pub fn memory(mut self, memory: impl Memory) -> Self {
174 self.memory = Some(Arc::new(memory));
175 self
176 }
177
178 #[must_use]
180 pub fn guard(mut self, guard: impl Guard) -> Self {
181 self.guards.push(Arc::new(guard));
182 self
183 }
184
185 #[must_use]
187 pub fn hint(mut self, hint: impl Hint) -> Self {
188 self.hints.push(Arc::new(hint));
189 self
190 }
191
192 #[must_use]
194 pub fn tracker(mut self, tracker: impl Tracker) -> Self {
195 self.tracker = Some(Arc::new(tracker));
196 self
197 }
198
199 #[must_use]
201 pub fn max_iterations(mut self, max: u32) -> Self {
202 self.config.max_iterations = max;
203 self
204 }
205
206 #[must_use]
208 pub fn max_tokens(mut self, max: u32) -> Self {
209 self.config.max_tokens = Some(max);
210 self
211 }
212
213 #[must_use]
215 pub fn temperature(mut self, temp: f32) -> Self {
216 self.config.temperature = Some(temp);
217 self
218 }
219
220 #[must_use]
222 pub fn token_budget(mut self, budget: usize) -> Self {
223 self.config.token_budget = Some(budget);
224 self
225 }
226
227 #[must_use]
232 pub fn context_manager(mut self, manager: impl ContextManager + 'static) -> Self {
233 self.context_manager = Some(Arc::new(manager));
234 self
235 }
236
237 #[must_use]
242 pub fn execution_strategy(mut self, strategy: impl ExecutionStrategy + 'static) -> Self {
243 self.execution_strategy = Some(Arc::new(strategy));
244 self
245 }
246
247 #[must_use]
252 pub fn output_transformer(mut self, transformer: impl OutputTransformer + 'static) -> Self {
253 self.output_transformer = Some(Arc::new(transformer));
254 self
255 }
256
257 #[must_use]
262 pub fn tool_registry(mut self, registry: impl ToolRegistry + 'static) -> Self {
263 self.tool_registry = Some(Arc::new(registry));
264 self
265 }
266
267 #[must_use]
272 pub fn strategy(mut self, strategy: impl AgentStrategy) -> Self {
273 self.strategy = Some(Box::new(strategy));
274 self
275 }
276
277 #[must_use]
294 pub fn hook(mut self, hook: impl AgentHook) -> Self {
295 self.hooks.push(Arc::new(hook));
296 self
297 }
298
299 pub fn build(self) -> Result<Agent> {
305 let provider = self.provider.ok_or_else(|| {
306 crate::Error::Config(
307 "AgentBuilder: no provider configured. Use .provider(my_provider) before .build()"
308 .into(),
309 )
310 })?;
311
312 let guards: Vec<Arc<dyn Guard>> = if self.guards.is_empty() {
314 vec![Arc::new(NoopGuard)]
315 } else {
316 self.guards
317 };
318
319 let hints: Vec<Arc<dyn Hint>> = if self.hints.is_empty() {
320 vec![Arc::new(NoopHint)]
321 } else {
322 self.hints
323 };
324
325 let tracker = self.tracker.unwrap_or_else(|| Arc::new(NoopTracker));
326
327 let default_ctx = RuleBasedCompressor::default();
328 let context_manager: Arc<dyn ContextManager> = self
330 .context_manager
331 .unwrap_or_else(|| Arc::new(default_ctx));
332
333 let execution_strategy = self
334 .execution_strategy
335 .unwrap_or_else(|| Arc::new(SequentialStrategy));
336
337 let default_out = BudgetAwareTruncator::default();
338 let output_transformer: Arc<dyn OutputTransformer> = self
340 .output_transformer
341 .unwrap_or_else(|| Arc::new(default_out));
342
343 let tool_registry: Arc<dyn ToolRegistry> = self
345 .tool_registry
346 .unwrap_or_else(|| Arc::new(SimpleRegistry::new(self.tools.clone())));
347
348 let memory = self
350 .memory
351 .unwrap_or_else(|| Arc::new(InMemoryMemory::new()));
352
353 let strategy = self.strategy.unwrap_or_else(|| Box::new(DefaultStrategy));
354
355 Ok(Agent::new(
356 provider,
357 self.tools,
358 memory,
359 guards,
360 hints,
361 tracker,
362 context_manager,
363 execution_strategy,
364 output_transformer,
365 tool_registry,
366 strategy,
367 self.hooks,
368 self.config,
369 ))
370 }
371}
372
373impl Default for AgentBuilder {
374 fn default() -> Self {
375 Self::new()
376 }
377}
378
379#[cfg(test)]
380mod tests {
381 use super::*;
382 use crate::types::completion::{CompletionRequest, CompletionResponse, ResponseContent, Usage};
383 use crate::types::model_info::{ModelInfo, ModelTier};
384 use crate::types::stream::CompletionStream;
385 use async_trait::async_trait;
386
387 struct FakeProvider {
388 info: ModelInfo,
389 }
390
391 impl FakeProvider {
392 fn new() -> Self {
393 Self {
394 info: ModelInfo::new("fake", ModelTier::Small, 4_096, false, false, false),
395 }
396 }
397 }
398
399 #[async_trait]
400 impl Provider for FakeProvider {
401 async fn complete(&self, _req: CompletionRequest) -> crate::Result<CompletionResponse> {
402 Ok(CompletionResponse {
403 content: ResponseContent::Text("ok".into()),
404 usage: Usage {
405 prompt_tokens: 1,
406 completion_tokens: 1,
407 total_tokens: 2,
408 },
409 })
410 }
411 async fn stream(&self, _req: CompletionRequest) -> crate::Result<CompletionStream> {
412 unimplemented!()
413 }
414 fn model_info(&self) -> &ModelInfo {
415 &self.info
416 }
417 }
418
419 #[test]
420 fn test_builder_without_provider_returns_error() {
421 let result = AgentBuilder::new().system("You are helpful").build();
423 assert!(result.is_err());
424 }
425
426 #[test]
427 fn test_builder_model_alias_ac1() {
428 let result = Agent::builder()
430 .model(FakeProvider::new())
431 .system("You are helpful")
432 .build();
433 assert!(result.is_ok());
434 }
435
436 #[test]
437 fn test_builder_accepts_str_and_string_ac3() {
438 let result_str = Agent::builder()
440 .model(FakeProvider::new())
441 .system("literal")
442 .build();
443 let result_string = Agent::builder()
444 .model(FakeProvider::new())
445 .system("owned".to_string())
446 .build();
447 assert!(result_str.is_ok());
448 assert!(result_string.is_ok());
449 }
450
451 #[test]
452 fn test_defaults_ac4() {
453 let config = AgentConfig::default();
455 assert_eq!(
456 config.max_iterations, 20,
457 "default max_iterations should be 20"
458 );
459 assert_eq!(
460 config.max_tokens,
461 Some(4096),
462 "default max_tokens should be 4096"
463 );
464 assert!(
465 (config.temperature.unwrap_or(0.0) - 0.7).abs() < f32::EPSILON,
466 "default temperature should be 0.7"
467 );
468 }
469}