1use std::sync::Arc;
2
3use runifold_core::CapabilitySet;
4use runifold_effect::{EffectExecutor, EffectRecoveryPolicy};
5use runifold_model::{
6 FeaturePolicy, GenerationOptions, Message, Model, ModelRef, OutputFormat, ProviderToolSpec,
7 ResponseMode,
8};
9use runifold_retrieval::{Document, RetrievalError, Retriever};
10use runifold_tool::{Tool, ToolRegistrationError};
11use schemars::JsonSchema;
12use thiserror::Error;
13
14use crate::agent::DynamicContext;
15use crate::{
16 Agent, AgentConfig, AgentDescriptor, AgentError, AgentFuture, AgentOutcome,
17 AgentRegistrationError, AgentRoute, GatewayMiddleware, StructuredAgent, ToolErrorPolicy,
18};
19
20#[derive(Clone, Debug, Error, Eq, PartialEq)]
22#[non_exhaustive]
23pub enum AgentBuildError {
24 #[error("agent Tool registration failed: {0}")]
26 Tool(#[from] ToolRegistrationError),
27 #[error("agent route registration failed: {0}")]
29 Route(#[from] AgentRegistrationError),
30 #[error("callable name `{0}` is registered as both a Tool and an Agent")]
32 CallableNameCollision(String),
33 #[error("agent name cannot be empty")]
35 EmptyName,
36 #[error("max_turns must be greater than zero")]
38 ZeroMaxTurns,
39 #[error("agent retrieval configuration failed: {0}")]
41 Retrieval(#[from] RetrievalError),
42}
43
44#[derive(Debug, Error)]
46#[non_exhaustive]
47pub enum AgentPromptError {
48 #[error("failed to build agent: {0}")]
50 Build(#[from] AgentBuildError),
51 #[error("agent prompt failed: {0}")]
53 Run(#[from] AgentError),
54}
55
56pub struct AgentBuilder {
62 agent: Agent,
63 error: Option<AgentBuildError>,
64}
65
66impl AgentBuilder {
67 pub fn new(name: impl Into<String>, model: Arc<dyn Model>, model_ref: ModelRef) -> Self {
69 Self {
70 agent: Agent::new(name, model, model_ref),
71 error: None,
72 }
73 }
74
75 #[must_use]
77 pub fn system(mut self, instruction: impl Into<String>) -> Self {
78 self.agent
79 .instructions
80 .push(Message::system(instruction.into()));
81 self
82 }
83
84 #[must_use]
89 pub fn context(self, text: impl Into<String>) -> Self {
90 let id = format!("static-context-{}", self.agent.context.len() + 1);
91 match Document::new(id, text) {
92 Ok(document) => self.context_document(document),
93 Err(error) => self.with_error(error.into()),
94 }
95 }
96
97 #[must_use]
99 pub fn context_document(mut self, document: Document) -> Self {
100 if self.error.is_none() {
101 self.agent.context.push(document);
102 }
103 self
104 }
105
106 #[must_use]
108 pub fn dynamic_context<R>(self, limit: usize, retriever: R) -> Self
109 where
110 R: Retriever + 'static,
111 {
112 self.shared_dynamic_context(limit, Arc::new(retriever))
113 }
114
115 #[must_use]
117 pub fn shared_dynamic_context(mut self, limit: usize, retriever: Arc<dyn Retriever>) -> Self {
118 if self.error.is_none() {
119 if limit == 0 {
120 self.error = Some(RetrievalError::ZeroLimit.into());
121 } else {
122 self.agent
123 .dynamic_context
124 .push(DynamicContext { limit, retriever });
125 }
126 }
127 self
128 }
129
130 #[must_use]
132 pub fn tool<T>(self, tool: T) -> Self
133 where
134 T: Tool + 'static,
135 {
136 self.shared_tool(Arc::new(tool))
137 }
138
139 #[must_use]
141 pub fn shared_tool(mut self, tool: Arc<dyn Tool>) -> Self {
142 if self.error.is_none()
143 && let Err(error) = self.agent.tools.register(tool)
144 {
145 self.error = Some(error.into());
146 }
147 self
148 }
149
150 fn with_error(mut self, error: AgentBuildError) -> Self {
151 if self.error.is_none() {
152 self.error = Some(error);
153 }
154 self
155 }
156
157 #[must_use]
159 pub fn child(
160 mut self,
161 descriptor: AgentDescriptor,
162 child: Arc<Agent>,
163 capabilities: CapabilitySet,
164 ) -> Self {
165 if self.error.is_none() {
166 let route = AgentRoute::new(descriptor, child).with_capabilities(capabilities);
167 if let Err(error) = self.agent.agents.register(route) {
168 self.error = Some(error.into());
169 }
170 }
171 self
172 }
173
174 #[must_use]
176 pub fn gateway_layer(mut self, middleware: Arc<dyn GatewayMiddleware>) -> Self {
177 self.agent.agents.push_middleware(middleware);
178 self
179 }
180
181 #[must_use]
183 pub fn max_delegation_depth(mut self, max_depth: u32) -> Self {
184 self.agent.agents = self.agent.agents.with_max_depth(max_depth);
185 self
186 }
187
188 #[must_use]
190 pub const fn max_turns(mut self, max_turns: u32) -> Self {
191 self.agent.config.max_turns = max_turns;
192 self
193 }
194
195 #[must_use]
197 pub const fn tool_error_policy(mut self, policy: ToolErrorPolicy) -> Self {
198 self.agent.config.tool_error_policy = policy;
199 self
200 }
201
202 #[must_use]
204 pub const fn feature_policy(mut self, policy: FeaturePolicy) -> Self {
205 self.agent.config.feature_policy = policy;
206 self
207 }
208
209 #[must_use]
211 pub fn output_format(mut self, output_format: OutputFormat) -> Self {
212 self.agent.output_format = output_format;
213 self
214 }
215
216 #[must_use]
218 pub fn structured_output<T>(self, name: impl Into<String>) -> Self
219 where
220 T: JsonSchema,
221 {
222 self.output_format(OutputFormat::typed::<T>(name))
223 }
224
225 #[must_use]
227 pub fn provider_tool(mut self, tool: ProviderToolSpec) -> Self {
228 self.agent.provider_tools.push(tool);
229 self
230 }
231
232 #[must_use]
234 pub fn generation(mut self, generation: GenerationOptions) -> Self {
235 self.agent.generation = generation;
236 self
237 }
238
239 #[must_use]
241 pub fn temperature(mut self, temperature: f64) -> Self {
242 self.agent.generation.temperature = Some(temperature);
243 self
244 }
245
246 #[must_use]
248 pub fn top_p(mut self, top_p: f64) -> Self {
249 self.agent.generation.top_p = Some(top_p);
250 self
251 }
252
253 #[must_use]
255 pub fn max_output_tokens(mut self, max_output_tokens: u64) -> Self {
256 self.agent.generation.max_output_tokens = Some(max_output_tokens);
257 self
258 }
259
260 #[must_use]
262 pub const fn response_mode(mut self, response_mode: ResponseMode) -> Self {
263 self.agent.response_mode = response_mode;
264 self
265 }
266
267 #[must_use]
269 pub fn provider_options(
270 mut self,
271 provider: impl Into<String>,
272 options: serde_json::Value,
273 ) -> Self {
274 self.agent.provider_options.insert(provider.into(), options);
275 self
276 }
277
278 #[must_use]
280 pub const fn config(mut self, config: AgentConfig) -> Self {
281 self.agent.config = config;
282 self
283 }
284
285 #[must_use]
287 pub fn effect_executor(mut self, effects: EffectExecutor) -> Self {
288 self.agent.effects = effects;
289 self
290 }
291
292 #[must_use]
294 pub const fn effect_recovery_policy(mut self, policy: EffectRecoveryPolicy) -> Self {
295 self.agent.effect_recovery = policy;
296 self
297 }
298
299 pub fn build(self) -> Result<Agent, AgentBuildError> {
306 if let Some(error) = self.error {
307 return Err(error);
308 }
309 if self.agent.name.trim().is_empty() {
310 return Err(AgentBuildError::EmptyName);
311 }
312 if self.agent.config.max_turns == 0 {
313 return Err(AgentBuildError::ZeroMaxTurns);
314 }
315 if let Some(collision) = self
316 .agent
317 .agents
318 .model_specs()
319 .into_iter()
320 .find(|spec| self.agent.tools.contains(&spec.name))
321 {
322 return Err(AgentBuildError::CallableNameCollision(collision.name));
323 }
324 Ok(self.agent)
325 }
326
327 pub fn prompt(
333 self,
334 input: impl Into<String> + Send + 'static,
335 ) -> AgentFuture<'static, Result<AgentOutcome, AgentPromptError>> {
336 let input = input.into();
337 Box::pin(async move {
338 let agent = self.build()?;
339 Ok(agent.prompt(input).await?)
340 })
341 }
342
343 pub fn prompt_text(
350 self,
351 input: impl Into<String> + Send + 'static,
352 ) -> AgentFuture<'static, Result<String, AgentPromptError>> {
353 let input = input.into();
354 Box::pin(async move {
355 let agent = self.build()?;
356 Ok(agent.prompt_text(input).await?)
357 })
358 }
359
360 pub fn build_structured<T>(
370 self,
371 name: impl Into<String>,
372 ) -> Result<StructuredAgent<T>, AgentBuildError>
373 where
374 T: JsonSchema,
375 {
376 self.structured_output::<T>(name)
377 .build()
378 .map(StructuredAgent::new)
379 }
380}
381
382impl std::fmt::Debug for AgentBuilder {
383 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
384 formatter
385 .debug_struct("AgentBuilder")
386 .field("agent", &self.agent.name)
387 .field("error", &self.error)
388 .finish_non_exhaustive()
389 }
390}
391
392#[cfg(test)]
393mod tests {
394 use std::{collections::BTreeMap, sync::Arc};
395
396 use runifold_core::{CapabilityId, CapabilitySet, EffectClass, RiskLevel};
397 use runifold_model::{
398 ContentPart, FinishReason, ModelRef, ModelStreamEvent, OutputFormat, ProviderToolSpec,
399 ResponseMode,
400 };
401 use runifold_testkit::ScriptedModel;
402 use runifold_tool::{Tool, ToolContext, ToolDescriptor, ToolError, ToolFuture, ToolOutput};
403 use schemars::JsonSchema;
404 use serde::Deserialize;
405 use serde_json::json;
406
407 use crate::{Agent, AgentBuildError, AgentDescriptor, AgentPromptError};
408
409 struct TestTool {
410 descriptor: ToolDescriptor,
411 }
412
413 #[derive(Deserialize, JsonSchema)]
414 struct TypedAnswer {
415 value: u32,
416 }
417
418 impl TestTool {
419 fn named(name: &str) -> Self {
420 Self {
421 descriptor: ToolDescriptor {
422 id: CapabilityId::new(),
423 name: name.into(),
424 version: "1".into(),
425 description: "test".into(),
426 input_schema: json!({"type": "object"}),
427 output_schema: json!({"type": "object"}),
428 effect: EffectClass::Pure,
429 risk: RiskLevel::Low,
430 metadata: BTreeMap::new(),
431 },
432 }
433 }
434 }
435
436 impl Tool for TestTool {
437 fn descriptor(&self) -> &ToolDescriptor {
438 &self.descriptor
439 }
440
441 fn invoke(
442 &self,
443 input: serde_json::Value,
444 _context: ToolContext,
445 ) -> ToolFuture<'_, Result<ToolOutput, ToolError>> {
446 Box::pin(async move { Ok(ToolOutput::model_visible(input)) })
447 }
448 }
449
450 #[test]
451 fn fluent_builder_assembles_the_canonical_agent() {
452 let model = Arc::new(ScriptedModel::new());
453 let agent = Agent::builder("worker", model, ModelRef::new("test", "scripted"))
454 .system("Be precise")
455 .tool(TestTool::named("lookup"))
456 .max_turns(4)
457 .build()
458 .unwrap();
459
460 assert_eq!(agent.name, "worker");
461 assert_eq!(agent.instructions.len(), 1);
462 assert!(agent.tools.contains("lookup"));
463 assert_eq!(agent.config.max_turns, 4);
464 assert_eq!(agent.callable_capabilities().len(), 1);
465 }
466
467 #[test]
468 fn builder_retains_generation_provider_and_delivery_controls() {
469 let provider_tool = ProviderToolSpec::new("ark", "web_search").unwrap();
470 let agent = Agent::builder(
471 "researcher",
472 Arc::new(ScriptedModel::new()),
473 ModelRef::new("ark", "doubao"),
474 )
475 .temperature(0.2)
476 .top_p(0.8)
477 .max_output_tokens(4_096)
478 .response_mode(ResponseMode::Complete)
479 .provider_tool(provider_tool)
480 .provider_options("ark", json!({"thinking": {"type": "enabled"}}))
481 .build()
482 .unwrap();
483
484 assert_eq!(agent.generation.temperature, Some(0.2));
485 assert_eq!(agent.generation.top_p, Some(0.8));
486 assert_eq!(agent.generation.max_output_tokens, Some(4_096));
487 assert_eq!(agent.response_mode, ResponseMode::Complete);
488 assert_eq!(agent.provider_tools[0].tool_type, "web_search");
489 assert_eq!(agent.provider_options["ark"]["thinking"]["type"], "enabled");
490 }
491
492 #[test]
493 fn build_rejects_tool_and_agent_name_collisions() {
494 let model = Arc::new(ScriptedModel::new());
495 let child = Arc::new(Agent::new(
496 "child",
497 model.clone(),
498 ModelRef::new("test", "child"),
499 ));
500 let error = Agent::builder("parent", model, ModelRef::new("test", "parent"))
501 .tool(TestTool::named("search"))
502 .child(
503 AgentDescriptor::new("search", "delegate search"),
504 child,
505 CapabilitySet::new(),
506 )
507 .build()
508 .unwrap_err();
509
510 assert_eq!(
511 error,
512 AgentBuildError::CallableNameCollision("search".into())
513 );
514 }
515
516 #[test]
517 fn builder_derives_a_strict_output_schema_from_a_rust_type() {
518 let example = TypedAnswer { value: 7 };
519 assert_eq!(example.value, 7);
520 let agent = Agent::builder(
521 "worker",
522 Arc::new(ScriptedModel::new()),
523 ModelRef::new("test", "scripted"),
524 )
525 .structured_output::<TypedAnswer>("typed_answer")
526 .build()
527 .unwrap();
528
529 let OutputFormat::JsonSchema {
530 name,
531 schema,
532 strict,
533 } = agent.output_format
534 else {
535 panic!("expected JSON-schema output");
536 };
537 assert_eq!(name, "typed_answer");
538 assert!(strict);
539 assert_eq!(schema["properties"]["value"]["type"], "integer");
540 }
541
542 #[test]
543 fn builder_prompt_text_is_a_single_use_golden_path() {
544 let model = ScriptedModel::new();
545 model.enqueue([
546 ModelStreamEvent::ResponseStarted {
547 id: Some("response-1".into()),
548 model: ModelRef::new("test", "scripted"),
549 },
550 ModelStreamEvent::ContentPartCompleted {
551 index: 0,
552 part: ContentPart::text("done"),
553 },
554 ModelStreamEvent::ResponseCompleted {
555 finish_reason: FinishReason::Stop,
556 provider_metadata: BTreeMap::new(),
557 },
558 ]);
559
560 let text = futures_executor::block_on(
561 Agent::builder("worker", Arc::new(model), ModelRef::new("test", "scripted"))
562 .system("Be precise")
563 .prompt_text("start"),
564 )
565 .unwrap();
566
567 assert_eq!(text, "done");
568 }
569
570 #[test]
571 fn builder_prompt_reports_build_failures_before_model_execution() {
572 let error = futures_executor::block_on(
573 Agent::builder(
574 "",
575 Arc::new(ScriptedModel::new()),
576 ModelRef::new("test", "scripted"),
577 )
578 .prompt("start"),
579 )
580 .unwrap_err();
581
582 assert!(matches!(
583 error,
584 AgentPromptError::Build(AgentBuildError::EmptyName)
585 ));
586 }
587}