1use 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
23pub 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, token_tracking_config: None,
57 }
58 }
59
60 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 pub fn with_planner(mut self, planner: Arc<dyn PlannerHandle>) -> Self {
69 self.planner = Some(planner);
70 self
71 }
72
73 pub fn with_tool(mut self, tool: ToolBox) -> Self {
75 self.tools.push(tool);
76 self
77 }
78
79 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 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 pub fn with_event_broadcaster(
159 mut self,
160 broadcaster: Arc<dyn agents_core::events::EventBroadcaster>,
161 ) -> Self {
162 if self.event_dispatcher.is_none() {
164 self.event_dispatcher = Some(Arc::new(agents_core::events::EventDispatcher::new()));
165 }
166
167 if let Some(dispatcher) = &self.event_dispatcher {
169 dispatcher.add_broadcaster(broadcaster);
170 }
171
172 self
173 }
174
175 pub fn with_event_broadcasters(
186 mut self,
187 broadcasters: Vec<Arc<dyn agents_core::events::EventBroadcaster>>,
188 ) -> Self {
189 if self.event_dispatcher.is_none() {
191 self.event_dispatcher = Some(Arc::new(agents_core::events::EventDispatcher::new()));
192 }
193
194 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 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 pub fn with_pii_sanitization(mut self, enabled: bool) -> Self {
239 self.enable_pii_sanitization = enabled;
240 self
241 }
242
243 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 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 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 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 let final_planner = if let Some(token_config) = token_tracking_config {
338 if token_config.enabled {
339 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}