tool_arg_defaults/
lib.rs

1/*!
2tool-arg-defaults: fill in missing kwargs on LLM-generated tool calls.
3
4LLMs often omit optional args. `ToolDefaults` merges per-tool default
5values into the LLM's args before execution. Caller-supplied keys always
6win — including explicit `null`.
7
8```rust
9use serde_json::json;
10use tool_arg_defaults::ToolDefaults;
11
12let mut defaults = ToolDefaults::new();
13defaults.register("search_web", json!({"timeout": 30, "max_results": 10}));
14
15// LLM only passed "q"
16let merged = defaults.apply("search_web", &json!({"q": "anthropic"}), false).unwrap();
17assert_eq!(merged["timeout"], json!(30));
18assert_eq!(merged["max_results"], json!(10));
19assert_eq!(merged["q"], json!("anthropic"));
20
21// Caller-supplied value wins
22let merged = defaults.apply("search_web", &json!({"q": "x", "timeout": 5}), false).unwrap();
23assert_eq!(merged["timeout"], json!(5));
24```
25*/
26
27use serde_json::{Map, Value};
28use std::collections::HashMap;
29
30// ---- error ----------------------------------------------------------------
31
32#[derive(Debug, Clone, PartialEq, Eq)]
33pub struct ToolNotRegisteredError(pub String);
34
35impl std::fmt::Display for ToolNotRegisteredError {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        write!(f, "tool {:?} is not registered", self.0)
38    }
39}
40
41impl std::error::Error for ToolNotRegisteredError {}
42
43// ---- ToolDefaults ---------------------------------------------------------
44
45/// Per-tool default kwargs for LLM-generated tool calls.
46///
47/// Defaults are stored as JSON objects (`serde_json::Map<String, Value>`).
48/// `apply` merges them into the caller-supplied args; caller-supplied keys win.
49/// `null` is a real value — passing `{"timeout": null}` will NOT be replaced
50/// by the registered default. To get the default, omit the key entirely.
51#[derive(Debug, Default, Clone)]
52pub struct ToolDefaults {
53    defaults: HashMap<String, Map<String, Value>>,
54}
55
56impl ToolDefaults {
57    pub fn new() -> Self {
58        Self::default()
59    }
60
61    // ---- registration -----------------------------------------------
62
63    /// Register (or overwrite) defaults for `tool_name`.
64    ///
65    /// `defaults` must be a JSON object (`Value::Object`); panics otherwise.
66    pub fn register(&mut self, tool_name: impl Into<String>, defaults: Value) {
67        let obj = match defaults {
68            Value::Object(m) => m,
69            other => panic!("defaults must be a JSON object, got {other:?}"),
70        };
71        self.defaults.insert(tool_name.into(), obj);
72    }
73
74    /// Register multiple tools at once from a JSON object of objects.
75    pub fn register_many(&mut self, map: Value) {
76        if let Value::Object(outer) = map {
77            for (name, val) in outer {
78                self.register(name, val);
79            }
80        }
81    }
82
83    /// Merge new defaults into an existing registration (or create one).
84    pub fn update(&mut self, tool_name: impl Into<String>, extra: Value) {
85        let name = tool_name.into();
86        let target = self.defaults.entry(name).or_default();
87        if let Value::Object(m) = extra {
88            target.extend(m);
89        }
90    }
91
92    /// Drop a tool's defaults. Returns true if it was registered.
93    pub fn unregister(&mut self, tool_name: &str) -> bool {
94        self.defaults.remove(tool_name).is_some()
95    }
96
97    // ---- inspection -------------------------------------------------
98
99    /// Sorted list of registered tool names.
100    pub fn tool_names(&self) -> Vec<&str> {
101        let mut names: Vec<&str> = self.defaults.keys().map(|s| s.as_str()).collect();
102        names.sort_unstable();
103        names
104    }
105
106    /// Copy of the defaults for `tool_name`.
107    pub fn defaults_for(&self, tool_name: &str) -> Option<Map<String, Value>> {
108        self.defaults.get(tool_name).cloned()
109    }
110
111    pub fn contains(&self, tool_name: &str) -> bool {
112        self.defaults.contains_key(tool_name)
113    }
114
115    // ---- core -------------------------------------------------------
116
117    /// Merge defaults into `args`. Caller-supplied keys win.
118    ///
119    /// `args` must be a JSON object or `Value::Null` (treated as empty).
120    /// Returns `Err(ToolNotRegisteredError)` if `strict=true` and tool not found.
121    /// If `strict=false` and tool not found, returns `args` unchanged.
122    pub fn apply(
123        &self,
124        tool_name: &str,
125        args: &Value,
126        strict: bool,
127    ) -> Result<Value, ToolNotRegisteredError> {
128        let supplied: Map<String, Value> = match args {
129            Value::Object(m) => m.clone(),
130            Value::Null => Map::new(),
131            _ => Map::new(),
132        };
133        let Some(tool_defaults) = self.defaults.get(tool_name) else {
134            if strict {
135                return Err(ToolNotRegisteredError(tool_name.to_owned()));
136            }
137            return Ok(Value::Object(supplied));
138        };
139        // Start from defaults, then overwrite with caller-supplied values.
140        let mut merged = tool_defaults.clone();
141        for (k, v) in supplied {
142            merged.insert(k, v);
143        }
144        Ok(Value::Object(merged))
145    }
146}
147
148// ---- tests ----------------------------------------------------------------
149
150#[cfg(test)]
151mod tests {
152    use super::*;
153    use serde_json::json;
154
155    #[test]
156    fn fill_missing_args() {
157        let mut d = ToolDefaults::new();
158        d.register("search", json!({"timeout": 30, "max_results": 10}));
159        let out = d.apply("search", &json!({"q": "hello"}), false).unwrap();
160        assert_eq!(out["q"], json!("hello"));
161        assert_eq!(out["timeout"], json!(30));
162        assert_eq!(out["max_results"], json!(10));
163    }
164
165    #[test]
166    fn caller_wins_on_conflict() {
167        let mut d = ToolDefaults::new();
168        d.register("search", json!({"timeout": 30}));
169        let out = d
170            .apply("search", &json!({"timeout": 5, "q": "x"}), false)
171            .unwrap();
172        assert_eq!(out["timeout"], json!(5));
173    }
174
175    #[test]
176    fn null_is_real_value() {
177        let mut d = ToolDefaults::new();
178        d.register("fetch", json!({"follow_redirects": true}));
179        let out = d
180            .apply("fetch", &json!({"follow_redirects": null}), false)
181            .unwrap();
182        assert_eq!(out["follow_redirects"], Value::Null);
183    }
184
185    #[test]
186    fn unknown_tool_non_strict() {
187        let d = ToolDefaults::new();
188        let out = d.apply("ghost", &json!({"x": 1}), false).unwrap();
189        assert_eq!(out["x"], json!(1));
190    }
191
192    #[test]
193    fn unknown_tool_strict() {
194        let d = ToolDefaults::new();
195        let err = d.apply("ghost", &json!({}), true).unwrap_err();
196        assert_eq!(err.0, "ghost");
197    }
198
199    #[test]
200    fn empty_args_fills_all_defaults() {
201        let mut d = ToolDefaults::new();
202        d.register("tool", json!({"a": 1, "b": 2}));
203        let out = d.apply("tool", &json!({}), false).unwrap();
204        assert_eq!(out["a"], json!(1));
205        assert_eq!(out["b"], json!(2));
206    }
207
208    #[test]
209    fn null_args_treated_as_empty() {
210        let mut d = ToolDefaults::new();
211        d.register("tool", json!({"x": 99}));
212        let out = d.apply("tool", &Value::Null, false).unwrap();
213        assert_eq!(out["x"], json!(99));
214    }
215
216    #[test]
217    fn register_overwrites() {
218        let mut d = ToolDefaults::new();
219        d.register("t", json!({"a": 1}));
220        d.register("t", json!({"a": 99, "b": 2}));
221        let out = d.apply("t", &json!({}), false).unwrap();
222        assert_eq!(out["a"], json!(99));
223        assert_eq!(out["b"], json!(2));
224    }
225
226    #[test]
227    fn update_merges() {
228        let mut d = ToolDefaults::new();
229        d.register("t", json!({"a": 1}));
230        d.update("t", json!({"b": 2}));
231        let out = d.apply("t", &json!({}), false).unwrap();
232        assert_eq!(out["a"], json!(1));
233        assert_eq!(out["b"], json!(2));
234    }
235
236    #[test]
237    fn update_creates_if_absent() {
238        let mut d = ToolDefaults::new();
239        d.update("new_tool", json!({"x": 42}));
240        let out = d.apply("new_tool", &json!({}), false).unwrap();
241        assert_eq!(out["x"], json!(42));
242    }
243
244    #[test]
245    fn unregister_returns_true_if_was_registered() {
246        let mut d = ToolDefaults::new();
247        d.register("t", json!({"a": 1}));
248        assert!(d.unregister("t"));
249        assert!(!d.unregister("t")); // already gone
250    }
251
252    #[test]
253    fn unregister_then_non_strict() {
254        let mut d = ToolDefaults::new();
255        d.register("t", json!({"a": 1}));
256        d.unregister("t");
257        let out = d.apply("t", &json!({"x": 7}), false).unwrap();
258        assert_eq!(out["x"], json!(7));
259        assert!(!out.as_object().unwrap().contains_key("a"));
260    }
261
262    #[test]
263    fn tool_names_sorted() {
264        let mut d = ToolDefaults::new();
265        d.register("zz", json!({}));
266        d.register("aa", json!({}));
267        d.register("mm", json!({}));
268        assert_eq!(d.tool_names(), vec!["aa", "mm", "zz"]);
269    }
270
271    #[test]
272    fn defaults_for_some() {
273        let mut d = ToolDefaults::new();
274        d.register("t", json!({"x": 1}));
275        let m = d.defaults_for("t").unwrap();
276        assert_eq!(m["x"], json!(1));
277    }
278
279    #[test]
280    fn defaults_for_none() {
281        let d = ToolDefaults::new();
282        assert!(d.defaults_for("missing").is_none());
283    }
284
285    #[test]
286    fn contains_check() {
287        let mut d = ToolDefaults::new();
288        assert!(!d.contains("t"));
289        d.register("t", json!({}));
290        assert!(d.contains("t"));
291        d.unregister("t");
292        assert!(!d.contains("t"));
293    }
294
295    #[test]
296    fn register_many() {
297        let mut d = ToolDefaults::new();
298        d.register_many(json!({
299            "a": {"x": 1},
300            "b": {"y": 2}
301        }));
302        assert_eq!(d.apply("a", &json!({}), false).unwrap()["x"], json!(1));
303        assert_eq!(d.apply("b", &json!({}), false).unwrap()["y"], json!(2));
304    }
305
306    #[test]
307    fn apply_does_not_mutate_args() {
308        let mut d = ToolDefaults::new();
309        d.register("t", json!({"a": 1}));
310        let args = json!({"b": 2});
311        let _ = d.apply("t", &args, false).unwrap();
312        assert!(!args.as_object().unwrap().contains_key("a"));
313    }
314
315    #[test]
316    fn nested_value_preserved() {
317        let mut d = ToolDefaults::new();
318        d.register("t", json!({"opts": {"retries": 3}}));
319        let out = d.apply("t", &json!({}), false).unwrap();
320        assert_eq!(out["opts"]["retries"], json!(3));
321    }
322
323    #[test]
324    fn bool_int_string_value_types() {
325        let mut d = ToolDefaults::new();
326        d.register(
327            "t",
328            json!({"flag": true, "count": 42, "label": "hi"}),
329        );
330        let out = d.apply("t", &json!({}), false).unwrap();
331        assert_eq!(out["flag"], json!(true));
332        assert_eq!(out["count"], json!(42));
333        assert_eq!(out["label"], json!("hi"));
334    }
335
336    #[test]
337    fn multiple_tools_independent() {
338        let mut d = ToolDefaults::new();
339        d.register("a", json!({"x": 1}));
340        d.register("b", json!({"y": 2}));
341        let a = d.apply("a", &json!({}), false).unwrap();
342        let b = d.apply("b", &json!({}), false).unwrap();
343        assert!(!a.as_object().unwrap().contains_key("y"));
344        assert!(!b.as_object().unwrap().contains_key("x"));
345    }
346}
tool_arg_defaults/lib.rs

tool_arg_defaults/
lib.rs
