1use std::borrow::Borrow;
2use std::collections::BTreeMap;
3use std::sync::Arc;
4
5use serde::de::DeserializeOwned;
6
7use crate::guardrails::{InputGuardrail, OutputGuardrail};
8use crate::handoffs::Handoff;
9use crate::model::ModelRef;
10use crate::model_settings::ModelSettings;
11use crate::output_validation::{
12 HostOutputValidator, OutputRepair, OutputRepairRequest, OutputValidationContext,
13 OutputValidationResult,
14};
15use crate::runtime::RuntimeHook;
16use crate::tools::common::trim_portable_whitespace;
17use crate::tools::{AgentToolBuilder, BackgroundAgentTaskBuilder};
18use crate::tools::{Tool, ToolPolicy};
19use crate::types::{Metadata, NoToolPolicy, SubAgentConfig};
20
21pub type InstructionProvider =
22 Arc<dyn Fn(&crate::context::RunContext, &Agent) -> String + Send + Sync>;
23pub type OutputValidator = Arc<dyn Fn(&str) -> Result<(), String> + Send + Sync>;
24
25#[derive(Clone)]
26pub struct Agent {
27 name: String,
28 instructions: String,
29 instruction_provider: Option<InstructionProvider>,
30 model: Option<ModelRef>,
31 model_settings: ModelSettings,
32 tools: Vec<Arc<dyn Tool>>,
33 handoffs: Vec<Handoff>,
34 input_guardrails: Vec<Arc<dyn InputGuardrail>>,
35 output_guardrails: Vec<Arc<dyn OutputGuardrail>>,
36 output_type_name: Option<&'static str>,
37 output_validator: Option<OutputValidator>,
38 output_validation_enabled: bool,
39 host_output_validator: Option<HostOutputValidator>,
40 output_repair: Option<OutputRepair>,
41 output_validation_max_repairs: u8,
42 output_repair_model: Option<ModelRef>,
43 output_repair_model_settings: Option<ModelSettings>,
44 hooks: Vec<Arc<dyn RuntimeHook>>,
45 max_cycles: Option<u32>,
46 no_tool_policy: Option<NoToolPolicy>,
47 tool_use_behavior: ToolUseBehavior,
48 tool_policy: ToolPolicy,
49 sub_agents: BTreeMap<String, SubAgentConfig>,
50 metadata: Metadata,
51}
52
53impl Agent {
54 pub fn builder(name: impl Into<String>) -> AgentBuilder {
55 AgentBuilder {
56 agent: Self {
57 name: name.into(),
58 instructions: String::new(),
59 instruction_provider: None,
60 model: None,
61 model_settings: ModelSettings::default(),
62 tools: Vec::new(),
63 handoffs: Vec::new(),
64 input_guardrails: Vec::new(),
65 output_guardrails: Vec::new(),
66 output_type_name: None,
67 output_validator: None,
68 output_validation_enabled: false,
69 host_output_validator: None,
70 output_repair: None,
71 output_validation_max_repairs: 1,
72 output_repair_model: None,
73 output_repair_model_settings: None,
74 hooks: Vec::new(),
75 max_cycles: None,
76 no_tool_policy: None,
77 tool_use_behavior: ToolUseBehavior::default(),
78 tool_policy: ToolPolicy::default(),
79 sub_agents: BTreeMap::new(),
80 metadata: Metadata::new(),
81 },
82 sub_agent_error: None,
83 output_validation_error: None,
84 }
85 }
86
87 pub fn name(&self) -> &str {
88 &self.name
89 }
90
91 pub fn instructions(&self) -> &str {
92 &self.instructions
93 }
94
95 pub(crate) fn has_dynamic_instructions(&self) -> bool {
96 self.instruction_provider.is_some()
97 }
98
99 pub fn resolve_instructions(&self, context: &crate::context::RunContext) -> String {
100 self.instruction_provider
101 .as_ref()
102 .map(|provider| provider(context, self))
103 .unwrap_or_else(|| self.instructions.clone())
104 }
105
106 pub fn model(&self) -> Option<&ModelRef> {
107 self.model.as_ref()
108 }
109
110 pub fn model_settings(&self) -> &ModelSettings {
111 &self.model_settings
112 }
113
114 pub fn tools(&self) -> &[Arc<dyn Tool>] {
115 &self.tools
116 }
117
118 pub fn handoffs(&self) -> &[Handoff] {
119 &self.handoffs
120 }
121
122 pub(crate) fn input_guardrails(&self) -> &[Arc<dyn InputGuardrail>] {
123 &self.input_guardrails
124 }
125
126 pub(crate) fn output_guardrails(&self) -> &[Arc<dyn OutputGuardrail>] {
127 &self.output_guardrails
128 }
129
130 pub fn output_type_name(&self) -> Option<&'static str> {
131 self.output_type_name
132 }
133
134 pub fn validate_output(&self, output: &str) -> Result<(), String> {
135 match self.output_validator.as_ref() {
136 Some(validator) => validator(output),
137 None => Ok(()),
138 }
139 }
140
141 pub fn output_validation_enabled(&self) -> bool {
142 self.output_validation_enabled
143 }
144
145 pub fn host_output_validator(&self) -> Option<&HostOutputValidator> {
146 self.host_output_validator.as_ref()
147 }
148
149 pub fn output_repair(&self) -> Option<&OutputRepair> {
150 self.output_repair.as_ref()
151 }
152
153 pub fn output_validation_max_repairs(&self) -> u8 {
154 self.output_validation_max_repairs
155 }
156
157 pub fn output_repair_model(&self) -> Option<&ModelRef> {
158 self.output_repair_model.as_ref()
159 }
160
161 pub fn output_repair_model_settings(&self) -> Option<&ModelSettings> {
162 self.output_repair_model_settings.as_ref()
163 }
164
165 pub fn hooks(&self) -> &[Arc<dyn RuntimeHook>] {
166 &self.hooks
167 }
168
169 pub fn max_cycles(&self) -> Option<u32> {
170 self.max_cycles
171 }
172
173 pub fn no_tool_policy(&self) -> Option<NoToolPolicy> {
174 self.no_tool_policy
175 }
176
177 pub fn tool_use_behavior(&self) -> &ToolUseBehavior {
178 &self.tool_use_behavior
179 }
180
181 pub fn tool_policy(&self) -> &ToolPolicy {
182 &self.tool_policy
183 }
184
185 pub fn sub_agents(&self) -> &BTreeMap<String, SubAgentConfig> {
186 &self.sub_agents
187 }
188
189 pub fn metadata(&self) -> &Metadata {
190 &self.metadata
191 }
192
193 pub fn as_tool(&self) -> AgentToolBuilder {
194 AgentToolBuilder::new(self.clone())
195 }
196
197 pub fn as_background_task(&self) -> BackgroundAgentTaskBuilder {
198 BackgroundAgentTaskBuilder::new(self.clone())
199 }
200}
201
202pub struct AgentBuilder {
203 agent: Agent,
204 sub_agent_error: Option<String>,
205 output_validation_error: Option<String>,
206}
207
208impl AgentBuilder {
209 pub fn instructions(mut self, instructions: impl Into<String>) -> Self {
210 self.agent.instructions = instructions.into();
211 self.agent.instruction_provider = None;
212 self
213 }
214
215 pub fn dynamic_instructions(
216 mut self,
217 provider: impl Fn(&crate::context::RunContext, &Agent) -> String + Send + Sync + 'static,
218 ) -> Self {
219 self.agent.instructions.clear();
220 self.agent.instruction_provider = Some(Arc::new(provider));
221 self
222 }
223
224 pub fn model(mut self, model: ModelRef) -> Self {
225 self.agent.model = Some(model);
226 self
227 }
228
229 pub fn model_settings(mut self, settings: ModelSettings) -> Self {
230 self.agent.model_settings = settings;
231 self
232 }
233
234 pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
235 self.agent.tools.push(Arc::new(tool));
236 self
237 }
238
239 pub fn tool_arc(mut self, tool: Arc<dyn Tool>) -> Self {
240 self.agent.tools.push(tool);
241 self
242 }
243
244 pub fn handoff(mut self, handoff: impl Into<Handoff>) -> Self {
245 self.agent.handoffs.push(handoff.into());
246 self
247 }
248
249 pub fn input_guardrail(mut self, guardrail: Arc<dyn InputGuardrail>) -> Self {
250 self.agent.input_guardrails.push(guardrail);
251 self
252 }
253
254 pub fn output_guardrail(mut self, guardrail: Arc<dyn OutputGuardrail>) -> Self {
255 self.agent.output_guardrails.push(guardrail);
256 self
257 }
258
259 pub fn output_type<T>(mut self) -> Self
260 where
261 T: DeserializeOwned + 'static,
262 {
263 self.agent.output_type_name = Some(std::any::type_name::<T>());
264 self.agent.output_validator = Some(Arc::new(|output| {
265 serde_json::from_str::<T>(output)
266 .map(|_| ())
267 .map_err(|error| error.to_string())
268 }));
269 self
270 }
271
272 pub fn output_validator(
273 mut self,
274 name: &'static str,
275 validator: impl Fn(&str) -> Result<(), String> + Send + Sync + 'static,
276 ) -> Self {
277 self.agent.output_type_name = Some(name);
278 self.agent.output_validator = Some(Arc::new(validator));
279 self
280 }
281
282 pub fn output_validation_enabled(mut self, enabled: bool) -> Self {
283 self.agent.output_validation_enabled = enabled;
284 self
285 }
286
287 pub fn host_output_validator(
288 mut self,
289 validator: impl Fn(&str, &OutputValidationContext) -> OutputValidationResult
290 + Send
291 + Sync
292 + 'static,
293 ) -> Self {
294 self.agent.host_output_validator = Some(Arc::new(validator));
295 self
296 }
297
298 pub fn output_repair(
299 mut self,
300 repair: impl Fn(&OutputRepairRequest) -> Result<String, String> + Send + Sync + 'static,
301 ) -> Self {
302 self.agent.output_repair = Some(Arc::new(repair));
303 self
304 }
305
306 pub fn output_validation_max_repairs(mut self, max_repairs: u8) -> Self {
307 if max_repairs > 1 {
308 self.output_validation_error =
309 Some("output_validation_max_repairs must be 0 or 1".to_string());
310 } else {
311 self.agent.output_validation_max_repairs = max_repairs;
312 }
313 self
314 }
315
316 pub fn output_repair_model(mut self, model: ModelRef) -> Self {
317 self.agent.output_repair_model = Some(model);
318 self
319 }
320
321 pub fn output_repair_model_settings(mut self, settings: ModelSettings) -> Self {
322 self.agent.output_repair_model_settings = Some(settings);
323 self
324 }
325
326 pub fn hook(mut self, hook: Arc<dyn RuntimeHook>) -> Self {
327 self.agent.hooks.push(hook);
328 self
329 }
330
331 pub fn max_cycles(mut self, max_cycles: u32) -> Self {
332 self.agent.max_cycles = Some(max_cycles);
333 self
334 }
335
336 pub fn no_tool_policy(mut self, policy: NoToolPolicy) -> Self {
337 self.agent.no_tool_policy = Some(policy);
338 self
339 }
340
341 pub fn tool_use_behavior(mut self, behavior: ToolUseBehavior) -> Self {
342 self.agent.tool_use_behavior = behavior;
343 self
344 }
345
346 pub fn tool_policy(mut self, tool_policy: ToolPolicy) -> Self {
347 self.agent.tool_policy = tool_policy;
348 self
349 }
350
351 pub fn sub_agent(mut self, id: impl AsRef<str>, config: impl Borrow<SubAgentConfig>) -> Self {
352 self.insert_sub_agent(id.as_ref(), config.borrow());
353 self
354 }
355
356 pub fn sub_agents<I, K, C>(mut self, sub_agents: I) -> Self
357 where
358 I: IntoIterator<Item = (K, C)>,
359 K: AsRef<str>,
360 C: Borrow<SubAgentConfig>,
361 {
362 for (id, config) in sub_agents {
363 self.insert_sub_agent(id.as_ref(), config.borrow());
364 }
365 self
366 }
367
368 pub fn metadata(mut self, key: impl Into<String>, value: serde_json::Value) -> Self {
369 self.agent.metadata.insert(key.into(), value);
370 self
371 }
372
373 pub fn build(self) -> Result<Agent, String> {
374 if self.agent.name.trim().is_empty() {
375 return Err("agent name cannot be empty".to_string());
376 }
377 if self.agent.instructions.trim().is_empty() && self.agent.instruction_provider.is_none() {
378 return Err("agent instructions cannot be empty".to_string());
379 }
380 if let Some(error) = self.sub_agent_error {
381 return Err(error);
382 }
383 if let Some(error) = self.output_validation_error {
384 return Err(error);
385 }
386 if self.agent.output_validation_enabled && self.agent.host_output_validator.is_none() {
387 return Err("enabled output validation requires a host_output_validator".to_string());
388 }
389 if self.agent.output_repair.is_some() && self.agent.host_output_validator.is_none() {
390 return Err("output_repair requires a host_output_validator".to_string());
391 }
392 Ok(self.agent)
393 }
394
395 fn insert_sub_agent(&mut self, id: &str, config: &SubAgentConfig) {
396 let normalized_id = trim_portable_whitespace(id);
397 if normalized_id.is_empty() {
398 self.sub_agent_error
399 .get_or_insert_with(|| "sub-agent id cannot be empty".to_string());
400 return;
401 }
402 if self.agent.sub_agents.contains_key(normalized_id) {
403 self.sub_agent_error.get_or_insert_with(|| {
404 format!("duplicate sub-agent id after normalization: {normalized_id}")
405 });
406 return;
407 }
408 self.agent
409 .sub_agents
410 .insert(normalized_id.to_string(), config.clone());
411 }
412}
413
414#[derive(Clone, Default)]
415pub enum ToolUseBehavior {
416 #[default]
417 RunLlmAgain,
418 StopOnFirstTool,
419 StopAtToolNames(Vec<String>),
420}