ferrum-sampler 0.8.8

//! Minimal JSON Schema → regex translator for structured output.
//!
//! Accepts a subset of JSON Schema sufficient for most structured-output
//! use cases — the classic "give me a typed object back" prompt flow
//! that OpenAI's `response_format = json_schema` enables. Coverage:
//!
//!   - `{ "type": "string" }` → `"[^"\\]*"`; `minLength`/`maxLength`
//!     become bounded repetitions. No escape handling yet — models very
//!     rarely emit backslashes in short structured answers.
//!   - `{ "type": "integer" }` → `-?\d+`
//!   - `{ "type": "number" }`  → `-?\d+(\.\d+)?([eE][+-]?\d+)?`
//!   - `{ "type": "boolean" }` → `true|false`
//!   - `{ "type": "null" }`    → `null`
//!   - `{ "const": value }`     → the exact JSON literal for `value`
//!   - `{ "enum": [...] }`     → alternation of JSON-quoted values
//!   - `{ "type": "array", "items": T }` → `\[\s*(T(\s*,\s*T)*)?\s*\]`
//!   - `{ "type": "object", "properties": {...}, "required": [...] }`
//!     → required fields in declared order: `\{ "k1": T1, "k2": T2, ... \}`
//!     Optional fields are skipped to keep the DFA small.
//!
//! Not supported (falls back to JsonObject-style unconstrained JSON):
//!   - `$ref`, `oneOf`, `anyOf`, `allOf`, `not`
//!   - nested-depth limits, numeric min/max constraints, string patterns
//!   - additionalProperties, propertyNames
//!
//! When the translator can't handle a schema, the caller should skip
//! guided decoding rather than silently produce wrong output.

use ferrum_types::{FerrumError, Result};
use serde_json::Value;

/// Compile a JSON Schema (serialised as a JSON string) into a regex
/// pattern suitable for feeding into `RegexGuidedProcessor`.
///
/// The returned pattern:
///   * allows a small finite amount of leading/trailing JSON whitespace so
///     BPE tokenisers that prepend a space to the first generated token still
///     transition cleanly, without letting the model generate whitespace
///     forever under a hard mask.
///   * is anchored to end internally by the processor (which wraps
///     with `^(?:...)\z`).
pub fn schema_to_regex(schema_json: &str) -> Result<String> {
    let schema: Value = serde_json::from_str(schema_json).map_err(|e| {
        FerrumError::invalid_request(format!("response_format.schema is not valid JSON: {e}"))
    })?;
    let inner = translate(&schema)?;
    Ok(format!(r"[ \t\r\n]{{0,8}}{inner}[ \t\r\n]{{0,8}}"))
}

