Skip to main content

tool_arg_coerce/
lib.rs

1//! # tool-arg-coerce
2//!
3//! Fix common type slips in LLM-generated tool arguments.
4//!
5//! LLMs often emit `"42"` instead of `42` for an integer argument,
6//! `["x"]` instead of `"x"` when the schema wants a scalar, or
7//! `"true"` for a bool. This crate gives you a `coerce_one` function
8//! that fixes the obvious cases for a single value, and `coerce_args`
9//! that walks an object against a typed schema.
10//!
11//! ## Example
12//!
13//! ```
14//! use tool_arg_coerce::{coerce_one, Type};
15//! use serde_json::json;
16//!
17//! assert_eq!(coerce_one(json!("42"), Type::Int), Some(json!(42)));
18//! assert_eq!(coerce_one(json!("true"), Type::Bool), Some(json!(true)));
19//! assert_eq!(coerce_one(json!(["only"]), Type::String), Some(json!("only")));
20//! ```
21
22#![deny(missing_docs)]
23
24use serde_json::{json, Value};
25
26/// Target type for coercion.
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub enum Type {
29    /// Signed 64-bit integer.
30    Int,
31    /// f64.
32    Float,
33    /// Bool.
34    Bool,
35    /// String.
36    String,
37}
38
39/// Coerce one value to `ty`. Returns `None` if no obvious fix exists.
40pub fn coerce_one(v: Value, ty: Type) -> Option<Value> {
41    // Single-element array -> scalar.
42    let v = if let Value::Array(ref arr) = v {
43        if arr.len() == 1 {
44            arr[0].clone()
45        } else {
46            v
47        }
48    } else {
49        v
50    };
51
52    match ty {
53        Type::Int => match v {
54            Value::Number(n) => n.as_i64().map(|i| json!(i)),
55            Value::String(s) => s.trim().parse::<i64>().ok().map(|i| json!(i)),
56            Value::Bool(b) => Some(json!(if b { 1 } else { 0 })),
57            _ => None,
58        },
59        Type::Float => match v {
60            Value::Number(n) => n.as_f64().map(|f| json!(f)),
61            Value::String(s) => s.trim().parse::<f64>().ok().map(|f| json!(f)),
62            _ => None,
63        },
64        Type::Bool => match v {
65            Value::Bool(b) => Some(json!(b)),
66            Value::String(s) => match s.trim().to_ascii_lowercase().as_str() {
67                "true" | "yes" | "y" | "1" => Some(json!(true)),
68                "false" | "no" | "n" | "0" => Some(json!(false)),
69                _ => None,
70            },
71            Value::Number(n) => n.as_i64().map(|i| json!(i != 0)),
72            _ => None,
73        },
74        Type::String => Some(match v {
75            Value::String(s) => json!(s),
76            Value::Number(n) => json!(n.to_string()),
77            Value::Bool(b) => json!(b.to_string()),
78            Value::Null => json!(""),
79            _ => json!(v.to_string()),
80        }),
81    }
82}
83
84/// Walk a JSON object, coercing fields named in `schema`.
85///
86/// Fields not in `schema` pass through unchanged.
87pub fn coerce_args(mut args: Value, schema: &[(&str, Type)]) -> Value {
88    if let Value::Object(map) = &mut args {
89        for (name, ty) in schema {
90            if let Some(slot) = map.remove(*name) {
91                if let Some(fixed) = coerce_one(slot.clone(), *ty) {
92                    map.insert((*name).to_string(), fixed);
93                } else {
94                    // Re-insert the original if we can't fix it.
95                    map.insert((*name).to_string(), slot);
96                }
97            }
98        }
99    }
100    args
101}