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

agents_runtime/agent/
builder.rs
