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

// src/agent/tools/astro/research/paper_reader_tools.rs
//
// 文献内容读取与大纲提取的智能体工具定义。

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

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

// ==========================================
// 1. GetPaperOutlineTool (获取文献大纲目录)
// ==========================================

#[derive(Default)]
pub struct GetPaperOutlineTool;

#[async_trait]
impl AgentTool for GetPaperOutlineTool {
    fn name(&self) -> &'static str {
        "get_paper_outline"
    }

    fn display_name(&self) -> &str {
        "文献提纲提取"
    }

    fn description(&self) -> &'static str {
        "获取文献的章节大纲目录，列出所有章节的序号和标题。\n\
         在读取文献内容前，强烈建议先使用此工具了解章节结构，以便按需读取，节约 token。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "bibcode": {
                    "type": "string",
                    "description": "文献的唯一标识符（bibcode，如 2006BaltA..15...69W）"
                }
            },
            "required": ["bibcode"]
        })
    }

    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!("[GetPaperOutline] 提取大纲: {}", bibcode);
        let state = &ctx.app_state;

        let mode = crate::services::paper::ReadMode::Outline { max_level: Some(3) };

        match crate::services::paper::read_paper_content(
            &state.db,
            &state.config.storage.library_dir,
            &bibcode,
            mode,
        )
        .await
        {
            Ok(res) => {
                let sections = res.outline.unwrap_or_default();
                if sections.is_empty() {
                    return ToolOutput::success(
                        "该文献未检测到章节标题（## 或 ### 格式）。请直接读取全文。",
                        json!({ "bibcode": bibcode, "sections": [] }),
                    );
                }

                let display_sections: Vec<serde_json::Value> = sections
                    .iter()
                    .map(|s| {
                        json!({
                            "index": s.index,
                            "level": s.level,
                            "heading": s.heading.clone(),
                        })
                    })
                    .collect();

                let header = format!("文献 {} 章节大纲（共 {} 节）:\n", bibcode, sections.len());
                let footer = "\n\n使用 get_paper_content 并传入 section_index 或 section_name 读取指定章节内容。";

                ToolOutput::success(
                    header + &res.content + footer,
                    json!({
                        "bibcode": bibcode,
                        "section_count": sections.len(),
                        "sections": display_sections
                    }),
                )
            }
            Err(e) => ToolOutput::error(e.to_string()),
        }
    }
}

// ==========================================
// 2. GetPaperContentTool (获取文献具体内容)
// ==========================================

#[derive(Default)]
pub struct GetPaperContentTool;

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

    fn display_name(&self) -> &str {
        "文献内容读取"
    }

    fn description(&self) -> &'static str {
        "获取文献的具体内容（全文或指定章节内容）。\n\
         可传入 section_index 或 section_name 指定仅读取某一章节。\n\
         先使用 get_paper_outline 获取章节结构后按需读取，可大幅减少 token 消耗。"
    }

    fn parameters(&self) -> serde_json::Value {
        json!({
            "type": "object",
            "properties": {
                "bibcode": {
                    "type": "string",
                    "description": "文献的唯一标识符（bibcode，如 2006BaltA..15...69W）"
                },
                "section_index": {
                    "type": "integer",
                    "description": "章节序号（0-based）。从 get_paper_outline 返回的 sections[index] 获取。与 section_name 二选一"
                },
                "section_name": {
                    "type": "string",
                    "description": "章节标题文本（支持模糊匹配，如 'introduction'）。与 section_index 二选一"
                }
            },
            "required": ["bibcode"]
        })
    }

    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 section_index = args
            .get("section_index")
            .and_then(|s| s.as_u64())
            .map(|n| n as usize);
        let section_name = args
            .get("section_name")
            .and_then(|s| s.as_str())
            .map(String::from);

        info!(
            "[GetPaperContent] bibcode={}, section_index={:?}, section_name={:?}",
            bibcode, section_index, section_name
        );

        let state = &ctx.app_state;

        let mode = if let Some(idx) = section_index {
            crate::services::paper::ReadMode::SectionIndex(idx)
        } else if let Some(name) = section_name {
            crate::services::paper::ReadMode::SectionName(name)
        } else {
            crate::services::paper::ReadMode::Full {
                include_translation: false,
            }
        };

        let label = match &mode {
            crate::services::paper::ReadMode::SectionIndex(idx) => Some(format!("#{}", idx)),
            crate::services::paper::ReadMode::SectionName(name) => Some(name.clone()),
            _ => None,
        };
        let is_full = matches!(mode, crate::services::paper::ReadMode::Full { .. });

        match crate::services::paper::read_paper_content(
            &state.db,
            &state.config.storage.library_dir,
            &bibcode,
            mode,
        )
        .await
        {
            Ok(res) => ToolOutput::success(
                res.content.clone(),
                json!({
                    "bibcode": bibcode,
                    "chars": res.content.len(),
                    "section": label,
                    "full_text": is_full
                }),
            ),
            Err(e) => ToolOutput::error(e.to_string()),
        }
    }
}