Skip to main content

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    /// Terminal validation and bounded repair happen inside the canonical
70    /// Agent loop before a completed checkpoint is committed. The final local
71    /// decode remains as a defensive invariant check.
72    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}