Skip to main content

machi_runtime/
schema.rs

1//! JSON Schema validation for structured agent outputs.
2
3use jsonschema::Validator;
4use machi_types::{ErrorCode, MachiError};
5use serde_json::Value;
6
7/// Max corrective re-samples when structured output fails validation.
8pub const STRUCTURED_OUTPUT_MAX_RETRIES: u32 = 3;
9
10/// Compile a JSON Schema once per turn (or once per agent definition).
11///
12/// # Errors
13///
14/// Returns [`ErrorCode::RuntimeStructuredOutput`] when the schema itself is invalid.
15pub fn compile_schema(schema: &Value) -> Result<Validator, MachiError> {
16    Validator::new(schema).map_err(|e| {
17        MachiError::new(
18            ErrorCode::RuntimeStructuredOutput,
19            format!("invalid output schema: {e}"),
20        )
21    })
22}
23
24/// Parse model text as JSON and validate against a compiled schema.
25///
26/// # Errors
27///
28/// Returns a human-readable validation error suitable for model feedback.
29pub fn validate_structured_output(validator: &Validator, raw: &str) -> Result<Value, String> {
30    let value: Value = serde_json::from_str(raw.trim())
31        .map_err(|e| format!("model output was not valid JSON: {e}"))?;
32    validator
33        .validate(&value)
34        .map_err(|e| format!("output does not match the required schema: {e}"))?;
35    Ok(value)
36}
37
38/// Build a corrective user reminder after a schema failure.
39#[must_use]
40pub fn schema_retry_reminder(error: &str) -> String {
41    format!(
42        "Your previous response failed structured-output validation:\n{error}\n\
43         Reply with JSON only that satisfies the required schema."
44    )
45}
46
47#[cfg(test)]
48mod tests {
49    use serde_json::json;
50
51    use super::*;
52
53    #[test]
54    fn accepts_valid() {
55        let schema = json!({
56            "type": "object",
57            "properties": { "ok": { "type": "boolean" } },
58            "required": ["ok"],
59            "additionalProperties": false
60        });
61        let v = compile_schema(&schema).expect("schema");
62        let out = validate_structured_output(&v, r#"{"ok": true}"#).expect("ok");
63        assert_eq!(out.get("ok").and_then(Value::as_bool), Some(true));
64    }
65
66    #[test]
67    fn rejects_invalid() {
68        let schema = json!({
69            "type": "object",
70            "properties": { "ok": { "type": "boolean" } },
71            "required": ["ok"]
72        });
73        let v = compile_schema(&schema).expect("schema");
74        let err = validate_structured_output(&v, r#"{"ok": "nope"}"#).expect_err("bad");
75        assert!(err.contains("schema") || err.contains("type"), "{err}");
76    }
77}