AstroResearch/src/agent/tools/astro/paper.rs

// src/agent/tools/paper.rs — 文献下载、解析、内容读取工具

use async_trait::async_trait;
use serde_json::json;
use tracing::info;

use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput};

// ── GetPaperContentTool ──

/// 获取文献内容工具：仅从本地读取已解析的 Markdown 全文
pub struct GetPaperContentTool;

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

    fn description(&self) -> &str {
        "读取并在本地库中获取已解析的文献 Markdown 完整文本内容。\
         注意：本工具仅能读取已在数据库注册且已解析的文献，不会自动触发下载或解析。\
         若文献未下载或未解析，本工具会返回详细指引。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "bibcode": {
                    "type": "string",
                    "description": "文献的唯一标识符，支持 ADS Bibcode、DOI 或 arXiv ID"
                }
            },
            "required": ["bibcode"]
        })
    }

    /// 纯读取操作，并发安全
    fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
            Some(b) => b.to_string(),
            None => return ToolOutput::error("缺少必需参数 'bibcode'"),
        };

        info!("[GetPaperContent] 获取文献内容: {}", bibcode);
        let state = &ctx.app_state;

        let paths = crate::api::helpers::check_paper_paths_in_db(
            &state.db,
            &state.config.library_dir,
            &bibcode,
        )
        .await;
        let md_opt = match paths {
            Ok(Some((_, _, md_opt, _))) => md_opt,
            Ok(None) => return ToolOutput::error(
                "获取文献内容失败：该文献未在本地数据库中注册，请先使用 search_papers 搜索该文献。",
            ),
            Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)),
        };

        let md_rel = match md_opt {
            Some(rel) => rel,
            None => {
                return ToolOutput::error(
                    "获取文献内容失败：该文献尚未完成结构化解析。如果未下载，请先调用 download_paper；如果已下载，请先调用 parse_paper 进行解析。",
                )
            }
        };

        let md_abs = state.config.library_dir.join(&md_rel);
        if !md_abs.exists() {
            return ToolOutput::error(
                "获取文献内容失败：文献本地 Markdown 文件已丢失，请重新调用 parse_paper 进行解析。",
            );
        }

        match std::fs::read_to_string(&md_abs) {
            Ok(content) => ToolOutput::success(
                content.clone(),
                json!({ "bibcode": bibcode, "chars": content.len() }),
            ),
            Err(e) => ToolOutput::error(format!("获取文献内容失败，读取本地文件错误: {}", e)),
        }
    }

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

// ── DownloadPaperTool ──

/// 下载文献全文资源工具
pub struct DownloadPaperTool;

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

    fn description(&self) -> &str {
        "下载指定文献的全文资源（PDF 或 HTML）至本地图书馆，为后续的结构化解析做好准备。\
         适用于：需要阅读或分析新搜寻到的、尚未下载的文献。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "bibcode": {
                    "type": "string",
                    "description": "文献的唯一标识符，支持 ADS Bibcode、DOI 或 arXiv ID"
                },
                "force": {
                    "type": "boolean",
                    "description": "是否强制重新下载（即使本地已下载该文献）"
                }
            },
            "required": ["bibcode"]
        })
    }

    /// 下载有副作用，中断时应阻塞以完成
    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Block
    }

    /// 下载失败时应中止兄弟并行执行（避免继续处理同一文献）
    fn causes_sibling_abort(&self) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
            Some(b) => b.to_string(),
            None => return ToolOutput::error("缺少必需参数 'bibcode'"),
        };
        let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);

        info!("[DownloadPaper] 下载文献: {}, 强制重下: {}", bibcode, force);
        let state = &ctx.app_state;

        match state
            .downloader
            .download_paper_service(&state.db, &state.config.library_dir, &bibcode, force)
            .await
        {
            Ok(paper) => ToolOutput::success(
                format!(
                    "文献 {} 全文资源下载成功。格式 - PDF: {}, HTML: {}",
                    bibcode, paper.has_pdf, paper.has_html
                ),
                json!({ "bibcode": bibcode, "has_pdf": paper.has_pdf, "has_html": paper.has_html }),
            ),
            Err(e) => ToolOutput::error(format!("文献 {} 下载失败: {}", bibcode, e)),
        }
    }

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

// ── ParsePaperTool ──

/// 结构化解析文献内容工具
pub struct ParsePaperTool;

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

    fn description(&self) -> &str {
        "将指定文献本地已下载的 HTML 或 PDF 资源解析为结构化的 Markdown 文本，并保存至本地 Markdown 文件夹。\
         注意：调用此工具前必须确保文献已被成功下载（已执行 download_paper）。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "bibcode": {
                    "type": "string",
                    "description": "文献的唯一标识符，支持 ADS Bibcode、DOI 或 arXiv ID"
                },
                "force": {
                    "type": "boolean",
                    "description": "是否强制重新解析（即使本地已解析过该文献）"
                }
            },
            "required": ["bibcode"]
        })
    }

    /// 解析有副作用（写文件），中断时应阻塞以完成
    fn interrupt_behavior(&self) -> InterruptBehavior {
        InterruptBehavior::Block
    }

    /// 解析失败时应中止兄弟并行执行
    fn causes_sibling_abort(&self) -> bool {
        true
    }

    async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput {
        let bibcode = match args.get("bibcode").and_then(|b| b.as_str()) {
            Some(b) => b.to_string(),
            None => return ToolOutput::error("缺少必需参数 'bibcode'"),
        };
        let force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false);

        info!("[ParsePaper] 解析文献: {}, 强制重析: {}", bibcode, force);
        let state = &ctx.app_state;

        match crate::services::parser::parse_paper_service(
            &state.db,
            &state.config.library_dir,
            &state.qiniu,
            &state.config,
            &bibcode,
            force,
        )
        .await
        {
            Ok(markdown) => ToolOutput::success(
                format!(
                    "文献 {} 结构化解析成功。解析后 Markdown 字符总数: {}",
                    bibcode,
                    markdown.len()
                ),
                json!({ "bibcode": bibcode, "chars": markdown.len() }),
            ),
            Err(e) => {
                let msg = e.to_string();
                if msg.contains("请先下载") {
                    ToolOutput::error(format!(
                        "文献 {} 解析失败：未检测到已下载的本地资源文件，请先调用 download_paper 工具进行下载。",
                        bibcode
                    ))
                } else {
                    ToolOutput::error(format!("文献 {} 解析失败: {}", bibcode, msg))
                }
            }
        }
    }

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