runifold_agent/
structured.rs

1use 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/// Failure while executing or locally decoding a typed Agent run.
11#[derive(Debug, Error)]
12#[non_exhaustive]
13pub enum StructuredAgentError {
14    /// Canonical Agent execution failed.
15    #[error(transparent)]
16    Agent(#[from] AgentError),
17    /// The terminal response did not satisfy the bound Rust output type.
18    #[error(transparent)]
19    Output(#[from] StructuredOutputError),
20}
21
22/// An Agent whose provider schema and local decoder are bound to the same type.
23#[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    /// Returns the underlying canonical Agent.
38    pub const fn agent(&self) -> &Agent {
39        &self.agent
40    }
41
42    /// Consumes the wrapper and returns the underlying canonical Agent.
43    pub fn into_agent(self) -> Agent {
44        self.agent
45    }
46
47    /// Streams the underlying canonical Agent lifecycle.
48    ///
49    /// The terminal `Completed` event retains an unparsed
50    /// [`crate::AgentOutcome`]. Use [`Self::run`] when the terminal item itself
51    /// must be typed.
52    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    /// Runs the canonical Agent and locally validates its terminal response.
66    ///
67    /// # Errors
68    ///
69    /// Returns [`StructuredAgentError::Agent`] for execution failures and
70    /// [`StructuredAgentError::Output`] for local decoding failures.
71    pub fn run<'a>(
72        &'a self,
73        input: impl Into<String> + Send + 'a,
74        run: &'a RunContext,
75    ) -> AgentFuture<'a, Result<StructuredAgentOutcome<T>, StructuredAgentError>> {
76        Box::pin(async move {
77            let outcome = self.agent.run(input, run).await?;
78            Ok(outcome.into_structured()?)
79        })
80    }
81}
82
83impl<T> std::fmt::Debug for StructuredAgent<T> {
84    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
85        formatter
86            .debug_tuple("StructuredAgent")
87            .field(&self.agent)
88            .finish()
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use std::{collections::BTreeMap, sync::Arc};
95
96    use runifold_core::{Budget, BudgetTracker, CapabilitySet, RunContext};
97    use runifold_model::{
98        ContentPart, FinishReason, ModelRef, ModelStreamEvent, OutputFormat,
99        StructuredOutputErrorKind,
100    };
101    use runifold_testkit::ScriptedModel;
102    use schemars::JsonSchema;
103    use serde::Deserialize;
104
105    use crate::{Agent, StructuredAgentError};
106
107    #[derive(Debug, Deserialize, Eq, JsonSchema, PartialEq)]
108    struct Answer {
109        value: u32,
110    }
111
112    fn events(text: &str) -> Vec<ModelStreamEvent> {
113        vec![
114            ModelStreamEvent::ResponseStarted {
115                id: Some("response".into()),
116                model: ModelRef::new("test", "scripted"),
117            },
118            ModelStreamEvent::ContentPartCompleted {
119                index: 0,
120                part: ContentPart::text(text),
121            },
122            ModelStreamEvent::ResponseCompleted {
123                finish_reason: FinishReason::Stop,
124                provider_metadata: BTreeMap::new(),
125            },
126        ]
127    }
128
129    fn run() -> RunContext {
130        RunContext::root(BudgetTracker::new(Budget::default()), CapabilitySet::new())
131    }
132
133    #[test]
134    fn typed_agent_uses_one_type_for_schema_and_decode() {
135        let model = ScriptedModel::new();
136        model.enqueue(events("{\"value\":42}"));
137        let agent = Agent::builder(
138            "typed",
139            Arc::new(model.clone()),
140            ModelRef::new("test", "scripted"),
141        )
142        .build_structured::<Answer>("answer")
143        .unwrap();
144        let run = run();
145
146        let typed = futures_executor::block_on(agent.run("answer", &run)).unwrap();
147
148        assert_eq!(typed.output, Answer { value: 42 });
149        let requests = model.recorded_requests();
150        let OutputFormat::JsonSchema { name, strict, .. } = &requests[0].output_format else {
151            panic!("expected JSON-schema output");
152        };
153        assert_eq!(name, "answer");
154        assert!(*strict);
155    }
156
157    #[test]
158    fn typed_agent_surfaces_local_decode_failure_separately() {
159        let model = ScriptedModel::new();
160        model.enqueue(events("{\"value\":\"wrong\"}"));
161        let agent = Agent::new("typed", Arc::new(model), ModelRef::new("test", "scripted"))
162            .into_structured::<Answer>("answer");
163        let run = run();
164
165        let error = futures_executor::block_on(agent.run("answer", &run)).unwrap_err();
166
167        assert!(matches!(
168            error,
169            StructuredAgentError::Output(ref output)
170                if output.kind == StructuredOutputErrorKind::InvalidOutput
171        ));
172    }
173}
runifold_agent/structured.rs

runifold_agent/
structured.rs
