1. 1//! In-Memory Polyglot Data Engine.
2. 2//!
3. 3//! Provides deterministic in-memory querying, filtering, projection,
4. 4//! grouping, sorting, and statistical aggregation over tabular data
5. 5//! (CSV, JSON arrays, or Markdown tables) with zero subprocess overhead.
6. 6
7. 7use std::collections::BTreeMap;
8. 8
9. 9use serde_json::Value;
10. 10
11. 11pub struct DataQueryTool;
12. 12
13. 13#[derive(Debug, Clone)]
14. 14struct TableData {
15. 15 headers: Vec<String>,
16. 16 rows: Vec<Vec<String>>,
17. 17}
18. 18
19. 19impl TableData {
20. 20 fn from_csv(text: &str) -> Result<Self, String> {
21. 21 let lines: Vec<&str> = text
22. 22 .lines()
23. 23 .map(str::trim)
24. 24 .filter(|l| !l.is_empty())
25. 25 .collect();
26. 26 if lines.is_empty() {
27. 27 return Err("empty CSV data".into());
28. 28 }
29. 29 let headers = parse_csv_line(lines[0]);
30. 30 let mut rows = Vec::new();
31. 31 for (idx, line) in lines.iter().skip(1).enumerate() {
32. 32 let cells = parse_csv_line(line);
33. 33 if cells.len() != headers.len() {
34. 34 return Err(format!(
35. 35 "CSV row {} has {} columns, expected {}",
36. 36 idx + 1,
37. 37 cells.len(),
38. 38 headers.len()
39. 39 ));
40. 40 }
41. 41 rows.push(cells);
42. 42 }
43. 43 Ok(TableData { headers, rows })
44. 44 }
45. 45
46. 46 fn from_json(val: &Value) -> Result<Self, String> {
47. 47 let arr = val
48. 48 .as_array()
49. 49 .ok_or_else(|| "JSON data must be an array of objects".to_string())?;
50. 50 if arr.is_empty() {
51. 51 return Ok(TableData {
52. 52 headers: Vec::new(),
53. 53 rows: Vec::new(),
54. 54 });
55. 55 }
56. 56 let mut header_set = Vec::new();
57. 57 for item in arr {
58. 58 if let Some(obj) = item.as_object() {
59. 59 for k in obj.keys() {
60. 60 if !header_set.contains(k) {
61. 61 header_set.push(k.clone());
62. 62 }
63. 63 }
64. 64 }
65. 65 }
66. 66 let mut rows = Vec::new();
67. 67 for item in arr {
68. 68 let mut row = Vec::new();
69. 69 if let Some(obj) = item.as_object() {
70. 70 for h in &header_set {
71. 71 let cell = match obj.get(h) {
72. 72 Some(Value::String(s)) => s.clone(),
73. 73 Some(Value::Null) | None => String::new(),
74. 74 Some(other) => other.to_string(),
75. 75 };
76. 76 row.push(cell);
77. 77 }
78. 78 }
79. 79 rows.push(row);
80. 80 }
81. 81 Ok(TableData {
82. 82 headers: header_set,
83. 83 rows,
84. 84 })
85. 85 }
86. 86
87. 87 fn from_markdown(text: &str) -> Result<Self, String> {
88. 88 let mut table_lines = Vec::new();
89. 89 for line in text.lines() {
90. 90 let trimmed = line.trim();
91. 91 if trimmed.starts_with('|') && trimmed.ends_with('|') && trimmed.len() > 1 {
92. 92 table_lines.push(trimmed);
93. 93 } else if !table_lines.is_empty() {
94. 94 break;
95. 95 }
96. 96 }
97. 97 if table_lines.len() < 2 {
98. 98 return Err("invalid markdown table (needs header and separator)".into());
99. 99 }
100. 100
101. 101 fn parse_md_cells(l: &str) -> Vec<String> {
102. 102 l.trim_matches('|')
103. 103 .split('|')
104. 104 .map(|c| c.trim().to_string())
105. 105 .collect()
106. 106 }
107. 107
108. 108 let headers = parse_md_cells(table_lines[0]);
109. 109 let mut rows = Vec::new();
110. 110 for (i, line) in table_lines.iter().skip(2).enumerate() {
111. 111 let cells = parse_md_cells(line);
112. 112 if cells.len() != headers.len() {
113. 113 return Err(format!(
114. 114 "markdown table row {} has {} columns, expected {}",
115. 115 i + 1,
116. 116 cells.len(),
117. 117 headers.len()
118. 118 ));
119. 119 }
120. 120 rows.push(cells);
121. 121 }
122. 122 Ok(TableData { headers, rows })
123. 123 }
124. 124
125. 125 fn col_index(&self, name: &str) -> Option<usize> {
126. 126 self.headers
127. 127 .iter()
128. 128 .position(|h| h.eq_ignore_ascii_case(name))
129. 129 }
130. 130
131. 131 fn to_markdown(&self) -> String {
132. 132 if self.headers.is_empty() {
133. 133 return String::new();
134. 134 }
135. 135 let mut out = format!("| {} |\n", self.headers.join(" | "));
136. 136 let sep: Vec<String> = self.headers.iter().map(|_| "---".to_string()).collect();
137. 137 out.push_str(&format!("| {} |\n", sep.join(" | ")));
138. 138 for row in &self.rows {
139. 139 out.push_str(&format!("| {} |\n", row.join(" | ")));
140. 140 }
141. 141 out
142. 142 }
143. 143
144. 144 fn to_json(&self) -> Value {
145. 145 let mut arr = Vec::new();
146. 146 for row in &self.rows {
147. 147 let mut obj = serde_json::Map::new();
148. 148 for (idx, header) in self.headers.iter().enumerate() {
149. 149 let cell = row.get(idx).cloned().unwrap_or_default();
150. 150 if let Ok(num) = cell.parse::<f64>() {
151. 151 if num.fract() == 0.0 && num >= i64::MIN as f64 && num <= i64::MAX as f64 {
152. 152 obj.insert(header.clone(), Value::from(num as i64));
153. 153 } else {
154. 154 obj.insert(header.clone(), Value::from(num));
155. 155 }
156. 156 } else if cell.eq_ignore_ascii_case("true") {
157. 157 obj.insert(header.clone(), Value::Bool(true));
158. 158 } else if cell.eq_ignore_ascii_case("false") {
159. 159 obj.insert(header.clone(), Value::Bool(false));
160. 160 } else {
161. 161 obj.insert(header.clone(), Value::String(cell));
162. 162 }
163. 163 }
164. 164 arr.push(Value::Object(obj));
165. 165 }
166. 166 Value::Array(arr)
167. 167 }
168. 168}
169. 169
170. 170fn parse_csv_line(line: &str) -> Vec<String> {
171. 171 let mut fields = Vec::new();
172. 172 let mut current = String::new();
173. 173 let mut in_quotes = false;
174. 174 let mut chars = line.chars().peekable();
175. 175
176. 176 while let Some(c) = chars.next() {
177. 177 match c {
178. 178 '"' if in_quotes => {
179. 179 if chars.peek() == Some(&'"') {
180. 180 current.push('"');
181. 181 chars.next();
182. 182 } else {
183. 183 in_quotes = false;
184. 184 }
185. 185 }
186. 186 '"' => {
187. 187 in_quotes = true;
188. 188 }
189. 189 ',' if !in_quotes => {
190. 190 fields.push(current.trim().to_string());
191. 191 current.clear();
192. 192 }
193. 193 other => current.push(other),
194. 194 }
195. 195 }
196. 196 fields.push(current.trim().to_string());
197. 197 fields
198. 198}
199. 199
200. 200#[async_trait::async_trait]
201. 201impl vak_tools::Tool for DataQueryTool {
202. 202 fn name(&self) -> &str {
203. 203 "data_query"
204. 204 }
205. 205
206. 206 fn serves(&self) -> &'static [&'static str] {
207. 207 &["documents"]
208. 208 }
209. 209
210. 210 fn description(&self) -> &str {
211. 211 "Query and analyze structured data (CSV, JSON arrays, or Markdown tables) in memory. \
212. 212 Supports column projection (select), row filtering (filter), grouping & aggregation \
213. 213 (count, sum, avg, min, max), sorting, and row limits."
214. 214 }
215. 215
216. 216 fn schema(&self) -> Value {
217. 217 serde_json::json!({
218. 218 "type": "object",
219. 219 "properties": {
220. 220 "data": {
221. 221 "type": "string",
222. 222 "description": "Raw tabular string data: CSV text, JSON array of objects, or Markdown table"
223. 223 },
224. 224 "format": {
225. 225 "type": "string",
226. 226 "enum": ["csv", "json", "markdown"],
227. 227 "description": "Optional format hint ('csv', 'json', 'markdown'). Auto-detected if omitted."
228. 228 },
229. 229 "select": {
230. 230 "type": "array",
231. 231 "items": { "type": "string" },
232. 232 "description": "Optional list of columns to include in the output"
233. 233 },
234. 234 "filter": {
235. 235 "type": "object",
236. 236 "properties": {
237. 237 "column": { "type": "string" },
238. 238 "op": { "type": "string", "enum": ["eq", "neq", "gt", "lt", "contains"] },
239. 239 "value": { "type": "string" }
240. 240 },
241. 241 "required": ["column", "op", "value"],
242. 242 "description": "Optional row filter predicate"
243. 243 },
244. 244 "group_by": {
245. 245 "type": "string",
246. 246 "description": "Optional column to group by before aggregation"
247. 247 },
248. 248 "aggregate": {
249. 249 "type": "object",
250. 250 "properties": {
251. 251 "column": { "type": "string" },
252. 252 "fn": { "type": "string", "enum": ["count", "sum", "avg", "min", "max"] }
253. 253 },
254. 254 "required": ["column", "fn"],
255. 255 "description": "Optional aggregation function"
256. 256 },
257. 257 "sort_by": {
258. 258 "type": "object",
259. 259 "properties": {
260. 260 "column": { "type": "string" },
261. 261 "descending": { "type": "boolean" }
262. 262 },
263. 263 "required": ["column"],
264. 264 "description": "Optional column to sort rows by"
265. 265 },
266. 266 "limit": {
267. 267 "type": "integer",
268. 268 "description": "Maximum number of rows to return"
269. 269 }
270. 270 },
271. 271 "required": ["data"]
272. 272 })
273. 273 }
274. 274
275. 275 #[allow(clippy::collapsible_if)]
276. 276 async fn execute(&self, args: &Value, _ctx: &vak_tools::ToolContext) -> vak_tools::ToolOutput {
277. 277 let Some(raw_data) = args.get("data").and_then(Value::as_str).map(str::trim) else {
278. 278 return vak_tools::ToolOutput::error("missing required argument 'data'");
279. 279 };
280. 280 if raw_data.is_empty() {
281. 281 return vak_tools::ToolOutput::error("'data' must not be empty");
282. 282 }
283. 283
284. 284 let format = args.get("format").and_then(Value::as_str);
285. 285 let table = if format == Some("json") || (format.is_none() && raw_data.starts_with('[')) {
286. 286 match serde_json::from_str::<Value>(raw_data) {
287. 287 Ok(v) => TableData::from_json(&v),
288. 288 Err(e) => return vak_tools::ToolOutput::error(format!("invalid JSON data: {e}")),
289. 289 }
290. 290 } else if format == Some("markdown") || (format.is_none() && raw_data.starts_with('|')) {
291. 291 TableData::from_markdown(raw_data)
292. 292 } else {
293. 293 TableData::from_csv(raw_data)
294. 294 };
295. 295
296. 296 let mut table = match table {
297. 297 Ok(t) => t,
298. 298 Err(e) => return vak_tools::ToolOutput::error(format!("could not parse data: {e}")),
299. 299 };
300. 300
301. 301 // 1. Filtering
302. 302 if let Some(filter) = args.get("filter") {
303. 303 let col = filter.get("column").and_then(Value::as_str).unwrap_or("");
304. 304 let op = filter.get("op").and_then(Value::as_str).unwrap_or("eq");
305. 305 let target_val = filter.get("value").and_then(Value::as_str).unwrap_or("");
306. 306
307. 307 if let Some(idx) = table.col_index(col) {
308. 308 let target_num = target_val.parse::<f64>().ok();
309. 309 table.rows.retain(|row| {
310. 310 let cell = row.get(idx).map(|s| s.as_str()).unwrap_or("");
311. 311 let cell_num = cell.parse::<f64>().ok();
312. 312 match op {
313. 313 "eq" => cell.eq_ignore_ascii_case(target_val),
314. 314 "neq" => !cell.eq_ignore_ascii_case(target_val),
315. 315 "contains" => cell
316. 316 .to_ascii_lowercase()
317. 317 .contains(&target_val.to_ascii_lowercase()),
318. 318 "gt" => match (cell_num, target_num) {
319. 319 (Some(c), Some(t)) => c > t,
320. 320 _ => cell > target_val,
321. 321 },
322. 322 "lt" => match (cell_num, target_num) {
323. 323 (Some(c), Some(t)) => c < t,
324. 324 _ => cell < target_val,
325. 325 },
326. 326 _ => true,
327. 327 }
328. 328 });
329. 329 } else {
330. 330 return vak_tools::ToolOutput::error(format!("filter column '{col}' not found"));
331. 331 }
332. 332 }
333. 333
334. 334 // 2. Group By & Aggregation
335. 335 if let Some(agg) = args.get("aggregate") {
336. 336 let agg_col = agg.get("column").and_then(Value::as_str).unwrap_or("");
337. 337 let agg_fn = agg.get("fn").and_then(Value::as_str).unwrap_or("count");
338. 338 let agg_idx = match table.col_index(agg_col) {
339. 339 Some(i) => i,
340. 340 None => {
341. 341 return vak_tools::ToolOutput::error(format!(
342. 342 "aggregate column '{agg_col}' not found"
343. 343 ));
344. 344 }
345. 345 };
346. 346
347. 347 if let Some(group_by_col) = args.get("group_by").and_then(Value::as_str) {
348. 348 let group_idx = match table.col_index(group_by_col) {
349. 349 Some(i) => i,
350. 350 None => {
351. 351 return vak_tools::ToolOutput::error(format!(
352. 352 "group_by column '{group_by_col}' not found"
353. 353 ));
354. 354 }
355. 355 };
356. 356
357. 357 let mut groups: BTreeMap<String, Vec<f64>> = BTreeMap::new();
358. 358 for row in &table.rows {
359. 359 let key = row.get(group_idx).cloned().unwrap_or_default();
360. 360 let val = row
361. 361 .get(agg_idx)
362. 362 .and_then(|s| s.parse::<f64>().ok())
363. 363 .unwrap_or(0.0);
364. 364 groups.entry(key).or_default().push(val);
365. 365 }
366. 366
367. 367 let mut agg_rows = Vec::new();
368. 368 for (k, vals) in groups {
369. 369 let result = match agg_fn {
370. 370 "count" => vals.len() as f64,
371. 371 "sum" => vals.iter().sum(),
372. 372 "avg" => {
373. 373 if vals.is_empty() {
374. 374 0.0
375. 375 } else {
376. 376 vals.iter().sum::<f64>() / vals.len() as f64
377. 377 }
378. 378 }
379. 379 "min" => vals.iter().cloned().fold(f64::INFINITY, f64::min),
380. 380 "max" => vals.iter().cloned().fold(f64::NEG_INFINITY, f64::max),
381. 381 _ => vals.len() as f64,
382. 382 };
383. 383 agg_rows.push(vec![
384. 384 k,
385. 385 format!("{result:.2}").trim_end_matches(".00").to_string(),
386. 386 ]);
387. 387 }
388. 388
389. 389 table.headers = vec![group_by_col.to_string(), format!("{agg_fn}_{agg_col}")];
390. 390 table.rows = agg_rows;
391. 391 } else {
392. 392 let vals: Vec<f64> = table
393. 393 .rows
394. 394 .iter()
395. 395 .map(|r| {
396. 396 r.get(agg_idx)
397. 397 .and_then(|s| s.parse::<f64>().ok())
398. 398 .unwrap_or(0.0)
399. 399 })
400. 400 .collect();
401. 401 let result = match agg_fn {
402. 402 "count" => vals.len() as f64,
403. 403 "sum" => vals.iter().sum(),
404. 404 "avg" => {
405. 405 if vals.is_empty() {
406. 406 0.0
407. 407 } else {
408. 408 vals.iter().sum::<f64>() / vals.len() as f64
409. 409 }
410. 410 }
411. 411 "min" => vals.iter().cloned().fold(f64::INFINITY, f64::min),
412. 412 "max" => vals.iter().cloned().fold(f64::NEG_INFINITY, f64::max),
413. 413 _ => vals.len() as f64,
414. 414 };
415. 415 table.headers = vec![format!("{agg_fn}_{agg_col}")];
416. 416 table.rows = vec![vec![
417. 417 format!("{result:.2}").trim_end_matches(".00").to_string(),
418. 418 ]];
419. 419 }
420. 420 }
421. 421
422. 422 // 3. Sorting
423. 423 if let Some(sort) = args.get("sort_by") {
424. 424 if let Some(sort_col) = sort.get("column").and_then(Value::as_str) {
425. 425 let desc = sort
426. 426 .get("descending")
427. 427 .and_then(Value::as_bool)
428. 428 .unwrap_or(false);
429. 429 if let Some(idx) = table.col_index(sort_col) {
430. 430 table.rows.sort_by(|a, b| {
431. 431 let val_a = a.get(idx).map(|s| s.as_str()).unwrap_or("");
432. 432 let val_b = b.get(idx).map(|s| s.as_str()).unwrap_or("");
433. 433 let num_a = val_a.parse::<f64>().ok();
434. 434 let num_b = val_b.parse::<f64>().ok();
435. 435 let ord = match (num_a, num_b) {
436. 436 (Some(na), Some(nb)) => {
437. 437 na.partial_cmp(&nb).unwrap_or(std::cmp::Ordering::Equal)
438. 438 }
439. 439 _ => val_a.cmp(val_b),
440. 440 };
441. 441 if desc { ord.reverse() } else { ord }
442. 442 });
443. 443 }
444. 444 }
445. 445 }
446. 446
447. 447 // 4. Projection (Select)
448. 448 if let Some(select_cols) = args.get("select").and_then(Value::as_array) {
449. 449 let requested: Vec<&str> = select_cols.iter().filter_map(Value::as_str).collect();
450. 450 if !requested.is_empty() {
451. 451 let indices: Vec<Option<usize>> =
452. 452 requested.iter().map(|c| table.col_index(c)).collect();
453. 453 let mut new_rows = Vec::new();
454. 454 for row in &table.rows {
455. 455 let new_row = indices
456. 456 .iter()
457. 457 .map(|opt| opt.and_then(|i| row.get(i).cloned()).unwrap_or_default())
458. 458 .collect();
459. 459 new_rows.push(new_row);
460. 460 }
461. 461 table.headers = requested.into_iter().map(String::from).collect();
462. 462 table.rows = new_rows;
463. 463 }
464. 464 }
465. 465
466. 466 // 5. Limit
467. 467 if let Some(limit) = args.get("limit").and_then(Value::as_u64) {
468. 468 table.rows.truncate(limit as usize);
469. 469 }
470. 470
471. 471 let md = table.to_markdown();
472. 472 let json_val = table.to_json();
473. 473 let response_text = format!(
474. 474 "Query Result ({} rows, {} columns):\n\n{}\n```json\n{}\n```",
475. 475 table.rows.len(),
476. 476 table.headers.len(),
477. 477 md,
478. 478 serde_json::to_string_pretty(&json_val).unwrap_or_default()
479. 479 );
480. 480
481. 481 vak_tools::ToolOutput::ok(response_text)
482. 482 }
483. 483
484. 484 fn claims(&self, _args: &Value) -> vak_tools::ResourceClaims {
485. 485 vak_tools::ResourceClaims {
486. 486 exclusive: false,
487. 487 read_only: true,
488. 488 paths: vec![],
489. 489 }
490. 490 }
491. 491}
492. 492
493. 493#[cfg(test)]
494. 494mod tests {
495. 495 #![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
496. 496 use super::*;
497. 497 use vak_tools::Tool;
498. 498
499. 499 #[tokio::test]
500. 500 async fn data_query_csv_filter_sort_aggregate() {
501. 501 let tool = DataQueryTool;
502. 502 let ctx = vak_tools::ToolContext {
503. 503 cwd: std::path::PathBuf::from("/tmp"),
504. 504 cancel: tokio_util::sync::CancellationToken::new(),
505. 505 sandbox: None,
506. 506 sandbox_sink: None,
507. 507 agent_id: None,
508. 508 new_documents: Vec::new(),
509. 509 };
510. 510
511. 511 let csv = "name,dept,salary\nAlice,Eng,120000\nBob,Sales,85000\nCharlie,Eng,140000\nDave,Sales,90000\nEve,Eng,110000\n";
512. 512
513. 513 // 1. Filter and sort
514. 514 let args = serde_json::json!({
515. 515 "data": csv,
516. 516 "filter": { "column": "dept", "op": "eq", "value": "Eng" },
517. 517 "sort_by": { "column": "salary", "descending": true },
518. 518 "select": ["name", "salary"]
519. 519 });
520. 520 let out = tool.execute(&args, &ctx).await;
521. 521 assert!(!out.is_error, "{}", out.content);
522. 522 assert!(out.content.contains("Charlie"), "{}", out.content);
523. 523 assert!(!out.content.contains("Bob"), "{}", out.content);
524. 524
525. 525 // 2. Group by and aggregation
526. 526 let agg_args = serde_json::json!({
527. 527 "data": csv,
528. 528 "group_by": "dept",
529. 529 "aggregate": { "column": "salary", "fn": "avg" }
530. 530 });
531. 531 let agg_out = tool.execute(&agg_args, &ctx).await;
532. 532 assert!(!agg_out.is_error, "{}", agg_out.content);
533. 533 assert!(agg_out.content.contains("Eng"), "{}", agg_out.content);
534. 534 assert!(
535. 535 agg_out.content.contains("123333.33") || agg_out.content.contains("123333"),
536. 536 "{}",
537. 537 agg_out.content
538. 538 );
539. 539 }
540. 540}
541. 541