Skip to main content

traitclaw_core/
agent_builder.rs

1//! Agent builder for fluent configuration.
2//!
3//! Use [`AgentBuilder`] to construct an [`Agent`] with progressive complexity.
4
5use 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
27/// Builder for constructing an [`Agent`] with a fluent API.
28///
29/// # Example
30///
31/// ```rust,no_run
32/// use traitclaw_core::prelude::*;
33/// use traitclaw_core::agent_builder::AgentBuilder;
34///
35/// # async fn example() -> traitclaw_core::Result<()> {
36/// // Minimal agent (provider required):
37/// // let agent = AgentBuilder::new()
38/// //     .provider(my_provider)
39/// //     .system("You are helpful")
40/// //     .build()?;
41/// # Ok(())
42/// # }
43/// ```
44pub 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    /// Create a new builder with default configuration.
62    #[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    /// Set the LLM provider (required).
82    ///
83    /// Prefer [`.model()`][Self::model] for the idiomatic fluent-API usage.
84    /// Both methods are equivalent; `.model()` matches the `agent.model()` pattern
85    /// described in the architecture docs.
86    #[must_use]
87    pub fn provider(mut self, provider: impl Provider) -> Self {
88        self.provider = Some(Arc::new(provider));
89        self
90    }
91
92    /// Set the LLM provider from a pre-wrapped `Arc<dyn Provider>`.
93    ///
94    /// Use this when you already hold a shared provider reference
95    /// (e.g., from [`AgentFactory`](crate::factory::AgentFactory)).
96    #[must_use]
97    pub fn provider_arc(mut self, provider: Arc<dyn Provider>) -> Self {
98        self.provider = Some(provider);
99        self
100    }
101
102    /// Set the LLM provider — preferred alias for [`.provider()`][Self::provider].
103    ///
104    /// Enables the idiomatic `Agent::builder().model(provider).system("...").build()` pattern.
105    #[must_use]
106    pub fn model(self, provider: impl Provider) -> Self {
107        self.provider(provider)
108    }
109
110    /// Wrap the configured provider with automatic retry and exponential backoff.
111    ///
112    /// Must be called **after** `.provider()` or `.model()`.
113    /// Uses [`RetryProvider`](crate::retry::RetryProvider) internally.
114    #[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    /// Set the system prompt.
125    #[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    /// Add a tool to the agent.
132    #[must_use]
133    pub fn tool(mut self, tool: impl ErasedTool) -> Self {
134        self.tools.push(Arc::new(tool));
135        self
136    }
137
138    /// Add a pre-wrapped `Arc<dyn ErasedTool>` directly.
139    ///
140    /// Use this when you already hold a shared tool instance that you want
141    /// to attach to multiple agents without cloning the underlying value.
142    #[must_use]
143    pub fn tool_arc(mut self, tool: Arc<dyn ErasedTool>) -> Self {
144        self.tools.push(tool);
145        self
146    }
147
148    /// Add multiple tools at once.
149    #[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    /// Add multiple pre-wrapped `Arc<dyn ErasedTool>` instances at once.
162    #[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    /// Set the memory backend.
172    #[must_use]
173    pub fn memory(mut self, memory: impl Memory) -> Self {
174        self.memory = Some(Arc::new(memory));
175        self
176    }
177
178    /// Add a guard to the agent.
179    #[must_use]
180    pub fn guard(mut self, guard: impl Guard) -> Self {
181        self.guards.push(Arc::new(guard));
182        self
183    }
184
185    /// Add a hint to the agent.
186    #[must_use]
187    pub fn hint(mut self, hint: impl Hint) -> Self {
188        self.hints.push(Arc::new(hint));
189        self
190    }
191
192    /// Set the tracker for runtime monitoring.
193    #[must_use]
194    pub fn tracker(mut self, tracker: impl Tracker) -> Self {
195        self.tracker = Some(Arc::new(tracker));
196        self
197    }
198
199    /// Set the maximum number of tool call iterations.
200    #[must_use]
201    pub fn max_iterations(mut self, max: u32) -> Self {
202        self.config.max_iterations = max;
203        self
204    }
205
206    /// Set the maximum tokens for LLM responses.
207    #[must_use]
208    pub fn max_tokens(mut self, max: u32) -> Self {
209        self.config.max_tokens = Some(max);
210        self
211    }
212
213    /// Set the sampling temperature.
214    #[must_use]
215    pub fn temperature(mut self, temp: f32) -> Self {
216        self.config.temperature = Some(temp);
217        self
218    }
219
220    /// Set the token budget for the entire run.
221    #[must_use]
222    pub fn token_budget(mut self, budget: usize) -> Self {
223        self.config.token_budget = Some(budget);
224        self
225    }
226
227    /// Set the async context window manager.
228    ///
229    /// Supports LLM-powered compression and accurate token counting.
230    /// Default: `RuleBasedCompressor`.
231    #[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    /// Set the tool execution strategy.
238    ///
239    /// Default: [`SequentialStrategy`] (one at a time).
240    /// Use [`ParallelStrategy`](crate::traits::execution_strategy::ParallelStrategy) for concurrent execution.
241    #[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    /// Set the async tool output transformer.
248    ///
249    /// Supports context-aware, async tool output processing.
250    /// Default: `BudgetAwareTruncator` (10,000 chars).
251    #[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    /// Set the dynamic tool registry (v0.3.0).
258    ///
259    /// Enables runtime tool activation/deactivation.
260    /// Default: `SimpleRegistry` wrapping configured tools.
261    #[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    /// Set the agent execution strategy.
268    ///
269    /// Default: [`DefaultStrategy`] (preserves v0.1.0 loop behavior).
270    /// Implement [`AgentStrategy`] for custom reasoning architectures.
271    #[must_use]
272    pub fn strategy(mut self, strategy: impl AgentStrategy) -> Self {
273        self.strategy = Some(Box::new(strategy));
274        self
275    }
276
277    /// Add a lifecycle hook for observability and interception.
278    ///
279    /// Multiple hooks can be registered and are called sequentially.
280    ///
281    /// # Example
282    ///
283    /// ```rust,no_run
284    /// use traitclaw_core::traits::hook::LoggingHook;
285    ///
286    /// # fn example() {
287    /// // Agent::builder()
288    /// //     .model(my_provider)
289    /// //     .hook(LoggingHook::new(tracing::Level::INFO))
290    /// //     .build()
291    /// # }
292    /// ```
293    #[must_use]
294    pub fn hook(mut self, hook: impl AgentHook) -> Self {
295        self.hooks.push(Arc::new(hook));
296        self
297    }
298
299    /// Build the agent. Returns an error if no provider is configured.
300    ///
301    /// # Errors
302    ///
303    /// Returns [`Error::Config`](crate::Error::Config) if no provider has been set.
304    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        // Default to Noop steering if none configured
313        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        // context_manager defaults to RuleBasedCompressor
329        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        // output_transformer defaults to BudgetAwareTruncator
339        let output_transformer: Arc<dyn OutputTransformer> = self
340            .output_transformer
341            .unwrap_or_else(|| Arc::new(default_out));
342
343        // tool_registry defaults to SimpleRegistry wrapping configured tools
344        let tool_registry: Arc<dyn ToolRegistry> = self
345            .tool_registry
346            .unwrap_or_else(|| Arc::new(SimpleRegistry::new(self.tools.clone())));
347
348        // Default to in-memory if no memory configured
349        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        // AC-2: .build() errors if no provider set
422        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        // AC-1: Agent::builder().model(provider).system("...").build() succeeds
429        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        // AC-3: system/other string params accept &str and String
439        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        // AC-4: optional settings have sensible defaults
454        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}