AstroResearch/src/agent/tools/astro/research/vizier.rs

// src/agent/tools/astro/research/vizier.rs
//
// QueryVizierTool —— VizieR TAP 星表查询（自由 ADQL + 便捷表查询）
// ConeSearchTool —— 锥形检索（按坐标查近邻天体）
//
// 对齐 QueryTargetTool 范式：单元结构体 + AgentTool 实现 + ctx.app_state 调 service

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::<Vec<_>>()
            .join(" | ")
    ));

    // 表体（限制预览行数）
    let show = result.rows.len().min(preview_rows);
    for row in result.rows.iter().take(show) {
        let cells: Vec<String> = 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
}

// ── QueryVizierTool ──

pub struct QueryVizierTool;

#[async_trait]
impl AgentTool for QueryVizierTool {
    fn name(&self) -> &str {
        "query_vizier"
    }

    fn description(&self) -> &str {
        "通过 VizieR TAP 服务查询天文星表数据，支持两种模式：\
         (1) 自由 ADQL 查询——传入 'adql' 参数执行标准 ADQL 语句；\
         (2) 便捷表查询——传入 'table_name' + 可选 'columns' + 'limit' 直接取行。\
         适用于：获取天体的精确测光/天体测量参数、查询星表中的近邻天体、交叉证认、批量拉取某类样本。\
         结果自动缓存 7 天。不确定表名时先用 search_catalogs 搜索。\
         ADQL 语法示例：SELECT TOP 10 ra, dec FROM \"I/355/gaiadr3\" WHERE parallax > 10"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "adql": {
                    "type": "string",
                    "description": "自由 ADQL 查询语句（与 table_name 二选一）。如 SELECT TOP 10 * FROM \"I/355/gaiadr3\""
                },
                "table_name": {
                    "type": "string",
                    "description": "VizieR 表名（便捷模式，与 adql 二选一），如 'I/355/gaiadr3'（Gaia DR3）、'J/AJ/165/8/table2'"
                },
                "columns": {
                    "type": "string",
                    "description": "需要返回的列名（逗号分隔），为空时返回所有列"
                },
                "limit": {
                    "type": "integer",
                    "description": "最大返回行数（默认 20，上限 2000）",
                    "default": 20
                }
            }
        })
    }

    fn group(&self) -> &str {
        "as:research"
    }

    fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
        true
    }

    fn is_readonly(&self) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let state = &ctx.app_state;

        // 解析参数：adql 优先，否则走 table_name 便捷模式
        let result = if let Some(adql) = args.get("adql").and_then(|v| v.as_str()) {
            let limit = args
                .get("limit")
                .and_then(|v| v.as_i64())
                .unwrap_or(20)
                .clamp(1, 2000);
            info!(
                "[QueryVizier] ADQL 查询 (limit={}): {}",
                limit,
                adql.chars().take(150).collect::<String>()
            );
            crate::services::cds::vizier::query_adql_cached(&state.db, &state.vizier, adql, limit)
                .await
        } else if let Some(table) = args.get("table_name").and_then(|v| v.as_str()) {
            let limit = args
                .get("limit")
                .and_then(|v| v.as_i64())
                .unwrap_or(20)
                .clamp(1, 2000);
            let columns: Vec<String> = args
                .get("columns")
                .and_then(|v| v.as_str())
                .map(|c| c.split(',').map(|s| s.trim().to_string()).collect())
                .unwrap_or_default();
            info!("[QueryVizier] 表查询 table={} limit={}", table, limit);
            crate::services::cds::vizier::query_table(
                &state.db,
                &state.vizier,
                table,
                &columns,
                limit,
            )
            .await
        } else {
            return ToolOutput::error(
                "需要提供 'adql'（自由 ADQL）或 'table_name'（便捷表查询）参数之一",
            );
        };

        match result {
            Ok(r) => {
                let content = render_result_table(&r, 20);
                ToolOutput::success(content, json!(r))
            }
            Err(e) => ToolOutput::error(format!("VizieR 查询失败: {}", e)),
        }
    }
}

// ── ConeSearchTool ──

pub struct ConeSearchTool;

#[async_trait]
impl AgentTool for ConeSearchTool {
    fn name(&self) -> &str {
        "cone_search"
    }

    fn description(&self) -> &str {
        "锥形检索（Cone Search）：按坐标在天文星表中检索近邻天体。\
         需要指定目标星表（table 参数）。\
         适用于：给定坐标找附近天体、获取某区域的星表数据、配合 query_target 解析名称后做区域查询。\
         坐标系统为 J2000（ICRS），单位为度。结果自动缓存 7 天。\
         不确定表名时先用 search_catalogs 搜索。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "ra": {
                    "type": "number",
                    "description": "赤经 RA（度，J2000/ICRS），范围 0~360"
                },
                "dec": {
                    "type": "number",
                    "description": "赤纬 Dec（度，J2000/ICRS），范围 -90~90"
                },
                "radius_deg": {
                    "type": "number",
                    "description": "检索半径（度），默认 0.1，范围 0~5",
                    "default": 0.1
                },
                "table": {
                    "type": "string",
                    "description": "目标星表（必填），如 'I/355/gaiadr3'（Gaia DR3）、'II/246/out'（2MASS）"
                },
                "max_records": {
                    "type": "integer",
                    "description": "最大返回行数（默认 50，上限 2000）",
                    "default": 50
                }
            },
            "required": ["ra", "dec", "table"]
        })
    }

    fn group(&self) -> &str {
        "as:research"
    }

    fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
        true
    }

    fn is_readonly(&self) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let state = &ctx.app_state;

        let ra = match args.get("ra").and_then(|v| v.as_f64()) {
            Some(v) => v,
            None => return ToolOutput::error("缺少必需参数 'ra'（赤经，度）"),
        };
        let dec = match args.get("dec").and_then(|v| v.as_f64()) {
            Some(v) => v,
            None => return ToolOutput::error("缺少必需参数 'dec'（赤纬，度）"),
        };
        let radius = args
            .get("radius_deg")
            .and_then(|v| v.as_f64())
            .unwrap_or(0.1);
        let table = match args.get("table").and_then(|v| v.as_str()) {
            Some(t) => t,
            None => return ToolOutput::error("缺少必需参数 'table'（目标星表）"),
        };
        let max_records = args
            .get("max_records")
            .and_then(|v| v.as_i64())
            .unwrap_or(50)
            .clamp(1, 2000);

        info!(
            "[ConeSearch] ra={} dec={} radius={}° table={}",
            ra, dec, radius, table
        );

        match crate::services::cds::vizier::cone_search(
            &state.db,
            &state.vizier,
            ra,
            dec,
            radius,
            table,
            max_records,
        )
        .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)),
        }
    }
}