runifold_agent/
structured.rs1use std::marker::PhantomData;
2
3use runifold_core::RunContext;
4use runifold_model::StructuredOutputError;
5use serde::de::DeserializeOwned;
6use thiserror::Error;
7
8use crate::{Agent, AgentError, AgentEventStream, AgentFuture, StructuredAgentOutcome};
9
10#[derive(Debug, Error)]
12#[non_exhaustive]
13pub enum StructuredAgentError {
14 #[error(transparent)]
16 Agent(#[from] AgentError),
17 #[error(transparent)]
19 Output(#[from] StructuredOutputError),
20}
21
22#[derive(Clone)]
24pub struct StructuredAgent<T> {
25 agent: Agent,
26 output: PhantomData<fn() -> T>,
27}
28
29impl<T> StructuredAgent<T> {
30 pub(crate) const fn new(agent: Agent) -> Self {
31 Self {
32 agent,
33 output: PhantomData,
34 }
35 }
36
37 pub const fn agent(&self) -> &Agent {
39 &self.agent
40 }
41
42 pub fn into_agent(self) -> Agent {
44 self.agent
45 }
46
47 pub fn stream<'a>(
53 &'a self,
54 input: impl Into<String> + Send + 'a,
55 run: &'a RunContext,
56 ) -> AgentEventStream<'a> {
57 self.agent.stream(input, run)
58 }
59}
60
61impl<T> StructuredAgent<T>
62where
63 T: DeserializeOwned + Send + 'static,
64{
65 pub fn run<'a>(
73 &'a self,
74 input: impl Into<String> + Send + 'a,
75 run: &'a RunContext,
76 ) -> AgentFuture<'a, Result<StructuredAgentOutcome<T>, StructuredAgentError>> {
77 Box::pin(async move {
78 let outcome = self.agent.run(input, run).await?;
79 Ok(outcome.into_structured()?)
80 })
81 }
82}
83
84impl<T> std::fmt::Debug for StructuredAgent<T> {
85 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86 formatter
87 .debug_tuple("StructuredAgent")
88 .field(&self.agent)
89 .finish()
90 }
91}
92
93#[cfg(test)]
94mod tests {
95 use std::{collections::BTreeMap, sync::Arc};
96
97 use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
98 use runifold_model::{
99 ContentPart, FinishReason, ModelRef, ModelStreamEvent, OutputFormat,
100 StructuredOutputErrorKind,
101 };
102 use runifold_testkit::ScriptedModel;
103 use schemars::JsonSchema;
104 use serde::Deserialize;
105
106 use crate::{Agent, StructuredAgentError};
107
108 #[derive(Debug, Deserialize, Eq, JsonSchema, PartialEq)]
109 struct Answer {
110 value: u32,
111 }
112
113 fn events(text: &str) -> Vec<ModelStreamEvent> {
114 vec![
115 ModelStreamEvent::ResponseStarted {
116 id: Some("response".into()),
117 model: ModelRef::new("test", "scripted"),
118 },
119 ModelStreamEvent::ContentPartCompleted {
120 index: 0,
121 part: ContentPart::text(text),
122 },
123 ModelStreamEvent::ResponseCompleted {
124 finish_reason: FinishReason::Stop,
125 provider_metadata: BTreeMap::new(),
126 },
127 ]
128 }
129
130 fn run() -> RunContext {
131 RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new())
132 }
133
134 #[test]
135 fn typed_agent_uses_one_type_for_schema_and_decode() {
136 let model = ScriptedModel::new();
137 model.enqueue(events("{\"value\":42}"));
138 let agent = Agent::builder(
139 "typed",
140 Arc::new(model.clone()),
141 ModelRef::new("test", "scripted"),
142 )
143 .build_structured::<Answer>("answer")
144 .unwrap();
145 let run = run();
146
147 let typed = futures_executor::block_on(agent.run("answer", &run)).unwrap();
148
149 assert_eq!(typed.output, Answer { value: 42 });
150 let requests = model.recorded_requests();
151 let OutputFormat::JsonSchema { name, strict, .. } = &requests[0].output_format else {
152 panic!("expected JSON-schema output");
153 };
154 assert_eq!(name, "answer");
155 assert!(*strict);
156 }
157
158 #[test]
159 fn typed_agent_fails_the_completion_requirement_before_returning_an_outcome() {
160 let model = ScriptedModel::new();
161 model.enqueue(events("{\"value\":\"wrong\"}"));
162 let agent = Agent::new("typed", Arc::new(model), ModelRef::new("test", "scripted"))
163 .into_structured::<Answer>("answer");
164 let run = run();
165
166 let error = futures_executor::block_on(agent.run("answer", &run)).unwrap_err();
167
168 assert!(matches!(
169 error,
170 StructuredAgentError::Agent(crate::AgentError::StructuredOutputUnsatisfied {
171 attempts: 0,
172 kind: StructuredOutputErrorKind::InvalidOutput,
173 ..
174 })
175 ));
176 }
177}