Skip to main content

yapi_ai/
validation.rs

1//! Tool argument validation: coerce what models commonly get wrong, then check the
2//! JSON schema with TypeBox's error messages.
3//!
4//! Port of `packages/ai/src/utils/validation.ts` in pi `v1.0.0`. The checker covers
5//! the JSON Schema keywords tool schemas use. `$ref`, `format` and `pattern` are
6//! accepted without checking.
7
8use serde_json::{Map, Value};
9use yapi_types::message::{ToolCall, ToolDeclaration};
10
11/// Validated, coerced arguments for `call`, or the error pi reports to the model.
12pub fn validate_tool_arguments(tool: &ToolDeclaration, call: &ToolCall) -> Result<Value, String> {
13    let mut args = Value::Object(call.arguments.clone());
14    normalize_optional_nulls(&mut args, &tool.parameters);
15    let args = coerce(args, &tool.parameters);
16    let mut errors = Vec::new();
17    check(&args, &tool.parameters, "", &mut errors);
18    if errors.is_empty() {
19        return Ok(args);
20    }
21    let lines: Vec<String> = errors
22        .iter()
23        .map(|(path, message)| format!("  - {path}: {message}"))
24        .collect();
25    let received = yapi_types::json::to_string_pretty(&call.arguments, "  ").unwrap_or_default();
26    Err(format!(
27        "Validation failed for tool \"{}\":\n{}\n\nReceived arguments:\n{received}",
28        call.name,
29        lines.join("\n")
30    ))
31}
32
33/// The types a schema's `type` names: one string or an array of them.
34pub(crate) fn schema_types(schema: &Value) -> Vec<&str> {
35    match schema.get("type") {
36        Some(Value::String(kind)) => vec![kind.as_str()],
37        Some(Value::Array(kinds)) => kinds.iter().filter_map(Value::as_str).collect(),
38        _ => Vec::new(),
39    }
40}
41
42/// The schemas of an array's items: positional for a tuple `items`, else the one
43/// `items` schema for every item.
44fn item_schemas(schema: &Value) -> impl Iterator<Item = &Value> {
45    let (tuple, single) = match schema.get("items") {
46        Some(Value::Array(schemas)) => (schemas.as_slice(), None),
47        single => (&[][..], single),
48    };
49    tuple.iter().chain(std::iter::from_fn(move || single))
50}
51
52fn matches_type(value: &Value, kind: &str) -> bool {
53    match kind {
54        "number" => value.is_number(),
55        "integer" => value.as_f64().is_some_and(|n| n.fract() == 0.0),
56        "boolean" => value.is_boolean(),
57        "string" => value.is_string(),
58        "null" => value.is_null(),
59        "array" => value.is_array(),
60        "object" => value.is_object(),
61        _ => false,
62    }
63}
64
65fn coerce_primitive(value: &Value, kind: &str) -> Option<Value> {
66    match (kind, value) {
67        ("number" | "integer", Value::Null) => Some(Value::from(0)),
68        ("number" | "integer", Value::String(text)) if !text.trim().is_empty() => {
69            let parsed: f64 = text.trim().parse().ok()?;
70            (parsed.is_finite() && (kind == "number" || parsed.fract() == 0.0))
71                .then(|| yapi_types::json::number(parsed))
72        }
73        ("number" | "integer", Value::Bool(flag)) => Some(Value::from(u8::from(*flag))),
74        ("boolean", Value::Null) => Some(Value::Bool(false)),
75        ("boolean", Value::String(text)) if text == "true" || text == "false" => {
76            Some(Value::Bool(text == "true"))
77        }
78        ("boolean", Value::Number(n)) if n.as_f64() == Some(1.0) || n.as_f64() == Some(0.0) => {
79            Some(Value::Bool(n.as_f64() == Some(1.0)))
80        }
81        ("string", Value::Null) => Some(Value::String(String::new())),
82        ("string", Value::Number(n)) => Some(Value::String(yapi_types::json::to_string(n).ok()?)),
83        ("string", Value::Bool(flag)) => Some(Value::String(flag.to_string())),
84        ("null", Value::String(text)) if text.is_empty() => Some(Value::Null),
85        ("null", Value::Bool(false)) => Some(Value::Null),
86        ("null", Value::Number(n)) if n.as_f64() == Some(0.0) => Some(Value::Null),
87        _ => None,
88    }
89}
90
91fn is_valid(value: &Value, schema: &Value) -> bool {
92    let mut errors = Vec::new();
93    check(value, schema, "", &mut errors);
94    errors.is_empty()
95}
96
97fn coerce_union(value: Value, schemas: &[Value]) -> Value {
98    if schemas.iter().any(|schema| is_valid(&value, schema)) {
99        return value;
100    }
101    for schema in schemas {
102        let candidate = coerce(value.clone(), schema);
103        if is_valid(&candidate, schema) {
104            return candidate;
105        }
106    }
107    value
108}
109
110fn coerce(value: Value, schema: &Value) -> Value {
111    let mut value = value;
112    if let Some(all_of) = schema.get("allOf").and_then(Value::as_array) {
113        for nested in all_of {
114            value = coerce(value, nested);
115        }
116    }
117    for key in ["anyOf", "oneOf"] {
118        if let Some(variants) = schema.get(key).and_then(Value::as_array) {
119            value = coerce_union(value, variants);
120        }
121    }
122    let types = schema_types(schema);
123    let matches_member = types.len() > 1 && types.iter().any(|kind| matches_type(&value, kind));
124    if !types.is_empty() && !matches_member {
125        for kind in &types {
126            if let Some(candidate) = coerce_primitive(&value, kind)
127                && candidate != value
128            {
129                value = candidate;
130                break;
131            }
132        }
133    }
134    if types.contains(&"object")
135        && let Value::Object(object) = &mut value
136    {
137        let properties = schema.get("properties").and_then(Value::as_object);
138        for (key, property_schema) in properties.into_iter().flatten() {
139            if let Some(property) = object.get_mut(key) {
140                *property = coerce(std::mem::take(property), property_schema);
141            }
142        }
143        if let Some(additional) = schema.get("additionalProperties").filter(|v| v.is_object()) {
144            for (key, property) in object.iter_mut() {
145                if !properties.is_some_and(|properties| properties.contains_key(key)) {
146                    *property = coerce(std::mem::take(property), additional);
147                }
148            }
149        }
150    }
151    if types.contains(&"array")
152        && let Value::Array(items) = &mut value
153    {
154        for (item, item_schema) in items.iter_mut().zip(item_schemas(schema)) {
155            *item = coerce(std::mem::take(item), item_schema);
156        }
157    }
158    value
159}
160
161/// Drops `null` for optional properties whose schema does not allow null, as
162/// strict-mode providers send them.
163fn normalize_optional_nulls(value: &mut Value, schema: &Value) {
164    match value {
165        Value::Array(items) => {
166            for (item, item_schema) in items.iter_mut().zip(item_schemas(schema)) {
167                normalize_optional_nulls(item, item_schema);
168            }
169        }
170        Value::Object(object) => {
171            let Some(properties) = schema.get("properties").and_then(Value::as_object) else {
172                return;
173            };
174            let required = required_keys(schema);
175            for (key, property_schema) in properties {
176                let Some(property) = object.get_mut(key) else {
177                    continue;
178                };
179                if property.is_null()
180                    && !required.contains(&key.as_str())
181                    && property_schema.get("$ref").is_none()
182                    && !is_valid(&Value::Null, property_schema)
183                {
184                    object.shift_remove(key);
185                } else {
186                    normalize_optional_nulls(property, property_schema);
187                }
188            }
189        }
190        _ => {}
191    }
192}
193
194fn required_keys(schema: &Value) -> Vec<&str> {
195    schema
196        .get("required")
197        .and_then(Value::as_array)
198        .map(|keys| keys.iter().filter_map(Value::as_str).collect())
199        .unwrap_or_default()
200}
201
202fn join_path(path: &str, key: &str) -> String {
203    if path.is_empty() {
204        key.to_owned()
205    } else {
206        format!("{path}.{key}")
207    }
208}
209
210fn display_path(path: &str) -> String {
211    if path.is_empty() {
212        "root".to_owned()
213    } else {
214        path.to_owned()
215    }
216}
217
218fn js_number(value: &Value) -> String {
219    yapi_types::json::stringify(value)
220}
221
222/// Appends `(path, message)` for every violation of `schema` by `value`.
223fn check(value: &Value, schema: &Value, path: &str, errors: &mut Vec<(String, String)>) {
224    let Some(schema_object) = schema.as_object() else {
225        if schema == &Value::Bool(false) {
226            errors.push((display_path(path), "schema is false".into()));
227        }
228        return;
229    };
230    let types = schema_types(schema);
231    if !types.is_empty() && !types.iter().any(|kind| matches_type(value, kind)) {
232        let message = if types.len() == 1 {
233            format!("must be {}", types[0])
234        } else {
235            format!("must be either {}", types.join(" or "))
236        };
237        errors.push((display_path(path), message));
238        return;
239    }
240    if let Some(constant) = schema_object.get("const")
241        && constant != value
242    {
243        errors.push((display_path(path), "must be equal to constant".into()));
244    }
245    if let Some(options) = schema_object.get("enum").and_then(Value::as_array)
246        && !options.contains(value)
247    {
248        errors.push((
249            display_path(path),
250            "must be equal to one of the allowed values".into(),
251        ));
252    }
253    if let Some(variants) = schema_object.get("anyOf").and_then(Value::as_array)
254        && !variants.iter().any(|variant| is_valid(value, variant))
255    {
256        errors.push((display_path(path), "must match a schema in anyOf".into()));
257    }
258    if let Some(variants) = schema_object.get("oneOf").and_then(Value::as_array)
259        && variants
260            .iter()
261            .filter(|variant| is_valid(value, variant))
262            .count()
263            != 1
264    {
265        errors.push((
266            display_path(path),
267            "must match exactly one schema in oneOf".into(),
268        ));
269    }
270    if let Some(all_of) = schema_object.get("allOf").and_then(Value::as_array) {
271        for nested in all_of {
272            check(value, nested, path, errors);
273        }
274    }
275    match value {
276        Value::Number(n) => {
277            let n = n.as_f64().unwrap_or_default();
278            let limit = |key: &str| schema_object.get(key).and_then(Value::as_f64);
279            let bound =
280                |key: &str, comparison: &str, ok: bool, errors: &mut Vec<(String, String)>| {
281                    if !ok {
282                        errors.push((
283                            display_path(path),
284                            format!("must be {comparison} {}", js_number(&schema_object[key])),
285                        ));
286                    }
287                };
288            if let Some(min) = limit("minimum") {
289                bound("minimum", ">=", n >= min, errors);
290            }
291            if let Some(max) = limit("maximum") {
292                bound("maximum", "<=", n <= max, errors);
293            }
294            if let Some(min) = limit("exclusiveMinimum") {
295                bound("exclusiveMinimum", ">", n > min, errors);
296            }
297            if let Some(max) = limit("exclusiveMaximum") {
298                bound("exclusiveMaximum", "<", n < max, errors);
299            }
300        }
301        Value::String(text) => {
302            let length = text.chars().count() as u64;
303            if let Some(min) = schema_object.get("minLength").and_then(Value::as_u64)
304                && length < min
305            {
306                errors.push((
307                    display_path(path),
308                    format!("must not have fewer than {min} characters"),
309                ));
310            }
311            if let Some(max) = schema_object.get("maxLength").and_then(Value::as_u64)
312                && length > max
313            {
314                errors.push((
315                    display_path(path),
316                    format!("must not have more than {max} characters"),
317                ));
318            }
319        }
320        Value::Array(items) => {
321            if let Some(min) = schema_object.get("minItems").and_then(Value::as_u64)
322                && (items.len() as u64) < min
323            {
324                errors.push((
325                    display_path(path),
326                    format!("must not have fewer than {min} items"),
327                ));
328            }
329            if let Some(max) = schema_object.get("maxItems").and_then(Value::as_u64)
330                && (items.len() as u64) > max
331            {
332                errors.push((
333                    display_path(path),
334                    format!("must not have more than {max} items"),
335                ));
336            }
337            for (index, (item, item_schema)) in items.iter().zip(item_schemas(schema)).enumerate() {
338                check(
339                    item,
340                    item_schema,
341                    &join_path(path, &index.to_string()),
342                    errors,
343                );
344            }
345        }
346        Value::Object(object) => check_object(object, schema_object, path, errors),
347        _ => {}
348    }
349}
350
351fn check_object(
352    object: &Map<String, Value>,
353    schema: &Map<String, Value>,
354    path: &str,
355    errors: &mut Vec<(String, String)>,
356) {
357    let missing: Vec<&str> = schema
358        .get("required")
359        .and_then(Value::as_array)
360        .into_iter()
361        .flatten()
362        .filter_map(Value::as_str)
363        .filter(|key| !object.contains_key(*key))
364        .collect();
365    if let Some(first) = missing.first() {
366        errors.push((
367            join_path(path, first),
368            format!("must have required properties {}", missing.join(", ")),
369        ));
370    }
371    let properties = schema.get("properties").and_then(Value::as_object);
372    for (key, property_schema) in properties.into_iter().flatten() {
373        if let Some(property) = object.get(key) {
374            check(property, property_schema, &join_path(path, key), errors);
375        }
376    }
377    let extra = object
378        .iter()
379        .filter(|(key, _)| !properties.is_some_and(|properties| properties.contains_key(*key)));
380    match schema.get("additionalProperties") {
381        Some(Value::Bool(false)) => {
382            if extra.count() > 0 {
383                errors.push((
384                    display_path(path),
385                    "must not have additional properties".into(),
386                ));
387            }
388        }
389        Some(additional @ Value::Object(_)) => {
390            for (key, property) in extra {
391                check(property, additional, &join_path(path, key), errors);
392            }
393        }
394        _ => {}
395    }
396}
397
398#[cfg(test)]
399mod tests {
400    use super::*;
401    use serde_json::json;
402
403    fn tool(parameters: Value) -> ToolDeclaration {
404        ToolDeclaration {
405            name: "read".into(),
406            description: String::new(),
407            parameters,
408            constrained_sampling: None,
409        }
410    }
411
412    fn call(arguments: Value) -> ToolCall {
413        ToolCall {
414            id: "1".into(),
415            name: "read".into(),
416            arguments: arguments.as_object().cloned().unwrap(),
417            thought_signature: None,
418            namespace: None,
419        }
420    }
421
422    fn read_schema() -> Value {
423        json!({"type":"object","required":["path"],"properties":{
424            "path":{"type":"string"},"offset":{"type":"number"},"limit":{"type":"number"}}})
425    }
426
427    #[test]
428    fn coerces_drops_strict_nulls() {
429        let args = validate_tool_arguments(
430            &tool(read_schema()),
431            &call(json!({"path": "a.txt", "offset": "5", "limit": null})),
432        )
433        .unwrap();
434        assert_eq!(args, json!({"path": "a.txt", "offset": 5}));
435    }
436
437    #[test]
438    fn reports_typebox_style_errors() {
439        let err = validate_tool_arguments(&tool(read_schema()), &call(json!({"offset": [1]})))
440            .unwrap_err();
441        assert_eq!(
442            err,
443            "Validation failed for tool \"read\":\n  - path: must have required properties path\n  - offset: must be number\n\nReceived arguments:\n{\n  \"offset\": [\n    1\n  ]\n}"
444        );
445    }
446    #[test]
447    fn coerces_tuple_list_items() {
448        let schema = json!({"type":"object","properties":{
449            "pair":{"type":"array","items":[{"type":"number"},{"type":"string"}]},
450            "list":{"type":"array","items":{"type":"number"}}}});
451        let args = validate_tool_arguments(
452            &tool(schema),
453            &call(json!({"pair": ["1", 2, "x"], "list": ["3", "4"]})),
454        )
455        .unwrap();
456        assert_eq!(args, json!({"pair": [1, "2", "x"], "list": [3, 4]}));
457    }
458}