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: String,
27 #[serde(default)]
29 data: Option<Value>,
30 #[serde(default)]
32 path: Option<String>,
33 #[serde(default)]
35 value: Option<Value>,
36 #[serde(default)]
38 data2: Option<Value>,
39 #[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}