rullama_reasoning/
schema_validator.rs

1//! JSON Schema validation for LLM output, with a retry helper for agents
2//! that want to enforce structured responses.
3//!
4//! The [`SchemaValidator`] wraps a compiled `jsonschema` validator and
5//! exposes a single `validate` method that returns a list of human-readable
6//! violation strings. The [`retry_until_valid`] free function is a thin
7//! orchestrator for the common pattern: ask the producer for output,
8//! validate it against the schema, and on violation re-invoke the producer
9//! with a corrective instruction injected into the prompt.
10//!
11//! Only available behind the `schema-validation` feature so the reasoning
12//! crate stays light when callers don't need it. The module is gated at its
13//! declaration in `lib.rs`, so no inner `#![cfg]` is needed here.
14
15use std::future::Future;
16
17use anyhow::{Result, anyhow};
18use jsonschema::Validator;
19use serde_json::Value;
20
21/// Compiled JSON Schema validator. Cheap to clone — internally `Arc`-shared.
22#[derive(Clone)]
23pub struct SchemaValidator {
24    validator: std::sync::Arc<Validator>,
25}
26
27impl SchemaValidator {
28    /// Compile a schema from a `serde_json::Value`. Returns an error if the
29    /// schema itself is invalid (unparseable or not a valid draft).
30    pub fn new(schema: &Value) -> Result<Self> {
31        let validator =
32            jsonschema::validator_for(schema).map_err(|e| anyhow!("invalid JSON Schema: {e}"))?;
33        Ok(Self {
34            validator: std::sync::Arc::new(validator),
35        })
36    }
37
38    /// Validate `value` against the schema. Returns `Ok(())` if it passes;
39    /// otherwise a list of one human-readable error per violating
40    /// instance (`instance_path: error`).
41    pub fn validate(&self, value: &Value) -> std::result::Result<(), Vec<String>> {
42        let errors: Vec<String> = self
43            .validator
44            .iter_errors(value)
45            .map(|e| format!("{}: {e}", e.instance_path))
46            .collect();
47        if errors.is_empty() {
48            Ok(())
49        } else {
50            Err(errors)
51        }
52    }
53}
54
55/// Call `producer` up to `max_attempts` times, validating each result
56/// against `validator`. On violation, `producer` is invoked again with the
57/// list of error strings — the closure is expected to surface those to the
58/// LLM (e.g. via a corrective user message) so the next attempt converges.
59///
60/// Returns `Ok(value)` on the first attempt that validates. Returns the
61/// final attempt's value wrapped in an error containing all the errors
62/// after `max_attempts` exhaustion.
63pub async fn retry_until_valid<F, Fut>(
64    validator: &SchemaValidator,
65    max_attempts: usize,
66    mut producer: F,
67) -> Result<Value>
68where
69    F: FnMut(Option<Vec<String>>) -> Fut,
70    Fut: Future<Output = Result<Value>>,
71{
72    let mut last_errors: Option<Vec<String>> = None;
73    for attempt in 0..max_attempts.max(1) {
74        let value = producer(last_errors.clone()).await?;
75        match validator.validate(&value) {
76            Ok(()) => return Ok(value),
77            Err(errors) => {
78                tracing::warn!(
79                    attempt = attempt + 1,
80                    max = max_attempts,
81                    "schema validation failed: {} violation(s)",
82                    errors.len()
83                );
84                last_errors = Some(errors);
85            }
86        }
87    }
88    Err(anyhow!(
89        "schema validation failed after {} attempts: {}",
90        max_attempts,
91        last_errors.unwrap_or_default().join("; ")
92    ))
93}
94
95#[cfg(test)]
96mod tests {
97    use super::*;
98    use serde_json::json;
99
100    fn person_schema() -> Value {
101        json!({
102            "type": "object",
103            "properties": {
104                "name": { "type": "string" },
105                "age": { "type": "integer", "minimum": 0 }
106            },
107            "required": ["name", "age"]
108        })
109    }
110
111    #[test]
112    fn valid_value_passes() {
113        let v = SchemaValidator::new(&person_schema()).unwrap();
114        assert!(v.validate(&json!({"name": "Ada", "age": 36})).is_ok());
115    }
116
117    #[test]
118    fn missing_required_field_fails() {
119        let v = SchemaValidator::new(&person_schema()).unwrap();
120        let result = v.validate(&json!({"name": "Ada"}));
121        assert!(result.is_err());
122        assert!(result.unwrap_err().iter().any(|e| e.contains("age")));
123    }
124
125    #[test]
126    fn wrong_type_fails() {
127        let v = SchemaValidator::new(&person_schema()).unwrap();
128        let result = v.validate(&json!({"name": "Ada", "age": "old"}));
129        assert!(result.is_err());
130    }
131
132    #[test]
133    fn invalid_schema_returns_error() {
134        // `type` must be a string or array — a number is invalid.
135        let result = SchemaValidator::new(&json!({"type": 42}));
136        assert!(result.is_err());
137    }
138
139    #[tokio::test]
140    async fn retry_converges_on_valid() {
141        let schema = SchemaValidator::new(&person_schema()).unwrap();
142        let mut call = 0;
143        let result = retry_until_valid(&schema, 3, |errors| {
144            call += 1;
145            let value = if errors.is_none() {
146                json!({"name": "Ada"}) // missing age
147            } else {
148                json!({"name": "Ada", "age": 36})
149            };
150            async move { Ok(value) }
151        })
152        .await
153        .unwrap();
154        assert_eq!(result, json!({"name": "Ada", "age": 36}));
155        assert_eq!(call, 2);
156    }
157
158    #[tokio::test]
159    async fn retry_exhausts_then_errors() {
160        let schema = SchemaValidator::new(&person_schema()).unwrap();
161        let result = retry_until_valid(&schema, 2, |_errors| async {
162            Ok(json!({"name": "Ada"})) // always missing age
163        })
164        .await;
165        assert!(result.is_err());
166    }
167}
rullama_reasoning/schema_validator.rs

rullama_reasoning/
schema_validator.rs
