Coverage for agentos/tools/data_tools.py: 19%

1"""数据处理工具 — JSON/CSV 解析、格式化、查询。""" 

2 

3from __future__ import annotations 

4 

5import csv 

6import json 

7import os 

8from io import StringIO 

9from typing import Any 

10 

11from agentos.tools.base import BaseTool, ToolResult 

12 

13 

14class JsonTool(BaseTool): 

15 """JSON 处理工具 — 解析、格式化、JSONPath 查询、验证。""" 

16 

17 name = "json_tool" 

18 description = "JSON 解析、格式化、JSONPath 查询、Schema 验证。输入 JSON 字符串或 .json 文件路径" 

19 

20 @property 

21 def parameters(self) -> dict: 

22 return { 

23 "type": "object", 

24 "properties": { 

25 "action": { 

26 "type": "string", 

27 "description": "操作类型：parse/format/query/validate", 

28 "enum": ["parse", "format", "query", "validate"], 

29 }, 

30 "input": {"type": "string", "description": "JSON 字符串或 .json 文件路径"}, 

31 "jsonpath": { 

32 "type": "string", 

33 "description": "JSONPath 查询表达式（仅 query），如 $.store.book[0].title", 

34 }, 

35 "indent": {"type": "integer", "description": "缩进空格数，默认 2"}, 

36 }, 

37 "required": ["action", "input"], 

38 } 

39 

40 async def execute(self, arguments: dict, sandbox=None) -> ToolResult: 

41 action = arguments.get("action", "parse") 

42 input_data = arguments.get("input", "") 

43 jsonpath = arguments.get("jsonpath", "$") 

44 indent = arguments.get("indent", 2) 

45 

46 # Read file if path 

47 if os.path.isfile(input_data): 

48 try: 

49 with open(input_data, encoding="utf-8") as f: 

50 data_str = f.read() 

51 except Exception as e: 

52 return ToolResult.fail(call_id="", error=f"File read error: {e}") 

53 else: 

54 data_str = input_data 

55 

56 # Parse 

57 try: 

58 data = json.loads(data_str) 

59 except json.JSONDecodeError as e: 

60 return ToolResult.fail(call_id="", error=f"JSON parse error: {e}") 

61 

62 if action == "parse": 

63 info = f"Type: {type(data).__name__}\n" 

64 if isinstance(data, dict): 

65 info += f"Keys: {list(data.keys())[:20]}\n" 

66 if isinstance(data, (list, dict)): 

67 info += f"Length: {len(data)}\n" 

68 info += f"Sample: {json.dumps(data, indent=indent, ensure_ascii=False)[:1000]}" 

69 return ToolResult.ok(call_id="", output=info) 

70 

71 elif action == "format": 

72 formatted = json.dumps(data, indent=indent, ensure_ascii=False) 

73 return ToolResult.ok(call_id="", output=formatted) 

74 

75 elif action == "query": 

76 result = self._jsonpath_query(data, jsonpath) 

77 output = ( 

78 json.dumps(result, indent=indent, ensure_ascii=False) 

79 if result is not None 

80 else "null" 

81 ) 

82 return ToolResult.ok(call_id="", output=output) 

83 

84 elif action == "validate": 

85 return ToolResult.ok( 

86 call_id="", 

87 output=f"Valid JSON. Type: {type(data).__name__}. " 

88 f"Size: {len(data_str)} chars. " 

89 f"{'Keys: ' + str(list(data.keys())[:20]) if isinstance(data, dict) else ''}", 

90 ) 

91 

92 return ToolResult.fail(call_id="", error=f"Unknown action: {action}") 

93 

94 def _jsonpath_query(self, data: Any, path: str) -> Any: 

95 if path == "$": 

96 return data 

97 parts = path.replace("[", ".").replace("]", "").split(".") 

98 current = data 

99 for part in parts: 

100 if not part or part == "$": 

101 continue 

102 if isinstance(current, dict): 

103 current = current.get(part) 

104 elif isinstance(current, list): 

105 try: 

106 current = current[int(part)] 

107 except (ValueError, IndexError): 

108 return None 

109 else: 

110 return None 

111 return current 

112 

113 

114class CsvTool(BaseTool): 

115 """CSV 处理工具 — 读取、查询、统计。""" 

116 

117 name = "csv_tool" 

118 description = "CSV 文件读取、列提取、基本统计。输入 .csv 文件路径" 

119 

120 @property 

121 def parameters(self) -> dict: 

122 return { 

123 "type": "object", 

124 "properties": { 

125 "action": { 

126 "type": "string", 

127 "description": "操作类型：read/stats/query", 

128 "enum": ["read", "stats", "query"], 

129 }, 

130 "input": {"type": "string", "description": ".csv 文件路径"}, 

131 "columns": {"type": "string", "description": "要提取的列名，逗号分隔（仅 query）"}, 

132 "limit": {"type": "integer", "description": "最大行数，默认 50"}, 

133 }, 

134 "required": ["action", "input"], 

135 } 

136 

137 async def execute(self, arguments: dict, sandbox=None) -> ToolResult: 

138 action = arguments.get("action", "read") 

139 input_data = arguments.get("input", "") 

140 columns = arguments.get("columns", "") 

141 limit = arguments.get("limit", 50) 

142 

143 # Try as file path 

144 if os.path.isfile(input_data): 

145 try: 

146 with open(input_data, encoding="utf-8", errors="ignore") as f: 

147 data_str = f.read() 

148 except Exception as e: 

149 return ToolResult.fail(call_id="", error=f"File read error: {e}") 

150 else: 

151 data_str = input_data 

152 

153 try: 

154 reader = csv.DictReader(StringIO(data_str)) 

155 col_names = reader.fieldnames or [] 

156 rows = [row for i, row in enumerate(reader) if i < limit] 

157 except Exception as e: 

158 return ToolResult.fail(call_id="", error=f"CSV parse error: {e}") 

159 

160 if action == "read": 

161 output = f"Columns: {col_names}\nRows: {len(rows)}\n\n" 

162 for row in rows[:20]: 

163 output += str(row) + "\n" 

164 return ToolResult.ok(call_id="", output=output) 

165 

166 elif action == "stats": 

167 output = f"Columns: {col_names}\nTotal rows loaded: {len(rows)}\n\n" 

168 for col in col_names: 

169 values = [row[col] for row in rows if row.get(col)] 

170 unique = len(set(values)) 

171 output += f" {col}: {unique} unique values, sample={values[:3]}\n" 

172 return ToolResult.ok(call_id="", output=output) 

173 

174 elif action == "query": 

175 target_cols = [c.strip() for c in columns.split(",")] if columns else col_names 

176 output = f"Columns: {target_cols}\n\n" 

177 for row in rows: 

178 output += ", ".join(f"{c}={row.get(c, '')}" for c in target_cols if c in row) + "\n" 

179 return ToolResult.ok(call_id="", output=output) 

180 

181 return ToolResult.fail(call_id="", error=f"Unknown action: {action}")