agents_runtime/agent/
builder.rs

1//! Fluent builder API for constructing Deep Agents
2//!
3//! This module provides the ConfigurableAgentBuilder that offers a fluent interface
4//! for building Deep Agents, mirroring the Python SDK's ergonomic construction patterns.
5
6use super::api::{
7    create_async_deep_agent_from_config, create_deep_agent_from_config, get_default_model,
8};
9use super::config::{DeepAgentConfig, SubAgentConfig, SummarizationConfig};
10use super::runtime::DeepAgent;
11use crate::middleware::{
12    token_tracking::{TokenTrackingConfig, TokenTrackingMiddleware},
13    HitlPolicy,
14};
15use crate::planner::LlmBackedPlanner;
16use agents_core::agent::PlannerHandle;
17use agents_core::llm::LanguageModel;
18use agents_core::persistence::Checkpointer;
19use agents_core::tools::ToolBox;
20use std::collections::{HashMap, HashSet};
21use std::sync::Arc;
22
23/// Builder API to assemble a DeepAgent in a single fluent flow, mirroring the Python
24/// `create_configurable_agent` experience. Prefer this for ergonomic construction.
25pub struct ConfigurableAgentBuilder {
26    instructions: String,
27    planner: Option<Arc<dyn PlannerHandle>>,
28    tools: Vec<ToolBox>,
29    subagents: Vec<SubAgentConfig>,
30    summarization: Option<SummarizationConfig>,
31    tool_interrupts: HashMap<String, HitlPolicy>,
32    builtin_tools: Option<HashSet<String>>,
33    auto_general_purpose: bool,
34    enable_prompt_caching: bool,
35    checkpointer: Option<Arc<dyn Checkpointer>>,
36    event_dispatcher: Option<Arc<agents_core::events::EventDispatcher>>,
37    enable_pii_sanitization: bool,
38    token_tracking_config: Option<TokenTrackingConfig>,
39}
40
41impl ConfigurableAgentBuilder {
42    pub fn new(instructions: impl Into<String>) -> Self {
43        Self {
44            instructions: instructions.into(),
45            planner: None,
46            tools: Vec::new(),
47            subagents: Vec::new(),
48            summarization: None,
49            tool_interrupts: HashMap::new(),
50            builtin_tools: None,
51            auto_general_purpose: true,
52            enable_prompt_caching: false,
53            checkpointer: None,
54            event_dispatcher: None,
55            enable_pii_sanitization: true, // Enabled by default for security
56            token_tracking_config: None,
57        }
58    }
59
60    /// Set the language model for the agent (mirrors Python's `model` parameter)
61    pub fn with_model(mut self, model: Arc<dyn LanguageModel>) -> Self {
62        let planner: Arc<dyn PlannerHandle> = Arc::new(LlmBackedPlanner::new(model));
63        self.planner = Some(planner);
64        self
65    }
66
67    /// Low-level planner API (for advanced use cases)
68    pub fn with_planner(mut self, planner: Arc<dyn PlannerHandle>) -> Self {
69        self.planner = Some(planner);
70        self
71    }
72
73    /// Add a tool to the agent
74    pub fn with_tool(mut self, tool: ToolBox) -> Self {
75        self.tools.push(tool);
76        self
77    }
78
79    /// Add multiple tools
80    pub fn with_tools<I>(mut self, tools: I) -> Self
81    where
82        I: IntoIterator<Item = ToolBox>,
83    {
84        self.tools.extend(tools);
85        self
86    }
87
88    pub fn with_subagent_config<I>(mut self, cfgs: I) -> Self
89    where
90        I: IntoIterator<Item = SubAgentConfig>,
91    {
92        self.subagents.extend(cfgs);
93        self
94    }
95
96    /// Convenience method: automatically create subagents from a list of tools.
97    /// Each tool becomes a specialized subagent with that single tool.
98    pub fn with_subagent_tools<I>(mut self, tools: I) -> Self
99    where
100        I: IntoIterator<Item = ToolBox>,
101    {
102        for tool in tools {
103            let tool_name = tool.schema().name.clone();
104            let subagent_config = SubAgentConfig::new(
105                format!("{}-agent", tool_name),
106                format!("Specialized agent for {} operations", tool_name),
107                format!(
108                    "You are a specialized agent. Use the {} tool to complete tasks efficiently.",
109                    tool_name
110                ),
111            )
112            .with_tools(vec![tool]);
113            self.subagents.push(subagent_config);
114        }
115        self
116    }
117
118    pub fn with_summarization(mut self, config: SummarizationConfig) -> Self {
119        self.summarization = Some(config);
120        self
121    }
122
123    pub fn with_tool_interrupt(mut self, tool_name: impl Into<String>, policy: HitlPolicy) -> Self {
124        self.tool_interrupts.insert(tool_name.into(), policy);
125        self
126    }
127
128    pub fn with_builtin_tools<I, S>(mut self, names: I) -> Self
129    where
130        I: IntoIterator<Item = S>,
131        S: Into<String>,
132    {
133        self.builtin_tools = Some(names.into_iter().map(|s| s.into()).collect());
134        self
135    }
136
137    pub fn with_auto_general_purpose(mut self, enabled: bool) -> Self {
138        self.auto_general_purpose = enabled;
139        self
140    }
141
142    pub fn with_prompt_caching(mut self, enabled: bool) -> Self {
143        self.enable_prompt_caching = enabled;
144        self
145    }
146
147    pub fn with_checkpointer(mut self, checkpointer: Arc<dyn Checkpointer>) -> Self {
148        self.checkpointer = Some(checkpointer);
149        self
150    }
151
152    /// Add a single event broadcaster to the agent
153    ///
154    /// Example:
155    /// ```ignore
156    /// builder.with_event_broadcaster(console_broadcaster)
157    /// ```
158    pub fn with_event_broadcaster(
159        mut self,
160        broadcaster: Arc<dyn agents_core::events::EventBroadcaster>,
161    ) -> Self {
162        // Create dispatcher if it doesn't exist
163        if self.event_dispatcher.is_none() {
164            self.event_dispatcher = Some(Arc::new(agents_core::events::EventDispatcher::new()));
165        }
166
167        // Add broadcaster to the dispatcher (uses interior mutability)
168        if let Some(dispatcher) = &self.event_dispatcher {
169            dispatcher.add_broadcaster(broadcaster);
170        }
171
172        self
173    }
174
175    /// Add multiple event broadcasters at once (cleaner API)
176    ///
177    /// Example:
178    /// ```ignore
179    /// builder.with_event_broadcasters(vec![
180    ///     console_broadcaster,
181    ///     whatsapp_broadcaster,
182    ///     dynamodb_broadcaster,
183    /// ])
184    /// ```
185    pub fn with_event_broadcasters(
186        mut self,
187        broadcasters: Vec<Arc<dyn agents_core::events::EventBroadcaster>>,
188    ) -> Self {
189        // Create dispatcher if it doesn't exist
190        if self.event_dispatcher.is_none() {
191            self.event_dispatcher = Some(Arc::new(agents_core::events::EventDispatcher::new()));
192        }
193
194        // Add all broadcasters
195        if let Some(dispatcher) = &self.event_dispatcher {
196            for broadcaster in broadcasters {
197                dispatcher.add_broadcaster(broadcaster);
198            }
199        }
200
201        self
202    }
203
204    /// Set the event dispatcher directly (replaces any existing dispatcher)
205    pub fn with_event_dispatcher(
206        mut self,
207        dispatcher: Arc<agents_core::events::EventDispatcher>,
208    ) -> Self {
209        self.event_dispatcher = Some(dispatcher);
210        self
211    }
212
213    /// Enable or disable PII sanitization in event data.
214    ///
215    /// **Enabled by default for security.**
216    ///
217    /// When enabled (default):
218    /// - Message previews are truncated to 100 characters
219    /// - Sensitive fields (passwords, tokens, api_keys, etc.) are redacted
220    /// - PII patterns (emails, phones, credit cards) are removed
221    ///
222    /// Disable only if you need raw data and have other security measures in place.
223    ///
224    /// # Example
225    ///
226    /// ```ignore
227    /// // Keep default (enabled)
228    /// let agent = DeepAgentBuilder::new("instructions")
229    ///     .with_model(model)
230    ///     .build()?;
231    ///
232    /// // Explicitly disable (not recommended for production)
233    /// let agent = DeepAgentBuilder::new("instructions")
234    ///     .with_model(model)
235    ///     .with_pii_sanitization(false)
236    ///     .build()?;
237    /// ```
238    pub fn with_pii_sanitization(mut self, enabled: bool) -> Self {
239        self.enable_pii_sanitization = enabled;
240        self
241    }
242
243    /// Enable token tracking for monitoring LLM usage and costs.
244    ///
245    /// This enables tracking of token usage, costs, and performance metrics
246    /// across all LLM requests made by the agent.
247    ///
248    /// # Example
249    ///
250    /// ```ignore
251    /// // Enable token tracking with default settings
252    /// let agent = ConfigurableAgentBuilder::new("instructions")
253    ///     .with_model(model)
254    ///     .with_token_tracking(true)
255    ///     .build()?;
256    ///
257    /// // Enable with custom configuration
258    /// let config = TokenTrackingConfig {
259    ///     enabled: true,
260    ///     emit_events: true,
261    ///     log_usage: true,
262    ///     custom_costs: Some(TokenCosts::openai_gpt4o_mini()),
263    /// };
264    /// let agent = ConfigurableAgentBuilder::new("instructions")
265    ///     .with_model(model)
266    ///     .with_token_tracking_config(config)
267    ///     .build()?;
268    /// ```
269    pub fn with_token_tracking(mut self, enabled: bool) -> Self {
270        self.token_tracking_config = Some(TokenTrackingConfig {
271            enabled,
272            emit_events: enabled,
273            log_usage: enabled,
274            custom_costs: None,
275        });
276        self
277    }
278
279    /// Configure token tracking with custom settings.
280    ///
281    /// This allows fine-grained control over token tracking behavior,
282    /// including custom cost models and event emission settings.
283    ///
284    /// # Example
285    ///
286    /// ```ignore
287    /// let config = TokenTrackingConfig {
288    ///     enabled: true,
289    ///     emit_events: true,
290    ///     log_usage: false, // Don't log to console
291    ///     custom_costs: Some(TokenCosts::openai_gpt4o_mini()),
292    /// };
293    /// let agent = ConfigurableAgentBuilder::new("instructions")
294    ///     .with_model(model)
295    ///     .with_token_tracking_config(config)
296    ///     .build()?;
297    /// ```
298    pub fn with_token_tracking_config(mut self, config: TokenTrackingConfig) -> Self {
299        self.token_tracking_config = Some(config);
300        self
301    }
302
303    pub fn build(self) -> anyhow::Result<DeepAgent> {
304        self.finalize(create_deep_agent_from_config)
305    }
306
307    /// Build an agent using the async constructor alias. This mirrors the Python
308    /// async_create_deep_agent entry point, while reusing the same runtime internals.
309    pub fn build_async(self) -> anyhow::Result<DeepAgent> {
310        self.finalize(create_async_deep_agent_from_config)
311    }
312
313    fn finalize(self, ctor: fn(DeepAgentConfig) -> DeepAgent) -> anyhow::Result<DeepAgent> {
314        let Self {
315            instructions,
316            planner,
317            tools,
318            subagents,
319            summarization,
320            tool_interrupts,
321            builtin_tools,
322            auto_general_purpose,
323            enable_prompt_caching,
324            checkpointer,
325            event_dispatcher,
326            enable_pii_sanitization,
327            token_tracking_config,
328        } = self;
329
330        let planner = planner.unwrap_or_else(|| {
331            // Use default model if no planner is set
332            let default_model = get_default_model().expect("Failed to get default model");
333            Arc::new(LlmBackedPlanner::new(default_model)) as Arc<dyn PlannerHandle>
334        });
335
336        // Wrap the planner with token tracking if enabled
337        let final_planner = if let Some(token_config) = token_tracking_config {
338            if token_config.enabled {
339                // Extract the underlying model from the planner
340                let planner_any = planner.as_any();
341                if let Some(llm_planner) = planner_any.downcast_ref::<LlmBackedPlanner>() {
342                    let model = llm_planner.model().clone();
343                    let tracked_model = Arc::new(TokenTrackingMiddleware::new(
344                        token_config,
345                        model,
346                        event_dispatcher.clone(),
347                    ));
348                    Arc::new(LlmBackedPlanner::new(tracked_model)) as Arc<dyn PlannerHandle>
349                } else {
350                    planner
351                }
352            } else {
353                planner
354            }
355        } else {
356            planner
357        };
358
359        let mut cfg = DeepAgentConfig::new(instructions, final_planner)
360            .with_auto_general_purpose(auto_general_purpose)
361            .with_prompt_caching(enable_prompt_caching)
362            .with_pii_sanitization(enable_pii_sanitization);
363
364        if let Some(ckpt) = checkpointer {
365            cfg = cfg.with_checkpointer(ckpt);
366        }
367        if let Some(dispatcher) = event_dispatcher {
368            cfg = cfg.with_event_dispatcher(dispatcher);
369        }
370        if let Some(sum) = summarization {
371            cfg = cfg.with_summarization(sum);
372        }
373        if let Some(selected) = builtin_tools {
374            cfg = cfg.with_builtin_tools(selected);
375        }
376        for (name, policy) in tool_interrupts {
377            cfg = cfg.with_tool_interrupt(name, policy);
378        }
379        for tool in tools {
380            cfg = cfg.with_tool(tool);
381        }
382        for sub_cfg in subagents {
383            cfg = cfg.with_subagent_config(sub_cfg);
384        }
385
386        Ok(ctor(cfg))
387    }
388}