// src/agent/tools/astro/research/vizier.rs // // CatalogOperationTool —— VizieR 星表统一操作工具 // search: 按关键词搜索星表目录 // describe: 查看表的列结构 // query: 执行 ADQL 或按表名查询,返回数据 // cone: 按坐标锥形检索 // export: 导出为 CSV 文件 // lookup: 通过 bibcode 查找关联数据表 use async_trait::async_trait; use serde_json::json; use tracing::info; use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; use crate::clients::cds::vizier::VizierQueryResult; /// 渲染查询结果为 Markdown 表格(前 N 行) fn render_result_table(result: &VizierQueryResult, preview_rows: usize) -> String { let mut content = String::new(); if let Some(ref table) = result.table_name { content.push_str(&format!("来源表: {}\n", table)); } content.push_str(&format!( "共 {} 行({})\n\n", result.row_count, if result.truncated { "已截断,存在更多结果" } else { "完整结果" } )); if result.fields.is_empty() || result.rows.is_empty() { content.push_str("(无数据行)\n"); return content; } let headers: Vec<&str> = result.fields.iter().map(|f| f.name.as_str()).collect(); content.push_str(&format!("| {} |\n", headers.join(" | "))); content.push_str(&format!( "| {} |\n", headers .iter() .map(|_| "---") .collect::>() .join(" | ") )); let show = result.rows.len().min(preview_rows); for row in result.rows.iter().take(show) { let cells: Vec = row .iter() .map(|v| match v { serde_json::Value::Null => "—".to_string(), serde_json::Value::String(s) => s.clone(), other => other.to_string(), }) .collect(); content.push_str(&format!("| {} |\n", cells.join(" | "))); } if result.rows.len() > preview_rows { content.push_str(&format!( "\n(已省略 {} 行,完整数据见结构化输出)\n", result.rows.len() - preview_rows )); } let units: Vec<&str> = result .fields .iter() .filter_map(|f| f.unit.as_deref()) .collect(); if !units.is_empty() { content.push_str(&format!("\n字段单位: {}\n", units.join(", "))); } content } // ── CatalogOperationTool ── pub struct CatalogOperationTool; #[async_trait] impl AgentTool for CatalogOperationTool { fn name(&self) -> &str { "catalog_operation" } fn display_name(&self) -> &str { "星表检索" } fn description(&self) -> &str { "VizieR 星表统一操作。通过 action 选择:\n\ search — 按关键词搜索星表目录;\n\ describe — 查看表的列结构;\n\ query — 执行 ADQL 或按表名查询数据;\n\ cone — 按坐标锥形检索近邻天体;\n\ export — 下载表数据保存为 CSV;\n\ lookup — 通过 bibcode 查找关联数据表。" } fn parameters(&self) -> serde_json::Value { json!({ "type": "object", "properties": { "action": { "type": "string", "enum": ["search", "describe", "query", "cone", "export", "lookup"], "description": "操作类型" }, "keyword": { "type": "string", "description": "search 时的搜索关键词,如 'Gaia DR3'" }, "table": { "type": "string", "description": "VizieR 表名,describe/query/cone/export 时使用,如 'I/355/gaiadr3'" }, "adql": { "type": "string", "description": "query/export 时的 ADQL 语句(与 table 二选一)" }, "columns": { "type": "string", "description": "query/export + table 模式下指定列(逗号分隔,默认 *)" }, "coords": { "type": "object", "description": "cone 时的坐标参数", "properties": { "ra": { "type": "number", "description": "赤经 RA(度)" }, "dec": { "type": "number", "description": "赤纬 Dec(度)" }, "radius_deg": { "type": "number", "description": "检索半径(度),默认 0.1", "default": 0.1 }, "strategy": { "type": "string", "enum": ["nearest", "all"], "description": "nearest(默认)/ all", "default": "nearest" } }, "required": ["ra", "dec"] }, "limit": { "type": "integer", "description": "最大返回行数(默认 50,上限 2000)", "default": 50 }, "bibcode": { "type": "string", "description": "lookup 时的 ADS bibcode" } }, "required": ["action"] }) } fn group(&self) -> &str { "as:research" } fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { true } fn is_readonly(&self) -> bool { false } async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { let action = match args.get("action").and_then(|v| v.as_str()) { Some(a) => a, None => return ToolOutput::error("缺少必需参数 'action'"), }; let state = &ctx.app_state; match action { "search" => do_catalog_search(state, &args).await, "describe" => do_catalog_describe(state, &args).await, "query" => do_catalog_query(state, &args).await, "cone" => do_catalog_cone(state, &args).await, "export" => do_catalog_export(state, &args).await, "lookup" => do_catalog_lookup(state, &args).await, other => ToolOutput::error(format!( "未知 action '{}',可选: search/describe/query/cone/export/lookup", other )), } } } async fn do_catalog_search(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput { let keyword = match args.get("keyword").and_then(|v| v.as_str()) { Some(k) => k, None => return ToolOutput::error("search 需要 'keyword' 参数"), }; let limit = args .get("limit") .and_then(|v| v.as_i64()) .unwrap_or(10) .clamp(1, 30) as usize; let catalog = crate::services::cds::vizier::VizierCatalog::new(&state.db, &state.sources.vizier); let results = match catalog.search(keyword, limit).await { Ok(r) => r, Err(e) => return ToolOutput::error(format!("搜索失败: {}", e)), }; if results.is_empty() { return ToolOutput::success(format!("未找到与 '{}' 匹配的星表", keyword), json!([])); } let mut content = format!( "搜索 '{}' 匹配到 {} 个星表(按数据量降序):\n\n", keyword, results.len() ); let items: Vec = results .iter() .map(|entry| { let nrows_str = entry.nrows.map(|n| format!("{} 行", n)).unwrap_or_else(|| "行数未知".into()); content.push_str(&format!("- `{}` — {} ({})\n", entry.table_name, entry.description, nrows_str)); json!({"table": entry.table_name, "description": entry.description, "nrows": entry.nrows}) }) .collect(); ToolOutput::success(content, json!(items)) } async fn do_catalog_describe(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput { let table = match args.get("table").and_then(|v| v.as_str()) { Some(t) => t, None => return ToolOutput::error("describe 需要 'table' 参数"), }; let catalog = crate::services::cds::vizier::VizierCatalog::new(&state.db, &state.sources.vizier); let columns = match catalog.describe(table).await { Ok(c) => c, Err(e) => return ToolOutput::error(format!("查询表结构失败: {}", e)), }; if columns.is_empty() { return ToolOutput::success(format!("表 '{}' 未找到列定义", table), json!([])); } let mut content = format!("表 `{}` 共 {} 列:\n\n", table, columns.len()); let items: Vec = columns .iter() .map(|col| { let unit_str = col.unit.as_deref().unwrap_or(""); let desc_str = col.description.as_deref().unwrap_or(""); content.push_str(&format!("- `{}` ({}) {} — {}\n", col.column_name, col.datatype, unit_str, desc_str)); json!({"name": col.column_name, "datatype": col.datatype, "unit": col.unit, "description": col.description}) }) .collect(); ToolOutput::success(content, json!(items)) } async fn do_catalog_query(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput { let limit = args .get("limit") .and_then(|v| v.as_i64()) .unwrap_or(50) .clamp(1, 2000); let result = if let Some(adql) = args.get("adql").and_then(|v| v.as_str()) { info!( "[CatalogOp:query] ADQL (limit={}): {}", limit, adql.chars().take(150).collect::() ); crate::services::cds::vizier::query_adql_cached( &state.db, &state.sources.vizier, adql, limit, ) .await } else if let Some(table) = args.get("table").and_then(|v| v.as_str()) { let columns: Vec = args .get("columns") .and_then(|v| v.as_str()) .map(|c| c.split(',').map(|s| s.trim().to_string()).collect()) .unwrap_or_default(); info!("[CatalogOp:query] table={} limit={}", table, limit); crate::services::cds::vizier::query_table( &state.db, &state.sources.vizier, table, &columns, limit, ) .await } else { return ToolOutput::error("query 需要 'adql' 或 'table' 参数之一"); }; match result { Ok(r) => { let content = render_result_table(&r, 20); ToolOutput::success(content, json!(r)) } Err(e) => ToolOutput::error(format!("查询失败: {}", e)), } } async fn do_catalog_cone(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput { let coords = match args.get("coords").and_then(|v| v.as_object()) { Some(c) => c, None => return ToolOutput::error("cone 需要 'coords' 参数(含 ra/dec)"), }; let ra = match coords.get("ra").and_then(|v| v.as_f64()) { Some(v) => v, None => return ToolOutput::error("coords 缺少 'ra'"), }; let dec = match coords.get("dec").and_then(|v| v.as_f64()) { Some(v) => v, None => return ToolOutput::error("coords 缺少 'dec'"), }; let radius = coords .get("radius_deg") .and_then(|v| v.as_f64()) .unwrap_or(0.1); let nearest = coords.get("strategy").and_then(|v| v.as_str()) != Some("all"); let table = match args.get("table").and_then(|v| v.as_str()) { Some(t) => t, None => return ToolOutput::error("cone 需要 'table' 参数"), }; let max_records = args .get("limit") .and_then(|v| v.as_i64()) .unwrap_or(50) .clamp(1, 2000); info!( "[CatalogOp:cone] ra={} dec={} radius={}° table={} strategy={}", ra, dec, radius, table, if nearest { "nearest" } else { "all" } ); match crate::services::cds::vizier::cone_search( &state.db, &state.sources.vizier, ra, dec, radius, table, max_records, nearest, ) .await { Ok(r) => { let mut content = format!( "Cone Search 结果 (ra={}, dec={}, radius={}°)\n\n", ra, dec, radius ); content.push_str(&render_result_table(&r, 20)); ToolOutput::success(content, json!(r)) } Err(e) => ToolOutput::error(format!("Cone Search 失败: {}", e)), } } async fn do_catalog_export(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput { let limit = args .get("limit") .and_then(|v| v.as_i64()) .unwrap_or(100) .clamp(1, 5000); let adql = args.get("adql").and_then(|v| v.as_str()); let table = args.get("table").and_then(|v| v.as_str()); let columns = args.get("columns").and_then(|v| v.as_str()); if adql.is_none() && table.is_none() { return ToolOutput::error("export 需要 'adql' 或 'table' 参数"); } let catalog = crate::services::cds::vizier::VizierCatalog::new(&state.db, &state.sources.vizier); let result = match catalog .export_to_file( state.config.storage.library_dir.to_str().unwrap_or("."), adql, table, columns, limit, ) .await { Ok(r) => r, Err(e) => return ToolOutput::error(format!("导出失败: {}", e)), }; ToolOutput::success( format!( "已导出 {} 行数据到 `{}`\n文件大小: {}", result.row_count, result.path.display(), result.size_bytes, ), json!({ "path": result.path.to_str(), "rows": result.row_count, "columns": result.column_count, }), ) } async fn do_catalog_lookup(state: &crate::api::AppState, args: &serde_json::Value) -> ToolOutput { let bibcode = match args.get("bibcode").and_then(|v| v.as_str()) { Some(b) => b, None => return ToolOutput::error("lookup 需要 'bibcode' 参数"), }; let limit = args .get("limit") .and_then(|v| v.as_i64()) .unwrap_or(10) .clamp(1, 30) as usize; let catalog = crate::services::cds::vizier::VizierCatalog::with_ads( &state.db, &state.sources.vizier, &state.sources.ads, ); let results = match catalog.lookup(bibcode, limit).await { Ok(r) => r, Err(e) => return ToolOutput::error(format!("查找失败: {}", e)), }; if results.is_empty() { return ToolOutput::success( format!("文献 '{}' 未在 CDS/VizieR 中找到关联数据表", bibcode), json!([]), ); } let mut content = format!( "文献 '{}' 关联 {} 个 VizieR 数据表:\n\n", bibcode, results.len() ); let items: Vec = results .iter() .map(|entry| { let nrows_str = entry.nrows.map(|n| format!("{} 行", n)).unwrap_or_default(); content.push_str(&format!("- `{}` — {} ({})\n", entry.table_name, entry.description, nrows_str)); json!({"table": entry.table_name, "description": entry.description, "nrows": entry.nrows}) }) .collect(); ToolOutput::success(content, json!(items)) }