fn translate(node: &Value) -> Result<String> {
    // `enum` takes precedence over `type`.
    if let Some(en) = node.get("enum").and_then(|v| v.as_array()) {
        return enum_pattern(en);
    }

    // `const` constrains the complete JSON value, regardless of a broader
    // (or inconsistent) `type` annotation on the same schema node.
    if let Some(value) = node.get("const") {
        return literal_pattern(value, "const");
    }

    let ty = node.get("type").and_then(|v| v.as_str());
    match ty {
        Some("string") => string_pattern(node),
        Some("integer") => Ok(r"-?\d+".to_string()),
        Some("number") => Ok(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?".to_string()),
        Some("boolean") => Ok("(?:true|false)".to_string()),
        Some("null") => Ok("null".to_string()),
        Some("array") => array_pattern(node),
        Some("object") => object_pattern(node),
        Some(other) => Err(FerrumError::invalid_request(format!(
            "unsupported JSON Schema type '{other}' in response_format"
        ))),
        None => Err(FerrumError::invalid_request(
            "JSON Schema node missing 'type' field and no 'enum' present",
        )),
    }
}

fn string_pattern(node: &Value) -> Result<String> {
    const MAX_SUPPORTED_STRING_LENGTH: u64 = 1024;
    let min = node.get("minLength").and_then(Value::as_u64).unwrap_or(0);
    let max = node.get("maxLength").and_then(Value::as_u64);
    if let Some(max) = max {
        if max > MAX_SUPPORTED_STRING_LENGTH {
            return Err(FerrumError::invalid_request(format!(
                "response_format string maxLength {max} exceeds supported limit {MAX_SUPPORTED_STRING_LENGTH}"
            )));
        }
        if min > max {
            return Err(FerrumError::invalid_request(format!(
                "response_format string minLength {min} exceeds maxLength {max}"
            )));
        }
        return Ok(format!(r#""[^"]{{{min},{max}}}""#));
    }
    if min > 0 {
        return Ok(format!(r#""[^"]{{{min},}}""#));
    }
    Ok(r#""[^"]*""#.to_string())
}

fn enum_pattern(values: &[Value]) -> Result<String> {
    if values.is_empty() {
        return Err(FerrumError::invalid_request(
            "response_format enum must have at least one value",
        ));
    }
    let mut alts: Vec<String> = Vec::with_capacity(values.len());
    for v in values {
        alts.push(literal_pattern(v, "enum")?);
    }
    Ok(format!("(?:{})", alts.join("|")))
}

fn literal_pattern(value: &Value, keyword: &str) -> Result<String> {
    // Keep constrained output independent of serde_json's workspace-unified
    // `preserve_order` feature. A canonical key order also bounds object const
    // to one compact, semantically valid JSON representation.
    let canonical = canonical_json_literal(value);
    let literal = serde_json::to_string(&canonical).map_err(|e| {
        FerrumError::invalid_request(format!("{keyword} value not JSON-serialisable: {e}"))
    })?;
    Ok(regex_escape(&literal))
}

fn canonical_json_literal(value: &Value) -> Value {
    match value {
        Value::Array(values) => Value::Array(
            values
                .iter()
                .map(canonical_json_literal)
                .collect::<Vec<_>>(),
        ),
        Value::Object(values) => {
            let mut keys = values.keys().collect::<Vec<_>>();
            keys.sort_unstable();
            let mut canonical = serde_json::Map::with_capacity(values.len());
            for key in keys {
                canonical.insert(key.clone(), canonical_json_literal(&values[key]));
            }
            Value::Object(canonical)
        }
        value => value.clone(),
    }
}

fn array_pattern(node: &Value) -> Result<String> {
    let items_schema = node
        .get("items")
        .ok_or_else(|| FerrumError::invalid_request("array schema missing 'items'"))?;
    let item_pat = translate(items_schema)?;
    // `\s*` between tokens so minor whitespace variation doesn't break the match.
    Ok(format!(r"\[\s*(?:{item_pat}(?:\s*,\s*{item_pat})*)?\s*\]"))
}

fn object_pattern(node: &Value) -> Result<String> {
    let props = node
        .get("properties")
        .and_then(|v| v.as_object())
        .ok_or_else(|| FerrumError::invalid_request("object schema missing 'properties'"))?;
    let required: Vec<&str> = node
        .get("required")
        .and_then(|v| v.as_array())
        .map(|a| a.iter().filter_map(|v| v.as_str()).collect())
        .unwrap_or_default();

    // Emit required fields in the order listed. Optional fields dropped
    // for pattern compactness — still produces a superset of valid outputs
    // but model will usually follow the constrained form.
    let keys: Vec<&str> = if required.is_empty() {
        // No `required` → walk properties in insertion order. `properties`
        // is a `serde_json::Map` which is ordered in our dep config, so
        // this is deterministic.
        props.keys().map(String::as_str).collect()
    } else {
        required
    };

    if keys.is_empty() {
        return Ok(r"\{\s*\}".to_string());
    }

    let mut fields: Vec<(String, String)> = Vec::with_capacity(keys.len());
    for key in keys {
        let sub = props.get(key).ok_or_else(|| {
            FerrumError::invalid_request(format!(
                "required property '{key}' missing from 'properties'"
            ))
        })?;
        let sub_pat = translate(sub)?;
        let key_literal = regex_escape(&format!("\"{key}\""));
        fields.push((key_literal, sub_pat));
    }

    let field_pattern = object_field_order_pattern(&fields);
    Ok(format!(r"\{{{field_pattern}\s*\}}"))
}

fn object_field_order_pattern(fields: &[(String, String)]) -> String {
    const MAX_PERMUTED_OBJECT_FIELDS: usize = 6;
    if fields.len() > MAX_PERMUTED_OBJECT_FIELDS {
        return object_fields_sequence_pattern(fields);
    }

    let mut orders = Vec::new();
    let mut indices = (0..fields.len()).collect::<Vec<_>>();
    permute_indices(0, &mut indices, &mut orders);
    let alternatives = orders
        .iter()
        .map(|order| {
            let ordered = order
                .iter()
                .map(|&idx| fields[idx].clone())
                .collect::<Vec<_>>();
            object_fields_sequence_pattern(&ordered)
        })
        .collect::<Vec<_>>();
    if alternatives.len() == 1 {
        alternatives[0].clone()
    } else {
        format!("(?:{})", alternatives.join("|"))
    }
}

fn object_fields_sequence_pattern(fields: &[(String, String)]) -> String {
    fields
        .iter()
        .map(|(key_literal, sub_pat)| format!(r"\s*{key_literal}\s*:\s*{sub_pat}"))
        .collect::<Vec<_>>()
        .join(r"\s*,")
}

fn permute_indices(start: usize, indices: &mut [usize], out: &mut Vec<Vec<usize>>) {
    if start == indices.len() {
        out.push(indices.to_vec());
        return;
    }
    for i in start..indices.len() {
        indices.swap(start, i);
        permute_indices(start + 1, indices, out);
        indices.swap(start, i);
    }
}

fn regex_escape(s: &str) -> String {
    // regex-syntax::escape would be nicer but pulling that crate in for
    // a few metachars is overkill; explicit table matches its behaviour
    // for our literal inputs.
    let mut out = String::with_capacity(s.len() + 4);
    for ch in s.chars() {
        match ch {
            '\\' | '.' | '*' | '+' | '?' | '(' | ')' | '[' | ']' | '{' | '}' | '|' | '^' | '$'
            | '/' => {
                out.push('\\');
                out.push(ch);
            }
            _ => out.push(ch),
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    // Use the std `regex` crate in tests — it ships with easier full-match
    // semantics than raw regex_automata DFAs and isn't in the prod hot path.
    use regex_lite::Regex;

    fn compile(pat: &str) -> Regex {
        Regex::new(&format!("^(?:{pat})$")).expect("regex compiles")
    }

    #[test]
    fn string_schema_matches_quoted() {
        let re = compile(&schema_to_regex(r#"{"type":"string"}"#).unwrap());
        assert!(re.is_match("\"hello\""));
        assert!(re.is_match("\"\""));
        assert!(!re.is_match("hello"));
        assert!(!re.is_match("\"un\"closed"));
    }

    #[test]
    fn string_schema_honors_min_and_max_length() {
        let re =
            compile(&schema_to_regex(r#"{"type":"string","minLength":1,"maxLength":3}"#).unwrap());
        assert!(re.is_match("\"a\""));
        assert!(re.is_match("\"abc\""));
        assert!(!re.is_match("\"\""));
        assert!(!re.is_match("\"abcd\""));
    }

    #[test]
    fn integer_schema() {
        let re = compile(&schema_to_regex(r#"{"type":"integer"}"#).unwrap());
        assert!(re.is_match("0"));
        assert!(re.is_match("42"));
        assert!(re.is_match("-7"));
        assert!(!re.is_match("1.5"));
        assert!(!re.is_match("abc"));
    }

    #[test]
    fn number_schema_accepts_decimal_and_exp() {
        let re = compile(&schema_to_regex(r#"{"type":"number"}"#).unwrap());
        assert!(re.is_match("3"));
        assert!(re.is_match("3.14"));
        assert!(re.is_match("-0.001"));
        assert!(re.is_match("1e5"));
        assert!(re.is_match("1.2e-3"));
    }

    #[test]
    fn boolean_schema() {
        let re = compile(&schema_to_regex(r#"{"type":"boolean"}"#).unwrap());
        assert!(re.is_match("true"));
        assert!(re.is_match("false"));
        assert!(!re.is_match("True"));
    }

    #[test]
    fn enum_schema() {
        let re = compile(&schema_to_regex(r#"{"enum":["red","green","blue"]}"#).unwrap());
        assert!(re.is_match("\"red\""));
        assert!(re.is_match("\"blue\""));
        assert!(!re.is_match("\"yellow\""));
    }

    #[test]
    fn string_const_takes_precedence_over_type() {
        let re = compile(&schema_to_regex(r#"{"type":"string","const":"a.b*"}"#).unwrap());
        assert!(re.is_match(r#""a.b*""#));
        assert!(!re.is_match(r#""anything else""#));
    }

    #[test]
    fn number_const_takes_precedence_over_type() {
        let re = compile(&schema_to_regex(r#"{"type":"number","const":3.14}"#).unwrap());
        assert!(re.is_match("3.14"));
        assert!(!re.is_match("3"));
        assert!(!re.is_match("3.140"));
    }

    #[test]
    fn boolean_const_takes_precedence_over_type() {
        let re = compile(&schema_to_regex(r#"{"type":"boolean","const":true}"#).unwrap());
        assert!(re.is_match("true"));
        assert!(!re.is_match("false"));
    }

    #[test]
    fn null_const_does_not_require_type() {
        let re = compile(&schema_to_regex(r#"{"const":null}"#).unwrap());
        assert!(re.is_match("null"));
        assert!(!re.is_match("false"));
    }

    #[test]
    fn object_const_matches_exact_json_literal() {
        let re = compile(
            &schema_to_regex(r#"{"type":"object","const":{"kind":"ok","count":2,"ready":true}}"#)
                .unwrap(),
        );
        assert!(re.is_match(r#"{"count":2,"kind":"ok","ready":true}"#));
        assert!(!re.is_match(r#"{"count":2, "kind":"ok","ready":true}"#));
        assert!(!re.is_match(r#"{"kind":"ok","count":2,"ready":true}"#));
        assert!(!re.is_match(r#"{"count":3,"kind":"ok","ready":true}"#));
    }

    #[test]
    fn object_const_pattern_is_independent_of_schema_key_order() {
        let left =
            schema_to_regex(r#"{"const":{"kind":"ok","nested":{"z":1,"a":2},"ready":true}}"#)
                .unwrap();
        let right =
            schema_to_regex(r#"{"const":{"ready":true,"nested":{"a":2,"z":1},"kind":"ok"}}"#)
                .unwrap();

        assert_eq!(left, right);
    }

    #[test]
    fn array_of_integers() {
        let re =
            compile(&schema_to_regex(r#"{"type":"array","items":{"type":"integer"}}"#).unwrap());
        assert!(re.is_match("[]"));
        assert!(re.is_match("[1]"));
        assert!(re.is_match("[1, 2, 3]"));
        assert!(re.is_match("[-1, 0, 2]"));
        assert!(!re.is_match("[1.5]"));
        assert!(!re.is_match("[1, \"two\"]"));
    }

    #[test]
    fn object_with_required_fields() {
        let schema = r#"{
            "type": "object",
            "properties": {
                "name": {"type": "string"},
                "age": {"type": "integer"}
            },
            "required": ["name", "age"]
        }"#;
        let re = compile(&schema_to_regex(schema).unwrap());
        assert!(re.is_match(r#"{"name": "Alice", "age": 30}"#));
        assert!(re.is_match(r#"{ "name":"Bob" , "age":7 }"#));
        assert!(re.is_match(r#"{"age": 30, "name": "Alice"}"#));
        assert!(!re.is_match(r#"{"name": "Alice"}"#));
        assert!(!re.is_match(r#"{"age": 30}"#));
    }

    #[test]
    fn nested_object_and_array() {
        let schema = r#"{
            "type": "object",
            "properties": {
                "tags": {"type": "array", "items": {"type": "string"}},
                "count": {"type": "integer"}
            },
            "required": ["tags", "count"]
        }"#;
        let re = compile(&schema_to_regex(schema).unwrap());
        assert!(re.is_match(r#"{"tags": ["a", "b"], "count": 2}"#));
        assert!(re.is_match(r#"{"tags": [], "count": 0}"#));
        assert!(!re.is_match(r#"{"tags": ["a"], "count": "two"}"#));
    }

    #[test]
    fn unsupported_type_errors_clearly() {
        let err = schema_to_regex(r#"{"type":"mystery"}"#).unwrap_err();
        assert!(
            err.to_string().contains("unsupported JSON Schema type"),
            "got: {err}"
        );
    }
}