Skip to main content

ai_agents_tools/builtin/
json.rs

1use async_trait::async_trait;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::generate_schema;
7use ai_agents_core::{Tool, ToolResult};
8
9pub struct JsonTool;
10
11impl JsonTool {
12    pub fn new() -> Self {
13        Self
14    }
15}
16
17impl Default for JsonTool {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23#[derive(Debug, Deserialize, JsonSchema)]
24struct JsonInput {
25    /// Operation to perform: parse, get, set, merge, stringify, keys, values
26    operation: String,
27    /// JSON string or data to operate on
28    #[serde(default)]
29    data: Option<Value>,
30    /// Path using dot notation (e.g., 'user.name', 'items.0.id')
31    #[serde(default)]
32    path: Option<String>,
33    /// Value to set (for set operation)
34    #[serde(default)]
35    value: Option<Value>,
36    /// Second JSON data (for merge operation)
37    #[serde(default)]
38    data2: Option<Value>,
39    /// Pretty print output (default: true)
40    #[serde(default)]
41    pretty: Option<bool>,
42}
43
44#[derive(Debug, Serialize, Deserialize)]
45struct ParseOutput {
46    parsed: Value,
47    valid: bool,
48}
49
50#[derive(Debug, Serialize, Deserialize)]
51struct GetOutput {
52    value: Value,
53    path: String,
54    found: bool,
55}
56
57#[derive(Debug, Serialize, Deserialize)]
58struct SetOutput {
59    result: Value,
60    path: String,
61}
62
63#[derive(Debug, Serialize, Deserialize)]
64struct MergeOutput {
65    result: Value,
66}
67
68#[derive(Debug, Serialize, Deserialize)]
69struct StringifyOutput {
70    result: String,
71}
72
73#[derive(Debug, Serialize, Deserialize)]
74struct KeysOutput {
75    keys: Vec<String>,
76    count: usize,
77}
78
79#[derive(Debug, Serialize, Deserialize)]
80struct ValuesOutput {
81    values: Vec<Value>,
82    count: usize,
83}
84
85#[async_trait]
86impl Tool for JsonTool {
87    fn id(&self) -> &str {
88        "json"
89    }
90
91    fn name(&self) -> &str {
92        "JSON Manipulation"
93    }
94
95    fn description(&self) -> &str {
96        "Parse, query, and manipulate JSON data. Operations: parse (validate JSON string), get (extract value by path), set (set value at path), merge (combine two JSON objects), stringify (convert to string), keys (get object keys), values (get object values)."
97    }
98
99    fn input_schema(&self) -> Value {
100        generate_schema::<JsonInput>()
101    }
102
103    async fn execute(&self, args: Value) -> ToolResult {
104        let input: JsonInput = match serde_json::from_value(args) {
105            Ok(input) => input,
106            Err(e) => return ToolResult::error(format!("Invalid input: {}", e)),
107        };
108
109        match input.operation.to_lowercase().as_str() {
110            "parse" => self.handle_parse(&input),
111            "get" => self.handle_get(&input),
112            "set" => self.handle_set(&input),
113            "merge" => self.handle_merge(&input),
114            "stringify" => self.handle_stringify(&input),
115            "keys" => self.handle_keys(&input),
116            "values" => self.handle_values(&input),
117            _ => ToolResult::error(format!(
118                "Unknown operation: {}. Valid operations: parse, get, set, merge, stringify, keys, values",
119                input.operation
120            )),
121        }
122    }
123}
124
125impl JsonTool {
126    fn handle_parse(&self, input: &JsonInput) -> ToolResult {
127        match &input.data {
128            Some(Value::String(s)) => match serde_json::from_str::<Value>(s) {
129                Ok(parsed) => {
130                    let output = ParseOutput {
131                        parsed,
132                        valid: true,
133                    };
134                    self.to_result(&output)
135                }
136                Err(e) => ToolResult::error(format!("Parse error: {}", e)),
137            },
138            Some(v) => {
139                let output = ParseOutput {
140                    parsed: v.clone(),
141                    valid: true,
142                };
143                self.to_result(&output)
144            }
145            None => ToolResult::error("'data' is required for parse operation"),
146        }
147    }
148
149    fn handle_get(&self, input: &JsonInput) -> ToolResult {
150        let data = match &input.data {
151            Some(d) => d,
152            None => return ToolResult::error("'data' is required for get operation"),
153        };
154
155        let path = match &input.path {
156            Some(p) => p,
157            None => return ToolResult::error("'path' is required for get operation"),
158        };
159
160        let value = self.get_by_path(data, path);
161        let found = !value.is_null();
162
163        let output = GetOutput {
164            value,
165            path: path.clone(),
166            found,
167        };
168
169        self.to_result(&output)
170    }
171
172    fn handle_set(&self, input: &JsonInput) -> ToolResult {
173        let data = match &input.data {
174            Some(d) => d.clone(),
175            None => return ToolResult::error("'data' is required for set operation"),
176        };
177
178        let path = match &input.path {
179            Some(p) => p,
180            None => return ToolResult::error("'path' is required for set operation"),
181        };
182
183        let value = match &input.value {
184            Some(v) => v.clone(),
185            None => return ToolResult::error("'value' is required for set operation"),
186        };
187
188        let result = match self.set_by_path(data, path, value) {
189            Ok(r) => r,
190            Err(e) => return ToolResult::error(e),
191        };
192
193        let output = SetOutput {
194            result,
195            path: path.clone(),
196        };
197
198        self.to_result(&output)
199    }
200
201    fn handle_merge(&self, input: &JsonInput) -> ToolResult {
202        let data1 = match &input.data {
203            Some(d) => d.clone(),
204            None => return ToolResult::error("'data' is required for merge operation"),
205        };
206
207        let data2 = match &input.data2 {
208            Some(d) => d.clone(),
209            None => return ToolResult::error("'data2' is required for merge operation"),
210        };
211
212        let result = self.merge_values(data1, data2);
213
214        let output = MergeOutput { result };
215
216        self.to_result(&output)
217    }
218
219    fn handle_stringify(&self, input: &JsonInput) -> ToolResult {
220        let data = match &input.data {
221            Some(d) => d,
222            None => return ToolResult::error("'data' is required for stringify operation"),
223        };
224
225        let pretty = input.pretty.unwrap_or(true);
226
227        let result = if pretty {
228            serde_json::to_string_pretty(data).unwrap_or_default()
229        } else {
230            serde_json::to_string(data).unwrap_or_default()
231        };
232
233        let output = StringifyOutput { result };
234
235        self.to_result(&output)
236    }
237
238    fn handle_keys(&self, input: &JsonInput) -> ToolResult {
239        let data = match &input.data {
240            Some(d) => d,
241            None => return ToolResult::error("'data' is required for keys operation"),
242        };
243
244        let keys: Vec<String> = match data {
245            Value::Object(obj) => obj.keys().cloned().collect(),
246            _ => return ToolResult::error("'data' must be a JSON object for keys operation"),
247        };
248
249        let count = keys.len();
250        let output = KeysOutput { keys, count };
251
252        self.to_result(&output)
253    }
254
255    fn handle_values(&self, input: &JsonInput) -> ToolResult {
256        let data = match &input.data {
257            Some(d) => d,
258            None => return ToolResult::error("'data' is required for values operation"),
259        };
260
261        let values: Vec<Value> = match data {
262            Value::Object(obj) => obj.values().cloned().collect(),
263            Value::Array(arr) => arr.clone(),
264            _ => {
265                return ToolResult::error(
266                    "'data' must be a JSON object or array for values operation",
267                );
268            }
269        };
270
271        let count = values.len();
272        let output = ValuesOutput { values, count };
273
274        self.to_result(&output)
275    }
276
277    fn get_by_path(&self, data: &Value, path: &str) -> Value {
278        let parts: Vec<&str> = path.split('.').collect();
279        let mut current = data;
280
281        for part in parts {
282            if part.is_empty() {
283                continue;
284            }
285
286            if let Ok(index) = part.parse::<usize>() {
287                match current.get(index) {
288                    Some(v) => current = v,
289                    None => return Value::Null,
290                }
291            } else {
292                match current.get(part) {
293                    Some(v) => current = v,
294                    None => return Value::Null,
295                }
296            }
297        }
298
299        current.clone()
300    }
301
302    fn set_by_path(&self, mut data: Value, path: &str, value: Value) -> Result<Value, String> {
303        let parts: Vec<&str> = path.split('.').filter(|s| !s.is_empty()).collect();
304
305        if parts.is_empty() {
306            return Ok(value);
307        }
308
309        let mut current = &mut data;
310
311        for (i, part) in parts.iter().enumerate() {
312            let is_last = i == parts.len() - 1;
313
314            if let Ok(index) = part.parse::<usize>() {
315                if !current.is_array() {
316                    *current = Value::Array(vec![]);
317                }
318
319                let arr = current.as_array_mut().unwrap();
320                while arr.len() <= index {
321                    arr.push(Value::Null);
322                }
323
324                if is_last {
325                    arr[index] = value.clone();
326                    break;
327                } else {
328                    current = &mut arr[index];
329                }
330            } else {
331                if !current.is_object() {
332                    *current = Value::Object(serde_json::Map::new());
333                }
334
335                let obj = current.as_object_mut().unwrap();
336
337                if is_last {
338                    obj.insert(part.to_string(), value.clone());
339                    break;
340                } else {
341                    if !obj.contains_key(*part) {
342                        obj.insert(part.to_string(), Value::Object(serde_json::Map::new()));
343                    }
344                    current = obj.get_mut(*part).unwrap();
345                }
346            }
347        }
348
349        Ok(data)
350    }
351
352    fn merge_values(&self, base: Value, overlay: Value) -> Value {
353        match (base, overlay) {
354            (Value::Object(mut base_obj), Value::Object(overlay_obj)) => {
355                for (key, value) in overlay_obj {
356                    let merged = if let Some(base_value) = base_obj.remove(&key) {
357                        self.merge_values(base_value, value)
358                    } else {
359                        value
360                    };
361                    base_obj.insert(key, merged);
362                }
363                Value::Object(base_obj)
364            }
365            (_, overlay) => overlay,
366        }
367    }
368
369    fn to_result<T: Serialize>(&self, output: &T) -> ToolResult {
370        match serde_json::to_string(output) {
371            Ok(json) => ToolResult::ok(json),
372            Err(e) => ToolResult::error(format!("Serialization error: {}", e)),
373        }
374    }
375}
376
377#[cfg(test)]
378mod tests {
379    use super::*;
380
381    #[tokio::test]
382    async fn test_parse_string() {
383        let tool = JsonTool::new();
384        let result = tool
385            .execute(serde_json::json!({
386                "operation": "parse",
387                "data": "{\"name\": \"test\", \"value\": 42}"
388            }))
389            .await;
390        assert!(result.success);
391
392        let output: ParseOutput = serde_json::from_str(&result.output).unwrap();
393        assert!(output.valid);
394        assert_eq!(output.parsed["name"], "test");
395        assert_eq!(output.parsed["value"], 42);
396    }
397
398    #[tokio::test]
399    async fn test_parse_invalid() {
400        let tool = JsonTool::new();
401        let result = tool
402            .execute(serde_json::json!({
403                "operation": "parse",
404                "data": "{invalid json}"
405            }))
406            .await;
407        assert!(!result.success);
408    }
409
410    #[tokio::test]
411    async fn test_get_simple() {
412        let tool = JsonTool::new();
413        let result = tool
414            .execute(serde_json::json!({
415                "operation": "get",
416                "data": {"user": {"name": "Alice", "age": 30}},
417                "path": "user.name"
418            }))
419            .await;
420        assert!(result.success);
421
422        let output: GetOutput = serde_json::from_str(&result.output).unwrap();
423        assert!(output.found);
424        assert_eq!(output.value, "Alice");
425    }
426
427    #[tokio::test]
428    async fn test_get_array_index() {
429        let tool = JsonTool::new();
430        let result = tool
431            .execute(serde_json::json!({
432                "operation": "get",
433                "data": {"items": ["a", "b", "c"]},
434                "path": "items.1"
435            }))
436            .await;
437        assert!(result.success);
438
439        let output: GetOutput = serde_json::from_str(&result.output).unwrap();
440        assert!(output.found);
441        assert_eq!(output.value, "b");
442    }
443
444    #[tokio::test]
445    async fn test_get_not_found() {
446        let tool = JsonTool::new();
447        let result = tool
448            .execute(serde_json::json!({
449                "operation": "get",
450                "data": {"a": 1},
451                "path": "b.c.d"
452            }))
453            .await;
454        assert!(result.success);
455
456        let output: GetOutput = serde_json::from_str(&result.output).unwrap();
457        assert!(!output.found);
458    }
459
460    #[tokio::test]
461    async fn test_set_simple() {
462        let tool = JsonTool::new();
463        let result = tool
464            .execute(serde_json::json!({
465                "operation": "set",
466                "data": {"user": {"name": "Alice"}},
467                "path": "user.age",
468                "value": 30
469            }))
470            .await;
471        assert!(result.success);
472
473        let output: SetOutput = serde_json::from_str(&result.output).unwrap();
474        assert_eq!(output.result["user"]["age"], 30);
475        assert_eq!(output.result["user"]["name"], "Alice");
476    }
477
478    #[tokio::test]
479    async fn test_set_nested() {
480        let tool = JsonTool::new();
481        let result = tool
482            .execute(serde_json::json!({
483                "operation": "set",
484                "data": {},
485                "path": "a.b.c",
486                "value": "deep"
487            }))
488            .await;
489        assert!(result.success);
490
491        let output: SetOutput = serde_json::from_str(&result.output).unwrap();
492        assert_eq!(output.result["a"]["b"]["c"], "deep");
493    }
494
495    #[tokio::test]
496    async fn test_merge() {
497        let tool = JsonTool::new();
498        let result = tool
499            .execute(serde_json::json!({
500                "operation": "merge",
501                "data": {"a": 1, "b": {"x": 10}},
502                "data2": {"b": {"y": 20}, "c": 3}
503            }))
504            .await;
505        assert!(result.success);
506
507        let output: MergeOutput = serde_json::from_str(&result.output).unwrap();
508        assert_eq!(output.result["a"], 1);
509        assert_eq!(output.result["b"]["x"], 10);
510        assert_eq!(output.result["b"]["y"], 20);
511        assert_eq!(output.result["c"], 3);
512    }
513
514    #[tokio::test]
515    async fn test_stringify() {
516        let tool = JsonTool::new();
517        let result = tool
518            .execute(serde_json::json!({
519                "operation": "stringify",
520                "data": {"name": "test"},
521                "pretty": false
522            }))
523            .await;
524        assert!(result.success);
525
526        let output: StringifyOutput = serde_json::from_str(&result.output).unwrap();
527        assert_eq!(output.result, "{\"name\":\"test\"}");
528    }
529
530    #[tokio::test]
531    async fn test_keys() {
532        let tool = JsonTool::new();
533        let result = tool
534            .execute(serde_json::json!({
535                "operation": "keys",
536                "data": {"a": 1, "b": 2, "c": 3}
537            }))
538            .await;
539        assert!(result.success);
540
541        let output: KeysOutput = serde_json::from_str(&result.output).unwrap();
542        assert_eq!(output.count, 3);
543        assert!(output.keys.contains(&"a".to_string()));
544        assert!(output.keys.contains(&"b".to_string()));
545        assert!(output.keys.contains(&"c".to_string()));
546    }
547
548    #[tokio::test]
549    async fn test_values() {
550        let tool = JsonTool::new();
551        let result = tool
552            .execute(serde_json::json!({
553                "operation": "values",
554                "data": [1, 2, 3]
555            }))
556            .await;
557        assert!(result.success);
558
559        let output: ValuesOutput = serde_json::from_str(&result.output).unwrap();
560        assert_eq!(output.count, 3);
561    }
562
563    #[tokio::test]
564    async fn test_invalid_operation() {
565        let tool = JsonTool::new();
566        let result = tool
567            .execute(serde_json::json!({"operation": "invalid"}))
568            .await;
569        assert!(!result.success);
570    }
571}