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}