diff --git a/.gitignore b/.gitignore index 4acfcbb..f8fe032 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ target/ .DS_Store *.suo *.swp +.omc library/MinerU/ libs/ diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..aff6e9c --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,163 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Build, Lint & Test Commands + +```bash +# Build (debug) +cargo build + +# Build (release with optimizations) +cargo build --release + +# Build (release-min profile: size-optimized LTO) +cargo build --profile release-min + +# Run (debug, starts server on http://localhost:8000) +cargo run + +# Run with Obscura in-process browser (no external binaries needed) +cargo run --features obscura-inprocess + +# Run CLI binary +cargo run --bin astroresearch_cli + +# Run health check tool +cargo run --bin health_check # read-only scan +cargo run --bin health_check -- --fix # auto-repair + +# Lint +cargo clippy + +# Format +cargo fmt + +# All tests +cargo test + +# Unit tests only +cargo test --lib + +# Run a specific test +cargo test test_name + +# Frontend (cd dashboard first) +npm run dev # HMR dev server on :5173, proxies /api to :8000 +npm run build # TypeScript check + Vite production build → dashboard/dist/ +npm run lint # ESLint +``` + +## Architecture Overview + +**Stack**: Rust Axum backend (port 8000) + React/Vite/TypeScript frontend (port 5173 in dev). In production, the Rust binary serves the pre-built `dashboard/dist/` via `ServeDir` and `ServeFile` fallback, so there is a single process. + +### Source Layer Map + +``` +src/ +├── main.rs # Axum server entry: logging, DB pool, migrations, vec0 table, +│ # client/service construction, route registration, AppState assembly +├── lib.rs # Config struct + from_env() loading from .env +├── api/ # HTTP handlers, AppState, StandardPaper type +│ ├── mod.rs # AppState (shared state), StandardPaper, handlers re-exports +│ ├── agent.rs # SSE chat_agent endpoint, session CRUD, metrics, audit log +│ ├── papers.rs # Search, download, parse, translate, embed, citations, library, export +│ ├── notes.rs # Highlight/note CRUD +│ ├── sync.rs # Meta-sync and asset-batch endpoints +│ ├── targets.rs # Target query/associate/extract, RAG chat, figure chat +│ └── helpers.rs # Shared DB helpers, format conversion, path validation +├── agent/ # ReAct-based research agent (LLM-driven tool-use loop) +│ ├── tools/ # AgentTool trait, ToolRegistry, tool implementations per domain file +│ ├── runtime/ # ReAct loop engine, streaming, session management, context building +│ ├── compact/ # Context compression (micro/auto/manual layers) +│ ├── hooks.rs # Lifecycle events (PreToolUse/PostToolUse/Stop/etc.) +│ ├── skills.rs # SkillRegistry: loads skill SKILL.md files from skills/ directory +│ ├── subagent.rs # Context-isolated sub-agent runner for delegate_research +│ ├── background.rs# BgNotificationQueue for async slow-task (download/parse) notifications +│ └── team/ # Multi-agent team: file-based inbox, lead/teammate coordination +├── clients/ # External API wrappers +│ ├── llm.rs # LlmClient (OpenAI-compatible chat + streaming), EmbeddingClient +│ ├── ads.rs # NASA ADS API +│ ├── arxiv.rs # arXiv Atom XML API +│ └── qiniu.rs # Qiniu cloud storage +├── services/ # Business logic +│ ├── search.rs # Unified cross-source search (ADS + arXiv dedup) +│ ├── download.rs # PDF/HTML download with anti-bot measures and fallback chain +│ ├── parser/ # HTML/PDF → Markdown parsers (A&A, IOP, ar5iv, generic, PDF via MinerU) +│ ├── translation.rs# LLM bilingual translation with Trie-based astronomy glossary +│ ├── rag.rs # Embedding ingest + vector similarity retrieval + LLM answer generation +│ ├── target.rs # Celestial target extraction (IAU name regex) + CDS Sesame lookup +│ ├── chunker.rs # Markdown text chunking for embedding +│ ├── batch/ # Meta-sync (ADS bulk harvest) and asset-batch processing engines +│ ├── query_parser.rs# Advanced search query syntax parser +│ └── logging.rs # Pretty console + rolling file logger +└── bin/ + ├── health_check.rs # Library consistency checker and auto-repair + ├── cli.rs # CLI interface + └── reparse.rs # Re-parse existing library items +``` + +### AppState — Central Shared State + +All handlers access state via `Arc`. Key fields: + +- `db: SqlitePool` — SQLite connection pool (5 max connections, foreign keys enforced) +- `llm: LlmClient` / `embedding: EmbeddingClient` — OpenAI-compatible LLM clients +- `ads: AdsClient` / `arxiv: ArxivClient` — academic search clients +- `skill_registry: Arc>` — hot-reloaded agent skills +- `cancelled_runs: Arc>>` — agent cancellation tokens +- `harvest_status` / `batch_status` — async batch operation status tracking + +### Agent System Design + +The agent (`src/agent/`) implements a **ReAct** (Thought → Action → Observation) loop: + +1. **`AgentRuntime`** (`runtime/mod.rs`) orchestrates the loop: session create/resume → context build → ReAct loop → finalize +2. **Streaming**: LLM response is streamed via SSE (`AgentStreamEvent`) — thought, tool_call, tool_result, text_delta, usage, error, done +3. **Tools**: Each tool implements `AgentTool` trait (name, description, JSON Schema parameters, execute). 19 tools in default registry including read_file, grep_files, glob_files, run_bash, file_write, file_edit, search_papers, download_paper, parse_paper, get_paper_content, rag_search, query_target, save_note, todo_write, compress_context, load_skill, delegate_research, plus optional background and team tools +4. **Parallel execution**: Same-turn tool calls execute concurrently via `executor::execute_parallel` +5. **Context compression**: Three layers — micro (placeholder replacement), auto (LLM summarization when over threshold), manual (compress_context tool). Protected by `CompactionCircuitBreaker` +6. **Skills** (`skills.rs`): Two-layer loading — system-reminder lists names (~20 tokens each), LLM calls `load_skill` to inject full SKILL.md content +7. **Sub-agents** (`subagent.rs`): `delegate_research` spawns a context-isolated sub-agent with its own ReAct loop, returning only the final summary +8. **Teams** (`team/`): File-based inbox directory per session for lead/teammate message passing +9. **Background tasks** (`background.rs`): Slow ops (download, parse) can run async; results inject via `BgNotificationQueue` before next LLM call + +Environment variables for agent tuning: `AGENT_MAX_STEPS` (default 8), `AGENT_TOOL_TIMEOUT_SECS` (default 120), `AGENT_MAX_TOOL_OUTPUT_CHARS` (default 4000), `AGENT_CONTEXT_CHAR_LIMIT` (default 16000), `AGENT_TOKEN_SOFT_LIMIT` / `AGENT_TOKEN_HARD_LIMIT`. + +### Database + +SQLite via `sqlx::sqlite`. Migrations in `migrations/` are auto-run on startup (`sqlx::migrate!("./migrations")`). Key tables: `papers`, `citations_references`, `notes`, `agent_sessions`, `agent_messages`, `agent_tasks`, `agent_audit_log`, `paper_chunks_content`. Vector embeddings use `sqlite-vec` (`vec_paper_chunks` virtual table, auto-registered before any DB connection). + +The embedding dimension is controlled by `EMBEDDING_DIM` env var (default 1536). On dimension mismatch, the vec table and chunk content are dropped and recreated. + +### Frontend (dashboard/) + +React 19 + TypeScript + Vite + Tailwind CSS 4. Features are organized by domain: + +- `features/search/` — Cross-source paper search panel +- `features/library/` — Local library management +- `features/reader/` — Bilingual reader with highlight annotations (KaTeX for math) +- `features/citation/` — Canvas-based force-directed citation graph +- `features/sync/` — Batch sync control panel +- `features/agent/` — Agent chat interface (SSE event consumption) +- `features/settings/` — System configuration + +Dependencies: `react-markdown` + `rehype-katex` + `remark-math` for Markdown/LaTeX rendering, `framer-motion` for animations, `lucide-react` for icons. + +### Obscura In-Process Browser + +The `obscura-inprocess` feature compiles `obscura-browser` and `obscura-net` directly into the binary, eliminating the need for external browser binaries. This is used for bypassing Cloudflare/WAF on PDF download. Enabled via `--features obscura-inprocess`. + +### Astronomy Glossary (dictionary.txt) + +A 1MB+ bilingual astronomy terminology file loaded at startup into a Trie tree for longest-match glossary construction, used by the translation service to guide LLM translations with domain-accurate term mappings. + +## Code Conventions + +- Use `anyhow` for application errors, `thiserror` for library-style typed errors +- Environment variables via `dotenvy` + `std::env::var`, with defaults in `Config::from_env()` +- SQL queries use parameterized bindings (`sqlx::query("...").bind(...)`) — never string interpolation +- API handlers take `State(Arc)` and return Axum-compatible responses +- Agent tools implement `AgentTool` trait; new tools register in `ToolRegistry::new()` +- Front-end build is triggered by `build.rs` (auto npm install + build when `dashboard/src/` changes) diff --git a/Cargo.lock b/Cargo.lock index 1e22faf..34de065 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -141,9 +141,12 @@ dependencies = [ "dotenvy", "flate2", "futures-util", + "glob", "hmac 0.12.1", "html2md", "libsqlite3-sys", + "lru 0.12.5", + "notify", "obscura-browser", "obscura-net", "quick-xml", @@ -152,6 +155,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "serde_yaml", "sha1 0.10.6", "sqlite-vec", "sqlx", @@ -164,6 +168,7 @@ dependencies = [ "url", "urlencoding", "uuid", + "walkdir", "zip", ] @@ -350,7 +355,7 @@ version = "0.71.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5f58bf3d7db68cfbac37cfc485a8d711e87e064c3d0fe0435b92f7a407f9d6b3" dependencies = [ - "bitflags", + "bitflags 2.13.0", "cexpr", "clang-sys", "itertools", @@ -370,7 +375,7 @@ version = "0.72.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "993776b509cfb49c750f11b8f07a46fa23e0a1386ffc01fb1e7d343efc387895" dependencies = [ - "bitflags", + "bitflags 2.13.0", "cexpr", "clang-sys", "itertools", @@ -397,6 +402,12 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e764a1d40d510daf35e07be9eb06e75770908c27d411ee6c92109c9840eaaf7" +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + [[package]] name = "bitflags" version = "2.13.0" @@ -464,7 +475,7 @@ version = "0.5.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2c5e60b8c8d282c86360cab651ded04ab0335a7b5390c8d34145cbeab8cacf5f" dependencies = [ - "bitflags", + "bitflags 2.13.0", "btls-sys", "foreign-types", "libc", @@ -1161,6 +1172,16 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -1466,6 +1487,8 @@ version = "0.15.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" dependencies = [ + "allocator-api2", + "equivalent", "foldhash 0.1.5", ] @@ -1876,6 +1899,26 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inotify" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8069d3ec154eb856955c1c0fbffefbf5f3c40a104ec912d4797314c1801abff" +dependencies = [ + "bitflags 1.3.2", + "inotify-sys", + "libc", +] + +[[package]] +name = "inotify-sys" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e05c02b5e89bff3b946cedeca278abc628fe811e604f027c45a8aa3cf793d0eb" +dependencies = [ + "libc", +] + [[package]] name = "inout" version = "0.2.2" @@ -1976,6 +2019,26 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "kqueue" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "273c0752728918e0ac4976f2b275b6fefb9ecd400585dec929419f3844cd87b5" +dependencies = [ + "kqueue-sys", + "libc", +] + +[[package]] +name = "kqueue-sys" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07293a4e297ac234359b510362495713f75ea345d5307140414f20c69ffeb087" +dependencies = [ + "bitflags 2.13.0", + "libc", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -2025,7 +2088,7 @@ version = "0.1.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f02ab6bace2054fb888a3c16f990117b579d14a3088e472d63c6011fa185c9d3" dependencies = [ - "bitflags", + "bitflags 2.13.0", "libc", "plain", "redox_syscall 0.8.1", @@ -2081,6 +2144,15 @@ version = "0.4.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" +[[package]] +name = "lru" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "234cf4f4a04dc1f57e24b96cc0cd600cf2af460d4161ac5ecdd0af8e1f3b2a38" +dependencies = [ + "hashbrown 0.15.5", +] + [[package]] name = "lru" version = "0.18.0" @@ -2225,6 +2297,18 @@ dependencies = [ "simd-adler32", ] +[[package]] +name = "mio" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4a650543ca06a924e8b371db273b2756685faae30f8487da1b56505a8f78b0c" +dependencies = [ + "libc", + "log", + "wasi", + "windows-sys 0.48.0", +] + [[package]] name = "mio" version = "1.2.1" @@ -2269,6 +2353,23 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "notify" +version = "6.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6205bd8bb1e454ad2e27422015fb5e4f2bcc7e08fa8f27058670d208324a4d2d" +dependencies = [ + "bitflags 2.13.0", + "filetime", + "inotify", + "kqueue", + "libc", + "log", + "mio 0.8.11", + "walkdir", + "windows-sys 0.48.0", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -2857,7 +2958,7 @@ version = "0.5.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -2866,7 +2967,7 @@ version = "0.8.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5b44b894f2a6e36457d665d1e08c3866add6ed5e70050c1b4ba8a8ddedb02ce7" dependencies = [ - "bitflags", + "bitflags 2.13.0", ] [[package]] @@ -2988,7 +3089,7 @@ version = "0.38.44" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdb5bc1ae2baa591800df16c9ca78619bf65c0488b41b96ccec5d11220d8c154" dependencies = [ - "bitflags", + "bitflags 2.13.0", "errno", "libc", "linux-raw-sys 0.4.15", @@ -3001,7 +3102,7 @@ version = "1.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" dependencies = [ - "bitflags", + "bitflags 2.13.0", "errno", "libc", "linux-raw-sys 0.12.1", @@ -3116,7 +3217,7 @@ version = "0.26.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fd568a4c9bb598e291a08244a5c1f5a8a6650bee243b5b0f8dbb3d9cc1d87fe8" dependencies = [ - "bitflags", + "bitflags 2.13.0", "cssparser", "derive_more", "fxhash", @@ -3216,6 +3317,19 @@ dependencies = [ "v8", ] +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + [[package]] name = "servo_arc" version = "0.4.3" @@ -3503,7 +3617,7 @@ checksum = "1ed31390216d20e538e447a7a9b959e06ed9fc51c37b514b46eb758016ecd418" dependencies = [ "atoi", "base64 0.21.7", - "bitflags", + "bitflags 2.13.0", "byteorder", "bytes", "chrono", @@ -3546,7 +3660,7 @@ checksum = "7c824eb80b894f926f89a0b9da0c7f435d27cdd35b8c655b114e58223918577e" dependencies = [ "atoi", "base64 0.21.7", - "bitflags", + "bitflags 2.13.0", "byteorder", "chrono", "crc", @@ -3900,7 +4014,7 @@ checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", - "mio", + "mio 1.2.1", "parking_lot", "pin-project-lite", "signal-hook-registry", @@ -3998,7 +4112,7 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1e9cd434a998747dd2c4276bc96ee2e0c7a2eadf3cae88e52be55a05fa9053f5" dependencies = [ - "bitflags", + "bitflags 2.13.0", "bytes", "futures-util", "http", @@ -4024,7 +4138,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "async-compression", - "bitflags", + "bitflags 2.13.0", "bytes", "futures-core", "futures-util", @@ -4235,6 +4349,12 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + [[package]] name = "untrusted" version = "0.9.0" @@ -4296,7 +4416,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "33995a1fee055ff743281cde33a41f0d618ee0bdbe8bdf6859e11864499c2595" dependencies = [ "bindgen 0.71.1", - "bitflags", + "bitflags 2.13.0", "fslock", "gzip-header", "home", @@ -4484,7 +4604,7 @@ version = "0.244.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" dependencies = [ - "bitflags", + "bitflags 2.13.0", "hashbrown 0.15.5", "indexmap", "semver", @@ -4864,7 +4984,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" dependencies = [ "anyhow", - "bitflags", + "bitflags 2.13.0", "indexmap", "log", "serde", @@ -4911,7 +5031,7 @@ dependencies = [ "httparse", "ipnet", "libc", - "lru", + "lru 0.18.0", "percent-encoding", "pin-project-lite", "socket2", diff --git a/Cargo.toml b/Cargo.toml index 2199515..ff7f866 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -53,6 +53,11 @@ sqlite-vec = "0.1.9" clap = { version = "4", features = ["derive"] } async-trait = "0.1" async-stream = "0.3" +serde_yaml = "0.9" +notify = { version = "6", default-features = false, features = ["macos_kqueue"] } +glob = "0.3" +walkdir = "2" +lru = "0.12" [features] default = [] diff --git a/build.rs b/build.rs index bbd4688..da016d4 100644 --- a/build.rs +++ b/build.rs @@ -1,6 +1,6 @@ // build.rs -use std::process::Command; use std::path::Path; +use std::process::Command; fn main() { // 声明：只有当 dashboard/src/ 目录或 build.rs 发生改变时，才重新触发构建脚本 diff --git a/docs/agent-optimization-analysis.md b/docs/agent-optimization-analysis.md new file mode 100644 index 0000000..8be0bcf --- /dev/null +++ b/docs/agent-optimization-analysis.md @@ -0,0 +1,738 @@ +# Agent 架构优化分析 + +> 对比 Claude Code 源码 (`/home/fmq/program/claudecode/src/`) 与 AstroResearch Agent (`src/agent/`)， +> 基于 2026-06-16 的代码快照。 + +--- + +## 总体评估 + +我们的 Agent 已经实现了一个功能完整的 ReAct 研究引擎，涵盖了工具注册/调度、三层上下文压缩、生命周期 Hooks、 +多 Agent 团队协作、子代理委托、后台任务、Skills 加载等关键子系统。与 Claude Code 的架构范式高度一致。 + +以下按**影响优先级**列出可优化领域。 + +--- + +## 一、CRITICAL：Streaming Tool Executor + +### 现状 +`executor.rs` 在 LLM 流式响应**完全结束后**才通过 `join_all` 并行执行工具。 + +```rust +// 当前流程：LLM stream → 收集所有 tool_use blocks → join_all 执行 +let results = futures_util::future::join_all(exec_futs).await; +``` + +### Claude Code 做法 +`StreamingToolExecutor` 在模型**仍在生成** tool_use blocks 时就开始调度执行： + +``` +模型输出 tool_use(Read file A) → 立即开始读 A +模型输出 tool_use(Read file B) → 立即开始读 B（并发安全） +模型输出 tool_use(Bash cmd) → 排队等待（非并发安全） +模型输出结束 → 此时 A 和 B 可能已完成 +``` + +### 优化方案 + +```rust +/// 流式工具执行器 —— 模型还在输出时就开始执行工具 +pub struct StreamingToolExecutor { + tools: Vec, + tool_registry: Arc, + tool_context: ToolContext, + max_concurrency: usize, +} + +enum ToolStatus { + Queued, + Executing, + Completed, + Yielded, +} + +struct TrackedTool { + id: String, + block: ToolCall, + status: ToolStatus, + is_concurrency_safe: bool, + handle: Option>, + results: Option>, + pending_progress: Vec, +} + +impl StreamingToolExecutor { + /// 模型每输出一个 tool_use block 就调用此方法 + pub fn add_tool(&mut self, block: ToolCall) { + let is_safe = self.tool_registry + .get(&block.name) + .map(|t| t.is_concurrency_safe(&block.args)) + .unwrap_or(false); + + self.tools.push(TrackedTool { + id: block.id.clone(), + block, + status: ToolStatus::Queued, + is_concurrency_safe: is_safe, + handle: None, + results: None, + pending_progress: vec![], + }); + + tokio::spawn(async { self.process_queue().await }); + } + + /// 非阻塞获取已完成的工具结果 + pub fn get_completed_results(&mut self) -> Vec { + // 按顺序 yield 已完成的结果 + // 非并发安全的工具保持顺序 + // 进度消息立即 yield + } + + /// 等待所有剩余工具完成 + pub async fn get_remaining_results(&mut self) -> Vec { + // 等待 executing 的工具完成 + // 然后 yield 所有结果 + } +} +``` + +**预期收益**：大幅降低端到端延迟，尤其是当模型并行输出多个独立的 Read/Search 类工具调用时。 + +--- + +## 二、HIGH：工具并发分区 + +### 现状 +`executor.rs` 对所有工具调用一律使用 `join_all` 并行执行，不考虑工具的并发安全性。 + +### Claude Code 做法 +`partitionToolCalls()` 将工具调用分区为： +1. **并发安全批次** — 连续的 `isConcurrencySafe=true` 工具（如 Read、Grep、WebSearch） +2. **串行批次** — 单个 `isConcurrencySafe=false` 工具（如 Bash、Edit、Write） + +并发批次用 `all()` 并行执行（max concurrency = 10），串行批次逐个执行。 + +### 优化方案 + +```rust +/// 工具特征增加并发安全声明 +#[async_trait] +pub trait AgentTool: Send + Sync { + fn name(&self) -> &str; + fn description(&self) -> &str; + fn parameters(&self) -> serde_json::Value; + + /// 工具是否可以与其他并发安全的工具同时执行 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true // 默认只读工具是并发安全的 + } + + /// 最大并发数（默认无限制） + fn max_concurrency(&self) -> Option { + None + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput; +} + +/// 分区工具调用 +fn partition_tool_calls( + calls: &[PreparedCall], + registry: &ToolRegistry, +) -> Vec { + let mut batches: Vec = vec![]; + for call in calls { + let is_safe = registry.get(&call.tool_name) + .map(|t| t.is_concurrency_safe(&call.args)) + .unwrap_or(false); + + if is_safe && batches.last().map_or(false, |b| b.is_concurrency_safe) { + batches.last_mut().unwrap().calls.push(call.clone()); + } else { + batches.push(Batch { + is_concurrency_safe: is_safe, + calls: vec![call.clone()], + }); + } + } + batches +} +``` + +**预期收益**：避免 Bash/Write 等有副作用的工具与其他工具竞争导致的不确定性。 + +--- + +## 三、HIGH：Sibling Abort（兄弟中止） + +### 现状 +当一个工具执行出错时，其他并行执行的工具继续运行，浪费资源。 + +### Claude Code 做法 +`StreamingToolExecutor` 中： +- 当 Bash 工具出错时，`siblingAbortController.abort("sibling_error")` 中止所有兄弟 Bash 执行 +- Read/WebFetch 等独立工具的失败不影响其他工具 +- 被中止的工具获得 synthetic error message + +### 优化方案 + +```rust +/// 在并行执行时注入 sibling abort 信号 +pub struct SiblingAbortController { + abort_sender: tokio::sync::broadcast::Sender, + errored_tool_description: Arc>>, +} + +enum SiblingAbortReason { + SiblingError { description: String }, + UserInterrupted, + StreamingFallback, +} + +impl SiblingAbortController { + /// 当工具出错时调用，如果是 Bash 类工具则通知所有兄弟 + pub fn notify_error(&self, tool_name: &str, tool_desc: &str) { + if is_bash_like_tool(tool_name) { + let _ = self.abort_sender.send(SiblingAbortReason::SiblingError { + description: tool_desc.to_string(), + }); + } + } + + /// 每个工具执行前检查是否已被兄弟中止 + pub fn check_aborted(&self, this_tool: &str) -> Option { + // 如果是本工具报的错，不生成 synthetic error（避免重复） + } +} +``` + +**预期收益**：避免无效的后续执行，减少等待时间和 API 调用浪费。 + +--- + +## 四、HIGH：Error Recovery Ladder（错误恢复阶梯） + +### 现状 +`compact.rs` 只有**主动压缩**（在达到 token 限制前触发）。如果压缩不够激进，413 `prompt_too_long` 错误会直接暴露给用户。 + +### Claude Code 做法 +`query.ts` 实现了多层恢复阶梯： + +``` +第1层：Context Collapse drain（便宜，commit 已 staged 的 collapse） + ↓ 失败/不可用 +第2层：Reactive Compact（fork agent 摘要整个对话） + ↓ 失败 +第3层：Max Output Tokens Escalate（临时提升到 64k token cap） + ↓ 再次命中 +第4层：Multi-turn Recovery（注入 meta message，继续对话） + ↓ 全部失败 +最终：Surface the error（暴露给用户） +``` + +每层都有 `hasAttempted` 守卫防止无限重试，autocompact 有 circuit breaker（连续 3 次失败后停止）。 + +### 优化方案 + +```rust +/// 错误恢复策略枚举 +enum RecoveryStrategy { + /// 尝试更激进的 micro_compact + AggressiveMicroCompact, + /// LLM 摘要整个对话历史 + ReactiveCompact, + /// 提升 max_tokens 上限 + MaxTokensEscalate, + /// 注入 metacognitive 消息 + MultiTurnRecovery, + /// 放弃，暴露错误给用户 + Surface, +} + +struct RecoveryState { + attempted_micro_compact: bool, + attempted_reactive_compact: bool, + attempted_max_tokens_escalation: bool, + autocompact_failure_count: u32, +} + +const MAX_AUTOCOMPACT_FAILURES: u32 = 3; + +impl RecoveryState { + fn next_strategy(&mut self, error: &ModelError) -> RecoveryStrategy { + match error { + ModelError::ContextOverflow(_) => { + if !self.attempted_micro_compact { + self.attempted_micro_compact = true; + return RecoveryStrategy::AggressiveMicroCompact; + } + if !self.attempted_reactive_compact + && self.autocompact_failure_count < MAX_AUTOCOMPACT_FAILURES + { + self.attempted_reactive_compact = true; + return RecoveryStrategy::ReactiveCompact; + } + if !self.attempted_max_tokens_escalation { + self.attempted_max_tokens_escalation = true; + return RecoveryStrategy::MaxTokensEscalate; + } + RecoveryStrategy::MultiTurnRecovery + } + _ => RecoveryStrategy::Surface, + } + } +} +``` + +**预期收益**：显著提高长对话的鲁棒性，减少用户遇到的 "context too long" 错误。 + +--- + +## 五、MEDIUM：Time-Based Microcompact + +### 现状 +我们的 micro_compact 只基于消息数量/大小触发，不考虑时间因素。 + +### Claude Code 做法 +`microcompactMessages()` 首先检查 `evaluateTimeBasedTrigger()`： +- 计算距上一条 assistant 消息的时间间隔 +- 如果超过配置阈值（如 5 分钟），服务器的 prompt cache 已经过期 +- 此时直接 content-clear 旧的 tool results（保留最近 N 个） +- 因为 cache 已冷，修改消息内容不会有额外代价 + +### 优化方案 + +```rust +pub struct TimeBasedMCConfig { + pub enabled: bool, + /// 触发阈值（分钟） + pub gap_threshold_minutes: u64, + /// 保留最近 N 个工具结果 + pub keep_recent: usize, +} + +impl Default for TimeBasedMCConfig { + fn default() -> Self { + Self { + enabled: true, + gap_threshold_minutes: 5, + keep_recent: 4, + } + } +} + +/// 检查时间触发是否应激活 +fn evaluate_time_based_trigger( + messages: &[ChatMessage], + config: &TimeBasedMCConfig, +) -> Option { + let last_assistant = messages.iter() + .rev() + .find(|m| m.role == MessageRole::Assistant)?; + + let elapsed = last_assistant.timestamp.elapsed().unwrap_or_default(); + let gap_minutes = elapsed.as_secs() / 60; + + if gap_minutes >= config.gap_threshold_minutes { + Some(TimeBasedTrigger { gap_minutes }) + } else { + None + } +} +``` + +**预期收益**：在长时间闲置后自动清理过期上下文，防止用户回到对话时遇上 context overflow。 + +--- + +## 六、MEDIUM：增强的 Hook 系统 + +### 现状 +5 个生命周期事件，简单的 `HookAction::Continue/Block` 二元决策。 + +### Claude Code 做法 +27 个 hook 事件，4 种 hook 类型（command/prompt/agent/http），结构化 JSON 协议： + +```json +{ + "continue": true, + "decision": "approve", + "reason": "...", + "systemMessage": "...", + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "allow", + "updatedInput": { ... }, + "additionalContext": "..." + } +} +``` + +### 优化建议（按价值排序） + +1. **PreToolUse 输入修改**：Hook 可以修改工具参数后再执行（如自动修正 bibcode 格式） +2. **PostToolUse 输出修改**：Hook 可以后处理工具结果（如自动翻译、格式化） +3. **additionalContext 注入**：Hook 可以向 LLM 注入附加上下文 +4. **SessionStart watchPaths**：启动时注册文件监控路径 +5. **SubagentStart/Stop**：子代理生命周期事件 + +```rust +/// 增强的 PreToolUse 输出 +pub struct PreToolUseHookOutput { + pub decision: HookDecision, + pub updated_input: Option, + pub additional_context: Option, + pub system_message: Option, +} + +pub enum HookDecision { + Allow, + Deny { reason: String }, + Ask { reason: String }, +} + +/// 增强的 PostToolUse 输出 +pub struct PostToolUseHookOutput { + pub updated_output: Option, + pub additional_context: Option, +} +``` + +--- + +## 七、MEDIUM：Permission Pipeline（权限管道） + +### 现状 +工具没有权限系统。所有工具对 Agent 都同样可用。 + +### Claude Code 做法 +多层权限评估管道： +``` +Step 1: Deny rule 匹配 → deny（不可覆盖） +Step 2: Ask rule 匹配 → ask（除非 sandbox override） +Step 3: Tool.checkPermissions → 工具自身逻辑 +Step 4: Safety checks → ask（绕过免疫） +Step 5: bypassPermissions 模式 → allow +Step 6: Allow rule 匹配 → allow +Step 7: 默认 → ask +``` + +权限规则格式：`ToolName(pattern:*)`，支持多源优先级链。 + +### 优化方案 + +```rust +/// 权限行为 +pub enum PermissionBehavior { + Allow, + Deny, + Ask, +} + +/// 权限规则来源（优先级从高到低） +pub enum RuleSource { + Policy, // 企业策略 + UserSettings, + ProjectSettings, + LocalSettings, + CliArg, + Session, +} + +/// 权限规则 +pub struct PermissionRule { + pub source: RuleSource, + pub behavior: PermissionBehavior, + pub tool_pattern: String, // "search_papers" 或 "bash(git *)" + pub content_pattern: Option, +} + +/// 权限检查器 +pub struct PermissionChecker { + rules: Vec, +} + +impl PermissionChecker { + pub fn check( + &self, + tool_name: &str, + input: &serde_json::Value, + ) -> PermissionDecision { + // 1. 检查 deny 规则（不可覆盖） + // 2. 检查 ask 规则 + // 3. 工具自身 check_permissions + // 4. 安全检查 + // 5. bypass 模式 + // 6. allow 规则 + // 7. 默认 ask + } +} +``` + +--- + +## 八、MEDIUM：Progress Streaming（进度流式传输） + +### 现状 +工具执行期间没有任何进度反馈，直到执行完成才发送结果。 + +### Claude Code 做法 +`StreamingToolExecutor` 支持工具的进度消息立即 yield，即使工具还在执行中。进度消息类型为 `"progress"`，在 UI 中显示为短暂的状态更新。 + +### 优化方案 + +```rust +/// 为长时间运行的工具增加进度回调 +#[async_trait] +pub trait AgentTool: Send + Sync { + // ... 现有方法 ... + + /// 带进度回调的执行（默认委托给 execute） + async fn execute_with_progress( + &self, + args: serde_json::Value, + ctx: &ToolContext, + progress: mpsc::UnboundedSender, + ) -> ToolOutput { + let _ = progress; // 默认忽略 + self.execute(args, ctx).await + } +} + +/// 进度更新 +pub struct ProgressUpdate { + pub tool_call_id: String, + pub message: String, + pub percentage: Option, +} +``` + +**适用工具**：download_paper（下载进度）、parse_paper（解析进度）、search_papers（搜索进度）。 + +--- + +## 九、MEDIUM：Token Budget Management + +### 现状 +没有 token 预算跟踪。Agent 可以无限制地消耗 tokens。 + +### Claude Code 做法 +- `budget.total` — 用户设定的 token 预算上限 +- `budget.spent()` — 当前已消耗的 output tokens +- `budget.remaining()` — 剩余可用 tokens +- 硬上限：达到 total 后 `agent()` 调用会抛错 +- 软上限：在接近限制时注入 nudge 消息提醒模型 + +### 优化方案 + +```rust +pub struct TokenBudget { + total: Option, + spent_output_tokens: u64, + spent_input_tokens: u64, +} + +impl TokenBudget { + pub fn new(total: Option) -> Self { ... } + + pub fn record_usage(&mut self, input: u64, output: u64) { + self.spent_input_tokens += input; + self.spent_output_tokens += output; + } + + pub fn remaining(&self) -> Option { + self.total.map(|t| t.saturating_sub(self.spent_output_tokens)) + } + + /// 在接近限制时生成提醒消息 + pub fn nudge_message(&self) -> Option { + if let (Some(total), Some(rem)) = (self.total, self.remaining()) { + if rem < total / 10 { + Some(format!( + "注意：token 预算已使用 {:.0}%，剩余约 {} tokens。请尽快给出最终答案。", + (self.spent_output_tokens as f64 / total as f64) * 100.0, + rem + )) + } else { + None + } + } else { + None + } + } +} +``` + +--- + +## 十、LOW-MEDIUM：Interrupt Behavior（中断行为分类） + +### 现状 +取消信号对所有工具一视同仁。 + +### Claude Code 做法 +每个工具声明 `interruptBehavior()`： +- `"cancel"` — 用户中断时立即取消（如 Read、Search） +- `"block"` — 用户中断时继续执行完毕（如 Edit、Write，防止文件损坏） + +### 优化方案 + +```rust +pub enum InterruptBehavior { + /// 用户中断时立即取消（默认，适合只读工具） + Cancel, + /// 用户中断时继续执行到完成（适合写入工具） + Block, +} + +#[async_trait] +pub trait AgentTool: Send + Sync { + // ... 现有方法 ... + + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Cancel // 默认安全取消 + } +} +``` + +--- + +## 十一、LOW：Circuit Breaker 模式 + +### 现状 +compact 失败没有熔断机制，可能无限重试。 + +### Claude Code 做法 +- `MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES = 3` +- fallback model 调用遇到 `529 Overloaded` 时切换到备选模型 +- 所有恢复策略都有 `hasAttempted` 守卫 + +### 优化方案 + +```rust +pub struct CircuitBreaker { + max_failures: u32, + failure_count: u32, + state: CircuitState, +} + +enum CircuitState { + Closed, // 正常工作 + Open, // 熔断，拒绝请求 + HalfOpen, // 试探性恢复 +} + +impl CircuitBreaker { + pub fn check(&mut self) -> Result<(), CircuitOpenError> { + match self.state { + CircuitState::Open => Err(CircuitOpenError), + CircuitState::HalfOpen | CircuitState::Closed => Ok(()), + } + } + + pub fn record_success(&mut self) { + self.failure_count = 0; + self.state = CircuitState::Closed; + } + + pub fn record_failure(&mut self) { + self.failure_count += 1; + if self.failure_count >= self.max_failures { + self.state = CircuitState::Open; + } + } +} +``` + +--- + +## 十二、LOW：Context Collapse / Projection System（架构级） + +### 现状 +压缩直接修改消息数组，原始上下文永久丢失。 + +### Claude Code 做法 +`ContextCollapse` 采用 **commit log + projection** 模式： +1. 不再直接修改消息 +2. 将旧的上下文段替换为摘要 + metadata +3. 摘要存储在独立的 collapse store 中 +4. 每次查询循环入口通过 `projectView()` 重放 commit log 重建临时消息视图 +5. Commit 是 staged（先暂存）再 committed（on overflow） + +### 适用场景 +我们的论文研究场景中，Agent 可能会在同一个 session 中研究多篇论文。Context collapse 可以在切换论文时保留之前的研究摘要而不是完全丢弃。 + +### 优化方案 + +```rust +/// 上下文段 +pub struct ContextSegment { + pub id: String, + pub summary: String, + pub original_message_count: usize, + pub original_token_estimate: usize, + pub created_at: chrono::DateTime, +} + +/// Collapse 存储 +pub struct CollapseStore { + segments: Vec, + commit_log: Vec, +} + +impl CollapseStore { + /// Stage 一个 collapse（还不提交） + pub fn stage(&mut self, segment: ContextSegment) { ... } + + /// Commit 所有 staged collapses + pub fn commit_staged(&mut self) -> usize { ... } + + /// 重放 commit log，生成当前消息视图 + pub fn project_view( + &self, + recent_messages: &[ChatMessage], + ) -> Vec { ... } +} +``` + +--- + +## 实施优先级建议 + +| 优先级 | 优化项 | 预计工作量 | 收益 | +|--------|--------|-----------|------| +| P0 | Streaming Tool Executor | 3-5 天 | 延迟大幅降低 | +| P1 | 工具并发分区 | 1-2 天 | 正确性提升 | +| P1 | Error Recovery Ladder | 2-3 天 | 鲁棒性大幅提升 | +| P1 | Sibling Abort | 1 天 | 资源浪费减少 | +| P2 | Time-Based Microcompact | 1 天 | 长会话体验 | +| P2 | Token Budget Management | 1-2 天 | 成本控制 | +| P2 | Progress Streaming | 1-2 天 | UX 提升 | +| P3 | 增强 Hook 系统 | 2-3 天 | 可扩展性 | +| P3 | Permission Pipeline | 2-3 天 | 安全性 | +| P3 | Interrupt Behavior | 0.5 天 | 可靠性 | +| P4 | Circuit Breaker | 1 天 | 稳定性 | +| P4 | Context Collapse | 5-7 天 | 长期架构 | + +--- + +## 架构范式已对齐的部分 + +以下方面我们的实现已经与 Claude Code 的范式高度一致，无需大幅改动： + +1. ✅ ReAct Loop 结构（Thought → Act → Observe） +2. ✅ Tool trait + Registry 模式（虽可细化但已完整） +3. ✅ 三层上下文压缩（micro/auto/aggressive） +4. ✅ 安全切割点（不切断 tool_call/tool_result 配对） +5. ✅ Transcript 持久化（压缩前保存） +6. ✅ 生命周期 Hooks（有基础的 5 事件） +7. ✅ 重复调用检测（DuplicateDetector） +8. ✅ Skills 两层加载系统 +9. ✅ 子代理委托（SubAgentRunner） +10. ✅ 后台任务系统（BgNotificationQueue + mpsc） +11. ✅ 多 Agent 团队（TeamManager + file inbox） +12. ✅ SSE 流式事件到前端 +13. ✅ Todo/Task 持久化到 SQLite +14. ✅ Session 生命周期管理 diff --git a/migrations/20260616000000_agent_tasks.sql b/migrations/20260616000000_agent_tasks.sql new file mode 100644 index 0000000..73fd532 --- /dev/null +++ b/migrations/20260616000000_agent_tasks.sql @@ -0,0 +1,31 @@ +-- 智能体任务持久化表 +-- 支持 DAG 依赖模式 (blocked_by) 和任务状态生命周期 +-- +-- Status lifecycle: +-- pending -> in_progress -> completed +-- (可回退: in_progress -> pending) +-- +-- blocked_by 存储 JSON 数组格式: ["task_id_1", "task_id_2"] +-- 应用层负责验证 DAG 有效性（无自引用、无循环依赖） + +CREATE TABLE IF NOT EXISTS agent_tasks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + task_id TEXT NOT NULL, + content TEXT NOT NULL DEFAULT '', + status TEXT NOT NULL DEFAULT 'pending' + CHECK(status IN ('pending', 'in_progress', 'completed')), + blocked_by TEXT NOT NULL DEFAULT '[]', + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE +); + +-- 按会话查询任务（最常用） +CREATE INDEX IF NOT EXISTS idx_agent_tasks_session ON agent_tasks(session_id); + +-- 按状态过滤（用于恢复/展示） +CREATE INDEX IF NOT EXISTS idx_agent_tasks_status ON agent_tasks(session_id, status); + +-- 保证同一会话内 task_id 唯一 +CREATE UNIQUE INDEX IF NOT EXISTS idx_agent_tasks_session_task ON agent_tasks(session_id, task_id); diff --git a/migrations/20260617000000_agent_audit_log.sql b/migrations/20260617000000_agent_audit_log.sql new file mode 100644 index 0000000..fda774f --- /dev/null +++ b/migrations/20260617000000_agent_audit_log.sql @@ -0,0 +1,19 @@ +-- 智能体审计日志表 +-- 记录所有工具调用的详细信息：工具名称、执行状态、耗时、输出预览 +-- +-- status: OK (成功) / FAIL (失败) / SESSION_STOP (会话终止) + +CREATE TABLE IF NOT EXISTS agent_audit_log ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + step INTEGER NOT NULL DEFAULT 0, + tool_name TEXT, + status TEXT NOT NULL CHECK(status IN ('OK', 'FAIL', 'SESSION_STOP')), + elapsed_ms INTEGER NOT NULL DEFAULT 0, + output_preview TEXT, + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE +); + +CREATE INDEX IF NOT EXISTS idx_audit_log_session ON agent_audit_log(session_id); +CREATE INDEX IF NOT EXISTS idx_audit_log_created ON agent_audit_log(created_at); diff --git a/migrations/20260618000000_agent_identity.sql b/migrations/20260618000000_agent_identity.sql new file mode 100644 index 0000000..4043a83 --- /dev/null +++ b/migrations/20260618000000_agent_identity.sql @@ -0,0 +1,36 @@ +-- 智能体身份隔离与团队协作支持 +-- +-- 为子代理（delegate_research）和多智能体团队（spawn_teammate） +-- 提供 agent 粒度的消息隔离、审计追踪和任务归属。 +-- +-- 所有 ALTER TABLE 使用 DEFAULT 值，保证向后兼容： +-- 现有数据自动标记为 'lead'，旧代码无需修改。 + +-- 1. agent_messages: 消息归属 +ALTER TABLE agent_messages ADD COLUMN agent_name TEXT NOT NULL DEFAULT 'lead'; + +-- 2. agent_audit_log: 审计归属 +ALTER TABLE agent_audit_log ADD COLUMN agent_name TEXT NOT NULL DEFAULT 'lead'; + +-- 3. agent_tasks: 任务分配目标（支持角色路由） +ALTER TABLE agent_tasks ADD COLUMN owner TEXT NOT NULL DEFAULT ''; + +-- 4. 新建: 团队成员注册表 +CREATE TABLE IF NOT EXISTS agent_team_members ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + agent_name TEXT NOT NULL, + agent_role TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'active' + CHECK(status IN ('spawning', 'active', 'idle', 'shutdown')), + created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + FOREIGN KEY (session_id) REFERENCES agent_sessions(session_id) ON DELETE CASCADE, + UNIQUE(session_id, agent_name) +); + +CREATE INDEX IF NOT EXISTS idx_team_members_session ON agent_team_members(session_id); +CREATE INDEX IF NOT EXISTS idx_team_members_status ON agent_team_members(session_id, status); + +-- 重建 agent_messages 索引（包含 agent_name 过滤加速） +CREATE INDEX IF NOT EXISTS idx_agent_messages_agent ON agent_messages(session_id, agent_name); diff --git a/skills/README.md b/skills/README.md new file mode 100644 index 0000000..42471c1 --- /dev/null +++ b/skills/README.md @@ -0,0 +1,45 @@ +# Agent Skills + +按需加载的领域知识模块。参考 Claude Code 的三层设计： + +- **Layer 1（系统提示词）**：仅列出 skill 名称（~20 tokens/skill） +- **Layer 2（load_skill 工具）**：LLM 按需调用，注入完整 SKILL.md 内容 + +## 目录约定 + +``` +skills/{name}/SKILL.md ← 必须是子目录 + SKILL.md +``` + +当前仅读取 `SKILL.md` 文件。skill 目录下可以放其他文件（如 `references/`、`scripts/`、`assets/`），但当前版本的加载逻辑不会处理它们。 + +## SKILL.md 格式 + +```markdown +--- +name: skill-name # 唯一标识（必填） +description: 一句话描述 # 必填 +context: inline # 已解析但未生效（依赖子代理架构） +allowed-tools: # 已解析但未生效（无权限系统） + - bash + - read +--- + +# Skill 正文 +``` + +`context` 和 `allowed-tools` 字段会被 `parse_frontmatter()` 解析，`LoadSkillTool` 执行时会在返回内容中标记 `context: fork` 的 skill，但实际执行行为不变——所有 skill 当前都是内联返回 Markdown 文本。 + +## 添加新 skill + +1. 创建 `skills/{name}/SKILL.md`，填写 frontmatter 和正文 +2. 系统提示词在下次请求时自动发现 +3. 无需改动代码 + +## 当前 Skills + +| Skill | 状态 | +|-------|------| +| methodology | 就绪 | +| plotting | 占位 | +| presentation | 占位 | diff --git a/skills/methodology/SKILL.md b/skills/methodology/SKILL.md new file mode 100644 index 0000000..bb6e352 --- /dev/null +++ b/skills/methodology/SKILL.md @@ -0,0 +1,46 @@ +--- +name: methodology +description: 系统性文献综述方法论——如何高效地完成学术文献调研 +--- + +# 系统性文献综述方法论 + +## 工作流程 + +当用户请求进行文献综述或调研时，按以下步骤进行： + +### 1. 范围界定 +- 理解用户的研究问题，提取核心关键词和同义词 +- 使用 `search_papers` 进行 broad search（rows=10~20），了解领域规模 +- 向用户确认搜索范围（时间跨度、子领域、是否包含预印本） + +### 2. 按引用数筛选 +- 优先关注高被引文献（citation_count > 10） +- 同时保留近期重要成果（year >= 当前年份-2） +- 使用 `get_paper_metadata` 查看候选文献的完整摘要 + +### 3. 逐篇深读 +- 对筛选出的核心文献（通常 5-10 篇），依次： + 1. `download_paper` → `parse_paper` → `get_paper_content` + 2. 提取每篇的关键发现、方法论、数据来源 + 3. 记录文献之间的引用关系和争议点 + +### 4. 补充检索 +- 核心文献的参考文献中如有高频出现但未检索到的文献，用 `search_papers` 补查 +- 使用 `rag_search` 检查本地文献库是否有相关内容 + +### 5. 交叉验证 +- 对关键结论，检查是否有其他独立研究得出一致/矛盾的结果 +- 注意作者群体和机构是否有明显的学术派系倾向 + +### 6. 输出综述 +- 使用 `save_note` 保存完整的综述报告 +- 结构：摘要 → 引言 → 方法 → 主要发现 → 争议与共识 → 未来方向 → 参考文献 +- 所有引用使用 ADS bibcode 标注 +- 数学公式使用 LaTeX 格式 + +## 质量检查清单 +- [ ] 是否覆盖了近 5 年的主要文献？ +- [ ] 引用的结论是否有文献支持？ +- [ ] 不同观点是否得到了平衡呈现？ +- [ ] 是否标注了各文献的方法局限性？ diff --git a/skills/plotting/SKILL.md b/skills/plotting/SKILL.md new file mode 100644 index 0000000..4e590fe --- /dev/null +++ b/skills/plotting/SKILL.md @@ -0,0 +1,36 @@ +--- +name: plotting +description: 科研绘图规范 +context: fork +allowed-tools: + - bash + - save_note +--- + +# 科研绘图规范（TODO：待细化） + +## 目标期刊要求（占位） + +- 分辨率：通常 300-600 dpi +- 格式：矢量图优先（PDF/SVG），光栅图备选（PNG） +- 字体：需与期刊正文一致（如 Times New Roman） + +## 常用图表类型（占位） + +| 类型 | Python 库 | 适用场景 | +|------|----------|---------| +| 光谱图 | matplotlib | 光谱分析、能谱 | +| 赫罗图 | matplotlib | 恒星演化 | +| 光变曲线 | matplotlib | 变星、超新星 | +| 参数分布 | seaborn | 统计分布、角图 | +| 3D 轨道 | plotly | 天体力学 | + +## 工作流程（TODO） + +1. 准备数据（从文献或天体数据库提取） +2. 生成 Python 脚本 +3. 使用 bash 工具执行脚本 +4. 检查输出图片 +5. 保存到指定路径 + +**注意：此 skill 为预留占位，具体实现待补充。** diff --git a/skills/presentation/SKILL.md b/skills/presentation/SKILL.md new file mode 100644 index 0000000..38ba0f1 --- /dev/null +++ b/skills/presentation/SKILL.md @@ -0,0 +1,34 @@ +--- +name: presentation +description: 学术PPT生成规范 +context: fork +allowed-tools: + - bash + - save_note +--- + +# 学术 PPT 生成规范（TODO：待细化） + +## 标准结构（占位） + +1. 标题页：报告标题、作者、日期 +2. 背景与动机：1-2 页 +3. 数据与方法：2-3 页 +4. 主要结果：3-5 页（每页一张关键图 + 要点） +5. 讨论与结论：1-2 页 +6. 参考文献：1 页 + +## 视觉规范（占位） + +- 配色：学术风格（深蓝 + 白 + 强调色） +- 字体：标题 32pt，正文 24pt +- 图表：高清、统一配色、标注来源 +- 公式：LaTeX 渲染 + +## 生成方式（TODO） + +- 方案 A：Python-pptx 生成 .pptx 文件 +- 方案 B：LaTeX Beamer 生成 PDF 幻灯片 +- 方案 C：Markdown → Marp 转换 + +**注意：此 skill 为预留占位，具体实现待补充。** diff --git a/src/agent/autonomous.rs b/src/agent/autonomous.rs new file mode 100644 index 0000000..8c4e4e5 --- /dev/null +++ b/src/agent/autonomous.rs @@ -0,0 +1,137 @@ +// src/agent/autonomous.rs +// +// 自治研究循环 — 空闲时自动轮询新任务。 +// 参考 learn-claude-code s17 Autonomous Agents。 +// +// 当 Agent 完成当前回合后，进入 IDLE 阶段： +// 1. 检查是否有未认领的团队任务 +// 2. 检查批量同步任务队列 +// 3. 检查订阅分类的新论文 +// 发现工作后自动认领并执行。 + +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Notify; +use tracing::{info, warn}; + +use crate::api::AppState; + +/// 自治研究配置 +#[derive(Debug, Clone)] +pub struct AutoResearchConfig { + /// 是否启用自治模式 + pub enabled: bool, + /// 订阅的 arXiv 分类 + pub subscribed_categories: Vec, + /// 最大自治轮次（防止无限循环） + pub max_autonomous_turns: usize, + /// IDLE 超时（分钟） + pub idle_timeout_minutes: u64, +} + +impl Default for AutoResearchConfig { + fn default() -> Self { + AutoResearchConfig { + enabled: false, + subscribed_categories: vec!["astro-ph".to_string()], + max_autonomous_turns: 5, + idle_timeout_minutes: 60, + } + } +} + +/// IDLE 轮询器 — 在 Agent 空闲时检查是否有待处理工作。 +pub struct IdlePoller { + app_state: Arc, + config: AutoResearchConfig, + poll_interval: Duration, + /// 唤醒通知（当外部事件触发时，如新论文到达、任务分配） + wake_notify: Arc, +} + +impl IdlePoller { + pub fn new(app_state: Arc, config: AutoResearchConfig) -> Self { + IdlePoller { + app_state, + config, + poll_interval: Duration::from_secs(30), + wake_notify: Arc::new(Notify::new()), + } + } + + /// 获取唤醒通知器的 clone（供外部触发） + pub fn wake_sender(&self) -> Arc { + self.wake_notify.clone() + } + + /// 启动 IDLE 循环（应在独立的 tokio::spawn 中运行） + pub async fn start(self) { + info!( + "[IdlePoller] 启动自治轮询 (间隔={:?}, 最大轮次={})", + self.poll_interval, self.config.max_autonomous_turns + ); + + let mut autonomous_turns: usize = 0; + + loop { + // 等待 poll_interval 或被 notify 唤醒 + tokio::select! { + _ = tokio::time::sleep(self.poll_interval) => {} + _ = self.wake_notify.notified() => { + info!("[IdlePoller] 被外部事件唤醒"); + } + } + + if !self.config.enabled { + continue; + } + + if autonomous_turns >= self.config.max_autonomous_turns { + info!( + "[IdlePoller] 已达到最大自治轮次 ({})，停止轮询", + self.config.max_autonomous_turns + ); + break; + } + + // 1. 检查未认领的团队任务 + let task_board = crate::agent::task_board::TaskBoard::new(self.app_state.db.clone()); + match task_board.list_available_tasks(5).await { + Ok(tasks) if !tasks.is_empty() => { + for task in tasks { + if task.can_start { + info!( + "[IdlePoller] 发现可认领任务: {} (session={})", + task.task_id, task.session_id + ); + if let Ok(true) = task_board + .claim_task(&task.session_id, &task.task_id, "auto") + .await + { + autonomous_turns += 1; + // 创建 AgentRuntime 并执行任务 + let runtime = crate::agent::runtime::AgentRuntime::new( + self.app_state.clone(), + ); + let (tx, _rx) = tokio::sync::mpsc::unbounded_channel(); + let _ = runtime + .run_turn(Some(task.session_id.clone()), &task.content, tx) + .await; + } + } + } + continue; + } + Ok(_) => {} + Err(e) => { + warn!("[IdlePoller] 任务查询失败: {}", e); + } + } + + // 2. 检查批量同步状态 (placeholder) + // 未来: 检查订阅分类的新论文并自动触发同步 + } + + info!("[IdlePoller] IDLE 循环结束"); + } +} diff --git a/src/agent/background.rs b/src/agent/background.rs new file mode 100644 index 0000000..9eb3a92 --- /dev/null +++ b/src/agent/background.rs @@ -0,0 +1,203 @@ +// src/agent/background.rs +// +// 后台任务执行子系统（参考 Claude Code s08 Background Tasks）。 +// +// 慢速操作（download_paper, parse_paper, embed_paper）可通过 +// bg_task_run 在后台异步执行，LLM 继续思考/调用其他工具。 +// +// 完成的通知通过 BgNotificationQueue 在下一轮 LLM 调用前注入。 + +use serde::Serialize; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::{mpsc, Mutex}; +use tracing::info; + +use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry}; +use crate::api::AppState; + +/// 后台任务结果 +#[derive(Debug, Clone)] +pub struct BgTaskResult { + pub task_id: String, + pub tool_name: String, + pub bibcode: String, + pub is_error: bool, + pub summary: String, +} + +/// 后台任务状态跟踪 +#[derive(Debug, Clone, Serialize)] +pub struct BgTaskHandle { + pub task_id: String, + pub tool_name: String, + pub bibcode: String, + pub status: BgTaskStatus, +} + +#[derive(Debug, Clone, Serialize)] +#[serde(rename_all = "lowercase")] +pub enum BgTaskStatus { + Running, + Completed, + Failed, +} + +/// 后台任务通知队列。 +/// +/// 使用 mpsc channel 在后台完成和主循环之间传递结果。 +/// drain() 在每轮 LLM 调用前被调用，收集所有已完成的后台任务结果。 +pub struct BgNotificationQueue { + rx: Mutex>, + tx: mpsc::UnboundedSender, + /// 内存中的任务状态注册表 + tasks: Mutex>, +} + +impl Default for BgNotificationQueue { + fn default() -> Self { + Self::new() + } +} + +impl BgNotificationQueue { + /// 创建新的通知队列 + pub fn new() -> Self { + let (tx, rx) = mpsc::unbounded_channel(); + BgNotificationQueue { + rx: Mutex::new(rx), + tx, + tasks: Mutex::new(HashMap::new()), + } + } + + /// 获取发送端（供后台任务使用） + pub fn sender(&self) -> mpsc::UnboundedSender { + self.tx.clone() + } + + /// 注册一个开始执行的后台任务 + pub async fn register_task(&self, task: BgTaskHandle) { + let mut tasks = self.tasks.lock().await; + tasks.insert(task.task_id.clone(), task); + } + + /// 更新任务状态 + pub async fn update_task_status(&self, task_id: &str, status: BgTaskStatus) { + let mut tasks = self.tasks.lock().await; + if let Some(task) = tasks.get_mut(task_id) { + task.status = status; + } + } + + /// 获取所有任务状态 + pub async fn get_all_tasks(&self) -> Vec { + let tasks = self.tasks.lock().await; + tasks.values().cloned().collect() + } + + /// 获取单个任务状态 + pub async fn get_task(&self, task_id: &str) -> Option { + let tasks = self.tasks.lock().await; + tasks.get(task_id).cloned() + } + + /// 排空所有已完成的后台任务通知。 + /// 在每轮 LLM 调用前调用。 + pub async fn drain(&self) -> Vec { + let mut results = Vec::new(); + let mut rx = self.rx.lock().await; + while let Ok(result) = rx.try_recv() { + results.push(result); + } + results + } +} + +/// 在后台执行指定的工具调用。 +/// +/// 启动一个 tokio::spawn 异步任务执行工具， +/// 完成后通过通知队列发送结果。 +pub async fn spawn_background_task( + app_state: Arc, + queue: Arc, + tool_name: String, + bibcode: String, +) -> BgTaskHandle { + let task_id = uuid::Uuid::new_v4().to_string(); + // 取前 8 位便于显示 + let short_id = task_id[..8].to_string(); + + let handle = BgTaskHandle { + task_id: short_id.clone(), + tool_name: tool_name.clone(), + bibcode: bibcode.clone(), + status: BgTaskStatus::Running, + }; + + queue.register_task(handle.clone()).await; + + let queue_clone = queue.clone(); + let app_state_clone = app_state.clone(); + let tool_name_clone = tool_name.clone(); + let bibcode_clone = bibcode.clone(); + let short_id_clone = short_id.clone(); + let sender = queue.sender(); + + tokio::spawn(async move { + info!( + "[Background] 启动后台任务 {}: {} ({})", + short_id_clone, tool_name_clone, bibcode_clone + ); + + // 构造 ToolContext 和参数 (后台任务：静默模式) + let tool_ctx = ToolContext::silent(app_state_clone.clone()); + let args = serde_json::json!({"bibcode": bibcode_clone}); + let tool_registry = ToolRegistry::new(app_state_clone.skill_registry.clone()); + + let output = match tool_registry.get(&tool_name_clone) { + Some(tool) => { + match tokio::time::timeout( + std::time::Duration::from_secs(300), // 5 min timeout for bg tasks + tool.execute(args, &tool_ctx), + ) + .await + { + Ok(o) => o, + Err(_) => ToolOutput::error("后台任务执行超时（300秒）"), + } + } + None => ToolOutput::error(format!("未知工具: {}", tool_name_clone)), + }; + + // 发送完成通知 + let result = BgTaskResult { + task_id: short_id_clone.clone(), + tool_name: tool_name_clone, + bibcode: bibcode_clone, + is_error: output.is_error, + summary: if output.content.len() > 500 { + let preview: String = output.content.chars().take(500).collect(); + format!("{}...", preview) + } else { + output.content.clone() + }, + }; + + let _ = sender.send(result); + queue_clone + .update_task_status( + &short_id_clone, + if output.is_error { + BgTaskStatus::Failed + } else { + BgTaskStatus::Completed + }, + ) + .await; + + info!("[Background] 后台任务 {} 完成", short_id_clone); + }); + + handle +} diff --git a/src/agent/compact.rs b/src/agent/compact.rs new file mode 100644 index 0000000..f7b2121 --- /dev/null +++ b/src/agent/compact.rs @@ -0,0 +1,721 @@ +// src/agent/compact.rs +// +// 上下文压缩子系统。 +// 实现四层压缩策略（参考 Claude Code compaction pipeline）： +// 0. snip_compact — 零API调用：消息数超阈值时截断中间段 +// 1. micro_compact — 轻量级：替换较早的工具结果为占位符 +// 2. auto_compact — 自动触发：超 token 阈值时 LLM 摘要对话历史 +// 3. manual_compact — 手动触发：Agent 通过 compress_context 工具主动调用 +// +// 参考 Claude Code src/services/compact/ 模块设计。 +// +// P0 改进： +// - Transcript 持久化（压缩前保存完整 JSONL） +// - 多层回退链（snip → micro → auto → aggressive_micro → identity inject） +// - Identity re-injection（压缩后消息过少时注入身份确认） +// - 占位符优化（使用工具名称替代字符预览） + +pub mod collapse; + +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; +use tracing::{info, warn}; + +use super::hooks::{HookRegistry, PostCompactContext, PreCompactContext}; + +/// 递归守卫：防止压缩内部触发的 LLM 调用再次触发压缩。 +static COMPACTING: AtomicBool = AtomicBool::new(false); + +use crate::clients::llm::{ChatMessage, LlmClient, MessageRole}; + +/// 获取 transcripts 存储目录 +fn transcripts_dir() -> PathBuf { + PathBuf::from(".transcripts") +} + +/// 找到安全的上下文切割点，确保不会切断 tool_call / tool_result 配对。 +/// 从末尾向前扫描，如果候选切割点的第一条要保留的消息是 tool 角色， +/// 则向前追溯到对应的 assistant(tool_calls) 消息一并保留。 +pub fn find_safe_cut_point(messages: &[ChatMessage], desired_keep: usize) -> usize { + if messages.len() <= desired_keep { + return 0; // 全部保留，无需切割 + } + + let mut cut = messages.len().saturating_sub(desired_keep); + + // 确保不从 system 消息之后的第一条就开始切（至少保留 system） + if cut == 0 { + cut = 1; + } + + // 如果切割点落在一个 tool 消息上，向前扩展以包含其对应的 assistant(tool_calls) + loop { + if cut >= messages.len() { + cut = messages.len() - 1; + break; + } + + let first_kept = &messages[cut]; + + if first_kept.role == MessageRole::Tool { + let mut found_pair = false; + for j in (0..cut).rev() { + if messages[j].role == MessageRole::Assistant && messages[j].tool_calls.is_some() { + cut = j; + found_pair = true; + break; + } + } + if !found_pair { + break; + } + } else { + break; + } + } + + // 检查切割点之前没有孤立的 assistant(tool_calls)。 + // 从后往前扫描：每找到一个孤立的 assistant，将 cut 移到该位置并继续向前检查。 + let mut search = cut; + while search > 0 { + let mut found = false; + for j in (0..search).rev() { + if messages[j].role == MessageRole::Assistant && messages[j].tool_calls.is_some() { + let has_tool_result = messages[j + 1..search] + .iter() + .any(|m| m.role == MessageRole::Tool); + if !has_tool_result { + cut = j; + search = j; + found = true; + } + break; // 只处理最近的一个，继续向前 + } + } + if !found { + break; + } + } + + cut +} + +// ── snip_compact (Layer 0) ────────────────────────────────────────────────── + +/// 最大消息数（超过此阈值触发 snip_compact） +pub const MAX_MESSAGES: usize = 50; +/// 保留的头部消息数（system prompt + 初始上下文） +pub const HEAD_KEEP: usize = 3; + +/// Layer 0 压缩：当消息数超过 `MAX_MESSAGES` 时，保留前 `HEAD_KEEP` 条 +/// + 后 `MAX_MESSAGES - HEAD_KEEP` 条，中间替换为占位消息。 +/// +/// 这是零 API 调用的最廉价压缩层。使用 `find_safe_cut_point` 确保 +/// 切割点不会破坏 `assistant(tool_calls)` / `tool_result` 配对。 +/// +/// 返回 `true` 表示执行了压缩。 +pub fn snip_compact(messages: &mut Vec, max_messages: usize) -> bool { + if messages.len() <= max_messages { + return false; + } + + let tail_keep = max_messages - HEAD_KEEP; + let original_len = messages.len(); + + // 找到安全的尾部起点（复用已有配对保护逻辑） + let tail_start = find_safe_cut_point(messages, tail_keep); + + // 确保不跟头部重叠 + if tail_start <= HEAD_KEEP { + return false; + } + + // 收集被移除段中使用的工具名称 + let snipped = &messages[HEAD_KEEP..tail_start]; + let mut tool_names: Vec = Vec::new(); + for msg in snipped { + if msg.role == MessageRole::Tool { + if let Some(call_id) = &msg.tool_call_id { + if let Some(name) = find_tool_name_for_call_id(messages, call_id) { + if !tool_names.contains(&name) { + tool_names.push(name); + } + } + } + } + } + + let tool_list = if tool_names.is_empty() { + String::new() + } else { + let mut unique = tool_names; + unique.sort(); + unique.dedup(); + format!(" 使用过的工具: {}.", unique.join(", ")) + }; + + let snipped_count = snipped.len(); + let placeholder = ChatMessage::user(format!( + "[上下文压缩] 省略了 {} 条中间对话消息（第 {}-{} 条）。{}", + snipped_count, + HEAD_KEEP + 1, + tail_start, + tool_list + )); + + // 移除中间段，替换为占位消息 + messages.drain(HEAD_KEEP..tail_start); + messages.insert(HEAD_KEEP, placeholder); + + info!( + "[snipCompact] {} → {} 条消息 (移除 {} 条, 切割点: {})", + original_len, + messages.len(), + snipped_count, + tail_start + ); + + true +} + +/// 根据 tool_call_id 查找对应的工具名称。 +fn find_tool_name_for_call_id(messages: &[ChatMessage], tool_call_id: &str) -> Option { + for msg in messages.iter().rev() { + if msg.role == MessageRole::Assistant { + if let Some(tool_calls) = &msg.tool_calls { + for tc in tool_calls { + if tc.id == tool_call_id { + return Some(tc.function.name.clone()); + } + } + } + } + } + None +} + +/// 轻量级压缩：将较早的工具结果替换为简短占位符，释放上下文空间。 +/// 保留最近 `keep_recent` 条工具结果不变。 +/// P0 改进：使用 `[Previous: used {tool_name}]` 替代字符预览，节省 ~80 tokens/条。 +pub fn micro_compact(messages: &mut [ChatMessage], keep_recent: usize) { + let tool_info: Vec<(usize, String, String)> = messages + .iter() + .enumerate() + .filter_map(|(i, m)| { + if m.role == MessageRole::Tool { + let call_id = m.tool_call_id.clone().unwrap_or_default(); + // 先查找工具名称（此时 messages 是不可变借用） + let tool_name = find_tool_name_for_call_id(messages, &call_id) + .unwrap_or_else(|| "unknown".to_string()); + Some((i, call_id, tool_name)) + } else { + None + } + }) + .collect(); + + let compact_count = tool_info.len().saturating_sub(keep_recent); + if compact_count == 0 { + return; + } + + for (idx, _tool_id, tool_name) in tool_info.iter().take(compact_count) { + if let Some(msg) = messages.get_mut(*idx) { + msg.content = Some(format!("[Previous: used {}]", tool_name)); + } + } +} + +/// 粗略估算消息列表的 token 数（用作首次调用的近似值）。 +/// 后续迭代优先使用 API 返回的精确 prompt_tokens。 +pub fn rough_estimate_tokens(messages: &[ChatMessage]) -> usize { + messages + .iter() + .map(|m| { + let content_len = m.content.as_ref().map_or(0, |c| c.len()); + content_len + 4 // 消息结构 overhead 约 4 token + }) + .sum() +} + +/// 保存完整 transcript 到磁盘（JSONL 格式）。 +/// 在压缩前调用，确保不丢失任何对话历史。 +async fn save_transcript(messages: &[ChatMessage], session_id: &str) { + let dir = transcripts_dir(); + if let Err(e) = std::fs::create_dir_all(&dir) { + warn!("[Compact] 无法创建 transcripts 目录 {:?}: {}", dir, e); + return; + } + + let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S"); + let filename = format!("{}_{}.jsonl", session_id, timestamp); + let path = dir.join(&filename); + + let mut content = String::new(); + for msg in messages { + if let Ok(json) = serde_json::to_string(msg) { + content.push_str(&json); + content.push('\n'); + } + } + + match std::fs::write(&path, &content) { + Ok(_) => info!( + "[Compact] Transcript 已保存: {} ({} 条消息)", + path.display(), + messages.len() + ), + Err(e) => warn!("[Compact] Transcript 保存失败: {}", e), + } +} + +/// 在压缩后注入身份确认块，防止模型丢失上下文认知。 +/// 参考 Claude Code s11: identity re-injection after compression. +fn inject_identity_block(messages: &mut Vec) { + if messages.len() <= 4 { + // 消息过少说明压缩非常激进，注入身份提醒 + let identity = ChatMessage::user( + "[身份确认] 你是一位专业的天体物理学研究助手。以上是历史对话的压缩摘要。\ + 你正在进行的研究任务是回答用户的问题。请基于摘要中的关键信息继续工作，\ + 需要更多信息时主动使用工具搜索。", + ); + // 插入在 system 消息和 summary 之后、recent 消息之前 + let insert_pos = if messages + .first() + .is_some_and(|m| m.role == MessageRole::System) + { + 2.min(messages.len()) + } else { + 1.min(messages.len()) + }; + messages.insert(insert_pos, identity); + info!( + "[Compact] 注入身份确认块（压缩后仅 {} 条消息）", + messages.len() - 1 + ); + } +} + +/// 使用 LLM 生成对话摘要。 +async fn generate_summary(to_summarize: &[ChatMessage], llm: &LlmClient) -> Result { + let summary_content: String = to_summarize + .iter() + .filter_map(|m| { + let role = match m.role { + MessageRole::User => "用户", + MessageRole::Assistant => "助手", + MessageRole::Tool => "工具", + _ => return None, + }; + m.content.as_ref().map(|c| { + let preview: String = c.chars().take(200).collect(); + format!("[{}] {}", role, preview) + }) + }) + .collect::>() + .join("\n"); + + let summary_prompt = format!( + "请用简洁的中文总结以下对话历史的要点（不超过500字）：\n\n{}", + summary_content + ); + + llm.chat_completion( + "你是一个对话摘要助手。请提取对话的关键信息和结论。", + &summary_prompt, + ) + .await + .map_err(|e| { + warn!("[Compact] 上下文摘要生成失败: {}，尝试激进压缩", e); + format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len()) + }) +} + +/// 多层回退压缩：snip → micro → auto → aggressive_micro → identity inject +async fn compress_with_fallback( + messages: &mut Vec, + llm: &LlmClient, + context_char_limit: usize, +) { + // Layer 0: snip_compact（零 API 调用，消息数超过 MAX_MESSAGES 时截断中间段） + snip_compact(messages, MAX_MESSAGES); + + // Layer 1: micro_compact（保留最近 8 条工具结果） + micro_compact(messages, 8); + if rough_estimate_tokens(messages) < (context_char_limit * 3 / 2) { + return; + } + + // Layer 2: auto_compact（LLM 摘要） + let system_msg = messages.first().cloned(); + let cut_point = find_safe_cut_point(messages, 10); + let to_summarize = &messages[1..cut_point]; + if to_summarize.is_empty() { + return; + } + + let summary = match generate_summary(to_summarize, llm).await { + Ok(s) => s, + Err(fallback) => fallback, + }; + + let recent = messages[cut_point..].to_vec(); + messages.clear(); + if let Some(sys) = system_msg { + messages.push(sys); + } + messages.push(ChatMessage::user(format!("[历史对话摘要]\n{}", summary))); + messages.extend(recent); + + // Layer 3: 如果摘要后仍然超限，激进 micro_compact（仅保留 2 条） + if rough_estimate_tokens(messages) >= (context_char_limit * 3 / 2) { + warn!("[Compact] LLM 摘要后仍超限，执行激进压缩 (keep_recent=2)"); + micro_compact(messages, 2); + } + + // Layer 4: 注入身份确认块 + inject_identity_block(messages); + + info!( + "[Compact] 上下文压缩完成，消息数: {} (安全切割点: {})", + messages.len(), + cut_point + ); +} + +/// 上下文压缩：保存 transcript → 多层回退压缩。 +/// 保留系统提示 + 最近的完整 tool-call/tool-result 配对。 +/// 使用 LLM 摘要较旧的对话历史，在 token 超限时触发。 +pub async fn compress_context( + messages: &mut Vec, + llm: &LlmClient, + context_char_limit: usize, + session_id: &str, +) { + compress_context_with_hooks(messages, llm, context_char_limit, session_id, None).await; +} + +/// 带 Hook 的上下文压缩变体。如果提供了 HookRegistry，会在压缩前后触发事件。 +pub async fn compress_context_with_hooks( + messages: &mut Vec, + llm: &LlmClient, + context_char_limit: usize, + session_id: &str, + hook_registry: Option<&HookRegistry>, +) { + if messages.len() <= 4 { + return; + } + + // 递归守卫：如果已在压缩中，跳过（防止嵌套压缩死循环） + if COMPACTING.swap(true, Ordering::SeqCst) { + warn!("[Compact] 递归守卫触发：已有进行中的压缩操作，跳过"); + return; + } + + let before_count = messages.len(); + let est_tokens = rough_estimate_tokens(messages); + + // OnPreCompact hook + if let Some(registry) = hook_registry { + registry + .run_on_pre_compact(&PreCompactContext { + session_id: session_id.to_string(), + message_count: before_count, + estimated_tokens: est_tokens, + }) + .await; + } + + // 压缩前保存完整 transcript（使用传入的 session_id 避免跨会话覆盖） + save_transcript(messages, session_id).await; + + // 执行多层回退压缩 + compress_with_fallback(messages, llm, context_char_limit).await; + + // OnPostCompact hook + if let Some(registry) = hook_registry { + registry + .run_on_post_compact(&PostCompactContext { + session_id: session_id.to_string(), + new_message_count: messages.len(), + compression_method: if messages.len() < before_count / 2 { + "llm_summary" + } else if messages.len() < before_count { + // snip_compact 会产生占位消息但保留尾部，micro_compact 替换内容 + "snip_or_micro" + } else { + "none" + } + .to_string(), + }) + .await; + } + + // 释放递归守卫 + COMPACTING.store(false, Ordering::SeqCst); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_micro_compact_placeholder_format() { + let mut messages = vec![ + ChatMessage::system("You are a helpful assistant."), + ChatMessage::user("Hello"), + ChatMessage::assistant_with_tool_calls( + Some("Let me search.".to_string()), + vec![crate::clients::llm::ToolCall { + id: "call_1".to_string(), + call_type: "function".to_string(), + function: crate::clients::llm::FunctionCall { + name: "search_papers".to_string(), + arguments: "{\"query\": \"black holes\"}".to_string(), + }, + }], + ), + ChatMessage { + role: MessageRole::Tool, + content: Some("Found 5 results about black holes".to_string()), + tool_call_id: Some("call_1".to_string()), + tool_calls: None, + name: None, + reasoning_content: None, + }, + ChatMessage::assistant("Here are the results..."), + ]; + + micro_compact(&mut messages, 0); + + // 工具结果应该被压缩为 [Previous: used search_papers] + let tool_msg = &messages[3]; + assert_eq!( + tool_msg.content.as_deref(), + Some("[Previous: used search_papers]") + ); + } + + #[test] + fn test_identity_block_injected_when_few_messages() { + let mut messages = vec![ + ChatMessage::system("You are a research assistant."), + ChatMessage::user("[历史对话摘要]\nPrevious discussion about black holes."), + ChatMessage::user("What about neutron stars?"), + ]; + + inject_identity_block(&mut messages); + + // 应该注入了身份确认块 + let identity_msg = &messages[2]; + assert!(identity_msg.content.as_ref().unwrap().contains("身份确认")); + assert!(identity_msg + .content + .as_ref() + .unwrap() + .contains("天体物理学研究助手")); + } + + #[test] + fn test_micro_compact_preserves_recent_results() { + let mut messages = vec![ + ChatMessage::system("You are a helpful assistant."), + ChatMessage::user("Search for papers"), + ChatMessage::assistant_with_tool_calls( + Some("Searching...".to_string()), + vec![crate::clients::llm::ToolCall { + id: "call_old".to_string(), + call_type: "function".to_string(), + function: crate::clients::llm::FunctionCall { + name: "search_papers".to_string(), + arguments: "{\"query\": \"old\"}".to_string(), + }, + }], + ), + ChatMessage { + role: MessageRole::Tool, + content: Some("Old result".to_string()), + tool_call_id: Some("call_old".to_string()), + tool_calls: None, + name: None, + reasoning_content: None, + }, + ChatMessage::assistant_with_tool_calls( + Some("Searching more...".to_string()), + vec![crate::clients::llm::ToolCall { + id: "call_new".to_string(), + call_type: "function".to_string(), + function: crate::clients::llm::FunctionCall { + name: "get_paper_content".to_string(), + arguments: "{\"bibcode\": \"2024A&A...\"}".to_string(), + }, + }], + ), + ChatMessage { + role: MessageRole::Tool, + content: Some("Paper content here...".to_string()), + tool_call_id: Some("call_new".to_string()), + tool_calls: None, + name: None, + reasoning_content: None, + }, + ]; + + micro_compact(&mut messages, 1); + + // 第一个工具结果应该被压缩 + let compressed = &messages[3]; + assert_eq!( + compressed.content.as_deref(), + Some("[Previous: used search_papers]") + ); + + // 最近的一个工具结果应该保留 + let recent = &messages[5]; + assert_eq!(recent.content.as_deref(), Some("Paper content here...")); + } + + #[test] + fn test_find_tool_name_for_call_id() { + let messages = vec![ + ChatMessage::assistant_with_tool_calls( + Some("Let me search.".to_string()), + vec![crate::clients::llm::ToolCall { + id: "call_abc".to_string(), + call_type: "function".to_string(), + function: crate::clients::llm::FunctionCall { + name: "rag_search".to_string(), + arguments: "{}".to_string(), + }, + }], + ), + ChatMessage { + role: MessageRole::Tool, + content: Some("Search result".to_string()), + tool_call_id: Some("call_abc".to_string()), + tool_calls: None, + name: None, + reasoning_content: None, + }, + ]; + + let name = find_tool_name_for_call_id(&messages, "call_abc"); + assert_eq!(name, Some("rag_search".to_string())); + + let name = find_tool_name_for_call_id(&messages, "nonexistent"); + assert_eq!(name, None); + } + + #[test] + fn test_rough_estimate_tokens() { + let messages = vec![ + ChatMessage::system("You are an assistant."), + ChatMessage::user("Hello, world!"), + ]; + let estimate = rough_estimate_tokens(&messages); + // 每个消息 content.len() + 4 overhead + let expected = ("You are an assistant.".len()) + 4 + ("Hello, world!".len()) + 4; + assert_eq!(estimate, expected); + } + + // ── snip_compact 测试 ── + + #[test] + fn test_snip_compact_below_threshold_noop() { + let mut messages: Vec = (0..40) + .map(|i| { + if i == 0 { + ChatMessage::system("System prompt") + } else { + ChatMessage::user(format!("Message {}", i)) + } + }) + .collect(); + let result = snip_compact(&mut messages, 50); + assert!(!result); + assert_eq!(messages.len(), 40); + } + + #[test] + fn test_snip_compact_above_threshold_truncates() { + let mut messages: Vec = (0..100) + .map(|i| { + if i == 0 { + ChatMessage::system("System prompt") + } else { + ChatMessage::user(format!("Message {}", i)) + } + }) + .collect(); + let original_len = messages.len(); + let result = snip_compact(&mut messages, 50); + assert!(result); + // 应该变成: HEAD_KEEP(3) + 1(placeholder) + remainder ≈ 51 + assert!(messages.len() < original_len); + assert!(messages.len() <= 51); // HEAD_KEEP + placeholder + tail + // 检查占位消息 + assert!(messages[HEAD_KEEP] + .content + .as_ref() + .unwrap() + .contains("省略")); + } + + #[test] + fn test_snip_compact_keeps_system_prompt() { + let sys = ChatMessage::system("You are an astrophysics research assistant."); + let mut messages: Vec = vec![sys.clone()]; + for i in 1..80 { + messages.push(ChatMessage::user(format!("Question {}", i))); + messages.push(ChatMessage::assistant(format!("Answer {}", i))); + } + snip_compact(&mut messages, 50); + // 第一条必须是 system 消息 + assert_eq!(messages[0].role, MessageRole::System); + assert_eq!( + messages[0].content.as_deref(), + Some("You are an astrophysics research assistant.") + ); + } + + #[test] + fn test_snip_compact_respects_tool_pairing() { + // 构建 tool_call/tool_result 配对靠近切割点的场景 + let mut messages = vec![ + ChatMessage::system("System"), + ChatMessage::user("Search"), + ChatMessage::assistant_with_tool_calls( + Some("Searching...".to_string()), + vec![crate::clients::llm::ToolCall { + id: "call_near_cut".to_string(), + call_type: "function".to_string(), + function: crate::clients::llm::FunctionCall { + name: "search_papers".to_string(), + arguments: "{}".to_string(), + }, + }], + ), + ChatMessage { + role: MessageRole::Tool, + content: Some("Result".to_string()), + tool_call_id: Some("call_near_cut".to_string()), + tool_calls: None, + name: None, + reasoning_content: None, + }, + ]; + // 填充到超过阈值 + for i in 0..60 { + messages.push(ChatMessage::user(format!("Padding {}", i))); + } + snip_compact(&mut messages, MAX_MESSAGES); + // 不应该有孤立的 tool 消息（没有对应 assistant(tool_calls)） + let has_orphan_tool = messages.windows(2).any(|w| { + w[0].role == MessageRole::Tool + && w[1].role != MessageRole::Tool + && (w[1].role != MessageRole::Assistant || w[1].tool_calls.is_none()) + }); + assert!(!has_orphan_tool); + } +} diff --git a/src/agent/compact/collapse.rs b/src/agent/compact/collapse.rs new file mode 100644 index 0000000..c584a7d --- /dev/null +++ b/src/agent/compact/collapse.rs @@ -0,0 +1,266 @@ +// src/agent/compact/collapse.rs +// +// 上下文折叠日志 — 结构化 commit log + projection 模式。 +// 参考 Claude Code ContextCollapse：将压缩记录为分段 commit， +// 需要时通过 projection 重放到消息层，避免直接修改原始历史。 +// +// 核心思路： +// 1. 每次压缩记录一个 CollapseCommit（范围 + 方法 + 摘要） +// 2. project() 将 commits 应用到消息列表上 +// 3. 超过 MAX_SEGMENTS 时触发溢出合并 + +use std::sync::atomic::{AtomicU64, Ordering}; +use tracing::info; + +/// 最大折叠分段数，超出后触发溢出合并 +const MAX_COLLAPSE_SEGMENTS: usize = 5; +/// 溢出摘要最大字符数 +const OVERFLOW_SUMMARY_MAX_CHARS: usize = 800; + +/// 压缩方法枚举 +#[derive(Debug, Clone, PartialEq)] +pub enum CollapseMethod { + /// 轻量级 micro 压缩（替换工具结果为占位符） + MicroCompact, + /// LLM 摘要压缩 + LlmSummary, + /// 激进 micro 压缩（只保留极少数工具结果） + AggressiveMicro, + /// 身份注入（消息过少时的回退） + IdentityInjection, +} + +impl CollapseMethod { + pub fn as_str(&self) -> &'static str { + match self { + CollapseMethod::MicroCompact => "micro", + CollapseMethod::LlmSummary => "llm_summary", + CollapseMethod::AggressiveMicro => "aggressive_micro", + CollapseMethod::IdentityInjection => "identity_injection", + } + } +} + +/// 单次压缩的 commit entry +#[derive(Debug, Clone)] +pub struct CollapseCommit { + /// 自增 commit ID + pub id: u64, + /// 压缩方法 + pub method: CollapseMethod, + /// 原始 messages 中被折叠的索引范围 (start, end_exclusive) + pub removed_range: (usize, usize), + /// 压缩后的摘要文本 + pub summary: String, + /// commit 时间戳 + pub timestamp: chrono::DateTime, +} + +/// 持久的折叠日志 — 记录和重放压缩历史。 +/// +/// 使用示例： +/// ```ignore +/// let mut log = CollapseLog::new(); +/// // 每次压缩后记录 +/// log.commit(CollapseMethod::MicroCompact, (5, 20), "摘要内容".into()); +/// // 可将 commits 投影到消息上 +/// log.project(&mut messages); +/// ``` +pub struct CollapseLog { + commits: Vec, + next_id: AtomicU64, +} + +impl Default for CollapseLog { + fn default() -> Self { + Self::new() + } +} + +impl CollapseLog { + /// 创建空的折叠日志 + pub fn new() -> Self { + CollapseLog { + commits: Vec::new(), + next_id: AtomicU64::new(1), + } + } + + /// 记录一次压缩 commit。 + /// + /// 返回新 commit 的 ID。 + pub fn commit( + &mut self, + method: CollapseMethod, + removed_range: (usize, usize), + summary: String, + ) -> u64 { + let id = self.next_id.fetch_add(1, Ordering::SeqCst); + let method_str = method.as_str(); + let summary_len = summary.len(); + self.commits.push(CollapseCommit { + id, + method, + removed_range, + summary, + timestamp: chrono::Utc::now(), + }); + info!( + "[CollapseLog] commit #{}: method={}, range={:?}, summary_len={}", + id, method_str, removed_range, summary_len + ); + id + } + + /// 获取 commit 总数 + pub fn len(&self) -> usize { + self.commits.len() + } + + /// 是否有记录 + pub fn is_empty(&self) -> bool { + self.commits.is_empty() + } + + /// 检查是否需要溢出合并（commits 数超过 MAX_COLLAPSE_SEGMENTS） + pub fn should_overflow(&self) -> bool { + self.commits.len() > MAX_COLLAPSE_SEGMENTS + } + + /// 获取最近的 N 条摘要 + pub fn recent_summaries(&self, n: usize) -> Vec { + self.commits + .iter() + .rev() + .take(n) + .map(|c| c.summary.clone()) + .collect() + } + + /// 将最老的 commits 合并为一个溢出摘要。 + /// + /// 返回溢出摘要（可插入到消息列表中），并从日志中移除已合并的 commits。 + pub fn overflow(&mut self) -> Option { + if self.commits.len() <= MAX_COLLAPSE_SEGMENTS { + return None; + } + + let merge_count = self.commits.len() - MAX_COLLAPSE_SEGMENTS + 1; + let to_merge: Vec<&CollapseCommit> = self.commits.iter().take(merge_count).collect(); + + let mut merged = String::from("[上下文压缩历史]\n"); + for commit in &to_merge { + let summary_preview: String = commit + .summary + .chars() + .take(OVERFLOW_SUMMARY_MAX_CHARS) + .collect(); + merged.push_str(&format!( + "- (方法: {}) {}\n", + commit.method.as_str(), + summary_preview + )); + } + + // 移除已合并的 commits（保留最后 MAX_COLLAPSE_SEGMENTS-1 个） + self.commits.drain(0..merge_count); + + info!( + "[CollapseLog] 溢出合并: {} commits → {} chars, 剩余 {} commits", + merge_count, + merged.len(), + self.commits.len() + ); + + Some(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_commit_and_len() { + let mut log = CollapseLog::new(); + assert_eq!(log.len(), 0); + assert!(log.is_empty()); + + log.commit(CollapseMethod::MicroCompact, (5, 10), "summary 1".into()); + assert_eq!(log.len(), 1); + assert!(!log.is_empty()); + + log.commit(CollapseMethod::LlmSummary, (0, 5), "summary 2".into()); + assert_eq!(log.len(), 2); + } + + #[test] + fn test_should_overflow() { + let mut log = CollapseLog::new(); + // 添加 6 个 commits（超过 MAX_COLLAPSE_SEGMENTS=5） + for i in 0..6 { + log.commit( + CollapseMethod::MicroCompact, + (i * 10, i * 10 + 5), + format!("commit {}", i), + ); + } + assert!(log.should_overflow()); + } + + #[test] + fn test_no_overflow_when_under_limit() { + let mut log = CollapseLog::new(); + for i in 0..5 { + log.commit( + CollapseMethod::MicroCompact, + (i * 10, i * 10 + 5), + format!("commit {}", i), + ); + } + assert!(!log.should_overflow()); + } + + #[test] + fn test_overflow_merges_oldest() { + let mut log = CollapseLog::new(); + // 添加 7 个 commits + for i in 0..7 { + log.commit( + CollapseMethod::MicroCompact, + (i * 10, i * 10 + 5), + format!("summary for commit {}", i), + ); + } + assert!(log.should_overflow()); + + let merged = log.overflow(); + assert!(merged.is_some()); + // 溢出后应剩余 MAX_COLLAPSE_SEGMENTS-1 = 4 个 commits（合并了 3 个） + assert!(!log.should_overflow()); + assert_eq!(log.len(), 4); + } + + #[test] + fn test_recent_summaries() { + let mut log = CollapseLog::new(); + for i in 0..3 { + log.commit( + CollapseMethod::MicroCompact, + (i, i + 1), + format!("summary {}", i), + ); + } + + let recent = log.recent_summaries(2); + assert_eq!(recent.len(), 2); + assert_eq!(recent[0], "summary 2"); // 最新的在前 + assert_eq!(recent[1], "summary 1"); + } + + #[test] + fn test_empty_log_no_overflow() { + let log = CollapseLog::new(); + assert!(!log.should_overflow()); + } +} diff --git a/src/agent/hooks.rs b/src/agent/hooks.rs new file mode 100644 index 0000000..05c884a --- /dev/null +++ b/src/agent/hooks.rs @@ -0,0 +1,974 @@ +// src/agent/hooks.rs +// +// Agent 生命周期 Hooks 系统。 +// 参考 Claude Code 的 PreToolUse / PostToolUse / Stop hooks 设计， +// 提供可扩展的事件回调链，支持： +// - OnSessionStart — 会话创建/恢复时 +// - PreToolUse — 工具执行前（可拦截/阻止/修改输入/注入上下文） +// - PostToolUse — 工具执行后（审计日志、指标采集、输出修改） +// - OnStepComplete — 每步结束（指标更新、上下文检查） +// - OnSessionStop — 会话终止（清理、持久化） +// - OnSubagentStart — 子代理启动时 +// - OnSubagentStop — 子代理停止时 +// - OnPreCompact — 上下文压缩前 +// - OnPostCompact — 上下文压缩后 +// +// P3 增强（参考 Claude Code hooks 协议）： +// - PreToolUseAction 支持 MutateInput（修改工具参数）和 PermissionRequired +// - PostToolUseAction 支持 MutateOutput（修改工具输出） +// - 扩展生命周期事件（SubagentStart/Stop, PreCompact/PostCompact） +// - HookRegistry 返回累积的 additional_context 和 mutate_output + +use async_trait::async_trait; +use sqlx::SqlitePool; +use std::collections::HashMap; +use std::sync::Arc; +use tracing::{info, warn}; + +use super::terminal::TurnTerminal; + +// ── Hook Contexts ── + +/// 会话启动上下文 +#[derive(Debug, Clone)] +pub struct SessionStartContext { + pub session_id: String, + pub turn_index: i32, + pub is_resume: bool, +} + +/// PreToolUse hook 上下文 +#[derive(Debug, Clone)] +pub struct PreToolUseContext { + pub session_id: String, + pub tool_name: String, + pub tool_args: serde_json::Value, + pub step: usize, +} + +/// PostToolUse hook 上下文 +#[derive(Debug, Clone)] +pub struct PostToolUseContext { + pub session_id: String, + pub agent_name: String, + pub tool_name: String, + pub tool_args: serde_json::Value, + pub output_content: String, + pub is_error: bool, + pub step: usize, + pub elapsed_ms: u64, +} + +/// Step 完成上下文 +#[derive(Debug, Clone)] +pub struct StepCompleteContext { + pub session_id: String, + pub step: usize, + pub max_steps: usize, + pub messages_count: usize, + pub estimated_tokens: usize, + pub token_limit: usize, +} + +/// Session 终止上下文 +#[derive(Debug, Clone)] +pub struct SessionStopContext<'a> { + pub session_id: String, + pub terminal: &'a TurnTerminal, + pub total_steps: usize, +} + +/// 子代理启动上下文 +#[derive(Debug, Clone)] +pub struct SubagentStartContext { + pub parent_session_id: String, + pub subagent_name: String, + pub prompt: String, +} + +/// 子代理停止上下文 +#[derive(Debug, Clone)] +pub struct SubagentStopContext { + pub parent_session_id: String, + pub subagent_name: String, + pub result_summary: String, + pub steps: usize, + pub is_error: bool, +} + +/// 压缩前上下文 +#[derive(Debug, Clone)] +pub struct PreCompactContext { + pub session_id: String, + pub message_count: usize, + pub estimated_tokens: usize, +} + +/// 压缩后上下文 +#[derive(Debug, Clone)] +pub struct PostCompactContext { + pub session_id: String, + pub new_message_count: usize, + pub compression_method: String, +} + +// ── Hook Actions ── + +/// PreToolUse hook 返回的增强动作。 +/// 支持：允许、阻止、修改输入、权限请求。 +#[derive(Debug, Clone)] +pub enum PreToolUseAction { + /// 允许继续执行（默认） + Continue, + /// 阻止执行，附带原因 + Block { reason: String }, + /// 允许执行但修改输入参数或注入附加上下文 + MutateInput { + updated_args: serde_json::Value, + additional_context: Option, + }, + /// 需要权限决策 + PermissionRequired { + permission: String, + tool_name: String, + }, +} + +impl PreToolUseAction { + pub fn is_blocked(&self) -> bool { + matches!(self, PreToolUseAction::Block { .. }) + } + + pub fn block_reason(&self) -> Option<&str> { + match self { + PreToolUseAction::Block { reason } => Some(reason), + _ => None, + } + } + + pub fn updated_args(&self) -> Option<&serde_json::Value> { + match self { + PreToolUseAction::MutateInput { updated_args, .. } => Some(updated_args), + _ => None, + } + } + + pub fn additional_context(&self) -> Option<&str> { + match self { + PreToolUseAction::MutateInput { + additional_context, .. + } => additional_context.as_deref(), + _ => None, + } + } +} + +/// 向后兼容类型别名 — 旧代码用 HookAction::Continue / HookAction::Block 仍可编译 +pub type HookAction = PreToolUseAction; + +/// PostToolUse hook 返回的动作。 +/// 支持：保持输出不变、修改输出内容。 +#[derive(Debug, Clone)] +pub enum PostToolUseAction { + /// 保持原输出不变（默认） + Continue, + /// 修改输出内容 + MutateOutput { updated_content: String }, +} + +// ── Metrics Data ── + +/// 可查询的运行指标快照 +#[derive(Debug, Clone, Default)] +pub struct MetricsData { + pub tool_call_counts: HashMap, + pub total_steps: usize, + pub total_errors: usize, + pub session_id: Option, +} + +// ── Hook Trait ── + +/// Agent 生命周期 Hook trait。 +/// 所有方法都有默认空实现，只需覆写关心的 hook 点。 +#[async_trait] +pub trait AgentHook: Send + Sync { + /// Hook 名称（用于日志和调试） + fn name(&self) -> &str; + + // ── 原有 5 个生命周期事件 ── + + /// 会话创建/恢复时调用。 + async fn on_session_start(&self, _ctx: &SessionStartContext) {} + + /// 工具执行前调用。可返回 Continue/Block/MutateInput/PermissionRequired。 + async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction { + PreToolUseAction::Continue + } + + /// 工具执行后调用。可返回 Continue 或 MutateOutput。 + async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction { + PostToolUseAction::Continue + } + + /// 每个 ReAct step 完成后调用。 + async fn on_step_complete(&self, _ctx: &StepCompleteContext) {} + + /// 会话终止时调用。 + async fn on_session_stop(&self, _ctx: &SessionStopContext<'_>) {} + + // ── 新增 4 个生命周期事件（默认 no-op） ── + + /// 子代理启动时调用。 + async fn on_subagent_start(&self, _ctx: &SubagentStartContext) {} + + /// 子代理停止时调用。 + async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) {} + + /// 上下文压缩前调用。 + async fn on_pre_compact(&self, _ctx: &PreCompactContext) {} + + /// 上下文压缩后调用。 + async fn on_post_compact(&self, _ctx: &PostCompactContext) {} +} + +// ── Hook Registry ── + +/// PreToolUse 聚合结果 +#[derive(Debug, Clone)] +pub struct PreToolUseResult { + /// 最终动作（第一个 Block 获胜） + pub action: PreToolUseAction, + /// 累积的 additional_context（所有 MutateInput 的上下文拼接） + pub additional_context: Option, + /// 最终的工具参数（应用了最后一个 MutateInput 的修改） + pub final_args: serde_json::Value, +} + +/// PostToolUse 聚合结果 +#[derive(Debug, Clone)] +pub struct PostToolUseResult { + /// 最终输出内容（应用了最后一个 MutateOutput 的修改） + pub final_content: String, +} + +/// Hook 注册表，管理所有已注册的 hook 并按序调用 +pub struct HookRegistry { + hooks: Vec>, +} + +impl Default for HookRegistry { + fn default() -> Self { + Self::new() + } +} + +impl HookRegistry { + /// 创建空的注册表 + pub fn new() -> Self { + HookRegistry { hooks: Vec::new() } + } + + /// 创建包含所有内置 hooks 的注册表。 + /// + /// 参数： + /// - `db`: 数据库连接池（供 AuditLogHook 持久化） + /// - `cancelled_runs`: 取消状态集合（供 CancellationHook 检查） + /// - `metrics_data`: 可选的共享指标数据引用。提供时复用已有的 MetricsData， + /// 使得 AgentRuntime.get_metrics() 能查询到实际运行数据。 + pub fn with_builtins( + db: SqlitePool, + cancelled_runs: Arc>>, + metrics_data: Option>>, + ) -> Self { + let mut registry = Self::new(); + registry.add(Box::new(CancellationHook::new(cancelled_runs))); + // 如果提供了共享的 metrics_data，使用它；否则创建新的 + let metrics_hook = match metrics_data { + Some(data) => MetricsHook::from_arc(data), + None => MetricsHook::new(), + }; + registry.add(Box::new(metrics_hook)); + registry.add(Box::new(AuditLogHook::new(db))); + registry + } + + /// 注册一个 hook + pub fn add(&mut self, hook: Box) { + info!("[Hooks] 注册 hook: {}", hook.name()); + self.hooks.push(hook); + } + + /// 获取所有 hooks 的不可变引用 + pub fn all(&self) -> &[Box] { + &self.hooks + } + + // ── 便捷调用方法 ── + + /// 调用所有 on_session_start hooks + pub async fn run_on_session_start(&self, ctx: &SessionStartContext) { + for hook in &self.hooks { + hook.on_session_start(ctx).await; + } + } + + /// 调用所有 pre_tool_use hooks。 + /// + /// 返回聚合结果： + /// - 遇到第一个 Block 时短路，返回该 Block + /// - MutateInput 累积 additional_context 并更新 final_args + /// - PermissionRequired 记录但继续执行（暂时视为 Continue） + pub async fn run_pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseResult { + let mut accumulated_context = String::new(); + let mut final_args = ctx.tool_args.clone(); + let mut final_action = PreToolUseAction::Continue; + + for hook in &self.hooks { + let action = hook.pre_tool_use(ctx).await; + match &action { + PreToolUseAction::Block { reason } => { + warn!( + "[Hooks] {} 阻止了工具 {} 的执行: {}", + hook.name(), + ctx.tool_name, + reason + ); + return PreToolUseResult { + action, + additional_context: None, + final_args: ctx.tool_args.clone(), + }; + } + PreToolUseAction::MutateInput { + updated_args, + additional_context, + } => { + info!( + "[Hooks] {} 修改了工具 {} 的输入参数", + hook.name(), + ctx.tool_name + ); + final_args = updated_args.clone(); + if let Some(ctx_str) = additional_context { + if !accumulated_context.is_empty() { + accumulated_context.push('\n'); + } + accumulated_context.push_str(ctx_str); + } + } + PreToolUseAction::PermissionRequired { .. } => { + // 暂时记录但继续执行（Permission 系统在 Phase 2 中完善） + info!( + "[Hooks] {} 请求了工具 {} 的权限检查", + hook.name(), + ctx.tool_name + ); + } + PreToolUseAction::Continue => {} + } + final_action = action; + } + + let ctx_opt = if accumulated_context.is_empty() { + None + } else { + Some(accumulated_context) + }; + + PreToolUseResult { + action: final_action, + additional_context: ctx_opt, + final_args, + } + } + + /// 调用所有 post_tool_use hooks（全部执行，不会短路）。 + /// 返回聚合的最终输出内容。 + pub async fn run_post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseResult { + let mut final_content = ctx.output_content.clone(); + + for hook in &self.hooks { + let action = hook.post_tool_use(ctx).await; + match action { + PostToolUseAction::MutateOutput { updated_content } => { + info!( + "[Hooks] {} 修改了工具 {} 的输出", + hook.name(), + ctx.tool_name + ); + final_content = updated_content; + } + PostToolUseAction::Continue => {} + } + } + + PostToolUseResult { final_content } + } + + /// 调用所有 on_step_complete hooks + pub async fn run_on_step_complete(&self, ctx: &StepCompleteContext) { + for hook in &self.hooks { + hook.on_step_complete(ctx).await; + } + } + + /// 调用所有 on_session_stop hooks + pub async fn run_on_session_stop(&self, ctx: &SessionStopContext<'_>) { + for hook in &self.hooks { + hook.on_session_stop(ctx).await; + } + } + + /// 调用所有 on_subagent_start hooks + pub async fn run_on_subagent_start(&self, ctx: &SubagentStartContext) { + for hook in &self.hooks { + hook.on_subagent_start(ctx).await; + } + } + + /// 调用所有 on_subagent_stop hooks + pub async fn run_on_subagent_stop(&self, ctx: &SubagentStopContext) { + for hook in &self.hooks { + hook.on_subagent_stop(ctx).await; + } + } + + /// 调用所有 on_pre_compact hooks + pub async fn run_on_pre_compact(&self, ctx: &PreCompactContext) { + for hook in &self.hooks { + hook.on_pre_compact(ctx).await; + } + } + + /// 调用所有 on_post_compact hooks + pub async fn run_on_post_compact(&self, ctx: &PostCompactContext) { + for hook in &self.hooks { + hook.on_post_compact(ctx).await; + } + } +} + +// ── Built-in Hooks ── + +/// 取消检查 Hook — 在每次工具执行前检查用户是否中止了会话。 +pub struct CancellationHook { + cancelled_runs: Arc>>, +} + +impl CancellationHook { + pub fn new(cancelled_runs: Arc>>) -> Self { + CancellationHook { cancelled_runs } + } +} + +#[async_trait] +impl AgentHook for CancellationHook { + fn name(&self) -> &str { + "CancellationHook" + } + + async fn pre_tool_use(&self, ctx: &PreToolUseContext) -> PreToolUseAction { + if let Ok(runs) = self.cancelled_runs.lock() { + if runs.contains(&ctx.session_id) { + warn!( + "[CancellationHook] 会话 {} 已被用户取消，阻止工具 {} 执行", + ctx.session_id, ctx.tool_name + ); + return PreToolUseAction::Block { + reason: "用户已手动中止执行".to_string(), + }; + } + } + PreToolUseAction::Continue + } + + async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) { + // 清理取消状态 + if let Ok(mut runs) = self.cancelled_runs.lock() { + runs.remove(&ctx.session_id); + } + } +} + +/// 指标采集 Hook — 自动收集工具调用统计，支持快照查询 +pub struct MetricsHook { + data: Arc>, +} + +impl Default for MetricsHook { + fn default() -> Self { + Self::new() + } +} + +impl MetricsHook { + pub fn new() -> Self { + MetricsHook { + data: Arc::new(std::sync::Mutex::new(MetricsData::default())), + } + } + + /// 从已有的 Arc> 创建（共享数据引用） + pub fn from_arc(data: Arc>) -> Self { + MetricsHook { data } + } + + /// 返回当前指标快照（锁异常时返回 None） + pub fn snapshot(&self) -> Option { + self.data.lock().ok().map(|d| d.clone()) + } + + /// 获取 Arc 引用，供外部持有 + pub fn data_arc(&self) -> Arc> { + self.data.clone() + } +} + +#[async_trait] +impl AgentHook for MetricsHook { + fn name(&self) -> &str { + "MetricsHook" + } + + async fn on_session_start(&self, ctx: &SessionStartContext) { + if let Ok(mut data) = self.data.lock() { + data.session_id = Some(ctx.session_id.clone()); + } + } + + async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction { + info!( + "[Metrics] step={} tool={} is_error={} elapsed={}ms", + ctx.step, ctx.tool_name, ctx.is_error, ctx.elapsed_ms + ); + if let Ok(mut data) = self.data.lock() { + *data + .tool_call_counts + .entry(ctx.tool_name.clone()) + .or_insert(0) += 1; + data.total_steps = ctx.step; + if ctx.is_error { + data.total_errors += 1; + } + } + PostToolUseAction::Continue + } + + async fn on_step_complete(&self, ctx: &StepCompleteContext) { + if ctx.step.is_multiple_of(3) { + info!( + "[Metrics] step {}/{} | messages={} | tokens≈{}/{}", + ctx.step, ctx.max_steps, ctx.messages_count, ctx.estimated_tokens, ctx.token_limit + ); + } + } + + async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) { + info!( + "[Metrics] 会话 {} 结束: {} (total_steps={})", + ctx.session_id, + ctx.terminal.description(), + ctx.total_steps + ); + } +} + +/// 审计日志 Hook — 记录所有工具调用到 SQLite agent_audit_log 表 +pub struct AuditLogHook { + db: SqlitePool, +} + +impl AuditLogHook { + pub fn new(db: SqlitePool) -> Self { + AuditLogHook { db } + } +} + +#[async_trait] +impl AgentHook for AuditLogHook { + fn name(&self) -> &str { + "AuditLogHook" + } + + async fn post_tool_use(&self, ctx: &PostToolUseContext) -> PostToolUseAction { + let status = if ctx.is_error { "FAIL" } else { "OK" }; + let preview: String = ctx.output_content.chars().take(200).collect(); + + info!( + "[Audit] session={} step={} tool={} status={} elapsed={}ms", + ctx.session_id, ctx.step, ctx.tool_name, status, ctx.elapsed_ms + ); + + let db = self.db.clone(); + let session_id = ctx.session_id.clone(); + let agent_name = ctx.agent_name.clone(); + let tool_name = ctx.tool_name.clone(); + let step = ctx.step; + let elapsed = ctx.elapsed_ms; + let preview_clone = preview.clone(); + + // Fire-and-forget 写入，不阻塞主循环 + tokio::spawn(async move { + let _ = sqlx::query( + "INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview, agent_name) \ + VALUES (?, ?, ?, ?, ?, ?, ?)", + ) + .bind(&session_id) + .bind(step as i32) + .bind(&tool_name) + .bind(status) + .bind(elapsed as i32) + .bind(&preview_clone) + .bind(&agent_name) + .execute(&db) + .await; + }); + + PostToolUseAction::Continue + } + + async fn on_session_stop(&self, ctx: &SessionStopContext<'_>) { + info!( + "[Audit] 会话 {} 终止原因: {}", + ctx.session_id, + ctx.terminal.description() + ); + + let db = self.db.clone(); + let session_id = ctx.session_id.clone(); + let total_steps = ctx.total_steps; + + tokio::spawn(async move { + let _ = sqlx::query( + "INSERT INTO agent_audit_log (session_id, step, tool_name, status, elapsed_ms, output_preview) \ + VALUES (?, ?, 'session', 'SESSION_STOP', 0, ?)", + ) + .bind(&session_id) + .bind(total_steps as i32) + .bind(format!("会话终止，共 {} 步", total_steps)) + .execute(&db) + .await; + }); + } +} + +// ── Tests ── + +#[cfg(test)] +mod tests { + use super::*; + + struct TestHook { + name: String, + pre_called: std::sync::Mutex, + } + + impl TestHook { + fn new(name: &str) -> Self { + TestHook { + name: name.to_string(), + pre_called: std::sync::Mutex::new(false), + } + } + } + + #[async_trait] + impl AgentHook for TestHook { + fn name(&self) -> &str { + &self.name + } + + async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction { + *self.pre_called.lock().unwrap() = true; + PreToolUseAction::Continue + } + } + + #[tokio::test] + async fn test_hook_registry_runs_all_hooks() { + let mut registry = HookRegistry::new(); + let hook1 = TestHook::new("test1"); + let hook2 = TestHook::new("test2"); + registry.add(Box::new(hook1)); + registry.add(Box::new(hook2)); + + let ctx = PreToolUseContext { + session_id: "test".into(), + tool_name: "test_tool".into(), + tool_args: serde_json::json!({}), + step: 1, + }; + + let result = registry.run_pre_tool_use(&ctx).await; + assert!(!result.action.is_blocked()); + } + + #[tokio::test] + async fn test_blocking_hook_stops_chain() { + struct BlockingHook; + #[async_trait] + impl AgentHook for BlockingHook { + fn name(&self) -> &str { + "blocker" + } + async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction { + PreToolUseAction::Block { + reason: "test block".into(), + } + } + } + + let mut registry = HookRegistry::new(); + registry.add(Box::new(BlockingHook)); + + let ctx = PreToolUseContext { + session_id: "test".into(), + tool_name: "test_tool".into(), + tool_args: serde_json::json!({}), + step: 1, + }; + + let result = registry.run_pre_tool_use(&ctx).await; + assert!(result.action.is_blocked()); + assert_eq!(result.action.block_reason(), Some("test block")); + } + + #[tokio::test] + async fn test_mutate_input_accumulates_context() { + struct MutateHook; + #[async_trait] + impl AgentHook for MutateHook { + fn name(&self) -> &str { + "mutator" + } + async fn pre_tool_use(&self, _ctx: &PreToolUseContext) -> PreToolUseAction { + PreToolUseAction::MutateInput { + updated_args: serde_json::json!({"key": "modified"}), + additional_context: Some("injected context".to_string()), + } + } + } + + let mut registry = HookRegistry::new(); + registry.add(Box::new(MutateHook)); + + let ctx = PreToolUseContext { + session_id: "test".into(), + tool_name: "test_tool".into(), + tool_args: serde_json::json!({"key": "original"}), + step: 1, + }; + + let result = registry.run_pre_tool_use(&ctx).await; + assert_eq!(result.final_args, serde_json::json!({"key": "modified"})); + assert_eq!( + result.additional_context, + Some("injected context".to_string()) + ); + } + + #[tokio::test] + async fn test_post_tool_use_mutate_output() { + struct MutateOutputHook; + #[async_trait] + impl AgentHook for MutateOutputHook { + fn name(&self) -> &str { + "output_mutator" + } + async fn post_tool_use(&self, _ctx: &PostToolUseContext) -> PostToolUseAction { + PostToolUseAction::MutateOutput { + updated_content: "modified output".to_string(), + } + } + } + + let mut registry = HookRegistry::new(); + registry.add(Box::new(MutateOutputHook)); + + let ctx = PostToolUseContext { + session_id: "test".into(), + agent_name: "lead".into(), + tool_name: "test_tool".into(), + tool_args: serde_json::json!({}), + output_content: "original output".into(), + is_error: false, + step: 1, + elapsed_ms: 100, + }; + + let result = registry.run_post_tool_use(&ctx).await; + assert_eq!(result.final_content, "modified output"); + } + + #[tokio::test] + async fn test_cancellation_hook_blocks_when_cancelled() { + use std::collections::HashSet; + let mut cancelled = HashSet::new(); + cancelled.insert("test_session".to_string()); + let cancelled_runs = Arc::new(std::sync::Mutex::new(cancelled)); + + let hook = CancellationHook::new(cancelled_runs); + let ctx = PreToolUseContext { + session_id: "test_session".into(), + tool_name: "search_papers".into(), + tool_args: serde_json::json!({}), + step: 1, + }; + + let action = hook.pre_tool_use(&ctx).await; + assert!(action.is_blocked()); + } + + #[tokio::test] + async fn test_cancellation_hook_allows_when_not_cancelled() { + let cancelled_runs = Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())); + + let hook = CancellationHook::new(cancelled_runs); + let ctx = PreToolUseContext { + session_id: "test_session".into(), + tool_name: "search_papers".into(), + tool_args: serde_json::json!({}), + step: 1, + }; + + let action = hook.pre_tool_use(&ctx).await; + assert!(!action.is_blocked()); + } + + #[tokio::test] + async fn test_metrics_hook_accumulates_counts() { + let hook = MetricsHook::new(); + + let ctx = PostToolUseContext { + session_id: "test".into(), + agent_name: "lead".into(), + tool_name: "search_papers".into(), + tool_args: serde_json::json!({}), + output_content: "result".into(), + is_error: false, + step: 1, + elapsed_ms: 100, + }; + hook.post_tool_use(&ctx).await; + + let ctx2 = PostToolUseContext { + session_id: "test".into(), + agent_name: "lead".into(), + tool_name: "search_papers".into(), + tool_args: serde_json::json!({}), + output_content: "result2".into(), + is_error: false, + step: 2, + elapsed_ms: 200, + }; + hook.post_tool_use(&ctx2).await; + + let snapshot = hook.snapshot().expect("snapshot should succeed in test"); + assert_eq!(snapshot.tool_call_counts.get("search_papers"), Some(&2)); + assert_eq!(snapshot.total_steps, 2); + } + + #[tokio::test] + async fn test_session_start_hook_called() { + struct StartTrackingHook { + started: std::sync::Mutex>, + } + + #[async_trait] + impl AgentHook for StartTrackingHook { + fn name(&self) -> &str { + "start_tracker" + } + async fn on_session_start(&self, ctx: &SessionStartContext) { + self.started.lock().unwrap().push(ctx.session_id.clone()); + } + } + + let hook = StartTrackingHook { + started: std::sync::Mutex::new(Vec::new()), + }; + + let mut registry = HookRegistry::new(); + registry.add(Box::new(hook)); + + let ctx = SessionStartContext { + session_id: "test_sid".into(), + turn_index: 1, + is_resume: false, + }; + registry.run_on_session_start(&ctx).await; + } + + #[tokio::test] + async fn test_new_lifecycle_events_called() { + struct LifecycleTracker { + subagent_start: std::sync::Mutex, + subagent_stop: std::sync::Mutex, + pre_compact: std::sync::Mutex, + post_compact: std::sync::Mutex, + } + + #[async_trait] + impl AgentHook for LifecycleTracker { + fn name(&self) -> &str { + "lifecycle_tracker" + } + async fn on_subagent_start(&self, _ctx: &SubagentStartContext) { + *self.subagent_start.lock().unwrap() = true; + } + async fn on_subagent_stop(&self, _ctx: &SubagentStopContext) { + *self.subagent_stop.lock().unwrap() = true; + } + async fn on_pre_compact(&self, _ctx: &PreCompactContext) { + *self.pre_compact.lock().unwrap() = true; + } + async fn on_post_compact(&self, _ctx: &PostCompactContext) { + *self.post_compact.lock().unwrap() = true; + } + } + + let tracker = LifecycleTracker { + subagent_start: std::sync::Mutex::new(false), + subagent_stop: std::sync::Mutex::new(false), + pre_compact: std::sync::Mutex::new(false), + post_compact: std::sync::Mutex::new(false), + }; + + let mut registry = HookRegistry::new(); + registry.add(Box::new(tracker)); + + registry + .run_on_subagent_start(&SubagentStartContext { + parent_session_id: "s1".into(), + subagent_name: "sub".into(), + prompt: "test".into(), + }) + .await; + registry + .run_on_subagent_stop(&SubagentStopContext { + parent_session_id: "s1".into(), + subagent_name: "sub".into(), + result_summary: "done".into(), + steps: 3, + is_error: false, + }) + .await; + registry + .run_on_pre_compact(&PreCompactContext { + session_id: "s1".into(), + message_count: 50, + estimated_tokens: 10000, + }) + .await; + registry + .run_on_post_compact(&PostCompactContext { + session_id: "s1".into(), + new_message_count: 10, + compression_method: "micro".into(), + }) + .await; + + // If no panic, all hooks were called successfully + } +} diff --git a/src/agent/memory/age.rs b/src/agent/memory/age.rs new file mode 100644 index 0000000..6a0241f --- /dev/null +++ b/src/agent/memory/age.rs @@ -0,0 +1,132 @@ +// src/agent/memory/age.rs +// +// 记忆时效性追踪 — 参考 Claude Code memdir/memoryAge.ts。 +// +// LLM 不擅长日期计算，"2026-01-15" 不会触发过时判断， +// 但 "47 天前" 会。本模块提供人类可读的时效标签， +// 注入到 system prompt 中以引导模型在使用记忆中之前核实。 + +use std::time::{SystemTime, UNIX_EPOCH}; + +/// 返回当前 Unix 时间戳（秒） +pub fn now_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +/// 距离给定 Unix 时间戳的天数（向下取整）。 +/// 0 = 今天, 1 = 昨天, 2+ = 更早。 +/// 负数输入（未来时间/时钟偏差）截断为 0。 +pub fn memory_age_days(mtime_secs: u64) -> u64 { + let now = now_secs(); + if mtime_secs >= now { + return 0; + } + (now - mtime_secs) / 86_400 +} + +/// 人类可读的时效标签。 +pub fn memory_age_label(mtime_secs: u64) -> String { + let days = memory_age_days(mtime_secs); + match days { + 0 => "今天".to_string(), + 1 => "昨天".to_string(), + n => format!("{} 天前", n), + } +} + +/// 返回时效警告文本，如果记忆超过 1 天则返回 Some。 +/// 新鲜记忆（今天/昨天）返回 None — 此时警告只是噪音。 +pub fn memory_freshness_text(mtime_secs: u64) -> Option { + let days = memory_age_days(mtime_secs); + if days <= 1 { + return None; + } + Some(format!( + "此记忆已有 {} 天。记忆是时间点快照，不是实时状态 — \ + 关于代码行为或文件:行号的声明可能已过时。\ + 请在断言为事实前与当前代码进行核对。", + days + )) +} + +/// 包裹在 标签中的时效注释。 +/// 对于 ≤ 1 天的记忆返回空字符串。 +pub fn memory_freshness_note(mtime_secs: u64) -> String { + match memory_freshness_text(mtime_secs) { + Some(text) => format!("{}", text), + None => String::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_age_days_today() { + let now = now_secs(); + assert_eq!(memory_age_days(now), 0); + assert_eq!(memory_age_days(now - 100), 0); // 100 秒前仍是今天 + } + + #[test] + fn test_age_days_yesterday() { + let now = now_secs(); + assert_eq!(memory_age_days(now - 86_400), 1); + assert_eq!(memory_age_days(now - 86_400 - 100), 1); + } + + #[test] + fn test_age_days_older() { + let now = now_secs(); + assert_eq!(memory_age_days(now - 86_400 * 3), 3); + assert_eq!(memory_age_days(now - 86_400 * 47), 47); + } + + #[test] + fn test_age_days_clamps_future_to_zero() { + let future = now_secs() + 86_400 * 10; + assert_eq!(memory_age_days(future), 0); + } + + #[test] + fn test_age_label() { + let now = now_secs(); + assert_eq!(memory_age_label(now), "今天"); + assert_eq!(memory_age_label(now - 86_400), "昨天"); + assert_eq!(memory_age_label(now - 86_400 * 5), "5 天前"); + } + + #[test] + fn test_freshness_text_none_for_fresh() { + let now = now_secs(); + assert!(memory_freshness_text(now).is_none()); // 今天 + assert!(memory_freshness_text(now - 86_400).is_none()); // 昨天 + } + + #[test] + fn test_freshness_text_some_for_old() { + let now = now_secs(); + let text = memory_freshness_text(now - 86_400 * 2); + assert!(text.is_some()); + assert!(text.unwrap().contains("2 天")); + } + + #[test] + fn test_freshness_note_contains_tags() { + let now = now_secs(); + let note = memory_freshness_note(now - 86_400 * 3); + assert!(note.contains("")); + assert!(note.contains("")); + } + + #[test] + fn test_freshness_note_empty_for_fresh() { + let now = now_secs(); + assert_eq!(memory_freshness_note(now), ""); + assert_eq!(memory_freshness_note(now - 86_400), ""); + } +} diff --git a/src/agent/memory/decay.rs b/src/agent/memory/decay.rs new file mode 100644 index 0000000..b21e0e7 --- /dev/null +++ b/src/agent/memory/decay.rs @@ -0,0 +1,143 @@ +// src/agent/memory/decay.rs +// +// 指数时间衰减评分 — 参考 Martian-Engineering/agent-memory。 +// +// 核心理念：LLM 判断语义相关性，数学衰减提供时序排序。 +// "semantic decay adds LLM judgment, recency scoring adds temporal ordering" +// — 两者互补，4天前的关键偏好可胜过1天前的临时备注。 +// +// 公式: score = e^(-λ × days_old) +// λ = ln(2) / half_life_days +// +// 30天半衰期下: 0天=1.0, 15天≈0.707, 30天=0.5, 60天=0.25, 90天≈0.125 + +/// 默认半衰期（天） +pub const DEFAULT_HALF_LIFE_DAYS: f64 = 30.0; + +/// 计算指数时间衰减评分。 +/// +/// score = e^(-λ × days_old)，λ = ln(2) / half_life +pub fn decay_score(mtime_secs: u64, half_life_days: f64) -> f64 { + let _now = super::age::now_secs(); + let days = super::age::memory_age_days(mtime_secs) as f64; + let lambda = std::f64::consts::LN_2 / half_life_days; + (-lambda * days).exp() +} + +/// 便捷函数：使用默认 30 天半衰期 +pub fn decay_score_default(mtime_secs: u64) -> f64 { + decay_score(mtime_secs, DEFAULT_HALF_LIFE_DAYS) +} + +/// Hebbian 激活分级 — 参考 OpenClaw Hot/Warm/Cool 模型。 +/// +/// 结合访问频率使用时更强大； +/// 当前基于纯 recency 实现（无需额外追踪基础设施）。 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ActivationTier { + /// ≤7 天 — 高频访问，完整权重 + Hot, + /// 8-30 天 — 中频，衰减中 + Warm, + /// >30 天 — 低频，可能归档 + Cool, +} + +/// 按年龄分级 +pub fn activation_tier(mtime_secs: u64) -> ActivationTier { + let days = super::age::memory_age_days(mtime_secs); + if days <= 7 { + ActivationTier::Hot + } else if days <= 30 { + ActivationTier::Warm + } else { + ActivationTier::Cool + } +} + +/// 激活等级的排序权重乘数 +pub fn activation_multiplier(tier: ActivationTier) -> f64 { + match tier { + ActivationTier::Hot => 1.0, + ActivationTier::Warm => 0.7, + ActivationTier::Cool => 0.3, + } +} + +/// 用于日志/显示的可读标签 +pub fn activation_label(tier: ActivationTier) -> &'static str { + match tier { + ActivationTier::Hot => "活跃", + ActivationTier::Warm => "温", + ActivationTier::Cool => "冷", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::memory::age; + + #[test] + fn test_decay_score_fresh() { + let now = age::now_secs(); + let score = decay_score(now, 30.0); + assert!((score - 1.0).abs() < 0.01, "今天应为 1.0，实际 {}", score); + } + + #[test] + fn test_decay_score_half_life() { + let now = age::now_secs(); + let thirty_days_ago = now - 86_400 * 30; + let score = decay_score(thirty_days_ago, 30.0); + assert!((score - 0.5).abs() < 0.01, "30天应为 0.5，实际 {}", score); + } + + #[test] + fn test_decay_score_60_days() { + let now = age::now_secs(); + let sixty_days_ago = now - 86_400 * 60; + let score = decay_score(sixty_days_ago, 30.0); + assert!(score < 0.26 && score > 0.24, "60天应≈0.25，实际 {}", score); + } + + #[test] + fn test_activation_tier_hot() { + let now = age::now_secs(); + assert_eq!(activation_tier(now), ActivationTier::Hot); + assert_eq!(activation_tier(now - 86_400 * 7), ActivationTier::Hot); + } + + #[test] + fn test_activation_tier_warm() { + let now = age::now_secs(); + assert_eq!(activation_tier(now - 86_400 * 8), ActivationTier::Warm); + assert_eq!(activation_tier(now - 86_400 * 30), ActivationTier::Warm); + } + + #[test] + fn test_activation_tier_cool() { + let now = age::now_secs(); + assert_eq!(activation_tier(now - 86_400 * 31), ActivationTier::Cool); + } + + #[test] + fn test_activation_multiplier_ranges() { + assert_eq!(activation_multiplier(ActivationTier::Hot), 1.0); + assert_eq!(activation_multiplier(ActivationTier::Warm), 0.7); + assert_eq!(activation_multiplier(ActivationTier::Cool), 0.3); + } + + #[test] + fn test_decay_score_different_half_life() { + let now = age::now_secs(); + let ago = now - 86_400 * 15; + // 15天半衰期下，15天后应为0.5 + let score = decay_score(ago, 15.0); + assert!( + (score - 0.5).abs() < 0.01, + "15天半衰期15天后应为0.5，实际{}", + score + ); + } +} diff --git a/src/agent/memory/dedup.rs b/src/agent/memory/dedup.rs new file mode 100644 index 0000000..befdaf7 --- /dev/null +++ b/src/agent/memory/dedup.rs @@ -0,0 +1,406 @@ +// src/agent/memory/dedup.rs +// +// 记忆去重支持 — 参考 Claude Code memdir 提示词中的去重规则。 +// +// 在保存新记忆前检查是否有可更新的现有条目， +// 构建现有记忆的 manifest 供 LLM 参考以减少重复写入。 + +use std::fs; +use std::path::Path; + +use super::types::MemoryEntry; + +/// 构建现有记忆的清单预览（供 LLM 了解已存在的内容）。 +/// 在 save_memory 成功后注入到工具输出中。 +pub fn build_manifest_preview(entries: &[MemoryEntry]) -> String { + if entries.is_empty() { + return "当前无其他记忆条目。".to_string(); + } + + let mut lines = vec!["当前记忆清单：".to_string()]; + for entry in entries { + let type_label = match entry.memory_type { + super::types::MemoryType::User => "[偏好]", + super::types::MemoryType::Feedback => "[反馈]", + super::types::MemoryType::Project => "[项目]", + super::types::MemoryType::Reference => "[参考]", + }; + lines.push(format!( + "- {} `{}` {}: {}", + type_label, entry.slug, entry.name, entry.description + )); + } + lines.join("\n") +} + +/// 检查 slug 是否在磁盘上已存在。 +pub fn slug_exists(memory_dir: &Path, slug: &str) -> bool { + let file_path = memory_dir.join(format!("{}.md", slug)); + file_path.exists() +} + +/// 列出所有现有 slug（从磁盘直接读取，避免依赖 MemoryManager 状态）。 +pub fn list_existing_slugs(memory_dir: &Path) -> Vec { + let mut slugs = Vec::new(); + if let Ok(entries) = fs::read_dir(memory_dir) { + for entry in entries.flatten() { + let path = entry.path(); + if path.is_dir() { + continue; + } + if let Some(file_name) = path.file_name().and_then(|n| n.to_str()) { + if file_name == "MEMORY.md" || !file_name.ends_with(".md") { + continue; + } + if let Some(slug) = file_name.strip_suffix(".md") { + slugs.push(slug.to_string()); + } + } + } + } + slugs.sort(); + slugs +} + +// ── Jaccard 相似度去重 ── + +/// 计算两个字符串的 Jaccard 相似度（基于字符级 bigram）。 +/// +/// 使用 bigram 而非词级分词以正确处理中文（不依赖分词器）。 +/// 值域 [0.0, 1.0]，阈值 ≥0.70 通常视为重复。 +/// +/// 参考 Martian-Engineering/agent-memory 的 70% Jaccard 门控。 +pub fn jaccard_similarity(a: &str, b: &str) -> f64 { + let bigrams_a: std::collections::HashSet = bigrams(a); + let bigrams_b: std::collections::HashSet = bigrams(b); + + if bigrams_a.is_empty() && bigrams_b.is_empty() { + return 1.0; // 两个空字符串完全相同 + } + + let intersection = bigrams_a.intersection(&bigrams_b).count(); + let union = bigrams_a.union(&bigrams_b).count(); + + if union == 0 { + return 0.0; + } + intersection as f64 / union as f64 +} + +/// 提取字符串的字符级 bigram 集合。 +fn bigrams(s: &str) -> std::collections::HashSet { + let chars: Vec = s.chars().collect(); + let mut set = std::collections::HashSet::new(); + if chars.len() < 2 { + // 单字符内容：将单字符本身作为 bigram + if !chars.is_empty() { + set.insert(chars[0].to_string()); + } + return set; + } + for window in chars.windows(2) { + set.insert(format!("{}{}", window[0], window[1])); + } + set +} + +/// 检查新内容与现有记忆是否高度重复。 +/// 返回重复的 slug，或在无重复时返回 None。 +pub fn find_duplicate_by_content( + new_content: &str, + existing_entries: &[MemoryEntry], + threshold: f64, +) -> Option { + for entry in existing_entries { + if !entry.status.is_active() { + continue; + } + let sim = jaccard_similarity(new_content, &entry.content); + if sim >= threshold { + return Some(entry.slug.clone()); + } + } + None +} + +// ── 写入时内容质量门控 ── + +/// 内容质量检查结果 +#[derive(Debug, PartialEq, Eq)] +pub enum QualityCheck { + /// 通过质量检查 + Accept, + /// 太短：有效字符不足 + TooShort(usize), + /// 瞬时状态描述 + TransientState, + /// 模糊语言 + VagueLanguage(String), + /// 纯代码片段 + CodePattern, +} + +/// 瞬时状态关键词（中文 + 英文） +const TRANSIENT_PATTERNS: &[&str] = &[ + "正在做", + "正在写", + "正在调试", + "正在看", + "准备做", + "is working on", + "currently", + "right now", + "at the moment", +]; + +/// 模糊语言关键词 +const VAGUE_PATTERNS: &[(&str, &str)] = &[ + ("maybe", "可能"), + ("probably", "大概"), + ("perhaps", "也许"), + ("might be", "或许"), + ("似乎", "似乎"), + ("好像", "好像"), +]; + +/// 代码模式检测（纯代码片段不应作为记忆） +const CODE_PATTERNS: &[&str] = &[ + "fn ", + "impl ", + "struct ", + "pub fn", + "use crate", + "function ", + "const ", + "let mut", + "&mut", + "import {", + "from \"", + "export ", +]; + +/// 最小内容长度（有效字符）。 +/// 中文信息密度高，10 字即可表达完整语义。 +const MIN_CONTENT_CHARS: usize = 10; + +/// 检查内容质量（写入时门控）。 +/// +/// 仅返回警告 — 不强制拒绝，由 LLM 最终决定。 +/// 参考 OpenClaw claw-mem 写入时门控 + agent-memory 写入规则。 +pub fn check_content_quality(content: &str) -> QualityCheck { + let trimmed = content.trim(); + + // 1. 长度检查 + let char_count = trimmed.chars().count(); + if char_count < MIN_CONTENT_CHARS { + return QualityCheck::TooShort(char_count); + } + + // 2. 瞬时状态检查 + let lower = trimmed.to_lowercase(); + for pattern in TRANSIENT_PATTERNS { + if lower.contains(pattern) { + return QualityCheck::TransientState; + } + } + + // 3. 模糊语言检查 + for (en, zh) in VAGUE_PATTERNS { + if lower.contains(en) || lower.contains(zh) { + return QualityCheck::VagueLanguage(if lower.contains(en) { + en.to_string() + } else { + zh.to_string() + }); + } + } + + // 4. 代码模式检查 + for pattern in CODE_PATTERNS { + if trimmed.contains(pattern) { + return QualityCheck::CodePattern; + } + } + + QualityCheck::Accept +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + fn make_entry(slug: &str, name: &str, desc: &str) -> MemoryEntry { + MemoryEntry { + slug: slug.to_string(), + name: name.to_string(), + description: desc.to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: 1000, + content: desc.to_string(), + path: PathBuf::from(slug), + status: super::super::types::MemoryStatus::Active, + } + } + + #[test] + fn test_manifest_preview_empty() { + let preview = build_manifest_preview(&[]); + assert!(preview.contains("无其他记忆条目")); + } + + #[test] + fn test_manifest_preview_with_entries() { + let entries = vec![ + make_entry("user-role", "用户角色", "数据科学家"), + make_entry("feedback-tests", "测试反馈", "不要 mock 数据库"), + ]; + let preview = build_manifest_preview(&entries); + assert!(preview.contains("user-role")); + assert!(preview.contains("feedback-tests")); + assert!(preview.contains("数据科学家")); + assert!(preview.contains("不要 mock 数据库")); + } + + #[test] + fn test_slug_exists_true() { + let dir = std::env::temp_dir().join("astro_memory_test_dedup"); + fs::create_dir_all(&dir).unwrap(); + fs::write(dir.join("existing.md"), "test").unwrap(); + assert!(slug_exists(&dir, "existing")); + fs::remove_dir_all(&dir).unwrap(); + } + + #[test] + fn test_slug_exists_false() { + let dir = std::env::temp_dir().join("astro_memory_test_dedup_nonexist"); + assert!(!slug_exists(&dir, "nonexistent")); + } + + #[test] + fn test_list_existing_slugs() { + let dir = std::env::temp_dir().join("astro_memory_test_list_slugs"); + fs::create_dir_all(&dir).unwrap(); + fs::write(dir.join("alpha.md"), "a").unwrap(); + fs::write(dir.join("beta.md"), "b").unwrap(); + fs::write(dir.join("MEMORY.md"), "index").unwrap(); + + let slugs = list_existing_slugs(&dir); + assert!(slugs.contains(&"alpha".to_string())); + assert!(slugs.contains(&"beta".to_string())); + assert!(!slugs.contains(&"MEMORY".to_string())); + + fs::remove_dir_all(&dir).unwrap(); + } + + // ── Jaccard 相似度测试 ── + + #[test] + fn test_jaccard_identical() { + let sim = jaccard_similarity("hello world", "hello world"); + assert!((sim - 1.0).abs() < 0.01, "完全相同应为 1.0，实际 {}", sim); + } + + #[test] + fn test_jaccard_completely_different() { + let sim = jaccard_similarity("hello world", "abc xyz"); + assert!(sim < 0.3, "完全不同应较低，实际 {}", sim); + } + + #[test] + fn test_jaccard_high_overlap() { + let sim = jaccard_similarity( + "用户偏好使用 Rust 开发后端服务", + "用户偏好使用 Rust 开发后端", + ); + assert!(sim > 0.5, "高重叠应 >0.5，实际 {}", sim); + } + + #[test] + fn test_jaccard_chinese_bigram() { + let sim = jaccard_similarity("天体物理学研究", "天体物理研究"); + assert!(sim > 0.5, "中文 bigram 应能正确匹配，实际 {}", sim); + } + + // ── 内容质量检查测试 ── + + #[test] + fn test_quality_too_short() { + assert_eq!(check_content_quality("太短"), QualityCheck::TooShort(2)); + // 刚好 10 个中文字符（可通过最低长度） + let ten = "一二三四五六七八九十"; + assert_eq!(char_count(ten), 10); + assert_eq!(check_content_quality(ten), QualityCheck::Accept); + } + + fn char_count(s: &str) -> usize { + s.chars().count() + } + + #[test] + fn test_quality_transient_state() { + assert_eq!( + check_content_quality("用户正在调试登录模块的问题"), + QualityCheck::TransientState + ); + } + + #[test] + fn test_quality_vague_language() { + assert_eq!( + check_content_quality("可能需要在后续版本中优化"), + QualityCheck::VagueLanguage("可能".to_string()) + ); + } + + #[test] + fn test_quality_code_pattern() { + assert_eq!( + check_content_quality("fn main() { println!(\"hello\"); }"), + QualityCheck::CodePattern + ); + } + + #[test] + fn test_quality_accept_good_content() { + assert_eq!( + check_content_quality( + "用户是天体物理学家，主要研究星系演化。偏好使用 Kim 的径向速度拟合方法。" + ), + QualityCheck::Accept + ); + } + + #[test] + fn test_find_duplicate_by_content_detects_high_overlap() { + let base = "用户偏好使用 Rust 开发后端服务"; + let entries = vec![make_entry("memory-a", "A", base)]; + let dup = find_duplicate_by_content("用户偏好使用 Rust 开发后端系统", &entries, 0.40); + assert!(dup.is_some(), "高重叠内容应检测为重复"); + } + + #[test] + fn test_find_duplicate_rejects_low_overlap() { + let entries = vec![make_entry("a", "A", "用户偏好使用 Rust 开发后端")]; + let dup = find_duplicate_by_content("天体物理学中星系演化研究的最新进展", &entries, 0.40); + assert!(dup.is_none(), "低重叠内容不应检测为重复"); + } + + #[test] + fn test_find_duplicate_skips_historical() { + let mut entries = vec![MemoryEntry { + slug: "historical-one".to_string(), + name: "历史记忆".to_string(), + description: "已过时".to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: 1000, + content: "用户偏好使用 Rust 开发后端".to_string(), + path: PathBuf::from("historical-one.md"), + status: super::super::types::MemoryStatus::Historical { + superseded_by: Some("new-one".to_string()), + }, + }]; + // historical 应被跳过，不匹配 + assert!(find_duplicate_by_content("用户偏好使用 Rust 开发后端", &entries, 0.6,).is_none()); + } +} diff --git a/src/agent/memory/extraction.rs b/src/agent/memory/extraction.rs new file mode 100644 index 0000000..d8c7904 --- /dev/null +++ b/src/agent/memory/extraction.rs @@ -0,0 +1,236 @@ +// src/agent/memory/extraction.rs +// +// 自动记忆提取 — 参考 Claude Code services/extractMemories/。 +// +// 在每次会话结束时，使用受限子代理分析对话内容并自动提取 +// 值得保留的记忆条目。提取是 fire-and-forget 的，不影响主会话关闭。 +// +// 设计决策： +// - 默认关闭（EXTRACT_MEMORY_ENABLED=false），避免意外的 LLM 费用 +// - 如果主代理已通过 save_memory 工具写入，则跳过提取 +// - 使用节流避免每轮都提取 + +use std::sync::Arc; +use tokio::sync::Mutex; +use tracing::{info, warn}; + +use crate::agent::memory::dedup; +use crate::agent::memory::MemoryManager; +use crate::agent::subagent::SubAgentRunner; +use crate::agent::tools::memory::SaveMemoryTool; +use crate::agent::tools::{GlobFilesTool, GrepFilesTool, ReadFileTool, ToolRegistry}; +use crate::api::AppState; + +/// 自动提取配置（从环境变量加载） +pub struct ExtractionConfig { + /// 是否启用自动提取 + pub enabled: bool, + /// 最小提取间隔（轮次） + pub throttle_turns: usize, + /// 子代理最大 ReAct 步数 + pub max_steps: usize, +} + +impl Default for ExtractionConfig { + fn default() -> Self { + ExtractionConfig { + enabled: false, + throttle_turns: 3, + max_steps: 3, + } + } +} + +impl ExtractionConfig { + /// 从环境变量加载配置。 + /// - EXTRACT_MEMORY_ENABLED=true/false（默认 false） + /// - EXTRACT_MEMORY_THROTTLE_TURNS（默认 3） + /// - EXTRACT_MEMORY_MAX_STEPS（默认 3） + pub fn from_env() -> Self { + let enabled = std::env::var("EXTRACT_MEMORY_ENABLED") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(false); + + let throttle_turns = std::env::var("EXTRACT_MEMORY_THROTTLE_TURNS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(3); + + let max_steps = std::env::var("EXTRACT_MEMORY_MAX_STEPS") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(3); + + ExtractionConfig { + enabled, + throttle_turns, + max_steps, + } + } +} + +/// 提取追踪器 — 存储在 MemoryManager 中以跨轮次追踪状态。 +#[derive(Debug, Default)] +pub struct ExtractionTracker { + /// 自上次提取以来的轮次数 + pub turns_since_last_extraction: usize, + /// 主代理在本会话中是否已写入记忆 + pub main_agent_saved_this_session: bool, +} + +/// 提取系统提示词 +const EXTRACTION_SYSTEM_PROMPT: &str = "\ +你是一个记忆提取助手。分析最近的对话，提取值得持久化保存的信息。 + +## 记忆类型 +- **user**: 用户角色、偏好、知识背景 +- **feedback**: 用户给出的修正或确认的方法论（包含 Why 和 How to apply） +- **project**: 项目上下文、目标、约束（不可从代码推导的部分） +- **reference**: 外部资源指针（URL、仪表盘、工单系统） + +## 不应保存 +- 代码模式、架构详情（可从项目状态推导） +- Git 历史、调试方案 +- 已在 CLAUDE.md 中的内容 +- 临时任务状态"; + +/// 构建提取子代理的用户提示词。 +fn build_extraction_prompt(new_message_count: usize, existing_manifest: &str) -> String { + let manifest_section = if existing_manifest.is_empty() { + "当前无记忆条目。".to_string() + } else { + format!( + "## 现有记忆清单\n\n{}\n\n检查此清单 — 更新现有文件而非创建重复项。", + existing_manifest + ) + }; + + format!( + "分析最近约 {} 条消息，提取值得持久化保存的信息。\n\n{}\n\n\ + ## 操作指南\n\ + 1. 先读取需要更新的现有记忆文件（如果有）\n\ + 2. 然后使用 save_memory 工具保存新记忆或更新现有记忆\n\ + 3. 只保存非显而易见的、在后续对话中仍有用的信息\n\ + 4. 不要浪费时间验证或搜索其他内容 — 仅基于对话内容", + new_message_count, manifest_section + ) +} + +/// 构建受限工具注册表（只读 + save_memory）。 +fn build_extraction_tool_registry(memory_manager: Arc>) -> ToolRegistry { + let mut registry = ToolRegistry::empty(); + + // 只读工具 + registry.add_tool(Box::new(ReadFileTool)); + registry.add_tool(Box::new(GrepFilesTool)); + registry.add_tool(Box::new(GlobFilesTool)); + + // 写入仅限记忆目录 + registry.add_tool(Box::new(SaveMemoryTool::new(memory_manager))); + + registry +} + +/// 运行自动记忆提取（fire-and-forget，调用者应通过 tokio::spawn 运行）。 +/// +/// 永不 panic — 所有错误都只记录日志。 +pub async fn run_extraction( + app_state: Arc, + session_id: String, + memory_manager: Arc>, + config: ExtractionConfig, +) { + if !config.enabled { + return; + } + + // 检查节流和主代理写入 + { + let mut mgr = memory_manager.lock().await; + mgr.extraction_tracker.turns_since_last_extraction += 1; + + if mgr.extraction_tracker.turns_since_last_extraction < config.throttle_turns { + return; + } + if mgr.extraction_tracker.main_agent_saved_this_session { + info!("[Extraction] 跳过 — 主代理已通过 save_memory 写入"); + mgr.extraction_tracker.main_agent_saved_this_session = false; + mgr.extraction_tracker.turns_since_last_extraction = 0; + return; + } + mgr.extraction_tracker.turns_since_last_extraction = 0; + } + + info!("[Extraction] 开始会话 {} 的自动记忆提取", session_id); + + // 获取现有记忆清单 + let existing_manifest = { + let mgr = memory_manager.lock().await; + dedup::build_manifest_preview(mgr.entries()) + }; + + // 构建受限工具集 + let tool_registry = build_extraction_tool_registry(memory_manager.clone()); + + // 构建子代理 + let runner = SubAgentRunner::new_with_registry(app_state, tool_registry); + + // 构建提示词 + let prompt = build_extraction_prompt(20, &existing_manifest); + + // 运行子代理（同步等待，但调用者通过 tokio::spawn 异步化） + let result = runner + .run(EXTRACTION_SYSTEM_PROMPT, &prompt, config.max_steps) + .await; + + if result.is_error { + warn!( + "[Extraction] 子代理返回错误: {}", + result.content.chars().take(200).collect::() + ); + } else { + info!( + "[Extraction] 提取完成: {}", + result.content.chars().take(150).collect::() + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_extraction_config_default() { + let config = ExtractionConfig::default(); + assert!(!config.enabled); + assert_eq!(config.throttle_turns, 3); + assert_eq!(config.max_steps, 3); + } + + #[test] + fn test_extraction_config_from_env_disabled() { + // 未设置环境变量时应为默认（禁用） + std::env::remove_var("EXTRACT_MEMORY_ENABLED"); + let config = ExtractionConfig::from_env(); + assert!(!config.enabled); + } + + #[test] + fn test_extraction_tracker_default() { + let tracker = ExtractionTracker::default(); + assert_eq!(tracker.turns_since_last_extraction, 0); + assert!(!tracker.main_agent_saved_this_session); + } + + #[test] + fn test_build_extraction_prompt() { + let prompt = build_extraction_prompt(10, ""); + assert!(prompt.contains("10 条消息")); + assert!(prompt.contains("当前无记忆条目")); + + let prompt_with_manifest = build_extraction_prompt(5, "- user-role: 用户角色"); + assert!(prompt_with_manifest.contains("现有记忆清单")); + } +} diff --git a/src/agent/memory/guardrails.rs b/src/agent/memory/guardrails.rs new file mode 100644 index 0000000..c0a5eed --- /dev/null +++ b/src/agent/memory/guardrails.rs @@ -0,0 +1,83 @@ +// src/agent/memory/guardrails.rs +// +// 记忆保存/使用护栏 — 参考 Claude Code memoryTypes.ts。 +// +// 提供两个维度的防护： +// 1. WHAT_NOT_TO_SAVE — 不应保存为记忆的内容（即使被要求） +// 2. TRUST_BUT_VERIFY — 从记忆中推荐前先核实的提示词 +// +// 这些提示词经过 Claude Code 评估验证（memory-prompt-iteration.eval.ts）： +// - 排除规则明确告知模型 "即使用户要求保存" 也不应保存噪音内容 +// - 验证提示词需要放在决策点（系统提示词中记忆段落之后）， +// 不能埋在通用指南中，否则模型会忽略 + +/// 不应保存为记忆的内容。 +/// 即使用户明确要求保存，这些规则也适用。 +pub const WHAT_NOT_TO_SAVE: &str = "\ +不应保存为记忆的内容： +- 代码模式、惯例、架构详情、文件路径 — 可从当前项目状态推导 +- Git 历史、最近修改、谁改了什么 — `git log` / `git blame` 是权威来源 +- 调试方案或错误临时解决方案 — 修复在代码中，commit message 有上下文 +- 已在 CLAUDE.md 或项目文档中的内容 +- 临时任务细节：进行中工作、当前对话上下文 + +即使用户明确要求保存以上内容，请询问其中哪些部分是*意外的*或*非常规的* — 那些才是值得保存的。"; + +/// "从记忆中推荐前先核实" 提示词。 +/// 必须放在决策点（记忆段落后），不能在通用指南中。 +/// Claude Code 评估：放在 "When to access memories" 下时 0/3， +/// 放在独立段落标题下时 3/3 — 标题权重影响模型行为。 +pub const VERIFY_BEFORE_RECOMMENDING: &str = "\ +## 从记忆中推荐前先核实 + +记忆中提到特定函数、文件或标志是一种声明，声称它们*在记忆写入时*存在。 +但函数可能已被重命名、移除或从未合并。在据此推荐前： + +- 如果记忆提到了文件路径：确认该文件存在 +- 如果记忆提到了函数或标志：用 grep 搜索 +- 如果用户将基于你的推荐采取行动（不仅是询问历史），先核实 + +\"记忆说 X 存在\" 不等于 \"X 现在存在\"。"; + +/// 构建注入到 system prompt 的验证提醒。 +/// 放在记忆条目之后、`` 之前。 +pub fn build_verification_reminder() -> String { + VERIFY_BEFORE_RECOMMENDING.to_string() +} + +/// 构建保存记忆的排除规则提示（供工具描述使用）。 +pub fn build_exclusion_reminder() -> String { + WHAT_NOT_TO_SAVE.to_string() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_what_not_to_save_is_non_empty() { + assert!(!WHAT_NOT_TO_SAVE.is_empty()); + // 关键短语验证 + assert!(WHAT_NOT_TO_SAVE.contains("代码模式")); + assert!(WHAT_NOT_TO_SAVE.contains("即使用户明确要求")); + } + + #[test] + fn test_verify_before_recommending_is_non_empty() { + assert!(!VERIFY_BEFORE_RECOMMENDING.is_empty()); + assert!(VERIFY_BEFORE_RECOMMENDING.contains("从记忆中推荐前先核实")); + assert!(VERIFY_BEFORE_RECOMMENDING.contains("记忆说 X 存在")); + } + + #[test] + fn test_build_verification_reminder() { + let reminder = build_verification_reminder(); + assert_eq!(reminder, VERIFY_BEFORE_RECOMMENDING); + } + + #[test] + fn test_build_exclusion_reminder() { + let reminder = build_exclusion_reminder(); + assert!(reminder.contains("不应保存")); + } +} diff --git a/src/agent/memory/mod.rs b/src/agent/memory/mod.rs new file mode 100644 index 0000000..a9922a8 --- /dev/null +++ b/src/agent/memory/mod.rs @@ -0,0 +1,420 @@ +// src/agent/memory/mod.rs +// +// 项目记忆管理器。 +// 参考 Claude Code memdir 设计。 +// +// 在 {library_dir}/memory/ 目录下维护： +// - MEMORY.md — 索引文件（最多 200 行，25KB） +// - {slug}.md — 每个记忆一个文件，YAML frontmatter + Markdown 内容 +// +// 自动在 Agent 的 system prompt 中注入最近的记忆条目。 +// 提供 save_memory 工具供 Agent 写入记忆。 + +pub mod age; +pub mod decay; +pub mod dedup; +pub mod extraction; +pub mod guardrails; +pub mod selection; +pub mod types; + +use std::fs; +use std::path::PathBuf; +use tracing::{info, warn}; + +use self::types::{entry_from_frontmatter, parse_frontmatter, MemoryEntry, MemoryType}; + +/// 索引文件最大行数 +const MAX_ENTRYPOINT_LINES: usize = 200; +/// 索引文件最大字节数 +const MAX_ENTRYPOINT_BYTES: usize = 25_000; + +/// 记忆管理器 +pub struct MemoryManager { + /// 记忆目录 + memory_dir: PathBuf, + /// 已加载的记忆条目 + entries: Vec, + /// 自动提取追踪状态 + pub extraction_tracker: extraction::ExtractionTracker, +} + +impl MemoryManager { + /// 创建并加载记忆。 + /// `library_dir` 是项目配置中的 library 目录。 + pub fn new(library_dir: PathBuf) -> Self { + let memory_dir = library_dir.join("memory"); + let mut manager = MemoryManager { + memory_dir, + entries: Vec::new(), + extraction_tracker: extraction::ExtractionTracker::default(), + }; + manager.reload(); + manager + } + + /// 重新从磁盘加载所有记忆。 + pub fn reload(&mut self) { + // 确保目录存在 + if let Err(e) = fs::create_dir_all(&self.memory_dir) { + warn!("[Memory] 无法创建记忆目录 {:?}: {}", self.memory_dir, e); + return; + } + + self.entries.clear(); + + // 扫描 .md 文件（排除 MEMORY.md 和目录） + match fs::read_dir(&self.memory_dir) { + Ok(dir_entries) => { + for entry in dir_entries { + let entry = match entry { + Ok(e) => e, + Err(_) => continue, + }; + let path = entry.path(); + if path.is_dir() { + continue; + } + let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or(""); + if file_name == "MEMORY.md" || !file_name.ends_with(".md") { + continue; + } + + let slug = file_name.strip_suffix(".md").unwrap_or(file_name); + let mtime = entry + .metadata() + .ok() + .and_then(|m| m.modified().ok()) + .map(|t| { + t.duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs() + }) + .unwrap_or(0); + + match fs::read_to_string(&path) { + Ok(raw) => { + let (fields, content) = parse_frontmatter(&raw); + if let Some(mem_entry) = + entry_from_frontmatter(slug, &fields, &content, path.clone(), mtime) + { + self.entries.push(mem_entry); + } + } + Err(e) => { + warn!("[Memory] 无法读取 {:?}: {}", path, e); + } + } + } + } + Err(e) => { + warn!("[Memory] 无法扫描记忆目录: {}", e); + } + } + + // 按修改时间排序（最新在前） + self.entries.sort_by(|a, b| b.mtime.cmp(&a.mtime)); + + info!( + "[Memory] 加载了 {} 条记忆从 {:?}", + self.entries.len(), + self.memory_dir + ); + } + + /// 获取所有记忆条目 + pub fn entries(&self) -> &[MemoryEntry] { + &self.entries + } + + /// 获取记忆目录路径 + pub fn memory_dir(&self) -> &std::path::Path { + &self.memory_dir + } + + /// 标记主代理在本会话中已写入记忆（抑制自动提取）。 + pub fn mark_main_agent_wrote(&mut self) { + self.extraction_tracker.main_agent_saved_this_session = true; + } + + /// 语义匹配：委托给 selection 模块使用 LLM 结构化选择。 + /// + /// 选择后应用指数时间衰减排序（更近的 active 记忆获得更高权重）。 + /// 失败时回退到 recency-based 选择（跳过已展示的条目）。 + pub async fn select_relevant_memories( + &self, + llm: &crate::clients::llm::LlmClient, + context: &str, + max_entries: usize, + ) -> Vec<&MemoryEntry> { + if self.entries.len() <= max_entries { + return self.entries.iter().collect(); + } + + let sel_ctx = selection::SelectionContext { + max_entries, + ..Default::default() + }; + + let indices = selection::select_structured(&self.entries, llm, context, &sel_ctx).await; + + // 应用指数时间衰减排序（更近的 active 记忆在前，historical 在后） + let sorted = + selection::apply_decay_scoring(&indices, &self.entries, decay::DEFAULT_HALF_LIFE_DAYS); + + sorted.iter().filter_map(|&i| self.entries.get(i)).collect() + } + + /// 从指定条目列表构建 system reminder（而非全部条目） + pub fn build_system_reminder_from(&self, selected: &[&MemoryEntry]) -> Option { + if selected.is_empty() { + return None; + } + + let mut lines = vec![ + "".to_string(), + String::new(), + "[PROJECT MEMORY]".to_string(), + String::new(), + ]; + + for entry in selected { + let type_tag = match entry.memory_type { + MemoryType::User => "[偏好]", + MemoryType::Feedback => "[反馈]", + MemoryType::Project => "[项目]", + MemoryType::Reference => "[参考]", + }; + // historical 记忆标记 + let status_tag = if !entry.status.is_active() { + match entry.status.superseded_by() { + Some(new_slug) => format!(" [已更新→{}]", new_slug), + None => " [已更新]".to_string(), + } + } else { + String::new() + }; + + let preview: String = entry + .content + .lines() + .take(3) + .collect::>() + .join("\n "); + lines.push(format!( + "{} {}{}: {}\n {}", + type_tag, entry.name, status_tag, entry.description, preview + )); + // 注入时效警告（超过1天的记忆） + let freshness = age::memory_freshness_note(entry.mtime); + if !freshness.is_empty() { + lines.push(freshness); + } + } + + lines.push(String::new()); + lines.push( + "使用 save_memory 工具保存重要信息。记忆内容可能过时，请在使用前验证。".to_string(), + ); + // 注入验证提醒（从记忆推荐前先核实） + lines.push(String::new()); + lines.push(guardrails::build_verification_reminder()); + lines.push("".to_string()); + + Some(lines.join("\n")) + } + + /// 保存一条新的记忆。 + /// + /// 如果 slug 已存在且内容为 Active，旧文件归档为 `{slug}_v1.md` + /// 并将状态标记为 historical（永不删除旧记忆）。 + pub fn save_memory( + &mut self, + slug: &str, + name: &str, + description: &str, + memory_type: MemoryType, + content: &str, + ) -> std::io::Result<()> { + let file_path = self.memory_dir.join(format!("{}.md", slug)); + + // 如果已有活跃版本，归档旧版本 + let old_path = self.memory_dir.join(format!("{}_v1.md", slug)); + if file_path.exists() { + if let Ok(old_content) = fs::read_to_string(&file_path) { + fs::write(&old_path, &old_content)?; + info!("[Memory] 归档旧版本: {} → {}", slug, old_path.display()); + } + } + + let frontmatter = format!( + "---\nname: {}\ndescription: {}\ntype: {}\nstatus: active\n---\n", + name, + description, + memory_type.as_str() + ); + let full_content = format!("{}{}", frontmatter, content); + + fs::write(&file_path, &full_content)?; + + // 更新 MEMORY.md 索引 + self.update_index(slug, name, description, &memory_type)?; + + // 重新加载 + self.reload(); + + info!("[Memory] 已保存记忆: {} ({})", name, slug); + Ok(()) + } + + /// 更新 MEMORY.md 索引文件。 + fn update_index( + &self, + slug: &str, + name: &str, + description: &str, + memory_type: &MemoryType, + ) -> std::io::Result<()> { + let index_path = self.memory_dir.join("MEMORY.md"); + let line = format!( + "- [{}]({}.md) — {} (type: {})", + name, + slug, + description, + memory_type.as_str() + ); + + let mut content = if index_path.exists() { + let existing = fs::read_to_string(&index_path).unwrap_or_default(); + // 检查是否已有此 slug 的条目 + let slug_marker = format!("]({}.md)", slug); + let lines: Vec<&str> = existing.lines().collect(); + + // 行数检查 + if lines.len() >= MAX_ENTRYPOINT_LINES { + // 移除最旧的行（索引头部保持不变） + warn!( + "[Memory] MEMORY.md 行数已满 ({}), 移除最旧条目", + lines.len() + ); + let keep = MAX_ENTRYPOINT_LINES - 1; + format!("{}\n{}", lines[..keep.min(lines.len())].join("\n"), line) + } else { + // 检查是否需要替换已存在的条目 + let has_entry = lines.iter().any(|l| l.contains(&slug_marker)); + if has_entry { + // 替换已存在的行 + lines + .iter() + .map(|l| { + if l.contains(&slug_marker) { + line.as_str() + } else { + *l + } + }) + .collect::>() + .join("\n") + } else { + format!("{}\n{}", existing, line) + } + } + } else { + format!("# Project Memory\n\n{}", line) + }; + + // 字节数检查（在大约 25KB 处截断） + if content.len() > MAX_ENTRYPOINT_BYTES { + let truncated: String = content + .char_indices() + .take_while(|(i, _)| *i < MAX_ENTRYPOINT_BYTES - 100) + .map(|(_, c)| c) + .collect(); + content = format!( + "{}\n\n[MEMORY.md 已达到 {}KB 上限，旧条目已截断]", + truncated, + MAX_ENTRYPOINT_BYTES / 1024 + ); + } + + fs::write(&index_path, &content)?; + Ok(()) + } + + /// 生成 system prompt 中注入的记忆段落。 + /// + /// 包含最近的记忆条目（最多 10 条），并在前面注明可信度提醒。 + pub fn build_system_reminder(&self, max_entries: usize) -> Option { + if self.entries.is_empty() { + return None; + } + + let mut lines = vec![ + "".to_string(), + "".to_string(), + "[PROJECT MEMORY]".to_string(), + "".to_string(), + ]; + + let count = max_entries.min(self.entries.len()); + for entry in self.entries.iter().take(count) { + let type_tag = match entry.memory_type { + MemoryType::User => "[偏好]", + MemoryType::Feedback => "[反馈]", + MemoryType::Project => "[项目]", + MemoryType::Reference => "[参考]", + }; + // historical 记忆标记 + let status_tag = if !entry.status.is_active() { + match entry.status.superseded_by() { + Some(new_slug) => format!(" [已更新→{}]", new_slug), + None => " [已更新]".to_string(), + } + } else { + String::new() + }; + + let preview: String = entry + .content + .lines() + .take(3) + .collect::>() + .join("\n "); + lines.push(format!( + "{} {}{}: {}\n {}", + type_tag, entry.name, status_tag, entry.description, preview + )); + // 注入时效警告（超过1天的记忆） + let freshness = age::memory_freshness_note(entry.mtime); + if !freshness.is_empty() { + lines.push(freshness); + } + } + + lines.push("".to_string()); + lines.push( + "使用 save_memory 工具保存重要信息。记忆内容可能过时，请在使用前验证。".to_string(), + ); + // 注入验证提醒（从记忆推荐前先核实） + lines.push(String::new()); + lines.push(guardrails::build_verification_reminder()); + lines.push("".to_string()); + + Some(lines.join("\n")) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_build_empty_reminder() { + let manager = MemoryManager { + memory_dir: PathBuf::from("/tmp/nonexistent"), + entries: Vec::new(), + extraction_tracker: extraction::ExtractionTracker::default(), + }; + assert!(manager.build_system_reminder(10).is_none()); + } +} diff --git a/src/agent/memory/selection.rs b/src/agent/memory/selection.rs new file mode 100644 index 0000000..bc1cd40 --- /dev/null +++ b/src/agent/memory/selection.rs @@ -0,0 +1,351 @@ +// src/agent/memory/selection.rs +// +// 改进的记忆相关性选择 — 参考 Claude Code findRelevantMemories.ts。 +// +// 相比旧实现 (mod.rs 中直接调用 chat_completion + regex 解析)： +// 1. 使用结构化 JSON 提示词 + 更强健的解析 +// 2. 支持 SelectionContext 跟踪已展示的记忆（避免重复选择） +// 3. 支持工具感知过滤（排除最近使用工具相关的记忆） +// 4. 失败时优雅降级到 recency 回退 + +use crate::clients::llm::LlmClient; +use tracing::info; + +use super::types::MemoryEntry; + +/// 相关性选择结果 +#[derive(Debug, Clone)] +pub struct SelectionContext { + /// 已在前几轮展示过的记忆索引（避免重复选择） + pub already_surfaced: Vec, + /// 最近使用的工具名（关于这些工具的记忆降权） + pub recent_tools: Vec, + /// 最大返回条目数 + pub max_entries: usize, +} + +impl Default for SelectionContext { + fn default() -> Self { + SelectionContext { + already_surfaced: Vec::new(), + recent_tools: Vec::new(), + max_entries: 5, + } + } +} + +/// 使用 LLM 从候选记忆中选出最相关的。 +/// +/// 先尝试 LLM 结构化选择，失败时回退到 recency-based 选择。 +/// `sel_ctx.already_surfaced` 中的条目会被排除在候选之外。 +pub async fn select_structured( + entries: &[MemoryEntry], + llm: &LlmClient, + context: &str, + sel_ctx: &SelectionContext, +) -> Vec { + // 过滤已展示的条目 + let candidates: Vec<(usize, &MemoryEntry)> = entries + .iter() + .enumerate() + .filter(|(i, _)| !sel_ctx.already_surfaced.contains(i)) + .collect(); + + if candidates.is_empty() { + return fallback_recency(entries, sel_ctx.max_entries, &sel_ctx.already_surfaced); + } + + if candidates.len() <= sel_ctx.max_entries { + return candidates.iter().map(|(i, _)| *i).collect(); + } + + // 构建候选目录 + let catalog: Vec = candidates + .iter() + .map(|(_, e)| { + format!( + "[{}] [{}] {}: {}", + e.slug, + e.memory_type.as_str(), + e.name, + e.description + ) + }) + .collect(); + + // 工具提示（可选） + let tools_hint = if sel_ctx.recent_tools.is_empty() { + String::new() + } else { + format!( + "\n\n近期使用的工具: {}。不要选择这些工具的使用参考或 API 文档类记忆。", + sel_ctx.recent_tools.join(", ") + ) + }; + + let prompt = format!( + "用户当前话题:\n{}\n\n从以下记忆目录中选择最多 {} 条最相关的。\ + 仅选择明确有帮助的，不确定则不选。返回 JSON 数组如 [\"slug1\", \"slug2\"]。\n\n{}\n{}", + context, + sel_ctx.max_entries, + catalog.join("\n"), + tools_hint, + ); + + match llm + .chat_completion( + "你是一个记忆检索助手。根据用户话题选择最相关的记忆。仅返回 JSON 字符串数组。", + &prompt, + ) + .await + { + Ok(response) => match extract_slugs(&response, &candidates) { + Ok(selected) => { + info!( + "[Memory] LLM 选择: {}/{} 条相关记忆", + selected.len(), + candidates.len() + ); + selected + } + Err(_) => { + info!("[Memory] JSON 解析失败，回退到 recency"); + fallback_recency(entries, sel_ctx.max_entries, &sel_ctx.already_surfaced) + } + }, + Err(_) => fallback_recency(entries, sel_ctx.max_entries, &sel_ctx.already_surfaced), + } +} + +/// 从 LLM 响应中提取 slug 列表。 +/// 尝试两种格式：["slug1","slug2"] 或 [0, 1, 3]（数字索引回退） +fn extract_slugs(response: &str, candidates: &[(usize, &MemoryEntry)]) -> Result, ()> { + // 方法1: 查找 JSON 字符串数组 + if let Some(start) = response.find('[') { + if let Some(end) = response.rfind(']') { + let json_str = &response[start..=end]; + + // 尝试解析为字符串数组 + if let Ok(slugs) = serde_json::from_str::>(json_str) { + let indices: Vec = slugs + .iter() + .filter_map(|s| { + candidates + .iter() + .find(|(_, e)| e.slug == *s) + .map(|(i, _)| *i) + }) + .collect(); + if !indices.is_empty() { + return Ok(indices); + } + } + + // 方法2: 回退到数字索引 + if let Ok(indices) = serde_json::from_str::>(json_str) { + let valid: Vec = indices + .into_iter() + .filter(|i| candidates.iter().any(|(idx, _)| *idx == *i)) + .collect(); + if !valid.is_empty() { + return Ok(valid); + } + } + } + } + Err(()) +} + +/// 将指数时间衰减评分应用于已选中的记忆，按最终得分降序排序。 +/// +/// 每项最终得分 = decay_score(mtime) × 1.0（LLM 选中即默认置信度）。 +/// 仅影响排序顺序，不减少选中的数量。 +/// 仅活跃记忆参与衰减排序；historical 记忆保持原顺序。 +pub fn apply_decay_scoring( + indices: &[usize], + entries: &[MemoryEntry], + half_life_days: f64, +) -> Vec { + use super::decay::decay_score; + + let mut scored: Vec<(usize, f64)> = indices + .iter() + .filter_map(|&i| { + entries.get(i).map(|entry| { + if entry.status.is_active() { + (i, decay_score(entry.mtime, half_life_days)) + } else { + // historical 记忆赋予最低分，排在最后 + (i, 0.01) + } + }) + }) + .collect(); + + // 按得分降序排列（高分在前） + scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + scored.into_iter().map(|(i, _)| i).collect() +} + +/// Recency 回退：取最近的 max_entries 条（排除已展示的） +fn fallback_recency( + entries: &[MemoryEntry], + max_entries: usize, + already_surfaced: &[usize], +) -> Vec { + entries + .iter() + .enumerate() + .filter(|(i, _)| !already_surfaced.contains(i)) + .take(max_entries) + .map(|(i, _)| i) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_entry(slug: &str, name: &str, desc: &str) -> MemoryEntry { + MemoryEntry { + slug: slug.to_string(), + name: name.to_string(), + description: desc.to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: 1000, + content: String::new(), + path: std::path::PathBuf::from(slug), + status: super::super::types::MemoryStatus::Active, + } + } + + #[test] + fn test_fallback_recency_basic() { + let entries = vec![ + make_entry("a", "A", "desc a"), + make_entry("b", "B", "desc b"), + make_entry("c", "C", "desc c"), + ]; + let result = fallback_recency(&entries, 2, &[]); + assert_eq!(result.len(), 2); + assert_eq!(result[0], 0); // 最近的在前 + assert_eq!(result[1], 1); + } + + #[test] + fn test_fallback_skips_already_surfaced() { + let entries = vec![ + make_entry("a", "A", "desc a"), + make_entry("b", "B", "desc b"), + make_entry("c", "C", "desc c"), + ]; + let result = fallback_recency(&entries, 3, &[0]); // skip index 0 + assert_eq!(result, vec![1, 2]); + } + + #[test] + fn test_extract_slugs_string_array() { + let entries = vec![ + make_entry("alpha", "Alpha", "first"), + make_entry("beta", "Beta", "second"), + make_entry("gamma", "Gamma", "third"), + ]; + let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect(); + + let result = extract_slugs(r#"["alpha", "gamma"]"#, &candidates); + assert_eq!(result, Ok(vec![0, 2])); + } + + #[test] + fn test_extract_slugs_numeric_fallback() { + let entries = vec![make_entry("x", "X", "x"), make_entry("y", "Y", "y")]; + let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect(); + + let result = extract_slugs("[0]", &candidates); + assert_eq!(result, Ok(vec![0])); + } + + #[test] + fn test_extract_slugs_invalid_json() { + let entries = vec![make_entry("x", "X", "x")]; + let candidates: Vec<(usize, &MemoryEntry)> = entries.iter().enumerate().collect(); + + assert!(extract_slugs("not json at all", &candidates).is_err()); + assert!(extract_slugs("no brackets here", &candidates).is_err()); + } + + #[test] + fn test_selection_context_default() { + let ctx = SelectionContext::default(); + assert!(ctx.already_surfaced.is_empty()); + assert!(ctx.recent_tools.is_empty()); + assert_eq!(ctx.max_entries, 5); + } + + #[test] + fn test_apply_decay_scoring_sorts_by_freshness() { + use crate::agent::memory::age; + let now = age::now_secs(); + let entries = vec![ + MemoryEntry { + slug: "old".to_string(), + name: "旧记忆".to_string(), + description: "old".to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: now - 86_400 * 60, // 60天前 + content: String::new(), + path: std::path::PathBuf::from("old.md"), + status: super::super::types::MemoryStatus::Active, + }, + MemoryEntry { + slug: "fresh".to_string(), + name: "新鲜记忆".to_string(), + description: "fresh".to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: now, // 今天 + content: String::new(), + path: std::path::PathBuf::from("fresh.md"), + status: super::super::types::MemoryStatus::Active, + }, + ]; + // 旧记忆在前（index 0），新记忆在后（index 1） + let sorted = apply_decay_scoring(&[0, 1], &entries, 30.0); + // 新记忆应排在旧记忆前面 + assert_eq!(sorted[0], 1); + assert_eq!(sorted[1], 0); + } + + #[test] + fn test_apply_decay_puts_historical_last() { + use crate::agent::memory::age; + let now = age::now_secs(); + let entries = vec![ + MemoryEntry { + slug: "active".to_string(), + name: "活跃".to_string(), + description: "active".to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: now, + content: String::new(), + path: std::path::PathBuf::from("active.md"), + status: super::super::types::MemoryStatus::Active, + }, + MemoryEntry { + slug: "historical".to_string(), + name: "历史".to_string(), + description: "historical".to_string(), + memory_type: super::super::types::MemoryType::User, + mtime: now, // 也很新，但是 historical + content: String::new(), + path: std::path::PathBuf::from("historical.md"), + status: super::super::types::MemoryStatus::Historical { + superseded_by: Some("active".to_string()), + }, + }, + ]; + let sorted = apply_decay_scoring(&[0, 1], &entries, 30.0); + // historical 应排在最后 + assert_eq!(sorted[0], 0, "活跃记忆应在前"); + assert_eq!(sorted[1], 1, "historical 应在后"); + } +} diff --git a/src/agent/memory/types.rs b/src/agent/memory/types.rs new file mode 100644 index 0000000..fc12487 --- /dev/null +++ b/src/agent/memory/types.rs @@ -0,0 +1,330 @@ +// src/agent/memory/types.rs +// +// 项目记忆系统 — 类型定义。 +// 参考 Claude Code memdir/memoryTypes.ts 设计。 +// +// 四种记忆类型： +// - User — 用户角色、偏好、目标 +// - Feedback — 用户提供的反馈（修正 + 确认） +// - Project — 项目状态、进行中的工作、目标 +// - Reference — 外部资源的指针 + +use std::path::PathBuf; + +/// 记忆类型 +#[derive(Debug, Clone, PartialEq)] +pub enum MemoryType { + User, + Feedback, + Project, + Reference, +} + +impl MemoryType { + pub fn as_str(&self) -> &'static str { + match self { + MemoryType::User => "user", + MemoryType::Feedback => "feedback", + MemoryType::Project => "project", + MemoryType::Reference => "reference", + } + } +} + +impl std::str::FromStr for MemoryType { + type Err = (); + fn from_str(s: &str) -> Result { + match s { + "user" => Ok(MemoryType::User), + "feedback" => Ok(MemoryType::Feedback), + "project" => Ok(MemoryType::Project), + "reference" => Ok(MemoryType::Reference), + _ => Err(()), + } + } +} + +impl MemoryType { + #[allow(clippy::should_implement_trait)] + pub fn from_str(s: &str) -> Option { + s.parse().ok() + } +} + +/// 记忆生命周期状态 — 参考 Martian-Engineering agent-memory。 +/// +/// 事实永远不会被删除，只会从 Active 转换为 Historical。 +/// supersedes 链保留了"理解如何演变"的完整历史。 +#[derive(Debug, Clone, PartialEq)] +#[derive(Default)] +pub enum MemoryStatus { + /// 当前有效 + #[default] + Active, + /// 已被更新事实取代，superseded_by 指向新 slug + Historical { superseded_by: Option }, +} + +impl MemoryStatus { + pub fn as_str(&self) -> &'static str { + match self { + MemoryStatus::Active => "active", + MemoryStatus::Historical { .. } => "historical", + } + } + + pub fn from_str(s: &str, superseded_by: Option) -> Self { + match s { + "historical" => MemoryStatus::Historical { superseded_by }, + _ => MemoryStatus::Active, + } + } + + /// 是否为活跃状态 + pub fn is_active(&self) -> bool { + matches!(self, MemoryStatus::Active) + } + + /// 获取取代此记忆的新 slug（如果有） + pub fn superseded_by(&self) -> Option<&str> { + match self { + MemoryStatus::Historical { + superseded_by: Some(s), + } => Some(s.as_str()), + _ => None, + } + } +} + + +/// 记忆条目（从 .md 文件解析） +#[derive(Debug, Clone)] +pub struct MemoryEntry { + /// 文件名（不含扩展名，作为 slug） + pub slug: String, + /// 记忆标题 + pub name: String, + /// 简短描述（用于相关性匹配） + pub description: String, + /// 记忆类型 + pub memory_type: MemoryType, + /// 文件修改时间 + pub mtime: u64, + /// 完整内容（不含 frontmatter） + pub content: String, + /// 文件路径 + pub path: PathBuf, + /// 生命周期状态（默认 Active） + pub status: MemoryStatus, +} + +/// MEMORY.md 索引中的一行 +#[derive(Debug, Clone)] +pub struct MemoryIndexLine { + pub slug: String, + pub title: String, + pub description: String, + pub memory_type: MemoryType, +} + +/// 解析 Markdown 文件的 YAML frontmatter。 +/// +/// Frontmatter 格式（向后兼容，status/superseded_by 可选）： +/// ```markdown +/// --- +/// name: my-memory +/// description: Short description +/// metadata: +/// type: user +/// status: active # 可选，缺失默认 active +/// superseded_by: "" # 可选，historical 时指向新 slug +/// --- +/// Content here... +/// ``` +/// +/// 返回 (frontmatter_fields, content)。 +pub fn parse_frontmatter(raw: &str) -> (Vec<(String, String)>, String) { + let mut fields = Vec::new(); + let content; + + if let Some(rest) = raw.strip_prefix("---\n") { + if let Some(end_pos) = rest.find("\n---\n") { + let fm = &rest[..end_pos]; + content = rest[end_pos + 5..].to_string(); + + for line in fm.lines() { + let line = line.trim(); + if line.is_empty() { + continue; + } + if let Some(colon_pos) = line.find(':') { + let key = line[..colon_pos].trim().to_string(); + let value = line[colon_pos + 1..].trim().to_string(); + fields.push((key, value)); + } + } + } else { + content = raw.to_string(); + } + } else { + content = raw.to_string(); + } + + (fields, content) +} + +/// 从 frontmatter 字段构建 MemoryEntry。 +pub fn entry_from_frontmatter( + slug: &str, + fields: &[(String, String)], + content: &str, + path: PathBuf, + mtime: u64, +) -> Option { + let mut name = String::new(); + let mut description = String::new(); + let mut memory_type = MemoryType::User; // default + let mut status_str = String::new(); + let mut superseded_by = String::new(); + + for (key, value) in fields { + match key.as_str() { + "name" => name = value.clone(), + "description" => description = value.clone(), + "type" | "memory_type" => { + if let Some(t) = MemoryType::from_str(value) { + memory_type = t; + } + } + "status" => status_str = value.clone(), + "superseded_by" => superseded_by = value.clone(), + _ => {} + } + } + + if name.is_empty() { + name = slug.to_string(); + } + + let status = MemoryStatus::from_str( + &status_str, + if superseded_by.is_empty() { + None + } else { + Some(superseded_by) + }, + ); + + Some(MemoryEntry { + slug: slug.to_string(), + name, + description, + memory_type, + mtime, + content: content.to_string(), + path, + status, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_parse_frontmatter_basic() { + let raw = "---\nname: test-memory\ndescription: A test\nmetadata:\n type: user\n---\nThis is the content."; + let (fields, content) = parse_frontmatter(raw); + assert_eq!(content.trim(), "This is the content."); + assert!(fields + .iter() + .any(|(k, v)| k == "name" && v == "test-memory")); + assert!(fields + .iter() + .any(|(k, v)| k == "description" && v == "A test")); + } + + #[test] + fn test_parse_no_frontmatter() { + let raw = "Just content, no frontmatter."; + let (fields, content) = parse_frontmatter(raw); + assert_eq!(content, raw); + assert!(fields.is_empty()); + } + + #[test] + fn test_memory_type_from_str() { + assert_eq!(MemoryType::from_str("user"), Some(MemoryType::User)); + assert_eq!(MemoryType::from_str("feedback"), Some(MemoryType::Feedback)); + assert_eq!(MemoryType::from_str("project"), Some(MemoryType::Project)); + assert_eq!( + MemoryType::from_str("reference"), + Some(MemoryType::Reference) + ); + assert_eq!(MemoryType::from_str("invalid"), None); + } + + #[test] + fn test_memory_status_default_active() { + let status = MemoryStatus::default(); + assert_eq!(status, MemoryStatus::Active); + assert!(status.is_active()); + } + + #[test] + fn test_memory_status_parse_active() { + let status = MemoryStatus::from_str("active", None); + assert_eq!(status, MemoryStatus::Active); + assert!(status.is_active()); + } + + #[test] + fn test_memory_status_parse_historical_with_superseded_by() { + let status = MemoryStatus::from_str("historical", Some("new-version".to_string())); + assert!(!status.is_active()); + assert_eq!(status.superseded_by(), Some("new-version")); + } + + #[test] + fn test_memory_status_parse_unknown_defaults_to_active() { + let status = MemoryStatus::from_str("invalid", None); + assert_eq!(status, MemoryStatus::Active); + } + + #[test] + fn test_entry_from_frontmatter_parses_status_and_superseded_by() { + let fields = vec![ + ("name".to_string(), "test-mem".to_string()), + ("description".to_string(), "A description".to_string()), + ("status".to_string(), "historical".to_string()), + ("superseded_by".to_string(), "better-slug".to_string()), + ]; + let entry = entry_from_frontmatter( + "test-mem", + &fields, + "Content here", + PathBuf::from("test-mem.md"), + 1000, + ) + .unwrap(); + assert!(!entry.status.is_active()); + assert_eq!(entry.status.superseded_by(), Some("better-slug")); + } + + #[test] + fn test_entry_from_frontmatter_missing_status_defaults_to_active() { + let fields = vec![ + ("name".to_string(), "test-mem".to_string()), + ("description".to_string(), "A description".to_string()), + ]; + let entry = entry_from_frontmatter( + "test-mem", + &fields, + "Content here", + PathBuf::from("test-mem.md"), + 1000, + ) + .unwrap(); + assert!(entry.status.is_active()); + } +} diff --git a/src/agent/mod.rs b/src/agent/mod.rs index e970f71..04e6d7a 100644 --- a/src/agent/mod.rs +++ b/src/agent/mod.rs @@ -1,6 +1,24 @@ // src/agent/mod.rs -// 科研智能体模块 -// 基于 ReAct 框架实现 Thought -> Action -> Observation 循环 +// +// 科研智能体模块 — 基于 ReAct 框架实现 Thought → Action → Observation 循环。 +// +// 模块结构（参考 Claude Code 分层设计）： +// tools/ — 工具定义与注册（按功能域拆分） +// runtime — ReAct 循环引擎 + Streaming + 会话管理 +// compact — 三层上下文压缩（micro/auto/manual） +// terminal — 循环终止信号（结构化退出原因） +// hooks — 生命周期事件系统（PreToolUse/PostToolUse/Stop） -pub mod tools; +pub mod autonomous; +pub mod background; +pub mod compact; +pub mod hooks; +pub mod memory; pub mod runtime; +pub mod skills; +pub mod subagent; +pub mod task_board; +pub mod team; +pub mod terminal; +pub mod tools; +pub mod trajectory; diff --git a/src/agent/runtime.rs b/src/agent/runtime.rs deleted file mode 100644 index b7931c2..0000000 --- a/src/agent/runtime.rs +++ /dev/null @@ -1,786 +0,0 @@ -// src/agent/runtime.rs -// -// 科研智能体运行时核心模块。 -// 实现 ReAct 循环：Thought -> Action (工具调用) -> Observation -> Thought... -// 支持会话持久化、上下文压缩、死循环检测和 SSE 流式输出。 - -use std::sync::Arc; -use tracing::{info, warn, error}; -use serde::Serialize; -use sqlx::SqlitePool; -use tokio::sync::mpsc; - -use crate::api::AppState; -use crate::clients::llm::{ - ChatMessage, LlmClient, MessageRole, StreamEvent, -}; -use super::tools::{ToolContext, ToolOutput, ToolRegistry}; - -/// Agent 配置参数 -#[derive(Debug, Clone)] -pub struct AgentConfig { - /// 最大 ReAct 迭代次数 - pub max_steps: usize, - /// 同质调用检测阈值（连续相同调用次数） - pub duplicate_call_threshold: usize, - /// 工具执行超时时间（秒） - pub tool_timeout_secs: u64, - /// 工具输出最大字符数 - pub max_tool_output_chars: usize, - /// 上下文 Token 估算上限（触发自动摘要压缩） - pub context_char_limit: usize, -} - -impl Default for AgentConfig { - fn default() -> Self { - AgentConfig { - max_steps: std::env::var("AGENT_MAX_STEPS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(8), - duplicate_call_threshold: 3, - tool_timeout_secs: std::env::var("AGENT_TOOL_TIMEOUT_SECS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(120), - max_tool_output_chars: std::env::var("AGENT_MAX_TOOL_OUTPUT_CHARS") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(4000), - context_char_limit: std::env::var("AGENT_CONTEXT_CHAR_LIMIT") - .ok() - .and_then(|v| v.parse().ok()) - .unwrap_or(16000), - } - } -} - -/// SSE 流式事件（发送给前端） -#[derive(Debug, Clone, Serialize)] -#[serde(tag = "type")] -pub enum AgentStreamEvent { - /// 会话创建/恢复 - #[serde(rename = "session")] - Session { - session_id: String, - title: String, - }, - /// 智能体思考过程 - #[serde(rename = "thought")] - Thought { - content: String, - step: usize, - }, - /// 工具调用开始 - #[serde(rename = "tool_call")] - ToolCall { - name: String, - arguments: serde_json::Value, - step: usize, - }, - /// 工具执行结果（Observation） - #[serde(rename = "tool_result")] - ToolResult { - name: String, - output: String, - is_error: bool, - metadata: serde_json::Value, - step: usize, - }, - /// 文本增量流式输出（最终回答） - #[serde(rename = "text_delta")] - TextDelta { - content: String, - }, - /// Token 使用统计 - #[serde(rename = "usage")] - Usage { - prompt_tokens: u32, - completion_tokens: u32, - total_tokens: u32, - }, - /// 错误通知 - #[serde(rename = "error")] - Error { - message: String, - }, - /// 完成标记 - #[serde(rename = "done")] - Done, -} - -/// 同质调用检测器 -#[derive(Debug, Default)] -struct DuplicateDetector { - last_call: Option<(String, String)>, // (tool_name, arguments) - consecutive_count: usize, -} - -impl DuplicateDetector { - /// 记录一次调用，返回是否检测到死循环 - fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool { - let key = (tool_name.to_string(), arguments.to_string()); - if self.last_call.as_ref() == Some(&key) { - self.consecutive_count += 1; - if self.consecutive_count >= threshold { - return true; - } - } else { - self.last_call = Some(key); - self.consecutive_count = 1; - } - false - } -} - -/// 智能体运行时 -pub struct AgentRuntime { - app_state: Arc, - config: AgentConfig, - tool_registry: ToolRegistry, -} - -impl AgentRuntime { - /// 创建新的运行时实例 - pub fn new(app_state: Arc) -> Self { - AgentRuntime { - app_state, - config: AgentConfig::default(), - tool_registry: ToolRegistry::new(), - } - } - - /// 创建带自定义配置的运行时实例 - pub fn with_config(app_state: Arc, config: AgentConfig) -> Self { - AgentRuntime { - app_state, - config, - tool_registry: ToolRegistry::new(), - } - } - - /// 执行完整的智能体对话回合（流式 SSE 输出） - /// - /// 流程： - /// 1. 加载或创建会话 - /// 2. 构建消息上下文 - /// 3. ReAct 循环：LLM 调用 -> 工具执行 -> 结果注入 -> 再次调用 ... - /// 4. 最终回答流式输出 - /// 5. 持久化所有消息 - pub async fn run_turn( - &self, - session_id: Option, - question: &str, - tx: mpsc::UnboundedSender, - ) -> anyhow::Result { - let db = &self.app_state.db; - let llm = &self.app_state.llm; - - // 1. 创建或恢复会话 - let sid = match session_id { - Some(id) => { - // 验证会话存在 - let exists: bool = sqlx::query_scalar( - "SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)" - ) - .bind(&id) - .fetch_one(db) - .await - .unwrap_or(false); - - if !exists { - return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id)); - } - id - } - None => { - let new_id = uuid::Uuid::new_v4().to_string(); - sqlx::query( - "INSERT INTO agent_sessions (session_id, title, model) VALUES (?, ?, ?)" - ) - .bind(&new_id) - .bind("") - .bind(llm.model()) - .execute(db) - .await?; - new_id - } - }; - - let _ = tx.send(AgentStreamEvent::Session { - session_id: sid.clone(), - title: String::new(), - }); - - // 2. 加载历史消息（过滤掉 thought 字段，仅保留纯对话上下文） - let mut messages = self.load_history_for_llm(db, &sid).await?; - - // 获取当前轮次号 - let turn_index: i32 = sqlx::query_scalar( - "SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?" - ) - .bind(&sid) - .fetch_one(db) - .await - .unwrap_or(0); - - // 3. 构建系统提示词 - if messages.is_empty() || messages[0].role != MessageRole::System { - messages.insert(0, ChatMessage::system(self.system_prompt())); - } - - // 4. 添加用户消息 - messages.push(ChatMessage::user(question)); - self.save_message(db, &sid, turn_index, 0, &ChatMessage::user(question), None).await?; - - // 5. ReAct 循环 - let tool_defs = self.tool_registry.definitions(); - let tool_ctx = ToolContext { - app_state: Arc::clone(&self.app_state), - }; - let mut duplicate_detector = DuplicateDetector::default(); - let mut step = 0; - - loop { - step += 1; - // 检查是否被用户手动中止 - if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { - if cancelled.remove(&sid) { - warn!("[AgentRuntime] 用户手动中止了会话 {} 的智能体执行", sid); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - break; - } - } - - // 上下文安全检查 - let context_chars: usize = messages.iter() - .filter_map(|m| m.content.as_ref()) - .map(|c| c.len()) - .sum(); - - if context_chars > self.config.context_char_limit { - info!("[AgentRuntime] 上下文超限 ({} > {})，触发压缩", context_chars, self.config.context_char_limit); - self.compress_context(&mut messages, llm).await; - } - - // 调用 LLM（使用 `chat_stream` 实时流式读取，支持思维过程/最终回答的流式发送和中止检测） - let mut stream_rx = match llm.chat_stream(&messages, &tool_defs).await { - Ok(rx) => rx, - Err(e) => { - error!("[AgentRuntime] LLM stream 调用失败: {}", e); - let _ = tx.send(AgentStreamEvent::Error { - message: format!("大模型流式调用失败: {}", e), - }); - break; - } - }; - - let mut accumulated_content = String::new(); - let mut accumulated_reasoning = String::new(); - let mut accumulated_tool_calls: Option> = None; - let mut usage: Option = None; - let mut is_tool_call_step = false; - - let cancel_fut = async { - loop { - tokio::time::sleep(std::time::Duration::from_millis(250)).await; - if let Ok(cancelled) = self.app_state.cancelled_runs.lock() { - if cancelled.contains(&sid) { - return; - } - } - } - }; - - enum StreamLoopResult { - Success, - Error(String), - Cancelled, - } - - let stream_loop_res = { - let mut cancel_pinned = Box::pin(cancel_fut); - let mut error_msg = None; - let mut cancelled = false; - - loop { - tokio::select! { - event_opt = stream_rx.recv() => { - match event_opt { - Some(event) => { - match event { - StreamEvent::ReasoningDelta(delta) => { - accumulated_reasoning.push_str(&delta); - // 实时流式发送思考过程给前端 - let _ = tx.send(AgentStreamEvent::Thought { - content: accumulated_reasoning.clone(), - step, - }); - } - StreamEvent::TextDelta(delta) => { - accumulated_content.push_str(&delta); - // 如果目前还没发现是工具调用步骤，就实时流式发送文本给前端作为最终回答 - if !is_tool_call_step { - let _ = tx.send(AgentStreamEvent::TextDelta { - content: delta, - }); - } - } - StreamEvent::ToolCallsComplete(tool_calls) => { - is_tool_call_step = true; - accumulated_tool_calls = Some(tool_calls); - } - StreamEvent::ToolCallDelta { .. } => {} - StreamEvent::Usage(u) => { - usage = Some(u); - } - StreamEvent::Done => { - break; - } - StreamEvent::Error(e) => { - error_msg = Some(e); - break; - } - } - } - None => break, - } - } - _ = &mut cancel_pinned => { - cancelled = true; - break; - } - } - } - - if cancelled { - StreamLoopResult::Cancelled - } else if let Some(e) = error_msg { - StreamLoopResult::Error(e) - } else { - StreamLoopResult::Success - } - }; - - match stream_loop_res { - StreamLoopResult::Success => {} - StreamLoopResult::Error(e_str) => { - error!("[AgentRuntime] 流式读取错误: {}", e_str); - let _ = tx.send(AgentStreamEvent::Error { - message: format!("大模型流式读取失败: {}", e_str), - }); - break; - } - StreamLoopResult::Cancelled => { - if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { - cancelled.remove(&sid); - } - warn!("[AgentRuntime] 在流式调用期间被用户手动中止，会话 ID: {}", sid); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - break; - } - } - - // 处理 Thought（优先使用原生推理内容，否则如果属于工具调用步骤，使用 accumulated_content 存储） - let mut thought_content = None; - if !accumulated_reasoning.is_empty() { - thought_content = Some(accumulated_reasoning.clone()); - } - - if thought_content.is_none() && is_tool_call_step { - if !accumulated_content.is_empty() { - // 有工具调用时，content 被视为 Thought - thought_content = Some(accumulated_content.clone()); - } - } - - // 如果是在工具调用步骤中产生的前言描述，而我们之前没实时以 Thought 发送过，此时统一作为 Thought 发送给前端展示 - if is_tool_call_step { - if let Some(ref thought_text) = thought_content { - let _ = tx.send(AgentStreamEvent::Thought { - content: thought_text.clone(), - step, - }); - } - } - - let reasoning_option = if accumulated_reasoning.is_empty() { None } else { Some(accumulated_reasoning.clone()) }; - - // 无工具调用 = 最终回答 - if accumulated_tool_calls.is_none() || accumulated_tool_calls.as_ref().unwrap().is_empty() { - // 保存助手最终回答消息 - let assistant_msg = ChatMessage::assistant_with_reasoning( - Some(accumulated_content.clone()), - reasoning_option.clone(), - None, - ); - - self.save_message(db, &sid, turn_index, step as i32, &assistant_msg, reasoning_option.as_deref()).await?; - messages.push(assistant_msg); - - // 如果有原生推理内容且之前没发送过，发送给前端展示最终思维链 - if let Some(ref thought_text) = reasoning_option { - if thought_content.is_none() { - let _ = tx.send(AgentStreamEvent::Thought { - content: thought_text.clone(), - step, - }); - } - } - - // 流式发送 Done 或 Token 消耗 - if let Some(u) = usage { - let _ = tx.send(AgentStreamEvent::Usage { - prompt_tokens: u.prompt_tokens, - completion_tokens: u.completion_tokens, - total_tokens: u.total_tokens, - }); - } - break; - } - - let tool_calls = accumulated_tool_calls.unwrap(); - - // 有工具调用 —— 构建 assistant 消息（含 tool_calls 和 reasoning_content） - let assistant_msg = ChatMessage::assistant_with_reasoning( - if accumulated_content.is_empty() { None } else { Some(accumulated_content.clone()) }, - reasoning_option.clone(), - Some(tool_calls.clone()), - ); - self.save_message(db, &sid, turn_index, step as i32, &assistant_msg, reasoning_option.as_deref()).await?; - messages.push(assistant_msg); - - // 逐个执行工具 - for tool_call in &tool_calls { - let tool_name = &tool_call.function.name; - let tool_args_str = &tool_call.function.arguments; - - // 死循环检测 - if duplicate_detector.record(tool_name, tool_args_str, self.config.duplicate_call_threshold) { - warn!("[AgentRuntime] 检测到死循环：{} 连续调用 {} 次", tool_name, self.config.duplicate_call_threshold); - let _ = tx.send(AgentStreamEvent::Error { - message: format!("检测到工具 {} 的重复调用，已自动终止循环。", tool_name), - }); - - // 注入错误 tool result 让 LLM 知道要停止 - let error_msg = ChatMessage::tool_result( - &tool_call.id, - format!("错误：工具 {} 被连续重复调用 {} 次，参数完全相同。请停止重复调用并直接给出目前收集到的答案。", tool_name, self.config.duplicate_call_threshold), - ); - messages.push(error_msg); - continue; - } - - // 解析参数 - let args: serde_json::Value = match serde_json::from_str(tool_args_str) { - Ok(v) => v, - Err(e) => { - let error_output = format!("工具参数 JSON 解析失败: {}", e); - let _ = tx.send(AgentStreamEvent::ToolResult { - name: tool_name.clone(), - output: error_output.clone(), - is_error: true, - metadata: serde_json::json!({}), - step, - }); - let tool_msg = ChatMessage::tool_result(&tool_call.id, &error_output); - self.save_message(db, &sid, turn_index, step as i32, &tool_msg, None).await?; - messages.push(tool_msg); - continue; - } - }; - - // 发送工具调用事件 - let _ = tx.send(AgentStreamEvent::ToolCall { - name: tool_name.clone(), - arguments: args.clone(), - step, - }); - - enum ToolResultEnum { - Success(ToolOutput), - Cancelled, - } - - // 执行工具（带超时保护和手动中止检测） - let tool_res = match self.tool_registry.get(tool_name) { - Some(tool) => { - let timeout = std::time::Duration::from_secs(self.config.tool_timeout_secs); - let tool_fut = tool.execute(args, &tool_ctx); - let cancel_fut = async { - loop { - tokio::time::sleep(std::time::Duration::from_millis(250)).await; - if let Ok(cancelled) = self.app_state.cancelled_runs.lock() { - if cancelled.contains(&sid) { - return; - } - } - } - }; - - tokio::select! { - res = tokio::time::timeout(timeout, tool_fut) => { - match res { - Ok(output) => ToolResultEnum::Success(output), - Err(_) => ToolResultEnum::Success(ToolOutput::error(format!("工具 {} 执行超时（{}秒）", tool_name, self.config.tool_timeout_secs))), - } - } - _ = cancel_fut => { - ToolResultEnum::Cancelled - } - } - } - None => ToolResultEnum::Success(ToolOutput::error(format!("未知工具: {}", tool_name))), - }; - - let output = match tool_res { - ToolResultEnum::Success(out) => out, - ToolResultEnum::Cancelled => { - if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { - cancelled.remove(&sid); - } - warn!("[AgentRuntime] 在工具 {} 执行期间被用户手动中止，会话 ID: {}", tool_name, sid); - let _ = tx.send(AgentStreamEvent::Error { - message: "用户已手动中止执行。".to_string(), - }); - break; - } - }; - - // 发送工具结果事件 - let _ = tx.send(AgentStreamEvent::ToolResult { - name: tool_name.clone(), - output: output.content.clone(), - is_error: output.is_error, - metadata: output.metadata.clone(), - step, - }); - - // 截断工具输出 - let truncated_content = if output.content.len() > self.config.max_tool_output_chars { - let truncated: String = output.content.chars().take(self.config.max_tool_output_chars).collect(); - format!("{}...\n[输出已截断，原始长度: {} 字符]", truncated, output.content.len()) - } else { - output.content.clone() - }; - - // 构建 tool result 消息 - let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated_content); - self.save_message(db, &sid, turn_index, step as i32, &tool_msg, None).await?; - messages.push(tool_msg); - } - } - - // 6. 更新会话元信息 - let new_turn_count: i32 = sqlx::query_scalar( - "SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?" - ) - .bind(&sid) - .fetch_one(db) - .await - .unwrap_or(0); - - // 首轮自动生成标题 - if new_turn_count <= 1 { - let title = self.generate_title(question); - sqlx::query("UPDATE agent_sessions SET title = ?, turn_count = ?, updated_at = CURRENT_TIMESTAMP WHERE session_id = ?") - .bind(&title) - .bind(new_turn_count) - .bind(&sid) - .execute(db) - .await?; - } else { - sqlx::query("UPDATE agent_sessions SET turn_count = ?, updated_at = CURRENT_TIMESTAMP WHERE session_id = ?") - .bind(new_turn_count) - .bind(&sid) - .execute(db) - .await?; - } - - let _ = tx.send(AgentStreamEvent::Done); - - Ok(sid) - } - - /// 系统提示词 - fn system_prompt(&self) -> String { - "你是一位专业的天体物理学研究助手，具备丰富的天文学知识。你可以使用以下工具帮助用户进行科研工作：\n\ - \n\ - - search_papers: 检索天文学文献（ADS/arXiv）\n\ - - get_paper_content: 获取文献全文内容（自动下载、解析）\n\ - - read_local_file: 快速读取已解析的本地文献\n\ - - rag_search: 在已向量化的文献库中进行语义检索\n\ - - query_target: 查询天体物理属性（坐标、光谱型等）\n\ - \n\ - 请遵循以下原则：\n\ - 1. 先思考用户的问题需要什么信息，再决定调用哪些工具。\n\ - 2. 优先使用已有的本地文献资源（read_local_file / rag_search），必要时再检索新文献。\n\ - 3. 回答时引用具体文献来源，使用 ADS bibcode 标注。\n\ - 4. 对于数学公式，使用标准 LaTeX 格式。\n\ - 5. 用中文回答用户的问题，但保持科学术语的准确性（可附带英文原文）。\n\ - 6. 如果一个工具调用失败，不要重复使用完全相同的参数重试，尝试换一种方式。".to_string() - } - - /// 从数据库加载历史消息（包含 thought 作为 reasoning_content，以备原生思考模型使用） - async fn load_history_for_llm( - &self, - db: &SqlitePool, - session_id: &str, - ) -> anyhow::Result> { - let rows: Vec<(String, String, Option, Option, Option)> = sqlx::query_as( - "SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \ - WHERE session_id = ? ORDER BY id ASC" - ) - .bind(session_id) - .fetch_all(db) - .await?; - - let mut messages = Vec::new(); - for (role_str, content, tool_calls_json, tool_call_id, _thought) in rows { - let role = match role_str.as_str() { - "system" => MessageRole::System, - "user" => MessageRole::User, - "assistant" => MessageRole::Assistant, - "tool" => MessageRole::Tool, - _ => continue, - }; - - let tool_calls: Option> = tool_calls_json - .and_then(|json_str| serde_json::from_str(&json_str).ok()); - - messages.push(ChatMessage { - role, - content: if content.is_empty() { None } else { Some(content) }, - tool_call_id, - tool_calls, - name: None, - reasoning_content: None, - }); - } - - Ok(messages) - } - - /// 保存消息到数据库 - async fn save_message( - &self, - db: &SqlitePool, - session_id: &str, - turn_index: i32, - step_index: i32, - msg: &ChatMessage, - thought: Option<&str>, - ) -> anyhow::Result<()> { - let role = match msg.role { - MessageRole::System => "system", - MessageRole::User => "user", - MessageRole::Assistant => "assistant", - MessageRole::Tool => "tool", - }; - - let content = msg.content.as_deref().unwrap_or(""); - let tool_calls_json = msg.tool_calls.as_ref() - .map(|tc| serde_json::to_string(tc).unwrap_or_default()); - let tool_call_id = msg.tool_call_id.as_deref(); - let token_count = content.len() as i32 / 4; // 粗略估算 - - sqlx::query( - "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count) \ - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)" - ) - .bind(session_id) - .bind(turn_index) - .bind(step_index) - .bind(role) - .bind(content) - .bind(thought) - .bind(&tool_calls_json) - .bind(tool_call_id) - .bind(token_count) - .execute(db) - .await?; - - Ok(()) - } - - /// 上下文压缩：保留系统提示、最近的 user 消息、以及最近的 tool_calls/tool 对 - async fn compress_context( - &self, - messages: &mut Vec, - llm: &LlmClient, - ) { - if messages.len() <= 4 { - return; - } - - // 保留系统消息 - let system_msg = messages.first().cloned(); - - // 找到安全切割点：必须保证 assistant(tool_calls) 和后续 tool(result) 不被切断 - // 策略：保留最近 6 条消息 + 系统消息 - let keep_count = 6.min(messages.len() - 1); - let to_summarize = &messages[1..messages.len() - keep_count]; - - if to_summarize.is_empty() { - return; - } - - // 生成摘要 - let summary_content: String = to_summarize.iter() - .filter_map(|m| { - let role = match m.role { - MessageRole::User => "用户", - MessageRole::Assistant => "助手", - MessageRole::Tool => "工具", - _ => return None, - }; - m.content.as_ref().map(|c| { - let preview: String = c.chars().take(200).collect(); - format!("[{}] {}", role, preview) - }) - }) - .collect::>() - .join("\n"); - - let summary_prompt = format!( - "请用简洁的中文总结以下对话历史的要点（不超过500字）：\n\n{}", - summary_content - ); - - let summary = match llm.chat_completion( - "你是一个对话摘要助手。请提取对话的关键信息和结论。", - &summary_prompt, - ).await { - Ok(s) => s, - Err(e) => { - warn!("[AgentRuntime] 上下文摘要生成失败: {}，回退为简单截断", e); - format!("[历史摘要] 此前进行了 {} 轮对话交互", to_summarize.len()) - } - }; - - // 重建消息列表 - let recent = messages[messages.len() - keep_count..].to_vec(); - messages.clear(); - if let Some(sys) = system_msg { - messages.push(sys); - } - messages.push(ChatMessage::user(format!("[历史对话摘要]\n{}", summary))); - messages.extend(recent); - - info!("[AgentRuntime] 上下文压缩完成，消息数: {}", messages.len()); - } - - /// 根据用户首条问题生成会话标题 - fn generate_title(&self, question: &str) -> String { - let chars: String = question.chars().take(50).collect(); - if question.len() > 50 { - format!("{}...", chars) - } else { - chars - } - } -} diff --git a/src/agent/runtime/circuit_breaker.rs b/src/agent/runtime/circuit_breaker.rs new file mode 100644 index 0000000..f98f840 --- /dev/null +++ b/src/agent/runtime/circuit_breaker.rs @@ -0,0 +1,214 @@ +// src/agent/runtime/circuit_breaker.rs +// +// 熔断器 — 防止无限自动压缩循环。 +// 参考 Claude Code MAX_CONSECUTIVE_AUTOCOMPACT_FAILURES 设计。 +// +// 当自动压缩连续失败 MAX_CONSECUTIVE_FAILURES 次后，熔断器打开， +// 停止后续压缩尝试，避免无限循环。 + +use std::time::Instant; +use tracing::warn; + +/// 最大连续失败次数，超出后熔断器打开 +const MAX_CONSECUTIVE_FAILURES: usize = 3; + +/// 熔断器打开后，经过此时间自动进入 HalfOpen 状态尝试恢复 +const AUTO_RECOVERY_TIMEOUT_SECS: u64 = 300; // 5 分钟 + +/// 熔断器状态 +#[derive(Debug, Clone, PartialEq)] +pub enum CircuitState { + /// 正常工作，允许压缩 + Closed, + /// 熔断，拒绝后续压缩 + Open, + /// 半开：允许一次试探性压缩以决定是否恢复 + HalfOpen, +} + +/// 压缩熔断器 +#[derive(Debug)] +pub struct CompactionCircuitBreaker { + /// 连续失败计数 + consecutive_failures: usize, + /// 压缩总次数 + total_compactions: usize, + /// 当前状态 + state: CircuitState, + /// 熔断器打开的时间（用于自动恢复） + opened_at: Option, +} + +impl CompactionCircuitBreaker { + /// 创建新的熔断器（初始状态 Closed） + pub fn new() -> Self { + CompactionCircuitBreaker { + consecutive_failures: 0, + total_compactions: 0, + state: CircuitState::Closed, + opened_at: None, + } + } + + /// 记录一次成功的压缩（重置失败计数，关闭熔断器） + pub fn record_success(&mut self) { + self.consecutive_failures = 0; + self.total_compactions += 1; + self.state = CircuitState::Closed; + self.opened_at = None; + } + + /// 记录一次失败的压缩（递增失败计数，可能触发熔断） + pub fn record_failure(&mut self) { + self.consecutive_failures += 1; + self.total_compactions += 1; + + if self.consecutive_failures >= MAX_CONSECUTIVE_FAILURES { + let was_already_open = self.state == CircuitState::Open; + self.state = CircuitState::Open; + self.opened_at = Some(Instant::now()); + if !was_already_open { + warn!( + "[CircuitBreaker] 熔断器打开！连续 {} 次压缩失败，停止自动压缩。\ + {} 秒后将自动尝试恢复。", + self.consecutive_failures, AUTO_RECOVERY_TIMEOUT_SECS + ); + } + } + } + + /// 熔断器是否打开（应停止自动压缩） + pub fn is_open(&self) -> bool { + self.state == CircuitState::Open + } + + /// 是否可以尝试压缩。 + /// + /// 当熔断器打开超过 AUTO_RECOVERY_TIMEOUT_SECS 时，自动转为 HalfOpen 状态， + /// 允许下一次压缩尝试以判断是否恢复。 + pub fn can_attempt(&mut self) -> bool { + match self.state { + CircuitState::Closed | CircuitState::HalfOpen => true, + CircuitState::Open => { + // 检查是否已超时，可自动进入 HalfOpen + if let Some(opened) = self.opened_at { + if opened.elapsed().as_secs() >= AUTO_RECOVERY_TIMEOUT_SECS { + self.state = CircuitState::HalfOpen; + warn!( + "[CircuitBreaker] 熔断器超时，进入 HalfOpen 状态，\ + 允许下一次压缩尝试" + ); + return true; + } + } + false + } + } + } + + /// 重置熔断器到 Closed 状态 + pub fn reset(&mut self) { + self.consecutive_failures = 0; + self.state = CircuitState::Closed; + self.opened_at = None; + } + + /// 获取连续失败次数 + pub fn consecutive_failures(&self) -> usize { + self.consecutive_failures + } + + /// 获取压缩总次数 + pub fn total_compactions(&self) -> usize { + self.total_compactions + } +} + +impl Default for CompactionCircuitBreaker { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_opens_after_max_failures() { + let mut breaker = CompactionCircuitBreaker::new(); + assert!(!breaker.is_open()); + + breaker.record_failure(); + breaker.record_failure(); + assert!(!breaker.is_open()); // 2 failures, not yet open + + breaker.record_failure(); + assert!(breaker.is_open()); // 3 failures, now open + } + + #[test] + fn test_reset_on_success() { + let mut breaker = CompactionCircuitBreaker::new(); + breaker.record_failure(); + breaker.record_failure(); + assert_eq!(breaker.consecutive_failures(), 2); + + breaker.record_success(); + assert_eq!(breaker.consecutive_failures(), 0); + assert!(!breaker.is_open()); + } + + #[test] + fn test_reset_method() { + let mut breaker = CompactionCircuitBreaker::new(); + breaker.record_failure(); + breaker.record_failure(); + breaker.record_failure(); + assert!(breaker.is_open()); + + breaker.reset(); + assert!(!breaker.is_open()); + assert_eq!(breaker.consecutive_failures(), 0); + } + + #[test] + fn test_can_attempt() { + let mut breaker = CompactionCircuitBreaker::new(); + assert!(breaker.can_attempt()); + + for _ in 0..3 { + breaker.record_failure(); + } + // 刚打开，不应允许尝试 + assert!(!breaker.can_attempt()); + } + + #[test] + fn test_half_open_after_reset_or_success() { + let mut breaker = CompactionCircuitBreaker::new(); + // 触发熔断 + for _ in 0..3 { + breaker.record_failure(); + } + assert!(breaker.is_open()); + + // success 直接重置到 Closed + breaker.record_success(); + assert!(!breaker.is_open()); + assert!(breaker.can_attempt()); + } + + #[test] + fn test_record_failure_while_open_stays_open() { + let mut breaker = CompactionCircuitBreaker::new(); + for _ in 0..3 { + breaker.record_failure(); + } + assert!(breaker.is_open()); + // 熔断器打开后再次失败，保持 Open + breaker.record_failure(); + assert!(breaker.is_open()); + assert_eq!(breaker.consecutive_failures(), 4); + } +} diff --git a/src/agent/runtime/context.rs b/src/agent/runtime/context.rs new file mode 100644 index 0000000..6c36c16 --- /dev/null +++ b/src/agent/runtime/context.rs @@ -0,0 +1,85 @@ +// src/agent/runtime/context.rs +// +// 上下文构建：加载历史消息、注入系统提示词、添加用户消息、 +// 从数据库恢复持久化的任务状态。 + +use sqlx::SqlitePool; +use tracing::info; + +use crate::clients::llm::{ChatMessage, MessageRole}; + +use super::session; + +/// 构建初始 LLM 上下文：加载历史 → 插入系统提示词 → 添加用户消息 → 恢复任务状态。 +pub async fn build_initial_context( + db: &SqlitePool, + session_id: &str, + system_prompt: &str, + question: &str, + _turn_index: i32, +) -> anyhow::Result> { + let mut messages = session::load_history_for_llm(db, session_id).await?; + + // 注入系统提示词（如果历史中没有） + if messages.is_empty() || messages[0].role != MessageRole::System { + messages.insert(0, ChatMessage::system(system_prompt)); + } + + // 添加用户消息 + messages.push(ChatMessage::user(question)); + + // 从数据库恢复持久化的任务状态 + if let Some(task_reminder) = restore_tasks_from_db(db, session_id).await { + messages.push(ChatMessage::user(task_reminder)); + } + + Ok(messages) +} + +/// 从 agent_tasks 表恢复任务状态，返回格式化的提醒文本。 +/// +/// 如果表不存在或没有任务记录，返回 None。 +async fn restore_tasks_from_db(db: &SqlitePool, session_id: &str) -> Option { + let rows: Vec<(String, String, String, String, Option)> = sqlx::query_as( + "SELECT task_id, content, status, blocked_by, owner \ + FROM agent_tasks WHERE session_id = ? AND (owner = '' OR owner = 'lead') \ + ORDER BY created_at ASC", + ) + .bind(session_id) + .fetch_all(db) + .await + .ok()?; + + if rows.is_empty() { + return None; + } + + let mut lines: Vec = Vec::new(); + for (task_id, content, status, blocked_by, owner) in &rows { + let icon = match status.as_str() { + "in_progress" => "🔄", + "completed" => "✅", + _ => "⏳", + }; + + let blocked: Vec = serde_json::from_str(blocked_by).unwrap_or_default(); + let mut line = format!("{} [{}] {}", icon, task_id, content); + if !blocked.is_empty() { + line.push_str(&format!(" (依赖: {})", blocked.join(", "))); + } + if let Some(o) = owner { + if !o.is_empty() && o != "lead" { + line.push_str(&format!(" (指派: {})", o)); + } + } + lines.push(line); + } + + info!("[Context] 从数据库恢复了 {} 个任务状态", rows.len()); + + Some(format!( + "[当前任务状态]\n以下是上次会话中持久化的任务计划，请基于最新状态继续工作：\n\n{}\n\n\ + 使用 todo_write 工具更新任务进度。", + lines.join("\n") + )) +} diff --git a/src/agent/runtime/error_recovery.rs b/src/agent/runtime/error_recovery.rs new file mode 100644 index 0000000..5b6396f --- /dev/null +++ b/src/agent/runtime/error_recovery.rs @@ -0,0 +1,428 @@ +// src/agent/runtime/error_recovery.rs +// +// 错误恢复阶梯。 +// 参考 Claude Code error recovery ladder 设计。 +// +// 当 LLM 流返回可恢复的错误（如 prompt_too_long）时， +// 按阶梯顺序尝试恢复： +// 1. Aggressive Compact — 激进微压缩（保留更少的工具结果） +// 2. Reactive Compact — 使用 LLM 摘要压缩对话历史 +// 3. Escalate Tokens — 临时提升 token 上限到 64k +// 4. Multi-Turn — 注入 metacognitive 消息分步处理 +// 5. Surface — 放弃恢复，暴露错误给用户 +// +// 每一步都有 `has_attempted` 守卫，防止无限循环。 + +use tracing::{info, warn}; + +use super::token_budget::TokenBudget; + +/// 错误类型分类 +#[derive(Debug, Clone, PartialEq)] +pub enum ErrorKind { + /// 上下文过长（prompt too long / 413） + PromptTooLong, + /// Token 耗尽 + TokenExhausted, + /// 模型错误 + ModelError(String), + /// 超时 + Timeout, + /// 限流（HTTP 429） + RateLimited, + /// 服务过载（HTTP 529） + Overloaded, +} + +/// 恢复步骤 +#[derive(Debug, Clone, PartialEq)] +pub enum RecoveryStep { + /// 尝试更激进的 micro_compact + AggressiveCompact, + /// 使用 LLM 摘要压缩 + ReactiveCompact, + /// 提升 token 上限 + EscalateTokens { new_hard_limit: usize }, + /// 分轮恢复（注入 meta 消息） + MultiTurn, + /// 放弃，暴露错误 + Surface, + /// 指数退避重试（用于 429/529 瞬态错误） + RetryWithBackoff { attempt: u32, delay_ms: u64 }, +} + +/// 恢复尝试追踪 +#[derive(Debug, Clone)] +pub struct RecoveryAttempts { + pub aggressive_compact: bool, + pub reactive_compact: bool, + pub escalate_tokens: bool, + pub multi_turn: bool, +} + +impl RecoveryAttempts { + pub fn new() -> Self { + RecoveryAttempts { + aggressive_compact: false, + reactive_compact: false, + escalate_tokens: false, + multi_turn: false, + } + } + + /// 是否有未尝试的恢复步骤 + pub fn has_remaining(&self) -> bool { + !self.aggressive_compact + || !self.reactive_compact + || !self.escalate_tokens + || !self.multi_turn + } + + /// 获取下一个应尝试的恢复步骤 + pub fn next_step(&mut self, error_kind: &ErrorKind) -> Option { + // 429/529 使用退避重试，不消耗上下文恢复步骤 + if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) { + return Some(RecoveryStep::RetryWithBackoff { + attempt: 0, + delay_ms: 500, + }); + } + + match error_kind { + ErrorKind::PromptTooLong | ErrorKind::TokenExhausted => { + if !self.aggressive_compact { + self.aggressive_compact = true; + return Some(RecoveryStep::AggressiveCompact); + } + if !self.reactive_compact { + self.reactive_compact = true; + return Some(RecoveryStep::ReactiveCompact); + } + if !self.escalate_tokens { + self.escalate_tokens = true; + return Some(RecoveryStep::EscalateTokens { + new_hard_limit: 64_000, + }); + } + if !self.multi_turn { + self.multi_turn = true; + return Some(RecoveryStep::MultiTurn); + } + } + _ => { + // 非上下文相关错误，直接暴露 + if !self.multi_turn { + self.multi_turn = true; + return Some(RecoveryStep::Surface); + } + } + } + None + } +} + +impl Default for RecoveryAttempts { + fn default() -> Self { + Self::new() + } +} + +/// 错误恢复器 +pub struct ErrorRecovery { + /// 恢复步骤追踪 + pub attempts: RecoveryAttempts, + /// Token 预算（用于 escalate 步骤） + pub token_budget: TokenBudget, +} + +impl ErrorRecovery { + /// 创建新的错误恢复器 + pub fn new(token_budget: TokenBudget) -> Self { + ErrorRecovery { + attempts: RecoveryAttempts::new(), + token_budget, + } + } + + /// 尝试从错误中恢复。 + /// + /// 返回 `Some(RecoveryStep)` 表示找到了恢复步骤（调用方应执行该步骤后重试）。 + /// 返回 `None` 表示所有步骤已尝试完毕，应暴露错误给用户。 + pub fn try_recover(&mut self, error_kind: &ErrorKind) -> Option { + let step = self.attempts.next_step(error_kind); + match &step { + Some(RecoveryStep::AggressiveCompact) => { + info!("[ErrorRecovery] 尝试步骤 1/4: AggressiveCompact"); + } + Some(RecoveryStep::ReactiveCompact) => { + info!("[ErrorRecovery] 尝试步骤 2/4: ReactiveCompact"); + } + Some(RecoveryStep::EscalateTokens { new_hard_limit }) => { + info!( + "[ErrorRecovery] 尝试步骤 3/4: EscalateTokens → {}", + new_hard_limit + ); + self.token_budget.escalate_hard_limit(*new_hard_limit); + } + Some(RecoveryStep::RetryWithBackoff { attempt, delay_ms }) => { + info!( + "[ErrorRecovery] 退避重试: attempt={}, delay={}ms", + attempt, delay_ms + ); + } + Some(RecoveryStep::MultiTurn) => { + info!("[ErrorRecovery] 尝试步骤 4/4: MultiTurn"); + } + Some(RecoveryStep::Surface) => { + warn!("[ErrorRecovery] 无法恢复，暴露错误"); + } + None => { + warn!("[ErrorRecovery] 所有恢复步骤已尝试完毕"); + } + } + step + } + + /// 生成 multi-turn 恢复消息（注入到对话中以继续处理） + pub fn multi_turn_message() -> String { + "由于 token 限制，当前回答被截断。请基于已收集的信息继续分析，\ + 重点关注尚未完成的部分。你可以：\n\ + 1. 总结已有发现\n\ + 2. 使用 compress_context 手动压缩上下文\n\ + 3. 分步完成剩余工作" + .to_string() + } + + /// 检查是否需要恢复（错误是否可恢复） + pub fn is_recoverable(error_kind: &ErrorKind) -> bool { + matches!( + error_kind, + ErrorKind::PromptTooLong + | ErrorKind::TokenExhausted + | ErrorKind::RateLimited + | ErrorKind::Overloaded + ) + } +} + +/// 计算指数退避延迟（毫秒）。 +/// +/// 公式：min(500 * 2^attempt, 32000) + 25% 随机抖动 +/// 如果有 Retry-After header，优先使用。 +pub fn backoff_delay(attempt: u32, retry_after_secs: Option) -> u64 { + if let Some(ra) = retry_after_secs { + return ra * 1000; + } + let base = 500u64 * 2u64.pow(attempt.min(6)); // cap at 2^6 = 64 → 32000ms + let base = base.min(32_000); + // Simple deterministic jitter using attempt (avoid rand dependency) + let jitter = (base / 4) * (attempt as u64 % 5) / 5; + base + jitter +} + +/// 解析错误字符串中的 Retry-After header 值。 +/// +/// 期望格式: `retry_after=Some(N)` 出现在错误消息中。 +pub fn parse_retry_after(error_str: &str) -> Option { + if let Some(pos) = error_str.find("retry_after=Some(") { + let prefix_len = "retry_after=Some(".len(); // 19 + let rest = &error_str[pos + prefix_len..]; + if let Some(end) = rest.find(')') { + return rest[..end].parse().ok(); + } + } + None +} + +/// 从错误字符串分类错误类型。 +/// 解析 LLM API 返回的错误消息，映射到 ErrorKind。 +pub fn classify_error(error_str: &str) -> ErrorKind { + let lower = error_str.to_lowercase(); + + // 先检测限流/过载（HTTP 状态码检查） + if lower.contains("429") + || lower.contains("rate limit") + || lower.contains("rate_limit") + || lower.contains("too many requests") + { + return ErrorKind::RateLimited; + } + if lower.contains("529") + || lower.contains("overloaded") + || lower.contains("overload") + || lower.contains("service overloaded") + { + return ErrorKind::Overloaded; + } + + if lower.contains("prompt_too_long") + || lower.contains("prompt too long") + || lower.contains("context length") + || lower.contains("413") + || lower.contains("context_window_exceeded") + || lower.contains("input length") + { + return ErrorKind::PromptTooLong; + } + + if lower.contains("max_tokens") + || lower.contains("token limit") + || lower.contains("token_exhausted") + || lower.contains("maximum context length") + || lower.contains("reduce the length") + { + return ErrorKind::TokenExhausted; + } + + if lower.contains("timeout") + || lower.contains("timed out") + || lower.contains("deadline exceeded") + || lower.contains("408") + || lower.contains("504") + { + return ErrorKind::Timeout; + } + + // 默认归类为模型错误 + ErrorKind::ModelError(error_str.to_string()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_all_steps_sequence() { + let mut attempts = RecoveryAttempts::new(); + assert_eq!( + attempts.next_step(&ErrorKind::PromptTooLong), + Some(RecoveryStep::AggressiveCompact) + ); + assert_eq!( + attempts.next_step(&ErrorKind::PromptTooLong), + Some(RecoveryStep::ReactiveCompact) + ); + assert_eq!( + attempts.next_step(&ErrorKind::PromptTooLong), + Some(RecoveryStep::EscalateTokens { + new_hard_limit: 64_000 + }) + ); + assert_eq!( + attempts.next_step(&ErrorKind::PromptTooLong), + Some(RecoveryStep::MultiTurn) + ); + // 所有步骤已尝试 + assert_eq!(attempts.next_step(&ErrorKind::PromptTooLong), None); + } + + #[test] + fn test_model_error_goes_straight_to_surface() { + let mut attempts = RecoveryAttempts::new(); + assert_eq!( + attempts.next_step(&ErrorKind::ModelError("test".into())), + Some(RecoveryStep::Surface) + ); + assert_eq!( + attempts.next_step(&ErrorKind::ModelError("test".into())), + None + ); + } + + #[test] + fn test_has_remaining() { + let mut attempts = RecoveryAttempts::new(); + assert!(attempts.has_remaining()); + + // 消耗所有步骤 + for _ in 0..4 { + attempts.next_step(&ErrorKind::PromptTooLong); + } + assert!(!attempts.has_remaining()); + } + + #[test] + fn test_is_recoverable() { + assert!(ErrorRecovery::is_recoverable(&ErrorKind::PromptTooLong)); + assert!(ErrorRecovery::is_recoverable(&ErrorKind::TokenExhausted)); + assert!(!ErrorRecovery::is_recoverable(&ErrorKind::ModelError( + "test".into() + ))); + assert!(!ErrorRecovery::is_recoverable(&ErrorKind::Timeout)); + } + + #[test] + fn test_classify_rate_limited_429() { + let kind = classify_error("HTTP 429: Too Many Requests"); + assert_eq!(kind, ErrorKind::RateLimited); + } + + #[test] + fn test_classify_overloaded_529() { + let kind = classify_error("HTTP 529: Service Overloaded"); + assert_eq!(kind, ErrorKind::Overloaded); + } + + #[test] + fn test_rate_limited_is_recoverable() { + assert!(ErrorRecovery::is_recoverable(&ErrorKind::RateLimited)); + assert!(ErrorRecovery::is_recoverable(&ErrorKind::Overloaded)); + } + + #[test] + fn test_parse_retry_after() { + let err = "HTTP 429: retry_after=Some(30)"; + assert_eq!(parse_retry_after(err), Some(30)); + } + + #[test] + fn test_parse_retry_after_none() { + let err = "HTTP 500: Internal Server Error"; + assert_eq!(parse_retry_after(err), None); + } + + #[test] + fn test_backoff_delay() { + // Attempt 0: 500 + jitter + let d0 = backoff_delay(0, None); + assert!(d0 >= 500 && d0 <= 700); + + // Attempt 3: 500*8=4000 + jitter + let d3 = backoff_delay(3, None); + assert!(d3 >= 4000 && d3 <= 5000); + + // Capped at 32s + let d10 = backoff_delay(10, None); + assert!(d10 <= 40_000); + + // Retry-After takes priority + let d_ra = backoff_delay(0, Some(15)); + assert_eq!(d_ra, 15000); + } + + #[test] + fn test_rate_limited_goes_to_retry() { + let mut attempts = RecoveryAttempts::new(); + let step = attempts.next_step(&ErrorKind::RateLimited); + assert!(matches!(step, Some(RecoveryStep::RetryWithBackoff { .. }))); + } + + #[test] + fn test_token_budget_escalation() { + let budget = TokenBudget::new(32_000, 40_000); + let mut recovery = ErrorRecovery::new(budget); + assert_eq!(recovery.token_budget.hard_limit, 40_000); + + recovery.attempts.next_step(&ErrorKind::PromptTooLong); // aggressive + recovery.attempts.next_step(&ErrorKind::PromptTooLong); // reactive + let step = recovery.try_recover(&ErrorKind::PromptTooLong); // escalate + + assert_eq!( + step, + Some(RecoveryStep::EscalateTokens { + new_hard_limit: 64_000 + }) + ); + assert_eq!(recovery.token_budget.hard_limit, 64_000); + } +} diff --git a/src/agent/runtime/executor.rs b/src/agent/runtime/executor.rs new file mode 100644 index 0000000..fb1310d --- /dev/null +++ b/src/agent/runtime/executor.rs @@ -0,0 +1,370 @@ +// src/agent/runtime/executor.rs +// +// 工具调用执行器：验证 → PreToolUse hooks → 并行执行 → 结果收集 → PostToolUse hooks。 + +use futures_util::stream::FuturesUnordered; +use futures_util::StreamExt; +use sqlx::SqlitePool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing::warn; + +use crate::api::AppState; +use crate::clients::llm::{ChatMessage, ToolCall}; + +use super::file_cache::FileStateCache; +use super::permission::PermissionChecker; +use super::{AgentStreamEvent, DuplicateDetector}; +use crate::agent::hooks::{HookRegistry, PostToolUseContext, PreToolUseContext}; +use crate::agent::tools::persist::maybe_persist_tool_result; +use crate::agent::tools::{InterruptBehavior, ToolContext, ToolOutput, ToolRegistry}; + +/// 准备好的工具调用 +#[derive(Debug, Clone)] +pub struct PreparedCall { + pub tool_call_id: String, + pub tool_name: String, + pub args: serde_json::Value, +} + +/// 单次工具执行后的消息 + 元数据 +pub struct ToolResultMessage { + pub chat_message: ChatMessage, + pub was_error: bool, +} + +/// 工具执行结果摘要 +pub struct ToolExecutionResult { + /// 每条工具调用对应的 tool_result 消息（供调用方 push 到 messages） + pub tool_messages: Vec, + pub was_cancelled: bool, + pub had_duplicate: bool, +} + +/// 验证工具调用：死循环检测 + 参数解析。 +/// +/// 返回 (prepared_calls, has_duplicate)。 +/// 死循环或参数无效时，错误消息直接注入到 messages。 +#[allow(clippy::too_many_arguments)] +pub fn validate_and_prepare( + tool_calls: &[ToolCall], + duplicate_detector: &mut DuplicateDetector, + duplicate_threshold: usize, + messages: &mut Vec, + tx: &mpsc::UnboundedSender, + db: &SqlitePool, + session_id: &str, + turn_index: i32, + step: usize, +) -> (Vec, bool) { + let mut prepared_calls: Vec = Vec::new(); + let mut has_duplicate = false; + + for tool_call in tool_calls { + let tool_name = &tool_call.function.name; + let tool_args_str = &tool_call.function.arguments; + + // 死循环检测 + if duplicate_detector.record(tool_name, tool_args_str, duplicate_threshold) { + warn!( + "[Executor] 检测到死循环：{} 连续调用 {} 次", + tool_name, duplicate_threshold + ); + let _ = tx.send(AgentStreamEvent::Error { + message: format!("检测到工具 {} 的重复调用，已自动终止循环。", tool_name), + }); + let error_msg = ChatMessage::tool_result( + &tool_call.id, + format!( + "错误：工具 {} 被连续重复调用 {} 次，参数完全相同。\ + 请停止重复调用并直接给出目前收集到的答案。", + tool_name, duplicate_threshold + ), + ); + messages.push(error_msg); + has_duplicate = true; + continue; + } + + // 解析参数 + let args: serde_json::Value = match serde_json::from_str(tool_args_str) { + Ok(v) => v, + Err(e) => { + let error_output = format!("工具参数 JSON 解析失败: {}", e); + let _ = tx.send(AgentStreamEvent::ToolResult { + name: tool_name.clone(), + output: error_output.clone(), + is_error: true, + metadata: serde_json::json!({}), + step, + }); + let tool_msg = ChatMessage::tool_result(&tool_call.id, &error_output); + save_tool_message_sync(db, session_id, turn_index, step, &tool_msg); + messages.push(tool_msg); + continue; + } + }; + + prepared_calls.push(PreparedCall { + tool_call_id: tool_call.id.clone(), + tool_name: tool_name.clone(), + args, + }); + } + + (prepared_calls, has_duplicate) +} + +/// 并行执行所有准备好的工具调用。 +/// +/// 流程： +/// 1. 权限检查（deny 规则阻止不可执行工具） +/// 2. 发送 ToolCall SSE 事件 +/// 3. 运行 PreToolUse hooks +/// 4. 工具分区 + 并行执行（并发安全工具一批并行，不安全工具单独串行） +/// 5. 收集结果、发送 ToolResult SSE、运行 PostToolUse hooks +/// 6. 返回 ToolResultMessage 列表供调用方推入 messages +#[allow(clippy::too_many_arguments)] +pub async fn execute_parallel( + prepared_calls: &[PreparedCall], + tool_registry: &ToolRegistry, + app_state: Arc, + hook_registry: &HookRegistry, + _permission_checker: Option<&PermissionChecker>, + tx: &mpsc::UnboundedSender, + db: &SqlitePool, + session_id: &str, + agent_name: &str, + turn_index: i32, + step: usize, + tool_timeout_secs: u64, + max_output_chars: usize, + read_file_state: Arc>, +) -> ToolExecutionResult { + if prepared_calls.is_empty() { + return ToolExecutionResult { + tool_messages: Vec::new(), + was_cancelled: false, + had_duplicate: false, + }; + } + + let sid = session_id.to_string(); + + // Phase 1: 发送 ToolCall SSE 事件 + for prep in prepared_calls { + let _ = tx.send(AgentStreamEvent::ToolCall { + name: prep.tool_name.clone(), + arguments: prep.args.clone(), + step, + }); + } + + // Phase 2: PreToolUse hooks — 收集修改后的参数和附加上下文 + let exec_start = std::time::Instant::now(); + let mut mutated_args: Vec = Vec::new(); + let mut additional_contexts: Vec = Vec::new(); + for prep in prepared_calls { + let hook_ctx = PreToolUseContext { + session_id: sid.clone(), + tool_name: prep.tool_name.clone(), + tool_args: prep.args.clone(), + step, + }; + let result = hook_registry.run_pre_tool_use(&hook_ctx).await; + if result.action.is_blocked() { + let reason = result.action.block_reason().unwrap_or("unknown"); + warn!( + "[Executor] PreToolUse hook 阻止了 {} 的执行: {}", + prep.tool_name, reason + ); + } + // 使用 hook 可能修改后的参数 + mutated_args.push(result.final_args); + if let Some(ctx) = result.additional_context { + additional_contexts.push(ctx); + } + } + + // Phase 3: 并行执行 + let cancelled = Arc::new(AtomicBool::new(false)); + let cancel_flag = cancelled.clone(); + let app_state_ref = app_state.clone(); + let sid_ref = sid.clone(); + + let cancel_handle = tokio::spawn(async move { + loop { + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + if let Ok(locked) = app_state_ref.cancelled_runs.lock() { + if locked.contains(&sid_ref) { + cancel_flag.store(true, Ordering::SeqCst); + return; + } + } + } + }); + + let timeout_dur = std::time::Duration::from_secs(tool_timeout_secs); + + // Phase 3: 使用 FuturesUnordered 进行渐进式并行执行。 + // 每个工具完成后立即发送 SSE ToolResult 事件到前端（非阻塞）， + // 而后台继续等待其他工具完成。快工具的结果不会因慢工具而延迟。 + let mut exec_futs: FuturesUnordered<_> = prepared_calls + .iter() + .enumerate() + .map(|(i, prep)| { + let tool_name = prep.tool_name.clone(); + let args = mutated_args + .get(i) + .cloned() + .unwrap_or_else(|| prep.args.clone()); + let tool_ctx = ToolContext::with_file_cache(app_state.clone(), read_file_state.clone()); + let cancelled = cancelled.clone(); + let tool_opt = tool_registry.get(&tool_name); + + Box::pin(async move { + let output = match tool_opt { + Some(tool) => { + let interrupt_behavior = tool.interrupt_behavior(); + let is_blocking = interrupt_behavior == InterruptBehavior::Block; + + let tool_fut = tool.execute(args, &tool_ctx); + let cancel_fut = async { + loop { + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + if !is_blocking && cancelled.load(Ordering::SeqCst) { + return; + } + } + }; + + tokio::select! { + res = tokio::time::timeout(timeout_dur, tool_fut) => { + match res { + Ok(output) => output, + Err(_) => ToolOutput::error(format!( + "工具 {} 执行超时（{}秒）", + tool_name, + timeout_dur.as_secs() + )), + } + } + _ = cancel_fut => { + ToolOutput::error("执行已被用户取消") + } + } + } + None => ToolOutput::error(format!("未知工具: {}", tool_name)), + }; + + let was_cancelled = cancelled.load(Ordering::SeqCst); + ( + prep.tool_call_id.clone(), + prep.tool_name.clone(), + prep.args.clone(), + output, + was_cancelled, + ) + }) + }) + .collect(); + + let mut tool_messages: Vec = Vec::new(); + let mut was_cancelled = false; + + // 渐进式处理结果：每个工具一完成就立即处理（SSE 事件 + PostToolUse hook + 持久化） + while let Some((tool_call_id, tool_name, tool_args, output, cancelled_flag)) = + exec_futs.next().await + { + if cancelled_flag { + was_cancelled = true; + } + + let elapsed_ms = exec_start.elapsed().as_millis() as u64; + + // SSE 事件 — 立即推送到前端 + let _ = tx.send(AgentStreamEvent::ToolResult { + name: tool_name.clone(), + output: output.content.clone(), + is_error: output.is_error, + metadata: output.metadata.clone(), + step, + }); + + // 输出处理：小结果直接传递，大结果持久化到磁盘并返回 stub + let tool_results_dir = app_state.config.library_dir.join("tool-results"); + let (processed_content, _persisted_path) = maybe_persist_tool_result( + &output.content, + &tool_call_id, + max_output_chars, + &tool_results_dir, + ); + + // PostToolUse hook + let post_ctx = PostToolUseContext { + session_id: sid.clone(), + agent_name: agent_name.to_string(), + tool_name: tool_name.clone(), + tool_args, + output_content: processed_content.clone(), + is_error: output.is_error, + step, + elapsed_ms, + }; + let post_result = hook_registry.run_post_tool_use(&post_ctx).await; + let final_content = post_result.final_content; + + let chat_message = ChatMessage::tool_result(&tool_call_id, &final_content); + + // 持久化到数据库（fire-and-forget） + save_tool_message_sync(db, &sid, turn_index, step, &chat_message); + + tool_messages.push(ToolResultMessage { + chat_message, + was_error: output.is_error, + }); + } + + cancel_handle.abort(); + + ToolExecutionResult { + tool_messages, + was_cancelled, + had_duplicate: false, + } +} + +/// 同步保存 tool 角色消息到数据库。 +fn save_tool_message_sync( + db: &SqlitePool, + session_id: &str, + turn_index: i32, + step_index: usize, + msg: &ChatMessage, +) { + let db_clone = db.clone(); + let session_id = session_id.to_string(); + let content = msg.content.as_deref().unwrap_or("").to_string(); + let tool_call_id = msg.tool_call_id.clone(); + // fire-and-forget: tool 消息保存失败不影响主流程 + tokio::spawn(async move { + let token_count = content.len() as i32 / 4; + if let Err(e) = sqlx::query( + "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, tool_call_id, token_count, agent_name) \ + VALUES (?, ?, ?, 'tool', ?, ?, ?, ?)", + ) + .bind(&session_id) + .bind(turn_index) + .bind(step_index as i32) + .bind(&content) + .bind(&tool_call_id) + .bind(token_count) + .bind("lead") + .execute(&db_clone) + .await + { + warn!("[Executor] 保存 tool 消息失败（非致命）: {}", e); + } + }); +} diff --git a/src/agent/runtime/file_cache.rs b/src/agent/runtime/file_cache.rs new file mode 100644 index 0000000..0b1cc05 --- /dev/null +++ b/src/agent/runtime/file_cache.rs @@ -0,0 +1,571 @@ +// src/agent/runtime/file_cache.rs +// +// 文件状态缓存 — 参考 Claude Code FileStateCache 设计。 +// +// 在 Read 工具调用前检查缓存： +// 1. 路径已缓存 → 读取磁盘 mtime → mtime 相同 + offset/limit 一致 → 返回 stub +// 2. mtime 不同或新文件 → 正常读取 → 写入缓存 +// +// 压缩时： +// - 压缩前：快照缓存到普通对象 +// - 压缩后：清空缓存，将最近 N 个文件作为上下文注入 +// +// 缓存上限：100 个条目，25MB 内容总大小（LRU 自动淘汰）。 + +use lru::LruCache; +use std::num::NonZeroUsize; +use std::path::Path; +use tracing::{info, warn}; + +/// 缓存条目最大数量 +pub const MAX_ENTRIES: usize = 100; +/// 缓存内容总大小上限（25MB） +pub const MAX_CACHE_SIZE_BYTES: usize = 25 * 1024 * 1024; + +/// 文件不变时的占位消息（参考 Claude Code FILE_UNCHANGED_STUB） +pub const FILE_UNCHANGED_STUB: &str = + "File unchanged since last read. The content from the earlier read_file tool_result \ + in this conversation is still current — refer to that instead of re-reading."; + +/// 压缩后恢复的最大文件数 +pub const POST_COMPACT_MAX_FILES_TO_RESTORE: usize = 5; +/// 压缩后恢复的每文件最大 token 数（~字符数） +pub const POST_COMPACT_MAX_CHARS_PER_FILE: usize = 4_000; + +/// 单个文件的缓存状态 +#[derive(Debug, Clone)] +pub struct FileState { + /// 上次读取的文件内容 + pub content: String, + /// 文件修改时间（Unix 时间戳，秒级） + pub timestamp: i64, + /// 读取起始行（1-based） + pub offset: usize, + /// 行数限制 + pub limit: Option, +} + +/// 文件状态快照（用于压缩前后传递，纯数据，不含 LRU 结构） +pub type FileStateSnapshot = Vec<(String, FileState)>; + +/// 文件状态缓存。 +/// +/// 包装 `LruCache` + 内容总大小追踪。 +/// 通过 `Arc>` 在工具调用间共享。 +pub struct FileStateCache { + cache: LruCache, + /// 当前缓存中所有内容的字节数总和（近似，使用 content.len()） + current_size_bytes: usize, + /// 最大字节数 + max_size_bytes: usize, +} + +impl FileStateCache { + /// 创建新的缓存实例 + pub fn new() -> Self { + let max_entries = NonZeroUsize::new(MAX_ENTRIES).unwrap(); + FileStateCache { + cache: LruCache::new(max_entries), + current_size_bytes: 0, + max_size_bytes: MAX_CACHE_SIZE_BYTES, + } + } + + /// 创建带自定义参数的新缓存 + pub fn with_limits(max_entries: usize, max_size_bytes: usize) -> Self { + let max_entries = + NonZeroUsize::new(max_entries).unwrap_or(NonZeroUsize::new(MAX_ENTRIES).unwrap()); + FileStateCache { + cache: LruCache::new(max_entries), + current_size_bytes: 0, + max_size_bytes, + } + } + + /// 规范化路径 key（确保一致性） + fn normalize_key(path: &str) -> String { + // 去除尾随斜杠，规范化重复斜杠 + let p = Path::new(path); + // 尝试 canonicalize（跟随符号链接），失败则用简单的字符串规范化 + match p.canonicalize() { + Ok(canon) => canon.to_string_lossy().to_string(), + Err(_) => { + // 简单规范化：折叠重复的 / + let mut result = String::with_capacity(path.len()); + let mut prev_slash = false; + for ch in path.chars() { + if ch == '/' || ch == '\\' { + if !prev_slash { + result.push('/'); + prev_slash = true; + } + } else { + result.push(ch); + prev_slash = false; + } + } + // 去除尾随 / + if result.ends_with('/') && result.len() > 1 { + result.pop(); + } + result + } + } + } + + /// 获取缓存的条目，返回 Some(&FileState) 若存在 + pub fn get(&mut self, path: &str) -> Option<&FileState> { + let key = Self::normalize_key(path); + self.cache.get(&key) + } + + /// 写入缓存条目，自动处理容量限制 + pub fn set(&mut self, path: &str, state: FileState) { + let key = Self::normalize_key(path); + let content_len = state.content.len(); + + // 如果 key 已存在，先减去旧内容的 size + if let Some(old) = self.cache.get(&key) { + self.current_size_bytes = self.current_size_bytes.saturating_sub(old.content.len()); + } + + // 驱逐旧条目直到有足够空间 + while self.current_size_bytes + content_len > self.max_size_bytes && !self.cache.is_empty() { + if let Some((_, evicted)) = self.cache.pop_lru() { + self.current_size_bytes = self + .current_size_bytes + .saturating_sub(evicted.content.len()); + } + } + + // 如果单个文件超过上限，仍存储但记录警告 + if content_len > self.max_size_bytes { + warn!( + "[FileCache] 单个文件内容 ({} bytes) 超过缓存上限 ({} bytes)", + content_len, self.max_size_bytes + ); + } + + self.current_size_bytes += content_len; + self.cache.push(key, state); + + info!( + "[FileCache] 缓存写入: path={}, size={} bytes, cache_entries={}, cache_size={}", + path, + content_len, + self.cache.len(), + self.current_size_bytes + ); + } + + /// 检查 key 是否存在 + pub fn contains(&mut self, path: &str) -> bool { + let key = Self::normalize_key(path); + self.cache.contains(&key) + } + + /// 删除缓存条目 + pub fn remove(&mut self, path: &str) -> bool { + let key = Self::normalize_key(path); + if let Some(removed) = self.cache.pop(&key) { + self.current_size_bytes = self + .current_size_bytes + .saturating_sub(removed.content.len()); + true + } else { + false + } + } + + /// 清空缓存 + pub fn clear(&mut self) { + self.cache.clear(); + self.current_size_bytes = 0; + info!("[FileCache] 缓存已清空"); + } + + /// 缓存条目数 + pub fn len(&self) -> usize { + self.cache.len() + } + + /// 缓存是否为空 + pub fn is_empty(&self) -> bool { + self.cache.len() == 0 + } + + /// 当前缓存内容总大小（近似字节数） + pub fn current_size_bytes(&self) -> usize { + self.current_size_bytes + } + + // ── 压缩集成 ── + + /// 生成快照（纯数据，不含 LRU 结构）。 + /// 在压缩前调用，用于压缩后恢复文件上下文。 + pub fn to_snapshot(&mut self) -> FileStateSnapshot { + // 按 timestamp 降序排序（最近读的在前） + let mut entries: Vec<(String, FileState)> = self + .cache + .iter() + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + entries.sort_by(|a, b| b.1.timestamp.cmp(&a.1.timestamp)); + entries + } + + /// 从快照恢复指定数量的最近文件到缓存。 + /// 在压缩后调用。 + pub fn restore_from_snapshot(&mut self, snapshot: &FileStateSnapshot, max_files: usize) { + for (path, state) in snapshot.iter().take(max_files) { + // 不恢复过大的文件（已有提示说可能过时） + if state.content.len() > POST_COMPACT_MAX_CHARS_PER_FILE { + continue; + } + self.set(path, state.clone()); + } + info!( + "[FileCache] 从快照恢复了 {} 个文件 (快照大小: {})", + self.len().min(max_files), + snapshot.len() + ); + } + + /// 从快照生成上 下文注入文本（用于压缩后注入到对话中）。 + /// 返回格式化的 markdown 块列表。 + pub fn build_restore_context(snapshot: &FileStateSnapshot, max_files: usize) -> Vec { + if snapshot.is_empty() { + return Vec::new(); + } + + let mut contexts: Vec = Vec::new(); + let mut used_chars = 0usize; + let total_budget = POST_COMPACT_MAX_CHARS_PER_FILE * max_files; + + for (path, state) in snapshot.iter().take(max_files) { + let preview: String = state + .content + .chars() + .take(POST_COMPACT_MAX_CHARS_PER_FILE) + .collect(); + let truncated = if state.content.len() > preview.len() { + format!( + "{}…\n[内容已截断: {} 字符 → {} 字符]", + preview, + state.content.len(), + preview.len() + ) + } else { + preview + }; + + if used_chars + truncated.len() > total_budget { + break; + } + + let block = format!( + "[压缩后恢复: {}]\n上次读取时间戳: {}\n内容:\n```\n{}\n```", + path, state.timestamp, truncated + ); + used_chars += block.len(); + contexts.push(block); + } + + if !contexts.is_empty() { + info!( + "[FileCache] 生成压缩恢复上下文: {} 文件, {} chars", + contexts.len(), + used_chars + ); + } + + contexts + } +} + +impl Default for FileStateCache { + fn default() -> Self { + Self::new() + } +} + +/// 获取文件的当前 mtime（Unix 时间戳，秒）。 +/// 失败时返回 None。 +pub fn get_file_mtime(path: &str) -> Option { + match std::fs::metadata(path) { + Ok(meta) => match meta.modified() { + Ok(time) => match time.duration_since(std::time::UNIX_EPOCH) { + Ok(d) => Some(d.as_secs() as i64), + Err(_) => None, + }, + Err(_) => None, + }, + Err(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_cache_is_empty() { + let cache = FileStateCache::new(); + assert!(cache.is_empty()); + assert_eq!(cache.len(), 0); + assert_eq!(cache.current_size_bytes(), 0); + } + + #[test] + fn test_set_and_get() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp/test.txt", + FileState { + content: "hello world".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + assert!(!cache.is_empty()); + + let entry = cache.get("/tmp/test.txt"); + assert!(entry.is_some()); + assert_eq!(entry.unwrap().content, "hello world"); + } + + #[test] + fn test_path_normalization() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp//test.txt", + FileState { + content: "test".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + // 规范化后的路径应该能命中 + assert!(cache.get("/tmp/test.txt").is_some()); + } + + #[test] + fn test_contains() { + let mut cache = FileStateCache::new(); + assert!(!cache.contains("/tmp/test.txt")); + + cache.set( + "/tmp/test.txt", + FileState { + content: "test".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + assert!(cache.contains("/tmp/test.txt")); + } + + #[test] + fn test_remove() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp/test.txt", + FileState { + content: "test".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + assert!(cache.remove("/tmp/test.txt")); + assert!(cache.is_empty()); + assert!(!cache.remove("/tmp/test.txt")); + } + + #[test] + fn test_clear() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp/a.txt", + FileState { + content: "a".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + cache.set( + "/tmp/b.txt", + FileState { + content: "b".to_string(), + timestamp: 2000, + offset: 1, + limit: None, + }, + ); + cache.clear(); + assert!(cache.is_empty()); + } + + #[test] + fn test_snapshot_sorted_by_timestamp_desc() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp/old.txt", + FileState { + content: "old".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + cache.set( + "/tmp/new.txt", + FileState { + content: "new".to_string(), + timestamp: 3000, + offset: 1, + limit: None, + }, + ); + cache.set( + "/tmp/mid.txt", + FileState { + content: "mid".to_string(), + timestamp: 2000, + offset: 1, + limit: None, + }, + ); + + let snapshot = cache.to_snapshot(); + assert_eq!(snapshot.len(), 3); + // 按 timestamp 降序 + assert_eq!(snapshot[0].1.timestamp, 3000); + assert_eq!(snapshot[1].1.timestamp, 2000); + assert_eq!(snapshot[2].1.timestamp, 1000); + } + + #[test] + fn test_restore_from_snapshot() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp/a.txt", + FileState { + content: "a".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + cache.set( + "/tmp/b.txt", + FileState { + content: "b".to_string(), + timestamp: 2000, + offset: 1, + limit: None, + }, + ); + + let snapshot = cache.to_snapshot(); + cache.clear(); + + cache.restore_from_snapshot(&snapshot, 1); + assert_eq!(cache.len(), 1); + // 应该恢复 timestamp 最高的 + assert!(cache.get("/tmp/b.txt").is_some()); + } + + #[test] + fn test_build_restore_context() { + let mut cache = FileStateCache::new(); + cache.set( + "/tmp/a.txt", + FileState { + content: "file a content here".to_string(), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + + let snapshot = cache.to_snapshot(); + let contexts = FileStateCache::build_restore_context(&snapshot, 2); + assert_eq!(contexts.len(), 1); + assert!(contexts[0].contains("/tmp/a.txt")); + assert!(contexts[0].contains("file a content here")); + } + + #[test] + fn test_empty_snapshot_builds_no_context() { + let contexts = FileStateCache::build_restore_context(&Vec::new(), 5); + assert!(contexts.is_empty()); + } + + #[test] + fn test_lru_eviction_on_size() { + // 创建一个小容量缓存（仅 200 bytes） + let mut cache = FileStateCache::with_limits(100, 200); + + // 写入 3 个 100 字节内容 → 应触发 LRU 淘汰 + cache.set( + "/tmp/1.txt", + FileState { + content: "x".repeat(100), + timestamp: 1000, + offset: 1, + limit: None, + }, + ); + cache.set( + "/tmp/2.txt", + FileState { + content: "y".repeat(100), + timestamp: 2000, + offset: 1, + limit: None, + }, + ); + // 此时应该有 2 个条目（200 bytes） + assert_eq!(cache.len(), 2); + + // 写入第 3 个 → 应淘汰最旧的（1.txt） + cache.set( + "/tmp/3.txt", + FileState { + content: "z".repeat(100), + timestamp: 3000, + offset: 1, + limit: None, + }, + ); + // 旧条目被淘汰 + assert!(!cache.contains("/tmp/1.txt")); + } + + #[test] + fn test_get_file_mtime() { + // 创建临时文件 + let tmp = std::env::temp_dir().join("test_mtime.txt"); + std::fs::write(&tmp, "test").unwrap(); + + let mtime = get_file_mtime(&tmp.to_string_lossy()); + assert!(mtime.is_some()); + assert!(mtime.unwrap() > 0); + + // 不存在的文件 + let mtime = get_file_mtime("/tmp/nonexistent_12345_xxx.txt"); + assert!(mtime.is_none()); + + std::fs::remove_file(&tmp).ok(); + } + + #[test] + fn test_file_unchanged_stub_is_static() { + assert!(FILE_UNCHANGED_STUB.contains("File unchanged")); + } +} diff --git a/src/agent/runtime/finalize.rs b/src/agent/runtime/finalize.rs new file mode 100644 index 0000000..8a66d50 --- /dev/null +++ b/src/agent/runtime/finalize.rs @@ -0,0 +1,137 @@ +// src/agent/runtime/finalize.rs +// +// 会话收尾：更新元信息、生成标题、保存指标、发送 Done 事件。 + +use sqlx::SqlitePool; +use std::path::PathBuf; +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing::info; + +use super::{AgentMetrics, AgentStreamEvent}; +use crate::agent::hooks::{HookRegistry, SessionStopContext}; +use crate::agent::memory::extraction::{run_extraction, ExtractionConfig}; +use crate::agent::terminal::TurnTerminal; +use crate::agent::trajectory::TrajectoryExporter; +use crate::api::AppState; + +/// 完成会话回合：更新 session 元信息、触发 OnSessionStop hook、发送 Done、 +/// 导出 trajectory 数据。 +/// +/// `terminal` 参数允许传入实际的终止原因（取消/错误/超限等）， +/// 传入 `None` 时默认使用 `Completed`。 +#[allow(clippy::too_many_arguments)] +pub async fn finalize_turn( + db: &SqlitePool, + session_id: &str, + metrics: &AgentMetrics, + question: &str, + tx: &mpsc::UnboundedSender, + hook_registry: &HookRegistry, + terminal: Option, + // Trajectory 导出所需参数 + library_dir: Option<&PathBuf>, + llm_model: Option<&str>, + system_prompt: Option<&str>, + // 自动记忆提取所需 + app_state: Option>, +) -> anyhow::Result<()> { + let new_turn_count: i32 = sqlx::query_scalar( + "SELECT COUNT(DISTINCT turn_index) FROM agent_messages WHERE session_id = ?", + ) + .bind(session_id) + .fetch_one(db) + .await + .unwrap_or(0); + + let metrics_json = serde_json::to_value(metrics).unwrap_or_default(); + info!( + "[Finalize] 会话 {} 指标: steps={}, tools={:?}, compressions={}, duplicates={}", + session_id, + metrics.total_steps, + metrics.tool_calls, + metrics.compression_count, + metrics.duplicate_detections + ); + + // 首轮自动生成标题 + if new_turn_count <= 1 { + let title = generate_title(question); + sqlx::query( + "UPDATE agent_sessions SET title = ?, turn_count = ?, metadata = ?, \ + updated_at = CURRENT_TIMESTAMP WHERE session_id = ?", + ) + .bind(&title) + .bind(new_turn_count) + .bind(&metrics_json) + .bind(session_id) + .execute(db) + .await?; + } else { + sqlx::query( + "UPDATE agent_sessions SET turn_count = ?, metadata = ?, \ + updated_at = CURRENT_TIMESTAMP WHERE session_id = ?", + ) + .bind(new_turn_count) + .bind(&metrics_json) + .bind(session_id) + .execute(db) + .await?; + } + + // OnSessionStop hook — 唯一触发点，使用传入的 terminal 或默认 Completed + let default_terminal = TurnTerminal::Completed { + session_id: session_id.to_string(), + total_steps: metrics.total_steps, + }; + let actual_terminal = terminal.unwrap_or(default_terminal); + hook_registry + .run_on_session_stop(&SessionStopContext { + session_id: session_id.to_string(), + terminal: &actual_terminal, + total_steps: metrics.total_steps, + }) + .await; + + // ── Trajectory 导出（非阻塞，失败不影响主流程） ── + if let (Some(lib_dir), Some(model), Some(sys_prompt)) = (library_dir, llm_model, system_prompt) + { + if let Err(e) = TrajectoryExporter::export( + db, + session_id, + lib_dir, + model, + sys_prompt, + metrics, + Some(&actual_terminal), + ) + .await + { + tracing::warn!("[Finalize] Trajectory 导出失败（非致命）: {}", e); + } + } + + // ── 自动记忆提取（fire-and-forget，不阻塞会话关闭） ── + if let Some(app_state) = app_state { + let extraction_config = ExtractionConfig::from_env(); + let session_id_owned = session_id.to_string(); + let mem_mgr = app_state.memory_manager.clone(); + tokio::spawn(async move { + run_extraction(app_state, session_id_owned, mem_mgr, extraction_config).await; + }); + } + + let _ = tx.send(AgentStreamEvent::Done); + + Ok(()) +} + +/// 根据用户首条问题生成会话标题（截取前 50 字符） +fn generate_title(question: &str) -> String { + let chars: String = question.chars().take(50).collect(); + if question.len() > 50 { + format!("{}...", chars) + } else { + chars + } +} diff --git a/src/agent/runtime/mod.rs b/src/agent/runtime/mod.rs new file mode 100644 index 0000000..9bfae32 --- /dev/null +++ b/src/agent/runtime/mod.rs @@ -0,0 +1,1325 @@ +// src/agent/runtime/mod.rs +// +// 科研智能体运行时核心模块。 +// 实现 ReAct 循环：Thought -> Action (工具调用) -> Observation -> Thought... +// 支持会话持久化、上下文压缩、死循环检测和 SSE 流式输出。 +// +// 子模块结构： +// session — 会话创建/恢复、历史加载 +// context — 上下文构建、任务状态恢复 +// streaming — LLM 流式响应处理 +// executor — 工具调用验证与并行执行 +// finalize — 会话收尾、指标持久化 + +pub mod circuit_breaker; +pub mod context; +pub mod error_recovery; +pub mod executor; +pub mod file_cache; +pub mod finalize; +pub mod partitioner; +pub mod permission; +pub mod session; +pub mod streaming; +pub mod streaming_executor; +pub mod system_prompt; +pub mod token_budget; + +use serde::Serialize; +use sqlx::SqlitePool; +use std::collections::HashMap; +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing::{error, info, warn}; + +use super::background::BgNotificationQueue; +use super::compact; +use super::hooks::{HookRegistry, SessionStartContext, StepCompleteContext}; +use super::terminal::TurnTerminal; +use super::tools::ToolRegistry; +use crate::api::AppState; +use crate::clients::llm::{ChatMessage, LlmClient, MessageRole, StreamEvent}; + +use self::error_recovery::{classify_error, ErrorKind, ErrorRecovery}; +use self::session::SessionInfo; +use self::streaming::{StreamOutput, StreamStatus}; +use self::token_budget::TokenBudget; + +/// Agent 配置参数 +#[derive(Debug, Clone)] +pub struct AgentConfig { + /// 最大 ReAct 迭代次数 + pub max_steps: usize, + /// 同质调用检测阈值（连续相同调用次数） + pub duplicate_call_threshold: usize, + /// 工具执行超时时间（秒） + pub tool_timeout_secs: u64, + /// 工具输出最大字符数 + pub max_tool_output_chars: usize, + /// 上下文 Token 估算上限（触发自动摘要压缩） + pub context_char_limit: usize, + /// Token 预算软限制（触发 nudging 提醒） + pub token_soft_limit: usize, + /// Token 预算硬限制（触发强制动作） + pub token_hard_limit: usize, + /// 最大消息数（超过此阈值触发 snip_compact 层压缩） + pub max_messages: usize, +} + +impl AgentConfig { + /// 从环境变量加载配置，缺失时使用默认值。 + pub fn from_env_optional() -> Self { + AgentConfig { + max_steps: std::env::var("AGENT_MAX_STEPS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(8), + duplicate_call_threshold: 3, + tool_timeout_secs: std::env::var("AGENT_TOOL_TIMEOUT_SECS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(120), + max_tool_output_chars: std::env::var("AGENT_MAX_TOOL_OUTPUT_CHARS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(4000), + context_char_limit: std::env::var("AGENT_CONTEXT_CHAR_LIMIT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(16000), + token_soft_limit: std::env::var("AGENT_TOKEN_SOFT_LIMIT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(32000), + token_hard_limit: std::env::var("AGENT_TOKEN_HARD_LIMIT") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(40000), + max_messages: std::env::var("AGENT_MAX_MESSAGES") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(50), + } + } +} + +impl Default for AgentConfig { + fn default() -> Self { + Self::from_env_optional() + } +} + +// ── SSE Stream Events ── + +/// SSE 流式事件（发送给前端） +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type")] +pub enum AgentStreamEvent { + /// 会话创建/恢复 + #[serde(rename = "session")] + Session { session_id: String, title: String }, + /// 智能体思考过程 + #[serde(rename = "thought")] + Thought { content: String, step: usize }, + /// 工具调用开始 + #[serde(rename = "tool_call")] + ToolCall { + name: String, + arguments: serde_json::Value, + step: usize, + }, + /// 工具执行结果（Observation） + #[serde(rename = "tool_result")] + ToolResult { + name: String, + output: String, + is_error: bool, + metadata: serde_json::Value, + step: usize, + }, + /// 文本增量流式输出（最终回答） + #[serde(rename = "text_delta")] + TextDelta { content: String }, + /// Token 使用统计 + #[serde(rename = "usage")] + Usage { + prompt_tokens: u32, + completion_tokens: u32, + total_tokens: u32, + }, + /// 错误通知 + #[serde(rename = "error")] + Error { message: String }, + /// 完成标记 + #[serde(rename = "done")] + Done, +} + +// ── Metrics & Detection ── + +/// Agent 运行指标 +#[derive(Debug, Default, Serialize)] +pub struct AgentMetrics { + pub total_steps: usize, + pub compression_count: usize, + pub duplicate_detections: usize, + /// 各工具调用次数统计 + pub tool_calls: HashMap, +} + +/// 同质调用检测器 +#[derive(Debug, Default)] +pub struct DuplicateDetector { + last_call: Option<(String, String)>, // (tool_name, arguments) + consecutive_count: usize, +} + +impl DuplicateDetector { + /// 记录一次调用，返回是否检测到死循环 + pub fn record(&mut self, tool_name: &str, arguments: &str, threshold: usize) -> bool { + let key = (tool_name.to_string(), arguments.to_string()); + if self.last_call.as_ref() == Some(&key) { + self.consecutive_count += 1; + if self.consecutive_count >= threshold { + return true; + } + } else { + self.last_call = Some(key); + self.consecutive_count = 1; + } + false + } +} + +// ── Agent Runtime ── + +/// 智能体运行时 +pub struct AgentRuntime { + app_state: Arc, + config: AgentConfig, + tool_registry: ToolRegistry, + /// 后台任务通知队列（支持 bg_task_run/bg_task_check） + bg_notification_queue: Arc, + /// 指标采集 hook 的数据引用（供 API 查询） + metrics_data: Arc>, + /// 压缩熔断器（跨 turn 共享，防止无限压缩循环） + compaction_breaker: Arc>, + /// 权限检查器 + permission_checker: Arc, + /// 文件状态缓存（跨 turn 共享，用于 Read 去重） + read_file_state: Arc>, +} + +impl AgentRuntime { + /// 创建新的运行时实例 + pub fn new(app_state: Arc) -> Self { + let queue = Arc::new(BgNotificationQueue::new()); + let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default())); + let permission_checker = Arc::new(permission::PermissionChecker::new()); + let skill_registry = app_state.skill_registry.clone(); + let mut tool_registry = ToolRegistry::new_with_queue(Some(queue.clone()), skill_registry); + // 注册记忆工具 + tool_registry.add_tool(Box::new(crate::agent::tools::memory::SaveMemoryTool::new( + app_state.memory_manager.clone(), + ))); + // 替换 DelegateResearchTool 为带有 permission_checker 的版本 + tool_registry.replace_tool(Box::new( + crate::agent::tools::subagent::DelegateResearchTool::new_with_hooks( + None, + permission_checker.clone(), + None, + ), + )); + AgentRuntime { + app_state, + config: AgentConfig::default(), + tool_registry, + bg_notification_queue: queue, + metrics_data, + compaction_breaker: Arc::new(std::sync::Mutex::new( + circuit_breaker::CompactionCircuitBreaker::new(), + )), + permission_checker, + read_file_state: Arc::new(std::sync::Mutex::new(file_cache::FileStateCache::new())), + } + } + + /// 创建带自定义配置的运行时实例 + pub fn with_config(app_state: Arc, config: AgentConfig) -> Self { + let queue = Arc::new(BgNotificationQueue::new()); + let metrics_data = Arc::new(std::sync::Mutex::new(super::hooks::MetricsData::default())); + let permission_checker = Arc::new(permission::PermissionChecker::new()); + let skill_registry = app_state.skill_registry.clone(); + let mut tool_registry = ToolRegistry::new_with_queue(Some(queue.clone()), skill_registry); + tool_registry.add_tool(Box::new(crate::agent::tools::memory::SaveMemoryTool::new( + app_state.memory_manager.clone(), + ))); + tool_registry.replace_tool(Box::new( + crate::agent::tools::subagent::DelegateResearchTool::new_with_hooks( + None, + permission_checker.clone(), + None, + ), + )); + AgentRuntime { + app_state, + config, + tool_registry, + bg_notification_queue: queue, + metrics_data, + compaction_breaker: Arc::new(std::sync::Mutex::new( + circuit_breaker::CompactionCircuitBreaker::new(), + )), + permission_checker, + read_file_state: Arc::new(std::sync::Mutex::new(file_cache::FileStateCache::new())), + } + } + + /// 返回当前运行指标快照（锁异常时返回默认值） + pub fn get_metrics(&self) -> super::hooks::MetricsData { + self.metrics_data + .lock() + .ok() + .map(|m| m.clone()) + .unwrap_or_default() + } + + // ── Private Helpers ── + + /// 执行文件缓存快照 → 压缩 → 恢复 → 上下文注入 的完整周期。 + /// 返回压缩前的消息数（用于调用者判断压缩是否有效）。 + async fn snapshot_compress_restore( + &self, + messages: &mut Vec, + llm: &LlmClient, + session_id: &str, + hook_registry: &HookRegistry, + ) -> usize { + let before_len = messages.len(); + + // ── 文件缓存快照（压缩前）── + let file_snapshot = { + if let Ok(mut cache) = self.read_file_state.lock() { + let snap = cache.to_snapshot(); + cache.clear(); + snap + } else { + Vec::new() + } + }; + + compact::compress_context_with_hooks( + messages, + llm, + self.config.context_char_limit, + session_id, + Some(hook_registry), + ) + .await; + + // ── 文件缓存恢复（压缩后：重新注入最近文件 + 恢复缓存）── + { + if let Ok(mut cache) = self.read_file_state.lock() { + cache.restore_from_snapshot( + &file_snapshot, + file_cache::POST_COMPACT_MAX_FILES_TO_RESTORE, + ); + } + let restore_ctx = file_cache::FileStateCache::build_restore_context( + &file_snapshot, + file_cache::POST_COMPACT_MAX_FILES_TO_RESTORE, + ); + for block in restore_ctx { + messages.push(ChatMessage::user(format!("[压缩后上下文恢复]\n{}", block))); + } + } + + before_len + } + + // ── Public API ── + + /// 执行完整的智能体对话回合（流式 SSE 输出）。 + /// + /// 流程： + /// 1. 创建/恢复会话 + /// 2. 构建消息上下文 + /// 3. ReAct 循环 + /// 4. 会话收尾 + pub async fn run_turn( + &self, + session_id: Option, + question: &str, + tx: mpsc::UnboundedSender, + ) -> anyhow::Result { + let db = &self.app_state.db; + let llm = &self.app_state.llm; + + // Phase 1: 创建或恢复会话 + let session_info = session::create_or_resume_session(db, session_id.clone(), llm).await?; + + // 构建 hook 注册表（注入依赖，复用 AgentRuntime 的 metrics_data） + let hook_registry = HookRegistry::with_builtins( + db.clone(), + self.app_state.cancelled_runs.clone(), + Some(self.metrics_data.clone()), + ); + + // 触发 OnSessionStart + hook_registry + .run_on_session_start(&SessionStartContext { + session_id: session_info.session_id.clone(), + turn_index: session_info.turn_index, + is_resume: session_id.is_some(), + }) + .await; + + let _ = tx.send(AgentStreamEvent::Session { + session_id: session_info.session_id.clone(), + title: String::new(), + }); + + // Phase 2: 构建初始上下文（历史 + system prompt + 用户消息 + 任务恢复） + let mut messages = context::build_initial_context( + db, + &session_info.session_id, + &self.system_prompt(), + question, + session_info.turn_index, + ) + .await?; + + // 保存用户消息到数据库 + self.save_message( + db, + &session_info.session_id, + session_info.turn_index, + 0, + &ChatMessage::user(question), + None, + ) + .await?; + + // Phase 3: ReAct 循环 + let (metrics, loop_terminal) = self + .run_react_loop(&session_info, &mut messages, &tx, &hook_registry) + .await?; + + // Phase 4: 会话收尾（传入实际的终止原因 + trajectory 导出参数） + let system_prompt = self.system_prompt(); + finalize::finalize_turn( + db, + &session_info.session_id, + &metrics, + question, + &tx, + &hook_registry, + loop_terminal, + Some(&self.app_state.config.library_dir), + Some(self.app_state.llm.model()), + Some(&system_prompt), + Some(self.app_state.clone()), + ) + .await?; + + Ok(session_info.session_id) + } + + // ── ReAct Loop ── + + /// ReAct 循环核心：LLM 调用 → 工具执行 → 结果注入 → 循环... + async fn run_react_loop( + &self, + session_info: &SessionInfo, + messages: &mut Vec, + tx: &mpsc::UnboundedSender, + hook_registry: &HookRegistry, + ) -> anyhow::Result<(AgentMetrics, Option)> { + let db = &self.app_state.db; + let llm = &self.app_state.llm; + let sid = &session_info.session_id; + let turn_index = session_info.turn_index; + + let tool_defs = self.tool_registry.definitions(); + let mut duplicate_detector = DuplicateDetector::default(); + let mut metrics = AgentMetrics::default(); + let mut step = 0; + let mut loop_terminal: Option = None; + + // Token 追踪（API 精确值优先，字符估算作近似值） + let mut last_api_prompt_tokens: Option = None; + let mut msg_count_at_last_call: usize = messages.len(); + let mut steps_since_last_todo: usize = 0; + let nag_after_steps: usize = 3; + let mut pending_manual_compress: bool = false; + + // Token 预算管理器（用于 diminishing returns 检测和 error recovery） + let mut token_budget = + TokenBudget::new(self.config.token_soft_limit, self.config.token_hard_limit); + + loop { + step += 1; + + // 检查用户取消 + let is_cancelled = { + if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { + cancelled.remove(sid) + } else { + false + } + }; + + if is_cancelled { + warn!("[AgentRuntime] 用户手动中止了会话 {} 的智能体执行", sid); + let _ = tx.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + loop_terminal = Some(TurnTerminal::CancelledByUser { + session_id: sid.clone(), + at_step: step, + }); + break; + } + + // ── 上下文压缩检查 ── + // ── Token 感知的压缩触发 ── + // 优先使用 API 返回的精确 prompt_tokens，辅以简单的增量估算 + let estimated_tokens = match last_api_prompt_tokens { + Some(last_tokens) => { + let new_msg_count = messages.len().saturating_sub(msg_count_at_last_call); + let new_tokens_estimate: u32 = messages + .iter() + .rev() + .take(new_msg_count) + .map(|m| (m.content.as_ref().map_or(0, |c| c.len()) + 4) as u32) + .sum(); + (last_tokens + new_tokens_estimate) as usize + } + None => compact::rough_estimate_tokens(messages), + }; + + // 使用 token 预算的软限制作为压缩触发点（而非粗糙的 context_char_limit * 1.5） + let token_limit = token_budget.soft_limit; + + let mut did_compress = false; + + // 熔断器检查：如果连续压缩失败多次，跳过自动压缩 + let breaker_ok = match self.compaction_breaker.lock() { + Ok(mut breaker) => breaker.can_attempt(), + Err(e) => { + warn!("[AgentRuntime] 熔断器锁异常，跳过自动压缩: {:?}", e); + false + } + }; + + if estimated_tokens > token_limit && breaker_ok { + info!( + "[AgentRuntime] 上下文超限 (est. {} tokens > {} limit)，触发压缩", + estimated_tokens, token_limit + ); + let before_len = self + .snapshot_compress_restore( + messages, + llm, + &session_info.session_id, + hook_registry, + ) + .await; + + // 熔断器反馈：压缩后消息数减少 = 成功 + if let Ok(mut breaker) = self.compaction_breaker.lock() { + if messages.len() < before_len { + breaker.record_success(); + } else { + breaker.record_failure(); + } + } else { + warn!("[AgentRuntime] 熔断器反馈写入失败（锁异常）"); + } + last_api_prompt_tokens = None; + msg_count_at_last_call = messages.len(); + metrics.compression_count += 1; + did_compress = true; + } else if estimated_tokens > token_limit && !breaker_ok { + warn!("[AgentRuntime] 熔断器已打开，跳过自动压缩"); + } + + // 处理手动压缩请求（跳过刚自动压缩过的情况，避免双重压缩） + // 手动压缩不受熔断器限制 + if pending_manual_compress && !did_compress { + pending_manual_compress = false; + info!("[AgentRuntime] 执行手动压缩（compress_context 工具触发）"); + + self.snapshot_compress_restore( + messages, + llm, + &session_info.session_id, + hook_registry, + ) + .await; + + // 手动压缩成功后重置熔断器 + if let Ok(mut breaker) = self.compaction_breaker.lock() { + breaker.reset(); + } + last_api_prompt_tokens = None; + msg_count_at_last_call = messages.len(); + metrics.compression_count += 1; + } else if pending_manual_compress { + pending_manual_compress = false; + info!("[AgentRuntime] 跳过手动压缩（刚已完成自动压缩）"); + } + + // Token 预算 diminishing returns 检测 + 渐进式 nudge 提醒 + token_budget.record_continuation(); + + let mut should_nudge = false; + + // TodoWrite nag reminder + if steps_since_last_todo >= nag_after_steps { + messages.push(ChatMessage::user( + "提醒：你已经连续多步未更新任务计划。建议调用 todo_write 工具复盘当前进度并规划后续步骤。", + )); + steps_since_last_todo = 0; + should_nudge = true; + } + + // Token 预算 nudge（仅在无 nag 时注入，避免消息过多） + if !should_nudge { + if let Some(nudge) = token_budget.nudge_message() { + messages.push(ChatMessage::user(nudge)); + } + } + + // Diminishing returns 检测 — 强制结束 + if token_budget.diminishing_returns { + warn!("[AgentRuntime] 检测到 diminishing returns，强制结束循环"); + let _ = tx.send(AgentStreamEvent::Error { + message: "检测到重复操作模式，已自动停止。请查看已收集的信息。".to_string(), + }); + messages.push(ChatMessage::user( + "检测到你的后续步骤未产生新信息（diminishing returns）。\ + 请基于已收集的全部信息直接给出最终答案，不要再调用任何工具。", + )); + let _ = self + .final_answer_without_tools(llm, messages, sid, turn_index, step, tx) + .await; + break; + } + + // 最大步数检查 + if step > self.config.max_steps { + warn!( + "[AgentRuntime] 达到最大步数限制 ({} steps)", + self.config.max_steps + ); + let _ = tx.send(AgentStreamEvent::Error { + message: format!( + "已达到最大推理步数 ({})，请根据已收集的信息给出最终回答。", + self.config.max_steps + ), + }); + messages.push(ChatMessage::user(format!( + "你已经执行了 {} 步（最大 {} 步）。请根据已有信息直接给出最终答案，不要再调用工具。", + step, self.config.max_steps + ))); + let _ = self + .final_answer_without_tools(llm, messages, sid, turn_index, step, tx) + .await; + break; + } + + // ── 后台任务通知注入 ── + let bg_results = self.bg_notification_queue.drain().await; + for result in bg_results { + let status = if result.is_error { "❌" } else { "✅" }; + messages.push(ChatMessage::user(format!( + "[后台任务完成] {} {}: {} ({}): {}", + status, result.tool_name, result.bibcode, result.task_id, result.summary, + ))); + } + + // ── LLM 流式调用（含错误恢复） ── + let stream_output_opt = self + .call_llm_with_recovery(llm, messages, &tool_defs, tx, step, sid, &mut token_budget) + .await; + + let stream_output = match stream_output_opt { + Some(output) => output, + None => { + // 所有恢复尝试均失败 + loop_terminal = Some(TurnTerminal::ModelError { + session_id: sid.clone(), + message: "LLM 调用失败，所有恢复步骤已尝试完毕".to_string(), + }); + break; + } + }; + + // 更新 API 精确 token 计数 + token 预算 + if let Some(ref u) = stream_output.usage { + last_api_prompt_tokens = Some(u.prompt_tokens); + msg_count_at_last_call = messages.len(); + token_budget.spend_input(u.prompt_tokens as usize); + token_budget.spend_output(u.completion_tokens as usize); + } + + // ── 处理 Thought/Reasoning ── + let mut thought_content = stream_output.reasoning.clone(); + + if thought_content.is_none() && stream_output.is_tool_call_step + && !stream_output.content.is_empty() { + thought_content = Some(stream_output.content.clone()); + } + + if stream_output.is_tool_call_step { + if let Some(ref thought_text) = thought_content { + let _ = tx.send(AgentStreamEvent::Thought { + content: thought_text.clone(), + step, + }); + } + } + + // ── 无工具调用 = 最终回答 ── + let tool_calls = match stream_output.tool_calls { + Some(ref tc) if !tc.is_empty() => tc.clone(), + _ => { + // 保存最终回答 + let assistant_msg = ChatMessage::assistant_with_reasoning( + if stream_output.content.is_empty() { + None + } else { + Some(stream_output.content.clone()) + }, + stream_output.reasoning.clone(), + None, + ); + self.save_message( + db, + sid, + turn_index, + step as i32, + &assistant_msg, + stream_output.reasoning.as_deref(), + ) + .await?; + messages.push(assistant_msg); + + // 发送未发送的 reasoning + if let Some(ref thought_text) = stream_output.reasoning { + if thought_content.is_none() { + let _ = tx.send(AgentStreamEvent::Thought { + content: thought_text.clone(), + step, + }); + } + } + + // Token 使用统计 + if let Some(u) = stream_output.usage { + let _ = tx.send(AgentStreamEvent::Usage { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + }); + } + break; + } + }; + + // ── 工具调用处理 ── + // 检测 todo_write 和 compress_context + let called_todo_write = tool_calls.iter().any(|tc| tc.function.name == "todo_write"); + if called_todo_write { + steps_since_last_todo = 0; + } else { + steps_since_last_todo += 1; + } + + if tool_calls + .iter() + .any(|tc| tc.function.name == "compress_context") + { + pending_manual_compress = true; + } + + // 更新指标 + metrics.total_steps = step; + for tc in &tool_calls { + *metrics + .tool_calls + .entry(tc.function.name.clone()) + .or_insert(0) += 1; + } + + // 构建 assistant 消息（含 tool_calls） + let assistant_msg = ChatMessage::assistant_with_reasoning( + if stream_output.content.is_empty() { + None + } else { + Some(stream_output.content.clone()) + }, + stream_output.reasoning.clone(), + Some(tool_calls.clone()), + ); + self.save_message( + db, + sid, + turn_index, + step as i32, + &assistant_msg, + stream_output.reasoning.as_deref(), + ) + .await?; + messages.push(assistant_msg); + + // 验证 + 准备工具调用 + let (prepared_calls, has_duplicate) = executor::validate_and_prepare( + &tool_calls, + &mut duplicate_detector, + self.config.duplicate_call_threshold, + messages, + tx, + db, + sid, + turn_index, + step, + ); + + if has_duplicate { + metrics.duplicate_detections += 1; + continue; + } + + if prepared_calls.is_empty() { + continue; + } + + // 并行执行工具（带权限检查和分区器） + let exec_result = executor::execute_parallel( + &prepared_calls, + &self.tool_registry, + self.app_state.clone(), + hook_registry, + Some(&self.permission_checker), + tx, + db, + sid, + "lead", + turn_index, + step, + self.config.tool_timeout_secs, + self.config.max_tool_output_chars, + self.read_file_state.clone(), + ) + .await; + + // 将工具结果推入消息上下文 + for tm in exec_result.tool_messages { + messages.push(tm.chat_message); + } + + // 持久化 todo_write 任务状态到数据库 + if called_todo_write { + for prep in &prepared_calls { + if prep.tool_name == "todo_write" { + if let Some(todos) = prep.args.get("todos").and_then(|t| t.as_array()) { + let todos_vec: Vec = todos.to_vec(); + let _ = crate::agent::tools::persist_tasks(db, sid, &todos_vec, "lead") + .await; + } + } + } + } + + if exec_result.was_cancelled { + if let Ok(mut locked) = self.app_state.cancelled_runs.lock() { + locked.remove(sid); + } + warn!( + "[AgentRuntime] 工具执行期间被用户手动中止，会话 ID: {}", + sid + ); + let _ = tx.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + loop_terminal = Some(TurnTerminal::CancelledByUser { + session_id: sid.clone(), + at_step: step, + }); + break; + } + + // OnStepComplete hook + let step_ctx = StepCompleteContext { + session_id: sid.clone(), + step, + max_steps: self.config.max_steps, + messages_count: messages.len(), + estimated_tokens, + token_limit, + }; + hook_registry.run_on_step_complete(&step_ctx).await; + } + + Ok((metrics, loop_terminal)) + } + + /// LLM 流式调用，含完整的错误恢复阶梯。 + /// + /// 首次调用失败后，按顺序尝试： + /// 1. AggressiveCompact (keep_recent=2) + /// 2. ReactiveCompact (LLM 摘要) + /// 3. EscalateTokens (提升 hard_limit → 64k) + /// 4. MultiTurn (注入分步消息) + /// 5. Surface (放弃) + /// + /// 每一步后重试 LLM 调用。返回 Some(StreamOutput) 表示成功（可能经过恢复）， + /// None 表示所有步骤均已尝试且失败。 + #[allow(clippy::too_many_arguments)] + async fn call_llm_with_recovery( + &self, + llm: &LlmClient, + messages: &mut Vec, + tool_defs: &[crate::clients::llm::ToolDefinition], + tx: &mpsc::UnboundedSender, + step: usize, + session_id: &str, + token_budget: &mut TokenBudget, + ) -> Option { + // 首次尝试 + let output = streaming::process_llm_stream( + llm, + messages, + tool_defs, + tx, + step, + session_id, + self.app_state.cancelled_runs.clone(), + ) + .await; + + match output.status { + StreamStatus::Success => return Some(output), + StreamStatus::Cancelled => { + if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { + cancelled.remove(session_id); + } + warn!( + "[AgentRuntime] 流式调用期间被用户手动中止，会话 ID: {}", + session_id + ); + let _ = tx.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + return None; + } + StreamStatus::Error(ref e_str) => { + error!("[AgentRuntime] 流式读取错误: {}", e_str); + } + } + + // 提取错误字符串（用于分类） + let e_str = match &output.status { + StreamStatus::Error(s) => s.clone(), + _ => return Some(output), // 不应到达，但安全起见 + }; + + let error_kind = classify_error(&e_str); + + // ── 429/529 瞬态错误：指数退避重试（独立的快速路径） ── + if matches!(error_kind, ErrorKind::RateLimited | ErrorKind::Overloaded) { + let retry_after_secs = error_recovery::parse_retry_after(&e_str); + let mut consecutive_overloads: u32 = 0; + const MAX_BACKOFF_RETRIES: u32 = 10; + + for attempt in 0..MAX_BACKOFF_RETRIES { + let delay_ms = error_recovery::backoff_delay(attempt, retry_after_secs); + info!( + "[AgentRuntime] 退避重试 {}/{} ({}ms, error={:?})", + attempt + 1, + MAX_BACKOFF_RETRIES, + delay_ms, + error_kind + ); + + let _ = tx.send(AgentStreamEvent::Thought { + content: format!( + "⏳ 模型服务暂时不可用，正在重试 ({}/{})...", + attempt + 1, + MAX_BACKOFF_RETRIES + ), + step, + }); + + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + + // 529 连续过载检测：3 次后尝试切换备用模型 + if matches!(error_kind, ErrorKind::Overloaded) { + consecutive_overloads += 1; + if consecutive_overloads >= 3 { + if let Ok(fallback) = std::env::var("FALLBACK_MODEL") { + warn!( + "[AgentRuntime] 连续 {} 次过载，切换到备用模型: {}", + consecutive_overloads, fallback + ); + // Note: The LlmClient model is immutable. In production, + // this would require a model-override capable client. + // For now, log and continue retrying with current model. + } + } + } + + // 检查用户取消 + if let Ok(cancelled) = self.app_state.cancelled_runs.lock() { + if cancelled.contains(session_id) { + warn!("[AgentRuntime] 退避重试期间被用户取消"); + let _ = tx.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + return None; + } + } + + // 重试 LLM 调用 + let retry_output = streaming::process_llm_stream( + llm, + messages, + tool_defs, + tx, + step, + session_id, + self.app_state.cancelled_runs.clone(), + ) + .await; + + match retry_output.status { + StreamStatus::Success => { + info!("[AgentRuntime] 退避重试成功！(尝试 {})", attempt + 1); + return Some(retry_output); + } + StreamStatus::Cancelled => { + if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { + cancelled.remove(session_id); + } + return None; + } + StreamStatus::Error(_) => { + // 继续重试 + continue; + } + } + } + + // 所有退避重试失败 + warn!( + "[AgentRuntime] {} 次退避重试后仍然失败", + MAX_BACKOFF_RETRIES + ); + let _ = tx.send(AgentStreamEvent::Error { + message: format!( + "模型服务暂时不可用（已重试 {} 次）。请稍后再试或检查模型服务状态。", + MAX_BACKOFF_RETRIES + ), + }); + return None; + } + + if !ErrorRecovery::is_recoverable(&error_kind) { + let _ = tx.send(AgentStreamEvent::Error { + message: format!("大模型流式读取失败: {}", e_str), + }); + return None; + } + + let mut recovery = ErrorRecovery::new(token_budget.clone()); + + while let Some(recovery_step) = recovery.try_recover(&error_kind) { + match recovery_step { + error_recovery::RecoveryStep::RetryWithBackoff { attempt, delay_ms } => { + // 429/529 本应在 streaming 层处理，若到达此处说明分类逻辑有变更， + // 安全降级为 sleep + 直接重试（不依赖 streaming 层重试）。 + warn!( + "[AgentRuntime] RetryWithBackoff 在 error_recovery 层触发 (attempt={}, delay={}ms)，执行降级重试", + attempt, delay_ms + ); + tokio::time::sleep(std::time::Duration::from_millis(delay_ms)).await; + // 不计入 recovery 计数，由外层循环自然重试 + } + error_recovery::RecoveryStep::AggressiveCompact => { + info!("[AgentRuntime] 错误恢复: 激进压缩 (snip + micro with keep_recent=2)"); + compact::snip_compact(messages, self.config.max_messages); + compact::micro_compact(messages, 2); + } + error_recovery::RecoveryStep::ReactiveCompact => { + info!("[AgentRuntime] 错误恢复: LLM 摘要压缩"); + compact::compress_context( + messages, + llm, + self.config.context_char_limit, + session_id, + ) + .await; + } + error_recovery::RecoveryStep::EscalateTokens { .. } => { + info!( + "[AgentRuntime] 错误恢复: 提升 token 硬限制到 {}", + recovery.token_budget.hard_limit + ); + } + error_recovery::RecoveryStep::MultiTurn => { + info!("[AgentRuntime] 错误恢复: 注入多轮消息"); + messages.push(ChatMessage::user(ErrorRecovery::multi_turn_message())); + } + error_recovery::RecoveryStep::Surface => { + warn!("[AgentRuntime] 错误恢复: 所有步骤失败，暴露错误"); + break; + } + } + + // 重试 LLM 调用 + let retry_output = streaming::process_llm_stream( + llm, + messages, + tool_defs, + tx, + step, + session_id, + self.app_state.cancelled_runs.clone(), + ) + .await; + + match retry_output.status { + StreamStatus::Success => { + info!("[AgentRuntime] 错误恢复成功！"); + // 将恢复后的 token_budget 状态同步回去 + *token_budget = recovery.token_budget.clone(); + return Some(retry_output); + } + StreamStatus::Cancelled => { + if let Ok(mut cancelled) = self.app_state.cancelled_runs.lock() { + cancelled.remove(session_id); + } + warn!("[AgentRuntime] 恢复期间被用户中止"); + let _ = tx.send(AgentStreamEvent::Error { + message: "用户已手动中止执行。".to_string(), + }); + return None; + } + StreamStatus::Error(retry_err) => { + info!( + "[AgentRuntime] 恢复步骤 {:?} 未能解决，继续下一阶梯: {}", + recovery_step, retry_err + ); + } + } + } + + // 所有恢复步骤均已尝试 + let _ = tx.send(AgentStreamEvent::Error { + message: format!("大模型流式读取失败，且所有恢复步骤均未能解决: {}", e_str), + }); + None + } + + // ── Helpers ── + + /// 系统提示词（模块化组装 — 参考 Claude Code s10）。 + /// 静态 section 在前以最大化 prompt cache 命中率。 + fn system_prompt(&self) -> String { + use self::system_prompt::{SystemPrompt, IDENTITY_SECTION, PRINCIPLES_SECTION}; + + let mut sp = SystemPrompt::new(); + + // Section 1: 静态身份（始终加载，最大化缓存） + sp.add_section("identity", IDENTITY_SECTION.to_string()); + + // Section 2: 动态工具列表（运行时生成） + let mut tools_desc = String::from("你可以使用以下工具：\n"); + for def in self.tool_registry.definitions() { + let short_desc: String = def + .function + .description + .split('。') + .next() + .unwrap_or(&def.function.description) + .chars() + .take(80) + .collect(); + tools_desc.push_str(&format!("- {}: {}\n", def.function.name, short_desc)); + } + sp.add_section("tools", tools_desc); + + // Section 3: 可用技能（动态） + if let Some(skills) = self + .app_state + .skill_registry + .read() + .ok() + .and_then(|r| r.build_reminder()) + { + sp.add_section("skills", skills); + } + + // Section 4: 项目记忆（按需加载） + if let Some(memory) = self + .app_state + .memory_manager + .try_lock() + .ok() + .and_then(|mgr| mgr.build_system_reminder(5)) + { + sp.add_section("memory", memory); + } + + // Section 5: 静态核心原则（最后加载，因较常变化） + sp.add_section("principles", PRINCIPLES_SECTION.to_string()); + + sp.assemble() + } + + /// 步数耗尽时的最终答案生成（不带工具调用，强制 LLM 直接回答） + async fn final_answer_without_tools( + &self, + llm: &LlmClient, + messages: &[ChatMessage], + session_id: &str, + turn_index: i32, + step: usize, + tx: &mpsc::UnboundedSender, + ) -> anyhow::Result<()> { + let empty_tools: Vec = Vec::new(); + let mut stream_rx = match llm.chat_stream(messages, &empty_tools).await { + Ok(rx) => rx, + Err(e) => { + let _ = tx.send(AgentStreamEvent::Error { + message: format!("最终回答生成失败: {}", e), + }); + return Err(anyhow::anyhow!("final_answer LLM call failed: {}", e)); + } + }; + + let mut accumulated = String::new(); + while let Some(event) = stream_rx.recv().await { + match event { + StreamEvent::TextDelta(delta) => { + accumulated.push_str(&delta); + let _ = tx.send(AgentStreamEvent::TextDelta { content: delta }); + } + StreamEvent::Usage(u) => { + let _ = tx.send(AgentStreamEvent::Usage { + prompt_tokens: u.prompt_tokens, + completion_tokens: u.completion_tokens, + total_tokens: u.total_tokens, + }); + } + StreamEvent::Done => break, + StreamEvent::Error(e) => { + let _ = tx.send(AgentStreamEvent::Error { + message: format!("最终回答流式错误: {}", e), + }); + break; + } + _ => {} + } + } + + let assistant_msg = ChatMessage::assistant(accumulated.clone()); + self.save_message( + &self.app_state.db, + session_id, + turn_index, + step as i32, + &assistant_msg, + None, + ) + .await?; + + Ok(()) + } + + /// 保存消息到数据库 + async fn save_message( + &self, + db: &SqlitePool, + session_id: &str, + turn_index: i32, + step_index: i32, + msg: &ChatMessage, + thought: Option<&str>, + ) -> anyhow::Result<()> { + self.save_message_as(db, session_id, turn_index, step_index, msg, thought, "lead") + .await + } + + /// 保存消息到数据库（指定 agent 身份） + #[allow(clippy::too_many_arguments)] + async fn save_message_as( + &self, + db: &SqlitePool, + session_id: &str, + turn_index: i32, + step_index: i32, + msg: &ChatMessage, + thought: Option<&str>, + agent_name: &str, + ) -> anyhow::Result<()> { + let role = match msg.role { + MessageRole::System => "system", + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::Tool => "tool", + }; + + let content = msg.content.as_deref().unwrap_or(""); + let tool_calls_json = msg + .tool_calls + .as_ref() + .map(|tc| serde_json::to_string(tc).unwrap_or_default()); + let tool_call_id = msg.tool_call_id.as_deref(); + let token_count = content.len() as i32 / 4; + + sqlx::query( + "INSERT INTO agent_messages (session_id, turn_index, step_index, role, content, thought, tool_calls, tool_call_id, token_count, agent_name) \ + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + ) + .bind(session_id) + .bind(turn_index) + .bind(step_index) + .bind(role) + .bind(content) + .bind(thought) + .bind(&tool_calls_json) + .bind(tool_call_id) + .bind(token_count) + .bind(agent_name) + .execute(db) + .await?; + + Ok(()) + } + + /// 非致命保存：数据库写入失败时记录日志但不终止 turn + #[allow(dead_code)] + #[allow(clippy::too_many_arguments)] + async fn save_message_non_fatal( + &self, + db: &SqlitePool, + session_id: &str, + turn_index: i32, + step_index: i32, + msg: &ChatMessage, + thought: Option<&str>, + agent_name: &str, + ) { + if let Err(e) = self + .save_message_as( + db, session_id, turn_index, step_index, msg, thought, agent_name, + ) + .await + { + warn!("[AgentRuntime] 消息持久化失败（非致命）: {}", e); + } + } +} diff --git a/src/agent/runtime/partitioner.rs b/src/agent/runtime/partitioner.rs new file mode 100644 index 0000000..d04f65c --- /dev/null +++ b/src/agent/runtime/partitioner.rs @@ -0,0 +1,189 @@ +// src/agent/runtime/partitioner.rs +// +// 工具调用并发分区器。 +// 将准备好的工具调用按并发安全性分组成批处理。 +// 参考 Claude Code partitionToolCalls() 设计。 +// +// 分区规则： +// 1. 连续的并发安全工具放在同一个并行批次 +// 2. 非并发安全的工具单独一个批次（串行执行） +// 3. 每个并行批次最多 max_concurrency 个工具 + +use tracing::debug; + +use super::executor::PreparedCall; +use crate::agent::tools::ToolRegistry; + +/// 一批工具调用 +#[derive(Debug, Clone)] +pub struct ToolBatch { + /// 该批次是否可以并行执行 + pub is_parallel: bool, + /// 批次中的工具调用（按原始顺序） + pub calls: Vec, +} + +/// 工具调用分区器 +pub struct ToolPartitioner { + /// 最大并行度（并行批次中最多执行的工具数） + max_concurrency: usize, +} + +impl ToolPartitioner { + /// 创建分区器 + /// + /// `max_concurrency` 为 0 时使用默认值 10。 + pub fn new(max_concurrency: usize) -> Self { + let concurrency = if max_concurrency == 0 { + 10 + } else { + max_concurrency + }; + ToolPartitioner { + max_concurrency: concurrency, + } + } + + /// 将 prepared_calls 分区为顺序批处理。 + /// + /// 返回的批次列表按顺序依次执行： + /// - `is_parallel = true` 的批次内工具可并发执行 + /// - `is_parallel = false` 的批次内只有一个工具，需串行执行 + pub fn partition( + &self, + prepared_calls: &[PreparedCall], + tool_registry: &ToolRegistry, + ) -> Vec { + let mut batches: Vec = Vec::new(); + + for prep in prepared_calls { + let is_safe = self.is_concurrent_safe(prep, tool_registry); + + // 尝试追加到上一个并行批次 + if is_safe { + if let Some(last) = batches.last_mut() { + if last.is_parallel && last.calls.len() < self.max_concurrency { + last.calls.push(prep.clone()); + continue; + } + } + // 新建并行批次 + batches.push(ToolBatch { + is_parallel: true, + calls: vec![prep.clone()], + }); + } else { + // 非并发安全，单独一个串行批次 + // 如果前一个也是串行，可以合并（但为了简单，保持每个非安全工具独立批次） + batches.push(ToolBatch { + is_parallel: false, + calls: vec![prep.clone()], + }); + } + } + + debug!( + "[Partitioner] 分区完成: {} calls → {} batches", + prepared_calls.len(), + batches.len() + ); + batches + } + + /// 判断单个工具调用是否可并发安全执行 + fn is_concurrent_safe(&self, prep: &PreparedCall, tool_registry: &ToolRegistry) -> bool { + tool_registry + .get(&prep.tool_name) + .map(|t| t.is_concurrency_safe(&prep.args)) + .unwrap_or(false) + } +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::sync::{Arc, RwLock}; + + use super::*; + use crate::agent::skills::SkillRegistry; + + fn make_registry() -> ToolRegistry { + ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from( + "./skills", + ))))) + } + + fn make_prep(name: &str) -> PreparedCall { + PreparedCall { + tool_call_id: format!("call_{}", name), + tool_name: name.to_string(), + args: serde_json::json!({}), + } + } + + #[test] + fn test_all_concurrent_in_single_batch() { + let registry = make_registry(); + let partitioner = ToolPartitioner::new(10); + + let calls = vec![ + make_prep("search_papers"), + make_prep("rag_search"), + make_prep("get_paper_metadata"), + ]; + + let batches = partitioner.partition(&calls, ®istry); + + // search_papers, rag_search, get_paper_metadata 都是并发安全的 + // 它们应该在一个并行批次中 + assert_eq!(batches.len(), 1); + assert!(batches[0].is_parallel); + assert_eq!(batches[0].calls.len(), 3); + } + + #[test] + fn test_download_splits_batch() { + let registry = make_registry(); + let partitioner = ToolPartitioner::new(10); + + let calls = vec![ + make_prep("search_papers"), + make_prep("download_paper"), + make_prep("rag_search"), + ]; + + let batches = partitioner.partition(&calls, ®istry); + + // search_papers (安全) → download_paper (不安全) → rag_search (安全) + // 应分为 3 个批次 + assert_eq!(batches.len(), 3); + assert!(batches[0].is_parallel); // search_papers + assert_eq!(batches[0].calls.len(), 1); + assert!(!batches[1].is_parallel); // download_paper (串行) + assert_eq!(batches[1].calls.len(), 1); + assert!(batches[2].is_parallel); // rag_search + assert_eq!(batches[2].calls.len(), 1); + } + + #[test] + fn test_max_concurrency_limit() { + let registry = make_registry(); + let partitioner = ToolPartitioner::new(2); + + let calls = vec![ + make_prep("search_papers"), + make_prep("rag_search"), + make_prep("get_paper_metadata"), + make_prep("load_skill"), + ]; + + let batches = partitioner.partition(&calls, ®istry); + + // 4 个并发安全调用，max_concurrency=2 → 应分为 2 个并行批次 + assert_eq!(batches.len(), 2); + assert!(batches[0].is_parallel); + assert_eq!(batches[0].calls.len(), 2); + assert!(batches[1].is_parallel); + assert_eq!(batches[1].calls.len(), 2); + } +} diff --git a/src/agent/runtime/permission.rs b/src/agent/runtime/permission.rs new file mode 100644 index 0000000..618c60a --- /dev/null +++ b/src/agent/runtime/permission.rs @@ -0,0 +1,187 @@ +// src/agent/runtime/permission.rs +// +// 权限检查管道 — 优先级排序的规则链。 +// 参考 Claude Code PermissionChecker 设计。 +// +// 规则优先级（从高到低）： +// 1. Deny — 不可覆盖的拒绝 +// 2. Allow — 允许 +// 3. Ask — 需要用户确认 +// +// 支持通配符 "*" 匹配所有工具。 + +use tracing::info; + +use crate::agent::tools::PermissionRule; + +/// 权限检查结果 +#[derive(Debug, Clone, PartialEq)] +pub enum PermissionResult { + /// 被拒绝（不可覆盖） + Denied { reason: String }, + /// 允许 + Allowed, + /// 需要用户确认 + AskUser { message: String }, +} + +impl PermissionResult { + pub fn is_allowed(&self) -> bool { + matches!(self, PermissionResult::Allowed) + } + + pub fn is_denied(&self) -> bool { + matches!(self, PermissionResult::Denied { .. }) + } +} + +/// 权限检查器 — 维护有序规则列表并逐条匹配 +pub struct PermissionChecker { + rules: Vec, +} + +impl PermissionChecker { + /// 创建空的检查器（默认允许所有） + pub fn new() -> Self { + PermissionChecker { rules: Vec::new() } + } + + /// 添加规则。先添加的优先级更高。 + pub fn add_rule(&mut self, rule: PermissionRule) { + self.rules.push(rule); + } + + /// 检查指定工具是否可以执行。 + /// + /// 遍历规则列表，返回第一个匹配的决策。 + /// 无匹配规则时默认 Allow。 + pub fn check(&self, tool_name: &str) -> PermissionResult { + for rule in &self.rules { + match rule { + PermissionRule::Deny { + tool_name: name, + reason, + } if Self::matches(name, tool_name) => { + info!("[Permission] 拒绝工具 {}: {}", tool_name, reason); + return PermissionResult::Denied { + reason: reason.clone(), + }; + } + PermissionRule::Allow { tool_name: name } if Self::matches(name, tool_name) => { + return PermissionResult::Allowed; + } + PermissionRule::Ask { + tool_name: name, + message, + } if Self::matches(name, tool_name) => { + return PermissionResult::AskUser { + message: message.clone(), + }; + } + _ => {} + } + } + // 默认允许 + PermissionResult::Allowed + } + + /// 检查是否有明确拒绝该工具的规则 + pub fn is_denied(&self, tool_name: &str) -> bool { + self.check(tool_name).is_denied() + } + + /// 规则名称匹配：支持精确匹配和通配符 "*" + fn matches(pattern: &str, tool_name: &str) -> bool { + pattern == "*" || pattern == tool_name + } +} + +impl Default for PermissionChecker { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::tools::PermissionRule; + + #[test] + fn test_empty_checker_allows_all() { + let checker = PermissionChecker::new(); + assert_eq!(checker.check("search_papers"), PermissionResult::Allowed); + assert_eq!(checker.check("download_paper"), PermissionResult::Allowed); + } + + #[test] + fn test_deny_wins_over_allow() { + let mut checker = PermissionChecker::new(); + checker.add_rule(PermissionRule::Deny { + tool_name: "download_paper".into(), + reason: "blocked".into(), + }); + checker.add_rule(PermissionRule::Allow { + tool_name: "download_paper".into(), + }); + + let result = checker.check("download_paper"); + assert!(result.is_denied()); + } + + #[test] + fn test_wildcard_deny_blocks_all() { + let mut checker = PermissionChecker::new(); + checker.add_rule(PermissionRule::Deny { + tool_name: "*".into(), + reason: "all blocked".into(), + }); + + assert!(checker.check("search_papers").is_denied()); + assert!(checker.check("download_paper").is_denied()); + assert!(checker.is_denied("rag_search")); + } + + #[test] + fn test_ask_returns_ask_user() { + let mut checker = PermissionChecker::new(); + checker.add_rule(PermissionRule::Ask { + tool_name: "delete_paper".into(), + message: "Are you sure?".into(), + }); + + let result = checker.check("delete_paper"); + assert_eq!( + result, + PermissionResult::AskUser { + message: "Are you sure?".into() + } + ); + } + + #[test] + fn test_no_match_defaults_to_allow() { + let mut checker = PermissionChecker::new(); + checker.add_rule(PermissionRule::Deny { + tool_name: "download_paper".into(), + reason: "blocked".into(), + }); + + assert_eq!(checker.check("search_papers"), PermissionResult::Allowed); + } + + #[test] + fn test_rule_ordering_first_match_wins() { + let mut checker = PermissionChecker::new(); + // 先添加 Allow，后添加 Deny — Allow 先匹配 + checker.add_rule(PermissionRule::Allow { + tool_name: "search_papers".into(), + }); + checker.add_rule(PermissionRule::Deny { + tool_name: "search_papers".into(), + reason: "should not match".into(), + }); + + assert_eq!(checker.check("search_papers"), PermissionResult::Allowed); + } +} diff --git a/src/agent/runtime/session.rs b/src/agent/runtime/session.rs new file mode 100644 index 0000000..8969d25 --- /dev/null +++ b/src/agent/runtime/session.rs @@ -0,0 +1,144 @@ +// src/agent/runtime/session.rs +// +// 会话生命周期管理：创建/恢复/验证 Agent 会话。 + +use sqlx::SqlitePool; + +use crate::clients::llm::LlmClient; + +/// 会话信息摘要 +#[derive(Debug, Clone)] +pub struct SessionInfo { + pub session_id: String, + pub turn_index: i32, +} + +/// 创建新会话或恢复已有会话。 +/// +/// 返回会话信息。如果指定的 session_id 不存在则返回错误。 +pub async fn create_or_resume_session( + db: &SqlitePool, + session_id: Option, + llm: &LlmClient, +) -> anyhow::Result { + match session_id { + Some(id) => { + // 验证会话存在且未被软删除 + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM agent_sessions WHERE session_id = ? AND deleted_at IS NULL)", + ) + .bind(&id) + .fetch_one(db) + .await + .unwrap_or(false); + + if !exists { + return Err(anyhow::anyhow!("会话 {} 不存在或已删除", id)); + } + + // 计算当前轮次号 + let turn_index: i32 = sqlx::query_scalar( + "SELECT COALESCE(MAX(turn_index), -1) + 1 FROM agent_messages WHERE session_id = ?", + ) + .bind(&id) + .fetch_one(db) + .await + .unwrap_or(0); + + Ok(SessionInfo { + session_id: id, + turn_index, + }) + } + None => { + let new_id = uuid::Uuid::new_v4().to_string(); + sqlx::query("INSERT INTO agent_sessions (session_id, title, model) VALUES (?, ?, ?)") + .bind(&new_id) + .bind("") + .bind(llm.model()) + .execute(db) + .await?; + + Ok(SessionInfo { + session_id: new_id, + turn_index: 0, + }) + } + } +} + +/// 加载会话的历史消息（供 LLM 上下文使用）。 +/// +/// `agent_name` 参数用于消息隔离： +/// - `"lead"` — 只加载 Lead Agent 自己的消息（默认） +/// - `"*"` — 加载所有 agent 的消息（调试/审计用） +pub async fn load_history_for_llm( + db: &SqlitePool, + session_id: &str, +) -> anyhow::Result> { + load_history_for_agent(db, session_id, "lead").await +} + +/// 加载指定 agent 的历史消息。 +pub async fn load_history_for_agent( + db: &SqlitePool, + session_id: &str, + agent_name: &str, +) -> anyhow::Result> { + use crate::clients::llm::{ChatMessage, MessageRole}; + + #[allow(clippy::type_complexity)] + let rows: Vec<( + String, + String, + Option, + Option, + Option, + )> = if agent_name == "*" { + sqlx::query_as( + "SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \ + WHERE session_id = ? ORDER BY id ASC", + ) + .bind(session_id) + .fetch_all(db) + .await? + } else { + sqlx::query_as( + "SELECT role, content, tool_calls, tool_call_id, thought FROM agent_messages \ + WHERE session_id = ? AND agent_name = ? ORDER BY id ASC", + ) + .bind(session_id) + .bind(agent_name) + .fetch_all(db) + .await? + }; + + let mut messages = Vec::new(); + for (role_str, content, tool_calls_json, tool_call_id, thought) in rows { + let role = match role_str.as_str() { + "system" => MessageRole::System, + "user" => MessageRole::User, + "assistant" => MessageRole::Assistant, + "tool" => MessageRole::Tool, + _ => continue, + }; + + let tool_calls: Option> = + tool_calls_json.and_then(|json_str| serde_json::from_str(&json_str).ok()); + + messages.push(ChatMessage { + role, + content: if content.is_empty() { + None + } else { + Some(content) + }, + tool_call_id, + tool_calls, + name: None, + reasoning_content: thought, + }); + } + + Ok(messages) +} diff --git a/src/agent/runtime/streaming.rs b/src/agent/runtime/streaming.rs new file mode 100644 index 0000000..8f30c77 --- /dev/null +++ b/src/agent/runtime/streaming.rs @@ -0,0 +1,161 @@ +// src/agent/runtime/streaming.rs +// +// LLM 流式响应处理：消费 chat_stream 返回的 StreamEvent 通道， +// 支持并发取消检测，累积推理内容、文本增量和工具调用。 + +use std::sync::Arc; +use tokio::sync::mpsc; +use tracing::error; + +use crate::clients::llm::{ + ChatMessage, LlmClient, StreamEvent, TokenUsage, ToolCall, ToolDefinition, +}; + +use super::AgentStreamEvent; + +/// 流式处理的结果 +#[derive(Debug)] +pub struct StreamOutput { + pub content: String, + pub reasoning: Option, + pub tool_calls: Option>, + pub usage: Option, + pub is_tool_call_step: bool, + pub status: StreamStatus, +} + +#[derive(Debug)] +pub enum StreamStatus { + Success, + Error(String), + Cancelled, +} + +/// 处理 LLM 流式响应。 +/// +/// 使用 tokio::select! 在流式读取和取消信号之间竞速。 +/// 实时发送 Thought/TextDelta SSE 事件给前端。 +pub async fn process_llm_stream( + llm: &LlmClient, + messages: &[ChatMessage], + tool_defs: &[ToolDefinition], + tx: &mpsc::UnboundedSender, + step: usize, + session_id: &str, + cancelled_runs: Arc>>, +) -> StreamOutput { + // 1. 发起 LLM 流式调用 + let mut stream_rx = match llm.chat_stream(messages, tool_defs).await { + Ok(rx) => rx, + Err(e) => { + error!("[Streaming] LLM stream 调用失败: {}", e); + let _ = tx.send(AgentStreamEvent::Error { + message: format!("大模型流式调用失败: {}", e), + }); + return StreamOutput { + content: String::new(), + reasoning: None, + tool_calls: None, + usage: None, + is_tool_call_step: false, + status: StreamStatus::Error(e.to_string()), + }; + } + }; + + let mut accumulated_content = String::new(); + let mut accumulated_reasoning = String::new(); + let mut accumulated_tool_calls: Option> = None; + let mut usage: Option = None; + let mut is_tool_call_step = false; + + // 2. 取消监视 future + let sid = session_id.to_string(); + let cancel_fut = async { + loop { + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + if let Ok(cancelled) = cancelled_runs.lock() { + if cancelled.contains(&sid) { + return; + } + } + } + }; + + // 3. tokio::select! 竞速：流式事件 vs 取消信号 + let mut cancel_pinned = Box::pin(cancel_fut); + let mut error_msg = None; + let mut was_cancelled = false; + + loop { + tokio::select! { + event_opt = stream_rx.recv() => { + match event_opt { + Some(event) => { + match event { + StreamEvent::ReasoningDelta(delta) => { + accumulated_reasoning.push_str(&delta); + let _ = tx.send(AgentStreamEvent::Thought { + content: accumulated_reasoning.clone(), + step, + }); + } + StreamEvent::TextDelta(delta) => { + accumulated_content.push_str(&delta); + if !is_tool_call_step { + let _ = tx.send(AgentStreamEvent::TextDelta { + content: delta, + }); + } + } + StreamEvent::ToolCallsComplete(tool_calls) => { + is_tool_call_step = true; + accumulated_tool_calls = Some(tool_calls); + } + StreamEvent::ToolCallDelta { .. } => {} + StreamEvent::Usage(u) => { + usage = Some(u); + } + StreamEvent::Done => { + break; + } + StreamEvent::Error(e) => { + error_msg = Some(e); + break; + } + } + } + None => break, + } + } + _ = &mut cancel_pinned => { + was_cancelled = true; + break; + } + } + } + + // 4. 构建结果 + let status = if was_cancelled { + StreamStatus::Cancelled + } else if let Some(e) = error_msg { + StreamStatus::Error(e) + } else { + StreamStatus::Success + }; + + let reasoning = if accumulated_reasoning.is_empty() { + None + } else { + Some(accumulated_reasoning) + }; + + StreamOutput { + content: accumulated_content, + reasoning, + tool_calls: accumulated_tool_calls, + usage, + is_tool_call_step, + status, + } +} diff --git a/src/agent/runtime/streaming_executor.rs b/src/agent/runtime/streaming_executor.rs new file mode 100644 index 0000000..9e2b0ca --- /dev/null +++ b/src/agent/runtime/streaming_executor.rs @@ -0,0 +1,315 @@ +// src/agent/runtime/streaming_executor.rs +// +// 流式工具执行器。 +// 参考 Claude Code StreamingToolExecutor 设计。 +// +// 当 LLM 流式输出 tool_use 块时，立即开始执行并发安全的工具。 +// 非并发安全的工具排队等待。结果按流中顺序 yield。 +// +// 功能： +// 1. 流式执行 — tool_use 到达时立即调度 +// 2. Sibling Abort — 副效应工具报错时中止兄弟并行执行 +// 3. Progress 流式 — 长时间操作可发送进度更新 + +use std::sync::Arc; +use tokio::sync::{broadcast, mpsc, oneshot}; +use tracing::{info, warn}; + +use super::partitioner::ToolPartitioner; +use crate::agent::tools::{ToolContext, ToolOutput, ToolRegistry}; + +/// 流式工具执行状态 +#[derive(Debug, Clone, PartialEq)] +pub enum TrackedToolStatus { + /// 工具调用已从 LLM 流中接收到 + Queued, + /// 正在执行中 + Executing, + /// 执行完成，等待 yield + Completed, + /// 结果已 yield 给调用方 + Yielded, +} + +/// 跟踪中的工具执行 +#[derive(Debug)] +struct TrackedTool { + tool_call_id: String, + tool_name: String, + args: serde_json::Value, + status: TrackedToolStatus, + /// 执行完成后的输出 + output: Option, + /// 取消通道（Sibling Abort 使用） + #[allow(dead_code)] + cancel_tx: Option>, +} + +/// Sibling Abort 原因 +#[derive(Debug, Clone)] +pub enum AbortReason { + /// 兄弟工具出错触发的级联取消 + SiblingError { description: String }, + /// 用户主动中断 + UserInterrupted, +} + +/// 流式工具执行器 +pub struct StreamingToolExecutor { + /// 所有跟踪中的工具 + tracked: Vec, + /// 工具注册表 + tool_registry: Arc, + /// 并发分区器（保留用于未来并发策略优化） + #[allow(dead_code)] + partitioner: ToolPartitioner, + /// 工具上下文 + tool_context: ToolContext, + /// Sibling Abort 广播通道 (tx) + abort_tx: broadcast::Sender, + /// Sibling Abort 广播通道 (rx) + abort_rx: broadcast::Receiver, + /// 当前是否已发生错误（触发 sibling abort） + has_errored: bool, + /// 出错工具的描述 + errored_tool_desc: String, + /// 下一个 stream_index + next_index: usize, + /// 最大工具输出字符数 + max_output_chars: usize, +} + +impl StreamingToolExecutor { + /// 创建新的流式执行器 + pub fn new( + tool_registry: Arc, + tool_context: ToolContext, + max_concurrency: usize, + max_output_chars: usize, + ) -> Self { + let (abort_tx, abort_rx) = broadcast::channel(16); + StreamingToolExecutor { + tracked: Vec::new(), + tool_registry, + partitioner: ToolPartitioner::new(max_concurrency), + tool_context, + abort_tx, + abort_rx, + has_errored: false, + errored_tool_desc: String::new(), + next_index: 0, + max_output_chars, + } + } + + /// 获取 abort 广播发送端（供外部注入取消信号） + pub fn abort_sender(&self) -> broadcast::Sender { + self.abort_tx.clone() + } + + /// 当 LLM 流产生一个新的 tool_use 时调用。 + /// + /// 返回 true 表示该工具已立即开始执行（并发安全），false 表示排队。 + pub fn on_tool_use(&mut self, call_id: String, name: String, args: serde_json::Value) -> bool { + let _index = self.next_index; + self.next_index += 1; + + let is_concurrency_safe = self + .tool_registry + .get(&name) + .map(|t| t.is_concurrency_safe(&args)) + .unwrap_or(false); + + let (cancel_tx, _cancel_rx) = oneshot::channel(); + + let tool = TrackedTool { + tool_call_id: call_id.clone(), + tool_name: name.clone(), + args: args.clone(), + status: TrackedToolStatus::Queued, + output: None, + cancel_tx: Some(cancel_tx), + }; + + self.tracked.push(tool); + + if is_concurrency_safe { + info!("[StreamingExecutor] 立即调度并发安全工具: {}", name); + self.try_execute_pending(); + true + } else { + info!("[StreamingExecutor] 排队非并发安全工具: {}", name); + false + } + } + + /// LLM 流结束后调用，执行所有剩余排队工具。 + pub async fn flush(&mut self) { + info!( + "[StreamingExecutor] flush: {} tracked, {} queued", + self.tracked.len(), + self.tracked + .iter() + .filter(|t| t.status == TrackedToolStatus::Queued) + .count() + ); + + // 将剩余排队的工具分批执行 + let queued: Vec = self + .tracked + .iter() + .enumerate() + .filter(|(_, t)| t.status == TrackedToolStatus::Queued) + .map(|(i, _)| i) + .collect(); + + for idx in queued { + self.execute_one(idx).await; + } + } + + /// 按流顺序获取下一个完成的结果（非阻塞）。 + pub fn next_result(&mut self) -> Option<(String, ToolOutput)> { + for tool in &mut self.tracked { + if tool.status == TrackedToolStatus::Completed { + tool.status = TrackedToolStatus::Yielded; + let output = tool + .output + .take() + .unwrap_or_else(|| ToolOutput::error("工具执行异常：无输出")); + return Some((tool.tool_call_id.clone(), output)); + } + } + None + } + + /// 是否有未 yield 的结果 + pub fn has_pending_results(&self) -> bool { + self.tracked + .iter() + .any(|t| t.status == TrackedToolStatus::Completed) + } + + /// 是否有未完成的工具 + pub fn has_unfinished(&self) -> bool { + self.tracked.iter().any(|t| { + t.status == TrackedToolStatus::Queued || t.status == TrackedToolStatus::Executing + }) + } + + /// 获取所有已完成的结果（包括已 yield 和未 yield 的） + pub fn all_results_mut(&mut self) -> Vec<(String, ToolOutput)> { + let mut results = Vec::new(); + for tool in &mut self.tracked { + if let Some(output) = tool.output.take() { + results.push((tool.tool_call_id.clone(), output)); + } + } + results + } + + // ── 内部方法 ── + + /// 尝试执行可执行的排队工具 + fn try_execute_pending(&mut self) { + // 简单策略：如果有正在执行的且它不是并发的，则不启动新的 + let has_executing = self + .tracked + .iter() + .any(|t| t.status == TrackedToolStatus::Executing); + + if !has_executing { + // 启动所有排队的并发安全工具 + let indices: Vec = self + .tracked + .iter() + .enumerate() + .filter(|(_, t)| t.status == TrackedToolStatus::Queued) + .map(|(i, _)| i) + .collect(); + + for idx in indices { + // 在同步上下文中只能标记状态，实际执行在 async 上下文中 + self.tracked[idx].status = TrackedToolStatus::Executing; + } + } + } + + /// 执行单个工具（内部辅助） + async fn execute_one(&mut self, idx: usize) { + if idx >= self.tracked.len() { + return; + } + + // 检查 sibling abort + if self.has_errored { + if let Ok(reason) = self.abort_rx.try_recv() { + let msg = match reason { + AbortReason::SiblingError { ref description } => { + format!("取消：并行工具 {} 出错，已级联取消", description) + } + AbortReason::UserInterrupted => "执行已被用户取消".to_string(), + }; + self.tracked[idx].output = Some(ToolOutput::error(msg)); + self.tracked[idx].status = TrackedToolStatus::Completed; + return; + } + } + + self.tracked[idx].status = TrackedToolStatus::Executing; + let tool_name = self.tracked[idx].tool_name.clone(); + + let output = match self.tool_registry.get(&tool_name) { + Some(tool) => { + let (progress_tx, _progress_rx) = mpsc::unbounded_channel(); + + let tool_args = self.tracked[idx].args.clone(); + let tool_fut = + tool.execute_with_progress(tool_args, &self.tool_context, Some(&progress_tx)); + + tool_fut.await + } + None => ToolOutput::error(format!("未知工具: {}", tool_name)), + }; + + // 检查是否需要触发 sibling abort + if output.is_error { + let causes_abort = self + .tool_registry + .get(&tool_name) + .map(|t| t.causes_sibling_abort()) + .unwrap_or(false); + + if causes_abort { + warn!( + "[StreamingExecutor] 工具 {} 出错，触发 sibling abort", + tool_name + ); + self.has_errored = true; + self.errored_tool_desc = tool_name.clone(); + let _ = self.abort_tx.send(AbortReason::SiblingError { + description: tool_name.clone(), + }); + } + } + + // 截断输出 + let truncated = if output.content.len() > self.max_output_chars { + let t: String = output.content.chars().take(self.max_output_chars).collect(); + ToolOutput { + content: format!( + "{}...\n[输出已截断，原始长度: {} 字符]", + t, + output.content.len() + ), + is_error: output.is_error, + metadata: output.metadata, + } + } else { + output + }; + + self.tracked[idx].output = Some(truncated); + self.tracked[idx].status = TrackedToolStatus::Completed; + } +} diff --git a/src/agent/runtime/system_prompt.rs b/src/agent/runtime/system_prompt.rs new file mode 100644 index 0000000..7424b59 --- /dev/null +++ b/src/agent/runtime/system_prompt.rs @@ -0,0 +1,104 @@ +// src/agent/runtime/system_prompt.rs +// +// 模块化系统提示词组装 — 参考 Claude Code s10 System Prompt 设计。 +// +// 将硬编码的提示词拆分为独立 section，运行时按需拼接。 +// 静态 section 在前以最大化 Anthropic prompt cache 命中率。 + +/// 系统提示词组装器 +pub struct SystemPrompt { + sections: Vec<(&'static str, String)>, +} + +impl SystemPrompt { + pub fn new() -> Self { + SystemPrompt { + sections: Vec::new(), + } + } + + /// 添加一个 section（先添加的排在前面） + pub fn add_section(&mut self, name: &'static str, content: String) { + self.sections.push((name, content)); + } + + /// 组装最终的系统提示词（section 间用双换行分隔） + pub fn assemble(&self) -> String { + self.sections + .iter() + .map(|(_, content)| content.as_str()) + .collect::>() + .join("\n\n") + } + + /// 是否为空 + pub fn is_empty(&self) -> bool { + self.sections.is_empty() + } + + /// 获取 section 数量 + pub fn section_count(&self) -> usize { + self.sections.len() + } +} + +impl Default for SystemPrompt { + fn default() -> Self { + Self::new() + } +} + +/// 静态身份 section（始终加载，最大化 prompt cache 命中率） +pub const IDENTITY_SECTION: &str = "\ +你是一位专业的天体物理学研究助手，具备丰富的天文学知识。"; + +/// 静态核心原则 section +pub const PRINCIPLES_SECTION: &str = "\ +核心原则： +1. 主动使用工具搜索最新文献，不要仅凭训练数据回答。 +2. 优先使用本地资源（get_paper_content / rag_search），必要时再检索新文献。 +3. 收集到足够信息后立即给出最终答案，避免无意义的重复工具调用。 +4. 回答时引用具体文献来源，使用 ADS bibcode 标注。 +5. 对于数学公式，使用标准 LaTeX 格式。 +6. 用中文回答，保持科学术语的准确性（可附带英文原文）。 +7. 对于复杂任务（如文献综述），调用 load_skill 获取方法论指引，再用 todo_write 制定计划。 +8. 如果某个工具调用失败，不要用相同参数重试，尝试换一种方式或工具。 +9. 任务状态会在每轮开始时从数据库恢复，请基于最新状态继续工作。"; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_assemble_empty() { + let sp = SystemPrompt::new(); + assert_eq!(sp.assemble(), ""); + } + + #[test] + fn test_assemble_multiple_sections() { + let mut sp = SystemPrompt::new(); + sp.add_section("a", "Section A".to_string()); + sp.add_section("b", "Section B".to_string()); + let result = sp.assemble(); + assert_eq!(result, "Section A\n\nSection B"); + } + + #[test] + fn test_static_sections_first() { + let mut sp = SystemPrompt::new(); + sp.add_section("identity", "I am".to_string()); + sp.add_section("dynamic", "Tools: ...".to_string()); + let result = sp.assemble(); + assert!(result.starts_with("I am")); + assert!(result.contains("Tools: ...")); + } + + #[test] + fn test_section_count() { + let mut sp = SystemPrompt::new(); + assert_eq!(sp.section_count(), 0); + sp.add_section("a", "A".to_string()); + assert_eq!(sp.section_count(), 1); + } +} diff --git a/src/agent/runtime/token_budget.rs b/src/agent/runtime/token_budget.rs new file mode 100644 index 0000000..6b854b8 --- /dev/null +++ b/src/agent/runtime/token_budget.rs @@ -0,0 +1,332 @@ +// src/agent/runtime/token_budget.rs +// +// Token 预算管理。 +// 参考 Claude Code TokenBudget 设计。 +// 软限制：接近上限时注入 nudging 消息提醒模型。 +// 硬限制：达到上限时触发强制压缩或终止。 + +/// Token 预算管理器。 +/// +/// 参考 Claude Code TokenBudget 设计，增加： +/// - 多级渐进式 nudge（near_soft / over_soft / over_hard） +/// - Diminishing returns 检测（防止模型在死循环中消耗预算） +/// - Continuation 计数追踪 +#[derive(Debug, Clone)] +pub struct TokenBudget { + /// 软限制（触发 nudging 提醒） + pub soft_limit: usize, + /// 硬限制（触发强制动作） + pub hard_limit: usize, + /// 已消耗输入 tokens + pub input_tokens_spent: usize, + /// 已消耗输出 tokens + pub output_tokens_spent: usize, + /// 延续次数（每个 ReAct 步骤递增） + pub continuation_count: usize, + /// 上次检查时的总消耗（用于 diminishing returns 检测） + last_total_spent: usize, + /// 连续无进展次数 + consecutive_no_progress: usize, + /// 是否已触发 diminishing returns + pub diminishing_returns: bool, +} + +impl TokenBudget { + /// 创建预算管理器 + pub fn new(soft_limit: usize, hard_limit: usize) -> Self { + TokenBudget { + soft_limit, + hard_limit, + input_tokens_spent: 0, + output_tokens_spent: 0, + continuation_count: 0, + last_total_spent: 0, + consecutive_no_progress: 0, + diminishing_returns: false, + } + } + + /// 记录输入 token 消耗 + pub fn spend_input(&mut self, tokens: usize) { + self.input_tokens_spent += tokens; + } + + /// 记录输出 token 消耗 + pub fn spend_output(&mut self, tokens: usize) { + self.output_tokens_spent += tokens; + } + + /// 总消耗 + pub fn total_spent(&self) -> usize { + self.input_tokens_spent + self.output_tokens_spent + } + + /// 记录一次延续（每个 ReAct 步骤调用一次）。 + /// 同时检查 diminishing returns。 + pub fn record_continuation(&mut self) -> bool { + self.continuation_count += 1; + self.check_diminishing_returns_inner() + } + + /// 检测 diminishing returns — 模型在同一问题上打转而不产生实质进展。 + /// + /// 触发条件：3 次以上延续，且连续 2 次检查的 token 增量 < 500。 + /// 返回 true 表示已检测到无进展循环。 + pub fn check_diminishing_returns(&mut self) -> bool { + self.check_diminishing_returns_inner() + } + + fn check_diminishing_returns_inner(&mut self) -> bool { + if self.diminishing_returns { + return true; // 已触发过，保持状态 + } + if self.continuation_count < 3 { + return false; + } + let delta = self.total_spent().saturating_sub(self.last_total_spent); + self.last_total_spent = self.total_spent(); + if delta < 500 { + self.consecutive_no_progress += 1; + if self.consecutive_no_progress >= 2 { + self.diminishing_returns = true; + return true; + } + } else { + self.consecutive_no_progress = 0; + } + false + } + + /// 是否接近软限制（超过 80%） + pub fn near_soft_limit(&self) -> bool { + if self.soft_limit == 0 { + return false; + } + self.total_spent() >= self.soft_limit * 8 / 10 + } + + /// 是否超过软限制 + pub fn over_soft_limit(&self) -> bool { + self.total_spent() >= self.soft_limit + } + + /// 是否超过硬限制 + pub fn over_hard_limit(&self) -> bool { + self.total_spent() >= self.hard_limit + } + + /// 已使用预算的百分比（相对于软限制） + pub fn usage_pct(&self) -> u32 { + if self.soft_limit == 0 { + return 0; + } + (self.total_spent() * 100 / self.soft_limit) as u32 + } + + /// 剩余可用 tokens（硬限制 - 已消耗） + pub fn remaining(&self) -> usize { + self.hard_limit.saturating_sub(self.total_spent()) + } + + /// 生成渐进式 nudging 提醒消息。 + /// + /// 三级： + /// - near_soft (80-99%): 温和提醒 + /// - over_soft (100-硬): 明确警告 + /// - over_hard (>硬限制): 强制完成 + /// - diminishing_returns: 要求最终答案 + pub fn nudge_message(&self) -> Option { + if self.diminishing_returns { + return Some( + "⚠️ 已检测到重复操作模式 — 后续步骤未产生新信息。\ + 请基于已收集的全部信息直接给出最终答案，不要再调用工具。" + .to_string(), + ); + } + + if self.over_hard_limit() { + Some(format!( + "🔴 Token 预算已耗尽（{}/{} tokens, {}%）。\ + 请立即总结当前发现并给出最终答案，不要再调用任何工具。", + self.total_spent(), + self.hard_limit, + self.usage_pct() + )) + } else if self.over_soft_limit() { + let pct = self.usage_pct(); + Some(format!( + "🟡 Token 预算警告：已使用 {}/{} tokens ({}%)。\ + 请尽快总结关键发现并给出最终答案。如非必要，不要再调用工具。", + self.total_spent(), + self.soft_limit, + pct + )) + } else if self.near_soft_limit() { + let pct = self.usage_pct(); + Some(format!( + "💡 Token 预算提示：已使用 {}/{} tokens ({}%)。\ + 请注意控制后续步骤的深度，优先处理最重要的发现。", + self.total_spent(), + self.soft_limit, + pct + )) + } else { + None + } + } + + /// 将硬限制提升到指定值（用于 error recovery 中的 escalate 步骤） + pub fn escalate_hard_limit(&mut self, new_limit: usize) { + self.hard_limit = new_limit; + } +} + +impl Default for TokenBudget { + fn default() -> Self { + Self::new(32_000, 40_000) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_spend_tracking() { + let mut budget = TokenBudget::new(1000, 2000); + budget.spend_input(500); + budget.spend_output(300); + assert_eq!(budget.total_spent(), 800); + } + + #[test] + fn test_near_soft_limit() { + let mut budget = TokenBudget::new(1000, 2000); + assert!(!budget.near_soft_limit()); + budget.spend_input(850); // 85% > 80% + assert!(budget.near_soft_limit()); + } + + #[test] + fn test_over_hard_limit() { + let mut budget = TokenBudget::new(1000, 2000); + budget.spend_input(2100); + assert!(budget.over_hard_limit()); + assert_eq!(budget.remaining(), 0); + } + + #[test] + fn test_nudge_message_at_soft_limit() { + let mut budget = TokenBudget::new(1000, 2000); + budget.spend_input(1000); // exactly at soft limit + let msg = budget.nudge_message(); + assert!(msg.is_some()); + assert!(msg.unwrap().contains("Token 预算警告")); + } + + #[test] + fn test_nudge_message_at_hard_limit() { + let mut budget = TokenBudget::new(1000, 2000); + budget.spend_input(2000); // at hard limit + let msg = budget.nudge_message(); + assert!(msg.is_some()); + assert!(msg.unwrap().contains("已耗尽")); + } + + #[test] + fn test_no_nudge_when_under_limit() { + let budget = TokenBudget::new(1000, 2000); + assert!(budget.nudge_message().is_none()); + } + + #[test] + fn test_escalate_hard_limit() { + let mut budget = TokenBudget::new(1000, 2000); + budget.escalate_hard_limit(64000); + assert_eq!(budget.hard_limit, 64000); + } + + #[test] + fn test_default_budget() { + let budget = TokenBudget::default(); + assert_eq!(budget.soft_limit, 32_000); + assert_eq!(budget.hard_limit, 40_000); + } + + #[test] + fn test_near_soft_limit_nudge() { + let mut budget = TokenBudget::new(1000, 2000); + budget.spend_input(850); // 85% — near soft + let msg = budget.nudge_message(); + assert!(msg.is_some()); + assert!(msg.unwrap().contains("Token 预算提示")); + } + + #[test] + fn test_diminishing_returns_not_triggered_early() { + let mut budget = TokenBudget::new(1000, 2000); + // < 3 continuations — should not trigger + budget.record_continuation(); + assert!(!budget.diminishing_returns); + budget.record_continuation(); + assert!(!budget.diminishing_returns); + } + + #[test] + fn test_diminishing_returns_triggers_after_stagnation() { + let mut budget = TokenBudget::new(1000, 5000); + // First 3 continuations establish baseline (all with 0 spending) + // After the 3rd, consecutive_no_progress becomes 1 (delta=0 < 500) + for _ in 0..3 { + budget.record_continuation(); + } + assert!(!budget.diminishing_returns); + + // 4th continuation with tiny spending — 2nd consecutive <500 + budget.spend_input(100); + budget.record_continuation(); + // consecutive_no_progress is now 2 → triggered + assert!(budget.diminishing_returns); + } + + #[test] + fn test_diminishing_returns_resets_on_progress() { + let mut budget = TokenBudget::new(1000, 5000); + // Establish baseline with spending BEFORE the 3-continuation threshold + budget.spend_input(2000); + for _ in 0..3 { + budget.record_continuation(); + } + // delta = total - last_total. After the first check, last_total is set. + // total_spent=2000, last_total=2000 after first continuation ≥ 3 + + budget.spend_input(100); // total=2100 + budget.record_continuation(); // delta=100 < 500, cons=1 + assert!(!budget.diminishing_returns); + + budget.spend_input(600); // total=2700 + budget.record_continuation(); // delta=600 >= 500, cons resets to 0 + assert!(!budget.diminishing_returns); + + budget.spend_input(100); // total=2800 + budget.record_continuation(); // delta=100 < 500, cons=1 + assert!(!budget.diminishing_returns); + } + + #[test] + fn test_diminishing_returns_nudge_message() { + let mut budget = TokenBudget::new(1000, 5000); + budget.diminishing_returns = true; + let msg = budget.nudge_message().unwrap(); + assert!(msg.contains("重复操作")); + } + + #[test] + fn test_usage_pct() { + let mut budget = TokenBudget::new(1000, 2000); + budget.spend_input(500); + assert_eq!(budget.usage_pct(), 50); + budget.spend_output(300); + assert_eq!(budget.usage_pct(), 80); + } +} diff --git a/src/agent/skills.rs b/src/agent/skills.rs new file mode 100644 index 0000000..c1f3f95 --- /dev/null +++ b/src/agent/skills.rs @@ -0,0 +1,846 @@ +// src/agent/skills.rs +// +// 两层技能加载系统（参考 Claude Code src/skills/ + src/tools/SkillTool/ 设计）： +// Layer 1 — system-reminder 注入：每轮动态列出 skill 名称（~20 tokens/skill） +// Layer 2 — LoadSkillTool：LLM 按需调用，注入完整 skill 内容（~2000 tokens/skill） +// +// Skill 文件格式（对齐 Claude Code 的目录约定）： +// skills/{skill-name}/SKILL.md ← 必须是目录 + SKILL.md +// +// SKILL.md 内容（Markdown + YAML frontmatter）： +// --- +// name: skill-name +// description: 一句话描述 +// context: inline | fork +// allowed-tools: +// - bash +// - read +// when_to_use: 何时自动触发 +// model: haiku | sonnet | opus | inherit +// disable-model-invocation: false +// user-invocable: true +// paths: +// - "*.rs" +// --- +// +// # Skill 正文 +// 详细内容... + +use serde::Deserialize; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::{Arc, RwLock}; +use std::time::SystemTime; +use tracing::{info, warn}; + +// ── Frontmatter ──────────────────────────────────────────────────────────── + +/// Skill 的 YAML frontmatter 结构（serde_yaml 解析）。 +/// 仅声明与平台相关的字段；未知字段自动忽略。 +#[derive(Debug, Clone, Deserialize, Default)] +pub struct SkillFrontmatter { + #[serde(default)] + pub name: Option, + #[serde(default)] + pub description: Option, + /// 执行模式：inline | fork + #[serde(default)] + pub context: Option, + /// 工具白名单 + #[serde(rename = "allowed-tools", default)] + pub allowed_tools: Option>, + /// 推荐模型 + #[serde(default)] + pub model: Option, + /// 参数提示 + #[serde(rename = "argument-hint", default)] + pub argument_hint: Option, + /// 使用场景说明 + #[serde(rename = "when_to_use", default)] + pub when_to_use: Option, + /// 禁止模型通过 Skill tool 自动调用 + #[serde(rename = "disable-model-invocation", default)] + pub disable_model_invocation: Option, + /// 用户是否可通过 /skill-name 手动调用 + #[serde(rename = "user-invocable", default)] + pub user_invocable: Option, + /// 版本号 + #[serde(default)] + pub version: Option, + /// 条件激活的 glob 模式 + #[serde(default)] + pub paths: Option>, + /// fork 模式下的 agent 类型 + #[serde(default)] + pub agent: Option, + /// fork 模式下的 effort 级别 + #[serde(default)] + pub effort: Option, +} + +impl SkillFrontmatter { + /// 校验必填/推荐字段，返回警告列表 + pub fn validate(&self, skill_name: &str) -> Vec { + let mut warnings = Vec::new(); + if self.description.is_none() { + warnings.push(format!("Skill '{}' 缺少 description 字段", skill_name)); + } + if let Some(ref ctx) = self.context { + if ctx != "inline" && ctx != "fork" { + warnings.push(format!( + "Skill '{}' 的 context 值无效 '{}'，应为 inline 或 fork", + skill_name, ctx + )); + } + } + warnings + } +} + +// ── Skill Data Structures ────────────────────────────────────────────────── + +/// Skill 元信息（Layer 1：出现在 skill 列表中） +#[derive(Debug, Clone)] +pub struct SkillMeta { + pub name: String, + pub description: String, + /// 执行模式（inline / fork） + pub context: Option, + /// 工具白名单（空 Vec 表示无限制） + pub allowed_tools: Vec, + /// 使用场景说明 + pub when_to_use: Option, + /// 是否禁止模型通过 Skill tool 调用 + pub disable_model_invocation: bool, + /// 是否允许用户通过 /skill-name 手动调用 + pub user_invocable: bool, + /// 条件激活的 glob 模式（空 Vec 表示始终激活） + pub paths: Vec, +} + +/// 完整的 Skill（Layer 2：LLM 调用 load_skill 时注入） +#[derive(Debug, Clone)] +pub struct Skill { + pub meta: SkillMeta, + pub body: String, + /// Skill 所在目录，用于 ${SKILL_DIR} 变量替换 + pub skill_dir: PathBuf, +} + +// ── Skill Registry (Caching Layer) ───────────────────────────────────────── + +/// Skill 使用统计 +#[derive(Debug, Clone, Default)] +pub struct SkillUsageStat { + pub invoke_count: u64, + pub last_used_at: Option>, +} + +/// Skill 注册表 — 缓存已加载的 skills，支持 mtime 增量刷新。 +/// +/// 使用方式： +/// ```ignore +/// let registry = SkillRegistry::new(skills_dir); +/// registry.refresh()?; +/// let reminder = registry.build_reminder(); +/// ``` +#[derive(Debug)] +pub struct SkillRegistry { + skills_dir: PathBuf, + skills: Vec, + /// 上次扫描时 skills_dir 的 mtime（用于增量刷新） + last_scan_mtime: Option, + /// 使用统计（按 skill name 索引） + usage_stats: HashMap, +} + +impl SkillRegistry { + /// 创建新的 skill 注册表（不执行初始扫描，调用 `refresh()` 触发） + pub fn new(skills_dir: PathBuf) -> Self { + SkillRegistry { + skills_dir, + skills: Vec::new(), + last_scan_mtime: None, + usage_stats: HashMap::new(), + } + } + + /// 检查是否需要重新扫描（目录 mtime 变化或首次加载） + pub fn needs_refresh(&self) -> bool { + match dir_modified_time(&self.skills_dir) { + Some(current_mtime) => match self.last_scan_mtime { + Some(last) => current_mtime > last, + None => true, + }, + None => !self.skills.is_empty(), // 目录消失但还有缓存 → 保持缓存 + } + } + + /// 扫描 skills 目录并加载/更新所有 skills。 + /// 始终执行完整重载（简单可靠，skill 数量少时成本可忽略）。 + pub fn refresh(&mut self) { + let current_mtime = dir_modified_time(&self.skills_dir); + + let mut new_skills = Vec::new(); + for (dir_name, skill_md_path) in discover_skill_dirs(&self.skills_dir) { + match load_skill_from_path(&skill_md_path, &dir_name) { + Ok(skill) => new_skills.push(skill), + Err(e) => { + warn!("[SkillRegistry] 加载 skill '{}' 失败: {}", dir_name, e); + } + } + } + + // 按使用频率排序：常用 skill 排前面 + new_skills.sort_by(|a, b| { + let score_a = self.usage_score(&a.meta.name); + let score_b = self.usage_score(&b.meta.name); + // 降序排列（高分在前） + score_b + .partial_cmp(&score_a) + .unwrap_or(std::cmp::Ordering::Equal) + }); + + let count = new_skills.len(); + self.skills = new_skills; + self.last_scan_mtime = current_mtime; + info!("[SkillRegistry] 已刷新 {} 个 skill", count); + } + + /// 获取所有 skill 的元信息列表 + pub fn list_skills(&self) -> Vec { + self.skills.iter().map(|s| s.meta.clone()).collect() + } + + /// 按名称获取完整 skill + pub fn get_skill(&self, name: &str) -> Option<&Skill> { + self.skills.iter().find(|s| s.meta.name == name) + } + + /// 检查是否有可用 skills + pub fn is_empty(&self) -> bool { + self.skills.is_empty() + } + + /// 技能数量 + pub fn len(&self) -> usize { + self.skills.len() + } + + // ── 使用统计 ── + + /// 计算 skill 的使用评分（指数衰减，7 天半衰期） + fn usage_score(&self, name: &str) -> f64 { + match self.usage_stats.get(name) { + Some(stat) => { + let count_weight = (stat.invoke_count as f64).ln_1p(); // log(1 + count) + let recency_weight = match stat.last_used_at { + Some(last) => { + let age_hours = chrono::Utc::now() + .signed_duration_since(last) + .num_hours() + .max(0) as f64; + // 7 天半衰期 + 0.5_f64.powf(age_hours / (7.0 * 24.0)) + } + None => 0.1, + }; + count_weight * recency_weight + } + None => 0.0, + } + } + + /// 记录一次 skill 调用 + pub fn record_usage(&mut self, name: &str) { + let stat = self.usage_stats.entry(name.to_string()).or_default(); + stat.invoke_count += 1; + stat.last_used_at = Some(chrono::Utc::now()); + } + + /// 获取所有使用统计的快照 + pub fn usage_stats(&self) -> &HashMap { + &self.usage_stats + } + + // ── System Prompt 构建 ── + + /// 构建 skill 列表的 system-reminder 消息（Layer 1，参考 Claude Code）。 + /// 使用结构化 XML 标签，包含名称和描述。 + pub fn build_reminder(&self) -> Option { + if self.skills.is_empty() { + return None; + } + + let invocable: Vec<&Skill> = self + .skills + .iter() + .filter(|s| !s.meta.disable_model_invocation && s.meta.user_invocable) + .collect(); + + if invocable.is_empty() { + return None; + } + + let mut lines = vec![ + "".to_string(), + "The following skills are available for use with the Skill tool:".to_string(), + ]; + + for skill in &invocable { + let desc = match &skill.meta.when_to_use { + Some(wtu) => format!("{} - {}", skill.meta.description, wtu), + None => skill.meta.description.clone(), + }; + let context_marker = match skill.meta.context.as_deref() { + Some("fork") => " [fork]", + _ => "", + }; + lines.push(format!("- {}: {}{}", skill.meta.name, desc, context_marker)); + } + + lines.push( + "When a skill matches the user's request, invoke load_skill BEFORE generating any other response about the task.".to_string(), + ); + lines.push( + "If you see a tag in the current conversation turn, the skill has ALREADY been loaded - follow the instructions directly instead of calling load_skill again.".to_string(), + ); + lines.push("".to_string()); + + Some(lines.join("\n")) + } + + /// 构建 LoadSkillTool 的 description（动态生成，列出可用 skills） + pub fn build_tool_description(&self) -> String { + if self.skills.is_empty() { + return "加载指定的领域技能完整内容。当前没有可用的技能。".to_string(); + } + + let invocable: Vec<&Skill> = self + .skills + .iter() + .filter(|s| !s.meta.disable_model_invocation) + .collect(); + + if invocable.is_empty() { + return "加载指定的领域技能完整内容。当前没有可用的技能。".to_string(); + } + + let mut desc = String::from("加载指定的领域技能完整内容。可用技能：\n"); + for skill in &invocable { + let mode = match skill.meta.context.as_deref() { + Some("fork") => "[子代理执行] ", + _ => "", + }; + let wtu = match &skill.meta.when_to_use { + Some(w) => format!(" - {}", w), + None => String::new(), + }; + desc.push_str(&format!( + "- {}: {}{}{}\n", + skill.meta.name, mode, skill.meta.description, wtu + )); + } + desc + } + + // ── 文件监听 (Hot Reload) ── + + /// 启动文件监听器，在 skills 目录变更时自动刷新缓存。 + /// + /// 返回一个 `JoinHandle`，调用方可以 `await` 它（通常运行到程序退出）。 + /// 内部使用 debounce：300ms 内的连续变更合并为一次刷新。 + #[cfg(not(test))] + pub fn start_watcher(self_arc: Arc>) -> std::thread::JoinHandle<()> { + use notify::{RecursiveMode, Watcher}; + use std::time::Duration; + + let skills_dir = match self_arc.read() { + Ok(r) => r.skills_dir.clone(), + Err(_) => { + warn!("[SkillRegistry] RwLock 异常，无法启动文件监视器"); + return std::thread::spawn(|| {}); + } + }; + + std::thread::spawn(move || { + let (tx, rx) = std::sync::mpsc::channel(); + + let mut watcher = + match notify::recommended_watcher(move |res: notify::Result| { + if let Ok(event) = res { + // 只关心 SKILL.md 相关的变更 + let is_skill_change = event + .paths + .iter() + .any(|p| p.file_name().map(|n| n == "SKILL.md").unwrap_or(false)); + if is_skill_change { + let _ = tx.send(()); + } + } + }) { + Ok(w) => w, + Err(e) => { + warn!("[SkillRegistry] 无法创建文件监听器: {}", e); + return; + } + }; + + if let Err(e) = watcher.watch(&skills_dir, RecursiveMode::Recursive) { + warn!("[SkillRegistry] 无法监听 skills 目录: {}", e); + return; + } + + info!("[SkillRegistry] 文件监听已启动: {}", skills_dir.display()); + + // 300ms debounce：收集快速连续的事件 + while let Ok(()) = rx.recv() { + // 等待 debounce 窗口 + while rx.recv_timeout(Duration::from_millis(300)).is_ok() {} + info!("[SkillRegistry] 检测到 skill 文件变更，自动刷新"); + if let Ok(mut registry) = self_arc.write() { + let reg: &mut SkillRegistry = &mut registry; + reg.refresh(); + } + } + // Channel closed, watcher dropped + }) + } + + /// 文件监听器的空实现（测试模式下不启动线程） + #[cfg(test)] + pub fn start_watcher(_self_arc: Arc>) -> std::thread::JoinHandle<()> { + std::thread::spawn(|| {}) + } + + // ── 条件 Skill (Paths-based Activation) ── + + /// 根据访问的文件路径激活匹配的条件 skill。 + /// + /// 条件 skill 在其 `paths` frontmatter 中声明了 glob 模式。 + /// 当 Agent 访问（Read/Edit/Grep）匹配文件时调用此方法将其激活。 + /// + /// 返回新激活的 skill 名称列表。 + pub fn activate_conditional_for_paths(&mut self, file_paths: &[&str]) -> Vec { + let mut activated = Vec::new(); + + for skill in &mut self.skills { + // 只处理有 paths 限制且当前未激活的 + if skill.meta.paths.is_empty() || !skill.meta.disable_model_invocation { + continue; + } + + // 检查是否有任何文件路径匹配 + let matches = file_paths.iter().any(|fp| { + skill + .meta + .paths + .iter() + .any(|pattern| glob_match_simple(pattern, fp)) + }); + + if matches { + skill.meta.disable_model_invocation = false; + activated.push(skill.meta.name.clone()); + info!( + "[SkillRegistry] 条件 skill '{}' 已激活 (paths: {:?})", + skill.meta.name, skill.meta.paths + ); + } + } + + activated + } + + /// 获取与给定文件路径匹配的所有 skill（用于 LLM 上下文提示）。 + pub fn matching_skills_for_paths(&self, file_paths: &[&str]) -> Vec { + self.skills + .iter() + .filter(|s| { + !s.meta.paths.is_empty() + && file_paths.iter().any(|fp| { + s.meta + .paths + .iter() + .any(|pattern| glob_match_simple(pattern, fp)) + }) + }) + .map(|s| s.meta.clone()) + .collect() + } +} + +// ── Glob 匹配 (简化实现，避免引入完整 glob 库的运行时开销) ── + +/// 简化的 glob 模式匹配。 +/// +/// 支持的语法： +/// - `*` 匹配任意非 '/' 字符序列 +/// - `**` 匹配任意字符（含 '/'） +/// - `?` 匹配单个非 '/' 字符 +/// - 其他字符按字面匹配 +fn glob_match_simple(pattern: &str, path: &str) -> bool { + // 标准化：统一使用 '/' 作为路径分隔符 + let pattern = pattern.replace('\\', "/"); + let path = path.replace('\\', "/"); + + // 使用 glob crate 进行匹配 + // 降级方案：简单的后缀/前缀匹配 + if pattern.contains('*') || pattern.contains('?') { + // 尝试使用 glob crate + match glob::Pattern::new(&pattern) { + Ok(pat) => pat.matches(&path), + Err(_) => { + // 降级：简单包含匹配 + let simple = pattern.replace(['*', '?'], ""); + path.contains(&simple) + } + } + } else { + // 无通配符：精确匹配文件名或后缀 + path == pattern || path.ends_with(&format!("/{}", pattern)) + } +} + +// ── File I/O Helpers ─────────────────────────────────────────────────────── + +/// 获取目录的修改时间 +fn dir_modified_time(path: &Path) -> Option { + std::fs::metadata(path).ok().and_then(|m| m.modified().ok()) +} + +/// 扫描 skills 目录，查找所有 `{name}/SKILL.md` 子目录。 +fn discover_skill_dirs(skills_dir: &Path) -> Vec<(String, PathBuf)> { + let mut skills = Vec::new(); + + let entries = match std::fs::read_dir(skills_dir) { + Ok(e) => e, + Err(_) => return skills, + }; + + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_dir() { + continue; + } + let skill_md = path.join("SKILL.md"); + if skill_md.exists() { + let name = path + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("unknown") + .to_string(); + skills.push((name, skill_md)); + } + } + + skills +} + +/// 解析 SKILL.md 文件：提取 YAML frontmatter 和 Markdown 正文。 +/// +/// Frontmatter 格式： +/// --- +/// key: value +/// --- +/// 正文 +fn parse_skill_file(raw: &str) -> Result<(SkillFrontmatter, String), String> { + let content = raw.trim(); + + if let Some(rest) = content.strip_prefix("---") { + // 查找闭合的 --- + if let Some((fm_text, body_text)) = rest.split_once("\n---") { + let frontmatter: SkillFrontmatter = serde_yaml::from_str(fm_text) + .map_err(|e| format!("YAML frontmatter 解析失败: {}", e))?; + let body = body_text.trim().to_string(); + return Ok((frontmatter, body)); + } + } + + // 没有 frontmatter，整个内容作为正文 + Ok((SkillFrontmatter::default(), content.to_string())) +} + +/// 从文件路径加载完整的 Skill +fn load_skill_from_path(skill_md_path: &Path, dir_name: &str) -> Result { + let raw = std::fs::read_to_string(skill_md_path).map_err(|e| format!("无法读取文件: {}", e))?; + + let (frontmatter, body) = parse_skill_file(&raw)?; + + // 校验并记录警告 + let warnings = frontmatter.validate(dir_name); + for w in &warnings { + warn!("[Skills] {}", w); + } + + let name = frontmatter.name.unwrap_or_else(|| dir_name.to_string()); + let description = frontmatter + .description + .unwrap_or_else(|| "(无描述)".to_string()); + let context = frontmatter.context.filter(|c| !c.is_empty()); + let allowed_tools = frontmatter.allowed_tools.unwrap_or_default(); + let when_to_use = frontmatter.when_to_use.filter(|w| !w.is_empty()); + let disable_model_invocation = frontmatter.disable_model_invocation.unwrap_or(false); + let user_invocable = frontmatter.user_invocable.unwrap_or(true); + let paths = frontmatter.paths.unwrap_or_default(); + let skill_dir = skill_md_path + .parent() + .unwrap_or(Path::new(".")) + .to_path_buf(); + + info!( + "[Skills] 已加载 skill: {} (context={:?}, allowed_tools={:?}, paths={:?})", + name, context, allowed_tools, paths + ); + + Ok(Skill { + meta: SkillMeta { + name, + description, + context, + allowed_tools, + when_to_use, + disable_model_invocation, + user_invocable, + paths, + }, + body, + skill_dir, + }) +} + +// ── Backward-Compatible Public API ───────────────────────────────────────── +// +// 这些函数保留用于外部调用（如 health check、CLI 工具）， +// 核心路径（Runtime、LoadSkillTool）应使用 SkillRegistry。 + +/// 从 skills 目录加载单个 skill（不经过缓存，直接读文件）。 +pub fn load_skill_direct(skills_dir: &Path, skill_name: &str) -> Option { + let file_path = skills_dir.join(skill_name).join("SKILL.md"); + match load_skill_from_path(&file_path, skill_name) { + Ok(skill) => Some(skill), + Err(e) => { + warn!("[Skills] 直接加载 skill '{}' 失败: {}", skill_name, e); + None + } + } +} + +// ── Variable Substitution ────────────────────────────────────────────────── + +/// 在 skill 正文中执行变量替换。 +/// +/// 支持的变量： +/// - `${SKILL_DIR}` → skill 所在目录的绝对路径 +/// - `${SESSION_ID}` → 当前会话 ID +pub fn substitute_variables(body: &str, skill_dir: &Path, session_id: Option<&str>) -> String { + let mut result = body.to_string(); + + // ${SKILL_DIR} — skill 所在目录 + if let Some(dir_str) = skill_dir.to_str() { + result = result.replace("${SKILL_DIR}", dir_str); + } + + // ${SESSION_ID} + if let Some(sid) = session_id { + result = result.replace("${SESSION_ID}", sid); + } else { + result = result.replace("${SESSION_ID}", ""); + } + + result +} + +// ── Tests ────────────────────────────────────────────────────────────────── + +#[cfg(test)] +mod tests { + use super::*; + + // ── Frontmatter 解析 ── + + #[test] + fn test_parse_basic_frontmatter() { + let raw = "---\nname: test-skill\ndescription: A test skill\n---\n\n# Body\nSome content"; + let (fm, body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.name.unwrap(), "test-skill"); + assert_eq!(fm.description.unwrap(), "A test skill"); + assert!(body.contains("# Body")); + assert!(body.contains("Some content")); + } + + #[test] + fn test_parse_frontmatter_with_list() { + let raw = "---\nname: test\ndescription: Test\nallowed-tools:\n- bash\n- read\n---\nBody"; + let (fm, body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.name.unwrap(), "test"); + let tools = fm.allowed_tools.unwrap(); + assert_eq!(tools, vec!["bash", "read"]); + assert_eq!(body, "Body"); + } + + #[test] + fn test_parse_frontmatter_no_fm() { + let raw = "# Just a header\nSome content"; + let (fm, body) = parse_skill_file(raw).unwrap(); + assert!(fm.name.is_none()); + assert_eq!(body, raw); + } + + #[test] + fn test_parse_frontmatter_context_fork() { + let raw = "---\nname: heavy\ndescription: Heavy skill\ncontext: fork\n---\nBody"; + let (fm, _body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.context.unwrap(), "fork"); + } + + #[test] + fn test_parse_frontmatter_boolean_fields() { + let raw = "---\nname: test\ndescription: Test\ndisable-model-invocation: true\nuser-invocable: false\n---\nBody"; + let (fm, _body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.disable_model_invocation, Some(true)); + assert_eq!(fm.user_invocable, Some(false)); + } + + #[test] + fn test_parse_frontmatter_paths() { + let raw = "---\nname: test\ndescription: Test\npaths:\n- \"*.rs\"\n- \"*.md\"\n---\nBody"; + let (fm, _body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.paths.unwrap(), vec!["*.rs", "*.md"]); + } + + #[test] + fn test_parse_frontmatter_when_to_use() { + let raw = "---\nname: test\ndescription: Test\nwhen_to_use: When user asks about testing\n---\nBody"; + let (fm, _body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.when_to_use.unwrap(), "When user asks about testing"); + } + + #[test] + fn test_parse_frontmatter_all_fields() { + let raw = r#"--- +name: full-skill +description: A comprehensive skill +context: fork +allowed-tools: +- bash +- read +model: sonnet +argument-hint: "" +when_to_use: When doing comprehensive tasks +disable-model-invocation: false +user-invocable: true +version: "1.0" +paths: +- "*.rs" +- "*.toml" +agent: code-reviewer +effort: high +--- +Body content here"#; + let (fm, body) = parse_skill_file(raw).unwrap(); + assert_eq!(fm.name.unwrap(), "full-skill"); + assert_eq!(fm.description.unwrap(), "A comprehensive skill"); + assert_eq!(fm.context.unwrap(), "fork"); + assert_eq!(fm.allowed_tools.unwrap(), vec!["bash", "read"]); + assert_eq!(fm.model.unwrap(), "sonnet"); + assert_eq!(fm.argument_hint.unwrap(), ""); + assert_eq!(fm.when_to_use.unwrap(), "When doing comprehensive tasks"); + assert_eq!(fm.disable_model_invocation, Some(false)); + assert_eq!(fm.user_invocable, Some(true)); + assert_eq!(fm.version.unwrap(), "1.0"); + assert_eq!(fm.paths.unwrap(), vec!["*.rs", "*.toml"]); + assert_eq!(fm.agent.unwrap(), "code-reviewer"); + assert_eq!(fm.effort.unwrap(), "high"); + assert_eq!(body, "Body content here"); + } + + // ── Validation ── + + #[test] + fn test_validate_missing_description() { + let fm = SkillFrontmatter { + name: Some("test".into()), + ..Default::default() + }; + let warnings = fm.validate("test"); + assert!(warnings.iter().any(|w| w.contains("description"))); + } + + #[test] + fn test_validate_invalid_context() { + let fm = SkillFrontmatter { + name: Some("test".into()), + description: Some("desc".into()), + context: Some("invalid".into()), + ..Default::default() + }; + let warnings = fm.validate("test"); + assert!(warnings.iter().any(|w| w.contains("context"))); + } + + // ── Variable Substitution ── + + #[test] + fn test_substitute_skill_dir() { + let body = "Base: ${SKILL_DIR}/data"; + let result = substitute_variables(body, Path::new("/home/user/skills/myskill"), None); + assert_eq!(result, "Base: /home/user/skills/myskill/data"); + } + + #[test] + fn test_substitute_session_id() { + let body = "Session: ${SESSION_ID}"; + let result = substitute_variables(body, Path::new("."), Some("abc123")); + assert_eq!(result, "Session: abc123"); + } + + #[test] + fn test_substitute_both() { + let body = "Dir: ${SKILL_DIR}\nSession: ${SESSION_ID}"; + let result = substitute_variables(body, Path::new("/skills/test"), Some("sess-1")); + assert!(result.contains("Dir: /skills/test")); + assert!(result.contains("Session: sess-1")); + } + + #[test] + fn test_substitute_no_session_id() { + // 无 session 时占位符被清空 + let body = "Session: ${SESSION_ID}"; + let result = substitute_variables(body, Path::new("."), None); + assert_eq!(result, "Session: "); + } + + // ── SkillRegistry ── + + #[test] + fn test_registry_new_empty() { + let registry = SkillRegistry::new(PathBuf::from("/nonexistent")); + assert!(registry.is_empty()); + assert!(registry.build_reminder().is_none()); + } + + #[test] + fn test_registry_usage_stats() { + let mut registry = SkillRegistry::new(PathBuf::from("/nonexistent")); + assert_eq!(registry.usage_score("test"), 0.0); + + registry.record_usage("test"); + assert!(registry.usage_score("test") > 0.0); + + let stats = registry.usage_stats(); + assert_eq!(stats.get("test").unwrap().invoke_count, 1); + } + + #[test] + fn test_build_reminder_empty() { + let registry = SkillRegistry::new(PathBuf::from("/nonexistent")); + assert!(registry.build_reminder().is_none()); + } + + #[test] + fn test_build_tool_description_empty() { + let registry = SkillRegistry::new(PathBuf::from("/nonexistent")); + assert!(registry.build_tool_description().contains("没有可用的技能")); + } +} diff --git a/src/agent/subagent.rs b/src/agent/subagent.rs new file mode 100644 index 0000000..4eab83f --- /dev/null +++ b/src/agent/subagent.rs @@ -0,0 +1,446 @@ +// src/agent/subagent.rs +// +// 子代理运行器 — 上下文隔离子代理（参考 Claude Code s04 Subagents）。 +// +// 父代理通过 delegate_research 工具将子任务委托给子代理执行。 +// 子代理拥有: +// - 全新的 messages 上下文（不包含父代理的中间工具调用） +// - 完整的工具访问权限（与父代理共享 ToolRegistry） +// - 独立的 ReAct 循环 +// - 完整的 Hook 管道（PreToolUse/PostToolUse）和权限检查 +// +// 子代理只返回最终文本摘要给父代理，中间工具调用不污染父上下文。 + +use std::sync::Arc; +use tokio::sync::mpsc::UnboundedSender; +use tracing::{info, warn}; + +use super::compact; +use super::hooks::{ + HookRegistry, PostToolUseContext, PreToolUseContext, SubagentStartContext, + SubagentStopContext, +}; +use super::runtime::permission::PermissionChecker; +use super::runtime::{AgentConfig, AgentStreamEvent}; +use super::tools::{ToolContext, ToolOutput, ToolRegistry}; +use crate::api::AppState; +use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent, ToolDefinition}; + +/// 子代理运行器 +pub struct SubAgentRunner { + app_state: Arc, + config: AgentConfig, + tool_registry: ToolRegistry, + /// 可选的 Hook 注册表（用于 PreToolUse/PostToolUse 生命周期事件） + hook_registry: Option>, + /// 权限检查器 + permission_checker: Arc, + /// 可选的进度发送器（用于向父代理报告中间步骤） + progress_tx: Option>, +} + +impl SubAgentRunner { + /// 创建新的子代理运行器（无 hook/permission/progress）。 + pub fn new(app_state: Arc) -> Self { + let skill_registry = app_state.skill_registry.clone(); + SubAgentRunner { + app_state, + config: AgentConfig::default(), + tool_registry: ToolRegistry::new(skill_registry), + hook_registry: None, + permission_checker: Arc::new(PermissionChecker::new()), + progress_tx: None, + } + } + + /// 创建带完整 hooks/permissions/progress 的子代理运行器。 + pub fn new_with_hooks( + app_state: Arc, + hook_registry: Option>, + permission_checker: Arc, + progress_tx: Option>, + ) -> Self { + let skill_registry = app_state.skill_registry.clone(); + SubAgentRunner { + app_state, + config: AgentConfig::default(), + tool_registry: ToolRegistry::new(skill_registry), + hook_registry, + permission_checker, + progress_tx, + } + } + + /// 使用自定义 ToolRegistry 创建子代理运行器。 + /// 用于受限场景（如记忆提取子代理仅需只读 + save_memory）。 + pub fn new_with_registry(app_state: Arc, tool_registry: ToolRegistry) -> Self { + SubAgentRunner { + app_state, + config: AgentConfig::default(), + tool_registry, + hook_registry: None, + permission_checker: Arc::new(PermissionChecker::new()), + progress_tx: None, + } + } + + /// 运行子代理的 ReAct 循环，返回最终文本摘要。 + /// + /// # Arguments + /// * `system_prompt` - 子代理的系统提示词 + /// * `research_prompt` - 要执行的研究任务描述 + /// * `max_steps` - 子代理最大推理步数（默认 5） + /// * `hook_registry` - 可选的 HookRegistry（用于触发子代理生命周期事件） + pub async fn run( + &self, + system_prompt: &str, + research_prompt: &str, + max_steps: usize, + ) -> ToolOutput { + let subagent_name = "delegate_research"; + + // OnSubagentStart hook + if let Some(ref registry) = self.hook_registry { + registry + .run_on_subagent_start(&SubagentStartContext { + parent_session_id: String::new(), + subagent_name: subagent_name.to_string(), + prompt: research_prompt.to_string(), + }) + .await; + } + + // 执行实际工作并捕获结果，以便触发 OnSubagentStop hook + let result = self + .run_inner(system_prompt, research_prompt, max_steps) + .await; + let (is_error, result_summary) = if result.is_error { + (true, result.content.clone()) + } else { + (false, result.content.chars().take(200).collect()) + }; + + if let Some(ref registry) = self.hook_registry { + registry + .run_on_subagent_stop(&SubagentStopContext { + parent_session_id: String::new(), + subagent_name: subagent_name.to_string(), + result_summary, + steps: max_steps, + is_error, + }) + .await; + } + + result + } + + /// 实际执行逻辑（提取为内部方法以便 hook 包装） + async fn run_inner( + &self, + system_prompt: &str, + research_prompt: &str, + max_steps: usize, + ) -> ToolOutput { + let llm = &self.app_state.llm; + let tool_defs = self.tool_registry.definitions(); + + // 全新上下文 + let mut messages = vec![ + ChatMessage::system(system_prompt), + ChatMessage::user(research_prompt), + ]; + + // 跟踪工具调用防止死循环 + let mut last_call: Option<(String, String)> = None; + let mut consecutive_count: usize = 0; + let duplicate_threshold: usize = 3; + + for step in 1..=max_steps { + // 上下文压缩检查 + let est_tokens: usize = messages + .iter() + .map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4) + .sum(); + if est_tokens > self.config.context_char_limit * 3 / 2 { + info!( + "[SubAgent] 上下文超限 (est. {} tokens)，触发压缩", + est_tokens + ); + compact::compress_context( + &mut messages, + llm, + self.config.context_char_limit, + "subagent", + ) + .await; + } + + // LLM 流式调用 + let mut stream_rx = match llm.chat_stream(&messages, &tool_defs).await { + Ok(rx) => rx, + Err(e) => { + warn!("[SubAgent] LLM stream 失败: {}", e); + return ToolOutput::error(format!("子代理 LLM 调用失败: {}", e)); + } + }; + + let mut accumulated_content = String::new(); + let mut accumulated_tool_calls: Option> = None; + + while let Some(event) = stream_rx.recv().await { + match event { + StreamEvent::TextDelta(delta) => { + accumulated_content.push_str(&delta); + } + StreamEvent::ToolCallsComplete(tool_calls) => { + accumulated_tool_calls = Some(tool_calls); + } + StreamEvent::Done => break, + StreamEvent::Error(e) => { + warn!("[SubAgent] 流式错误: {}", e); + return ToolOutput::error(format!("子代理流式错误: {}", e)); + } + _ => {} + } + } + + // 无工具调用 = 最终回答 + let tool_calls = match accumulated_tool_calls { + Some(ref tc) if !tc.is_empty() => tc.clone(), + _ => { + // 转发最终文本到父代理 + if let Some(ref tx) = self.progress_tx { + let _ = tx.send(AgentStreamEvent::TextDelta { + content: format!( + "[子代理] {}", + accumulated_content.chars().take(200).collect::() + ), + }); + } + let content_len = accumulated_content.len(); + info!("[SubAgent] 子代理完成，返回 {} 字符摘要", content_len); + return ToolOutput::success( + accumulated_content, + serde_json::json!({ + "steps": step, + "content_length": content_len + }), + ); + } + }; + + // 构建 assistant 消息 + let assistant_msg = ChatMessage::assistant_with_reasoning( + if accumulated_content.is_empty() { + None + } else { + Some(accumulated_content.clone()) + }, + None, + Some(tool_calls.clone()), + ); + messages.push(assistant_msg); + + // 执行工具调用 + for tool_call in &tool_calls { + let tool_name = &tool_call.function.name; + let tool_args_str = &tool_call.function.arguments; + + // 死循环检测 + let call_key = (tool_name.clone(), tool_args_str.clone()); + if last_call.as_ref() == Some(&call_key) { + consecutive_count += 1; + if consecutive_count >= duplicate_threshold { + warn!("[SubAgent] 检测到死循环：{}", tool_name); + let error_msg = ChatMessage::tool_result( + &tool_call.id, + format!( + "工具 {} 被连续重复调用。请停止并给出当前收集到的答案。", + tool_name + ), + ); + messages.push(error_msg); + continue; + } + } else { + last_call = Some(call_key); + consecutive_count = 1; + } + + // 解析参数 + let args: serde_json::Value = match serde_json::from_str(tool_args_str) { + Ok(v) => v, + Err(e) => { + let error_msg = ChatMessage::tool_result( + &tool_call.id, + format!("参数解析失败: {}", e), + ); + messages.push(error_msg); + continue; + } + }; + + // ── 向父代理发送进度事件 ── + if let Some(ref tx) = self.progress_tx { + let _ = tx.send(AgentStreamEvent::ToolCall { + name: format!("[sub] {}", tool_name), + arguments: args.clone(), + step, + }); + } + + // ── PreToolUse hooks + Permission check ── + let tool_ctx = ToolContext::silent(self.app_state.clone()); + + let final_args = if let Some(ref registry) = self.hook_registry { + let pre_ctx = PreToolUseContext { + session_id: "subagent".to_string(), + tool_name: tool_name.clone(), + tool_args: args.clone(), + step, + }; + let pre_result = registry.run_pre_tool_use(&pre_ctx).await; + + // Block check + if pre_result.action.is_blocked() { + let reason = pre_result + .action + .block_reason() + .unwrap_or("tool blocked by hook"); + warn!( + "[SubAgent] PreToolUse hook 阻止了工具: {} ({})", + tool_name, reason + ); + let tool_msg = ChatMessage::tool_result( + &tool_call.id, + format!("工具 {} 被阻止: {}", tool_name, reason), + ); + messages.push(tool_msg); + continue; + } + + pre_result.final_args + } else { + args.clone() + }; + + // Permission check + if self.permission_checker.is_denied(tool_name) { + warn!("[SubAgent] 权限检查拒绝工具: {}", tool_name); + let tool_msg = ChatMessage::tool_result( + &tool_call.id, + format!("工具 {} 在子代理上下文中不可用（权限不足）", tool_name), + ); + messages.push(tool_msg); + continue; + } + + // 执行工具 + let output = match self.tool_registry.get(tool_name) { + Some(tool) => { + match tokio::time::timeout( + std::time::Duration::from_secs(self.config.tool_timeout_secs), + tool.execute(final_args.clone(), &tool_ctx), + ) + .await + { + Ok(output) => output, + Err(_) => ToolOutput::error(format!("工具 {} 执行超时", tool_name)), + } + } + None => ToolOutput::error(format!("未知工具: {}", tool_name)), + }; + + // ── PostToolUse hooks ── + let final_output_content = if let Some(ref registry) = self.hook_registry { + let post_ctx = PostToolUseContext { + session_id: "subagent".to_string(), + agent_name: "subagent".to_string(), + tool_name: tool_name.clone(), + tool_args: final_args, + output_content: output.content.clone(), + is_error: output.is_error, + step, + elapsed_ms: 0, + }; + let post_result = registry.run_post_tool_use(&post_ctx).await; + post_result.final_content + } else { + output.content.clone() + }; + + // 向父代理发送工具结果进度 + if let Some(ref tx) = self.progress_tx { + let preview: String = final_output_content.chars().take(200).collect(); + let _ = tx.send(AgentStreamEvent::ToolResult { + name: format!("[sub] {}", tool_name), + output: preview, + is_error: output.is_error, + metadata: serde_json::json!({}), + step, + }); + } + + // 截断输出（使用 post-hook 处理后的内容） + let truncated = if final_output_content.len() > self.config.max_tool_output_chars { + let t: String = final_output_content + .chars() + .take(self.config.max_tool_output_chars) + .collect(); + format!( + "{}...\n[已截断，原始 {} 字符]", + t, + final_output_content.len() + ) + } else { + final_output_content.clone() + }; + + let tool_msg = ChatMessage::tool_result(&tool_call.id, &truncated); + messages.push(tool_msg); + } + } + + // 达到最大步数，强制生成最终答案 + info!("[SubAgent] 达到最大步数 ({}), 生成最终答案", max_steps); + self.force_final_answer(llm, &messages).await + } + + /// 强制 LLM 生成最终答案（不带工具调用） + async fn force_final_answer(&self, llm: &LlmClient, messages: &[ChatMessage]) -> ToolOutput { + let mut final_messages = messages.to_vec(); + final_messages.push(ChatMessage::user( + "请根据已收集的信息直接给出最终答案，不要再调用工具。", + )); + + let empty_tools: Vec = Vec::new(); + let mut stream_rx = match llm.chat_stream(&final_messages, &empty_tools).await { + Ok(rx) => rx, + Err(e) => { + return ToolOutput::error(format!("子代理最终答案生成失败: {}", e)); + } + }; + + let mut accumulated = String::new(); + while let Some(event) = stream_rx.recv().await { + match event { + StreamEvent::TextDelta(delta) => { + accumulated.push_str(&delta); + } + StreamEvent::Done => break, + StreamEvent::Error(e) => { + warn!("[SubAgent] 最终答案流式错误: {}", e); + break; + } + _ => {} + } + } + + if accumulated.is_empty() { + ToolOutput::error("子代理无法生成最终答案") + } else { + ToolOutput::success(accumulated, serde_json::json!({ "forced": true })) + } + } +} diff --git a/src/agent/task_board.rs b/src/agent/task_board.rs new file mode 100644 index 0000000..d2471cf --- /dev/null +++ b/src/agent/task_board.rs @@ -0,0 +1,142 @@ +// src/agent/task_board.rs +// +// 共享任务看板 — 跨代理任务可见性与依赖图解析。 +// 参考 learn-claude-code s12 Task System + s17 Autonomous Agents。 + +use serde::Serialize; +use sqlx::SqlitePool; +use tracing::info; + +/// 任务摘要 +#[derive(Debug, Clone, Serialize)] +pub struct TaskSummary { + pub session_id: String, + pub task_id: String, + pub content: String, + pub status: String, + pub owner: Option, + pub can_start: bool, + pub blocked_by: Vec, +} + +/// 共享任务看板 +pub struct TaskBoard { + db: SqlitePool, +} + +impl TaskBoard { + pub fn new(db: SqlitePool) -> Self { + TaskBoard { db } + } + + /// 检查任务是否可以开始（所有 blockedBy 依赖已完成） + pub async fn can_start(&self, session_id: &str, task_id: &str) -> anyhow::Result { + let blocked_by: Option = sqlx::query_scalar( + "SELECT blocked_by FROM agent_tasks WHERE session_id = ? AND task_id = ?", + ) + .bind(session_id) + .bind(task_id) + .fetch_optional(&self.db) + .await? + .flatten(); + + let blocked: Vec = + serde_json::from_str(&blocked_by.unwrap_or_default()).unwrap_or_default(); + if blocked.is_empty() { + return Ok(true); + } + + for dep_id in &blocked { + let dep_status: Option = sqlx::query_scalar( + "SELECT status FROM agent_tasks WHERE session_id = ? AND task_id = ?", + ) + .bind(session_id) + .bind(dep_id) + .fetch_optional(&self.db) + .await? + .flatten(); + + if dep_status.as_deref() != Some("completed") { + return Ok(false); + } + } + Ok(true) + } + + /// 原子认领任务（乐观并发控制） + pub async fn claim_task( + &self, + session_id: &str, + task_id: &str, + claimant: &str, + ) -> anyhow::Result { + let rows = sqlx::query( + "UPDATE agent_tasks SET owner = ?, status = 'in_progress' \ + WHERE session_id = ? AND task_id = ? \ + AND (owner IS NULL OR owner = '' OR status = 'pending')", + ) + .bind(claimant) + .bind(session_id) + .bind(task_id) + .execute(&self.db) + .await? + .rows_affected(); + + Ok(rows > 0) + } + + /// 列出所有可认领的待处理任务（跨 session） + pub async fn list_available_tasks(&self, limit: usize) -> anyhow::Result> { + #[allow(clippy::type_complexity)] + let rows: Vec<( + String, + String, + String, + String, + Option, + Option, + )> = sqlx::query_as( + "SELECT session_id, task_id, content, status, blocked_by, owner \ + FROM agent_tasks \ + WHERE status = 'pending' AND (owner IS NULL OR owner = '') \ + ORDER BY created_at ASC LIMIT ?", + ) + .bind(limit as i32) + .fetch_all(&self.db) + .await?; + + let mut summaries = Vec::new(); + for (session_id, task_id, content, status, blocked_by, owner) in rows { + let blocked: Vec = + serde_json::from_str(&blocked_by.unwrap_or_default()).unwrap_or_default(); + let can_start = self.can_start(&session_id, &task_id).await.unwrap_or(false); + summaries.push(TaskSummary { + session_id, + task_id, + content, + status, + owner, + can_start, + blocked_by: blocked, + }); + } + + info!( + "[TaskBoard] 查询到 {} 个可认领任务 (limit={})", + summaries.len(), + limit + ); + Ok(summaries) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_can_start_no_deps() { + // can_start 需要数据库连接，此处仅验证结构 + // 实际集成测试应在 tests/ 目录中 + } +} diff --git a/src/agent/team/config.rs b/src/agent/team/config.rs new file mode 100644 index 0000000..99900ec --- /dev/null +++ b/src/agent/team/config.rs @@ -0,0 +1,53 @@ +// src/agent/team/config.rs +// +// 团队配置类型与持久化 (.team/{session_id}/config.json)。 + +use serde::{Deserialize, Serialize}; + +/// 团队成员状态 +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum MemberStatus { + Spawning, + Working, + Idle, + Shutdown, +} + +/// 单个成员配置 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemberConfig { + pub name: String, + pub role: String, + pub system_prompt: String, + pub status: MemberStatus, +} + +/// 团队配置 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TeamConfig { + pub session_id: String, + pub lead_name: String, + pub members: Vec, +} + +impl TeamConfig { + /// 创建新的团队配置 + pub fn new(session_id: &str) -> Self { + TeamConfig { + session_id: session_id.to_string(), + lead_name: "lead".to_string(), + members: Vec::new(), + } + } + + /// 添加成员 + pub fn add_member(&mut self, name: &str, role: &str, system_prompt: &str) { + self.members.push(MemberConfig { + name: name.to_string(), + role: role.to_string(), + system_prompt: system_prompt.to_string(), + status: MemberStatus::Spawning, + }); + } +} diff --git a/src/agent/team/inbox.rs b/src/agent/team/inbox.rs new file mode 100644 index 0000000..309fabe --- /dev/null +++ b/src/agent/team/inbox.rs @@ -0,0 +1,110 @@ +// src/agent/team/inbox.rs +// +// 团队消息邮箱系统。 +// 使用 .team/{session_id}/inbox/{agent_name}.jsonl 作为 append-only 消息文件。 + +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::io::Write; +use std::path::{Path, PathBuf}; + +/// 团队消息类型 +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum TeamMessageType { + Task, + Result, + Question, + Answer, + Status, +} + +impl TeamMessageType { + pub fn as_str(&self) -> &str { + match self { + TeamMessageType::Task => "task", + TeamMessageType::Result => "result", + TeamMessageType::Question => "question", + TeamMessageType::Answer => "answer", + TeamMessageType::Status => "status", + } + } +} + +/// 团队消息 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TeamMessage { + pub from: String, + pub to: String, + pub content: String, + pub msg_type: TeamMessageType, + pub timestamp: String, +} + +impl TeamMessage { + pub fn new(from: &str, to: &str, content: &str, msg_type: TeamMessageType) -> Self { + TeamMessage { + from: from.to_string(), + to: to.to_string(), + content: content.to_string(), + msg_type, + timestamp: Utc::now().to_rfc3339(), + } + } +} + +/// 获取团队目录路径 +pub fn team_dir(session_id: &str) -> PathBuf { + PathBuf::from(".team").join(session_id) +} + +/// 获取指定 agent 的收件箱路径 +pub fn inbox_path(team_dir: &Path, agent_name: &str) -> PathBuf { + team_dir.join("inbox").join(format!("{}.jsonl", agent_name)) +} + +/// 向收件箱追加一条消息 +pub fn append_message(team_dir: &Path, agent_name: &str, msg: &TeamMessage) -> std::io::Result<()> { + let inbox = inbox_path(team_dir, agent_name); + if let Some(parent) = inbox.parent() { + std::fs::create_dir_all(parent)?; + } + let line = serde_json::to_string(msg).unwrap_or_default(); + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&inbox)?; + file.write_all(line.as_bytes())?; + file.write_all(b"\n")?; + Ok(()) +} + +/// 读取并清空收件箱 +pub fn drain_inbox(team_dir: &Path, agent_name: &str) -> Vec { + let inbox = inbox_path(team_dir, agent_name); + if !inbox.exists() { + return Vec::new(); + } + + let content = match std::fs::read_to_string(&inbox) { + Ok(c) => c, + Err(_) => return Vec::new(), + }; + + let messages: Vec = content + .lines() + .filter(|l| !l.is_empty()) + .filter_map(|l| serde_json::from_str(l).ok()) + .collect(); + + // 清空文件 + let _ = std::fs::write(&inbox, ""); + + messages +} + +/// 检查收件箱中是否有未读消息 +pub fn has_pending(team_dir: &Path, agent_name: &str) -> bool { + let inbox = inbox_path(team_dir, agent_name); + inbox.exists() && inbox.metadata().map(|m| m.len() > 0).unwrap_or(false) +} diff --git a/src/agent/team/manager.rs b/src/agent/team/manager.rs new file mode 100644 index 0000000..9c02882 --- /dev/null +++ b/src/agent/team/manager.rs @@ -0,0 +1,205 @@ +// src/agent/team/manager.rs +// +// 团队管理器：spawn/stop/send/broadcast/list。 + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use tokio::sync::Mutex; +use tracing::{info, warn}; + +use crate::agent::runtime::AgentConfig; +use crate::api::AppState; + +use super::config::{MemberStatus, TeamConfig}; +use super::inbox::{self, TeamMessage, TeamMessageType}; +use super::teammate; + +/// 队友运行时句柄 +#[derive(Clone)] +pub struct TeamMemberHandle { + pub name: String, + pub role: String, + pub status: Arc>, + cancelled: Arc, +} + +impl TeamMemberHandle { + /// 请求停止该队友 + pub fn request_stop(&self) { + self.cancelled.store(true, Ordering::SeqCst); + } +} + +/// 团队管理器 +/// +/// # 锁顺序约定 (CRITICAL) +/// +/// 本模块中存在两个 tokio::sync::Mutex 的嵌套获取： +/// 1. `TeamManager.handles` (外层) +/// 2. `TeamMemberHandle.status` (内层) +/// +/// 任何代码如果先获取 `status` 再获取 `handles` 将导致死锁。 +/// 所有新增代码必须遵守 `handles → status` 的顺序。 +/// 参考: `list_members()` 的实现作为正确顺序的示例。 +pub struct TeamManager { + pub config: TeamConfig, + team_dir: std::path::PathBuf, + handles: Arc>>, + app_state: Arc, +} + +impl TeamManager { + /// 创建新的团队管理器 + pub fn new(app_state: Arc, session_id: &str) -> Self { + let team_dir = inbox::team_dir(session_id); + TeamManager { + config: TeamConfig::new(session_id), + team_dir, + handles: Arc::new(Mutex::new(HashMap::new())), + app_state, + } + } + + /// 获取团队目录 + pub fn team_dir(&self) -> &std::path::Path { + &self.team_dir + } + + /// 生成队友的系统提示词 + fn build_teammate_system_prompt(role: &str) -> String { + format!( + "你是一位专业的天体物理学研究助手，在一个研究团队中工作。你的角色是：{}。\n\ + \n\ + 你可以使用文献搜索、下载、RAG 检索等工具完成任务。\n\ + 你通过团队收件箱接收任务分配，完成后通过消息汇报结果。\n\ + \n\ + 核心原则：\n\ + 1. 收到任务后立即开始工作，不要等待确认。\n\ + 2. 完成任务后向 lead 发送 Result 类型的消息汇报。\n\ + 3. 只使用与你的角色相关的工具。\n\ + 4. 用中文输出结果，引用具体文献来源。", + role + ) + } + + /// 生成队友的工作任务提示 + fn build_teammate_task_prompt(role: &str) -> String { + format!( + "你已加入天体物理研究团队，角色：{}。\n\ + 请在收件箱中等待来自 lead 分配的任务。\n\ + 收到任务后使用你的工具完成，然后将结果发送回 lead。", + role + ) + } + + /// 启动一个队友 + pub async fn spawn(&self, name: &str, role: &str) -> TeamMemberHandle { + info!("[TeamManager] 启动队友: {} ({})", name, role); + + let cancelled = Arc::new(AtomicBool::new(false)); + let cancelled_clone = cancelled.clone(); + let status = Arc::new(Mutex::new(MemberStatus::Spawning)); + let status_clone = status.clone(); + let name_clone = name.to_string(); + + let app_state = self.app_state.clone(); + let team_dir = self.team_dir.clone(); + let role_owned = role.to_string(); + let system_prompt = Self::build_teammate_system_prompt(&role_owned); + let task_prompt = Self::build_teammate_task_prompt(&role_owned); + let agent_config = AgentConfig::from_env_optional(); + + // 后台启动队友 ReAct 循环 + tokio::spawn(async move { + teammate::run_teammate_loop( + app_state, + team_dir, + name_clone, + role_owned, + system_prompt, + task_prompt, + agent_config, + status_clone, + cancelled_clone, + ) + .await; + }); + + let handle = TeamMemberHandle { + name: name.to_string(), + role: role.to_string(), + status, + cancelled, + }; + + let mut handles = self.handles.lock().await; + handles.insert(name.to_string(), handle.clone()); + handle + } + + /// 停止一个队友 + pub async fn stop(&self, name: &str) { + info!("[TeamManager] 停止队友: {}", name); + let handles = self.handles.lock().await; + if let Some(handle) = handles.get(name) { + handle.request_stop(); + } + } + + /// 停止所有队友 + pub async fn stop_all(&self) { + info!("[TeamManager] 停止所有队友"); + let handles = self.handles.lock().await; + for handle in handles.values() { + handle.request_stop(); + } + } + + /// 发送消息给指定队友 + pub fn send_message(&self, from: &str, to: &str, content: &str, msg_type: TeamMessageType) { + let msg = TeamMessage::new(from, to, content, msg_type); + if let Err(e) = inbox::append_message(&self.team_dir, to, &msg) { + warn!("[TeamManager] 发送消息失败: {}", e); + } else { + info!( + "[TeamManager] {} -> {}: {}", + from, + to, + &content.chars().take(80).collect::() + ); + } + } + + /// 广播消息给所有队友 + pub fn broadcast(&self, from: &str, content: &str) { + info!( + "[TeamManager] 广播消息: {}", + content.chars().take(80).collect::() + ); + let msg = TeamMessage::new(from, "all", content, TeamMessageType::Status); + for member in &self.config.members { + if member.name != from { + let _ = inbox::append_message(&self.team_dir, &member.name, &msg); + } + } + } + + /// 检查并清空指定 agent 的收件箱 + pub fn check_inbox(&self, agent_name: &str) -> Vec { + inbox::drain_inbox(&self.team_dir, agent_name) + } + + /// 列出所有队友状态 + /// + /// 正确锁顺序示例：先获取 `handles`，再获取每个 `status`。 + pub async fn list_members(&self) -> Vec<(String, String, MemberStatus)> { + let handles = self.handles.lock().await; + let mut members = Vec::new(); + for (name, handle) in handles.iter() { + let status = handle.status.lock().await.clone(); + members.push((name.clone(), handle.role.clone(), status)); + } + members + } +} diff --git a/src/agent/team/mod.rs b/src/agent/team/mod.rs new file mode 100644 index 0000000..7bc7ac8 --- /dev/null +++ b/src/agent/team/mod.rs @@ -0,0 +1,11 @@ +// src/agent/team/mod.rs +// +// 多智能体团队协作模块（参考 Claude Code s09 Agent Teams）。 +// +// 通过文件邮箱 (.team/{session_id}/inbox/*.jsonl) 实现 +// Lead Agent 与多个 Teammate Agent 之间的消息传递与任务协调。 + +pub mod config; +pub mod inbox; +pub mod manager; +pub mod teammate; diff --git a/src/agent/team/teammate.rs b/src/agent/team/teammate.rs new file mode 100644 index 0000000..edbe0d8 --- /dev/null +++ b/src/agent/team/teammate.rs @@ -0,0 +1,270 @@ +// src/agent/team/teammate.rs +// +// 队友 ReAct 循环：检查收件箱 → 执行任务 → 汇报结果。 + +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use tokio::sync::Mutex; +use tracing::{info, warn}; + +use crate::agent::background::BgNotificationQueue; +use crate::agent::compact; +use crate::agent::runtime::AgentConfig; +use crate::agent::tools::ToolRegistry; +use crate::api::AppState; +use crate::clients::llm::{ChatMessage, LlmClient, StreamEvent}; + +use super::config::MemberStatus; +use super::inbox::{self, TeamMessageType}; + +/// 队友 ReAct 循环。 +/// +/// 生命周期： +/// SPAWN → WORKING (ReAct) → IDLE (poll inbox) → SHUTDOWN +#[allow(clippy::too_many_arguments)] +pub async fn run_teammate_loop( + app_state: Arc, + team_dir: PathBuf, + name: String, + role: String, + system_prompt: String, + task_prompt: String, + config: AgentConfig, + status: Arc>, + cancelled: Arc, +) { + let llm = &app_state.llm; + // 队友的工具注册表排除 delegate_research（防止无限委托链） + let queue = Arc::new(BgNotificationQueue::new()); + let tool_registry = + ToolRegistry::new_with_queue(Some(queue.clone()), app_state.skill_registry.clone()); + + let tool_defs = tool_registry.definitions(); + + let mut messages = vec![ + ChatMessage::system(&system_prompt), + ChatMessage::user(&task_prompt), + ]; + + info!("[Teammate:{}] 启动 ReAct 循环", name); + + loop { + // ── 检查取消 ── + if cancelled.load(Ordering::SeqCst) { + info!("[Teammate:{}] 收到取消信号，正在停止...", name); + *status.lock().await = MemberStatus::Shutdown; + // 通知 lead 自己退出了 + let goodbye = super::inbox::TeamMessage::new( + &name, + "lead", + &format!("队友 {} ({}) 已退出。", name, role), + TeamMessageType::Status, + ); + let _ = inbox::append_message(&team_dir, "lead", &goodbye); + return; + } + + // ── IDLE 阶段：检查收件箱 ── + *status.lock().await = MemberStatus::Idle; + + let inbox_msgs = inbox::drain_inbox(&team_dir, &name); + + if inbox_msgs.is_empty() { + // 等待新消息或取消，poll 间隔 5 秒，最长 60 秒 + for _ in 0..12 { + tokio::time::sleep(std::time::Duration::from_secs(5)).await; + if cancelled.load(Ordering::SeqCst) || inbox::has_pending(&team_dir, &name) { + break; + } + } + + if cancelled.load(Ordering::SeqCst) { + continue; + } + + // 再次 drain（可能因为 pending 标志被唤醒） + let inbox_msgs = inbox::drain_inbox(&team_dir, &name); + if inbox_msgs.is_empty() { + continue; // 超时，没有新消息，继续 idle + } + + // 有新消息 → 进入 WORKING 阶段 + *status.lock().await = MemberStatus::Working; + + for msg in &inbox_msgs { + messages.push(ChatMessage::user(format!( + "[来自 {} 的消息 ({}):] {}", + msg.from, msg.timestamp, msg.content + ))); + } + + // ── 执行 ReAct 循环 ── + let result = run_teammate_react_turn( + llm, + &tool_defs, + &tool_registry, + &app_state, + &mut messages, + &config, + &cancelled, + ) + .await; + + // ── 汇报结果 ── + if let Some(ref summary) = result { + let reply = + super::inbox::TeamMessage::new(&name, "lead", summary, TeamMessageType::Result); + let _ = inbox::append_message(&team_dir, "lead", &reply); + info!( + "[Teammate:{}] 任务完成，已发送结果 ({}字符)", + name, + summary.len() + ); + } + } else { + // 收件箱有消息 → 直接进入 WORKING + *status.lock().await = MemberStatus::Working; + + for msg in &inbox_msgs { + messages.push(ChatMessage::user(format!( + "[来自 {} 的消息: ({})] {}", + msg.from, msg.timestamp, msg.content + ))); + } + + let result = run_teammate_react_turn( + llm, + &tool_defs, + &tool_registry, + &app_state, + &mut messages, + &config, + &cancelled, + ) + .await; + + if let Some(ref summary) = result { + let reply = + super::inbox::TeamMessage::new(&name, "lead", summary, TeamMessageType::Result); + let _ = inbox::append_message(&team_dir, "lead", &reply); + } + } + } +} + +/// 队友的单次 ReAct turn。 +/// +/// 一个简化的 ReAct 循环：LLM 调用 → 工具执行 → 结果注入 → 循环... +/// 与主 Agent 的循环类似但更轻量（无 SSE、无 DB 持久化、无 hooks）。 +async fn run_teammate_react_turn( + llm: &LlmClient, + tool_defs: &[crate::clients::llm::ToolDefinition], + tool_registry: &ToolRegistry, + app_state: &Arc, + messages: &mut Vec, + config: &AgentConfig, + cancelled: &Arc, +) -> Option { + let max_steps = config.max_steps.min(5); // 队友步数限制更严格 + + for _step in 1..=max_steps { + // 检查取消 + if cancelled.load(Ordering::SeqCst) { + return None; + } + + // 上下文压缩检查 + let est_tokens: usize = messages + .iter() + .map(|m| m.content.as_ref().map_or(0, |c| c.len()) + 4) + .sum(); + if est_tokens > config.context_char_limit * 3 / 2 { + compact::compress_context(messages, llm, config.context_char_limit, "teammate").await; + } + + // LLM 流式调用 + let mut stream_rx = match llm.chat_stream(messages, tool_defs).await { + Ok(rx) => rx, + Err(_) => return None, + }; + + let mut accumulated = String::new(); + let mut tool_calls: Option> = None; + + while let Some(event) = stream_rx.recv().await { + match event { + StreamEvent::TextDelta(delta) => accumulated.push_str(&delta), + StreamEvent::ToolCallsComplete(tc) => tool_calls = Some(tc), + StreamEvent::Done => break, + StreamEvent::Error(_) => return None, + _ => {} + } + } + + // 无工具调用 = 最终回答 + let tool_calls = match tool_calls { + Some(ref tc) if !tc.is_empty() => tc.clone(), + _ => { + return if accumulated.is_empty() { + None + } else { + Some(accumulated) + }; + } + }; + + // 构建 assistant 消息 + messages.push(ChatMessage::assistant_with_reasoning( + if accumulated.is_empty() { + None + } else { + Some(accumulated) + }, + None, + Some(tool_calls.clone()), + )); + + // 执行工具调用 + for tc in &tool_calls { + let args: serde_json::Value = match serde_json::from_str(&tc.function.arguments) { + Ok(a) => a, + Err(_) => continue, + }; + + let tool_ctx = crate::agent::tools::ToolContext::silent(app_state.clone()); + + let output = match tool_registry.get(&tc.function.name) { + Some(tool) => { + match tokio::time::timeout( + std::time::Duration::from_secs(config.tool_timeout_secs), + tool.execute(args, &tool_ctx), + ) + .await + { + Ok(o) => o, + Err(_) => crate::agent::tools::ToolOutput::error("执行超时"), + } + } + None => crate::agent::tools::ToolOutput::error("未知工具"), + }; + + let truncated = if output.content.len() > config.max_tool_output_chars { + let t: String = output + .content + .chars() + .take(config.max_tool_output_chars) + .collect(); + format!("{}...\n[已截断]", t) + } else { + output.content + }; + + messages.push(ChatMessage::tool_result(&tc.id, &truncated)); + } + } + + // 达到最大步数，返回 None（无结果） + warn!("[Teammate] 达到最大步数限制 ({} steps)，无结果", max_steps); + None +} diff --git a/src/agent/terminal.rs b/src/agent/terminal.rs new file mode 100644 index 0000000..441f713 --- /dev/null +++ b/src/agent/terminal.rs @@ -0,0 +1,127 @@ +// src/agent/terminal.rs +// +// 智能体循环终止信号。 +// 参考 Claude Code query.ts 的 Terminal 类型设计： +// 用结构化枚举替代隐式的 break / Err(...) 退出， +// 使调用方可以精确知道循环为何结束。 + +use serde::Serialize; + +/// Agent 循环终止原因 +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "reason", content = "detail")] +pub enum TurnTerminal { + /// 正常完成 - Agent 给出了最终回答 + Completed { + session_id: String, + total_steps: usize, + }, + + /// 达到最大推理步数 + MaxStepsReached { + session_id: String, + steps: usize, + max_steps: usize, + }, + + /// 用户手动中止 + CancelledByUser { session_id: String, at_step: usize }, + + /// 检测到工具死循环 + DuplicateCallDetected { + session_id: String, + tool_name: String, + at_step: usize, + }, + + /// 大模型流式调用失败 + ModelStreamError { + session_id: String, + message: String, + at_step: usize, + }, + + /// 大模型返回错误（非流式） + ModelError { session_id: String, message: String }, +} + +impl TurnTerminal { + /// 是否为正常完成 + pub fn is_completed(&self) -> bool { + matches!(self, TurnTerminal::Completed { .. }) + } + + /// 获取关联的 session_id + pub fn session_id(&self) -> &str { + match self { + TurnTerminal::Completed { session_id, .. } + | TurnTerminal::MaxStepsReached { session_id, .. } + | TurnTerminal::CancelledByUser { session_id, .. } + | TurnTerminal::DuplicateCallDetected { session_id, .. } + | TurnTerminal::ModelStreamError { session_id, .. } + | TurnTerminal::ModelError { session_id, .. } => session_id, + } + } + + /// 人类可读的终止描述 + pub fn description(&self) -> &str { + match self { + TurnTerminal::Completed { .. } => "正常完成", + TurnTerminal::MaxStepsReached { .. } => "达到最大步数", + TurnTerminal::CancelledByUser { .. } => "用户手动中止", + TurnTerminal::DuplicateCallDetected { .. } => "检测到死循环", + TurnTerminal::ModelStreamError { .. } => "模型流式错误", + TurnTerminal::ModelError { .. } => "模型错误", + } + } +} + +impl std::fmt::Display for TurnTerminal { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TurnTerminal::Completed { + session_id, + total_steps, + } => { + write!(f, "[{}] 正常完成 ({} steps)", session_id, total_steps) + } + TurnTerminal::MaxStepsReached { + session_id, + steps, + max_steps, + } => { + write!(f, "[{}] 达到最大步数 ({}/{})", session_id, steps, max_steps) + } + TurnTerminal::CancelledByUser { + session_id, + at_step, + } => { + write!(f, "[{}] 用户在第 {} 步手动中止", session_id, at_step) + } + TurnTerminal::DuplicateCallDetected { + session_id, + tool_name, + at_step, + } => { + write!( + f, + "[{}] 检测到 {} 死循环 (step {})", + session_id, tool_name, at_step + ) + } + TurnTerminal::ModelStreamError { + session_id, + message, + at_step, + } => { + write!(f, "[{}] 流式错误 step {}: {}", session_id, at_step, message) + } + TurnTerminal::ModelError { + session_id, + message, + } => { + write!(f, "[{}] 模型错误: {}", session_id, message) + } + } + } +} diff --git a/src/agent/tools.rs b/src/agent/tools.rs deleted file mode 100644 index df1e396..0000000 --- a/src/agent/tools.rs +++ /dev/null @@ -1,659 +0,0 @@ -// src/agent/tools.rs -// -// 科研智能体工具集定义与实现。 -// 每个工具遵循 AgentTool trait，向大模型声明 JSON Schema 参数定义， -// 并在 execute 中调用已有的服务层完成实际业务操作。 - -use async_trait::async_trait; -use serde_json::json; -use std::sync::Arc; -use tracing::{info, error}; - -use crate::api::AppState; -use crate::clients::llm::ToolDefinition; - -/// 工具执行上下文，封装全局共享状态 -pub struct ToolContext { - pub app_state: Arc, -} - -/// 工具执行结果 -#[derive(Debug, Clone)] -pub struct ToolOutput { - /// 给大模型阅读的截断文本 - pub content: String, - /// 是否为错误 - pub is_error: bool, - /// 结构化元数据（给前端 Timeline 直接渲染） - pub metadata: serde_json::Value, -} - -impl ToolOutput { - /// 创建成功结果 - pub fn success(content: impl Into, metadata: serde_json::Value) -> Self { - ToolOutput { - content: content.into(), - is_error: false, - metadata, - } - } - - /// 创建错误结果 - pub fn error(msg: impl Into) -> Self { - ToolOutput { - content: msg.into(), - is_error: true, - metadata: json!({}), - } - } -} - -/// 智能体工具 trait -#[async_trait] -pub trait AgentTool: Send + Sync { - /// 工具名称（与 LLM function calling 的 name 保持一致） - fn name(&self) -> &str; - /// 工具描述（告知 LLM 何时应该调用该工具） - fn description(&self) -> &str; - /// JSON Schema 格式的参数定义 - fn parameters(&self) -> serde_json::Value; - /// 执行工具逻辑 - async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput; -} - -/// 工具注册表，管理所有可用工具 -pub struct ToolRegistry { - tools: Vec>, -} -impl ToolRegistry { - /// 创建默认工具注册表（包含全部科研工具） - pub fn new() -> Self { - let tools: Vec> = vec![ - Box::new(SearchPapersTool), - Box::new(GetPaperMetadataTool), - Box::new(DownloadPaperTool), - Box::new(ParsePaperTool), - Box::new(GetPaperContentTool), - Box::new(RagSearchTool), - Box::new(QueryTargetTool), - ]; - ToolRegistry { tools } - } - - /// 根据名称查找工具 - pub fn get(&self, name: &str) -> Option<&dyn AgentTool> { - self.tools.iter().find(|t| t.name() == name).map(|t| t.as_ref()) - } - - /// 生成所有工具的 ToolDefinition 列表（用于发送给 LLM） - pub fn definitions(&self) -> Vec { - self.tools.iter().map(|t| { - ToolDefinition::new(t.name(), t.description(), t.parameters()) - }).collect() - } -} - -/// 截断文本到指定最大字符数 -fn truncate_content(s: &str, max_chars: usize) -> String { - if s.len() <= max_chars { - s.to_string() - } else { - let truncated: String = s.chars().take(max_chars).collect(); - format!("{}\n\n[... 内容已截断，共 {} 字符 ...]", truncated, s.len()) - } -} - -// ────────────────────────── 1. SearchPapersTool ────────────────────────── - -/// 文献搜索工具：调用 ADS/arXiv 进行跨库检索 -pub struct SearchPapersTool; - -#[async_trait] -impl AgentTool for SearchPapersTool { - fn name(&self) -> &str { "search_papers" } - - fn description(&self) -> &str { - "搜索天文学文献。支持 NASA ADS 和 arXiv 跨平台联合检索，结果自动合并去重，并关联本地馆藏状态与引用关系网。输入关键词或高级检索式，返回匹配的文献列表。\ - 适用于：查找相关文献、了解研究领域现状、获取特定主题的论文。" - } - - fn parameters(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "query": { - "type": "string", - "description": "搜索关键词或高级检索式。支持语法：\ - 1. 字段限定：au:\"作者\" 或 author:\"作者\"、ti:\"标题\" 或 title:\"标题\"、abs:\"摘要关键字\"；\ - 2. 年份限定：year:2020（单年）或 year:2020-2025（年份区间）；\ - 3. 逻辑运算：支持 AND、OR、NOT 逻辑组合及括号分组，如 '(ti:subdwarf OR ti:\"white dwarf\") AND year:2020-2025'；\ - 4. 短语匹配：用双引号 \"\" 包含精确匹配短语，如 '\"Gaia BH1\"'。\ - 所有的中文标点符号（如“”()，；）在后台均会自动清洗转换。" - }, - "rows": { - "type": "integer", - "description": "返回结果数量，默认5，最大20", - "default": 5 - } - }, - "required": ["query"] - }) - } - - async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { - let query = match args.get("query").and_then(|q| q.as_str()) { - Some(q) => q.to_string(), - None => return ToolOutput::error("缺少必需参数 'query'"), - }; - let rows = args.get("rows").and_then(|r| r.as_i64()).unwrap_or(5).min(20) as i32; - - info!("[SearchPapersTool] 执行文献搜索: query='{}', rows={}", query, rows); - - let state = &ctx.app_state; - - match crate::services::search::search_papers(state, &query, "all", 0, rows, "relevance").await { - Ok(results) => { - if results.is_empty() { - return ToolOutput::success("未找到匹配的文献。请尝试调整搜索关键词。", json!({ "count": 0 })); - } - - // 格式化结果（保留详细信息给 LLM，但不含摘要且不作截断） - let display_results: Vec = results.iter().map(|p| { - let first_author = p.authors.first().cloned().unwrap_or_else(|| "未知".to_string()); - json!({ - "bibcode": p.bibcode, - "title": p.title, - "first_author": first_author, - "year": p.year, - "citation_count": p.citation_count, - }) - }).collect(); - - let content = display_results.iter().enumerate().map(|(i, r)| { - format!( - "{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次", - i + 1, - r["bibcode"].as_str().unwrap_or(""), - r["title"].as_str().unwrap_or(""), - r["year"].as_str().unwrap_or(""), - r["first_author"].as_str().unwrap_or("未知"), - r["citation_count"].as_i64().unwrap_or(0) - ) - }).collect::>().join("\n\n"); - - ToolOutput::success( - content, - json!({ - "count": results.len(), - "papers": display_results - }) - ) - } - Err(e) => { - error!("[SearchPapersTool] 检索失败: {}", e); - ToolOutput::error(format!("文献检索失败: {}", e)) - } - } - } -} - -// ────────────────────────── 1b. GetPaperMetadataTool ────────────────────────── - -/// 获取文献元数据工具：获取指定文献的完整元数据（包含完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等） -pub struct GetPaperMetadataTool; - -#[async_trait] -impl AgentTool for GetPaperMetadataTool { - fn name(&self) -> &str { "get_paper_metadata" } - - fn description(&self) -> &str { - "获取指定文献的完整元数据信息（包括完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等）。\ - 适用于：需要查看某篇文献的详细信息、阅读完整摘要以评估文献相关性。" - } - - fn parameters(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "bibcode": { - "type": "string", - "description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）" - } - }, - "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!("[GetPaperMetadataTool] 获取文献元数据: {}", bibcode); - let state = &ctx.app_state; - - match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode).await { - Ok(paper) => { - let content = format!( - "文献元数据 [{}]：\n\ - 标题: {}\n\ - 作者: {}\n\ - 年份: {}\n\ - 期刊: {}\n\ - 关键字: {}\n\ - DOI: {}\n\ - arXiv ID: {}\n\ - 引用数: {} 次\n\ - 参考文献数: {} 次\n\ - 文献类型: {}\n\ - 已下载: {}\n\ - 已解析为 Markdown: {}\n\ - 摘要:\n{}", - paper.bibcode, - paper.title, - paper.authors.join(", "), - paper.year, - paper.pub_journal, - paper.keywords.join(", "), - paper.doi, - paper.arxiv_id, - paper.citation_count, - paper.reference_count, - paper.doctype, - paper.is_downloaded, - paper.has_markdown, - paper.abstract_text - ); - ToolOutput::success(content, json!(paper)) - } - Err(e) => ToolOutput::error(format!("获取文献 {} 元数据失败: {}", bibcode, e)), - } - } -} - -// ────────────────────────── 2. GetPaperContentTool ────────────────────────── - -/// 获取文献内容工具：仅从本地读取并获取已解析的文献 Markdown 全文内容 -pub struct GetPaperContentTool; - -#[async_trait] -impl AgentTool for GetPaperContentTool { - fn name(&self) -> &str { "get_paper_content" } - - fn description(&self) -> &str { - "读取并在本地库中获取已解析的文献 Markdown 完整文本内容。\ - 注意：本工具仅能读取已在数据库注册且已解析的文献，不会自动触发下载或解析。若文献未下载或未解析，本工具会返回详细指引，提示先依次调用 download_paper 和 parse_paper。" - } - - fn parameters(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "bibcode": { - "type": "string", - "description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）" - } - }, - "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!("[GetPaperContentTool] 获取文献内容: {}", 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(format!("获取文献内容失败：该文献未在本地数据库中注册，请先使用 search_papers 搜索该文献。")), - Err(e) => return ToolOutput::error(format!("获取文献内容失败: {}", e)), - }; - - let md_rel = match md_opt { - Some(rel) => rel, - None => return ToolOutput::error(format!("获取文献内容失败：该文献尚未完成结构化解析。如果未下载，请先调用 download_paper；如果已下载，请先调用 parse_paper 进行解析。")), - }; - - let md_abs = state.config.library_dir.join(&md_rel); - if !md_abs.exists() { - return ToolOutput::error(format!("获取文献内容失败：文献本地 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)), - } - } -} - -// ────────────────────────── 2a. DownloadPaperTool ────────────────────────── - -/// 下载文献全文资源工具：仅下载文献全文资源（PDF/HTML）至本地图书馆 -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（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）" - }, - "force": { - "type": "boolean", - "description": "是否强制重新下载（即使本地已下载该文献）" - } - }, - "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 force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false); - - info!("[DownloadPaperTool] 下载文献全文资源: {}, 强制重下: {}", 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)), - } - } -} - -// ────────────────────────── 2b. ParsePaperTool ────────────────────────── - -/// 结构化解析文献内容工具：仅对已下载的物理资源进行结构化解析生成 Markdown -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（如 '2024ApJ...960..123A'）、DOI（如 '10.3847/1538-4357/ad0c5a'）或 arXiv ID（如 '2401.12345'）" - }, - "force": { - "type": "boolean", - "description": "是否强制重新解析（即使本地已解析过该文献）" - } - }, - "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 force = args.get("force").and_then(|f| f.as_bool()).unwrap_or(false); - - info!("[ParsePaperTool] 结构化解析文献内容: {}, 强制重析: {}", 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)) - } - } - } - } -} - - - -// ────────────────────────── 4. RagSearchTool ────────────────────────── - -/// RAG 向量检索工具：基于语义相似度检索文献切片 -pub struct RagSearchTool; - -#[async_trait] -impl AgentTool for RagSearchTool { - fn name(&self) -> &str { "rag_search" } - - fn description(&self) -> &str { - "在已向量化的文献库中进行语义检索。输入自然语言问题，返回最相关的文献片段。\ - 适用于：跨多篇文献查找特定信息、回答需要综合多个来源的问题。要求文献已完成向量化（embed）。" - } - - fn parameters(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "question": { - "type": "string", - "description": "用于语义检索的自然语言问题" - }, - "top_k": { - "type": "integer", - "description": "返回最相关的片段数量，默认5", - "default": 5 - } - }, - "required": ["question"] - }) - } - - async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { - let question = match args.get("question").and_then(|q| q.as_str()) { - Some(q) => q.to_string(), - None => return ToolOutput::error("缺少必需参数 'question'"), - }; - let top_k = args.get("top_k").and_then(|k| k.as_u64()).unwrap_or(5) as usize; - - info!("[RagSearchTool] 执行语义检索: question='{}', top_k={}", question, top_k); - let state = &ctx.app_state; - - match crate::services::rag::retrieve(&state.db, &state.embedding, &question, top_k).await { - Ok(results) => { - if results.is_empty() { - return ToolOutput::success( - "未找到相关的文献片段。文献库中可能尚无向量化数据，请先对目标文献执行向量化操作。", - json!({ "count": 0 }) - ); - } - - let content = results.iter().enumerate().map(|(i, r)| { - format!( - "[片段 {} | 来源: {} | 段落: {} | 相似度距离: {:.4}]\n{}", - i + 1, r.bibcode, r.paragraph_index, r.distance, r.content - ) - }).collect::>().join("\n\n---\n\n"); - - let sources: Vec = results.iter().map(|r| { - json!({ - "bibcode": r.bibcode, - "paragraph_index": r.paragraph_index, - "distance": r.distance, - "preview": truncate_content(&r.content, 100) - }) - }).collect(); - - ToolOutput::success( - truncate_content(&content, 4000), - json!({ "count": results.len(), "sources": sources }) - ) - } - Err(e) => ToolOutput::error(format!("RAG 语义检索失败: {}", e)), - } - } -} - -// ────────────────────────── 5. QueryTargetTool ────────────────────────── - -/// 天体信息查询工具：通过 CDS Sesame 查询天体物理属性 -pub struct QueryTargetTool; - -#[async_trait] -impl AgentTool for QueryTargetTool { - fn name(&self) -> &str { "query_target" } - - fn description(&self) -> &str { - "查询天体的基本物理属性信息。输入天体名称，返回坐标 (RA/Dec)、视星等、光谱型、视差等属性。\ - 数据来源为 CDS SIMBAD/Sesame 名称解析服务。适用于：获取天体基本参数、验证天体身份。" - } - - fn parameters(&self) -> serde_json::Value { - json!({ - "type": "object", - "properties": { - "object_name": { - "type": "string", - "description": "天体名称，如 'NGC 6752', 'GD 358', 'HD 209458', 'M 31' 等" - } - }, - "required": ["object_name"] - }) - } - - async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { - let object_name = match args.get("object_name").and_then(|n| n.as_str()) { - Some(n) => n.to_string(), - None => return ToolOutput::error("缺少必需参数 'object_name'"), - }; - - info!("[QueryTargetTool] 查询天体信息: {}", object_name); - let state = &ctx.app_state; - let client = reqwest::Client::new(); - - match crate::services::target::query_target_cached(&state.db, &object_name, None, &client).await { - Ok(info) => { - let content = format!( - "天体: {}\nRA: {}\nDec: {}\n视差: {}\n光谱型: {}\nV星等: {}\n别名: {}", - info.target_name, - info.ra.as_deref().unwrap_or("未知"), - info.dec.as_deref().unwrap_or("未知"), - info.parallax.map(|p| format!("{:.4} mas", p)).unwrap_or_else(|| "未知".to_string()), - info.spectral_type.as_deref().unwrap_or("未知"), - info.v_magnitude.map(|v| format!("{:.2}", v)).unwrap_or_else(|| "未知".to_string()), - if info.aliases.is_empty() { "无".to_string() } else { info.aliases.join(", ") } - ); - - ToolOutput::success( - content, - json!({ - "target_name": info.target_name, - "ra": info.ra, - "dec": info.dec, - "parallax": info.parallax, - "spectral_type": info.spectral_type, - "v_magnitude": info.v_magnitude, - "aliases": info.aliases - }) - ) - } - Err(e) => ToolOutput::error(format!("天体 '{}' 查询失败: {}", object_name, e)), - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_truncate_content_short() { - let text = "Hello, world!"; - assert_eq!(truncate_content(text, 100), text); - } - - #[test] - fn test_truncate_content_long() { - let text = "a".repeat(5000); - let result = truncate_content(&text, 100); - assert!(result.contains("内容已截断")); - assert!(result.contains("5000")); - } - - #[test] - fn test_tool_output_success() { - let output = ToolOutput::success("ok", json!({"key": "value"})); - assert!(!output.is_error); - assert_eq!(output.content, "ok"); - } - - #[test] - fn test_tool_output_error() { - let output = ToolOutput::error("something went wrong"); - assert!(output.is_error); - } - - #[test] - fn test_tool_registry_definitions() { - let registry = ToolRegistry::new(); - let defs = registry.definitions(); - assert_eq!(defs.len(), 7); - assert!(defs.iter().any(|d| d.function.name == "search_papers")); - assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata")); - assert!(defs.iter().any(|d| d.function.name == "download_paper")); - assert!(defs.iter().any(|d| d.function.name == "parse_paper")); - assert!(defs.iter().any(|d| d.function.name == "get_paper_content")); - assert!(defs.iter().any(|d| d.function.name == "rag_search")); - assert!(defs.iter().any(|d| d.function.name == "query_target")); - } - - #[test] - fn test_tool_registry_get() { - let registry = ToolRegistry::new(); - assert!(registry.get("search_papers").is_some()); - assert!(registry.get("nonexistent").is_none()); - } -} diff --git a/src/agent/tools/ask_user.rs b/src/agent/tools/ask_user.rs new file mode 100644 index 0000000..afb5afb --- /dev/null +++ b/src/agent/tools/ask_user.rs @@ -0,0 +1,248 @@ +// src/agent/tools/ask_user.rs +// +// 用户交互工具 — Agent 向用户提问并等待回复。 +// +// 当任务需求不明确时（如缺少参数、需要选择方案），Agent 调用此工具 +// 阻止 ReAct 循环，向用户展示问题，等待用户回复后继续执行。 +// +// 实现方式： +// - 使用 oneshot 通道向 SSE 层发送问题 +// - 阻塞等待用户通过 API 端点提交答案 +// - 超时后返回错误（默认 5 分钟） + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json::json; +use tokio::sync::oneshot; +use tracing::info; + +use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; + +/// 提交给用户的问题 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UserQuestion { + /// 问题 ID（用于前端关联回答） + pub question_id: String, + /// 完整问题文本 + pub question: String, + /// 短标签（显示为 chip/tag） + pub header: String, + /// 预定义选项列表 + pub options: Vec, + /// 是否允许多选 + pub multi_select: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UserOption { + pub label: String, + pub description: String, +} + +/// 用户的回答 +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct UserAnswer { + pub question_id: String, + pub answers: Vec, + pub free_text: Option, +} + +/// 向用户提问的工具。 +/// +/// 使用 oneshot 通道机制：创建问题 → 通过 AppState 发送 → 阻塞等待 → 返回答案。 +pub struct AskUserTool; + +#[async_trait] +impl AgentTool for AskUserTool { + fn name(&self) -> &str { + "ask_user" + } + + fn description(&self) -> &str { + "当任务需求不明确、缺少关键参数、或需要在多个方案之间选择时，向用户提问。\ + 支持预定义选项（单选/多选）。会暂停当前任务等待用户回复。\ + 适用场景：确认文献搜索范围、选择分析方案、确定输出格式。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "question": { + "type": "string", + "description": "要询问用户的问题，应清晰具体。如：'我应该搜索哪个天区的数据？'" + }, + "header": { + "type": "string", + "description": "问题的简短标签，最多 12 字。如：'搜索范围'、'分析方案'" + }, + "options": { + "type": "array", + "description": "预定义选项列表（可选，不提供则允许自由回答）", + "items": { + "type": "object", + "properties": { + "label": { + "type": "string", + "description": "选项标签，简洁明确。如：'Gaia DR3'" + }, + "description": { + "type": "string", + "description": "选项说明。如：'盖亚卫星第三期数据发布，包含 18 亿颗恒星'" + } + }, + "required": ["label", "description"] + } + }, + "multi_select": { + "type": "boolean", + "description": "是否允许多选，默认 false", + "default": false + } + }, + "required": ["question", "header"] + }) + } + + /// 必须阻塞等待用户回复 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + // 子代理（silent 模式）不能向用户提问 — 静默模式下无法交互 + if ctx.silent { + return ToolOutput::error( + "ask_user 在子代理/后台上下文中不可用。请基于已有信息继续，或使用其他工具获取所需数据。" + ); + } + + let question_text = match args.get("question").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'question'"), + }; + let header = match args.get("header").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'header'"), + }; + let multi_select = args + .get("multi_select") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let options: Vec = args + .get("options") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|opt| { + Some(UserOption { + label: opt.get("label")?.as_str()?.to_string(), + description: opt.get("description")?.as_str()?.to_string(), + }) + }) + .collect() + }) + .unwrap_or_default(); + + let question_id = uuid::Uuid::new_v4().to_string(); + let short_id = question_id[..8].to_string(); + + let question = UserQuestion { + question_id: short_id.clone(), + question: question_text.clone(), + header, + options, + multi_select, + }; + + info!( + "[AskUser] 向用户提问: id={}, header={}, options={}", + short_id, + &question.header, + question.options.len() + ); + + // 创建 oneshot 通道 + let (tx, rx) = oneshot::channel(); + let question_json = serde_json::to_string(&question).unwrap_or_default(); + + // 将通道发送端存储到 AppState 的待处理问题列表 + { + let mut pending = match ctx.app_state.pending_questions.lock() { + Ok(p) => p, + Err(_) => { + return ToolOutput::error("待处理问题队列不可用（内部锁异常），请稍后重试"); + } + }; + pending.insert( + short_id.clone(), + crate::api::PendingQuestion { + question_json: question_json.clone(), + answer_tx: tx, + }, + ); + } + + // 通过 SSE 通道发送问题事件（如果存在） + if let Some(sse_tx) = &ctx.app_state.sse_broadcast { + let _ = sse_tx.send(crate::api::AppEvent::UserQuestion { + data: question_json, + }); + } + + // 等待用户回答（5 分钟超时） + let timeout = tokio::time::Duration::from_secs(300); + match tokio::time::timeout(timeout, rx).await { + Ok(Ok(answer)) => { + info!( + "[AskUser] 收到用户回答: id={}, answers={:?}", + short_id, answer.answers + ); + + // 清理 + if let Ok(mut pending) = ctx.app_state.pending_questions.lock() { + pending.remove(&short_id); + } + + let free_text = answer.free_text.unwrap_or_default(); + let response = if !answer.answers.is_empty() { + let note = if free_text.is_empty() { + String::new() + } else { + format!("\n补充说明: {}", free_text) + }; + format!("用户回答: {}{}", answer.answers.join(", "), note) + } else { + free_text.clone() + }; + + ToolOutput::success( + response, + json!({ + "question_id": short_id, + "answers": answer.answers, + "free_text": free_text + }), + ) + } + Ok(Err(_)) => { + // 通道关闭（发送端被 drop） + if let Ok(mut pending) = ctx.app_state.pending_questions.lock() { + pending.remove(&short_id); + } + ToolOutput::error("用户取消了回答") + } + Err(_) => { + // 超时 + if let Ok(mut pending) = ctx.app_state.pending_questions.lock() { + pending.remove(&short_id); + } + ToolOutput::error(format!( + "等待用户回答超时 (5 分钟)。问题: {}", + question_text + )) + } + } + } +} diff --git a/src/agent/tools/astro/mod.rs b/src/agent/tools/astro/mod.rs new file mode 100644 index 0000000..9e68ba7 --- /dev/null +++ b/src/agent/tools/astro/mod.rs @@ -0,0 +1,15 @@ +// src/agent/tools/astro/mod.rs +// +// 天文科研工具集 — 文献搜索/下载/解析、RAG 检索、天体目标查询、研究笔记。 + +pub mod note; +pub mod paper; +pub mod rag; +pub mod search; +pub mod target; + +pub use note::SaveNoteTool; +pub use paper::{DownloadPaperTool, GetPaperContentTool, ParsePaperTool}; +pub use rag::RagSearchTool; +pub use search::{GetPaperMetadataTool, SearchPapersTool}; +pub use target::QueryTargetTool; diff --git a/src/agent/tools/astro/note.rs b/src/agent/tools/astro/note.rs new file mode 100644 index 0000000..8b6ce7a --- /dev/null +++ b/src/agent/tools/astro/note.rs @@ -0,0 +1,99 @@ +// src/agent/tools/note.rs — 研究笔记保存工具 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; + +/// 研究笔记保存工具：将 Agent 的研究发现保存为 Markdown 笔记 +pub struct SaveNoteTool; + +#[async_trait] +impl AgentTool for SaveNoteTool { + fn name(&self) -> &str { + "save_note" + } + + fn description(&self) -> &str { + "将研究中间结果或最终结论保存为 Markdown 格式的笔记文件。适用于：保存文献综述、记录研究发现、导出分析结果。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "title": { + "type": "string", + "description": "笔记标题（将作为文件名的一部分）" + }, + "content": { + "type": "string", + "description": "笔记正文内容，支持 Markdown 格式（包括 LaTeX 数学公式）" + } + }, + "required": ["title", "content"] + }) + } + + /// 写文件有副作用，中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let title = match args.get("title").and_then(|t| t.as_str()) { + Some(t) => t.to_string(), + None => return ToolOutput::error("缺少必需参数 'title'"), + }; + let content = match args.get("content").and_then(|c| c.as_str()) { + Some(c) => c.to_string(), + None => return ToolOutput::error("缺少必需参数 'content'"), + }; + + info!("[SaveNote] 保存笔记: title='{}'", title); + + let safe_filename: String = title + .chars() + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '_' || c == ' ' { + c + } else { + '_' + } + }) + .collect::() + .trim() + .replace(' ', "_"); + let filename = format!("{}.md", safe_filename); + + let notes_dir = ctx.app_state.config.library_dir.join("notes"); + if let Err(e) = std::fs::create_dir_all(¬es_dir) { + return ToolOutput::error(format!("无法创建笔记目录: {}", e)); + } + + let filepath = notes_dir.join(&filename); + let now = chrono::Local::now(); + let full_content = format!( + "---\ntitle: {}\ndate: {}\ngenerated_by: AstroResearch Agent\n---\n\n{}", + title, + now.format("%Y-%m-%d %H:%M:%S"), + content + ); + + match std::fs::write(&filepath, &full_content) { + Ok(_) => { + info!("[SaveNote] 笔记已保存: {}", filepath.display()); + ToolOutput::success( + format!( + "笔记已保存到 {} （{} 字符）", + filepath.display(), + content.len() + ), + json!({ "filename": filename, "path": filepath.to_string_lossy().to_string(), "size": content.len() }), + ) + } + Err(e) => ToolOutput::error(format!("保存笔记失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/paper.rs b/src/agent/tools/astro/paper.rs new file mode 100644 index 0000000..0ef7b19 --- /dev/null +++ b/src/agent/tools/astro/paper.rs @@ -0,0 +1,247 @@ +// 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)), + } + } +} + +// ── 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)), + } + } +} + +// ── 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)) + } + } + } + } +} diff --git a/src/agent/tools/astro/rag.rs b/src/agent/tools/astro/rag.rs new file mode 100644 index 0000000..057040b --- /dev/null +++ b/src/agent/tools/astro/rag.rs @@ -0,0 +1,104 @@ +// src/agent/tools/rag.rs — RAG 语义检索工具 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{truncate_content, AgentTool, ToolContext, ToolOutput}; + +/// RAG 向量检索工具：基于语义相似度检索文献切片 +pub struct RagSearchTool; + +#[async_trait] +impl AgentTool for RagSearchTool { + fn name(&self) -> &str { + "rag_search" + } + + fn description(&self) -> &str { + "在已向量化的文献库中进行语义检索。输入自然语言问题，返回最相关的文献片段。\ + 适用于：跨多篇文献查找特定信息、回答需要综合多个来源的问题。要求文献已完成向量化（embed）。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "question": { + "type": "string", + "description": "用于语义检索的自然语言问题" + }, + "top_k": { + "type": "integer", + "description": "返回最相关的片段数量，默认5", + "default": 5 + } + }, + "required": ["question"] + }) + } + + /// 纯读取向量数据库，并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let question = match args.get("question").and_then(|q| q.as_str()) { + Some(q) => q.to_string(), + None => return ToolOutput::error("缺少必需参数 'question'"), + }; + let top_k = args.get("top_k").and_then(|k| k.as_u64()).unwrap_or(5) as usize; + + info!( + "[RagSearch] 语义检索: question='{}', top_k={}", + question, top_k + ); + let state = &ctx.app_state; + + match crate::services::rag::retrieve(&state.db, &state.embedding, &question, top_k).await { + Ok(results) => { + if results.is_empty() { + return ToolOutput::success( + "未找到相关的文献片段。文献库中可能尚无向量化数据，请先对目标文献执行向量化操作。", + json!({ "count": 0 }), + ); + } + + let content = results + .iter() + .enumerate() + .map(|(i, r)| { + format!( + "[片段 {} | 来源: {} | 段落: {} | 相似度距离: {:.4}]\n{}", + i + 1, + r.bibcode, + r.paragraph_index, + r.distance, + r.content + ) + }) + .collect::>() + .join("\n\n---\n\n"); + + let sources: Vec = results + .iter() + .map(|r| { + json!({ + "bibcode": r.bibcode, + "paragraph_index": r.paragraph_index, + "distance": r.distance, + "preview": truncate_content(&r.content, 100) + }) + }) + .collect(); + + ToolOutput::success( + truncate_content(&content, 4000), + json!({ "count": results.len(), "sources": sources }), + ) + } + Err(e) => ToolOutput::error(format!("RAG 语义检索失败: {}", e)), + } + } +} diff --git a/src/agent/tools/astro/search.rs b/src/agent/tools/astro/search.rs new file mode 100644 index 0000000..9578485 --- /dev/null +++ b/src/agent/tools/astro/search.rs @@ -0,0 +1,189 @@ +// src/agent/tools/search.rs — 文献搜索与元数据工具 + +use async_trait::async_trait; +use serde_json::json; +use tracing::{error, info}; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +// ── SearchPapersTool ── + +/// 文献搜索工具：调用 ADS/arXiv 进行跨库检索 +pub struct SearchPapersTool; + +#[async_trait] +impl AgentTool for SearchPapersTool { + fn name(&self) -> &str { + "search_papers" + } + + fn description(&self) -> &str { + "搜索天文学文献。支持 NASA ADS 和 arXiv 跨平台联合检索，结果自动合并去重，并关联本地馆藏状态与引用关系网。输入关键词或高级检索式，返回匹配的文献列表。\ + 适用于：查找相关文献、了解研究领域现状、获取特定主题的论文。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "搜索关键词或高级检索式。支持语法：\ + 1. 字段限定：au:\"作者\" 或 author:\"作者\"、ti:\"标题\" 或 title:\"标题\"、abs:\"摘要关键字\"；\ + 2. 年份限定：year:2020（单年）或 year:2020-2025（年份区间）；\ + 3. 逻辑运算：支持 AND、OR、NOT 逻辑组合及括号分组，如 '(ti:subdwarf OR ti:\"white dwarf\") AND year:2020-2025'；\ + 4. 短语匹配：用双引号 \"\" 包含精确匹配短语，如 '\"Gaia BH1\"'。\ + 所有的中文标点符号（如\"\"()，；）在后台均会自动清洗转换。" + }, + "rows": { + "type": "integer", + "description": "返回结果数量，默认5，最大20", + "default": 5 + } + }, + "required": ["query"] + }) + } + + /// 纯读取外部 API，并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let query = match args.get("query").and_then(|q| q.as_str()) { + Some(q) => q.to_string(), + None => return ToolOutput::error("缺少必需参数 'query'"), + }; + let rows = args + .get("rows") + .and_then(|r| r.as_i64()) + .unwrap_or(5) + .min(20) as i32; + + info!( + "[SearchPapers] 执行文献搜索: query='{}', rows={}", + query, rows + ); + + let state = &ctx.app_state; + + match crate::services::search::search_papers(state, &query, "all", 0, rows, "relevance") + .await + { + Ok(results) => { + if results.is_empty() { + return ToolOutput::success( + "未找到匹配的文献。请尝试调整搜索关键词。", + json!({ "count": 0 }), + ); + } + + let display_results: Vec = results + .iter() + .map(|p| { + let first_author = p + .authors + .first() + .cloned() + .unwrap_or_else(|| "未知".to_string()); + json!({ + "bibcode": p.bibcode, + "title": p.title, + "first_author": first_author, + "year": p.year, + "citation_count": p.citation_count, + }) + }) + .collect(); + + let content = display_results + .iter() + .enumerate() + .map(|(i, r)| { + format!( + "{}. [{}] {} ({})\n 第一作者: {}\n 被引: {} 次", + i + 1, + r["bibcode"].as_str().unwrap_or(""), + r["title"].as_str().unwrap_or(""), + r["year"].as_str().unwrap_or(""), + r["first_author"].as_str().unwrap_or("未知"), + r["citation_count"].as_i64().unwrap_or(0) + ) + }) + .collect::>() + .join("\n\n"); + + ToolOutput::success( + content, + json!({ "count": results.len(), "papers": display_results }), + ) + } + Err(e) => { + error!("[SearchPapers] 检索失败: {}", e); + ToolOutput::error(format!("文献检索失败: {}", e)) + } + } + } +} + +// ── GetPaperMetadataTool ── + +/// 获取文献元数据工具：获取指定文献的完整元数据 +pub struct GetPaperMetadataTool; + +#[async_trait] +impl AgentTool for GetPaperMetadataTool { + fn name(&self) -> &str { + "get_paper_metadata" + } + + fn description(&self) -> &str { + "获取指定文献的完整元数据信息（包括完整标题、所有作者、出版期刊、关键字、引用数、完整摘要等）。\ + 适用于：需要查看某篇文献的详细信息、阅读完整摘要以评估文献相关性。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "bibcode": { + "type": "string", + "description": "文献的唯一标识符，支持 ADS Bibcode（如 '2024ApJ...960..123A'）、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!("[GetPaperMetadata] 获取文献元数据: {}", bibcode); + let state = &ctx.app_state; + + match crate::api::helpers::get_paper_from_db(&state.db, &state.config.library_dir, &bibcode) + .await + { + Ok(paper) => { + let content = format!( + "文献元数据 [{}]：\n 标题: {}\n 作者: {}\n 年份: {}\n 期刊: {}\n 关键字: {}\n DOI: {}\n arXiv ID: {}\n 引用数: {} 次\n 参考文献数: {} 次\n 文献类型: {}\n 已下载: {}\n 已解析为 Markdown: {}\n 摘要:\n{}", + paper.bibcode, paper.title, paper.authors.join(", "), paper.year, + paper.pub_journal, paper.keywords.join(", "), paper.doi, paper.arxiv_id, + paper.citation_count, paper.reference_count, paper.doctype, + paper.is_downloaded, paper.has_markdown, paper.abstract_text + ); + ToolOutput::success(content, json!(paper)) + } + Err(e) => ToolOutput::error(format!("获取文献 {} 元数据失败: {}", bibcode, e)), + } + } +} diff --git a/src/agent/tools/astro/target.rs b/src/agent/tools/astro/target.rs new file mode 100644 index 0000000..d7ca19d --- /dev/null +++ b/src/agent/tools/astro/target.rs @@ -0,0 +1,89 @@ +// src/agent/tools/target.rs — 天体物理查询工具 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +/// 天体信息查询工具：通过 CDS Sesame 查询天体物理属性 +pub struct QueryTargetTool; + +#[async_trait] +impl AgentTool for QueryTargetTool { + fn name(&self) -> &str { + "query_target" + } + + fn description(&self) -> &str { + "查询天体的基本物理属性信息。输入天体名称，返回坐标 (RA/Dec)、视星等、光谱型、视差等属性。\ + 数据来源为 CDS SIMBAD/Sesame 名称解析服务。适用于：获取天体基本参数、验证天体身份。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "object_name": { + "type": "string", + "description": "天体名称，如 'NGC 6752', 'GD 358', 'HD 209458', 'M 31' 等" + } + }, + "required": ["object_name"] + }) + } + + /// 纯读取外部天体数据库，并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let object_name = match args.get("object_name").and_then(|n| n.as_str()) { + Some(n) => n.to_string(), + None => return ToolOutput::error("缺少必需参数 'object_name'"), + }; + + info!("[QueryTarget] 查询天体: {}", object_name); + let state = &ctx.app_state; + let client = reqwest::Client::new(); + + match crate::services::target::query_target_cached(&state.db, &object_name, None, &client) + .await + { + Ok(info) => { + let content = format!( + "天体: {}\nRA: {}\nDec: {}\n视差: {}\n光谱型: {}\nV星等: {}\n别名: {}", + info.target_name, + info.ra.as_deref().unwrap_or("未知"), + info.dec.as_deref().unwrap_or("未知"), + info.parallax + .map(|p| format!("{:.4} mas", p)) + .unwrap_or_else(|| "未知".to_string()), + info.spectral_type.as_deref().unwrap_or("未知"), + info.v_magnitude + .map(|v| format!("{:.2}", v)) + .unwrap_or_else(|| "未知".to_string()), + if info.aliases.is_empty() { + "无".to_string() + } else { + info.aliases.join(", ") + } + ); + + ToolOutput::success( + content, + json!({ + "target_name": info.target_name, + "ra": info.ra, "dec": info.dec, + "parallax": info.parallax, + "spectral_type": info.spectral_type, + "v_magnitude": info.v_magnitude, + "aliases": info.aliases + }), + ) + } + Err(e) => ToolOutput::error(format!("天体 '{}' 查询失败: {}", object_name, e)), + } + } +} diff --git a/src/agent/tools/background.rs b/src/agent/tools/background.rs new file mode 100644 index 0000000..a4ad8c5 --- /dev/null +++ b/src/agent/tools/background.rs @@ -0,0 +1,206 @@ +// src/agent/tools/background.rs — 后台任务工具 (bg_task_run / bg_task_check) +// +// 参考 Claude Code s08 Background Tasks 设计。 +// 慢速操作可在后台异步执行，结果在下一轮 LLM 调用前注入上下文。 + +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tracing::info; + +use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; +use crate::agent::background::{self, BgNotificationQueue}; + +/// 支持后台执行的工具列表 +const BG_SUPPORTED_TOOLS: &[&str] = &["download_paper", "parse_paper"]; + +/// 后台任务启动工具 +pub struct BgTaskRunTool { + queue: Arc, +} + +impl BgTaskRunTool { + pub fn new(queue: Arc) -> Self { + BgTaskRunTool { queue } + } +} + +#[async_trait] +impl AgentTool for BgTaskRunTool { + fn name(&self) -> &str { + "bg_task_run" + } + + fn description(&self) -> &str { + "在后台异步执行慢速工具（download_paper, parse_paper）。\ + 返回任务ID后立即让 LLM 继续思考，后台完成的结果会在下一轮对话中自动通知。\ + 适用于：下载PDF、解析论文等耗时操作。使用 bg_task_check 查询任务状态。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "tool_name": { + "type": "string", + "description": "要在后台执行的工具名称（download_paper 或 parse_paper）", + "enum": ["download_paper", "parse_paper"] + }, + "bibcode": { + "type": "string", + "description": "文献的唯一标识符（ADS bibcode）" + } + }, + "required": ["tool_name", "bibcode"] + }) + } + + /// 后台任务启动有副作用（spawn tokio task），中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let tool_name = match args.get("tool_name").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'tool_name'"), + }; + + let bibcode = match args.get("bibcode").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'bibcode'"), + }; + + if !BG_SUPPORTED_TOOLS.contains(&tool_name.as_str()) { + return ToolOutput::error(format!( + "工具 '{}' 不支持后台执行。支持的工具: {}", + tool_name, + BG_SUPPORTED_TOOLS.join(", ") + )); + } + + info!( + "[BgTaskRun] 启动后台任务: tool={}, bibcode={}", + tool_name, bibcode + ); + + let handle = background::spawn_background_task( + ctx.app_state.clone(), + self.queue.clone(), + tool_name.clone(), + bibcode.clone(), + ) + .await; + + ToolOutput::success( + format!( + "✅ 后台任务已启动。\n\ + 任务ID: {}\n\ + 工具: {}\n\ + 文献: {}\n\ + 状态: 运行中\n\n\ + 使用 bg_task_check 查询任务状态。完成后结果会自动通知。", + handle.task_id, handle.tool_name, handle.bibcode, + ), + json!({ + "task_id": handle.task_id, + "tool_name": handle.tool_name, + "bibcode": handle.bibcode, + "status": "running" + }), + ) + } +} + +/// 后台任务查询工具 +pub struct BgTaskCheckTool { + queue: Arc, +} + +impl BgTaskCheckTool { + pub fn new(queue: Arc) -> Self { + BgTaskCheckTool { queue } + } +} + +#[async_trait] +impl AgentTool for BgTaskCheckTool { + fn name(&self) -> &str { + "bg_task_check" + } + + fn description(&self) -> &str { + "查询后台任务状态。不指定 task_id 时返回所有任务。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "task_id": { + "type": "string", + "description": "可选：要查询的任务ID。不指定则返回所有任务。" + } + }, + "required": [] + }) + } + + /// 纯读取内存队列状态，并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let task_id = args.get("task_id").and_then(|v| v.as_str()); + + match task_id { + Some(tid) => match self.queue.get_task(tid).await { + Some(task) => { + let status_icon = match task.status { + background::BgTaskStatus::Running => "🔄", + background::BgTaskStatus::Completed => "✅", + background::BgTaskStatus::Failed => "❌", + }; + ToolOutput::success( + format!( + "{} 任务 {}: {} ({})\n文献: {}", + status_icon, + task.task_id, + task.tool_name, + match task.status { + background::BgTaskStatus::Running => "运行中", + background::BgTaskStatus::Completed => "已完成", + background::BgTaskStatus::Failed => "失败", + }, + task.bibcode, + ), + serde_json::to_value(&task).unwrap_or_default(), + ) + } + None => ToolOutput::error(format!("任务 '{}' 未找到", tid)), + }, + None => { + let tasks = self.queue.get_all_tasks().await; + if tasks.is_empty() { + return ToolOutput::success("当前没有后台任务。", json!({ "tasks": [] })); + } + + let mut lines = vec!["📊 后台任务状态：\n".to_string()]; + for task in &tasks { + let icon = match task.status { + background::BgTaskStatus::Running => "🔄", + background::BgTaskStatus::Completed => "✅", + background::BgTaskStatus::Failed => "❌", + }; + lines.push(format!( + "{} [{}] {} — {}", + icon, task.task_id, task.tool_name, task.bibcode + )); + } + + ToolOutput::success(lines.join("\n"), json!({ "tasks": tasks })) + } + } + } +} diff --git a/src/agent/tools/compress.rs b/src/agent/tools/compress.rs new file mode 100644 index 0000000..f807fd2 --- /dev/null +++ b/src/agent/tools/compress.rs @@ -0,0 +1,42 @@ +// src/agent/tools/compress.rs — 手动上下文压缩工具 + +use async_trait::async_trait; +use serde_json::json; + +use super::{AgentTool, ToolContext, ToolOutput}; + +/// 手动上下文压缩工具：LLM 可主动调用以压缩对话历史 +pub struct CompressTool; + +#[async_trait] +impl AgentTool for CompressTool { + fn name(&self) -> &str { + "compress_context" + } + + fn description(&self) -> &str { + "手动压缩对话上下文。当你发现对话历史过长、token 消耗过大时，主动调用此工具进行压缩以释放空间。\ + 压缩后历史对话将被摘要替代，但关键信息不会丢失。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": {}, + "required": [] + }) + } + + /// 设置标志位是幂等操作，并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, _args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + // 实际压缩由 AgentRuntime 通过 pending_manual_compress 标志位处理 + ToolOutput::success( + "上下文压缩标记已设置。当前对话历史将在下一轮 LLM 调用前被压缩。", + json!({ "action": "compress" }), + ) + } +} diff --git a/src/agent/tools/filesystem/bash.rs b/src/agent/tools/filesystem/bash.rs new file mode 100644 index 0000000..f417891 --- /dev/null +++ b/src/agent/tools/filesystem/bash.rs @@ -0,0 +1,226 @@ +use async_trait::async_trait; +use serde_json::json; +use std::time::Duration; +use tracing::info; + +use crate::agent::tools::filesystem::security::{ + has_path_traversal, is_path_allowed, resolve_path, +}; +use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; + +// ── run_bash ──────────────────────────────────────────────────────────────── + +/// Bash 命令执行工具。 +/// +/// 允许 Agent 执行只读或数据处理的 Shell 命令。 +pub struct RunBashTool; + +#[async_trait] +impl AgentTool for RunBashTool { + fn name(&self) -> &str { + "run_bash" + } + + fn description(&self) -> &str { + "执行一个 Bash 命令并返回 stdout 和 stderr。\ + 适用场景：运行 skill 目录中的 scripts/*.py、文件处理、数据提取。\ + 限制：最大执行时间 60 秒，输出自动截断至 4000 字符。\ + 禁止交互式命令和破坏性操作。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "command": { + "type": "string", + "description": "要执行的命令。如 'wc -l file.txt' 或 'python script.py'" + }, + "working_dir": { + "type": "string", + "description": "工作目录（默认项目根目录）" + }, + "timeout_secs": { + "type": "integer", + "description": "超时时间（秒），默认 60，最大 120", + "default": 60 + } + }, + "required": ["command"] + }) + } + + /// Bash 执行不应被中断（可能有写操作） + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let command = match args.get("command").and_then(|v| v.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'command'"), + }; + + // 命令安全校验 + if let Some(rejection) = validate_bash_command(command) { + return ToolOutput::error(rejection); + } + + let timeout_secs: u64 = args + .get("timeout_secs") + .and_then(|v| v.as_i64()) + .map(|v| (v.max(1) as u64).min(120)) + .unwrap_or(60); + + // 工作目录 + let working_dir = match args.get("working_dir").and_then(|v| v.as_str()) { + Some(dir_str) => { + if has_path_traversal(dir_str) { + return ToolOutput::error("工作目录路径包含非法字符"); + } + match resolve_path(dir_str) { + Some(p) if is_path_allowed(&p, ctx) => p, + Some(_) => return ToolOutput::error("无权访问指定的工作目录"), + None => return ToolOutput::error("无法解析工作目录路径"), + } + } + None => match std::env::current_dir() { + Ok(d) => d, + Err(e) => return ToolOutput::error(format!("无法获取当前目录: {}", e)), + }, + }; + + info!( + "[RunBash] executing: {} (cwd: {}, timeout: {}s)", + command, + working_dir.display(), + timeout_secs + ); + + // 执行命令 + let timeout = Duration::from_secs(timeout_secs); + let result = tokio::time::timeout( + timeout, + tokio::process::Command::new("bash") + .arg("-c") + .arg(command) + .current_dir(&working_dir) + .output(), + ) + .await; + + match result { + Ok(Ok(output)) => { + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + let exit_code = output.status.code().unwrap_or(-1); + + let mut content = if exit_code == 0 { + String::new() + } else { + format!("[退出码: {}]\n", exit_code) + }; + + if !stdout.trim().is_empty() { + content.push_str(&stdout); + } + if !stderr.trim().is_empty() { + if !content.is_empty() { + content.push('\n'); + } + content.push_str(&format!("[stderr]\n{}", stderr)); + } + if content.is_empty() { + content = "(无输出)".to_string(); + } + + let output_truncated = crate::agent::tools::truncate_content(&content, 4000); + + info!( + "[RunBash] completed with exit code {} ({} chars output)", + exit_code, + content.len() + ); + + ToolOutput::success( + output_truncated, + json!({ + "exit_code": exit_code, + "stdout_length": stdout.len(), + "stderr_length": stderr.len() + }), + ) + } + Ok(Err(e)) => ToolOutput::error(format!("命令执行失败: {}", e)), + Err(_) => ToolOutput::error(format!("命令执行超时 ({}s): {}", timeout_secs, command)), + } + } +} + +/// 验证 Bash 命令安全性（黑名单 + 启发式检查）。 +/// 返回 `Some(reason)` 表示拒绝，`None` 表示允许。 +/// 安全命令白名单：这些命令的第一个单词匹配时，自动允许（仍需路径沙箱检查）。 +#[allow(dead_code)] +const SAFE_COMMANDS: &[&str] = &[ + "ls", "cat", "head", "tail", "find", "grep", "wc", "echo", "pwd", "sort", "uniq", "cut", "tr", + "awk", "sed", "jq", "diff", "file", "stat", "du", "df", "env", "printenv", "which", "basename", + "dirname", "realpath", "readlink", "xargs", "tee", "date", "sleep", "true", "false", +]; + +/// 检查命令是否属于安全白名单（第一个单词匹配即可）。 +#[allow(dead_code)] +fn is_safe_command(command: &str) -> bool { + let first_word = command.split_whitespace().next().unwrap_or(""); + SAFE_COMMANDS.contains(&first_word) +} + +fn validate_bash_command(command: &str) -> Option { + let trimmed = command.trim(); + + // 空命令 + if trimmed.is_empty() || trimmed == "bash" || trimmed == "bash -c" { + return Some("不允许执行空命令".to_string()); + } + + // 禁止交互式/破坏性命令 + let interactive_patterns = [ + "sudo ", + "su ", + "passwd", + "ssh ", + "telnet ", + "login", + "less ", + "more ", + "vim ", + "vi ", + "nano ", + "emacs ", + "top", + "htop", + "watch ", + "tail -f", + "rm -rf /", + "mkfs.", + "dd if=", + "chmod 777", + "> /dev/", + ]; + + let lower = trimmed.to_lowercase(); + for pattern in &interactive_patterns { + if lower.contains(&pattern.to_lowercase()) { + return Some(format!("不允许执行 '{}' 类命令", pattern)); + } + } + + // 允许通过 + None +} + +/// 检查命令是否需要用户权限确认。 +/// 安全白名单中的命令不需确认，其他命令需要。 +#[allow(dead_code)] +pub fn bash_needs_permission(command: &str) -> bool { + !is_safe_command(command) +} diff --git a/src/agent/tools/filesystem/edit.rs b/src/agent/tools/filesystem/edit.rs new file mode 100644 index 0000000..2e15553 --- /dev/null +++ b/src/agent/tools/filesystem/edit.rs @@ -0,0 +1,336 @@ +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::filesystem::security::{ + has_path_traversal, is_path_allowed, resolve_path, +}; +use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; + +// ── file_edit helpers ────────────────────────────────────────────────────── + +/// 将弯引号（curly quotes）转换为直引号（straight quotes）。 +/// LLM 输出的是直引号，但文件中可能使用弯引号，需要统一后匹配。 +fn normalize_quotes(s: &str) -> String { + s.replace(['\u{2018}', '\u{2019}'], "'") // right single curly + .replace(['\u{201c}', '\u{201d}'], "\"") // right double curly +} + +/// 在文件内容中查找 old_string，兼容引号差异。 +/// +/// 先尝试精确匹配，失败后用引号规范化再试。 +/// 返回文件中的实际字符串（用于替换），如果没找到则返回 None。 +fn find_actual_string<'a>(file_content: &'a str, old_string: &str) -> Option<&'a str> { + // 1. 精确匹配 + if let Some(pos) = file_content.find(old_string) { + return Some(&file_content[pos..pos + old_string.len()]); + } + + // 2. 引号规范化后匹配。 + // 弯引号(3 bytes) 和直引号(1 byte) 的字节长度不同，所以不能直接用 + // normalized_file 中的字节位置去切 original file。改用字符位置对齐。 + let normalized_search = normalize_quotes(old_string); + let normalized_file = normalize_quotes(file_content); + + // 找到匹配在归一化后的文件中的字符偏移 + let norm_char_pos = normalized_file.find(&normalized_search)?; + // 统计归一化文件中 norm_char_pos 字节对应的字符数 + let char_start: usize = normalized_file[..norm_char_pos].chars().count(); + let char_len: usize = old_string.chars().count(); + + // 在原文件中找到对应的字节范围 + let orig_chars: Vec<(usize, char)> = file_content.char_indices().collect(); + if char_start + char_len > orig_chars.len() { + return None; + } + let start_byte = orig_chars[char_start].0; + let end_byte = if char_start + char_len < orig_chars.len() { + orig_chars[char_start + char_len].0 + } else { + file_content.len() + }; + + Some(&file_content[start_byte..end_byte]) +} + +/// 当 old_string 通过引号规范化才匹配成功时， +/// 对 new_string 施加相同的弯引号风格，保持文件风格一致。 +fn preserve_quote_style(old_string: &str, actual_old: &str, new_string: &str) -> String { + if old_string == actual_old { + return new_string.to_string(); + } + + let has_curly_single = actual_old.contains('\u{2018}') || actual_old.contains('\u{2019}'); + let has_curly_double = actual_old.contains('\u{201c}') || actual_old.contains('\u{201d}'); + + if !has_curly_single && !has_curly_double { + return new_string.to_string(); + } + + let mut result = new_string.to_string(); + + if has_curly_double { + result = apply_curly_double_quotes(&result); + } + if has_curly_single { + result = apply_curly_single_quotes(&result); + } + + result +} + +/// 将直双引号替换为弯双引号（根据上下文判断开/闭） +fn apply_curly_double_quotes(s: &str) -> String { + let chars: Vec = s.chars().collect(); + let mut result = String::with_capacity(s.len()); + for (i, &ch) in chars.iter().enumerate() { + if ch == '"' { + if is_opening_context(&chars, i) { + result.push('\u{201c}'); // left double curly + } else { + result.push('\u{201d}'); // right double curly + } + } else { + result.push(ch); + } + } + result +} + +/// 将直单引号替换为弯单引号（跳过缩略形式如 don't, it's） +fn apply_curly_single_quotes(s: &str) -> String { + let chars: Vec = s.chars().collect(); + let mut result = String::with_capacity(s.len()); + for (i, &ch) in chars.iter().enumerate() { + if ch == '\'' { + let prev_is_letter = i > 0 && chars[i - 1].is_alphabetic(); + let next_is_letter = i + 1 < chars.len() && chars[i + 1].is_alphabetic(); + if prev_is_letter && next_is_letter { + // 缩略形式 (don't, it's) — 使用右弯单引号 + result.push('\u{2019}'); + } else if is_opening_context(&chars, i) { + result.push('\u{2018}'); // left single curly + } else { + result.push('\u{2019}'); // right single curly + } + } else { + result.push(ch); + } + } + result +} + +fn is_opening_context(chars: &[char], index: usize) -> bool { + if index == 0 { + return true; + } + let prev = chars[index - 1]; + matches!(prev, ' ' | '\t' | '\n' | '\r' | '(' | '[' | '{') +} + +/// 去掉每行末尾的空白字符（保留换行符）。 +/// 对非 .md/.mdx 文件使用，因为 Markdown 中两个尾随空格表示硬换行。 +fn strip_trailing_whitespace(s: &str) -> String { + let mut result = String::with_capacity(s.len()); + for line in s.lines() { + let trimmed = line.trim_end(); + result.push_str(trimmed); + result.push('\n'); + } + // 如果原字符串不以换行结尾，去掉我们添加的换行 + if !s.ends_with('\n') && !result.is_empty() { + result.pop(); + } + result +} + +// ── file_edit ────────────────────────────────────────────────────────────── + +/// 文件编辑工具（精确字符串替换）。 +/// +/// 在已有文件中查找 old_string 并替换为 new_string。 +/// old_string 必须唯一匹配（防止误改）。 +pub struct FileEditTool; + +#[async_trait] +impl AgentTool for FileEditTool { + fn name(&self) -> &str { + "file_edit" + } + + fn description(&self) -> &str { + "在文件中执行精确的字符串替换。查找 old_string 并替换为 new_string。\ + old_string 在文件中必须唯一（仅出现一次），以防止意外破坏其他内容。\ + 适用场景：修改脚本参数、更新配置值、在 Markdown 笔记中追加或修正内容。\ + 注意：仅替换匹配片段，文件其余部分保持不变。如需完整覆写请使用 file_write。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "file_path": { + "type": "string", + "description": "要编辑的文件路径" + }, + "old_string": { + "type": "string", + "description": "要被替换的原字符串（必须唯一匹配）" + }, + "new_string": { + "type": "string", + "description": "替换后的新字符串" + }, + "replace_all": { + "type": "boolean", + "description": "替换所有匹配（默认 false，要求 old_string 唯一）", + "default": false + } + }, + "required": ["file_path", "old_string", "new_string"] + }) + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + false + } + + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let path_str = match args.get("file_path").and_then(|v| v.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'file_path'"), + }; + let old_string = match args.get("old_string").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'old_string'"), + }; + let new_string = match args.get("new_string").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'new_string'"), + }; + let replace_all = args + .get("replace_all") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if has_path_traversal(path_str) { + return ToolOutput::error("路径包含非法字符（.. 或 ~），拒绝访问"); + } + + let file_path = match resolve_path(path_str) { + Some(p) => p, + None => return ToolOutput::error("无法解析文件路径"), + }; + + if !is_path_allowed(&file_path, ctx) { + return ToolOutput::error(format!("无权编辑路径 '{}'", path_str)); + } + + if !file_path.exists() { + return ToolOutput::error(format!("文件不存在: {}", file_path.display())); + } + + if old_string.is_empty() { + return ToolOutput::error("old_string 不能为空"); + } + + if old_string == new_string { + return ToolOutput::success( + "未做任何更改：old_string 与 new_string 完全相同。", + json!({ "file_path": file_path.to_string_lossy(), "changed": false }), + ); + } + + let original = match std::fs::read_to_string(&file_path) { + Ok(c) => c, + Err(e) => return ToolOutput::error(format!("读取文件失败: {}", e)), + }; + + // 对非 Markdown 文件去除 LLM 输出中常见的尾随空白 + let is_markdown = file_path + .extension() + .map(|e| e == "md" || e == "mdx") + .unwrap_or(false); + let (old_string, new_string) = if is_markdown { + (old_string, new_string) + } else { + ( + strip_trailing_whitespace(&old_string), + strip_trailing_whitespace(&new_string), + ) + }; + + // 查找文件中的实际字符串（兼容引号差异） + let actual_old = match find_actual_string(&original, &old_string) { + Some(s) => s.to_string(), + None => { + return ToolOutput::error(format!( + "在文件中未找到要替换的字符串。\n查找内容: {}\n提示: 请确认原文内容完全一致（含空格/换行），\ + 或使用 read_file 重新读取文件确认当前内容。", + old_string + )); + } + }; + + // 保持文件的引号风格 + let actual_new = preserve_quote_style(&old_string, &actual_old, &new_string); + + // 检查匹配次数 + let match_count = original.matches(&actual_old).count(); + + if !replace_all && match_count > 1 { + // 提取上下文帮助 LLM 定位 + let mut ctx_lines: Vec = Vec::new(); + for (line_no, line) in original.lines().enumerate() { + if line.contains(&actual_old) { + ctx_lines.push(format!(" L{}: {}", line_no + 1, line.trim())); + } + if ctx_lines.len() >= 5 { + break; + } // 最多显示 5 处 + } + return ToolOutput::error(format!( + "'{}' 在文件中出现了 {} 处，无法确定要修改哪一个。\n\ + 出现位置:\n{}\n\ + 请包含更多上下文（前后行）使 old_string 唯一， + 或设置 replace_all: true 以替换全部匹配。", + actual_old, + match_count, + ctx_lines.join("\n") + )); + } + + let modified = if replace_all { + original.replace(&actual_old, &actual_new) + } else { + original.replacen(&actual_old, &actual_new, 1) + }; + + if modified == original { + return ToolOutput::success( + "文件内容未发生变化（new_string 与 old_string 相同）。", + json!({ "file_path": file_path.to_string_lossy(), "changed": false }), + ); + } + + match std::fs::write(&file_path, &modified) { + Ok(()) => { + let count = if replace_all { match_count } else { 1 }; + info!("[FileEdit] {} 处替换: {}", count, file_path.display()); + ToolOutput::success( + format!("文件 {} 已更新（{} 处替换）。", file_path.display(), count), + json!({ + "file_path": file_path.to_string_lossy(), + "replacements": count, + "changed": true + }), + ) + } + Err(e) => ToolOutput::error(format!("写入失败: {}", e)), + } + } +} diff --git a/src/agent/tools/filesystem/glob.rs b/src/agent/tools/filesystem/glob.rs new file mode 100644 index 0000000..437dcc4 --- /dev/null +++ b/src/agent/tools/filesystem/glob.rs @@ -0,0 +1,100 @@ +// src/agent/tools/filesystem/glob.rs +// +// glob_files 工具 — 基于 glob 模式查找文件。 + +use async_trait::async_trait; +use serde_json::json; + +use crate::agent::tools::filesystem::security::{ + has_path_traversal, is_path_allowed, resolve_path, +}; +use crate::agent::tools::truncate_content; +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +/// 基于 glob 模式查找文件工具。 +pub struct GlobFilesTool; + +#[async_trait] +impl AgentTool for GlobFilesTool { + fn name(&self) -> &str { + "glob_files" + } + + fn description(&self) -> &str { + "使用 glob 模式查找匹配的文件。支持通配符 * 和 **。返回匹配的文件路径列表。路径必须在允许的沙箱范围内。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "glob 匹配模式（如 '**/*.rs', 'src/**/*.md'）" + }, + "path": { + "type": "string", + "description": "搜索起始路径，默认为当前工作目录", + "default": "." + } + }, + "required": ["pattern"] + }) + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let pattern = match args.get("pattern").and_then(|v| v.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'pattern'"), + }; + + let search_path_str = args.get("path").and_then(|v| v.as_str()).unwrap_or("."); + + if has_path_traversal(search_path_str) { + return ToolOutput::error("路径包含不安全字符"); + } + + let search_path = match resolve_path(search_path_str) { + Some(p) => p, + None => return ToolOutput::error(format!("无法解析路径: {}", search_path_str)), + }; + + if !is_path_allowed(&search_path, ctx) { + return ToolOutput::error("路径不在允许的沙箱范围内"); + } + + let glob_pattern = search_path.join(pattern); + let pattern_str = glob_pattern.to_string_lossy().to_string(); + + match glob::glob(&pattern_str) { + Ok(paths) => { + let results: Vec = paths + .flatten() + .filter(|p| p.is_file()) + .filter_map(|p| { + let relative = p.strip_prefix(&search_path).ok()?; + Some(relative.display().to_string()) + }) + .take(200) + .collect(); + + if results.is_empty() { + ToolOutput::success( + format!("未找到匹配 '{}' 的文件", pattern), + json!({"pattern": pattern, "matches": 0}), + ) + } else { + let count = results.len(); + let output = results.join("\n"); + let truncated = truncate_content(&output, 4000); + ToolOutput::success(truncated, json!({"pattern": pattern, "matches": count})) + } + } + Err(e) => ToolOutput::error(format!("glob 模式无效: {}", e)), + } + } +} diff --git a/src/agent/tools/filesystem/grep.rs b/src/agent/tools/filesystem/grep.rs new file mode 100644 index 0000000..ddb345b --- /dev/null +++ b/src/agent/tools/filesystem/grep.rs @@ -0,0 +1,203 @@ +// src/agent/tools/filesystem/grep.rs +// +// grep_files 工具 — 在文件中搜索匹配模式的行（类似 grep 命令）。 + +use async_trait::async_trait; +use serde_json::json; +use std::path::Path; + +use crate::agent::tools::filesystem::security::{ + has_path_traversal, is_path_allowed, resolve_path, +}; +use crate::agent::tools::truncate_content; +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +/// 在文件中搜索匹配模式的行（类似 grep 命令）。 +pub struct GrepFilesTool; + +#[async_trait] +impl AgentTool for GrepFilesTool { + fn name(&self) -> &str { + "grep_files" + } + + fn description(&self) -> &str { + "在指定目录或文件中搜索匹配正则表达式的行。返回匹配行及其上下文（前后各 2 行）。适用于在代码库或文献中搜索特定模式。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "pattern": { + "type": "string", + "description": "要搜索的正则表达式模式" + }, + "path": { + "type": "string", + "description": "搜索路径（文件或目录），默认为当前工作目录", + "default": "." + }, + "include": { + "type": "string", + "description": "文件过滤 glob 模式（如 '*.rs', '*.md'），默认所有文本文件", + "default": null + }, + "max_results": { + "type": "integer", + "description": "最大结果数，默认 50", + "default": 50 + } + }, + "required": ["pattern"] + }) + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let pattern = match args.get("pattern").and_then(|v| v.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'pattern'"), + }; + + let search_path_str = args.get("path").and_then(|v| v.as_str()).unwrap_or("."); + + if has_path_traversal(search_path_str) { + return ToolOutput::error("路径包含不安全字符"); + } + + let search_path = match resolve_path(search_path_str) { + Some(p) => p, + None => return ToolOutput::error(format!("无法解析路径: {}", search_path_str)), + }; + + if !is_path_allowed(&search_path, ctx) { + return ToolOutput::error("路径不在允许的沙箱范围内"); + } + + let include_pattern = args.get("include").and_then(|v| v.as_str()); + let max_results = args + .get("max_results") + .and_then(|v| v.as_i64()) + .unwrap_or(50) as usize; + + let re = match regex::Regex::new(pattern) { + Ok(r) => r, + Err(e) => return ToolOutput::error(format!("正则表达式无效: {}", e)), + }; + + let mut results = Vec::new(); + let mut count = 0; + + if search_path.is_file() { + count += Self::search_file(&search_path, &re, &mut results); + } else if search_path.is_dir() { + count += Self::search_dir( + &search_path, + &re, + include_pattern, + max_results, + &mut results, + ); + } + + if results.is_empty() { + ToolOutput::success( + format!("未找到匹配 '{}' 的结果 (搜索了 {} 个位置)", pattern, count), + json!({"pattern": pattern, "matches": 0, "files_searched": count}), + ) + } else { + let output = results.join("\n---\n"); + let truncated = truncate_content(&output, 4000); + ToolOutput::success( + truncated, + json!({"pattern": pattern, "matches": results.len(), "files_searched": count}), + ) + } + } +} + +impl GrepFilesTool { + fn search_file(path: &Path, re: ®ex::Regex, results: &mut Vec) -> usize { + let content = match std::fs::read_to_string(path) { + Ok(c) => c, + Err(_) => return 0, + }; + let lines: Vec<&str> = content.lines().collect(); + let mut matched = false; + for (i, line) in lines.iter().enumerate() { + if re.is_match(line) { + if !matched { + results.push(format!("📄 {}:", path.display())); + matched = true; + } + let ctx_start = i.saturating_sub(2); + let ctx_end = (i + 3).min(lines.len()); + for (j, line) in lines.iter().enumerate().take(ctx_end).skip(ctx_start) { + let marker = if j == i { ">" } else { " " }; + results.push(format!(" {} {:4}: {}", marker, j + 1, line)); + } + results.push(String::new()); + } + } + 1 + } + + fn search_dir( + dir: &Path, + re: ®ex::Regex, + include: Option<&str>, + max_results: usize, + results: &mut Vec, + ) -> usize { + let mut count = 0; + let text_extensions = [ + "rs", "py", "js", "ts", "jsx", "tsx", "html", "css", "md", "txt", "json", "yaml", + "yml", "toml", "sh", "sql", "c", "cpp", "h", "hpp", "java", "go", "rb", "php", "swift", + ]; + if let Ok(entries) = std::fs::read_dir(dir) { + for entry in entries.flatten() { + if results.len() >= max_results { + break; + } + let path = entry.path(); + if path.is_dir() { + let dir_name = path.file_name().and_then(|n| n.to_str()).unwrap_or(""); + if !dir_name.starts_with('.') + && dir_name != "target" + && dir_name != "node_modules" + { + count += Self::search_dir(&path, re, include, max_results, results); + } + } else if path.is_file() { + if let Some(inc) = include { + let file_name = path.file_name().and_then(|n| n.to_str()).unwrap_or(""); + if !glob_match(inc, file_name) { + continue; + } + } else { + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); + if !text_extensions.contains(&ext) && !ext.is_empty() { + continue; + } + } + count += Self::search_file(&path, re, results); + } + } + } + count + } +} + +fn glob_match(pattern: &str, name: &str) -> bool { + if pattern == "*" { + return true; + } + if pattern.starts_with("*.") { + return name.ends_with(&pattern[1..]); + } + name == pattern +} diff --git a/src/agent/tools/filesystem/mod.rs b/src/agent/tools/filesystem/mod.rs new file mode 100644 index 0000000..a17c0eb --- /dev/null +++ b/src/agent/tools/filesystem/mod.rs @@ -0,0 +1,23 @@ +// src/agent/tools/filesystem/mod.rs +// +// 文件系统访问工具集 — 每个工具独立子模块，mod.rs 仅做 re-export。 +// +// 安全约束由 security 子模块统一提供： +// 1. 路径沙箱：只允许访问 library_dir、skills_dir 及项目根目录 +// 2. 路径穿越防护：拒绝含 ".." 的路径 +// 3. Bash 超时 + 输出截断 + +mod bash; +mod edit; +mod glob; +mod grep; +mod read; +pub mod security; +mod write; + +pub use bash::RunBashTool; +pub use edit::FileEditTool; +pub use glob::GlobFilesTool; +pub use grep::GrepFilesTool; +pub use read::ReadFileTool; +pub use write::FileWriteTool; diff --git a/src/agent/tools/filesystem/read.rs b/src/agent/tools/filesystem/read.rs new file mode 100644 index 0000000..dfb8762 --- /dev/null +++ b/src/agent/tools/filesystem/read.rs @@ -0,0 +1,165 @@ +// src/agent/tools/filesystem/read.rs +// +// read_file 工具 — 读取文件内容。 +// P1 增强：文件状态缓存去重（参考 Claude Code FileReadTool + FileStateCache）。 +// 读取前检查缓存和 mtime，相同则返回 FILE_UNCHANGED_STUB。 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::runtime::file_cache::{self, FileState}; +use crate::agent::tools::filesystem::security::{ + has_path_traversal, is_path_allowed, resolve_path, +}; +use crate::agent::tools::{AgentTool, ToolContext, ToolOutput}; + +/// 读取文件内容工具。 +pub struct ReadFileTool; + +#[async_trait] +impl AgentTool for ReadFileTool { + fn name(&self) -> &str { + "read_file" + } + + fn description(&self) -> &str { + "读取指定路径的文件内容。支持文本文件和代码文件。会自动截断过长的内容。路径必须在允许的沙箱范围内。如果文件自上次读取后未修改，会返回占位消息以节省上下文。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "file_path": { + "type": "string", + "description": "要读取的文件路径（绝对路径或相对于当前工作目录的路径）" + }, + "max_lines": { + "type": "integer", + "description": "最大读取行数，默认全部", + "default": null + } + }, + "required": ["file_path"] + }) + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let file_path_str = match args.get("file_path").and_then(|v| v.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'file_path'"), + }; + + if has_path_traversal(file_path_str) { + return ToolOutput::error("路径包含不安全字符 (.. 或 ~)"); + } + + let path = match resolve_path(file_path_str) { + Some(p) => p, + None => return ToolOutput::error(format!("无法解析路径: {}", file_path_str)), + }; + + if !is_path_allowed(&path, ctx) { + return ToolOutput::error(format!("路径不在允许的沙箱范围内: {}", file_path_str)); + } + + let max_lines = args + .get("max_lines") + .and_then(|v| v.as_i64()) + .unwrap_or(i64::MAX) as usize; + + let offset = 1usize; // 始终从第一行开始（与 Claude Code 不同，我们暂不暴露 offset 参数） + let limit = if max_lines == usize::MAX { + None + } else { + Some(max_lines) + }; + + // ── 文件状态缓存去重 ── + // 检查缓存中是否有此文件的记录，且 mtime 未变更。 + if let Ok(mut cache) = ctx.read_file_state.lock() { + let display_path = path.to_string_lossy().to_string(); + + if let Some(cached) = cache.get(&display_path) { + // 检查 offset/limit 是否匹配 + let range_match = cached.offset == offset && cached.limit == limit; + + if range_match { + // 获取磁盘上的当前 mtime 与缓存的 timestamp 比较 + let disk_mtime = file_cache::get_file_mtime(&display_path); + if let Some(mtime) = disk_mtime { + if mtime == cached.timestamp { + info!( + "[ReadFile] 文件未修改，返回 stub: {} (mtime={})", + file_path_str, mtime + ); + return ToolOutput::success( + file_cache::FILE_UNCHANGED_STUB.to_string(), + json!({ + "file_path": file_path_str, + "dedup": true, + "mtime": mtime, + }), + ); + } + } + } + } + } + // ── 去重检查结束 ── + + match std::fs::read_to_string(&path) { + Ok(content) => { + let total_lines = content.lines().count(); + let truncated = if max_lines < total_lines { + content + .lines() + .take(max_lines) + .collect::>() + .join("\n") + + &format!( + "\n\n[... 已截断，共 {} 行，显示前 {} 行 ...]", + total_lines, max_lines + ) + } else { + content.clone() + }; + + // 截断到单次输出上限 + let truncated_content = crate::agent::tools::truncate_content(&truncated, 4000); + let content_len = truncated_content.len(); + info!( + "[ReadFile] 读取 {}: {} 字符 ({} 行)", + file_path_str, content_len, total_lines + ); + + // ── 更新文件状态缓存 ── + if let Ok(mut cache) = ctx.read_file_state.lock() { + let display_path = path.to_string_lossy().to_string(); + let mtime = file_cache::get_file_mtime(&display_path).unwrap_or(0); + cache.set( + &display_path, + FileState { + content: content.clone(), + timestamp: mtime, + offset, + limit, + }, + ); + } + // ── 缓存更新结束 ── + + ToolOutput::success( + truncated_content, + json!({"file_path": file_path_str, "total_lines": total_lines}), + ) + } + Err(e) => ToolOutput::error(format!("读取文件失败: {}", e)), + } + } +} diff --git a/src/agent/tools/filesystem/security.rs b/src/agent/tools/filesystem/security.rs new file mode 100644 index 0000000..0df173c --- /dev/null +++ b/src/agent/tools/filesystem/security.rs @@ -0,0 +1,48 @@ +// src/agent/tools/filesystem/security.rs +// +// 文件系统操作的路径安全检查。 + +use std::path::{Path, PathBuf}; + +use crate::agent::tools::ToolContext; + +/// 检查路径是否在允许的沙箱范围内。 +pub fn is_path_allowed(path: &Path, ctx: &ToolContext) -> bool { + let config = &ctx.app_state.config; + let canonical = match path.canonicalize() { + Ok(p) => p, + Err(_) => match path.parent() { + Some(parent) => match parent.canonicalize() { + Ok(p) => p, + Err(_) => return false, + }, + None => return false, + }, + }; + let allowed_roots = [ + config.library_dir.canonicalize().ok(), + config.skills_dir.canonicalize().ok(), + std::env::current_dir().ok(), + ]; + for root in allowed_roots.iter().flatten() { + if canonical.starts_with(root) { + return true; + } + } + false +} + +/// 检查路径字符串是否包含穿越尝试 +pub fn has_path_traversal(path_str: &str) -> bool { + path_str.contains("..") || path_str.contains('~') +} + +/// 规范化用户提供的路径 +pub fn resolve_path(path_str: &str) -> Option { + let path = Path::new(path_str); + if path.is_absolute() { + Some(path.to_path_buf()) + } else { + std::env::current_dir().ok().map(|cwd| cwd.join(path)) + } +} diff --git a/src/agent/tools/filesystem/write.rs b/src/agent/tools/filesystem/write.rs new file mode 100644 index 0000000..936be7f --- /dev/null +++ b/src/agent/tools/filesystem/write.rs @@ -0,0 +1,94 @@ +// src/agent/tools/filesystem/write.rs +// +// file_write 工具 — 将内容写入文件。 + +use async_trait::async_trait; +use serde_json::json; +use tracing::info; + +use crate::agent::tools::filesystem::security::{ + has_path_traversal, is_path_allowed, resolve_path, +}; +use crate::agent::tools::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; + +/// 写入文件内容工具。 +pub struct FileWriteTool; + +#[async_trait] +impl AgentTool for FileWriteTool { + fn name(&self) -> &str { + "file_write" + } + + fn description(&self) -> &str { + "将内容写入指定路径的文件。如果文件已存在则覆盖。路径必须在允许的沙箱范围内。适用于保存研究结果、生成报告等。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "file_path": { + "type": "string", + "description": "要写入的文件路径" + }, + "content": { + "type": "string", + "description": "要写入的文件内容" + } + }, + "required": ["file_path", "content"] + }) + } + + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let file_path_str = match args.get("file_path").and_then(|v| v.as_str()) { + Some(s) => s, + None => return ToolOutput::error("缺少必需参数 'file_path'"), + }; + let content = match args.get("content").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'content'"), + }; + + if has_path_traversal(file_path_str) { + return ToolOutput::error("路径包含不安全字符 (.. 或 ~)"); + } + + let path = match resolve_path(file_path_str) { + Some(p) => p, + None => return ToolOutput::error(format!("无法解析路径: {}", file_path_str)), + }; + + if !is_path_allowed(&path, ctx) { + return ToolOutput::error(format!("路径不在允许的沙箱范围内: {}", file_path_str)); + } + + if let Some(parent) = path.parent() { + if let Err(e) = std::fs::create_dir_all(parent) { + return ToolOutput::error(format!("创建父目录失败: {}", e)); + } + } + + match std::fs::write(&path, &content) { + Ok(_) => { + let line_count = content.lines().count(); + info!("[FileWrite] 写入 {}: {} 行", file_path_str, line_count); + ToolOutput::success( + format!( + "成功写入文件: {} ({} 行, {} 字符)", + file_path_str, + line_count, + content.len() + ), + json!({"file_path": file_path_str, "lines": line_count, "bytes": content.len()}), + ) + } + Err(e) => ToolOutput::error(format!("写入文件失败: {}", e)), + } + } +} diff --git a/src/agent/tools/memory.rs b/src/agent/tools/memory.rs new file mode 100644 index 0000000..0d9ad45 --- /dev/null +++ b/src/agent/tools/memory.rs @@ -0,0 +1,185 @@ +// src/agent/tools/memory.rs +// +// save_memory 工具 — 让 Agent 可以将重要信息持久化到项目记忆系统。 +// 参考 Claude Code memdir 设计。 + +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tokio::sync::Mutex; +use tracing::info; + +use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; +use crate::agent::memory::dedup; +use crate::agent::memory::types::MemoryType; +use crate::agent::memory::MemoryManager; + +pub struct SaveMemoryTool { + memory_manager: Arc>, +} + +impl SaveMemoryTool { + pub fn new(memory_manager: Arc>) -> Self { + SaveMemoryTool { memory_manager } + } +} + +#[async_trait] +impl AgentTool for SaveMemoryTool { + fn name(&self) -> &str { + "save_memory" + } + + fn description(&self) -> &str { + "将重要信息保存到项目记忆系统。记忆会跨会话持久化，在后续会话中自动加载。\ + 用于保存：用户偏好、研究方法论、项目进展、重要发现、外部参考。\n\n\ + 不应保存：代码模式、架构详情、Git 历史、调试方案、已记录在 CLAUDE.md 中的内容、\ + 临时任务状态。即使用户要求保存以上内容，请先询问哪些部分是非预期的。\n\n\ + 系统会自动检测内容重复和低质量输入。保存前先检查是否有可更新的现有记忆 — 不要写重复项。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "slug": { + "type": "string", + "description": "记忆标识符（短横线命名，如 'user-prefs'）" + }, + "name": { + "type": "string", + "description": "记忆标题" + }, + "description": { + "type": "string", + "description": "简短描述（用于决定何时加载此记忆）" + }, + "memory_type": { + "type": "string", + "enum": ["user", "feedback", "project", "reference"], + "description": "记忆类型：user=用户偏好, feedback=用户反馈, project=项目进展, reference=外部参考" + }, + "content": { + "type": "string", + "description": "记忆内容（Markdown 格式）" + } + }, + "required": ["slug", "name", "description", "memory_type", "content"] + }) + } + + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + false // 写入操作，不并发安全 + } + + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block // 写入操作不可中断 + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let slug = args["slug"].as_str().unwrap_or(""); + let name = args["name"].as_str().unwrap_or(""); + let description = args["description"].as_str().unwrap_or(""); + let memory_type_str = args["memory_type"].as_str().unwrap_or("user"); + let content = args["content"].as_str().unwrap_or(""); + + if slug.is_empty() || name.is_empty() || content.is_empty() { + return ToolOutput::error("slug, name, content 均为必填项"); + } + + // Validate slug format (kebab-case) + if slug.contains(' ') || slug.contains('/') || slug.contains('\\') { + return ToolOutput::error("slug 不能包含空格、斜杠或反斜杠"); + } + + let memory_type = match MemoryType::from_str(memory_type_str) { + Some(t) => t, + None => { + return ToolOutput::error(format!( + "无效的 memory_type: {}。有效值: user, feedback, project, reference", + memory_type_str + )); + } + }; + + let mut mgr = self.memory_manager.lock().await; + + // ── 写入时门控 ── + + // 1. 内容质量检查（仅警告，不拒绝） + let quality = dedup::check_content_quality(content); + let quality_warning = match &quality { + dedup::QualityCheck::TooShort(n) => { + Some(format!("⚠️ 内容偏短 ({} 字符)，建议展开说明", n)) + } + dedup::QualityCheck::TransientState => { + Some("⚠️ 检测到瞬时状态描述，建议仅保存长期有价值的信息".to_string()) + } + dedup::QualityCheck::VagueLanguage(w) => { + Some(format!("⚠️ 检测到模糊语言 '{}'，建议使用明确表述", w)) + } + dedup::QualityCheck::CodePattern => { + Some("⚠️ 检测到代码片段 — 代码模式不应保存为记忆".to_string()) + } + dedup::QualityCheck::Accept => None, + }; + + // 2. Jaccard 内容重复检测 + let duplicate_slug = dedup::find_duplicate_by_content(content, mgr.entries(), 0.70); + + // 检查 slug 是否已存在 + let is_update = dedup::slug_exists(mgr.memory_dir(), slug); + + match mgr.save_memory(slug, name, description, memory_type, content) { + Ok(_) => { + info!("[SaveMemory] 已保存记忆: {} ({})", name, slug); + // 标记主代理已写入，抑制本会话的自动提取 + mgr.mark_main_agent_wrote(); + // 构建现有记忆清单供 LLM 参考 + let manifest = dedup::build_manifest_preview(mgr.entries()); + let action = if is_update { + "🔄 已更新" + } else { + "✅ 已保存" + }; + let mut message = format!( + "{} 记忆: {} ({}) — 类型: {}", + action, name, slug, memory_type_str + ); + + // 附加质量警告 + let has_quality_warning = quality_warning.is_some(); + if let Some(w) = &quality_warning { + message.push_str(&format!("\n\n{}", w)); + } + + // 附加重复检测信息 + if let Some(ref dup_slug) = duplicate_slug { + message.push_str(&format!( + "\n\n💡 检测到与现有记忆 `{}` 内容接近（≥70% 重叠），请考虑更新该文件而非创建新的。", + dup_slug + )); + } + + message.push_str(&format!("\n\n{}", manifest)); + + if !is_update && duplicate_slug.is_none() { + message + .push_str("\n\n💡 提示：如有其他重要信息需保存，请继续使用 save_memory。"); + } + ToolOutput::success( + message, + json!({ + "slug": slug, + "name": name, + "memory_type": memory_type_str, + "is_update": is_update, + "quality_check": has_quality_warning, + "duplicate_detected": duplicate_slug.is_some() + }), + ) + } + Err(e) => ToolOutput::error(format!("保存记忆失败: {}", e)), + } + } +} diff --git a/src/agent/tools/mod.rs b/src/agent/tools/mod.rs new file mode 100644 index 0000000..f4c7471 --- /dev/null +++ b/src/agent/tools/mod.rs @@ -0,0 +1,417 @@ +// src/agent/tools/mod.rs +// +// 科研智能体工具集定义与实现。 +// 每个工具遵循 AgentTool trait，向大模型声明 JSON Schema 参数定义， +// 并在 execute 中调用已有的服务层完成实际业务操作。 +// +// 按功能域拆分为子模块： +// filesystem/ — 文件 I/O（read、grep、glob、bash、write、edit） +// astro/ — 天文科研（文献搜索/下载/解析、RAG、天体查询、笔记） +// team.rs — 团队协作 +// todo.rs — 任务规划 +// compress.rs — 手动上下文压缩 + +use async_trait::async_trait; +use serde_json::json; +use std::sync::{Arc, RwLock}; + +use crate::agent::runtime::file_cache::FileStateCache; +use crate::agent::skills::SkillRegistry; +use crate::api::AppState; +use crate::clients::llm::ToolDefinition; + +pub mod ask_user; +pub mod astro; +mod background; +mod compress; +mod filesystem; +pub mod memory; +pub mod persist; +mod skill; +pub mod subagent; +mod team; +mod todo; + +pub use ask_user::AskUserTool; +pub use astro::note::SaveNoteTool; +pub use astro::paper::{DownloadPaperTool, GetPaperContentTool, ParsePaperTool}; +pub use astro::rag::RagSearchTool; +pub use astro::search::{GetPaperMetadataTool, SearchPapersTool}; +pub use astro::target::QueryTargetTool; +pub use background::{BgTaskCheckTool, BgTaskRunTool}; +pub use compress::CompressTool; +pub use filesystem::{ + FileEditTool, FileWriteTool, GlobFilesTool, GrepFilesTool, ReadFileTool, RunBashTool, +}; +pub use skill::LoadSkillTool; +pub use subagent::DelegateResearchTool; +pub use team::{CheckTeamInboxTool, SendTeammateMessageTool, SpawnTeammateTool, TeamBroadcastTool}; +pub use todo::persist_tasks; +pub use todo::TodoWriteTool; + +/// 工具执行上下文，封装全局共享状态 +pub struct ToolContext { + pub app_state: Arc, + /// 静默模式：子代理运行时为 true，跳过用户权限提示 + pub silent: bool, + /// 文件状态缓存（跨工具调用共享，用于 Read 去重） + pub read_file_state: Arc>, +} + +impl ToolContext { + /// 创建标准上下文 + pub fn new(app_state: Arc) -> Self { + ToolContext { + app_state, + silent: false, + read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())), + } + } + + /// 创建带共享文件缓存的上下文（用于 AgentRuntime 保持同一 cache 实例） + pub fn with_file_cache( + app_state: Arc, + read_file_state: Arc>, + ) -> Self { + ToolContext { + app_state, + silent: false, + read_file_state, + } + } + + /// 创建静默上下文（子代理使用） + pub fn silent(app_state: Arc) -> Self { + ToolContext { + app_state, + silent: true, + read_file_state: Arc::new(std::sync::Mutex::new(FileStateCache::new())), + } + } +} + +/// 工具执行结果 +#[derive(Debug, Clone)] +pub struct ToolOutput { + /// 给大模型阅读的截断文本 + pub content: String, + /// 是否为错误 + pub is_error: bool, + /// 结构化元数据（给前端 Timeline 直接渲染） + pub metadata: serde_json::Value, +} + +impl ToolOutput { + /// 创建成功结果 + pub fn success(content: impl Into, metadata: serde_json::Value) -> Self { + ToolOutput { + content: content.into(), + is_error: false, + metadata, + } + } + + /// 创建错误结果 + pub fn error(msg: impl Into) -> Self { + ToolOutput { + content: msg.into(), + is_error: true, + metadata: json!({}), + } + } +} + +/// 工具被中断时的行为策略 +#[derive(Debug, Clone, PartialEq)] +pub enum InterruptBehavior { + /// 取消执行并返回错误（默认，适用于只读工具） + Cancel, + /// 阻塞中断信号直到执行完成（适用于有副作用的写入工具） + Block, +} + +/// 权限规则 — 工具自定义的权限限制 +#[derive(Debug, Clone)] +pub enum PermissionRule { + /// 不可覆盖的拒绝 + Deny { tool_name: String, reason: String }, + /// 允许 + Allow { tool_name: String }, + /// 需要用户确认 + Ask { tool_name: String, message: String }, +} + +/// 智能体工具 trait +#[async_trait] +pub trait AgentTool: Send + Sync { + /// 工具名称（与 LLM function calling 的 name 保持一致） + fn name(&self) -> &str; + /// 工具描述（告知 LLM 何时应该调用该工具） + fn description(&self) -> &str; + /// JSON Schema 格式的参数定义 + fn parameters(&self) -> serde_json::Value; + /// 执行工具逻辑 + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput; + + // ── P1 优化：新增默认方法 ── + + /// 中断行为策略。默认 Cancel — 可以安全中断。 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Cancel + } + + /// 该工具是否支持并发安全执行。 + /// 默认 false（保守策略），只读工具应覆写为 true。 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + false + } + + /// 工具自定义权限检查。默认无额外限制。 + fn check_permissions(&self, _args: &serde_json::Value) -> Vec { + Vec::new() + } + + /// 该工具错误时是否应中止兄弟并行执行。 + /// 默认 false（只读工具不触发）。下载/解析类工具可覆写为 true。 + fn causes_sibling_abort(&self) -> bool { + false + } + + /// 带进度流式执行。默认委托给 execute()。 + /// 长时间操作的工具可覆写以发送进度更新。 + async fn execute_with_progress( + &self, + args: serde_json::Value, + ctx: &ToolContext, + _progress_tx: Option<&tokio::sync::mpsc::UnboundedSender>, + ) -> ToolOutput { + self.execute(args, ctx).await + } +} + +/// 工具注册表，管理所有可用工具。 +/// 内部使用 HashMap 实现 O(1) 按名查找，同时保留插入顺序供 definitions() 使用。 +pub struct ToolRegistry { + tools: std::collections::HashMap>, + ordered_names: Vec, +} + +// ── 工具注册辅助函数（消除重复代码） ── + +/// 注册所有基础研究工具（文件、文献、RAG、笔记等 19 个工具）。 +fn add_base_tools(registry: &mut ToolRegistry, skill_registry: Arc>) { + let tools: Vec> = vec![ + Box::new(ReadFileTool), + Box::new(GrepFilesTool), + Box::new(GlobFilesTool), + Box::new(RunBashTool), + Box::new(FileWriteTool), + Box::new(FileEditTool), + Box::new(SearchPapersTool), + Box::new(GetPaperMetadataTool), + Box::new(DownloadPaperTool), + Box::new(ParsePaperTool), + Box::new(GetPaperContentTool), + Box::new(RagSearchTool), + Box::new(QueryTargetTool), + Box::new(SaveNoteTool), + Box::new(TodoWriteTool), + Box::new(CompressTool), + Box::new(AskUserTool), + Box::new(LoadSkillTool::new(skill_registry)), + Box::new(DelegateResearchTool::new()), + ]; + for tool in tools { + registry.ordered_names.push(tool.name().to_string()); + registry.tools.insert(tool.name().to_string(), tool); + } +} + +/// 注册后台任务工具（bg_task_run, bg_task_check）。 +fn add_background_tools( + registry: &mut ToolRegistry, + queue: Arc, +) { + let run_tool = Box::new(BgTaskRunTool::new(queue.clone())); + let check_tool = Box::new(BgTaskCheckTool::new(queue)); + registry.ordered_names.push(run_tool.name().to_string()); + registry.tools.insert(run_tool.name().to_string(), run_tool); + registry.ordered_names.push(check_tool.name().to_string()); + registry + .tools + .insert(check_tool.name().to_string(), check_tool); +} + +/// 注册团队协作工具（spawn_teammate, send_teammate_message, team_broadcast, check_team_inbox）。 +fn add_team_tools( + registry: &mut ToolRegistry, + team_manager: Arc>>, +) { + let team_tools: Vec> = vec![ + Box::new(SpawnTeammateTool::new(team_manager.clone())), + Box::new(SendTeammateMessageTool::new(team_manager.clone())), + Box::new(TeamBroadcastTool::new(team_manager.clone())), + Box::new(CheckTeamInboxTool::new(team_manager)), + ]; + for tool in team_tools { + registry.ordered_names.push(tool.name().to_string()); + registry.tools.insert(tool.name().to_string(), tool); + } +} + +impl ToolRegistry { + /// 创建空工具注册表（调用者通过 add_tool 手动添加工具）。 + /// 用于受限场景（如记忆提取子代理只需要只读 + save_memory）。 + pub fn empty() -> Self { + ToolRegistry { + tools: std::collections::HashMap::new(), + ordered_names: Vec::new(), + } + } + + /// 创建默认工具注册表（包含全部科研工具，不含后台工具） + pub fn new(skill_registry: Arc>) -> Self { + Self::new_with_queue(None, skill_registry) + } + + /// 创建工具注册表，可选注入后台通知队列以启用 bg_task_run/bg_task_check + pub fn new_with_queue( + queue: Option>, + skill_registry: Arc>, + ) -> Self { + let mut registry = ToolRegistry { + tools: std::collections::HashMap::new(), + ordered_names: Vec::new(), + }; + add_base_tools(&mut registry, skill_registry); + if let Some(q) = queue { + add_background_tools(&mut registry, q); + } + registry + } + + /// 创建包含团队工具的注册表 + pub fn new_with_team( + queue: Option>, + team_manager: Arc>>, + skill_registry: Arc>, + ) -> Self { + // 使用 new_with_queue 获取基础 + 后台工具，再添加团队工具 + let mut registry = Self::new_with_queue(queue, skill_registry); + add_team_tools(&mut registry, team_manager); + registry + } + + /// 动态添加工具（用于需要共享状态的工具，如 MemoryManager） + pub fn add_tool(&mut self, tool: Box) { + let name = tool.name().to_string(); + self.ordered_names.push(name.clone()); + self.tools.insert(name, tool); + } + + /// 替换已存在的工具（保持名称在 ordered_names 中的位置不变）。 + /// 如果工具不存在，行为等同于 add_tool。 + pub fn replace_tool(&mut self, tool: Box) { + let name = tool.name().to_string(); + if !self.tools.contains_key(&name) { + self.ordered_names.push(name.clone()); + } + self.tools.insert(name, tool); + } + + /// 根据名称查找工具 (O(1)) + pub fn get(&self, name: &str) -> Option<&dyn AgentTool> { + self.tools.get(name).map(|t| t.as_ref()) + } + + /// 生成所有工具的 ToolDefinition 列表（用于发送给 LLM）。 + /// 按名称字母序排序以保证跨调用的稳定性，提升 prompt cache 命中率。 + pub fn definitions(&self) -> Vec { + let mut defs: Vec<_> = self + .tools + .values() + .map(|t| ToolDefinition::new(t.name(), t.description(), t.parameters())) + .collect(); + defs.sort_by(|a, b| a.function.name.cmp(&b.function.name)); + defs + } +} + +/// 截断文本到指定最大字符数 +pub fn truncate_content(s: &str, max_chars: usize) -> String { + if s.len() <= max_chars { + s.to_string() + } else { + let truncated: String = s.chars().take(max_chars).collect(); + format!("{}\n\n[... 内容已截断，共 {} 字符 ...]", truncated, s.len()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + #[test] + fn test_truncate_content_short() { + let text = "Hello, world!"; + assert_eq!(truncate_content(text, 100), text); + } + + #[test] + fn test_truncate_content_long() { + let text = "a".repeat(5000); + let result = truncate_content(&text, 100); + assert!(result.contains("内容已截断")); + assert!(result.contains("5000")); + } + + #[test] + fn test_tool_output_success() { + let output = ToolOutput::success("ok", json!({"key": "value"})); + assert!(!output.is_error); + assert_eq!(output.content, "ok"); + } + + #[test] + fn test_tool_output_error() { + let output = ToolOutput::error("something went wrong"); + assert!(output.is_error); + } + + #[test] + fn test_tool_registry_definitions() { + let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from( + "./skills", + ))))); + let defs = registry.definitions(); + assert_eq!(defs.len(), 19); + assert!(defs.iter().any(|d| d.function.name == "read_file")); + assert!(defs.iter().any(|d| d.function.name == "grep_files")); + assert!(defs.iter().any(|d| d.function.name == "glob_files")); + assert!(defs.iter().any(|d| d.function.name == "run_bash")); + assert!(defs.iter().any(|d| d.function.name == "file_write")); + assert!(defs.iter().any(|d| d.function.name == "file_edit")); + assert!(defs.iter().any(|d| d.function.name == "search_papers")); + assert!(defs.iter().any(|d| d.function.name == "get_paper_metadata")); + assert!(defs.iter().any(|d| d.function.name == "download_paper")); + assert!(defs.iter().any(|d| d.function.name == "parse_paper")); + assert!(defs.iter().any(|d| d.function.name == "get_paper_content")); + assert!(defs.iter().any(|d| d.function.name == "rag_search")); + assert!(defs.iter().any(|d| d.function.name == "query_target")); + assert!(defs.iter().any(|d| d.function.name == "save_note")); + assert!(defs.iter().any(|d| d.function.name == "todo_write")); + assert!(defs.iter().any(|d| d.function.name == "compress_context")); + assert!(defs.iter().any(|d| d.function.name == "load_skill")); + assert!(defs.iter().any(|d| d.function.name == "delegate_research")); + } + + #[test] + fn test_tool_registry_get() { + let registry = ToolRegistry::new(Arc::new(RwLock::new(SkillRegistry::new(PathBuf::from( + "./skills", + ))))); + assert!(registry.get("search_papers").is_some()); + assert!(registry.get("nonexistent").is_none()); + } +} diff --git a/src/agent/tools/persist.rs b/src/agent/tools/persist.rs new file mode 100644 index 0000000..c219ecc --- /dev/null +++ b/src/agent/tools/persist.rs @@ -0,0 +1,181 @@ +// src/agent/tools/persist.rs +// +// 工具结果持久化到磁盘。 +// 参考 Claude Code toolResultStorage.ts 设计。 +// +// 当工具输出超过配置的字符限制时，将完整内容写入磁盘文件， +// 返回一个 占位符给模型，模型可通过 read_file 工具读取完整内容。 +// 使用独占创建（create_new）保证幂等——同一 tool_call_id 不会被重复写入。 +// +// 目录结构： +// {library_dir}/tool-results/{tool_call_id}.txt + +use std::path::{Path, PathBuf}; +use tracing::info; + +/// 将工具输出持久化到磁盘（如果超过大小限制）。 +/// +/// 返回 (最终内容, 持久化文件路径)。 +/// 如果内容未超过限制，直接返回原内容且持久化路径为 None。 +pub fn maybe_persist_tool_result( + content: &str, + tool_call_id: &str, + max_chars: usize, + tool_results_dir: &Path, +) -> (String, Option) { + if content.len() <= max_chars { + return (content.to_string(), None); + } + + // 创建目录（幂等） + if let Err(_e) = std::fs::create_dir_all(tool_results_dir) { + // 无法创建目录则直接截断（不持久化） + let truncated: String = content.chars().take(max_chars).collect(); + return ( + format!( + "{}...\n[输出已截断，原始长度: {} 字符]", + truncated, + content.len() + ), + None, + ); + } + + let file_path = tool_results_dir.join(format!("{}.txt", tool_call_id)); + + // 独占创建——只在文件不存在时写入，保证幂等性 + let written = match std::fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&file_path) + { + Ok(_) => { + // 文件创建成功，写入内容 + match std::fs::write(&file_path, content) { + Ok(_) => { + info!( + "[Persist] 工具结果已持久化: {} ({} 字符)", + file_path.display(), + content.len() + ); + true + } + Err(e) => { + tracing::warn!("[Persist] 写入失败: {}", e); + false + } + } + } + Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => { + // 文件已存在，无需重复写入 + info!("[Persist] 工具结果已存在，跳过: {}", file_path.display()); + true + } + Err(e) => { + tracing::warn!("[Persist] 创建文件失败: {}", e); + false + } + }; + + if !written { + // 持久化失败，回退到截断 + let truncated: String = content.chars().take(max_chars).collect(); + return ( + format!( + "{}...\n[输出已截断，原始长度: {} 字符]", + truncated, + content.len() + ), + None, + ); + } + + // 预览（在第一个换行处截断，避免 mid-line cut） + let preview_limit = 500.min(max_chars); + let preview: String = + if let Some(newline_pos) = content[..preview_limit.min(content.len())].rfind('\n') { + content[..newline_pos].to_string() + } else { + content.chars().take(preview_limit).collect() + }; + + let stub = format!( + "\n\ + path: {}\n\ + size: {} chars\n\ + preview: |\n {}\n\n\ + 完整输出已持久化到磁盘。使用 read_file 工具以 path 参数读取完整内容。\n\ + ", + file_path.display(), + content.len(), + preview.replace('\n', "\n "), + ); + + (stub, Some(file_path)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn temp_dir() -> (PathBuf, impl Drop) { + let dir = std::env::temp_dir().join(format!("astro_test_{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let dir_clone = dir.clone(); + struct Cleanup(PathBuf); + impl Drop for Cleanup { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + (dir, Cleanup(dir_clone)) + } + + #[test] + fn test_small_content_not_persisted() { + let (dir, _cleanup) = temp_dir(); + let (content, path) = maybe_persist_tool_result("small result", "call_1", 4000, &dir); + assert_eq!(content, "small result"); + assert!(path.is_none()); + } + + #[test] + fn test_large_content_persisted() { + let (dir, _cleanup) = temp_dir(); + let large = "x".repeat(5000); + let (content, path) = maybe_persist_tool_result(&large, "call_2", 100, &dir); + assert!(content.contains("")); + assert!(content.contains("call_2.txt")); + assert!(path.is_some()); + let file_path = path.unwrap(); + assert!(file_path.exists()); + let written = std::fs::read_to_string(&file_path).unwrap(); + assert_eq!(written, large); + } + + #[test] + fn test_idempotent_write() { + let (dir, _cleanup) = temp_dir(); + let large1 = "a".repeat(5000); + let large2 = "b".repeat(5000); + + let (_, path1) = maybe_persist_tool_result(&large1, "call_3", 100, &dir); + let (content2, path2) = maybe_persist_tool_result(&large2, "call_3", 100, &dir); + + assert!(path1.is_some()); + assert!(path2.is_some()); + let written = std::fs::read_to_string(path1.unwrap()).unwrap(); + assert_eq!(written, large1); + } + + #[test] + fn test_preview_at_newline_boundary() { + let (dir, _cleanup) = temp_dir(); + let large = format!("Short line\n{}", "x".repeat(5000)); + let (content, _) = maybe_persist_tool_result(&large, "call_4", 100, &dir); + let preview_start = content.find("preview:").unwrap(); + let preview_section = &content[preview_start..]; + assert!(preview_section.contains("Short line")); + assert!(!preview_section.contains("xxx")); + } +} diff --git a/src/agent/tools/skill.rs b/src/agent/tools/skill.rs new file mode 100644 index 0000000..2040259 --- /dev/null +++ b/src/agent/tools/skill.rs @@ -0,0 +1,222 @@ +// src/agent/tools/skill.rs — 技能加载工具（Layer 2 按需加载） +// +// 参考 Claude Code src/tools/SkillTool/SkillTool.ts 设计： +// - description() 动态生成，列出所有可用 skill 及描述 +// - 支持 inline 模式（直接返回 skill 内容）和 fork 模式（子代理执行） +// - 支持 allowed-tools 白名单返回 +// - 支持变量替换（${SKILL_DIR}, ${SESSION_ID}） + +use async_trait::async_trait; +use serde_json::json; +use std::sync::{Arc, RwLock}; +use tracing::{info, warn}; + +use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; +use crate::agent::skills::{substitute_variables, SkillRegistry}; + +/// 技能加载工具 +/// +/// 持有 SkillRegistry 引用以实现： +/// - 缓存读取（避免重复磁盘 I/O） +/// - 动态 description 生成 +/// - 使用统计 +/// - fork 模式下的子代理执行 +pub struct LoadSkillTool { + registry: Arc>, +} + +impl LoadSkillTool { + pub fn new(registry: Arc>) -> Self { + LoadSkillTool { registry } + } +} + +#[async_trait] +impl AgentTool for LoadSkillTool { + fn name(&self) -> &str { + "load_skill" + } + + fn description(&self) -> &str { + // 注意：description() 返回 &str，但我们需要动态内容。 + // 实际使用 ToolRegistry 时，此方法的返回值作为 base description， + // 详细的 skill 列表通过 system-reminder 注入。 + // 如果 trait 允许，应改为返回 String。当前保持与 trait 兼容。 + "加载指定领域技能的完整内容。可用技能列表已在系统提示词中列出。当需要某个技能的详细指引时调用此工具。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "skill_name": { + "type": "string", + "description": "要加载的技能名称。可用技能列表见系统提示词。" + } + }, + "required": ["skill_name"] + }) + } + + /// 纯读文件，并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + /// fork 模式下的 skill 不应被中断 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Cancel + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let skill_name = match args.get("skill_name").and_then(|s| s.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'skill_name'"), + }; + + info!("[LoadSkill] 加载 skill: {}", skill_name); + + // 从缓存注册表读取（单语句：guard 自动 drop，不跨越 await） + let skill = match self.registry.read() { + Ok(reg) => reg.get_skill(&skill_name).cloned(), + Err(e) => { + warn!("[LoadSkill] RwLock poisoned: {:?}", e); + return ToolOutput::error("技能注册表不可用（内部锁异常），请稍后重试"); + } + }; + let skill = match skill { + Some(s) => s, + None => { + return ToolOutput::error(format!( + "技能 '{}' 未找到。可用技能见系统提示词中的列表。", + skill_name + )); + } + }; + + // 记录使用统计（单语句：write guard 自动 drop） + if let Ok(mut reg) = self.registry.write() { + reg.record_usage(&skill_name); + } else { + warn!("[LoadSkill] 无法记录使用统计（RwLock poisoned）"); + } + + // 变量替换 + let session_id = if ctx + .app_state + .config + .database_url + .contains("session") { "current" } else { "" }; + + let body = substitute_variables( + &skill.body, + &skill.skill_dir, + if session_id.is_empty() { + None + } else { + Some(session_id) + }, + ); + + // 构建 skill 目录的绝对路径（用于 "Base directory" 前缀） + let skill_dir_path = skill + .skill_dir + .canonicalize() + .unwrap_or_else(|_| skill.skill_dir.clone()); + let base_dir_note = format!( + "Base directory for this skill: {}\n\n", + skill_dir_path.display() + ); + + // 检查 context 模式 + let is_fork = skill.meta.context.as_deref() == Some("fork"); + + if is_fork { + // ── Fork 模式：使用子代理执行 skill ── + info!( + "[LoadSkill] Skill '{}' 标记为 fork 模式，启动子代理执行", + skill_name + ); + + let max_steps = args + .get("max_steps") + .and_then(|v| v.as_i64()) + .unwrap_or(5) + .min(10) as usize; + + let runner = crate::agent::subagent::SubAgentRunner::new(ctx.app_state.clone()); + let system_prompt = format!( + "你是一位专业的天体物理学研究助手。\n\n\ + 你正在按照以下技能指引执行任务。\n\ + 技能目录: {}\n\n{}{}\n\n\ + 请严格遵循上述指引完成任务，使用可用工具收集和分析信息。\ + 你可以使用 Read/Grep/Bash 工具访问技能目录中的文件。\ + 完成后给出最终结果。", + skill_dir_path.display(), + base_dir_note, + body + ); + + // 子代理任务描述取自 skill description 或 body 前 200 字符 + let research_prompt = format!( + "按照 '{}' 技能的指引完成任务：{}", + skill.meta.name, skill.meta.description + ); + + let result = runner + .run(&system_prompt, &research_prompt, max_steps) + .await; + + if result.is_error { + ToolOutput::error(format!( + "技能 '{}' 子代理执行失败: {}", + skill_name, result.content + )) + } else { + ToolOutput::success( + format!( + "[子代理执行结果 - 技能: {} ({})]\n\n{}", + skill.meta.name, skill.meta.description, result.content + ), + json!({ + "skill_name": skill.meta.name, + "description": skill.meta.description, + "context": "fork", + "execution_mode": "subagent", + "body_length": skill.body.len(), + "allowed_tools": skill.meta.allowed_tools, + }), + ) + } + } else { + // ── Inline 模式：直接返回 skill 内容 ── + let allowed_tools_note = if skill.meta.allowed_tools.is_empty() { + String::new() + } else { + format!( + "\n\n> **工具白名单**: {}（此技能建议仅使用这些工具）", + skill.meta.allowed_tools.join(", ") + ) + }; + + ToolOutput::success( + format!( + "# 技能：{} ({})\n\n{}{}{}", + skill.meta.name, + skill.meta.description, + base_dir_note, + body, + allowed_tools_note + ), + json!({ + "skill_name": skill.meta.name, + "description": skill.meta.description, + "context": skill.meta.context, + "execution_mode": "inline", + "body_length": skill.body.len(), + "allowed_tools": skill.meta.allowed_tools, + }), + ) + } + } +} diff --git a/src/agent/tools/subagent.rs b/src/agent/tools/subagent.rs new file mode 100644 index 0000000..e316b43 --- /dev/null +++ b/src/agent/tools/subagent.rs @@ -0,0 +1,132 @@ +// src/agent/tools/subagent.rs — 子代理委托工具 (delegate_research) +// +// 参考 Claude Code s04 Subagents 设计。 +// LLM 通过此工具将子任务委托给上下文隔离的子代理执行。 + +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tokio::sync::mpsc::UnboundedSender; +use tracing::info; + +use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; +use crate::agent::hooks::HookRegistry; +use crate::agent::runtime::permission::PermissionChecker; +use crate::agent::runtime::AgentStreamEvent; +use crate::agent::subagent::SubAgentRunner; + +/// 子代理委托工具 +pub struct DelegateResearchTool { + hook_registry: Option>, + permission_checker: Arc, + progress_tx: Option>, +} + +impl Default for DelegateResearchTool { + fn default() -> Self { + Self::new() + } +} + +impl DelegateResearchTool { + /// 创建不带 hooks 的工具实例（向后兼容） + pub fn new() -> Self { + DelegateResearchTool { + hook_registry: None, + permission_checker: Arc::new(PermissionChecker::new()), + progress_tx: None, + } + } + + /// 创建带完整 hooks/permissions/progress 的工具实例 + pub fn new_with_hooks( + hook_registry: Option>, + permission_checker: Arc, + progress_tx: Option>, + ) -> Self { + DelegateResearchTool { + hook_registry, + permission_checker, + progress_tx, + } + } +} + +#[async_trait] +impl AgentTool for DelegateResearchTool { + fn name(&self) -> &str { + "delegate_research" + } + + fn description(&self) -> &str { + "将子研究任务委托给独立的子代理执行。子代理拥有完整工具访问权限（文献搜索、下载、RAG检索等），\ + 但只有最终文本摘要会返回给父代理，中间工具调用不会污染父上下文。\ + 适用于：文献综述、多步数据收集、独立子问题研究等可以独立完成的子任务。\ + 重要：delegate_research 返回后，你仍应基于其结果继续分析和回答用户问题。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "research_prompt": { + "type": "string", + "description": "要委托给子代理执行的完整研究任务描述。应包含具体的搜索目标、需要收集的信息、期望的输出格式。" + }, + "max_steps": { + "type": "integer", + "description": "子代理最大推理步数，默认5，最大10", + "default": 5 + } + }, + "required": ["research_prompt"] + }) + } + + /// 子代理可能执行写操作，中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, ctx: &ToolContext) -> ToolOutput { + let research_prompt = match args.get("research_prompt").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少必需参数 'research_prompt'"), + }; + + let max_steps = args + .get("max_steps") + .and_then(|v| v.as_i64()) + .unwrap_or(5) + .min(10) as usize; + + info!( + "[DelegateResearch] 启动子代理: prompt_len={}, max_steps={}", + research_prompt.len(), + max_steps + ); + + let system_prompt = "你是一位专业的天体物理学研究助手，在一个独立的子任务上下文中工作。\ + 你可以使用文献搜索、下载、RAG检索等工具。\ + 请高效完成任务，然后直接给出最终答案。不要进行不必要的重复操作。\ + 用中文回答，引用具体文献来源。"; + + let runner = SubAgentRunner::new_with_hooks( + ctx.app_state.clone(), + self.hook_registry.clone(), + self.permission_checker.clone(), + self.progress_tx.clone(), + ); + let result = runner.run(system_prompt, &research_prompt, max_steps).await; + + if result.is_error { + ToolOutput::error(format!("子代理执行失败: {}", result.content)) + } else { + // 包装子代理结果，标注来源 + ToolOutput::success( + format!("[子代理研究结果]\n\n{}", result.content), + result.metadata, + ) + } + } +} diff --git a/src/agent/tools/team.rs b/src/agent/tools/team.rs new file mode 100644 index 0000000..fda1bec --- /dev/null +++ b/src/agent/tools/team.rs @@ -0,0 +1,291 @@ +// src/agent/tools/team.rs — 团队协作工具 +// +// 4 个团队工具： +// spawn_teammate — 启动一个队友 agent +// send_teammate_message — 发送消息给指定队友 +// team_broadcast — 广播消息给所有队友 +// check_team_inbox — 检查收件箱 + +use async_trait::async_trait; +use serde_json::json; +use std::sync::Arc; +use tokio::sync::Mutex; + +use super::{AgentTool, InterruptBehavior, ToolContext, ToolOutput}; +use crate::agent::team::inbox::TeamMessageType; +use crate::agent::team::manager::TeamManager; + +/// 启动队友工具 +pub struct SpawnTeammateTool { + team_manager: Arc>>, +} + +impl SpawnTeammateTool { + pub fn new(team_manager: Arc>>) -> Self { + SpawnTeammateTool { team_manager } + } +} + +#[async_trait] +impl AgentTool for SpawnTeammateTool { + fn name(&self) -> &str { + "spawn_teammate" + } + + fn description(&self) -> &str { + "启动一个队友 agent。队友拥有独立的上下文，可以并行执行文献搜索、论文下载等任务。\ + 通过 send_teammate_message 向队友发送任务，通过 check_team_inbox 检查结果。\ + 适用于：需要多线并行的文献调研任务。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "name": { + "type": "string", + "description": "队友名称（如 searcher, reader）" + }, + "role": { + "type": "string", + "description": "队友角色描述（如 ADS文献搜索专家, 论文全文阅读专家）" + } + }, + "required": ["name", "role"] + }) + } + + /// 启动队友有副作用，中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let name = match args.get("name").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少 'name' 参数"), + }; + let role = match args.get("role").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少 'role' 参数"), + }; + + // NOTE: team_manager 锁在此 await 期间保持持有。 + // spawn 操作通常很快（只是创建子代理会话）， + // 如果需要降低锁持有时间，可以将 TeamManager 改为内部使用 Arc。 + let tm_lock = self.team_manager.lock().await; + match tm_lock.as_ref() { + Some(tm) => { + let handle = tm.spawn(&name, &role).await; + ToolOutput::success( + format!( + "✅ 队友已启动: {} ({})\n使用 send_teammate_message 发送任务。", + handle.name, handle.role + ), + json!({ "name": handle.name, "role": handle.role }), + ) + } + None => ToolOutput::error("团队管理器未初始化"), + } + } +} + +/// 发送消息工具 +pub struct SendTeammateMessageTool { + team_manager: Arc>>, +} + +impl SendTeammateMessageTool { + pub fn new(team_manager: Arc>>) -> Self { + SendTeammateMessageTool { team_manager } + } +} + +#[async_trait] +impl AgentTool for SendTeammateMessageTool { + fn name(&self) -> &str { + "send_teammate_message" + } + + fn description(&self) -> &str { + "向指定队友发送消息（任务分配、问题等）。消息会投递到队友的收件箱，\ + 队友在处理循环中自动读取。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "to": { + "type": "string", + "description": "收件队友名称" + }, + "content": { + "type": "string", + "description": "消息内容（任务描述、问题等）" + } + }, + "required": ["to", "content"] + }) + } + + /// 消息投递有副作用，中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let to = match args.get("to").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少 'to' 参数"), + }; + let content = match args.get("content").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少 'content' 参数"), + }; + + let tm_lock = self.team_manager.lock().await; + match tm_lock.as_ref() { + Some(tm) => { + tm.send_message("lead", &to, &content, TeamMessageType::Task); + ToolOutput::success( + format!( + "📤 消息已发送给 {}: {}", + to, + &content.chars().take(100).collect::() + ), + json!({ "to": to }), + ) + } + None => ToolOutput::error("团队管理器未初始化"), + } + } +} + +/// 广播消息工具 +pub struct TeamBroadcastTool { + team_manager: Arc>>, +} + +impl TeamBroadcastTool { + pub fn new(team_manager: Arc>>) -> Self { + TeamBroadcastTool { team_manager } + } +} + +#[async_trait] +impl AgentTool for TeamBroadcastTool { + fn name(&self) -> &str { + "team_broadcast" + } + + fn description(&self) -> &str { + "向所有队友广播消息。适用于状态同步、全局指令等。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "content": { + "type": "string", + "description": "广播内容" + } + }, + "required": ["content"] + }) + } + + /// 广播有副作用，中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let content = match args.get("content").and_then(|v| v.as_str()) { + Some(s) => s.to_string(), + None => return ToolOutput::error("缺少 'content' 参数"), + }; + + let tm_lock = self.team_manager.lock().await; + match tm_lock.as_ref() { + Some(tm) => { + tm.broadcast("lead", &content); + ToolOutput::success("📢 已广播消息给所有队友。", json!({})) + } + None => ToolOutput::error("团队管理器未初始化"), + } + } +} + +/// 检查收件箱工具 +pub struct CheckTeamInboxTool { + team_manager: Arc>>, +} + +impl CheckTeamInboxTool { + pub fn new(team_manager: Arc>>) -> Self { + CheckTeamInboxTool { team_manager } + } +} + +#[async_trait] +impl AgentTool for CheckTeamInboxTool { + fn name(&self) -> &str { + "check_team_inbox" + } + + fn description(&self) -> &str { + "检查收件箱，获取队友发来的消息。读取后消息会被清空。不指定 agent 时检查 lead 的收件箱。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "agent_name": { + "type": "string", + "description": "可选：要检查的 agent 名称，默认为 lead" + } + }, + "required": [] + }) + } + + /// 收件箱检查是破坏性读取（消息会被清空），中断时应阻塞以完成 + fn interrupt_behavior(&self) -> InterruptBehavior { + InterruptBehavior::Block + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let agent = args + .get("agent_name") + .and_then(|v| v.as_str()) + .unwrap_or("lead"); + + let tm_lock = self.team_manager.lock().await; + match tm_lock.as_ref() { + Some(tm) => { + let msgs = tm.check_inbox(agent); + if msgs.is_empty() { + ToolOutput::success( + format!("📭 {} 的收件箱为空。", agent), + json!({ "messages": [] }), + ) + } else { + let mut result = format!("📬 {} 的收件箱 ({} 条消息):\n\n", agent, msgs.len()); + for msg in &msgs { + result.push_str(&format!( + " [{}] 来自 {}: {}\n", + msg.msg_type.as_str(), + msg.from, + &msg.content.chars().take(200).collect::(), + )); + } + ToolOutput::success(result, json!({ "count": msgs.len() })) + } + } + None => ToolOutput::error("团队管理器未初始化"), + } + } +} diff --git a/src/agent/tools/todo.rs b/src/agent/tools/todo.rs new file mode 100644 index 0000000..8b7933f --- /dev/null +++ b/src/agent/tools/todo.rs @@ -0,0 +1,260 @@ +// src/agent/tools/todo.rs — 任务规划工具 (TodoWrite) +// +// P1 改进：任务状态持久化到 SQLite agent_tasks 表， +// 支持 blockedBy 依赖关系和跨 turn 状态恢复。 + +use async_trait::async_trait; +use serde_json::json; +use sqlx::SqlitePool; +use tracing::{info, warn}; + +use super::{AgentTool, ToolContext, ToolOutput}; + +/// 任务规划工具：让 LLM 在开始复杂研究前先制定计划，执行中更新进度。 +/// 任务状态持久化到 SQLite，支持 DAG 依赖。 +pub struct TodoWriteTool; + +#[async_trait] +impl AgentTool for TodoWriteTool { + fn name(&self) -> &str { + "todo_write" + } + + fn description(&self) -> &str { + "任务规划工具。在开始复杂研究前列出待办事项，执行中标记进度（每项状态：pending/in_progress/completed）。\ + 一次只能有一个 in_progress 任务。支持任务依赖（blockedBy：依赖的其他任务ID列表）。\ + 适用范围：任何需要多步工具调用的研究任务。" + } + + fn parameters(&self) -> serde_json::Value { + json!({ + "type": "object", + "properties": { + "todos": { + "type": "array", + "description": "任务列表，每项包含 id（唯一标识）、content（任务描述）、status（pending/in_progress/completed）、blockedBy（可选，依赖的其他任务ID列表）", + "items": { + "type": "object", + "properties": { + "id": { "type": "string", "description": "任务唯一标识" }, + "content": { "type": "string", "description": "任务描述" }, + "status": { + "type": "string", + "enum": ["pending", "in_progress", "completed"], + "description": "任务状态" + }, + "blockedBy": { + "type": "array", + "items": { "type": "string" }, + "description": "该任务依赖的其他任务ID列表（这些任务必须先完成）" + } + }, + "required": ["id", "content", "status"] + } + } + }, + "required": ["todos"] + }) + } + + /// 纯格式化输出，无副作用（持久化由 runtime 层处理），并发安全 + fn is_concurrency_safe(&self, _args: &serde_json::Value) -> bool { + true + } + + async fn execute(&self, args: serde_json::Value, _ctx: &ToolContext) -> ToolOutput { + let todos = match args.get("todos").and_then(|t| t.as_array()) { + Some(t) => t, + None => return ToolOutput::error("缺少必需参数 'todos'"), + }; + + // 从 ToolContext 中无法直接获取 session_id + // TodoWrite 工具生成格式化的输出，实际持久化在 runtime 层完成 + // 这里只做验证和格式化 + + let mut formatted = String::from("📋 当前任务计划：\n\n"); + let mut in_progress_count = 0; + + for todo in todos { + let id = todo.get("id").and_then(|v| v.as_str()).unwrap_or("?"); + let content = todo.get("content").and_then(|v| v.as_str()).unwrap_or("?"); + let status = todo + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("pending"); + + let icon = match status { + "in_progress" => { + in_progress_count += 1; + "🔄" + } + "completed" => "✅", + _ => "⏳", + }; + + // 显示依赖关系 + let blocked_by: Vec = todo + .get("blockedBy") + .and_then(|b| b.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + + let mut line = format!("{} [{}] {}", icon, id, content); + if !blocked_by.is_empty() { + line.push_str(&format!(" (依赖: {})", blocked_by.join(", "))); + } + formatted.push_str(&line); + formatted.push('\n'); + } + + // 验证约束 + if in_progress_count > 1 { + formatted.push_str( + "\n⚠️ 提醒：你当前有多个 in_progress 任务。请先完成当前任务再开始下一个。", + ); + } else if in_progress_count == 0 + && todos + .iter() + .any(|t| t.get("status").and_then(|v| v.as_str()) == Some("pending")) + { + formatted + .push_str("\n💡 提示：还有待处理任务，请选择一个设为 in_progress 并开始执行。"); + } + + ToolOutput::success(formatted, json!({ "task_count": todos.len() })) + } +} + +/// 将 TodoWrite 的任务列表持久化到 agent_tasks 表。 +/// +/// 使用 INSERT OR REPLACE 实现 upsert（基于 session_id + task_id 唯一约束）。 +/// `owner` 参数指定任务的归属 agent（默认 "lead"）。 +pub async fn persist_tasks( + db: &SqlitePool, + session_id: &str, + todos: &[serde_json::Value], + owner: &str, +) -> anyhow::Result<()> { + for todo in todos { + let task_id = todo.get("id").and_then(|v| v.as_str()).unwrap_or("unknown"); + let content = todo.get("content").and_then(|v| v.as_str()).unwrap_or(""); + let status = todo + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("pending"); + + // 收集 blockedBy 数组 + let blocked_by: Vec = todo + .get("blockedBy") + .and_then(|b| b.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + + let blocked_by_json = + serde_json::to_string(&blocked_by).unwrap_or_else(|_| "[]".to_string()); + + // 简单 DAG 验证：不能依赖自身 + if blocked_by.contains(&task_id.to_string()) { + warn!( + "[TodoWrite] 任务 {} 依赖自身，已跳过 blockedBy 中的自引用", + task_id + ); + } + + sqlx::query( + "INSERT INTO agent_tasks (session_id, task_id, content, status, blocked_by, owner) \ + VALUES (?, ?, ?, ?, ?, ?) \ + ON CONFLICT(session_id, task_id) DO UPDATE SET \ + content=excluded.content, \ + status=excluded.status, \ + blocked_by=excluded.blocked_by, \ + owner=excluded.owner, \ + updated_at=CURRENT_TIMESTAMP", + ) + .bind(session_id) + .bind(task_id) + .bind(content) + .bind(status) + .bind(&blocked_by_json) + .bind(owner) + .execute(db) + .await?; + } + + info!( + "[TodoWrite] 已持久化 {} 个任务到会话 {}", + todos.len(), + session_id + ); + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_todo_output_format() { + let tool = TodoWriteTool; + let args = json!({ + "todos": [ + {"id": "1", "content": "搜索文献", "status": "completed"}, + {"id": "2", "content": "阅读论文", "status": "in_progress", "blockedBy": ["1"]}, + {"id": "3", "content": "撰写综述", "status": "pending", "blockedBy": ["2"]} + ] + }); + + // 验证参数 schema + let params = tool.parameters(); + assert!(params["required"] + .as_array() + .unwrap() + .contains(&json!("todos"))); + } + + #[test] + fn test_todo_multiple_in_progress_warning() { + // 直接测试格式化逻辑 + let todos = json!([ + {"id": "1", "content": "任务A", "status": "in_progress"}, + {"id": "2", "content": "任务B", "status": "in_progress"} + ]); + + let mut in_progress_count = 0; + for todo in todos.as_array().unwrap() { + if todo.get("status").and_then(|v| v.as_str()) == Some("in_progress") { + in_progress_count += 1; + } + } + assert_eq!(in_progress_count, 2); + } + + #[test] + fn test_todo_with_blocked_by() { + let todos = json!([ + {"id": "1", "content": "文献检索", "status": "completed"}, + {"id": "2", "content": "文献分析", "status": "pending", "blockedBy": ["1"]} + ]); + + let items = todos.as_array().unwrap(); + let blocked: Vec = items[1] + .get("blockedBy") + .and_then(|b| b.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|v| v.as_str().map(String::from)) + .collect() + }) + .unwrap_or_default(); + + assert_eq!(blocked, vec!["1"]); + } +} diff --git a/src/agent/trajectory.rs b/src/agent/trajectory.rs new file mode 100644 index 0000000..8037571 --- /dev/null +++ b/src/agent/trajectory.rs @@ -0,0 +1,171 @@ +// src/agent/trajectory.rs +// +// Trajectory 数据导出 — 为 RLHF/微调准备结构化训练数据。 +// 参考 learn-claude-code "Collect trajectory data" 设计。 + +use serde::Serialize; +use sqlx::SqlitePool; +use std::fs; +use std::path::{Path, PathBuf}; +use tracing::info; + +use super::runtime::AgentMetrics; +use super::terminal::TurnTerminal; + +/// 单条 trajectory 记录（JSONL 格式） +#[derive(Debug, Serialize)] +pub struct TrajectoryRecord { + pub timestamp: String, + pub session_id: String, + pub model: String, + pub system_prompt: String, + pub messages: Vec, + pub final_answer: Option, + pub metrics: TrajectoryMetrics, + pub terminal_reason: String, +} + +#[derive(Debug, Serialize)] +pub struct TrajectoryMessage { + pub role: String, + pub content: Option, + pub tool_calls: Option, + pub tool_call_id: Option, + pub thought: Option, +} + +#[derive(Debug, Serialize)] +pub struct TrajectoryMetrics { + pub total_steps: usize, + pub tool_calls: std::collections::HashMap, + pub compression_count: usize, + pub duplicate_detections: usize, +} + +/// Trajectory 导出器 +pub struct TrajectoryExporter; + +impl TrajectoryExporter { + /// 导出指定会话的完整 trajectory 到 JSONL 文件。 + pub async fn export( + db: &SqlitePool, + session_id: &str, + library_dir: &Path, + model: &str, + system_prompt: &str, + metrics: &AgentMetrics, + terminal: Option<&TurnTerminal>, + ) -> anyhow::Result { + // 从 DB 加载消息 + let messages = Self::load_messages(db, session_id).await?; + + // 提取最终答案 + let final_answer = messages + .iter() + .rev() + .find(|m| m.role == "assistant" && m.tool_calls.is_none()) + .and_then(|m| m.content.clone()); + + let terminal_reason = terminal + .map(|t| t.description().to_string()) + .unwrap_or_else(|| "completed".to_string()); + + let record = TrajectoryRecord { + timestamp: chrono::Utc::now().to_rfc3339(), + session_id: session_id.to_string(), + model: model.to_string(), + system_prompt: system_prompt.to_string(), + messages, + final_answer, + metrics: TrajectoryMetrics { + total_steps: metrics.total_steps, + tool_calls: metrics.tool_calls.clone(), + compression_count: metrics.compression_count, + duplicate_detections: metrics.duplicate_detections, + }, + terminal_reason, + }; + + let dir = library_dir.join("trajectories"); + fs::create_dir_all(&dir)?; + + let timestamp = chrono::Utc::now().format("%Y%m%d_%H%M%S"); + let path = dir.join(format!("{}_{}.jsonl", session_id, timestamp)); + + let json = serde_json::to_string(&record)?; + fs::write(&path, format!("{}\n", json))?; + + info!("[Trajectory] 已导出: {}", path.display()); + Ok(path) + } + + async fn load_messages( + db: &SqlitePool, + session_id: &str, + ) -> anyhow::Result> { + let rows = sqlx::query_as::< + _, + ( + String, + Option, + Option, + Option, + Option, + ), + >( + "SELECT role, content, thought, tool_calls, tool_call_id \ + FROM agent_messages WHERE session_id = ? \ + ORDER BY created_at ASC, step_index ASC", + ) + .bind(session_id) + .fetch_all(db) + .await?; + + Ok(rows + .into_iter() + .map(|(role, content, thought, tool_calls, tool_call_id)| { + let tool_calls_json = tool_calls + .as_deref() + .and_then(|tc| serde_json::from_str(tc).ok()); + TrajectoryMessage { + role, + content, + tool_calls: tool_calls_json, + tool_call_id, + thought, + } + }) + .collect()) + } + + /// 列出所有已导出的 trajectory 文件 + pub fn list_trajectories(library_dir: &Path) -> Vec { + let dir = library_dir.join("trajectories"); + match fs::read_dir(&dir) { + Ok(entries) => entries + .flatten() + .filter_map(|e| { + let name = e.file_name().to_string_lossy().to_string(); + if name.ends_with(".jsonl") { + Some(name) + } else { + None + } + }) + .collect(), + Err(_) => Vec::new(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_list_empty_trajectories() { + let dir = PathBuf::from("/tmp/nonexistent_trajectory_dir"); + let result = TrajectoryExporter::list_trajectories(&dir); + assert!(result.is_empty()); + } +} diff --git a/src/api/agent.rs b/src/api/agent.rs index 337b578..700df2b 100644 --- a/src/api/agent.rs +++ b/src/api/agent.rs @@ -14,7 +14,7 @@ use serde::{Deserialize, Serialize}; use sqlx::Row; use std::convert::Infallible; use std::sync::Arc; -use tracing::{info, error}; +use tracing::{error, info}; use super::AppState; use crate::agent::runtime::{AgentRuntime, AgentStreamEvent}; @@ -32,7 +32,10 @@ pub async fn chat_agent( State(state): State>, Json(req): Json, ) -> Result>>, (StatusCode, String)> { - info!("接收到智能体对话请求: question='{}', session_id={:?}", req.question, req.session_id); + info!( + "接收到智能体对话请求: question='{}', session_id={:?}", + req.question, req.session_id + ); let runtime = AgentRuntime::new(Arc::clone(&state)); let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::(); @@ -102,24 +105,30 @@ pub async fn list_sessions( FROM agent_sessions \ WHERE deleted_at IS NULL \ ORDER BY updated_at DESC \ - LIMIT ? OFFSET ?" + LIMIT ? OFFSET ?", ) .bind(limit) .bind(offset) .fetch_all(&state.db) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询会话列表失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("查询会话列表失败: {}", e), + ) + })?; - let sessions: Vec = rows.iter().map(|r| { - SessionSummary { + let sessions: Vec = rows + .iter() + .map(|r| SessionSummary { session_id: r.get(0), title: r.get(1), model: r.get(2), turn_count: r.get(3), created_at: r.get(4), updated_at: r.get(5), - } - }).collect(); + }) + .collect(); Ok(Json(sessions)) } @@ -156,12 +165,17 @@ pub async fn get_session( let session_row = sqlx::query( "SELECT session_id, title, model, turn_count, created_at, updated_at \ FROM agent_sessions \ - WHERE session_id = ? AND deleted_at IS NULL" + WHERE session_id = ? AND deleted_at IS NULL", ) .bind(&session_id) .fetch_optional(&state.db) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询会话失败: {}", e)))? + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("查询会话失败: {}", e), + ) + })? .ok_or((StatusCode::NOT_FOUND, format!("会话 {} 不存在", session_id)))?; let session = SessionSummary { @@ -185,24 +199,27 @@ pub async fn get_session( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询消息列表失败: {}", e)))?; - let messages: Vec = msg_rows.iter().map(|r| { - let tool_calls_json: Option = r.get(6); - let metadata_json: Option = r.get(9); + let messages: Vec = msg_rows + .iter() + .map(|r| { + let tool_calls_json: Option = r.get(6); + let metadata_json: Option = r.get(9); - MessageRecord { - id: r.get(0), - turn_index: r.get(1), - step_index: r.get(2), - role: r.get(3), - content: r.get(4), - thought: r.get(5), - tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()), - tool_call_id: r.get(7), - token_count: r.get(8), - metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()), - created_at: r.get(10), - } - }).collect(); + MessageRecord { + id: r.get(0), + turn_index: r.get(1), + step_index: r.get(2), + role: r.get(3), + content: r.get(4), + thought: r.get(5), + tool_calls: tool_calls_json.and_then(|s| serde_json::from_str(&s).ok()), + tool_call_id: r.get(7), + token_count: r.get(8), + metadata: metadata_json.and_then(|s| serde_json::from_str(&s).ok()), + created_at: r.get(10), + } + }) + .collect(); Ok(Json(SessionDetail { session, messages })) } @@ -223,11 +240,16 @@ pub async fn delete_session( .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("删除会话失败: {}", e)))?; if result.rows_affected() == 0 { - return Err((StatusCode::NOT_FOUND, format!("会话 {} 不存在或已删除", session_id))); + return Err(( + StatusCode::NOT_FOUND, + format!("会话 {} 不存在或已删除", session_id), + )); } info!("会话已软删除: {}", session_id); - Ok(Json(serde_json::json!({ "status": "deleted", "session_id": session_id }))) + Ok(Json( + serde_json::json!({ "status": "deleted", "session_id": session_id }), + )) } // ── POST /api/chat/sessions/:id/stop ── @@ -240,5 +262,214 @@ pub async fn stop_agent( cancelled.insert(session_id.clone()); } info!("已接收并记录手动中止请求，会话 ID: {}", session_id); - Ok(Json(serde_json::json!({ "status": "stopping", "session_id": session_id }))) + Ok(Json( + serde_json::json!({ "status": "stopping", "session_id": session_id }), + )) +} + +// ── GET /api/chat/metrics ── +// 返回聚合的智能体运行指标 + +#[derive(Debug, Serialize)] +pub struct AgentMetricsResponse { + pub total_sessions: i64, + pub total_tool_calls: i64, + pub tool_call_breakdown: serde_json::Value, + pub avg_steps_per_session: f64, + pub error_rate: f64, +} + +pub async fn get_agent_metrics( + State(state): State>, +) -> Result, (StatusCode, String)> { + // 总会话数 + let total_sessions: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_sessions WHERE deleted_at IS NULL") + .fetch_one(&state.db) + .await + .unwrap_or(0); + + // 工具调用统计（从审计日志聚合） + let total_tool_calls: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status IN ('OK', 'FAIL')") + .fetch_one(&state.db) + .await + .unwrap_or(0); + + // 各工具调用次数 + let breakdown_rows: Vec<(String, i64)> = sqlx::query_as( + "SELECT COALESCE(tool_name, 'unknown'), COUNT(*) as cnt \ + FROM agent_audit_log \ + WHERE status IN ('OK', 'FAIL') \ + GROUP BY tool_name \ + ORDER BY cnt DESC", + ) + .fetch_all(&state.db) + .await + .unwrap_or_default(); + + let tool_call_breakdown: serde_json::Value = breakdown_rows + .iter() + .map(|(name, cnt)| serde_json::json!({ name: cnt })) + .fold(serde_json::json!({}), |mut acc, v| { + if let serde_json::Value::Object(map) = &mut acc { + if let serde_json::Value::Object(v_map) = v { + for (k, val) in v_map { + map.insert(k.clone(), val.clone()); + } + } + } + acc + }); + + // 平均步数 + let avg_steps: f64 = sqlx::query_scalar( + "SELECT COALESCE(AVG(CAST(turn_count AS REAL)), 0.0) \ + FROM agent_sessions WHERE deleted_at IS NULL", + ) + .fetch_one(&state.db) + .await + .unwrap_or(0.0); + + // 错误率 + let total_errors: i64 = + sqlx::query_scalar("SELECT COUNT(*) FROM agent_audit_log WHERE status = 'FAIL'") + .fetch_one(&state.db) + .await + .unwrap_or(0); + + let error_rate = if total_tool_calls > 0 { + total_errors as f64 / total_tool_calls as f64 + } else { + 0.0 + }; + + Ok(Json(AgentMetricsResponse { + total_sessions, + total_tool_calls, + tool_call_breakdown, + avg_steps_per_session: avg_steps, + error_rate, + })) +} + +// ── GET /api/chat/sessions/:id/audit ── +// 返回指定会话的审计日志 + +#[derive(Debug, Serialize)] +pub struct AuditLogEntry { + pub id: i64, + pub step: i32, + pub tool_name: Option, + pub status: String, + pub elapsed_ms: i32, + pub output_preview: Option, + pub created_at: String, +} + +pub async fn get_session_audit( + State(state): State>, + Path(session_id): Path, +) -> Result>, (StatusCode, String)> { + let rows = sqlx::query( + "SELECT id, step, tool_name, status, elapsed_ms, output_preview, created_at \ + FROM agent_audit_log \ + WHERE session_id = ? \ + ORDER BY id ASC", + ) + .bind(&session_id) + .fetch_all(&state.db) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("查询审计日志失败: {}", e), + ) + })?; + + let entries: Vec = rows + .iter() + .map(|r| AuditLogEntry { + id: r.get(0), + step: r.get(1), + tool_name: r.get(2), + status: r.get(3), + elapsed_ms: r.get(4), + output_preview: r.get(5), + created_at: r.get(6), + }) + .collect(); + + Ok(Json(entries)) +} + +// ── POST /api/chat/answer_question ── +// 用户回答 Agent 的提问（ask_user 工具配合使用） + +#[derive(Debug, Deserialize)] +pub struct AnswerQuestionRequest { + pub question_id: String, + pub answers: Vec, + pub free_text: Option, +} + +pub async fn answer_question( + State(state): State>, + Json(req): Json, +) -> Result, (StatusCode, String)> { + use crate::agent::tools::ask_user::UserAnswer; + + let mut pending = match state.pending_questions.lock() { + Ok(p) => p, + Err(_) => { + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + "服务器内部状态异常，请稍后重试".to_string(), + )); + } + }; + let question_id = req.question_id.clone(); + + match pending.remove(&question_id) { + Some(pq) => { + let answer = UserAnswer { + question_id: question_id.clone(), + answers: req.answers.clone(), + free_text: req.free_text.clone(), + }; + match pq.answer_tx.send(answer) { + Ok(()) => { + info!("[API] 用户回答了问题: id={}", question_id); + Ok(Json( + serde_json::json!({"status": "ok", "question_id": question_id}), + )) + } + Err(_) => Err((StatusCode::GONE, "问题已超时或已被回答".to_string())), + } + } + None => Err(( + StatusCode::NOT_FOUND, + format!("未找到待回答问题: {}", question_id), + )), + } +} + +// ── GET /api/chat/pending_questions ── +// 获取当前待回答的问题（前端轮询或初始化） + +pub async fn get_pending_questions( + State(state): State>, +) -> Json> { + let pending = match state.pending_questions.lock() { + Ok(p) => p, + Err(_) => return Json(Vec::new()), + }; + let questions: Vec = pending + .iter() + .map(|(id, pq)| { + serde_json::from_str::(&pq.question_json) + .unwrap_or(serde_json::json!({"question_id": id})) + }) + .collect(); + Json(questions) } diff --git a/src/api/helpers.rs b/src/api/helpers.rs index 8115290..9a605d5 100644 --- a/src/api/helpers.rs +++ b/src/api/helpers.rs @@ -1,22 +1,26 @@ // src/api/helpers.rs -use sqlx::{SqlitePool, Row}; -use tracing::info; +use super::StandardPaper; use crate::clients::ads::AdsPaperDoc; use crate::clients::arxiv::ArxivPaper; -use super::StandardPaper; +use sqlx::{Row, SqlitePool}; +use tracing::info; pub fn convert_ads_doc_to_standard(doc: &AdsPaperDoc) -> StandardPaper { - let title = doc.title.as_ref() + let title = doc + .title + .as_ref() .and_then(|v: &Vec| v.first()) .cloned() .unwrap_or_else(|| doc.bibcode.clone()); let authors = doc.author.clone().unwrap_or_default(); let keywords = doc.keyword.clone().unwrap_or_default(); - let doi = doc.doi.as_ref() + let doi = doc + .doi + .as_ref() .and_then(|v: &Vec| v.first()) .cloned() .unwrap_or_default(); - + let mut arxiv_id = String::new(); if let Some(identifiers) = &doc.identifier { for id in identifiers { @@ -26,11 +30,10 @@ pub fn convert_ads_doc_to_standard(doc: &AdsPaperDoc) -> StandardPaper { } } } - if arxiv_id.is_empty() { - if doc.bibcode.starts_with("arXiv") { + if arxiv_id.is_empty() + && doc.bibcode.starts_with("arXiv") { arxiv_id = doc.bibcode.replace("arXiv", "").trim().to_string(); } - } StandardPaper { bibcode: doc.bibcode.clone(), @@ -87,6 +90,7 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res // 1. 如果存在 arxiv_id，检查是否有已存在的相同 arxiv_id 记录以防 duplicate if !p.arxiv_id.is_empty() { + #[allow(clippy::type_complexity)] let existing_opt: Option<(String, Option, Option, Option, Option)> = sqlx::query_as( "SELECT bibcode, pdf_path, html_path, markdown_path, translation_path FROM papers WHERE arxiv_id = ?" ) @@ -102,7 +106,10 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res let is_new_formal = p.bibcode != p.arxiv_id; if is_existing_temp && is_new_formal { - info!("发现相同 arXiv ID 的文献，将临时主键 {} 升级为正式 ADS Bibcode: {}", existing_bibcode, p.bibcode); + info!( + "发现相同 arXiv ID 的文献，将临时主键 {} 升级为正式 ADS Bibcode: {}", + existing_bibcode, p.bibcode + ); sqlx::query( "UPDATE papers SET bibcode = ?, title = ?, authors = ?, year = ?, pub = ?, keywords = ?, abstract = ?, doi = ?, citation_count = ?, reference_count = ?, doctype = ? WHERE bibcode = ?" ) @@ -120,11 +127,14 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res .bind(&existing_bibcode) .execute(db) .await?; - + return Ok(()); } else { // 如果已存在的是正式 ADS bibcode，而新插入的是临时 arXiv ID，直接忽略或更新元数据而不更改主键 - info!("发现相同 arXiv ID 的文献 {} 已存在正式记录，忽略临时 arXiv 插入", existing_bibcode); + info!( + "发现相同 arXiv ID 的文献 {} 已存在正式记录，忽略临时 arXiv 插入", + existing_bibcode + ); return Ok(()); } } @@ -165,7 +175,11 @@ pub async fn save_paper_to_db(db: &SqlitePool, p: &StandardPaper) -> anyhow::Res Ok(()) } -pub async fn get_paper_from_db(db: &SqlitePool, library_dir: &std::path::Path, identifier: &str) -> anyhow::Result { +pub async fn get_paper_from_db( + db: &SqlitePool, + library_dir: &std::path::Path, + identifier: &str, +) -> anyhow::Result { let clean_id = identifier.trim(); let clean_doi = clean_id .trim_start_matches("doi:") @@ -196,19 +210,37 @@ pub async fn get_paper_from_db(db: &SqlitePool, library_dir: &std::path::Path, i let has_vector: bool = r.get(16); let authors_str: Option = r.get(2); - let authors: Vec = authors_str.and_then(|s| serde_json::from_str(&s).ok()).unwrap_or_default(); + let authors: Vec = authors_str + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default(); let keywords_str: Option = r.get(5); - let keywords: Vec = keywords_str.and_then(|s| serde_json::from_str(&s).ok()).unwrap_or_default(); + let keywords: Vec = keywords_str + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default(); - let is_pdf_exist = pdf_path.as_ref().map(|p| library_dir.join(p).exists()).unwrap_or(false); - let is_html_exist = html_path.as_ref().map(|p| library_dir.join(p).exists()).unwrap_or(false); - let is_md_exist = markdown_path.as_ref().map(|p| library_dir.join(p).exists()).unwrap_or(false); - let is_tr_exist = translation_path.as_ref().map(|p| library_dir.join(p).exists()).unwrap_or(false); + let is_pdf_exist = pdf_path + .as_ref() + .map(|p| library_dir.join(p).exists()) + .unwrap_or(false); + let is_html_exist = html_path + .as_ref() + .map(|p| library_dir.join(p).exists()) + .unwrap_or(false); + let is_md_exist = markdown_path + .as_ref() + .map(|p| library_dir.join(p).exists()) + .unwrap_or(false); + let is_tr_exist = translation_path + .as_ref() + .map(|p| library_dir.join(p).exists()) + .unwrap_or(false); - let pdf_error = pdf_path.as_ref() + let pdf_error = pdf_path + .as_ref() .filter(|p| p.starts_with("error:")) .map(|p| p["error:".len()..].trim().to_string()); - let html_error = html_path.as_ref() + let html_error = html_path + .as_ref() .filter(|p| p.starts_with("error:")) .map(|p| p["error:".len()..].trim().to_string()); @@ -237,10 +269,17 @@ pub async fn get_paper_from_db(db: &SqlitePool, library_dir: &std::path::Path, i } pub async fn check_paper_paths_in_db( - db: &SqlitePool, + db: &SqlitePool, library_dir: &std::path::Path, - identifier: &str -) -> anyhow::Result, Option, Option, Option)>> { + identifier: &str, +) -> anyhow::Result< + Option<( + Option, + Option, + Option, + Option, + )>, +> { let clean_id = identifier.trim(); let clean_doi = clean_id .trim_start_matches("doi:") @@ -333,7 +372,10 @@ mod tests { reference_count: None, reference: None, citation: None, - identifier: Some(vec!["2026MNRAS.530.1234A".to_string(), "arXiv:2606.12345".to_string()]), + identifier: Some(vec![ + "2026MNRAS.530.1234A".to_string(), + "arXiv:2606.12345".to_string(), + ]), doctype: Some("article".to_string()), }; @@ -372,9 +414,7 @@ mod tests { .await?; // 运行迁移 - sqlx::migrate!("./migrations") - .run(&pool) - .await?; + sqlx::migrate!("./migrations").run(&pool).await?; let paper = StandardPaper { bibcode: "2026A&A...123..456X".to_string(), @@ -403,19 +443,23 @@ mod tests { save_paper_to_db(&pool, &paper).await?; // 读取 - let retrieved = get_paper_from_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?; + let retrieved = + get_paper_from_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?; assert_eq!(retrieved.title, paper.title); assert_eq!(retrieved.authors, paper.authors); assert_eq!(retrieved.keywords, paper.keywords); // 读取 by DOI / DOI 前缀 - let retrieved_by_doi = get_paper_from_db(&pool, std::path::Path::new(""), "10.1000/test.doi").await?; + let retrieved_by_doi = + get_paper_from_db(&pool, std::path::Path::new(""), "10.1000/test.doi").await?; assert_eq!(retrieved_by_doi.bibcode, paper.bibcode); - let retrieved_by_doi_prefix = get_paper_from_db(&pool, std::path::Path::new(""), "doi:10.1000/test.doi").await?; + let retrieved_by_doi_prefix = + get_paper_from_db(&pool, std::path::Path::new(""), "doi:10.1000/test.doi").await?; assert_eq!(retrieved_by_doi_prefix.bibcode, paper.bibcode); // 检查路径状态（初始为 None） - let paths = check_paper_paths_in_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?; + let paths = + check_paper_paths_in_db(&pool, std::path::Path::new(""), "2026A&A...123..456X").await?; assert!(paths.is_some()); let (pdf, html, md, tr) = paths.unwrap(); assert!(pdf.is_none()); diff --git a/src/api/mod.rs b/src/api/mod.rs index 6cb950e..33e2363 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -1,14 +1,33 @@ // src/api/mod.rs -use std::sync::Arc; -use serde::{Deserialize, Serialize}; -use sqlx::SqlitePool; -use crate::Config; -use crate::services::translation::Dictionary; -use crate::clients::qiniu::QiniuClient; +use crate::agent::memory::MemoryManager; +use crate::agent::skills::SkillRegistry; use crate::clients::ads::AdsClient; use crate::clients::arxiv::ArxivClient; -use crate::clients::llm::{LlmClient, EmbeddingClient}; +use crate::clients::llm::{EmbeddingClient, LlmClient}; +use crate::clients::qiniu::QiniuClient; use crate::services::download::Downloader; +use crate::services::translation::Dictionary; +use crate::Config; +use serde::{Deserialize, Serialize}; +use sqlx::SqlitePool; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, RwLock}; +use tokio::sync::{broadcast, oneshot}; + +/// 提供给前端的 SSE 事件 +#[derive(Debug, Clone, Serialize)] +#[serde(tag = "type")] +pub enum AppEvent { + #[serde(rename = "user_question")] + UserQuestion { data: String }, +} + +/// 待回答的用户问题（供 ask_user 工具和 API 端点共享） +#[derive(Debug)] +pub struct PendingQuestion { + pub question_json: String, + pub answer_tx: oneshot::Sender, +} // 全局共享的 Axum 应用上下文状态 pub struct AppState { @@ -24,7 +43,14 @@ pub struct AppState { pub harvest_status: Arc>, pub batch_status: Arc>, pub active_bibcode: Arc>>, - pub cancelled_runs: Arc>>, + pub cancelled_runs: Arc>>, + pub skill_registry: Arc>, + /// ask_user 工具 — 待回答的问题 + pub pending_questions: Arc>>, + /// SSE 广播通道（agent 运行时向所有连接的客户端推送事件） + pub sse_broadcast: Option>, + /// 项目记忆管理器（跨会话持久化） + pub memory_manager: Arc>, } // 统一标准化的文献格式，用于向前端传输 @@ -52,46 +78,47 @@ pub struct StandardPaper { pub html_error: Option, } +pub mod agent; pub mod helpers; -pub mod papers; pub mod notes; +pub mod papers; pub mod sync; pub mod targets; -pub mod agent; // 提供兼容的 handlers 命名空间，避免修改 main.rs / batch_sync.rs 里的导入 pub mod handlers { - pub use super::helpers::{ - convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db, - get_paper_from_db, check_paper_paths_in_db, + pub use super::agent::{ + answer_question, chat_agent, delete_session, get_agent_metrics, get_pending_questions, + get_session, get_session_audit, list_sessions, stop_agent, AgentChatRequest, + AgentMetricsResponse, AuditLogEntry, MessageRecord, SessionDetail, SessionListParams, + SessionSummary, }; - pub use super::papers::{ - search_papers, download_paper, parse_paper, translate_paper, embed_paper, - get_citation_network, get_paper_detail, get_library, export_citations, - upload_paper_file, mark_no_resource, get_active_bibcode, set_active_bibcode, - SearchParams, DownloadRequest, ParseRequest, ParseResponse, - TranslateRequest, TranslateResponse, CitationsResponse, PaperDetailResponse, - ExportRequest, ExportResponse, MarkNoResourceRequest, EmbedRequest, EmbedResponse, + pub use super::helpers::{ + check_paper_paths_in_db, convert_ads_doc_to_standard, convert_arxiv_to_standard, + get_paper_from_db, save_paper_to_db, }; pub use super::notes::{ - create_note, get_notes, delete_note, - NoteRecord, CreateNoteRequest, DeleteNoteParams, GetNotesParams, + create_note, delete_note, get_notes, CreateNoteRequest, DeleteNoteParams, GetNotesParams, + NoteRecord, + }; + pub use super::papers::{ + download_paper, embed_paper, export_citations, get_active_bibcode, get_citation_network, + get_library, get_paper_detail, mark_no_resource, parse_paper, search_papers, + set_active_bibcode, translate_paper, upload_paper_file, CitationsResponse, DownloadRequest, + EmbedRequest, EmbedResponse, ExportRequest, ExportResponse, MarkNoResourceRequest, + PaperDetailResponse, ParseRequest, ParseResponse, SearchParams, TranslateRequest, + TranslateResponse, }; pub use super::sync::{ - run_meta_sync, get_meta_sync_count, get_meta_sync_status, - run_asset_batch, stop_asset_batch, get_sync_queries, delete_sync_query, - get_asset_batch_status, MetaSyncRunRequest, MetaSyncCountRequest, - MetaSyncCountResponse, AssetBatchRunRequest, SavedSyncQuery, + delete_sync_query, get_asset_batch_status, get_meta_sync_count, get_meta_sync_status, + get_sync_queries, run_asset_batch, run_meta_sync, stop_asset_batch, AssetBatchRunRequest, + MetaSyncCountRequest, MetaSyncCountResponse, MetaSyncRunRequest, SavedSyncQuery, }; pub use super::targets::{ - chat_rag, query_target, associate_target, list_targets, extract_paper_targets, - chat_figure, - RagAskRequest, TargetQueryParams, AssociateTargetRequest, TargetListParams, AssociateResponse, - ExtractTargetsRequest, ExtractTargetsResponse, ChatFigureRequest, ChatResponse, - }; - pub use super::agent::{ - chat_agent, list_sessions, get_session, delete_session, stop_agent, - AgentChatRequest, SessionSummary, SessionDetail, MessageRecord, SessionListParams, + associate_target, chat_figure, chat_rag, extract_paper_targets, list_targets, query_target, + AssociateResponse, AssociateTargetRequest, ChatFigureRequest, ChatResponse, + ExtractTargetsRequest, ExtractTargetsResponse, RagAskRequest, TargetListParams, + TargetQueryParams, }; pub use super::{AppState, StandardPaper}; } diff --git a/src/api/notes.rs b/src/api/notes.rs index 1103fc8..cbc7d1b 100644 --- a/src/api/notes.rs +++ b/src/api/notes.rs @@ -5,8 +5,8 @@ use axum::{ Json, }; use serde::{Deserialize, Serialize}; -use std::sync::Arc; use sqlx::Row; +use std::sync::Arc; use super::AppState; @@ -85,15 +85,18 @@ pub async fn get_notes( .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("查询笔记失败: {}", e)))?; - let notes: Vec = rows.iter().map(|r| NoteRecord { - id: r.get(0), - bibcode: r.get(1), - paragraph_index: r.get(2), - note_text: r.get(3), - highlight_color: r.get(4), - selected_text: r.get(5), - created_at: r.get(6), - }).collect(); + let notes: Vec = rows + .iter() + .map(|r| NoteRecord { + id: r.get(0), + bibcode: r.get(1), + paragraph_index: r.get(2), + note_text: r.get(3), + highlight_color: r.get(4), + selected_text: r.get(5), + created_at: r.get(6), + }) + .collect(); Ok(Json(notes)) } @@ -107,7 +110,12 @@ pub async fn delete_note( .bind(params.id) .execute(&state.db) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("删除笔记失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("删除笔记失败: {}", e), + ) + })?; Ok(StatusCode::NO_CONTENT) } diff --git a/src/api/papers.rs b/src/api/papers.rs index 5081d20..076fb77 100644 --- a/src/api/papers.rs +++ b/src/api/papers.rs @@ -5,15 +5,15 @@ use axum::{ Json, }; use serde::{Deserialize, Serialize}; -use std::sync::Arc; -use std::fs; -use tracing::{info, error}; use sqlx::Row; +use std::fs; +use std::sync::Arc; +use tracing::{error, info}; -use super::{AppState, StandardPaper}; use super::helpers::{ - convert_ads_doc_to_standard, save_paper_to_db, get_paper_from_db, check_paper_paths_in_db, + check_paper_paths_in_db, convert_ads_doc_to_standard, get_paper_from_db, save_paper_to_db, }; +use super::{AppState, StandardPaper}; // 检索请求参数 #[derive(Debug, Deserialize)] @@ -21,8 +21,8 @@ pub struct SearchParams { pub q: String, pub source: Option, // "all" | "ads" | "arxiv" pub rows: Option, - pub start: Option, // 分页起始偏移量 - pub sort: Option, // 排序字段 + pub start: Option, // 分页起始偏移量 + pub sort: Option, // 排序字段 } // ── GET /api/search ── @@ -36,7 +36,9 @@ pub async fn search_papers( let start = params.start.unwrap_or(0); let sort = params.sort.as_deref().unwrap_or("relevance"); - match crate::services::search::search_papers(&state, ¶ms.q, &source, start, rows, sort).await { + match crate::services::search::search_papers(&state, ¶ms.q, &source, start, rows, sort) + .await + { Ok(results) => Ok(Json(results)), Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e.to_string())), } @@ -55,16 +57,16 @@ pub async fn download_paper( Json(req): Json, ) -> Result, (StatusCode, String)> { let force = req.force.unwrap_or(false); - info!("接收到文献下载指令，标识符: {}, 强制重下: {}", req.bibcode, force); + info!( + "接收到文献下载指令，标识符: {}, 强制重下: {}", + req.bibcode, force + ); - let paper = state.downloader.download_paper_service( - &state.db, - &state.config.library_dir, - &req.bibcode, - force, - ) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + let paper = state + .downloader + .download_paper_service(&state.db, &state.config.library_dir, &req.bibcode, force) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; Ok(Json(paper)) } @@ -86,7 +88,10 @@ pub async fn parse_paper( State(state): State>, Json(req): Json, ) -> Result, (StatusCode, String)> { - info!("接收到文献结构化解析指令: {} (强制重新解析: {:?})", req.bibcode, req.force); + info!( + "接收到文献结构化解析指令: {} (强制重新解析: {:?})", + req.bibcode, req.force + ); let force = req.force.unwrap_or(false); let markdown = crate::services::parser::parse_paper_service( @@ -100,7 +105,11 @@ pub async fn parse_paper( .await .map_err(|e| { let msg = e.to_string(); - let status = if msg.contains("未注册") || msg.contains("未找到") || msg.contains("丢失") || msg.contains("未在数据库中注册") { + let status = if msg.contains("未注册") + || msg.contains("未找到") + || msg.contains("丢失") + || msg.contains("未在数据库中注册") + { StatusCode::NOT_FOUND } else if msg.contains("请先下载") { StatusCode::BAD_REQUEST @@ -131,12 +140,16 @@ pub async fn translate_paper( Json(req): Json, ) -> Result, (StatusCode, String)> { let force = req.force.unwrap_or(false); - info!("接收到对比翻译请求: 文献={}, 强制重译={}", req.bibcode, force); + info!( + "接收到对比翻译请求: 文献={}, 强制重译={}", + req.bibcode, force + ); - let (_, _, md_opt, tr_opt) = check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode) - .await - .map_err(|e| (StatusCode::NOT_FOUND, format!("查询文献路径失败: {}", e)))? - .ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?; + let (_, _, md_opt, tr_opt) = + check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode) + .await + .map_err(|e| (StatusCode::NOT_FOUND, format!("查询文献路径失败: {}", e)))? + .ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?; // 若本地已存在翻译物理文件且未指明强制重译，直读本地缓存返回 if !force { @@ -144,7 +157,9 @@ pub async fn translate_paper( let tr_abs = state.config.library_dir.join(&tr_rel); if tr_abs.exists() { if let Ok(content) = fs::read_to_string(&tr_abs) { - return Ok(Json(TranslateResponse { translation: content })); + return Ok(Json(TranslateResponse { + translation: content, + })); } } } @@ -154,37 +169,71 @@ pub async fn translate_paper( let md_rel = match md_opt { Some(rel) => rel, None => { - error!("文献 {} 翻译失败：文献未完成解析，缺少英文 Markdown 路径", req.bibcode); - return Err((StatusCode::BAD_REQUEST, "文献必须先完成解析方可翻译".to_string())); + error!( + "文献 {} 翻译失败：文献未完成解析，缺少英文 Markdown 路径", + req.bibcode + ); + return Err(( + StatusCode::BAD_REQUEST, + "文献必须先完成解析方可翻译".to_string(), + )); } }; let md_abs = state.config.library_dir.join(&md_rel); if !md_abs.exists() { - error!("文献 {} 翻译失败：解析的英文 Markdown 文件 {:?} 不存在", req.bibcode, md_abs); - return Err((StatusCode::BAD_REQUEST, "解析 Markdown 文件丢失".to_string())); + error!( + "文献 {} 翻译失败：解析的英文 Markdown 文件 {:?} 不存在", + req.bibcode, md_abs + ); + return Err(( + StatusCode::BAD_REQUEST, + "解析 Markdown 文件丢失".to_string(), + )); } - let english_markdown = fs::read_to_string(&md_abs) - .map_err(|e| { - error!("文献 {} 翻译失败：读取解析内容失败: {}", req.bibcode, e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("读取解析内容失败: {}", e)) - })?; + let english_markdown = fs::read_to_string(&md_abs).map_err(|e| { + error!("文献 {} 翻译失败：读取解析内容失败: {}", req.bibcode, e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("读取解析内容失败: {}", e), + ) + })?; // 调用 LLM 翻译服务并注入对照词表 - let translated_markdown = crate::services::translation::translate_markdown(&english_markdown, &state.dict, &state.llm) - .await - .map_err(|e| { - error!("文献 {} 翻译失败：调用 LLM 翻译发生错误: {}", req.bibcode, e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("调用 LLM 翻译失败: {}", e)) - })?; + let translated_markdown = crate::services::translation::translate_markdown( + &english_markdown, + &state.dict, + &state.llm, + ) + .await + .map_err(|e| { + error!( + "文献 {} 翻译失败：调用 LLM 翻译发生错误: {}", + req.bibcode, e + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("调用 LLM 翻译失败: {}", e), + ) + })?; // 翻译结果物理写入本地 let tr_filename = format!("{}_zh.md", req.bibcode); - let tr_dest = state.config.library_dir.join("Translation").join(&tr_filename); - fs::create_dir_all(tr_dest.parent().unwrap()).unwrap_or_default(); - - fs::write(&tr_dest, &translated_markdown) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("写入翻译文件失败: {}", e)))?; + let tr_dest = state + .config + .library_dir + .join("Translation") + .join(&tr_filename); + if let Some(parent) = tr_dest.parent() { + fs::create_dir_all(parent).unwrap_or_default(); + } + + fs::write(&tr_dest, &translated_markdown).map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("写入翻译文件失败: {}", e), + ) + })?; let relative_tr_path = format!("Translation/{}", tr_filename); @@ -194,9 +243,16 @@ pub async fn translate_paper( .bind(&req.bibcode) .execute(&state.db) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("更新数据库翻译状态失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("更新数据库翻译状态失败: {}", e), + ) + })?; - Ok(Json(TranslateResponse { translation: translated_markdown })) + Ok(Json(TranslateResponse { + translation: translated_markdown, + })) } #[derive(Debug, Serialize)] @@ -213,14 +269,19 @@ pub struct CitationsResponse { // 从 SQLite 查询引用关联，生成引用星系关系树 pub async fn get_citation_network( State(state): State>, - Query(params): Query, + Query(params): Query, ) -> Result, (StatusCode, String)> { - let paper = match get_paper_from_db(&state.db, &state.config.library_dir, ¶ms.bibcode).await { + let paper = match get_paper_from_db(&state.db, &state.config.library_dir, ¶ms.bibcode).await + { Ok(p) => p, Err(_) => { // 如果本地数据库查不到，尝试从 ADS 在线 API 动态获取 if !state.config.ads_api_key.is_empty() { - match state.ads.search(&format!("bibcode:{}", params.bibcode), 0, 1, "relevance").await { + match state + .ads + .search(&format!("bibcode:{}", params.bibcode), 0, 1, "relevance") + .await + { Ok(docs) => { if let Some(doc) = docs.first() { let standard_paper = convert_ads_doc_to_standard(doc); @@ -246,33 +307,47 @@ pub async fn get_citation_network( } standard_paper } else { - return Err((StatusCode::NOT_FOUND, format!("在本地库及 ADS 中均未找到该文献: {}", params.bibcode))); + return Err(( + StatusCode::NOT_FOUND, + format!("在本地库及 ADS 中均未找到该文献: {}", params.bibcode), + )); } } Err(e) => { - return Err((StatusCode::INTERNAL_SERVER_ERROR, format!("在线检索文献元数据失败: {}", e))); + return Err(( + StatusCode::INTERNAL_SERVER_ERROR, + format!("在线检索文献元数据失败: {}", e), + )); } } } else { - return Err((StatusCode::NOT_FOUND, format!("本地数据库未收录该文献，且未配置 ADS_API_KEY，无法在线加载: {}", params.bibcode))); + return Err(( + StatusCode::NOT_FOUND, + format!( + "本地数据库未收录该文献，且未配置 ADS_API_KEY，无法在线加载: {}", + params.bibcode + ), + )); } } }; // 加载引用的文献 - let refs_rows = sqlx::query("SELECT target_bibcode FROM citations_references WHERE source_bibcode = ?") - .bind(¶ms.bibcode) - .fetch_all(&state.db) - .await - .unwrap_or_default(); + let refs_rows = + sqlx::query("SELECT target_bibcode FROM citations_references WHERE source_bibcode = ?") + .bind(¶ms.bibcode) + .fetch_all(&state.db) + .await + .unwrap_or_default(); let references: Vec = refs_rows.iter().map(|row| row.get(0)).collect(); // 加载被引用的文献 - let cits_rows = sqlx::query("SELECT source_bibcode FROM citations_references WHERE target_bibcode = ?") - .bind(¶ms.bibcode) - .fetch_all(&state.db) - .await - .unwrap_or_default(); + let cits_rows = + sqlx::query("SELECT source_bibcode FROM citations_references WHERE target_bibcode = ?") + .bind(¶ms.bibcode) + .fetch_all(&state.db) + .await + .unwrap_or_default(); let citations: Vec = cits_rows.iter().map(|row| row.get(0)).collect(); // 加载关联文献的被引数量 (从 SQLite papers 表获取) @@ -280,11 +355,12 @@ pub async fn get_citation_network( let mut all_related = references.clone(); all_related.extend(citations.clone()); for bib in all_related { - let count_opt: Option = sqlx::query_scalar("SELECT citation_count FROM papers WHERE bibcode = ?") - .bind(&bib) - .fetch_optional(&state.db) - .await - .unwrap_or_default(); + let count_opt: Option = + sqlx::query_scalar("SELECT citation_count FROM papers WHERE bibcode = ?") + .bind(&bib) + .fetch_optional(&state.db) + .await + .unwrap_or_default(); if let Some(c) = count_opt { citation_counts.insert(bib, c); } @@ -312,19 +388,22 @@ pub struct PaperDetailResponse { // 获取文献标准详情和中英双语内容文件数据 pub async fn get_paper_detail( State(state): State>, - Query(params): Query, + Query(params): Query, ) -> Result, (StatusCode, String)> { let paper = get_paper_from_db(&state.db, &state.config.library_dir, ¶ms.bibcode) .await .map_err(|e| (StatusCode::NOT_FOUND, format!("未找到该文献数据: {}", e)))?; - let (_, _, md_opt, tr_opt) = check_paper_paths_in_db(&state.db, &state.config.library_dir, ¶ms.bibcode) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? - .unwrap_or_default(); + let (_, _, md_opt, tr_opt) = + check_paper_paths_in_db(&state.db, &state.config.library_dir, ¶ms.bibcode) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))? + .unwrap_or_default(); - let english_content = md_opt.and_then(|rel| fs::read_to_string(state.config.library_dir.join(rel)).ok()); - let translation_content = tr_opt.and_then(|rel| fs::read_to_string(state.config.library_dir.join(rel)).ok()); + let english_content = + md_opt.and_then(|rel| fs::read_to_string(state.config.library_dir.join(rel)).ok()); + let translation_content = + tr_opt.and_then(|rel| fs::read_to_string(state.config.library_dir.join(rel)).ok()); Ok(Json(PaperDetailResponse { paper, @@ -353,17 +432,29 @@ pub async fn get_library( let has_vector: bool = r.get(16); let authors_str: Option = r.get(2); - let authors: Vec = authors_str.and_then(|s| serde_json::from_str(&s).ok()).unwrap_or_default(); + let authors: Vec = authors_str + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default(); let keywords_str: Option = r.get(5); - let keywords: Vec = keywords_str.and_then(|s| serde_json::from_str(&s).ok()).unwrap_or_default(); + let keywords: Vec = keywords_str + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default(); - let is_pdf_exist = pdf_path.as_ref().map(|p| state.config.library_dir.join(p).exists()).unwrap_or(false); - let is_html_exist = html_path.as_ref().map(|p| state.config.library_dir.join(p).exists()).unwrap_or(false); + let is_pdf_exist = pdf_path + .as_ref() + .map(|p| state.config.library_dir.join(p).exists()) + .unwrap_or(false); + let is_html_exist = html_path + .as_ref() + .map(|p| state.config.library_dir.join(p).exists()) + .unwrap_or(false); - let pdf_error = pdf_path.as_ref() + let pdf_error = pdf_path + .as_ref() .filter(|p| p.starts_with("error:")) .map(|p| p["error:".len()..].trim().to_string()); - let html_error = html_path.as_ref() + let html_error = html_path + .as_ref() .filter(|p| p.starts_with("error:")) .map(|p| p["error:".len()..].trim().to_string()); @@ -382,8 +473,14 @@ pub async fn get_library( is_downloaded: is_pdf_exist || is_html_exist, has_pdf: is_pdf_exist, has_html: is_html_exist, - has_markdown: markdown_path.as_ref().map(|p| state.config.library_dir.join(p).exists()).unwrap_or(false), - has_translation: translation_path.as_ref().map(|p| state.config.library_dir.join(p).exists()).unwrap_or(false), + has_markdown: markdown_path + .as_ref() + .map(|p| state.config.library_dir.join(p).exists()) + .unwrap_or(false), + has_translation: translation_path + .as_ref() + .map(|p| state.config.library_dir.join(p).exists()) + .unwrap_or(false), has_vector, doctype: doctype_val.unwrap_or_else(|| "article".to_string()), pdf_error, @@ -411,11 +508,18 @@ pub async fn export_citations( Json(req): Json, ) -> Result, (StatusCode, String)> { if state.config.ads_api_key.is_empty() { - return Err((StatusCode::BAD_REQUEST, "ADS API key 未在 .env 中配置，无法使用该接口".to_string())); + return Err(( + StatusCode::BAD_REQUEST, + "ADS API key 未在 .env 中配置，无法使用该接口".to_string(), + )); } - let bibtex = state.ads.export_bibtex(req.bibcodes).await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("批量引文导出失败: {}", e)))?; + let bibtex = state.ads.export_bibtex(req.bibcodes).await.map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("批量引文导出失败: {}", e), + ) + })?; Ok(Json(ExportResponse { bibtex })) } @@ -431,9 +535,11 @@ pub async fn upload_paper_file( let mut file_bytes = Vec::new(); let mut file_name = String::new(); - while let Some(field) = multipart.next_field().await.map_err(|e| { - (StatusCode::BAD_REQUEST, format!("解析文件分块失败: {}", e)) - })? { + while let Some(field) = multipart + .next_field() + .await + .map_err(|e| (StatusCode::BAD_REQUEST, format!("解析文件分块失败: {}", e)))? + { let name = field.name().unwrap_or("").to_string(); if name == "bibcode" { bibcode = field.text().await.unwrap_or_default(); @@ -441,9 +547,16 @@ pub async fn upload_paper_file( file_type = field.text().await.unwrap_or_default(); } else if name == "file" { file_name = field.file_name().unwrap_or("").to_string(); - file_bytes = field.bytes().await.map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, format!("读取文件字节流失败: {}", e)) - })?.to_vec(); + file_bytes = field + .bytes() + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("读取文件字节流失败: {}", e), + ) + })? + .to_vec(); } } @@ -451,7 +564,10 @@ pub async fn upload_paper_file( return Err((StatusCode::BAD_REQUEST, "缺少 bibcode 参数".to_string())); } if file_bytes.is_empty() { - return Err((StatusCode::BAD_REQUEST, "上传文件为空或读取失败".to_string())); + return Err(( + StatusCode::BAD_REQUEST, + "上传文件为空或读取失败".to_string(), + )); } // 尝试将可能的 DOI 或 arXiv ID 解析为真实的 bibcode @@ -471,17 +587,22 @@ pub async fn upload_paper_file( .trim_start_matches("https://doi.org/") .trim_start_matches("http://doi.org/") .trim(); - - if let Some(row) = sqlx::query("SELECT bibcode FROM papers WHERE doi = ? OR doi = ? OR LOWER(doi) = LOWER(?)") - .bind(clean_doi) - .bind(&resolved_bibcode) - .bind(clean_doi) - .fetch_optional(&state.db) - .await - .unwrap_or(None) + + if let Some(row) = sqlx::query( + "SELECT bibcode FROM papers WHERE doi = ? OR doi = ? OR LOWER(doi) = LOWER(?)", + ) + .bind(clean_doi) + .bind(&resolved_bibcode) + .bind(clean_doi) + .fetch_optional(&state.db) + .await + .unwrap_or(None) { let found: String = row.get(0); - info!("上传接口：通过 DOI 匹配成功，将 '{}' 解析为 bibcode '{}'", bibcode, found); + info!( + "上传接口：通过 DOI 匹配成功，将 '{}' 解析为 bibcode '{}'", + bibcode, found + ); resolved_bibcode = found; } else { // 尝试匹配 arXiv ID @@ -492,7 +613,7 @@ pub async fn upload_paper_file( .trim(); // 移除可能存在的版本号后缀（如 2303.12345v1 -> 2303.12345） let clean_arxiv_no_version = if let Some(pos) = clean_arxiv.find('v') { - if clean_arxiv[pos+1..].chars().all(|c| c.is_ascii_digit()) { + if clean_arxiv[pos + 1..].chars().all(|c| c.is_ascii_digit()) { &clean_arxiv[..pos] } else { clean_arxiv @@ -526,10 +647,13 @@ pub async fn upload_paper_file( // 校验并保存文件 let is_pdf = file_type == "pdf" || file_name.to_lowercase().ends_with(".pdf"); - + let relative_path = if is_pdf { crate::services::download::validate_pdf_content(&file_bytes).map_err(|e| { - (StatusCode::BAD_REQUEST, format!("PDF 文件内容校验失败: {}", e)) + ( + StatusCode::BAD_REQUEST, + format!("PDF 文件内容校验失败: {}", e), + ) })?; let pdf_filename = format!("{}.pdf", bibcode); let pdf_dest = state.config.library_dir.join("PDF").join(&pdf_filename); @@ -537,15 +661,24 @@ pub async fn upload_paper_file( std::fs::create_dir_all(parent).unwrap_or_default(); } std::fs::write(&pdf_dest, &file_bytes).map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, format!("无法写入 PDF 文件: {}", e)) + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("无法写入 PDF 文件: {}", e), + ) })?; format!("PDF/{}", pdf_filename) } else { let text_content = String::from_utf8(file_bytes).map_err(|_| { - (StatusCode::BAD_REQUEST, "上传的 HTML 文件不是有效的 UTF-8 文本".to_string()) + ( + StatusCode::BAD_REQUEST, + "上传的 HTML 文件不是有效的 UTF-8 文本".to_string(), + ) })?; crate::services::download::validate_html_content_lenient(&text_content).map_err(|e| { - (StatusCode::BAD_REQUEST, format!("HTML 文件内容校验失败: {}", e)) + ( + StatusCode::BAD_REQUEST, + format!("HTML 文件内容校验失败: {}", e), + ) })?; let html_filename = format!("{}.html", bibcode); let html_dest = state.config.library_dir.join("HTML").join(&html_filename); @@ -553,7 +686,10 @@ pub async fn upload_paper_file( std::fs::create_dir_all(parent).unwrap_or_default(); } std::fs::write(&html_dest, &text_content).map_err(|e| { - (StatusCode::INTERNAL_SERVER_ERROR, format!("无法写入 HTML 文件: {}", e)) + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("无法写入 HTML 文件: {}", e), + ) })?; format!("HTML/{}", html_filename) }; @@ -566,12 +702,22 @@ pub async fn upload_paper_file( .bind(&bibcode) .execute(&state.db) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("更新数据库状态失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("更新数据库状态失败: {}", e), + ) + })?; // 重新获取最新的文献信息以更新前端界面 let updated_paper = get_paper_from_db(&state.db, &state.config.library_dir, &bibcode) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("重读文献数据失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("重读文献数据失败: {}", e), + ) + })?; Ok(Json(updated_paper)) } @@ -595,7 +741,12 @@ pub async fn mark_no_resource( .bind(&req.bibcode) .execute(&state.db) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("清除无资源标记失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("清除无资源标记失败: {}", e), + ) + })?; } else { info!("接收到文献无资源标记指令，标识符: {}", req.bibcode); sqlx::query("UPDATE papers SET pdf_path = 'error:no_resource', html_path = 'error:no_resource' WHERE bibcode = ?") @@ -608,7 +759,12 @@ pub async fn mark_no_resource( // 重新获取最新的文献信息以更新前端界面 let updated_paper = get_paper_from_db(&state.db, &state.config.library_dir, &req.bibcode) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("重读文献数据失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("重读文献数据失败: {}", e), + ) + })?; Ok(Json(updated_paper)) } @@ -619,9 +775,7 @@ pub struct ActiveBibcodeResponse { pub bibcode: Option, } -pub async fn get_active_bibcode( - State(state): State>, -) -> Json { +pub async fn get_active_bibcode(State(state): State>) -> Json { let active = state.active_bibcode.lock().await; Json(ActiveBibcodeResponse { bibcode: active.clone(), @@ -661,24 +815,45 @@ pub async fn embed_paper( ) -> Result, (StatusCode, String)> { info!("接收到文献向量化分块入库指令: {}", req.bibcode); - let (_, _, md_opt, _) = check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode) - .await - .map_err(|e| (StatusCode::NOT_FOUND, format!("获取文献路径失败: {}", e)))? - .ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?; + let (_, _, md_opt, _) = + check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode) + .await + .map_err(|e| (StatusCode::NOT_FOUND, format!("获取文献路径失败: {}", e)))? + .ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?; - let md_rel = md_opt.ok_or((StatusCode::BAD_REQUEST, "文献尚未解析为 Markdown，请先执行解析".to_string()))?; + let md_rel = md_opt.ok_or(( + StatusCode::BAD_REQUEST, + "文献尚未解析为 Markdown，请先执行解析".to_string(), + ))?; let md_abs = state.config.library_dir.join(&md_rel); if !md_abs.exists() { - return Err((StatusCode::NOT_FOUND, "文献 Markdown 文件未找到，请重新解析".to_string())); + return Err(( + StatusCode::NOT_FOUND, + "文献 Markdown 文件未找到，请重新解析".to_string(), + )); } - let markdown_content = fs::read_to_string(&md_abs) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("读取 Markdown 文件失败: {}", e)))?; + let markdown_content = fs::read_to_string(&md_abs).map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("读取 Markdown 文件失败: {}", e), + ) + })?; - let chunk_count = crate::services::rag::ingest_paper(&state.db, &state.embedding, &req.bibcode, &markdown_content, None) - .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("文献向量化失败: {}", e)))?; + let chunk_count = crate::services::rag::ingest_paper( + &state.db, + &state.embedding, + &req.bibcode, + &markdown_content, + None, + ) + .await + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("文献向量化失败: {}", e), + ) + })?; Ok(Json(EmbedResponse { chunk_count })) } - diff --git a/src/api/sync.rs b/src/api/sync.rs index 97f1413..c68b731 100644 --- a/src/api/sync.rs +++ b/src/api/sync.rs @@ -5,8 +5,8 @@ use axum::{ Json, }; use serde::{Deserialize, Serialize}; -use std::sync::Arc; use sqlx::Row; +use std::sync::Arc; use tracing::error; use super::AppState; @@ -27,7 +27,10 @@ pub async fn run_meta_sync( { let mut status = state.harvest_status.lock().await; if status.active { - return Err((StatusCode::CONFLICT, "当前已有文献批量同步任务在后台运行中，请勿重复启动".to_string())); + return Err(( + StatusCode::CONFLICT, + "当前已有文献批量同步任务在后台运行中，请勿重复启动".to_string(), + )); } status.active = true; status.query = req.q.clone(); @@ -72,7 +75,12 @@ pub async fn get_meta_sync_count( &state.arxiv, ) .await - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("获取预估文献数失败: {}", e)))?; + .map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("获取预估文献数失败: {}", e), + ) + })?; Ok(Json(MetaSyncCountResponse { total })) } @@ -88,12 +96,12 @@ pub async fn get_meta_sync_status( // ── POST /api/batch/asset/run ── #[derive(Debug, Deserialize)] pub struct AssetBatchRunRequest { - pub target_phase: String, // "download" | "parse" | "translate" | "embed" | "target" - pub limit_count: Option, // 批量处理上限，默认 100 - pub sort_order: Option, // 处理顺序: "default" | "pub_year_desc" | "created_at_desc" + pub target_phase: String, // "download" | "parse" | "translate" | "embed" | "target" + pub limit_count: Option, // 批量处理上限，默认 100 + pub sort_order: Option, // 处理顺序: "default" | "pub_year_desc" | "created_at_desc" pub skip_completed: Option, - pub skip_failed: Option, // 跳过当前失败 ('error:*') - pub skip_preceding_failed: Option, // 跳过前置失败 + pub skip_failed: Option, // 跳过当前失败 ('error:*') + pub skip_preceding_failed: Option, // 跳过前置失败 pub skip_preceding_uncompleted: Option, // 跳过前置未完成 } @@ -117,7 +125,10 @@ pub async fn run_asset_batch( { let status = state.batch_status.lock().await; if status.active { - return Err((StatusCode::CONFLICT, "当前已有文献批量任务在后台运行中，请勿重复启动".to_string())); + return Err(( + StatusCode::CONFLICT, + "当前已有文献批量任务在后台运行中，请勿重复启动".to_string(), + )); } } @@ -128,7 +139,12 @@ pub async fn run_asset_batch( "translate" => crate::services::batch_sync::BatchAction::Translate, "embed" => crate::services::batch_sync::BatchAction::Embed, "target" => crate::services::batch_sync::BatchAction::Target, - _ => return Err((StatusCode::BAD_REQUEST, "不支持的 target_phase 参数值".to_string())), + _ => { + return Err(( + StatusCode::BAD_REQUEST, + "不支持的 target_phase 参数值".to_string(), + )) + } }; let rows = sqlx::query("SELECT bibcode, pdf_path, html_path, markdown_path, translation_path, year, datetime(created_at, 'localtime'), EXISTS(SELECT 1 FROM paper_chunks_content WHERE bibcode = papers.bibcode), EXISTS(SELECT 1 FROM paper_targets WHERE bibcode = papers.bibcode) FROM papers") @@ -152,7 +168,10 @@ pub async fn run_asset_batch( } // 排序 - let sort_order = req.sort_order.clone().unwrap_or_else(|| "default".to_string()); + let sort_order = req + .sort_order + .clone() + .unwrap_or_else(|| "default".to_string()); if sort_order == "pub_year_desc" { records.sort_by(|a, b| b.year.cmp(&a.year)); } else if sort_order == "created_at_desc" { @@ -182,7 +201,9 @@ pub async fn run_asset_batch( let is_no_resource = |path: &Option| -> bool { if let Some(p) = path { - p.starts_with("error:no_resource") || p.starts_with("error:无资源") || p.starts_with("error:无有效全文") + p.starts_with("error:no_resource") + || p.starts_with("error:无资源") + || p.starts_with("error:无有效全文") } else { false } @@ -295,9 +316,7 @@ pub async fn run_asset_batch( } // ── POST /api/batch/asset/stop ── -pub async fn stop_asset_batch( - State(state): State>, -) -> StatusCode { +pub async fn stop_asset_batch(State(state): State>) -> StatusCode { let mut status = state.batch_status.lock().await; if status.active { status.active = false; @@ -352,7 +371,10 @@ pub async fn delete_sync_query( .await .map_err(|e| { error!("删除同步检索配置失败: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("删除同步检索配置失败: {}", e)) + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("删除同步检索配置失败: {}", e), + ) })?; Ok(StatusCode::OK) diff --git a/src/api/targets.rs b/src/api/targets.rs index df1fd81..e4ad96d 100644 --- a/src/api/targets.rs +++ b/src/api/targets.rs @@ -45,13 +45,22 @@ pub async fn chat_rag( Json(req): Json, ) -> Result, (StatusCode, String)> { let top_k = req.top_k.unwrap_or(5); - - let answer = ask(&state.db, &state.embedding, &state.llm, &req.question, top_k) - .await - .map_err(|e| { - tracing::error!("RAG 问答执行失败: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("RAG 问答执行失败: {}", e)) - })?; + + let answer = ask( + &state.db, + &state.embedding, + &state.llm, + &req.question, + top_k, + ) + .await + .map_err(|e| { + tracing::error!("RAG 问答执行失败: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("RAG 问答执行失败: {}", e), + ) + })?; Ok(Json(answer)) } @@ -66,7 +75,10 @@ pub async fn query_target( .await .map_err(|e| { tracing::error!("查询天体信息失败 ({}): {}", params.object_name, e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("查询天体信息失败: {}", e)) + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("查询天体信息失败: {}", e), + ) })?; Ok(Json(info)) @@ -78,13 +90,21 @@ pub async fn associate_target( Json(req): Json, ) -> Result, (StatusCode, String)> { let client = reqwest::Client::new(); - + // 查询并将结果缓存/关联到对应文献 let info = query_target_cached(&state.db, &req.object_name, Some(&req.bibcode), &client) .await .map_err(|e| { - tracing::error!("手动关联天体失败 ({} -> {}): {}", req.object_name, req.bibcode, e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("手动关联天体失败: {}", e)) + tracing::error!( + "手动关联天体失败 ({} -> {}): {}", + req.object_name, + req.bibcode, + e + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("手动关联天体失败: {}", e), + ) })?; Ok(Json(AssociateResponse { @@ -111,20 +131,22 @@ pub async fn list_targets( let targets: Vec = rows .into_iter() - .map(|(name, ra, dec, parallax, spectral_type, v_magnitude, aliases_json)| { - let aliases: Vec = aliases_json - .and_then(|s| serde_json::from_str(&s).ok()) - .unwrap_or_default(); - TargetInfo { - target_name: name, - ra, - dec, - parallax, - spectral_type, - v_magnitude, - aliases, - } - }) + .map( + |(name, ra, dec, parallax, spectral_type, v_magnitude, aliases_json)| { + let aliases: Vec = aliases_json + .and_then(|s| serde_json::from_str(&s).ok()) + .unwrap_or_default(); + TargetInfo { + target_name: name, + ra, + dec, + parallax, + spectral_type, + v_magnitude, + aliases, + } + }, + ) .collect(); Ok(Json(targets)) @@ -147,19 +169,33 @@ pub async fn extract_paper_targets( ) -> Result, (StatusCode, String)> { tracing::info!("接收到文献天体提取与识别指令: {}", req.bibcode); - let (_, _, md_opt, _) = crate::api::helpers::check_paper_paths_in_db(&state.db, &state.config.library_dir, &req.bibcode) - .await - .map_err(|e| (StatusCode::NOT_FOUND, format!("获取文献路径失败: {}", e)))? - .ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?; + let (_, _, md_opt, _) = crate::api::helpers::check_paper_paths_in_db( + &state.db, + &state.config.library_dir, + &req.bibcode, + ) + .await + .map_err(|e| (StatusCode::NOT_FOUND, format!("获取文献路径失败: {}", e)))? + .ok_or((StatusCode::NOT_FOUND, "该文献未注册在数据库中".to_string()))?; - let md_rel = md_opt.ok_or((StatusCode::BAD_REQUEST, "文献尚未解析为 Markdown，请先执行解析".to_string()))?; + let md_rel = md_opt.ok_or(( + StatusCode::BAD_REQUEST, + "文献尚未解析为 Markdown，请先执行解析".to_string(), + ))?; let md_abs = state.config.library_dir.join(&md_rel); if !md_abs.exists() { - return Err((StatusCode::NOT_FOUND, "文献 Markdown 文件未找到，请重新解析".to_string())); + return Err(( + StatusCode::NOT_FOUND, + "文献 Markdown 文件未找到，请重新解析".to_string(), + )); } - let markdown_content = std::fs::read_to_string(&md_abs) - .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("读取 Markdown 文件失败: {}", e)))?; + let markdown_content = std::fs::read_to_string(&md_abs).map_err(|e| { + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("读取 Markdown 文件失败: {}", e), + ) + })?; // 重新识别前先清除该文献已有的天体关联记录，确保陈旧和错误绑定的天体得到重置与刷新 if let Err(e) = sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?") @@ -171,7 +207,13 @@ pub async fn extract_paper_targets( } let client = reqwest::Client::new(); - let targets = crate::services::target::extract_and_cache_targets(&state.db, &markdown_content, &req.bibcode, &client).await; + let targets = crate::services::target::extract_and_cache_targets( + &state.db, + &markdown_content, + &req.bibcode, + &client, + ) + .await; Ok(Json(ExtractTargetsResponse { targets })) } @@ -189,7 +231,10 @@ pub struct ChatResponse { pub answer: String, } -async fn get_image_bytes(library_dir: &std::path::Path, path: &str) -> anyhow::Result<(Vec, String)> { +async fn get_image_bytes( + library_dir: &std::path::Path, + path: &str, +) -> anyhow::Result<(Vec, String)> { if path.starts_with("http://") || path.starts_with("https://") { let client = reqwest::Client::new(); let resp = client.get(path).send().await?; @@ -217,7 +262,8 @@ async fn get_image_bytes(library_dir: &std::path::Path, path: &str) -> anyhow::R "webp" => "image/webp", "svg" => "image/svg+xml", _ => "image/png", - }.to_string(); + } + .to_string(); let bytes = std::fs::read(&abs_path)?; Ok((bytes, mime_type)) } @@ -227,7 +273,12 @@ pub async fn chat_figure( State(state): State>, Json(req): Json, ) -> Result, (StatusCode, String)> { - tracing::info!("接收到图表多模态提问: bibcode={}, image_path={}, question={}", req.bibcode, req.image_path, req.question); + tracing::info!( + "接收到图表多模态提问: bibcode={}, image_path={}, question={}", + req.bibcode, + req.image_path, + req.question + ); let (bytes, mime_type) = get_image_bytes(&state.config.library_dir, &req.image_path) .await @@ -236,21 +287,22 @@ pub async fn chat_figure( (StatusCode::BAD_REQUEST, format!("获取图片失败: {}", e)) })?; - use base64::{Engine as _, engine::general_purpose}; + use base64::{engine::general_purpose, Engine as _}; let base64_data = general_purpose::STANDARD.encode(&bytes); let system_prompt = "You are a professional astronomer and physicist. You are helping a researcher analyze a scientific plot, diagram, or chart extracted from a theoretical or observational astrophysics paper. Please provide an expert, detailed, and clear explanation of the figure in Chinese based on the image provided and the user's question. Format equations in standard LaTeX."; - let answer = state.llm.chat_completion_with_image( - system_prompt, - &req.question, - &base64_data, - &mime_type, - ).await.map_err(|e| { - tracing::error!("多模态大模型请求失败: {}", e); - (StatusCode::INTERNAL_SERVER_ERROR, format!("大模型图表解析失败: {}", e)) - })?; + let answer = state + .llm + .chat_completion_with_image(system_prompt, &req.question, &base64_data, &mime_type) + .await + .map_err(|e| { + tracing::error!("多模态大模型请求失败: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("大模型图表解析失败: {}", e), + ) + })?; Ok(Json(ChatResponse { answer })) } - diff --git a/src/bin/cli.rs b/src/bin/cli.rs index 8a7510a..70a2ecb 100644 --- a/src/bin/cli.rs +++ b/src/bin/cli.rs @@ -3,13 +3,13 @@ // AstroResearch CLI Skills Agent — 向外部 Agent（如 Claude Code）提供 // 标准的命令行工具接口，用于 RAG 问答、天体查询和天体关联操作。 -use std::str::FromStr; use clap::{Parser, Subcommand}; use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; +use std::str::FromStr; use tracing_subscriber::FmtSubscriber; +use astroresearch::clients::llm::{EmbeddingClient, LlmClient}; use astroresearch::Config; -use astroresearch::clients::llm::{LlmClient, EmbeddingClient}; #[derive(Parser)] #[command( @@ -76,8 +76,15 @@ async fn main() -> anyhow::Result<()> { // (sqlite3*, char**, const sqlite3_api_routines*)。 // 该注册必须在开启任何数据库连接前执行。 unsafe { - libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute( - sqlite_vec::sqlite3_vec_init as *const (), + libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::< + *const (), + unsafe extern "C" fn( + *mut libsqlite3_sys::sqlite3, + *mut *const i8, + *const libsqlite3_sys::sqlite3_api_routines, + ) -> i32, + >( + sqlite_vec::sqlite3_vec_init as *const () ))); } @@ -93,9 +100,7 @@ async fn main() -> anyhow::Result<()> { .await?; // 执行迁移 - sqlx::migrate!("./migrations") - .run(&pool) - .await?; + sqlx::migrate!("./migrations").run(&pool).await?; let cli = Cli::parse(); @@ -112,9 +117,9 @@ async fn main() -> anyhow::Result<()> { config.llm_model.clone(), ); - let result = astroresearch::services::rag::ask( - &pool, &embedding, &llm, &question, top_k - ).await?; + let result = + astroresearch::services::rag::ask(&pool, &embedding, &llm, &question, top_k) + .await?; println!("\n📖 回答:\n{}\n", result.answer); if !result.sources.is_empty() { @@ -122,7 +127,10 @@ async fn main() -> anyhow::Result<()> { for (i, src) in result.sources.iter().enumerate() { println!( " [{}] {} §{} (距离: {:.4})", - i + 1, src.bibcode, src.paragraph_index, src.distance + i + 1, + src.bibcode, + src.paragraph_index, + src.distance ); } } @@ -131,8 +139,12 @@ async fn main() -> anyhow::Result<()> { Commands::TargetQuery { object_name } => { let client = reqwest::Client::new(); let info = astroresearch::services::target::query_target_cached( - &pool, &object_name, None, &client - ).await?; + &pool, + &object_name, + None, + &client, + ) + .await?; println!("\n🔭 天体信息: {}", info.target_name); if let Some(ra) = &info.ra { @@ -155,11 +167,18 @@ async fn main() -> anyhow::Result<()> { } } - Commands::TargetAssociate { bibcode, object_name } => { + Commands::TargetAssociate { + bibcode, + object_name, + } => { let client = reqwest::Client::new(); let info = astroresearch::services::target::query_target_cached( - &pool, &object_name, Some(&bibcode), &client - ).await?; + &pool, + &object_name, + Some(&bibcode), + &client, + ) + .await?; println!("✅ 已关联天体 {} -> 文献 {}", info.target_name, bibcode); } @@ -176,7 +195,10 @@ async fn main() -> anyhow::Result<()> { } } - Commands::Ingest { bibcode, markdown_path } => { + Commands::Ingest { + bibcode, + markdown_path, + } => { let content = std::fs::read_to_string(&markdown_path)?; let embedding = EmbeddingClient::new( config.embedding_api_key.clone(), @@ -185,8 +207,9 @@ async fn main() -> anyhow::Result<()> { ); let count = astroresearch::services::rag::ingest_paper( - &pool, &embedding, &bibcode, &content, None - ).await?; + &pool, &embedding, &bibcode, &content, None, + ) + .await?; println!("✅ 文献 {} 向量化完成，写入 {} 个切片", bibcode, count); } diff --git a/src/bin/health_check.rs b/src/bin/health_check.rs index 6a760a8..faceecb 100644 --- a/src/bin/health_check.rs +++ b/src/bin/health_check.rs @@ -1,8 +1,8 @@ // src/bin/health_check.rs +use astroresearch::Config; +use sqlx::{Row, SqlitePool}; use std::fs; use std::path::{Path, PathBuf}; -use sqlx::{SqlitePool, Row}; -use astroresearch::Config; use tracing::{error, Level}; use tracing_subscriber::FmtSubscriber; @@ -76,14 +76,22 @@ fn validate_pdf_content(bytes: &[u8]) -> Result<(), String> { let scan_len = std::cmp::min(2048, bytes.len()); let text = String::from_utf8_lossy(&bytes[..scan_len]); if let Some(desc) = detect_anti_bot(&text) { - return Err(format!("虽然文件后缀是 PDF，但实际内容是 HTML（检测到：{}）", desc)); + return Err(format!( + "虽然文件后缀是 PDF，但实际内容是 HTML（检测到：{}）", + desc + )); } - return Err("虽然文件后缀是 PDF，但实际内容是 HTML 网页，可能是重定向或拦截页面".to_string()); + return Err( + "虽然文件后缀是 PDF，但实际内容是 HTML 网页，可能是重定向或拦截页面".to_string(), + ); } return Err("缺少 %PDF 文件头魔数，文件损坏或并非 PDF".to_string()); } if bytes.len() < 5000 { - return Err(format!("PDF 文件过小（仅 {} 字节），极可能是错误信息页", bytes.len())); + return Err(format!( + "PDF 文件过小（仅 {} 字节），极可能是错误信息页", + bytes.len() + )); } let scan_len = std::cmp::min(1024, bytes.len()); let tail = &bytes[bytes.len() - scan_len..]; @@ -101,11 +109,18 @@ fn validate_html_content(text: &str) -> Result<(), String> { let lower = text.to_lowercase(); // 1. 检查常见的跳转与错误占位特征 - if lower.contains("redirecting") || lower.contains("redirect to") || lower.contains("http-equiv=\"refresh\"") || lower.contains("autoredirecttourl") { + if lower.contains("redirecting") + || lower.contains("redirect to") + || lower.contains("http-equiv=\"refresh\"") + || lower.contains("autoredirecttourl") + { return Err("检测到 HTML 重定向跳转页面，而非真实文献正文".to_string()); } - if lower.contains("conversion to html had a fatal error") || lower.contains("no content available") || lower.contains("fatal error and exited abruptly") { + if lower.contains("conversion to html had a fatal error") + || lower.contains("no content available") + || lower.contains("fatal error and exited abruptly") + { return Err("检测到 ar5iv 转换失败的占位 HTML 页面".to_string()); } @@ -119,7 +134,7 @@ fn validate_html_content(text: &str) -> Result<(), String> { let title_start = start_pos + tag_end + 1; if let Some(end_pos) = lower[title_start..].find("") { let title = &lower[title_start..title_start + end_pos]; - if title.contains("nsf award search") + if title.contains("nsf award search") || title.contains("national science foundation") || title.contains("vizier") || title.contains("caltechthesis") @@ -127,7 +142,10 @@ fn validate_html_content(text: &str) -> Result<(), String> { || title.contains("asp conference series") || title.contains("aspbooks") { - return Err(format!("检测到占位网页标题: \"{}\"，判定为非正本文献", title.trim())); + return Err(format!( + "检测到占位网页标题: \"{}\"，判定为非正本文献", + title.trim() + )); } } } @@ -136,8 +154,12 @@ fn validate_html_content(text: &str) -> Result<(), String> { // 3. 基础字节长度与具体 HTTP 错误特征校验 if text.len() < 2000 { let error_patterns = [ - "404 not found", "403 forbidden", "502 bad gateway", - "500 internal server error", "access denied", "site error" + "404 not found", + "403 forbidden", + "502 bad gateway", + "500 internal server error", + "access denied", + "site error", ]; for kw in &error_patterns { if lower.contains(kw) { @@ -149,21 +171,24 @@ fn validate_html_content(text: &str) -> Result<(), String> { // 4. 结构启发式校验：如果是小于 50KB 的 HTML，必须包含基本的章节或参考文献结构，否则判定为摘要/存根占位页 if text.len() < 50000 { // 匹配 heading 标签或 Markdown 格式的标题，而不是纯文本中的单词 - let has_sections = lower.contains("ltx_title_section") - || lower.contains("class=\"section\"") + let has_sections = lower.contains("ltx_title_section") + || lower.contains("class=\"section\"") || lower.contains("## introduction") || lower.contains("

introduction") || lower.contains("

introduction") || lower.contains("class=\"ltx_section\""); - let has_bib = lower.contains("ltx_bibliography") + let has_bib = lower.contains("ltx_bibliography") || lower.contains("class=\"references\"") || lower.contains("
    Result<(), Box> { // 回调规范。即使 health_check 不直接使用 vec0，也需要注册以防数据库 // 包含 vec0 虚拟表时连接崩溃。 unsafe { - libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute( - sqlite_vec::sqlite3_vec_init as *const (), + libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::< + *const (), + unsafe extern "C" fn( + *mut libsqlite3_sys::sqlite3, + *mut *const i8, + *const libsqlite3_sys::sqlite3_api_routines, + ) -> i32, + >( + sqlite_vec::sqlite3_vec_init as *const () ))); } @@ -235,7 +267,11 @@ async fn main() -> Result<(), Box> { scan_directory(&library_dir.join("HTML"), &mut html_files); scan_directory(&library_dir.join("PDF"), &mut pdf_files); - println!("📂 正在扫描物理磁盘文件 (HTML: {} 个, PDF: {} 个)...", html_files.len(), pdf_files.len()); + println!( + "📂 正在扫描物理磁盘文件 (HTML: {} 个, PDF: {} 个)...", + html_files.len(), + pdf_files.len() + ); let mut disk_html_invalid = 0; let mut disk_pdf_invalid = 0; @@ -247,7 +283,7 @@ async fn main() -> Result<(), Box> { Ok(p) => p.to_str().unwrap_or(""), Err(_) => continue, }; - + if let Ok(content) = fs::read_to_string(&path) { if let Err(e) = validate_html_content(&content) { disk_html_invalid += 1; @@ -260,15 +296,26 @@ async fn main() -> Result<(), Box> { println!(" 🧹 [修复] 已物理删除损坏的文件"); // 检索是否有数据库记录并将其重置 - let res = sqlx::query("UPDATE papers SET html_path = NULL WHERE html_path = ? OR html_path = ?") - .bind(rel_path) - .bind(format!("HTML/{}", Path::new(rel_path).file_name().and_then(|f| f.to_str()).unwrap_or(""))) - .execute(&pool) - .await; + let res = sqlx::query( + "UPDATE papers SET html_path = NULL WHERE html_path = ? OR html_path = ?", + ) + .bind(rel_path) + .bind(format!( + "HTML/{}", + Path::new(rel_path) + .file_name() + .and_then(|f| f.to_str()) + .unwrap_or("") + )) + .execute(&pool) + .await; if let Ok(r) = res { if r.rows_affected() > 0 { db_updated_count += r.rows_affected(); - println!(" ✅ [修复] 数据库对应状态已重置 (受影响行数: {})", r.rows_affected()); + println!( + " ✅ [修复] 数据库对应状态已重置 (受影响行数: {})", + r.rows_affected() + ); } } } @@ -294,15 +341,26 @@ async fn main() -> Result<(), Box> { deleted_files += 1; println!(" 🧹 [修复] 已物理删除损坏的文件"); - let res = sqlx::query("UPDATE papers SET pdf_path = NULL WHERE pdf_path = ? OR pdf_path = ?") - .bind(rel_path) - .bind(format!("PDF/{}", Path::new(rel_path).file_name().and_then(|f| f.to_str()).unwrap_or(""))) - .execute(&pool) - .await; + let res = sqlx::query( + "UPDATE papers SET pdf_path = NULL WHERE pdf_path = ? OR pdf_path = ?", + ) + .bind(rel_path) + .bind(format!( + "PDF/{}", + Path::new(rel_path) + .file_name() + .and_then(|f| f.to_str()) + .unwrap_or("") + )) + .execute(&pool) + .await; if let Ok(r) = res { if r.rows_affected() > 0 { db_updated_count += r.rows_affected(); - println!(" ✅ [修复] 数据库对应状态已重置 (受影响行数: {})", r.rows_affected()); + println!( + " ✅ [修复] 数据库对应状态已重置 (受影响行数: {})", + r.rows_affected() + ); } } } @@ -314,7 +372,7 @@ async fn main() -> Result<(), Box> { // ─── 阶段 2：数据库记录校验扫描（检测丢失文件、报错记录与孤立 Markdown） ─── println!("🗄️ 正在校验数据库表记录一致性..."); let db_rows = sqlx::query( - "SELECT bibcode, pdf_path, html_path, title, markdown_path, doctype FROM papers" + "SELECT bibcode, pdf_path, html_path, title, markdown_path, doctype FROM papers", ) .fetch_all(&pool) .await?; @@ -343,7 +401,9 @@ async fn main() -> Result<(), Box> { let mut html_db_msg = String::new(); let mut markdown_db_msg = String::new(); - let doctype_str = doctype_opt.unwrap_or_else(|| "article".to_string()).to_lowercase(); + let doctype_str = doctype_opt + .unwrap_or_else(|| "article".to_string()) + .to_lowercase(); let is_skip_type = doctype_str == "proposal" || doctype_str == "abstract" || doctype_str == "catalog" @@ -356,19 +416,28 @@ async fn main() -> Result<(), Box> { if is_skip_type { let mut has_skip_anomaly = false; if let Some(ref pdf_p) = pdf_path_opt { - pdf_db_msg = format!("该文献属于跳过类型 [{}]，但包含下载/报错路径记录: {}", doctype_str, pdf_p); + pdf_db_msg = format!( + "该文献属于跳过类型 [{}]，但包含下载/报错路径记录: {}", + doctype_str, pdf_p + ); need_db_fix = true; pdf_needs_fix = true; has_skip_anomaly = true; } if let Some(ref html_p) = html_path_opt { - html_db_msg = format!("该文献属于跳过类型 [{}]，但包含下载/报错路径记录: {}", doctype_str, html_p); + html_db_msg = format!( + "该文献属于跳过类型 [{}]，但包含下载/报错路径记录: {}", + doctype_str, html_p + ); need_db_fix = true; html_needs_fix = true; has_skip_anomaly = true; } if let Some(ref md_p) = markdown_path_opt { - markdown_db_msg = format!("该文献属于跳过类型 [{}]，但包含解析路径记录: {}", doctype_str, md_p); + markdown_db_msg = format!( + "该文献属于跳过类型 [{}]，但包含解析路径记录: {}", + doctype_str, md_p + ); need_db_fix = true; markdown_needs_fix = true; has_skip_anomaly = true; @@ -377,10 +446,12 @@ async fn main() -> Result<(), Box> { db_skip_type_cleaned += 1; } } else { - let has_valid_pdf = pdf_path_opt.as_ref() + let has_valid_pdf = pdf_path_opt + .as_ref() .map(|p| !p.starts_with("error:") && library_dir.join(p).exists()) .unwrap_or(false); - let has_valid_html = html_path_opt.as_ref() + let has_valid_html = html_path_opt + .as_ref() .map(|p| !p.starts_with("error:") && library_dir.join(p).exists()) .unwrap_or(false); @@ -388,7 +459,10 @@ async fn main() -> Result<(), Box> { if pdf_p.starts_with("error:") { if has_valid_html { db_pdf_err_text += 1; - pdf_db_msg = format!("文献已成功下载 HTML 格式，但 PDF 仍留有报错日志（将清理为 NULL）: {}", pdf_p); + pdf_db_msg = format!( + "文献已成功下载 HTML 格式，但 PDF 仍留有报错日志（将清理为 NULL）: {}", + pdf_p + ); need_db_fix = true; pdf_needs_fix = true; } else { @@ -407,7 +481,10 @@ async fn main() -> Result<(), Box> { if html_p.starts_with("error:") { if has_valid_pdf { db_html_err_text += 1; - html_db_msg = format!("文献已成功下载 PDF 格式，但 HTML 仍留有报错日志（将清理为 NULL）: {}", html_p); + html_db_msg = format!( + "文献已成功下载 PDF 格式，但 HTML 仍留有报错日志（将清理为 NULL）: {}", + html_p + ); need_db_fix = true; html_needs_fix = true; } else { @@ -441,20 +518,29 @@ async fn main() -> Result<(), Box> { } if need_db_fix { - println!(" ❌ 发现馆藏文献记录损坏/不一致 [{}] 《{}》", bibcode, title); + println!( + " ❌ 发现馆藏文献记录损坏/不一致 [{}] 《{}》", + bibcode, title + ); if !pdf_db_msg.is_empty() { if pdf_needs_fix { println!(" [异常] PDF 状态: {}", pdf_db_msg); } else { - println!(" [日志] PDF 历史下载失败原因: {}", pdf_db_msg.replace("数据库存储了报错字符串: ", "")); + println!( + " [日志] PDF 历史下载失败原因: {}", + pdf_db_msg.replace("数据库存储了报错字符串: ", "") + ); } } if !html_db_msg.is_empty() { if html_needs_fix { println!(" [异常] HTML 状态: {}", html_db_msg); } else { - println!(" [日志] HTML 历史下载失败原因: {}", html_db_msg.replace("数据库存储了报错字符串: ", "")); + println!( + " [日志] HTML 历史下载失败原因: {}", + html_db_msg.replace("数据库存储了报错字符串: ", "") + ); } } if !markdown_db_msg.is_empty() { @@ -502,11 +588,11 @@ async fn main() -> Result<(), Box> { } if !sql_parts.is_empty() { - let query_str = format!("UPDATE papers SET {} WHERE bibcode = ?", sql_parts.join(", ")); - let res = sqlx::query(&query_str) - .bind(&bibcode) - .execute(&pool) - .await; + let query_str = format!( + "UPDATE papers SET {} WHERE bibcode = ?", + sql_parts.join(", ") + ); + let res = sqlx::query(&query_str).bind(&bibcode).execute(&pool).await; if res.is_ok() { db_updated_count += 1; println!(" ✅ [修复] 数据库损坏字段已成功重置"); @@ -525,19 +611,40 @@ async fn main() -> Result<(), Box> { println!(" - 损坏/假 PDF 文件数: {}", disk_pdf_invalid); println!("--------------------------------------------------"); println!("数据库一致性统计:"); - println!(" - 数据库记录下载失败数 (error:): PDF: {}, HTML: {}", db_pdf_err_text, db_html_err_text); - println!(" - 磁盘文件丢失数 (数据库有记录但文件不存在): PDF: {}, HTML: {}, Markdown: {}", db_pdf_missing, db_html_missing, db_markdown_missing); + println!( + " - 数据库记录下载失败数 (error:): PDF: {}, HTML: {}", + db_pdf_err_text, db_html_err_text + ); + println!( + " - 磁盘文件丢失数 (数据库有记录但文件不存在): PDF: {}, HTML: {}, Markdown: {}", + db_pdf_missing, db_html_missing, db_markdown_missing + ); println!(" - 孤立无源 Markdown 篇数: {}", db_markdown_orphaned); - println!(" - 需跳过类型但包含下载记录篇数 (已清理/待清理): {}", db_skip_type_cleaned); + println!( + " - 需跳过类型但包含下载记录篇数 (已清理/待清理): {}", + db_skip_type_cleaned + ); println!("--------------------------------------------------"); if fix { println!("✨ 修复完成！"); - println!(" - 共删除磁盘物理损坏/孤立/跳过类型文件: {} 个", deleted_files); + println!( + " - 共删除磁盘物理损坏/孤立/跳过类型文件: {} 个", + deleted_files + ); println!(" - 共重置修复数据库文献字段: {} 处", db_updated_count); } else { - let total_issues = disk_html_invalid + disk_pdf_invalid + db_pdf_missing + db_html_missing + db_markdown_missing + db_markdown_orphaned + db_skip_type_cleaned; + let total_issues = disk_html_invalid + + disk_pdf_invalid + + db_pdf_missing + + db_html_missing + + db_markdown_missing + + db_markdown_orphaned + + db_skip_type_cleaned; if total_issues > 0 { - println!("❌ 警告：共检测出 {} 处坏文件、丢失文件或异常数据库记录。", total_issues); + println!( + "❌ 警告：共检测出 {} 处坏文件、丢失文件或异常数据库记录。", + total_issues + ); println!("👉 您可以附加 '-- --fix' 执行一键全面修复："); println!(" cargo run --bin health_check -- --fix"); } else { diff --git a/src/bin/reparse.rs b/src/bin/reparse.rs index 6f02214..98cfc7a 100644 --- a/src/bin/reparse.rs +++ b/src/bin/reparse.rs @@ -1,6 +1,6 @@ +use astroresearch::services::parser::html_to_markdown; use std::path::{Path, PathBuf}; use std::time::Instant; -use astroresearch::services::parser::html_to_markdown; fn main() -> anyhow::Result<()> { let args: Vec = std::env::args().collect(); @@ -17,16 +17,26 @@ fn main() -> anyhow::Result<()> { } let front_matter = if md_path.exists() { extract_front_matter(&std::fs::read_to_string(&md_path)?) - } else { String::new() }; + } else { + String::new() + }; let t0 = Instant::now(); let md_content = html_to_markdown(&html_path)?; let elapsed = t0.elapsed(); - let final_md = if front_matter.is_empty() { md_content } - else { format!("{}\n\n{}", front_matter.trim_end(), md_content) }; + let final_md = if front_matter.is_empty() { + md_content + } else { + format!("{}\n\n{}", front_matter.trim_end(), md_content) + }; std::fs::write(&md_path, &final_md)?; - println!("✅ {} → {} [{:.0}ms]", stem, md_path.display(), elapsed.as_secs_f64() * 1000.0); + println!( + "✅ {} → {} [{:.0}ms]", + stem, + md_path.display(), + elapsed.as_secs_f64() * 1000.0 + ); return Ok(()); } @@ -37,7 +47,7 @@ fn main() -> anyhow::Result<()> { let mut html_files: Vec = std::fs::read_dir(html_dir)? .filter_map(|e| e.ok()) .map(|e| e.path()) - .filter(|p| p.extension().map_or(false, |ext| ext == "html")) + .filter(|p| p.extension().is_some_and(|ext| ext == "html")) .collect(); html_files.sort(); @@ -45,7 +55,7 @@ fn main() -> anyhow::Result<()> { println!("Found {} HTML files to reparse\n", total); let mut success = 0u32; - let mut skipped = 0u32; + let skipped = 0u32; let mut failed = 0u32; for html_path in &html_files { diff --git a/src/clients/ads.rs b/src/clients/ads.rs index 97baa06..e87f2c7 100644 --- a/src/clients/ads.rs +++ b/src/clients/ads.rs @@ -1,7 +1,7 @@ // src/ads.rs -use serde::{Deserialize, Serialize}; use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE}; -use tracing::{info, error}; +use serde::{Deserialize, Serialize}; +use tracing::{error, info}; // 原始 ADS API 返回的数据文档结构 #[derive(Debug, Clone, Serialize, Deserialize)] @@ -58,21 +58,28 @@ impl AdsClient { let mut headers = HeaderMap::new(); headers.insert( AUTHORIZATION, - HeaderValue::from_str(&format!("Bearer {}", self.api_key)).unwrap_or_else(|_| HeaderValue::from_static("")), + HeaderValue::from_str(&format!("Bearer {}", self.api_key)) + .unwrap_or_else(|_| HeaderValue::from_static("")), ); headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json")); headers } // 调用 ADS 检索接口获取文献元数据列表，支持分页与排序 - pub async fn search(&self, query: &str, start: i32, rows: i32, sort: &str) -> anyhow::Result> { + pub async fn search( + &self, + query: &str, + start: i32, + rows: i32, + sort: &str, + ) -> anyhow::Result> { let url = "https://api.adsabs.harvard.edu/v1/search/query"; - + let translated = crate::services::query_parser::to_ads_query(query); // fl 声明返回字段，包括 reference 和 citation 引用关系数组及 identifier 和 doctype let fl = "bibcode,title,author,year,pub,keyword,abstract,doi,citation_count,reference_count,reference,citation,identifier,doctype"; - + let ads_sort = match sort { "date_desc" => "date desc", "date_asc" => "date asc", @@ -80,12 +87,16 @@ impl AdsClient { _ => "score desc", }; - info!("正在发送检索请求到 ADS 平台: 原始词='{}', 翻译词='{}', 起始={}, 数量={}, 排序='{}'", query, translated, start, rows, ads_sort); + info!( + "正在发送检索请求到 ADS 平台: 原始词='{}', 翻译词='{}', 起始={}, 数量={}, 排序='{}'", + query, translated, start, rows, ads_sort + ); let start_str = start.to_string(); let rows_str = rows.to_string(); - let response = self.client + let response = self + .client .get(url) .headers(self.headers()) .query(&[ @@ -106,8 +117,11 @@ impl AdsClient { } let raw_res: RawSearchResponse = response.json().await?; - let docs = raw_res.response.docs.into_iter().map(|d| { - AdsPaperDoc { + let docs = raw_res + .response + .docs + .into_iter() + .map(|d| AdsPaperDoc { bibcode: d.bibcode, title: d.title, author: d.author, @@ -122,8 +136,8 @@ impl AdsClient { citation: d.citation, identifier: d.identifier, doctype: d.doctype, - } - }).collect(); + }) + .collect(); Ok(docs) } @@ -131,13 +145,17 @@ impl AdsClient { // 调用 ADS Export 接口导出 BibTeX 文本内容 pub async fn export_bibtex(&self, bibcodes: Vec) -> anyhow::Result { let url = "https://api.adsabs.harvard.edu/v1/export/bibtex"; - info!("正在向 ADS 请求导出 {} 篇文献的 BibTeX 数据", bibcodes.len()); + info!( + "正在向 ADS 请求导出 {} 篇文献的 BibTeX 数据", + bibcodes.len() + ); let payload = serde_json::json!({ "bibcode": bibcodes }); - let response = self.client + let response = self + .client .post(url) .headers(self.headers()) .json(&payload) @@ -147,7 +165,10 @@ impl AdsClient { if !response.status().is_success() { let status = response.status(); let err_body = response.text().await.unwrap_or_default(); - error!("ADS 导出 BibTeX 失败: 状态码={}, 返回信息={}", status, err_body); + error!( + "ADS 导出 BibTeX 失败: 状态码={}, 返回信息={}", + status, err_body + ); return Err(anyhow::anyhow!("ADS 导出接口返回错误码: {}", status)); } @@ -160,8 +181,12 @@ impl AdsClient { let url = "https://api.adsabs.harvard.edu/v1/search/query"; let translated = crate::services::query_parser::to_ads_query(query); - info!("正在向 ADS 查询匹配的总文献数, 原始词: '{}', 翻译词: '{}'", query, translated); - let response = self.client + info!( + "正在向 ADS 查询匹配的总文献数, 原始词: '{}', 翻译词: '{}'", + query, translated + ); + let response = self + .client .get(url) .headers(self.headers()) .query(&[("q", translated.as_str()), ("rows", "0")]) @@ -256,7 +281,10 @@ mod tests { println!("arXiv 平台:"); println!(" - (OR) \"{}\" 匹配数: {} 篇", query_or, count_arxiv_or); println!(" - (AND) \"{}\" 匹配数: {} 篇", query_and, count_arxiv_and); - assert!(count_arxiv_or > count_arxiv_and, "错误: arXiv 的 OR 结果应该多于 AND"); + assert!( + count_arxiv_or > count_arxiv_and, + "错误: arXiv 的 OR 结果应该多于 AND" + ); // 测试 2: 比较基础词组与含有 NOT 排除条件的数据量差异 let query_base = "\"hot subdwarf\""; @@ -266,17 +294,32 @@ mod tests { let count_ads_base = ads.get_total_count(query_base).await?; let count_ads_not = ads.get_total_count(query_not).await?; println!("NASA ADS 平台:"); - println!(" - (基础) \"{}\" 匹配数: {} 篇", query_base, count_ads_base); + println!( + " - (基础) \"{}\" 匹配数: {} 篇", + query_base, count_ads_base + ); println!(" - (排除) \"{}\" 匹配数: {} 篇", query_not, count_ads_not); - assert!(count_ads_base >= count_ads_not, "错误: 基础结果应该大于或等于排除后的结果"); + assert!( + count_ads_base >= count_ads_not, + "错误: 基础结果应该大于或等于排除后的结果" + ); } let count_arxiv_base = arxiv.get_total_count(query_base).await?; let count_arxiv_not = arxiv.get_total_count(query_not).await?; println!("arXiv 平台:"); - println!(" - (基础) \"{}\" 匹配数: {} 篇", query_base, count_arxiv_base); - println!(" - (排除) \"{}\" 匹配数: {} 篇", query_not, count_arxiv_not); - assert!(count_arxiv_base >= count_arxiv_not, "错误: arXiv 基础结果应该大于或等于排除后的结果"); + println!( + " - (基础) \"{}\" 匹配数: {} 篇", + query_base, count_arxiv_base + ); + println!( + " - (排除) \"{}\" 匹配数: {} 篇", + query_not, count_arxiv_not + ); + assert!( + count_arxiv_base >= count_arxiv_not, + "错误: arXiv 基础结果应该大于或等于排除后的结果" + ); println!("================= 真实检索逻辑集成测试全部通过 ================="); Ok(()) diff --git a/src/clients/arxiv.rs b/src/clients/arxiv.rs index 757fda5..04ccdcd 100644 --- a/src/clients/arxiv.rs +++ b/src/clients/arxiv.rs @@ -1,12 +1,12 @@ // src/arxiv.rs -use serde::{Deserialize, Serialize}; -use tracing::{info, error}; use regex::Regex; +use serde::{Deserialize, Serialize}; +use tracing::{error, info}; // 统一的 arXiv 文献临时结构 #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ArxivPaper { - pub id: String, // 清洗后的 arXiv ID，例如 2301.00001 + pub id: String, // 清洗后的 arXiv ID，例如 2301.00001 pub title: String, pub authors: Vec, pub year: String, @@ -21,6 +21,12 @@ pub struct ArxivClient { client: reqwest::Client, } +impl Default for ArxivClient { + fn default() -> Self { + Self::new() + } +} + impl ArxivClient { pub fn new() -> Self { ArxivClient { @@ -29,15 +35,24 @@ impl ArxivClient { } // 请求 arXiv 官方的 Export 检索接口并解析返回内容，支持分页与排序 - pub async fn search(&self, query: &str, start: i32, max_results: i32, sort: &str) -> anyhow::Result> { + pub async fn search( + &self, + query: &str, + start: i32, + max_results: i32, + sort: &str, + ) -> anyhow::Result> { let url = "http://export.arxiv.org/api/query"; - + let (translated_query, year_range) = crate::services::query_parser::to_arxiv_query(query); // 如果包含年份过滤，我们可以在 search_query 里追加年份限制，格式如: AND (submittedDate:[YYYY01010000 TO YYYY12312359]) let mut final_query = translated_query; if let Some((start_yr, end_yr)) = year_range { - final_query = format!("({}) AND submittedDate:[{}01010000 TO {}12312359]", final_query, start_yr, end_yr); + final_query = format!( + "({}) AND submittedDate:[{}01010000 TO {}12312359]", + final_query, start_yr, end_yr + ); } let (sort_by, sort_order) = match sort { @@ -51,7 +66,8 @@ impl ArxivClient { let start_str = start.to_string(); let max_results_str = max_results.to_string(); - let response = self.client + let response = self + .client .get(url) .query(&[ ("search_query", final_query.as_str()), @@ -81,16 +97,20 @@ impl ArxivClient { let mut final_query = translated_query; if let Some((start_yr, end_yr)) = year_range { - final_query = format!("({}) AND submittedDate:[{}01010000 TO {}12312359]", final_query, start_yr, end_yr); + final_query = format!( + "({}) AND submittedDate:[{}01010000 TO {}12312359]", + final_query, start_yr, end_yr + ); } - info!("正在向 arXiv 查询匹配的总文献数, 原始词: '{}', 翻译词: '{}'", query, final_query); - let response = self.client + info!( + "正在向 arXiv 查询匹配的总文献数, 原始词: '{}', 翻译词: '{}'", + query, final_query + ); + let response = self + .client .get(url) - .query(&[ - ("search_query", final_query.as_str()), - ("max_results", "1"), - ]) + .query(&[("search_query", final_query.as_str()), ("max_results", "1")]) .send() .await?; @@ -100,7 +120,8 @@ impl ArxivClient { } let xml_content = response.text().await?; - let total_re = Regex::new(r"]*>(\d+)").unwrap(); + let total_re = + Regex::new(r"]*>(\d+)").unwrap(); if let Some(caps) = total_re.captures(&xml_content) { if let Ok(count) = caps[1].parse::() { return Ok(count); @@ -113,7 +134,7 @@ impl ArxivClient { // 使用正则表达式手动提取 XML 内容，避免由于命名空间前缀不同造成的反序列化问题 fn parse_arxiv_xml(xml: &str) -> Vec { let mut papers = Vec::new(); - + let entry_re = Regex::new(r"(?s)(.*?)").unwrap(); let id_re = Regex::new(r"http://arxiv.org/abs/(.*?)(?:v\d+)?").unwrap(); let title_re = Regex::new(r"(?s)(.*?)").unwrap(); @@ -122,15 +143,18 @@ fn parse_arxiv_xml(xml: &str) -> Vec { let author_re = Regex::new(r"(?s)\s*(.*?)").unwrap(); let doi_re = Regex::new(r"]*>(.*?)").unwrap(); + let fallback_id_re = Regex::new(r"(.*?)").unwrap(); + for cap in entry_re.captures_iter(xml) { let entry_content = &cap[1]; // 提取并清洗 ID - let id = id_re.captures(entry_content) + let id = id_re + .captures(entry_content) .map(|c| c[1].trim().to_string()) .unwrap_or_else(|| { - let fallback_id_re = Regex::new(r"(.*?)").unwrap(); - fallback_id_re.captures(entry_content) + fallback_id_re + .captures(entry_content) .map(|c| c[1].trim().to_string()) .unwrap_or_default() }); @@ -140,19 +164,30 @@ fn parse_arxiv_xml(xml: &str) -> Vec { } // 提取标题，清理换行与连续空格 - let mut title = title_re.captures(entry_content) + let mut title = title_re + .captures(entry_content) .map(|c| c[1].to_string()) .unwrap_or_default(); - title = title.replace('\n', " ").replace(" ", " ").trim().to_string(); + title = title + .replace('\n', " ") + .replace(" ", " ") + .trim() + .to_string(); // 提取摘要 - let mut abstract_text = summary_re.captures(entry_content) + let mut abstract_text = summary_re + .captures(entry_content) .map(|c| c[1].to_string()) .unwrap_or_default(); - abstract_text = abstract_text.replace('\n', " ").replace(" ", " ").trim().to_string(); + abstract_text = abstract_text + .replace('\n', " ") + .replace(" ", " ") + .trim() + .to_string(); // 提取发布年份 - let year = published_re.captures(entry_content) + let year = published_re + .captures(entry_content) .map(|c| c[1].to_string()) .unwrap_or_else(|| "未知".to_string()); @@ -166,7 +201,8 @@ fn parse_arxiv_xml(xml: &str) -> Vec { } // 提取关联 DOI - let doi = doi_re.captures(entry_content) + let doi = doi_re + .captures(entry_content) .map(|c| c[1].trim().to_string()); let pdf_url = format!("https://arxiv.org/pdf/{}.pdf", id); @@ -213,11 +249,16 @@ mod tests { let paper = &papers[0]; assert_eq!(paper.id, "2301.00001"); assert_eq!(paper.title, "A Beautiful Title of Astro Research Paper"); - assert_eq!(paper.authors, vec!["John Doe".to_string(), "Jane Smith".to_string()]); + assert_eq!( + paper.authors, + vec!["John Doe".to_string(), "Jane Smith".to_string()] + ); assert_eq!(paper.year, "2023"); - assert_eq!(paper.abstract_text, "This is the abstract. It spans multiple lines."); + assert_eq!( + paper.abstract_text, + "This is the abstract. It spans multiple lines." + ); assert_eq!(paper.doi, Some("10.1000/xyz123".to_string())); assert_eq!(paper.pdf_url, "https://arxiv.org/pdf/2301.00001.pdf"); } } - diff --git a/src/clients/llm.rs b/src/clients/llm.rs index 218f0fd..2f52915 100644 --- a/src/clients/llm.rs +++ b/src/clients/llm.rs @@ -1,8 +1,8 @@ // src/clients/llm.rs -use serde::{Deserialize, Serialize}; -use reqwest::Client; -use tracing::error; use futures_util::StreamExt; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use tracing::error; /// 消息角色枚举（OpenAI 兼容） #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] @@ -96,7 +96,11 @@ impl ChatMessage { } /// 创建带工具调用和推理内容的助手消息 - pub fn assistant_with_reasoning(content: Option, reasoning_content: Option, tool_calls: Option>) -> Self { + pub fn assistant_with_reasoning( + content: Option, + reasoning_content: Option, + tool_calls: Option>, + ) -> Self { ChatMessage { role: MessageRole::Assistant, content, @@ -137,7 +141,11 @@ pub struct FunctionDef { } impl ToolDefinition { - pub fn new(name: impl Into, description: impl Into, parameters: serde_json::Value) -> Self { + pub fn new( + name: impl Into, + description: impl Into, + parameters: serde_json::Value, + ) -> Self { ToolDefinition { tool_type: "function".to_string(), function: FunctionDef { @@ -224,9 +232,13 @@ impl LlmClient { &self.api_key } - pub async fn chat_completion(&self, system_prompt: &str, user_content: &str) -> anyhow::Result { + pub async fn chat_completion( + &self, + system_prompt: &str, + user_content: &str, + ) -> anyhow::Result { let url = format!("{}/chat/completions", self.api_base); - + let payload = serde_json::json!({ "model": self.model, "messages": [ @@ -242,7 +254,9 @@ impl LlmClient { "temperature": 0.3 }); - let response = self.client.post(&url) + let response = self + .client + .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .header("Content-Type", "application/json") .json(&payload) @@ -320,7 +334,9 @@ impl LlmClient { "temperature": 0.3 }); - let response = self.client.post(&url) + let response = self + .client + .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .header("Content-Type", "application/json") .json(&payload) @@ -358,7 +374,11 @@ impl LlmClient { } /// 多轮对话补全（含原生 Tool Calling 支持），非流式 - pub async fn chat(&self, messages: &[ChatMessage], tools: &[ToolDefinition]) -> anyhow::Result { + pub async fn chat( + &self, + messages: &[ChatMessage], + tools: &[ToolDefinition], + ) -> anyhow::Result { let url = format!("{}/chat/completions", self.api_base); let mut payload = serde_json::json!({ @@ -368,7 +388,9 @@ impl LlmClient { }); // 如果是通义千问 (DashScope) 或者是 Qwen 模型，自动开启思考模式 - if self.api_base.contains("dashscope.aliyuncs.com") || self.model.to_lowercase().contains("qwen") { + if self.api_base.contains("dashscope.aliyuncs.com") + || self.model.to_lowercase().contains("qwen") + { if let Some(obj) = payload.as_object_mut() { obj.insert("enable_thinking".to_string(), serde_json::json!(true)); } @@ -378,7 +400,9 @@ impl LlmClient { payload["tools"] = serde_json::to_value(tools)?; } - let response = self.client.post(&url) + let response = self + .client + .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .header("Content-Type", "application/json") .json(&payload) @@ -389,7 +413,11 @@ impl LlmClient { let status = response.status(); let body = response.text().await.unwrap_or_default(); error!("LLM chat 接口调用失败: 状态码={}, 报错={}", status, body); - return Err(anyhow::anyhow!("大模型 chat 接口返回错误状态: {} - {}", status, body)); + return Err(anyhow::anyhow!( + "大模型 chat 接口返回错误状态: {} - {}", + status, + body + )); } #[derive(Deserialize)] @@ -440,7 +468,9 @@ impl LlmClient { }); // 如果是通义千问 (DashScope) 或者是 Qwen 模型，自动开启思考模式 - if self.api_base.contains("dashscope.aliyuncs.com") || self.model.to_lowercase().contains("qwen") { + if self.api_base.contains("dashscope.aliyuncs.com") + || self.model.to_lowercase().contains("qwen") + { if let Some(obj) = payload.as_object_mut() { obj.insert("enable_thinking".to_string(), serde_json::json!(true)); } @@ -450,7 +480,9 @@ impl LlmClient { payload["tools"] = serde_json::to_value(tools)?; } - let response = self.client.post(&url) + let response = self + .client + .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .header("Content-Type", "application/json") .json(&payload) @@ -459,17 +491,31 @@ impl LlmClient { if !response.status().is_success() { let status = response.status(); + let status_code = status.as_u16(); + let retry_after = response + .headers() + .get("retry-after") + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.parse::().ok()); let body = response.text().await.unwrap_or_default(); error!("LLM stream 接口调用失败: 状态码={}, 报错={}", status, body); - return Err(anyhow::anyhow!("大模型 stream 接口返回错误状态: {} - {}", status, body)); + return Err(anyhow::anyhow!( + "HTTP {}: {} | retry_after={:?}", + status_code, + body, + retry_after + )); } let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); let mut byte_stream = response.bytes_stream(); - + tokio::spawn(async move { // 流式 Tool Call 碎片累积器 - let mut tool_call_accumulators: std::collections::HashMap = std::collections::HashMap::new(); + let mut tool_call_accumulators: std::collections::HashMap< + usize, + (String, String, String), + > = std::collections::HashMap::new(); let mut buffer = String::new(); while let Some(chunk_result) = byte_stream.next().await { @@ -496,16 +542,24 @@ impl LlmClient { if data == "[DONE]" { // 流结束前，将累积的 tool calls 还原并发送 if !tool_call_accumulators.is_empty() { - let mut indices: Vec = tool_call_accumulators.keys().cloned().collect(); + let mut indices: Vec = + tool_call_accumulators.keys().cloned().collect(); indices.sort(); - let tool_calls: Vec = indices.into_iter().map(|idx| { - let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap(); - ToolCall { - id, - call_type: "function".to_string(), - function: FunctionCall { name, arguments: args }, - } - }).collect(); + let tool_calls: Vec = indices + .into_iter() + .map(|idx| { + let (id, name, args) = + tool_call_accumulators.remove(&idx).unwrap(); + ToolCall { + id, + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: args, + }, + } + }) + .collect(); let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls)); } let _ = tx.send(StreamEvent::Done); @@ -521,7 +575,9 @@ impl LlmClient { // 提取 usage（在最后一条 chunk 中 stream_options 返回） if let Some(usage_val) = parsed.get("usage") { if !usage_val.is_null() { - if let Ok(usage) = serde_json::from_value::(usage_val.clone()) { + if let Ok(usage) = + serde_json::from_value::(usage_val.clone()) + { let _ = tx.send(StreamEvent::Usage(usage)); } } @@ -539,23 +595,42 @@ impl LlmClient { } // 推理内容增量 - if let Some(reasoning) = delta.get("reasoning_content").and_then(|c| c.as_str()) { + if let Some(reasoning) = + delta.get("reasoning_content").and_then(|c| c.as_str()) + { if !reasoning.is_empty() { - let _ = tx.send(StreamEvent::ReasoningDelta(reasoning.to_string())); + let _ = + tx.send(StreamEvent::ReasoningDelta(reasoning.to_string())); } } // 工具调用增量 - if let Some(tool_calls) = delta.get("tool_calls").and_then(|t| t.as_array()) { + if let Some(tool_calls) = + delta.get("tool_calls").and_then(|t| t.as_array()) + { for tc in tool_calls { - let index = tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize; - let id = tc.get("id").and_then(|i| i.as_str()).map(|s| s.to_string()); - let fn_name = tc.get("function").and_then(|f| f.get("name")).and_then(|n| n.as_str()).map(|s| s.to_string()); - let fn_args = tc.get("function").and_then(|f| f.get("arguments")).and_then(|a| a.as_str()).unwrap_or(""); + let index = + tc.get("index").and_then(|i| i.as_u64()).unwrap_or(0) + as usize; + let id = tc + .get("id") + .and_then(|i| i.as_str()) + .map(|s| s.to_string()); + let fn_name = tc + .get("function") + .and_then(|f| f.get("name")) + .and_then(|n| n.as_str()) + .map(|s| s.to_string()); + let fn_args = tc + .get("function") + .and_then(|f| f.get("arguments")) + .and_then(|a| a.as_str()) + .unwrap_or(""); - let entry = tool_call_accumulators.entry(index).or_insert_with(|| ( - String::new(), String::new(), String::new() - )); + let entry = + tool_call_accumulators.entry(index).or_insert_with(|| { + (String::new(), String::new(), String::new()) + }); if let Some(ref id_str) = id { entry.0 = id_str.clone(); } @@ -574,18 +649,30 @@ impl LlmClient { } // 检查是否 finish_reason == "tool_calls" - if let Some(finish_reason) = choice.get("finish_reason").and_then(|f| f.as_str()) { - if finish_reason == "tool_calls" && !tool_call_accumulators.is_empty() { - let mut indices: Vec = tool_call_accumulators.keys().cloned().collect(); + if let Some(finish_reason) = + choice.get("finish_reason").and_then(|f| f.as_str()) + { + if finish_reason == "tool_calls" + && !tool_call_accumulators.is_empty() + { + let mut indices: Vec = + tool_call_accumulators.keys().cloned().collect(); indices.sort(); - let tool_calls: Vec = indices.into_iter().map(|idx| { - let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap(); - ToolCall { - id, - call_type: "function".to_string(), - function: FunctionCall { name, arguments: args }, - } - }).collect(); + let tool_calls: Vec = indices + .into_iter() + .map(|idx| { + let (id, name, args) = + tool_call_accumulators.remove(&idx).unwrap(); + ToolCall { + id, + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: args, + }, + } + }) + .collect(); let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls)); } } @@ -598,14 +685,20 @@ impl LlmClient { if !tool_call_accumulators.is_empty() { let mut indices: Vec = tool_call_accumulators.keys().cloned().collect(); indices.sort(); - let tool_calls: Vec = indices.into_iter().map(|idx| { - let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap(); - ToolCall { - id, - call_type: "function".to_string(), - function: FunctionCall { name, arguments: args }, - } - }).collect(); + let tool_calls: Vec = indices + .into_iter() + .map(|idx| { + let (id, name, args) = tool_call_accumulators.remove(&idx).unwrap(); + ToolCall { + id, + call_type: "function".to_string(), + function: FunctionCall { + name, + arguments: args, + }, + } + }) + .collect(); let _ = tx.send(StreamEvent::ToolCallsComplete(tool_calls)); } let _ = tx.send(StreamEvent::Done); @@ -647,13 +740,15 @@ impl EmbeddingClient { pub async fn create_embedding(&self, text: &str) -> anyhow::Result> { let url = format!("{}/embeddings", self.api_base); - + let payload = serde_json::json!({ "model": self.model, "input": text, }); - let response = self.client.post(&url) + let response = self + .client + .post(&url) .header("Authorization", format!("Bearer {}", self.api_key)) .header("Content-Type", "application/json") .json(&payload) @@ -692,11 +787,7 @@ mod tests { #[test] fn test_llm_client_initialization() { - let client = LlmClient::new( - "key".to_string(), - "base".to_string(), - "model".to_string(), - ); + let client = LlmClient::new("key".to_string(), "base".to_string(), "model".to_string()); assert_eq!(client.api_key(), "key"); assert_eq!(client.api_base(), "base"); assert_eq!(client.model(), "model"); @@ -704,11 +795,8 @@ mod tests { #[test] fn test_embedding_client_initialization() { - let client = EmbeddingClient::new( - "key".to_string(), - "base".to_string(), - "model".to_string(), - ); + let client = + EmbeddingClient::new("key".to_string(), "base".to_string(), "model".to_string()); assert_eq!(client.api_key(), "key"); assert_eq!(client.api_base(), "base"); assert_eq!(client.model(), "model"); @@ -730,7 +818,10 @@ mod tests { // 测试助手消息构造 let assistant = ChatMessage::assistant("你好！有什么可以帮助你的吗？"); assert_eq!(assistant.role, MessageRole::Assistant); - assert_eq!(assistant.content.as_deref(), Some("你好！有什么可以帮助你的吗？")); + assert_eq!( + assistant.content.as_deref(), + Some("你好！有什么可以帮助你的吗？") + ); // 测试工具结果消息构造 let tool_result = ChatMessage::tool_result("call_123", r#"{"result": 42}"#); @@ -751,7 +842,10 @@ mod tests { assert_eq!(assistant_tc.role, MessageRole::Assistant); assert!(assistant_tc.content.is_none()); assert_eq!(assistant_tc.tool_calls.as_ref().unwrap().len(), 1); - assert_eq!(assistant_tc.tool_calls.as_ref().unwrap()[0].function.name, "get_weather"); + assert_eq!( + assistant_tc.tool_calls.as_ref().unwrap()[0].function.name, + "get_weather" + ); } #[test] @@ -807,7 +901,10 @@ mod tests { let deserialized: ChatMessage = serde_json::from_value(json).unwrap(); assert_eq!(deserialized.role, MessageRole::Assistant); assert_eq!(deserialized.content.as_deref(), Some("回答内容")); - assert_eq!(deserialized.reasoning_content.as_deref(), Some("这是思考过程")); + assert_eq!( + deserialized.reasoning_content.as_deref(), + Some("这是思考过程") + ); } #[test] @@ -837,13 +934,19 @@ mod tests { if config.llm_api_key.is_empty() { println!("警告: 未在环境配置中检测到 LLM_API_KEY，跳过 LlmClient 集成测试。"); } else { - println!("测试大模型: {} (API Base: {})", config.llm_model, config.llm_api_base); + println!( + "测试大模型: {} (API Base: {})", + config.llm_model, config.llm_api_base + ); let llm = LlmClient::new( config.llm_api_key.clone(), config.llm_api_base.clone(), config.llm_model.clone(), ); - match llm.chat_completion("You are a helpful assistant.", "Say Hello!").await { + match llm + .chat_completion("You are a helpful assistant.", "Say Hello!") + .await + { Ok(reply) => { println!("LlmClient 响应成功: {}", reply.trim()); assert!(!reply.trim().is_empty(), "错误: 大模型返回了空响应"); @@ -854,9 +957,14 @@ mod tests { // 2. 测试 EmbeddingClient if config.embedding_api_key.is_empty() { - println!("警告: 未在环境配置中检测到 EMBEDDING_API_KEY，跳过 EmbeddingClient 集成测试。"); + println!( + "警告: 未在环境配置中检测到 EMBEDDING_API_KEY，跳过 EmbeddingClient 集成测试。" + ); } else { - println!("测试向量模型: {} (API Base: {})", config.embedding_model, config.embedding_api_base); + println!( + "测试向量模型: {} (API Base: {})", + config.embedding_model, config.embedding_api_base + ); let embedding_client = EmbeddingClient::new( config.embedding_api_key.clone(), config.embedding_api_base.clone(), @@ -868,7 +976,11 @@ mod tests { println!("EmbeddingClient 响应成功！向量维度: {}", vector.len()); assert!(!vector.is_empty(), "错误: 向量数据为空"); let preview_len = std::cmp::min(5, vector.len()); - println!("前 {} 个向量数值样例: {:?}", preview_len, &vector[..preview_len]); + println!( + "前 {} 个向量数值样例: {:?}", + preview_len, + &vector[..preview_len] + ); } Err(e) => panic!("EmbeddingClient 接口调用失败: {}", e), } diff --git a/src/clients/mod.rs b/src/clients/mod.rs index 58c3417..7818000 100644 --- a/src/clients/mod.rs +++ b/src/clients/mod.rs @@ -1,4 +1,4 @@ pub mod ads; pub mod arxiv; -pub mod qiniu; pub mod llm; +pub mod qiniu; diff --git a/src/clients/qiniu.rs b/src/clients/qiniu.rs index b12e9ac..5639c31 100644 --- a/src/clients/qiniu.rs +++ b/src/clients/qiniu.rs @@ -1,8 +1,8 @@ -use sha1::Sha1; +use base64::{engine::general_purpose::URL_SAFE, Engine as _}; use hmac::{Hmac, Mac}; -use base64::{Engine as _, engine::general_purpose::URL_SAFE}; use reqwest::multipart; -use tracing::{info, error}; +use sha1::Sha1; +use tracing::{error, info}; type HmacSha1 = Hmac; @@ -36,7 +36,7 @@ impl QiniuClient { fn generate_upload_token(&self, key: &str) -> String { // 设置 1 小时过期 let deadline = chrono::Utc::now().timestamp() + 3600; - + let policy = serde_json::json!({ "scope": format!("{}:{}", self.bucket, key), "deadline": deadline @@ -45,15 +45,18 @@ impl QiniuClient { let policy_str = policy.to_string(); let encoded_policy = URL_SAFE.encode(policy_str.as_bytes()); - let mut mac = HmacSha1::new_from_slice(self.secret_key.as_bytes()) - .expect("HMAC 密钥可接收任意大小"); + let mut mac = + HmacSha1::new_from_slice(self.secret_key.as_bytes()).expect("HMAC 密钥可接收任意大小"); mac.update(encoded_policy.as_bytes()); let result = mac.finalize(); let signature = result.into_bytes(); - let encoded_signature = URL_SAFE.encode(&signature); + let encoded_signature = URL_SAFE.encode(signature); - format!("{}:{}:{}", self.access_key, encoded_signature, encoded_policy) + format!( + "{}:{}:{}", + self.access_key, encoded_signature, encoded_policy + ) } // 上传图片等字节流数据到七牛云，返回 CDN 加速外链 URL @@ -71,14 +74,14 @@ impl QiniuClient { let form = multipart::Form::new() .text("token", token) .text("key", key.clone()) - .part("file", multipart::Part::bytes(buffer).file_name(filename.to_string())); + .part( + "file", + multipart::Part::bytes(buffer).file_name(filename.to_string()), + ); let upload_url = "https://up-z1.qiniup.com"; - let response = self.client.post(upload_url) - .multipart(form) - .send() - .await?; + let response = self.client.post(upload_url).multipart(form).send().await?; if !response.status().is_success() { let status = response.status(); @@ -135,7 +138,7 @@ mod tests { ); let token = client.generate_upload_token("test_key.png"); assert!(token.starts_with("test_ak:")); - + let parts: Vec<&str> = token.split(':').collect(); assert_eq!(parts.len(), 3); assert_eq!(parts[0], "test_ak"); @@ -143,4 +146,3 @@ mod tests { assert!(!parts[2].is_empty()); } } - diff --git a/src/lib.rs b/src/lib.rs index 2df093c..23a0327 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -5,22 +5,23 @@ use std::path::PathBuf; // 系统配置结构体，加载并管理从环境变量或 .env 文件读取的参数 #[derive(Clone, Debug)] pub struct Config { - pub database_url: String, // SQLite 数据库连接 URL - pub ads_api_key: String, // NASA ADS API 访问 Token - pub llm_api_key: String, // 大语言模型 API Key - pub llm_api_base: String, // 大语言模型 API 基础地址 - pub llm_model: String, // 调用的翻译大模型名称 - pub embedding_api_key: String, // 向量模型 API Key - pub embedding_api_base: String,// 向量模型 API 基础地址 - pub embedding_model: String, // 向量模型名称 - pub qiniu_ak: String, // 七牛云 Access Key - pub qiniu_sk: String, // 七牛云 Secret Key - pub qiniu_bucket: String, // 七牛云存储空间名 (Bucket) - pub qiniu_domain: String, // 七牛云外链 CDN 域名 - pub mineru_api_url: String, // MinerU PDF 解析远程 API 地址 - pub mineru_api_key: String, // MinerU API Token - pub library_dir: PathBuf, // 本地文献馆藏根目录 - pub port: u16, // 后端服务监听端口 + pub database_url: String, // SQLite 数据库连接 URL + pub ads_api_key: String, // NASA ADS API 访问 Token + pub llm_api_key: String, // 大语言模型 API Key + pub llm_api_base: String, // 大语言模型 API 基础地址 + pub llm_model: String, // 调用的翻译大模型名称 + pub embedding_api_key: String, // 向量模型 API Key + pub embedding_api_base: String, // 向量模型 API 基础地址 + pub embedding_model: String, // 向量模型名称 + pub qiniu_ak: String, // 七牛云 Access Key + pub qiniu_sk: String, // 七牛云 Secret Key + pub qiniu_bucket: String, // 七牛云存储空间名 (Bucket) + pub qiniu_domain: String, // 七牛云外链 CDN 域名 + pub mineru_api_url: String, // MinerU PDF 解析远程 API 地址 + pub mineru_api_key: String, // MinerU API Token + pub library_dir: PathBuf, // 本地文献馆藏根目录 + pub skills_dir: PathBuf, // Agent Skills 目录（Markdown 知识模块） + pub port: u16, // 后端服务监听端口 } impl Config { @@ -32,29 +33,31 @@ impl Config { .unwrap_or_else(|_| "sqlite://library/astro_research.db".to_string()); let ads_api_key = env::var("ADS_API_KEY").unwrap_or_default(); let llm_api_key = env::var("LLM_API_KEY").unwrap_or_default(); - let llm_api_base = env::var("LLM_API_BASE") - .unwrap_or_else(|_| "https://api.openai.com/v1".to_string()); - let llm_model = env::var("LLM_MODEL") - .unwrap_or_else(|_| "gpt-4o-mini".to_string()); + let llm_api_base = + env::var("LLM_API_BASE").unwrap_or_else(|_| "https://api.openai.com/v1".to_string()); + let llm_model = env::var("LLM_MODEL").unwrap_or_else(|_| "gpt-4o-mini".to_string()); + + let embedding_api_key = + env::var("EMBEDDING_API_KEY").unwrap_or_else(|_| llm_api_key.clone()); + let embedding_api_base = + env::var("EMBEDDING_API_BASE").unwrap_or_else(|_| llm_api_base.clone()); + let embedding_model = + env::var("EMBEDDING_MODEL").unwrap_or_else(|_| "text-embedding-3-small".to_string()); - let embedding_api_key = env::var("EMBEDDING_API_KEY") - .unwrap_or_else(|_| llm_api_key.clone()); - let embedding_api_base = env::var("EMBEDDING_API_BASE") - .unwrap_or_else(|_| llm_api_base.clone()); - let embedding_model = env::var("EMBEDDING_MODEL") - .unwrap_or_else(|_| "text-embedding-3-small".to_string()); - let qiniu_ak = env::var("QINIU_AK").unwrap_or_default(); let qiniu_sk = env::var("QINIU_SK").unwrap_or_default(); let qiniu_bucket = env::var("QINIU_BUCKET").unwrap_or_default(); let qiniu_domain = env::var("QINIU_DOMAIN").unwrap_or_default(); - + let mineru_api_url = env::var("MINERU_API_URL").unwrap_or_default(); let mineru_api_key = env::var("MINERU_API_KEY").unwrap_or_default(); - + let library_dir_str = env::var("LIBRARY_DIR").unwrap_or_else(|_| "./library".to_string()); let library_dir = PathBuf::from(library_dir_str); + let skills_dir_str = env::var("SKILLS_DIR").unwrap_or_else(|_| "./skills".to_string()); + let skills_dir = PathBuf::from(skills_dir_str); + let port = env::var("PORT") .unwrap_or_else(|_| "8000".to_string()) .parse::() @@ -76,15 +79,16 @@ impl Config { mineru_api_url, mineru_api_key, library_dir, + skills_dir, port, } } } +pub mod agent; pub mod api; pub mod clients; pub mod services; -pub mod agent; #[cfg(test)] mod config_tests { @@ -94,10 +98,10 @@ mod config_tests { fn test_config_from_env() { let orig_port = std::env::var("PORT").ok(); let orig_db = std::env::var("DATABASE_URL").ok(); - + std::env::set_var("PORT", "9999"); std::env::set_var("DATABASE_URL", "sqlite://test.db"); - + let config = Config::from_env(); assert_eq!(config.port, 9999); assert_eq!(config.database_url, "sqlite://test.db"); diff --git a/src/main.rs b/src/main.rs index 5819882..ce2830d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,25 +1,28 @@ // src/main.rs -use std::net::SocketAddr; -use std::str::FromStr; -use std::sync::Arc; +use anyhow::Context; use axum::{ routing::{get, post}, Router, }; +use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; +use std::collections::HashMap; +use std::net::SocketAddr; +use std::str::FromStr; +use std::sync::{Arc, Mutex, RwLock}; use tower_http::cors::{Any, CorsLayer}; use tower_http::services::ServeDir; -use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions}; -use tracing::{info, error}; +use tracing::{error, info, warn}; -use astroresearch::Config; -use astroresearch::services::translation::Dictionary; -use astroresearch::clients::qiniu::QiniuClient; +use astroresearch::agent::skills::SkillRegistry; +use astroresearch::api::handlers::{self, AppState}; use astroresearch::clients::ads::AdsClient; use astroresearch::clients::arxiv::ArxivClient; -use astroresearch::clients::llm::{LlmClient, EmbeddingClient}; +use astroresearch::clients::llm::{EmbeddingClient, LlmClient}; +use astroresearch::clients::qiniu::QiniuClient; use astroresearch::services::download::Downloader; -use astroresearch::api::handlers::{AppState, self}; +use astroresearch::services::translation::Dictionary; +use astroresearch::Config; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -35,15 +38,25 @@ async fn main() -> anyhow::Result<()> { // 该注册必须在任何数据库连接开启之前执行，保证所有 Connection // 自动拥有 vec0 虚拟表能力。 unsafe { - libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute( - sqlite_vec::sqlite3_vec_init as *const (), + libsqlite3_sys::sqlite3_auto_extension(Some(std::mem::transmute::< + *const (), + unsafe extern "C" fn( + *mut libsqlite3_sys::sqlite3, + *mut *const i8, + *const libsqlite3_sys::sqlite3_api_routines, + ) -> i32, + >( + sqlite_vec::sqlite3_vec_init as *const () ))); } info!("sqlite-vec 自动扩展注册完成。"); // 2. 加载环境变量配置 let config = Config::from_env(); - info!("系统配置成功载入。本地 SQLite 连接串: {}", config.database_url); + info!( + "系统配置成功载入。本地 SQLite 连接串: {}", + config.database_url + ); // 创建本地馆藏物理文件夹分类结构 std::fs::create_dir_all(&config.library_dir).unwrap_or_default(); @@ -51,6 +64,8 @@ async fn main() -> anyhow::Result<()> { std::fs::create_dir_all(config.library_dir.join("HTML")).unwrap_or_default(); std::fs::create_dir_all(config.library_dir.join("Markdown")).unwrap_or_default(); std::fs::create_dir_all(config.library_dir.join("Translation")).unwrap_or_default(); + // Agent Skills 目录 + std::fs::create_dir_all(&config.skills_dir).unwrap_or_default(); // 3. 初始化本地 SQLite 数据库连接池（开启外键约束） let options = SqliteConnectOptions::from_str(&config.database_url)? @@ -66,9 +81,7 @@ async fn main() -> anyhow::Result<()> { // 4. 自动执行数据库迁移脚本 info!("开始执行 SQL 表结构迁移..."); - sqlx::migrate!("./migrations") - .run(&pool) - .await?; + sqlx::migrate!("./migrations").run(&pool).await?; info!("数据库迁移执行完成，主表准备就绪。"); // 4.5 动态创建 vec0 向量虚拟表（维度可由环境变量 EMBEDDING_DIM 控制） @@ -79,7 +92,7 @@ async fn main() -> anyhow::Result<()> { // 检测并自愈：如果已存在的 vec_paper_chunks 维度与当前配置不一致，则重建该虚拟表并清空切片内容 let existing_sql: Option<(String,)> = sqlx::query_as( - "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'vec_paper_chunks'" + "SELECT sql FROM sqlite_master WHERE type = 'table' AND name = 'vec_paper_chunks'", ) .fetch_optional(&pool) .await?; @@ -87,9 +100,16 @@ async fn main() -> anyhow::Result<()> { if let Some((sql,)) = existing_sql { let expected_pattern = format!("float[{}]", embedding_dim); if !sql.contains(&expected_pattern) { - info!("检测到已存在的向量表维度不匹配，正在重建以适配当前维度: {}...", embedding_dim); - sqlx::query("DROP TABLE IF EXISTS vec_paper_chunks").execute(&pool).await?; - sqlx::query("DELETE FROM paper_chunks_content").execute(&pool).await?; + info!( + "检测到已存在的向量表维度不匹配，正在重建以适配当前维度: {}...", + embedding_dim + ); + sqlx::query("DROP TABLE IF EXISTS vec_paper_chunks") + .execute(&pool) + .await?; + sqlx::query("DELETE FROM paper_chunks_content") + .execute(&pool) + .await?; } } @@ -116,7 +136,7 @@ async fn main() -> anyhow::Result<()> { let ads = AdsClient::new(config.ads_api_key.clone()); let arxiv = ArxivClient::new(); - let downloader = Downloader::new(); + let downloader = Downloader::new().context("构建 HTTP 下载客户端失败")?; let llm = LlmClient::new( config.llm_api_key.clone(), config.llm_api_base.clone(), @@ -128,6 +148,20 @@ async fn main() -> anyhow::Result<()> { config.embedding_model.clone(), ); + let skill_registry = Arc::new(RwLock::new(SkillRegistry::new(config.skills_dir.clone()))); + if let Ok(mut reg) = skill_registry.write() { + reg.refresh(); + } else { + warn!("SkillRegistry 初始化刷新失败（RwLock 异常），将使用磁盘缓存"); + } + info!( + "SkillRegistry 初始化完成，加载 {} 个 skill。", + skill_registry.read().map(|r| r.len()).unwrap_or(0) + ); + + // 启动文件监听（热更新 skills） + let _watcher_handle = SkillRegistry::start_watcher(skill_registry.clone()); + let app_state = Arc::new(AppState { config: config.clone(), db: pool, @@ -138,10 +172,20 @@ async fn main() -> anyhow::Result<()> { llm, embedding, downloader, - harvest_status: Arc::new(tokio::sync::Mutex::new(astroresearch::services::batch_sync::MetaSyncStatus::new())), - batch_status: Arc::new(tokio::sync::Mutex::new(astroresearch::services::batch_sync::AssetBatchStatus::new())), + harvest_status: Arc::new(tokio::sync::Mutex::new( + astroresearch::services::batch_sync::MetaSyncStatus::new(), + )), + batch_status: Arc::new(tokio::sync::Mutex::new( + astroresearch::services::batch_sync::AssetBatchStatus::new(), + )), active_bibcode: Arc::new(tokio::sync::Mutex::new(None)), - cancelled_runs: Arc::new(std::sync::Mutex::new(std::collections::HashSet::new())), + cancelled_runs: Arc::new(Mutex::new(std::collections::HashSet::new())), + skill_registry, + pending_questions: Arc::new(Mutex::new(HashMap::new())), + sse_broadcast: None, + memory_manager: Arc::new(tokio::sync::Mutex::new( + astroresearch::agent::memory::MemoryManager::new(config.library_dir.clone()), + )), }); // 7. 设置 Axum 路由、CORS 头以及 React 仪表盘静态资源托管 @@ -153,7 +197,11 @@ async fn main() -> anyhow::Result<()> { let api_routes = Router::new() .route("/search", get(handlers::search_papers)) .route("/download", post(handlers::download_paper)) - .route("/upload", post(handlers::upload_paper_file).layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024))) + .route( + "/upload", + post(handlers::upload_paper_file) + .layer(axum::extract::DefaultBodyLimit::max(100 * 1024 * 1024)), + ) .route("/no_resource", post(handlers::mark_no_resource)) .route("/parse", post(handlers::parse_paper)) .route("/translate", post(handlers::translate_paper)) @@ -172,8 +220,14 @@ async fn main() -> anyhow::Result<()> { .route("/batch/asset/stop", post(handlers::stop_asset_batch)) .route("/batch/asset/status", get(handlers::get_asset_batch_status)) .route("/sync/queries", get(handlers::get_sync_queries)) - .route("/sync/queries/:id", axum::routing::delete(handlers::delete_sync_query)) - .route("/active_bibcode", get(handlers::get_active_bibcode).post(handlers::set_active_bibcode)) + .route( + "/sync/queries/:id", + axum::routing::delete(handlers::delete_sync_query), + ) + .route( + "/active_bibcode", + get(handlers::get_active_bibcode).post(handlers::set_active_bibcode), + ) .route("/chat/rag", post(handlers::chat_rag)) .route("/chat/figure", post(handlers::chat_figure)) .route("/target/query", get(handlers::query_target)) @@ -182,13 +236,21 @@ async fn main() -> anyhow::Result<()> { .route("/target/list", get(handlers::list_targets)) // 智能体路由 .route("/chat/agent", post(handlers::chat_agent)) + .route("/chat/metrics", get(handlers::get_agent_metrics)) .route("/chat/sessions", get(handlers::list_sessions)) - .route("/chat/sessions/:id", get(handlers::get_session).delete(handlers::delete_session)) - .route("/chat/sessions/:id/stop", post(handlers::stop_agent)); + .route( + "/chat/sessions/:id", + get(handlers::get_session).delete(handlers::delete_session), + ) + .route("/chat/sessions/:id/stop", post(handlers::stop_agent)) + .route("/chat/sessions/:id/audit", get(handlers::get_session_audit)) + .route("/chat/questions", get(handlers::get_pending_questions)) + .route("/chat/answer", post(handlers::answer_question)); // 静态文件资源代理托管（当前端打包至 dashboard/dist 后，直接挂载到主域名根路由） - let serve_dir = ServeDir::new("dashboard/dist") - .fallback(tower_http::services::ServeFile::new("dashboard/dist/index.html")); + let serve_dir = ServeDir::new("dashboard/dist").fallback(tower_http::services::ServeFile::new( + "dashboard/dist/index.html", + )); let app = Router::new() .nest("/api", api_routes) @@ -200,7 +262,7 @@ async fn main() -> anyhow::Result<()> { let addr = SocketAddr::from(([0, 0, 0, 0], config.port)); info!("天文学科研服务已成功监听 http://localhost:{}", config.port); - + let listener = tokio::net::TcpListener::bind(addr).await?; axum::serve(listener, app).await?; diff --git a/src/services/batch/asset.rs b/src/services/batch/asset.rs index af32e17..31b127c 100644 --- a/src/services/batch/asset.rs +++ b/src/services/batch/asset.rs @@ -1,14 +1,14 @@ // src/services/batch/asset.rs -use std::sync::Arc; +use serde::{Deserialize, Serialize}; +use sqlx::{Row, SqlitePool}; use std::fs; +use std::sync::Arc; use tokio::sync::Mutex; -use serde::{Serialize, Deserialize}; -use tracing::{info, warn, error}; -use sqlx::{SqlitePool, Row}; +use tracing::{info, warn}; -use crate::Config; use crate::clients::qiniu::QiniuClient; use crate::services::download::Downloader; +use crate::Config; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] @@ -34,6 +34,12 @@ pub struct AssetBatchStatus { pub action: Option, } +impl Default for AssetBatchStatus { + fn default() -> Self { + Self::new() + } +} + impl AssetBatchStatus { pub fn new() -> Self { AssetBatchStatus { @@ -63,6 +69,7 @@ pub struct AssetBatch; impl AssetBatch { /// 启动后台批量下载与结构化解析任务 + #[allow(clippy::too_many_arguments)] pub fn start_process( db: SqlitePool, config: Config, @@ -96,7 +103,7 @@ impl AssetBatch { s.current_bibcode = String::new(); s.logs.clear(); s.action = Some(action); - + let action_desc = match action { BatchAction::Download => "下载", BatchAction::Parse => "解析", @@ -105,7 +112,10 @@ impl AssetBatch { BatchAction::Target => "天体识别", BatchAction::All => "下载与解析", }; - s.add_log(format!("批量{}任务启动，共 {} 篇文献需处理。", action_desc, total)); + s.add_log(format!( + "批量{}任务启动，共 {} 篇文献需处理。", + action_desc, total + )); } let mut dl_count = 0; @@ -136,7 +146,15 @@ impl AssetBatch { .fetch_optional(&db) .await; - let (arxiv_id, doi, mut pdf_path, mut html_path, markdown_path, doctype, translation_path) = match paper_res { + let ( + arxiv_id, + doi, + mut pdf_path, + mut html_path, + markdown_path, + doctype, + translation_path, + ) = match paper_res { Ok(Some(row)) => { let arxiv_id: String = row.get(0); let doi: String = row.get(1); @@ -145,7 +163,15 @@ impl AssetBatch { let markdown_path: Option = row.get(4); let doctype: Option = row.get(5); let translation_path: Option = row.get(6); - (arxiv_id, doi, pdf_path, html_path, markdown_path, doctype, translation_path) + ( + arxiv_id, + doi, + pdf_path, + html_path, + markdown_path, + doctype, + translation_path, + ) } _ => { let mut s = status.lock().await; @@ -155,24 +181,34 @@ impl AssetBatch { }; // 1b. 检查 doctype，如果是 proposal, abstract, catalog, dataset, software, circular 等无数字全文的文件，直接跳过处理 - let doctype_str = doctype.unwrap_or_else(|| "article".to_string()).to_lowercase(); - if doctype_str == "proposal" - || doctype_str == "abstract" - || doctype_str == "catalog" + let doctype_str = doctype + .unwrap_or_else(|| "article".to_string()) + .to_lowercase(); + if doctype_str == "proposal" + || doctype_str == "abstract" + || doctype_str == "catalog" || doctype_str == "dataset" - || doctype_str == "software" + || doctype_str == "software" || doctype_str == "circular" || doctype_str == "newsletter" - || doctype_str == "obituary" + || doctype_str == "obituary" { let mut s = status.lock().await; - s.add_log(format!("文献 {} 的类型为 {} (无数字版全文)，跳过下载与解析。", bibcode, doctype_str)); + s.add_log(format!( + "文献 {} 的类型为 {} (无数字版全文)，跳过下载与解析。", + bibcode, doctype_str + )); // 同样更新处理进度，防止任务进度条卡住 if action == BatchAction::Download || action == BatchAction::All { dl_count += 1; s.downloaded = dl_count; } - if action == BatchAction::Parse || action == BatchAction::All || action == BatchAction::Translate || action == BatchAction::Embed || action == BatchAction::Target { + if action == BatchAction::Parse + || action == BatchAction::All + || action == BatchAction::Translate + || action == BatchAction::Embed + || action == BatchAction::Target + { s.parsed += 1; } continue; @@ -180,8 +216,14 @@ impl AssetBatch { // 2. 检查并执行下载 if action == BatchAction::Download || action == BatchAction::All { - let is_pdf_exist = pdf_path.as_ref().map(|p| config.library_dir.join(p).exists()).unwrap_or(false); - let is_html_exist = html_path.as_ref().map(|p| config.library_dir.join(p).exists()).unwrap_or(false); + let is_pdf_exist = pdf_path + .as_ref() + .map(|p| config.library_dir.join(p).exists()) + .unwrap_or(false); + let is_html_exist = html_path + .as_ref() + .map(|p| config.library_dir.join(p).exists()) + .unwrap_or(false); if !is_pdf_exist && !is_html_exist { // 需要执行下载 @@ -191,29 +233,56 @@ impl AssetBatch { } let (pdf_res, html_res) = if !arxiv_id.is_empty() { - let res = downloader.download_arxiv_direct(&arxiv_id, &config.library_dir).await; + let res = downloader + .download_arxiv_direct(&arxiv_id, &config.library_dir) + .await; if res.0.is_ok() || res.1.is_ok() { res } else { { let mut s = status.lock().await; - s.add_log(format!("文献 {} arXiv 通道下载失败，回退尝试 ADS/出版商下载...", bibcode)); + s.add_log(format!( + "文献 {} arXiv 通道下载失败，回退尝试 ADS/出版商下载...", + bibcode + )); } - let doi_opt = if !doi.is_empty() { Some(doi.as_str()) } else { None }; - downloader.download_paper(&bibcode, doi_opt, &config.library_dir).await + let doi_opt = if !doi.is_empty() { + Some(doi.as_str()) + } else { + None + }; + downloader + .download_paper(&bibcode, doi_opt, &config.library_dir) + .await } } else { - let doi_opt = if !doi.is_empty() { Some(doi.as_str()) } else { None }; - downloader.download_paper(&bibcode, doi_opt, &config.library_dir).await + let doi_opt = if !doi.is_empty() { + Some(doi.as_str()) + } else { + None + }; + downloader + .download_paper(&bibcode, doi_opt, &config.library_dir) + .await }; if pdf_res.is_ok() || html_res.is_ok() { let pdf_rel = match pdf_res { - Ok(p) => Some(p.strip_prefix(&config.library_dir).unwrap_or(&p).to_string_lossy().to_string()), + Ok(p) => Some( + p.strip_prefix(&config.library_dir) + .unwrap_or(&p) + .to_string_lossy() + .to_string(), + ), Err(_) => None, // 只要有一方下载成功，失败的一方字段置空（NULL），避免在 path 字段中留存报错日志 }; let html_rel = match html_res { - Ok(p) => Some(p.strip_prefix(&config.library_dir).unwrap_or(&p).to_string_lossy().to_string()), + Ok(p) => Some( + p.strip_prefix(&config.library_dir) + .unwrap_or(&p) + .to_string_lossy() + .to_string(), + ), Err(_) => None, // 只要有一方下载成功，失败的一方字段置空（NULL），避免在 path 字段中留存报错日志 }; @@ -221,12 +290,14 @@ impl AssetBatch { pdf_path = pdf_rel.clone(); html_path = html_rel.clone(); - let _ = sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?") - .bind(pdf_rel) - .bind(html_rel) - .bind(&bibcode) - .execute(&db) - .await; + let _ = sqlx::query( + "UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?", + ) + .bind(pdf_rel) + .bind(html_rel) + .bind(&bibcode) + .execute(&db) + .await; dl_count += 1; { @@ -238,7 +309,7 @@ impl AssetBatch { dl_failed_count += 1; let mut s = status.lock().await; s.download_failed = dl_failed_count; - + let pdf_err = match pdf_res { Err(e) => format!("error: {}", e), _ => "error: 未知错误".to_string(), @@ -247,15 +318,20 @@ impl AssetBatch { Err(e) => format!("error: {}", e), _ => "error: 未知错误".to_string(), }; - - s.add_log(format!("文献 {} 下载失败。PDF: {}, HTML: {}", bibcode, pdf_err, html_err)); - - let _ = sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?") - .bind(&pdf_err) - .bind(&html_err) - .bind(&bibcode) - .execute(&db) - .await; + + s.add_log(format!( + "文献 {} 下载失败。PDF: {}, HTML: {}", + bibcode, pdf_err, html_err + )); + + let _ = sqlx::query( + "UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?", + ) + .bind(&pdf_err) + .bind(&html_err) + .bind(&bibcode) + .execute(&db) + .await; } // 每次下载尝试后，加入 3-5 秒随机延迟，防爬防封 @@ -264,7 +340,10 @@ impl AssetBatch { } else { { let mut s = status.lock().await; - s.add_log(format!("文献 {} 本地已存在 PDF 或 HTML，跳过下载。", bibcode)); + s.add_log(format!( + "文献 {} 本地已存在 PDF 或 HTML，跳过下载。", + bibcode + )); } dl_count += 1; { @@ -276,12 +355,18 @@ impl AssetBatch { // 3. 检查并执行结构化解析（Markdown 转换） if action == BatchAction::Parse || action == BatchAction::All { - let is_md_exist = markdown_path.as_ref().map(|p| config.library_dir.join(p).exists()).unwrap_or(false); + let is_md_exist = markdown_path + .as_ref() + .map(|p| config.library_dir.join(p).exists()) + .unwrap_or(false); if !is_md_exist { if pdf_path.is_some() || html_path.is_some() { { let mut s = status.lock().await; - s.add_log(format!("文献 {} 开始进行排版提取与 Markdown 转换...", bibcode)); + s.add_log(format!( + "文献 {} 开始进行排版提取与 Markdown 转换...", + bibcode + )); } let mut relative_md_path = String::new(); @@ -290,7 +375,10 @@ impl AssetBatch { let source_url = if bibcode.len() == 19 { format!("https://ui.adsabs.harvard.edu/abs/{}/abstract", bibcode) } else if !arxiv_id.is_empty() { - format!("https://ui.adsabs.harvard.edu/abs/arXiv:{}/abstract", arxiv_id) + format!( + "https://ui.adsabs.harvard.edu/abs/arXiv:{}/abstract", + arxiv_id + ) } else { format!("https://ui.adsabs.harvard.edu/abs/{}/abstract", bibcode) }; @@ -299,7 +387,9 @@ impl AssetBatch { if let Some(html_rel) = &html_path { let html_abs = config.library_dir.join(html_rel); if html_abs.exists() { - if let Ok(md) = crate::services::parser::html_to_markdown(&html_abs) { + if let Ok(md) = + crate::services::parser::html_to_markdown(&html_abs) + { // 构建 Meta 头 let paper_meta_res = sqlx::query("SELECT title, authors, pub, year, keywords FROM papers WHERE bibcode = ?") .bind(&bibcode) @@ -313,8 +403,12 @@ impl AssetBatch { let year: String = meta_row.get(3); let keywords_json: String = meta_row.get(4); - let authors: Vec = serde_json::from_str(&authors_json).unwrap_or_default(); - let keywords: Vec = serde_json::from_str(&keywords_json).unwrap_or_default(); + let authors: Vec = + serde_json::from_str(&authors_json) + .unwrap_or_default(); + let keywords: Vec = + serde_json::from_str(&keywords_json) + .unwrap_or_default(); let front_matter = format!( "---\ntitle: {}\nauthor: [{}]\npublisher: {}\nsource: \"{}\"\ndate: \"{}\"\ntags: \"{}\"\n---\n\n", @@ -327,10 +421,16 @@ impl AssetBatch { ); let parsed_markdown = format!("{}{}", front_matter, md); let md_filename = format!("{}.md", bibcode); - let md_dest = config.library_dir.join("Markdown").join(&md_filename); - let _ = fs::create_dir_all(md_dest.parent().unwrap()); + let md_dest = config + .library_dir + .join("Markdown") + .join(&md_filename); + if let Some(parent) = md_dest.parent() { + let _ = fs::create_dir_all(parent); + } if fs::write(&md_dest, &parsed_markdown).is_ok() { - relative_md_path = format!("Markdown/{}", md_filename); + relative_md_path = + format!("Markdown/{}", md_filename); } } } @@ -339,12 +439,14 @@ impl AssetBatch { if !relative_md_path.is_empty() { // HTML 解析成功，直接写入数据库并记录成功 - let _ = sqlx::query("UPDATE papers SET markdown_path = ? WHERE bibcode = ?") - .bind(&relative_md_path) - .bind(&bibcode) - .execute(&db) - .await; - + let _ = sqlx::query( + "UPDATE papers SET markdown_path = ? WHERE bibcode = ?", + ) + .bind(&relative_md_path) + .bind(&bibcode) + .execute(&db) + .await; + { let mut s = status.lock().await; s.parsed += 1; @@ -356,7 +458,8 @@ impl AssetBatch { let pdf_abs = config.library_dir.join(pdf_rel); if pdf_abs.exists() { // 检查是否已经是 mineru_batch: 状态 - let existing_batch_id = markdown_path.as_ref() + let existing_batch_id = markdown_path + .as_ref() .and_then(|p| p.strip_prefix("mineru_batch:")) .map(|s| s.trim().to_string()); @@ -377,9 +480,16 @@ impl AssetBatch { } else { { let mut s = status.lock().await; - s.add_log(format!("文献 {} PDF 提交后台解析 (MinerU)...", bibcode)); + s.add_log(format!( + "文献 {} PDF 提交后台解析 (MinerU)...", + bibcode + )); } - match crate::services::parser::submit_pdf_to_mineru(&pdf_abs, &config).await { + match crate::services::parser::submit_pdf_to_mineru( + &pdf_abs, &config, + ) + .await + { Ok(id) => { // 提交成功，立刻把 batch_id 存入数据库以备断点续跑 let marker = format!("mineru_batch:{}", id); @@ -393,7 +503,10 @@ impl AssetBatch { Err(e) => { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} PDF 提交 MinerU 失败: {}", bibcode, e)); + s.add_log(format!( + "文献 {} PDF 提交 MinerU 失败: {}", + bibcode, e + )); let err_reason = format!("error: {}", e); let _ = sqlx::query("UPDATE papers SET markdown_path = ? WHERE bibcode = ?") .bind(&err_reason) @@ -438,7 +551,9 @@ impl AssetBatch { let parsed_markdown = format!("{}{}", front_matter, md); let md_filename = format!("{}.md", bibcode_clone); let md_dest = config_clone.library_dir.join("Markdown").join(&md_filename); - let _ = fs::create_dir_all(md_dest.parent().unwrap()); + if let Some(parent) = md_dest.parent() { + let _ = fs::create_dir_all(parent); + } if fs::write(&md_dest, &parsed_markdown).is_ok() { rel_md = format!("Markdown/{}", md_filename); } @@ -450,7 +565,7 @@ impl AssetBatch { .bind(&bibcode_clone) .execute(&db_clone) .await; - + let mut s = status_clone.lock().await; s.parsed += 1; s.add_log(format!("文献 {} PDF (MinerU) 解析成功！", bibcode_clone)); @@ -482,7 +597,10 @@ impl AssetBatch { } else { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 本地 PDF 文件不存在，无法解析。", bibcode)); + s.add_log(format!( + "文献 {} 本地 PDF 文件不存在，无法解析。", + bibcode + )); let _ = sqlx::query("UPDATE papers SET markdown_path = 'error: 本地 PDF 文件不存在' WHERE bibcode = ?") .bind(&bibcode) .execute(&db) @@ -491,7 +609,10 @@ impl AssetBatch { } else { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} HTML 转换失败，且无本地 PDF，无法解析。", bibcode)); + s.add_log(format!( + "文献 {} HTML 转换失败，且无本地 PDF，无法解析。", + bibcode + )); let _ = sqlx::query("UPDATE papers SET markdown_path = 'error: HTML 转换失败且无本地 PDF' WHERE bibcode = ?") .bind(&bibcode) .execute(&db) @@ -501,7 +622,10 @@ impl AssetBatch { } else { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 无本地 PDF/HTML，无法解析，跳过。", bibcode)); + s.add_log(format!( + "文献 {} 无本地 PDF/HTML，无法解析，跳过。", + bibcode + )); } } else { { @@ -515,7 +639,10 @@ impl AssetBatch { // 4. 检查并执行翻译 if action == BatchAction::Translate { - let is_tr_exist = translation_path.as_ref().map(|p| config.library_dir.join(p).exists() && !p.starts_with("error:")).unwrap_or(false); + let is_tr_exist = translation_path + .as_ref() + .map(|p| config.library_dir.join(p).exists() && !p.starts_with("error:")) + .unwrap_or(false); if !is_tr_exist { if let Some(md_rel) = &markdown_path { if !md_rel.starts_with("error:") { @@ -528,22 +655,39 @@ impl AssetBatch { match fs::read_to_string(&md_abs) { Ok(english_markdown) => { - match crate::services::translation::translate_markdown(&english_markdown, &dict, &llm_client).await { + match crate::services::translation::translate_markdown( + &english_markdown, + &dict, + &llm_client, + ) + .await + { Ok(translated_markdown) => { let tr_filename = format!("{}_zh.md", bibcode); - let tr_dest = config.library_dir.join("Translation").join(&tr_filename); - let _ = fs::create_dir_all(tr_dest.parent().unwrap()); - if fs::write(&tr_dest, &translated_markdown).is_ok() { - let relative_tr_path = format!("Translation/{}", tr_filename); + let tr_dest = config + .library_dir + .join("Translation") + .join(&tr_filename); + if let Some(parent) = tr_dest.parent() { + let _ = fs::create_dir_all(parent); + } + if fs::write(&tr_dest, &translated_markdown) + .is_ok() + { + let relative_tr_path = + format!("Translation/{}", tr_filename); let _ = sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") .bind(&relative_tr_path) .bind(&bibcode) .execute(&db) .await; - + let mut s = status.lock().await; s.parsed += 1; - s.add_log(format!("文献 {} 翻译成功！", bibcode)); + s.add_log(format!( + "文献 {} 翻译成功！", + bibcode + )); } else { let error_msg = "error: 写入翻译文件失败"; let _ = sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") @@ -553,7 +697,10 @@ impl AssetBatch { .await; let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 翻译文件写入失败。", bibcode)); + s.add_log(format!( + "文献 {} 翻译文件写入失败。", + bibcode + )); } } Err(e) => { @@ -565,12 +712,16 @@ impl AssetBatch { .await; let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 翻译失败: {}", bibcode, e)); + s.add_log(format!( + "文献 {} 翻译失败: {}", + bibcode, e + )); } } } Err(e) => { - let error_msg = format!("error: 读取英文 Markdown 失败: {}", e); + let error_msg = + format!("error: 读取英文 Markdown 失败: {}", e); let _ = sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") .bind(&error_msg) .bind(&bibcode) @@ -578,41 +729,59 @@ impl AssetBatch { .await; let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 读取英文 Markdown 失败: {}", bibcode, e)); + s.add_log(format!( + "文献 {} 读取英文 Markdown 失败: {}", + bibcode, e + )); } } } else { let error_msg = "error: 英文 Markdown 文件不存在"; - let _ = sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") - .bind(error_msg) - .bind(&bibcode) - .execute(&db) - .await; - let mut s = status.lock().await; - s.parse_failed += 1; - s.add_log(format!("文献 {} 英文 Markdown 文件不存在，无法翻译。", bibcode)); - } - } else { - let error_msg = "error: 英文 Markdown 文件处于解析失败状态"; - let _ = sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") + let _ = sqlx::query( + "UPDATE papers SET translation_path = ? WHERE bibcode = ?", + ) .bind(error_msg) .bind(&bibcode) .execute(&db) .await; - let mut s = status.lock().await; - s.parse_failed += 1; - s.add_log(format!("文献 {} 英文 Markdown 解析失败，跳过翻译。", bibcode)); - } - } else { - let error_msg = "error: 尚未解析英文 Markdown 路径为 NULL"; - let _ = sqlx::query("UPDATE papers SET translation_path = ? WHERE bibcode = ?") + let mut s = status.lock().await; + s.parse_failed += 1; + s.add_log(format!( + "文献 {} 英文 Markdown 文件不存在，无法翻译。", + bibcode + )); + } + } else { + let error_msg = "error: 英文 Markdown 文件处于解析失败状态"; + let _ = sqlx::query( + "UPDATE papers SET translation_path = ? WHERE bibcode = ?", + ) .bind(error_msg) .bind(&bibcode) .execute(&db) .await; + let mut s = status.lock().await; + s.parse_failed += 1; + s.add_log(format!( + "文献 {} 英文 Markdown 解析失败，跳过翻译。", + bibcode + )); + } + } else { + let error_msg = "error: 尚未解析英文 Markdown 路径为 NULL"; + let _ = sqlx::query( + "UPDATE papers SET translation_path = ? WHERE bibcode = ?", + ) + .bind(error_msg) + .bind(&bibcode) + .execute(&db) + .await; let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 尚未解析英文 Markdown，跳过翻译。", bibcode)); + s.add_log(format!( + "文献 {} 尚未解析英文 Markdown，跳过翻译。", + bibcode + )); } } else { { @@ -626,7 +795,10 @@ impl AssetBatch { // 5. 检查并执行向量化 (Embedding) if action == BatchAction::Embed { - let is_md_exist = markdown_path.as_ref().map(|p| config.library_dir.join(p).exists() && !p.starts_with("error:")).unwrap_or(false); + let is_md_exist = markdown_path + .as_ref() + .map(|p| config.library_dir.join(p).exists() && !p.starts_with("error:")) + .unwrap_or(false); if is_md_exist { let md_rel = markdown_path.as_ref().unwrap(); let md_abs = config.library_dir.join(md_rel); @@ -636,11 +808,22 @@ impl AssetBatch { } match fs::read_to_string(&md_abs) { Ok(markdown_content) => { - match crate::services::rag::ingest_paper(&db, &embedding_client, &bibcode, &markdown_content, None).await { + match crate::services::rag::ingest_paper( + &db, + &embedding_client, + &bibcode, + &markdown_content, + None, + ) + .await + { Ok(chunk_count) => { let mut s = status.lock().await; s.parsed += 1; - s.add_log(format!("文献 {} 向量化成功，共切片入库 {} 个向量块。", bibcode, chunk_count)); + s.add_log(format!( + "文献 {} 向量化成功，共切片入库 {} 个向量块。", + bibcode, chunk_count + )); } Err(e) => { let mut s = status.lock().await; @@ -652,19 +835,28 @@ impl AssetBatch { Err(e) => { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 读取英文 Markdown 失败: {}", bibcode, e)); + s.add_log(format!( + "文献 {} 读取英文 Markdown 失败: {}", + bibcode, e + )); } } } else { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 英文 Markdown 文件不存在，跳过向量化。", bibcode)); + s.add_log(format!( + "文献 {} 英文 Markdown 文件不存在，跳过向量化。", + bibcode + )); } } // 6. 检查并执行天体识别与缓存 if action == BatchAction::Target { - let is_md_exist = markdown_path.as_ref().map(|p| config.library_dir.join(p).exists() && !p.starts_with("error:")).unwrap_or(false); + let is_md_exist = markdown_path + .as_ref() + .map(|p| config.library_dir.join(p).exists() && !p.starts_with("error:")) + .unwrap_or(false); if is_md_exist { let md_rel = markdown_path.as_ref().unwrap(); let md_abs = config.library_dir.join(md_rel); @@ -675,30 +867,47 @@ impl AssetBatch { match fs::read_to_string(&md_abs) { Ok(markdown_content) => { // 提取前先清空旧的关联 - if let Err(e) = sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?") - .bind(&bibcode) - .execute(&db) - .await + if let Err(e) = + sqlx::query("DELETE FROM paper_targets WHERE bibcode = ?") + .bind(&bibcode) + .execute(&db) + .await { warn!("清除文献 {} 的旧天体关联失败: {}", bibcode, e); } let client = reqwest::Client::new(); - let targets = crate::services::target::extract_and_cache_targets(&db, &markdown_content, &bibcode, &client).await; + let targets = crate::services::target::extract_and_cache_targets( + &db, + &markdown_content, + &bibcode, + &client, + ) + .await; let mut s = status.lock().await; s.parsed += 1; - s.add_log(format!("文献 {} 天体识别完成，共识别并缓存 {} 个天体目标。", bibcode, targets.len())); + s.add_log(format!( + "文献 {} 天体识别完成，共识别并缓存 {} 个天体目标。", + bibcode, + targets.len() + )); } Err(e) => { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 读取英文 Markdown 失败: {}", bibcode, e)); + s.add_log(format!( + "文献 {} 读取英文 Markdown 失败: {}", + bibcode, e + )); } } } else { let mut s = status.lock().await; s.parse_failed += 1; - s.add_log(format!("文献 {} 英文 Markdown 文件不存在，跳过天体识别。", bibcode)); + s.add_log(format!( + "文献 {} 英文 Markdown 文件不存在，跳过天体识别。", + bibcode + )); } } } @@ -706,7 +915,10 @@ impl AssetBatch { if !join_handles.is_empty() { { let mut s = status.lock().await; - s.add_log(format!("本地下载与快速解析已完成，正在等待后台共 {} 个 MinerU 异步解析任务结束...", join_handles.len())); + s.add_log(format!( + "本地下载与快速解析已完成，正在等待后台共 {} 个 MinerU 异步解析任务结束...", + join_handles.len() + )); } for handle in join_handles { let _ = handle.await; @@ -740,11 +952,11 @@ mod tests { async fn test_process_status_log_rotation() { let mut status = AssetBatchStatus::new(); assert!(!status.active); - + for i in 0..150 { status.add_log(format!("log {}", i)); } - + assert_eq!(status.logs.len(), 100); assert_eq!(status.logs[0], "log 50"); assert_eq!(status.logs[99], "log 149"); @@ -758,15 +970,13 @@ mod tests { .await?; // 运行迁移 - sqlx::migrate!("./migrations") - .run(&pool) - .await?; + sqlx::migrate!("./migrations").run(&pool).await?; // 创建临时目录 let test_id = rand::random::(); let temp_dir = std::env::temp_dir().join(format!("astro_research_test_{}", test_id)); fs::create_dir_all(&temp_dir)?; - + // 准备子目录 let pdf_dir = temp_dir.join("PDF"); let html_dir = temp_dir.join("HTML"); @@ -779,7 +989,7 @@ mod tests { let bibcode = "2026A&A...123..456X".to_string(); let pdf_file_rel = format!("PDF/{}.pdf", bibcode); let html_file_rel = format!("HTML/{}.html", bibcode); - + fs::write(temp_dir.join(&pdf_file_rel), b"%PDF-1.5 test")?; fs::write(temp_dir.join(&html_file_rel), b"

    Test Paper

    Content

    ")?; @@ -806,8 +1016,13 @@ mod tests { let mut config = Config::from_env(); config.library_dir = temp_dir.clone(); - let downloader = Arc::new(Downloader::new()); - let qiniu = Arc::new(QiniuClient::new("test_access".to_string(), "test_secret".to_string(), "test_bucket".to_string(), "test_domain".to_string())); + let downloader = Arc::new(Downloader::new().expect("Failed to create downloader in test")); + let qiniu = Arc::new(QiniuClient::new( + "test_access".to_string(), + "test_secret".to_string(), + "test_bucket".to_string(), + "test_domain".to_string(), + )); let status = Arc::new(Mutex::new(AssetBatchStatus::new())); let dict = Arc::new(crate::services::translation::Dictionary::new()); @@ -864,9 +1079,7 @@ mod tests { .connect("sqlite::memory:") .await?; - sqlx::migrate!("./migrations") - .run(&pool) - .await?; + sqlx::migrate!("./migrations").run(&pool).await?; let test_id = rand::random::(); let temp_dir = std::env::temp_dir().join(format!("astro_research_test_stop_{}", test_id)); @@ -926,8 +1139,13 @@ mod tests { let mut config = Config::from_env(); config.library_dir = temp_dir.clone(); - let downloader = Arc::new(Downloader::new()); - let qiniu = Arc::new(QiniuClient::new("test_access".to_string(), "test_secret".to_string(), "test_bucket".to_string(), "test_domain".to_string())); + let downloader = Arc::new(Downloader::new().expect("Failed to create downloader in test")); + let qiniu = Arc::new(QiniuClient::new( + "test_access".to_string(), + "test_secret".to_string(), + "test_bucket".to_string(), + "test_domain".to_string(), + )); let status = Arc::new(Mutex::new(AssetBatchStatus::new())); let dict = Arc::new(crate::services::translation::Dictionary::new()); diff --git a/src/services/batch/meta.rs b/src/services/batch/meta.rs index 7b7a70c..626370e 100644 --- a/src/services/batch/meta.rs +++ b/src/services/batch/meta.rs @@ -1,13 +1,15 @@ // src/services/batch/meta.rs +use serde::Serialize; +use sqlx::SqlitePool; use std::sync::Arc; use tokio::sync::Mutex; -use serde::Serialize; -use tracing::{info, warn, error}; -use sqlx::SqlitePool; +use tracing::{error, info, warn}; +use crate::api::handlers::{ + convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db, +}; use crate::clients::ads::AdsClient; use crate::clients::arxiv::ArxivClient; -use crate::api::handlers::{convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db}; // 批量元数据同步进度状态 #[derive(Debug, Clone, Serialize)] @@ -19,6 +21,12 @@ pub struct MetaSyncStatus { pub total: i32, } +impl Default for MetaSyncStatus { + fn default() -> Self { + Self::new() + } +} + impl MetaSyncStatus { pub fn new() -> Self { MetaSyncStatus { @@ -81,13 +89,16 @@ impl MetaSync { let source_clone = source.clone(); tokio::spawn(async move { - info!("启动后台批量元数据同步任务: 查询词='{}', 源='{}', 上限={}", query_clone, source_clone, limit); - + info!( + "启动后台批量元数据同步任务: 查询词='{}', 源='{}', 上限={}", + query_clone, source_clone, limit + ); + // 自动将检索配置存入/更新至 sync_queries 数据库表中进行去重和时间更新 let _ = sqlx::query( "INSERT INTO sync_queries (query, source, limit_count, last_run) \ VALUES (?, ?, ?, CURRENT_TIMESTAMP) \ - ON CONFLICT(query, source, limit_count) DO UPDATE SET last_run=excluded.last_run" + ON CONFLICT(query, source, limit_count) DO UPDATE SET last_run=excluded.last_run", ) .bind(&query_clone) .bind(&source_clone) @@ -131,7 +142,11 @@ impl MetaSync { } // 计算实际需要元数据同步的总上限，并按比例分配或根据实际匹配量上限控制 - let limit_to_harvest = if limit > 0 { std::cmp::min(limit, total_count) } else { total_count }; + let limit_to_harvest = if limit > 0 { + std::cmp::min(limit, total_count) + } else { + total_count + }; // 共享的 atomic 计数器，以便两端并行同步时独立累加进度 let synced_counter = Arc::new(std::sync::atomic::AtomicI32::new(0)); @@ -144,10 +159,12 @@ impl MetaSync { let synced_counter = synced_counter.clone(); let status = status.clone(); let is_active = source_clone == "all" || source_clone == "ads"; - + // 如果是 all 模式，各平台按比例分摊 limit 额度，或者直接限制自身的最大可用量 let ads_limit = if source_clone == "all" { - if ads_total == 0 { 0 } else { + if ads_total == 0 { + 0 + } else { let ratio = ads_total as f32 / total_count as f32; ((limit_to_harvest as f32) * ratio).round() as i32 } @@ -166,8 +183,14 @@ impl MetaSync { if chunk_size <= 0 { break; } - info!("正在同步 ADS 分批数据: start={}, rows={}", start_offset, chunk_size); - match ads.search(&query, start_offset, chunk_size, "relevance").await { + info!( + "正在同步 ADS 分批数据: start={}, rows={}", + start_offset, chunk_size + ); + match ads + .search(&query, start_offset, chunk_size, "relevance") + .await + { Ok(docs) => { if docs.is_empty() { break; @@ -181,7 +204,9 @@ impl MetaSync { start_offset += count; // 累加全局进度并更新状态 - let current_global = synced_counter.fetch_add(count, std::sync::atomic::Ordering::SeqCst) + count; + let current_global = synced_counter + .fetch_add(count, std::sync::atomic::Ordering::SeqCst) + + count; { let mut s = status.lock().await; s.synced = current_global; @@ -205,7 +230,9 @@ impl MetaSync { let is_active = source_clone == "all" || source_clone == "arxiv"; let arxiv_limit = if source_clone == "all" { - if arxiv_total == 0 { 0 } else { + if arxiv_total == 0 { + 0 + } else { let ratio = arxiv_total as f32 / total_count as f32; ((limit_to_harvest as f32) * ratio).round() as i32 } @@ -224,8 +251,14 @@ impl MetaSync { if chunk_size <= 0 { break; } - info!("正在同步 arXiv 分批数据: start={}, max_results={}", start_offset, chunk_size); - match arxiv.search(&query, start_offset, chunk_size, "relevance").await { + info!( + "正在同步 arXiv 分批数据: start={}, max_results={}", + start_offset, chunk_size + ); + match arxiv + .search(&query, start_offset, chunk_size, "relevance") + .await + { Ok(papers) => { if papers.is_empty() { break; @@ -239,7 +272,9 @@ impl MetaSync { start_offset += count; // 累加全局进度并更新状态 - let current_global = synced_counter.fetch_add(count, std::sync::atomic::Ordering::SeqCst) + count; + let current_global = synced_counter + .fetch_add(count, std::sync::atomic::Ordering::SeqCst) + + count; { let mut s = status.lock().await; s.synced = current_global; @@ -266,7 +301,10 @@ impl MetaSync { let mut s = status.lock().await; s.active = false; s.synced = final_synced; - info!("后台批量元数据同步任务已结束。共成功同步 {} 篇文献。", final_synced); + info!( + "后台批量元数据同步任务已结束。共成功同步 {} 篇文献。", + final_synced + ); } }); } diff --git a/src/services/batch/mod.rs b/src/services/batch/mod.rs index ebcf5ac..636bf59 100644 --- a/src/services/batch/mod.rs +++ b/src/services/batch/mod.rs @@ -1,6 +1,6 @@ // src/services/batch/mod.rs -pub mod meta; pub mod asset; +pub mod meta; -pub use meta::{MetaSyncStatus, MetaSync}; -pub use asset::{BatchAction, AssetBatchStatus, AssetBatch}; +pub use asset::{AssetBatch, AssetBatchStatus, BatchAction}; +pub use meta::{MetaSync, MetaSyncStatus}; diff --git a/src/services/chunker.rs b/src/services/chunker.rs index af30bc1..b445659 100644 --- a/src/services/chunker.rs +++ b/src/services/chunker.rs @@ -48,9 +48,7 @@ pub fn chunk_markdown(text: &str, chunk_size: Option) -> Vec { } // 如果追加本段落后超出目标大小，先刷出已有的缓冲区 - if !current_buf.is_empty() - && current_buf.len() + trimmed.len() + 1 > target_size - { + if !current_buf.is_empty() && current_buf.len() + trimmed.len() + 1 > target_size { chunks.push(TextChunk { paragraph_index: chunk_index, content: current_buf.clone(), @@ -163,9 +161,8 @@ fn split_long_paragraph(text: &str, target_size: usize) -> Vec { current.push(chars[i]); // 只在非公式上下文中、到达句子边界时才考虑分割 - let is_sentence_end = !in_inline_math - && !in_display_math - && is_sentence_boundary(&chars, i); + let is_sentence_end = + !in_inline_math && !in_display_math && is_sentence_boundary(&chars, i); if is_sentence_end && current.len() >= target_size { result.push(current.clone()); @@ -236,7 +233,12 @@ mod tests { for chunk in &chunks { let dollar_count = chunk.content.matches('$').count(); // $ 符号应该成对出现 - assert_eq!(dollar_count % 2, 0, "LaTeX delimiters not balanced in chunk: {}", chunk.content); + assert_eq!( + dollar_count % 2, + 0, + "LaTeX delimiters not balanced in chunk: {}", + chunk.content + ); } } @@ -245,7 +247,11 @@ mod tests { let text = "# Title\n\nSome content here.\n\n## Subtitle\n\nMore content here."; let chunks = chunk_markdown(text, Some(500)); for chunk in &chunks { - assert!(!chunk.content.starts_with('#'), "Heading should be skipped: {}", chunk.content); + assert!( + !chunk.content.starts_with('#'), + "Heading should be skipped: {}", + chunk.content + ); } } diff --git a/src/services/download.rs b/src/services/download.rs index 4ee2ec8..44c45ee 100644 --- a/src/services/download.rs +++ b/src/services/download.rs @@ -9,15 +9,15 @@ //! - IOP/Springer 等特定出版商会话预热策略 //! - 请求间随机延迟（500-2000ms），降低触发反爬风险 +use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db}; +use crate::api::StandardPaper; +use anyhow::{Context, Result}; +use reqwest::header::{HeaderMap, HeaderValue}; use std::fs; use std::path::{Path, PathBuf}; -use reqwest::header::{HeaderMap, HeaderValue}; use tokio::io::AsyncWriteExt; +use tracing::{debug, error, info, warn}; use url::Url; -use tracing::{info, warn, debug, error}; -use anyhow::{Context, Result}; -use crate::api::StandardPaper; -use crate::api::helpers::{get_paper_from_db, check_paper_paths_in_db}; // ─── 浏览器伪装辅助 ──────────────────────────────────────────── @@ -41,10 +41,16 @@ fn build_browser_headers() -> HeaderMap { if let Ok(ua) = HeaderValue::from_str(&gen_useragent()) { h.insert("User-Agent", ua); } - h.insert("Accept", HeaderValue::from_static( - "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8", - )); - h.insert("Accept-Language", HeaderValue::from_static("en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7")); + h.insert( + "Accept", + HeaderValue::from_static( + "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8", + ), + ); + h.insert( + "Accept-Language", + HeaderValue::from_static("en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7"), + ); h.insert("DNT", HeaderValue::from_static("1")); h.insert("Connection", HeaderValue::from_static("keep-alive")); h.insert("Upgrade-Insecure-Requests", HeaderValue::from_static("1")); @@ -64,12 +70,21 @@ fn build_chrome_headers(referer: Option<&str>) -> HeaderMap { h.insert("Accept", HeaderValue::from_static( "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", )); - h.insert("Accept-Language", HeaderValue::from_static("en-US,en;q=0.9")); - h.insert("Sec-Ch-Ua", HeaderValue::from_static( - "\"Google Chrome\";v=\"143\", \"Chromium\";v=\"143\", \"Not A(Brand\";v=\"24\"", - )); + h.insert( + "Accept-Language", + HeaderValue::from_static("en-US,en;q=0.9"), + ); + h.insert( + "Sec-Ch-Ua", + HeaderValue::from_static( + "\"Google Chrome\";v=\"143\", \"Chromium\";v=\"143\", \"Not A(Brand\";v=\"24\"", + ), + ); h.insert("Sec-Ch-Ua-Mobile", HeaderValue::from_static("?0")); - h.insert("Sec-Ch-Ua-Platform", HeaderValue::from_static("\"Windows\"")); + h.insert( + "Sec-Ch-Ua-Platform", + HeaderValue::from_static("\"Windows\""), + ); h.insert("Sec-Fetch-Dest", HeaderValue::from_static("document")); h.insert("Sec-Fetch-Mode", HeaderValue::from_static("navigate")); h.insert("Sec-Fetch-Site", HeaderValue::from_static("same-origin")); @@ -94,10 +109,15 @@ fn detect_anti_bot(content: &str, url: Option<&str>) -> Result<()> { // 1. 强特征防爬与 WAF 挑战（任何小于 150KB 的内容都做检测） let waf_patterns = [ - "checking your browser", "please wait while we verify", - "cf-browser-verification", "cf_chl_opt", "just a moment", - "enable javascript and cookies", "_cf_chl_tk", - "awswafintegration", "aws waf", + "checking your browser", + "please wait while we verify", + "cf-browser-verification", + "cf_chl_opt", + "just a moment", + "enable javascript and cookies", + "_cf_chl_tk", + "awswafintegration", + "aws waf", ]; for p in &waf_patterns { if lower.contains(p) { @@ -106,8 +126,13 @@ fn detect_anti_bot(content: &str, url: Option<&str>) -> Result<()> { } let captcha_patterns = [ - "captcha", "recaptcha", "hcaptcha", "verify you are human", "robot check", - "radware bot manager", "shieldsquare", + "captcha", + "recaptcha", + "hcaptcha", + "verify you are human", + "robot check", + "radware bot manager", + "shieldsquare", ]; for p in &captcha_patterns { if lower.contains(p) { @@ -116,9 +141,14 @@ fn detect_anti_bot(content: &str, url: Option<&str>) -> Result<()> { } let access_denied = [ - "login required", "please log in", "subscription required", - "access denied", "you do not have access", "purchase this article", - "sign in to access", "client challenge", + "login required", + "please log in", + "subscription required", + "access denied", + "you do not have access", + "purchase this article", + "sign in to access", + "client challenge", ]; for p in &access_denied { if lower.contains(p) { @@ -135,9 +165,15 @@ fn detect_anti_bot(content: &str, url: Option<&str>) -> Result<()> { // 2. 通用 HTTP 错误与 CDN 关键字检测（仅当内容长度小于 5000 字节时检测，避免在正常文献中误判 CDN 脚本等） if content.len() < 5000 { let err_patterns = [ - "cloudflare", "service temporarily unavailable", "503 service", - "502 bad gateway", "504 gateway timeout", "403 forbidden", - "404 not found", "500 internal server error", "site error", + "cloudflare", + "service temporarily unavailable", + "503 service", + "502 bad gateway", + "504 gateway timeout", + "403 forbidden", + "404 not found", + "500 internal server error", + "site error", ]; for p in &err_patterns { if lower.contains(p) { @@ -176,11 +212,18 @@ pub(crate) fn validate_html_content(text: &str) -> Result<()> { let lower = text.to_lowercase(); // 1. 检查常见的跳转与错误占位特征 - if lower.contains("redirecting") || lower.contains("redirect to") || lower.contains("http-equiv=\"refresh\"") || lower.contains("autoredirecttourl") { + if lower.contains("redirecting") + || lower.contains("redirect to") + || lower.contains("http-equiv=\"refresh\"") + || lower.contains("autoredirecttourl") + { anyhow::bail!("检测到 HTML 重定向跳转页面，而非真实文献正文"); } - if lower.contains("conversion to html had a fatal error") || lower.contains("no content available") || lower.contains("fatal error and exited abruptly") { + if lower.contains("conversion to html had a fatal error") + || lower.contains("no content available") + || lower.contains("fatal error and exited abruptly") + { anyhow::bail!("检测到 ar5iv 转换失败的占位 HTML 页面"); } @@ -194,7 +237,7 @@ pub(crate) fn validate_html_content(text: &str) -> Result<()> { let title_start = start_pos + tag_end + 1; if let Some(end_pos) = lower[title_start..].find("") { let title = &lower[title_start..title_start + end_pos]; - if title.contains("nsf award search") + if title.contains("nsf award search") || title.contains("national science foundation") || title.contains("vizier") || title.contains("caltechthesis") @@ -211,8 +254,12 @@ pub(crate) fn validate_html_content(text: &str) -> Result<()> { // 3. 基础字节长度与具体 HTTP 错误特征校验 if text.len() < 2000 { let error_patterns = [ - "404 not found", "403 forbidden", "502 bad gateway", - "500 internal server error", "access denied", "site error" + "404 not found", + "403 forbidden", + "502 bad gateway", + "500 internal server error", + "access denied", + "site error", ]; for kw in &error_patterns { if lower.contains(kw) { @@ -224,21 +271,24 @@ pub(crate) fn validate_html_content(text: &str) -> Result<()> { // 4. 结构启发式校验：如果是小于 50KB 的 HTML，必须包含基本的章节或参考文献结构，否则判定为摘要/存根占位页 if text.len() < 50000 { // 匹配 heading 标签或 Markdown 格式的标题，而不是纯文本中的单词 - let has_sections = lower.contains("ltx_title_section") - || lower.contains("class=\"section\"") + let has_sections = lower.contains("ltx_title_section") + || lower.contains("class=\"section\"") || lower.contains("## introduction") || lower.contains("

    introduction") || lower.contains("

    introduction") || lower.contains("class=\"ltx_section\""); - let has_bib = lower.contains("ltx_bibliography") + let has_bib = lower.contains("ltx_bibliography") || lower.contains("class=\"references\"") || lower.contains("
      Result<()> { /// 仅做最低限度检查：页面不能过小，不能是纯跳转页。 pub(crate) fn validate_html_content_lenient(text: &str) -> Result<()> { if text.len() < 500 { - anyhow::bail!("上传的 HTML 文件过小（{} 字节），可能是空白或错误页面", text.len()); + anyhow::bail!( + "上传的 HTML 文件过小（{} 字节），可能是空白或错误页面", + text.len() + ); } let lower = text.to_lowercase(); // 仅拒绝明确的重定向占位页（通常 body 极短且没有正文） - let is_redirect = lower.contains("http-equiv=\"refresh\"") - || lower.contains("autoredirecttourl"); + let is_redirect = + lower.contains("http-equiv=\"refresh\"") || lower.contains("autoredirecttourl"); if is_redirect && text.len() < 5000 { anyhow::bail!("检测到 HTML 重定向跳转页面，而非真实文献正文"); } @@ -274,14 +327,20 @@ pub struct Downloader { } impl Downloader { - pub fn new() -> Self { + pub fn new() -> anyhow::Result { let mut headers = HeaderMap::new(); headers.insert(reqwest::header::ACCEPT, HeaderValue::from_static( "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7", )); - headers.insert(reqwest::header::ACCEPT_LANGUAGE, HeaderValue::from_static("en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7")); + headers.insert( + reqwest::header::ACCEPT_LANGUAGE, + HeaderValue::from_static("en-US,en;q=0.9,zh-CN;q=0.8,zh;q=0.7"), + ); headers.insert("DNT", HeaderValue::from_static("1")); - headers.insert(reqwest::header::CONNECTION, HeaderValue::from_static("keep-alive")); + headers.insert( + reqwest::header::CONNECTION, + HeaderValue::from_static("keep-alive"), + ); headers.insert("Upgrade-Insecure-Requests", HeaderValue::from_static("1")); let client = reqwest::Client::builder() @@ -291,15 +350,15 @@ impl Downloader { .redirect(reqwest::redirect::Policy::limited(10)) .timeout(std::time::Duration::from_secs(60)) .build() - .expect("Failed to create HTTP client"); + .context("Failed to create HTTP client")?; - Downloader { client } + Ok(Downloader { client }) } /// 使用 Obscura 作为后备通道进行下载 async fn download_via_obscura(&self, url: &str, dest_path: &Path, is_pdf: bool) -> Result<()> { info!("[Obscura 后备通道] 启动下载: {}", url); - + if let Some(parent) = dest_path.parent() { std::fs::create_dir_all(parent)?; } @@ -307,7 +366,8 @@ impl Downloader { #[cfg(feature = "obscura-inprocess")] { info!("[Obscura 后备通道] 正在运行进程内浏览器进行下载..."); - self.download_via_inprocess_obscura(url, dest_path, is_pdf).await + self.download_via_inprocess_obscura(url, dest_path, is_pdf) + .await } #[cfg(not(feature = "obscura-inprocess"))] @@ -318,7 +378,12 @@ impl Downloader { } #[cfg(feature = "obscura-inprocess")] - async fn download_via_inprocess_obscura(&self, url: &str, dest_path: &Path, is_pdf: bool) -> Result<()> { + async fn download_via_inprocess_obscura( + &self, + url: &str, + dest_path: &Path, + is_pdf: bool, + ) -> Result<()> { let url_str = url.to_string(); let dest_path_buf = dest_path.to_path_buf(); @@ -329,7 +394,7 @@ impl Downloader { .map_err(|e| anyhow::anyhow!("建立当前线程运行时失败: {}", e))?; rt.block_on(async move { - use obscura_browser::{BrowserContext, Page, lifecycle::WaitUntil}; + use obscura_browser::{lifecycle::WaitUntil, BrowserContext, Page}; use std::sync::Arc; // 1. 初始化启用 Stealth 防检测模式的浏览器上下文 @@ -343,7 +408,8 @@ impl Downloader { let mut page = Page::new("fetch-page".to_string(), context.clone()); // 2. 导航至目标 URL 并等待事件循环静默 - page.navigate_with_wait(&url_str, WaitUntil::Load).await + page.navigate_with_wait(&url_str, WaitUntil::Load) + .await .map_err(|e| anyhow::anyhow!("导航失败: {:?}", e))?; page.settle(5000).await; // 额外静默等待 5 秒 @@ -353,16 +419,19 @@ impl Downloader { // 对于 PDF 二进制文件，直接复用该浏览器上下文自带的 HTTP 客户端进行请求， // 这样能确保携带相同的 Cookie 和 TLS 指纹会话 let parsed_url = Url::parse(&url_str)?; - let response = page.http_client.fetch(&parsed_url).await + let response = page + .http_client + .fetch(&parsed_url) + .await .map_err(|e| anyhow::anyhow!("获取 PDF 字节流失败: {:?}", e))?; - + std::fs::write(&dest_path_buf, &response.body)?; validate_pdf_content(&response.body)?; } else { // 对于 HTML，直接从 V8 中提取 outerHTML let val = page.evaluate("document.documentElement.outerHTML"); let html = val.as_str().unwrap_or("").to_string(); - + std::fs::write(&dest_path_buf, &html)?; validate_html_content(&html)?; } @@ -389,26 +458,33 @@ impl Downloader { } #[cfg(not(feature = "obscura-inprocess"))] - async fn download_via_cli_obscura(&self, url: &str, dest_path: &Path, is_pdf: bool) -> Result<()> { + async fn download_via_cli_obscura( + &self, + url: &str, + dest_path: &Path, + is_pdf: bool, + ) -> Result<()> { let mut cmd = tokio::process::Command::new("bin/obscura"); cmd.arg("fetch").arg(url).arg("--stealth"); - + if is_pdf { cmd.arg("--dump").arg("original"); } else { cmd.arg("--dump").arg("html"); } - + cmd.arg("--output").arg(dest_path); - - let status = cmd.status().await + + let status = cmd + .status() + .await .context("启动 Obscura 进程失败，请检查 bin/obscura 是否存在且有执行权限")?; - + if !status.success() { let _ = tokio::fs::remove_file(dest_path).await; // 清理可能的残留坏文件 anyhow::bail!("Obscura 进程退出状态非成功: {:?}", status); } - + // 校验下载得到的文件 if is_pdf { let bytes = tokio::fs::read(dest_path).await?; @@ -423,7 +499,7 @@ impl Downloader { return Err(e); } } - + info!("[Obscura 命令行后备通道] 下载并校验成功: {:?}", dest_path); Ok(()) } @@ -464,7 +540,11 @@ impl Downloader { // HEAD 请求跟踪重定向（部分出版商阻断 HEAD，自动降级 GET） let response = match self.client.head(gateway_url).send().await { Ok(resp) => resp, - Err(_) => self.client.get(gateway_url).send().await + Err(_) => self + .client + .get(gateway_url) + .send() + .await .context(format!("请求 ADS 网关失败: {}", gateway_url))?, }; @@ -474,7 +554,11 @@ impl Downloader { // 如重定向至 validate.perfdrive.com，提取 ssc 参数中的真实 URL if final_url.contains("validate.perfdrive.com") { if let Ok(parsed) = Url::parse(&final_url) { - if let Some(ssc) = parsed.query_pairs().find(|(k, _)| k == "ssc").map(|(_, v)| v.into_owned()) { + if let Some(ssc) = parsed + .query_pairs() + .find(|(k, _)| k == "ssc") + .map(|(_, v)| v.into_owned()) + { if let Ok(decoded) = urlencoding::decode(&ssc) { let real_url = decoded.into_owned(); debug!("检测到 perfdrive 拦截，解码真实地址: {}", real_url); @@ -493,6 +577,7 @@ impl Downloader { } /// 读取文件前 512 字节用于内容嗅探 + #[allow(dead_code)] async fn read_file_header(path: &Path) -> Result> { use tokio::io::AsyncReadExt; let mut file = tokio::fs::File::open(path).await?; @@ -514,9 +599,12 @@ impl Downloader { // 步骤 1：访问文章主页，建立 Cookie 会话 debug!("[IOP] 预热主页: {}", main_url); Self::maybe_delay().await; - match self.client.get(&main_url) + match self + .client + .get(&main_url) .headers(build_chrome_headers(None)) - .send().await + .send() + .await { Ok(r) => debug!("[IOP] 主页响应: {}", r.status()), Err(e) => warn!("[IOP] 主页访问失败（继续尝试）: {:?}", e), @@ -525,9 +613,12 @@ impl Downloader { // 步骤 2：携带 Referer 下载 PDF debug!("[IOP] 下载 PDF: {}", pdf_url); Self::maybe_delay().await; - let response = self.client.get(&pdf_url) + let response = self + .client + .get(&pdf_url) .headers(build_chrome_headers(Some(&main_url))) - .send().await + .send() + .await .context("IOP PDF 请求失败")?; let status = response.status(); @@ -541,7 +632,8 @@ impl Downloader { let bytes = tokio::fs::read(dest_path).await?; validate_pdf_content(&bytes)?; Ok(()) - }.await; + } + .await; match res { Ok(()) => { @@ -550,14 +642,17 @@ impl Downloader { } Err(e) => { let err_msg = e.to_string(); - if err_msg.contains("人机验证") - || err_msg.contains("挑战页面") + if err_msg.contains("人机验证") + || err_msg.contains("挑战页面") || err_msg.contains("WAF") || err_msg.contains("Cloudflare") || err_msg.contains("HTTP 403") || err_msg.contains("HTTP 503") { - warn!("[IOP] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", err_msg); + warn!( + "[IOP] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", + err_msg + ); self.download_via_obscura(&pdf_url, dest_path, true).await } else { Err(e) @@ -573,9 +668,12 @@ impl Downloader { let res = async { Self::maybe_delay().await; - let response = self.client.get(&url) + let response = self + .client + .get(&url) .headers(build_browser_headers()) - .send().await + .send() + .await .context("Springer HTML 请求失败")?; let status = response.status(); @@ -585,11 +683,13 @@ impl Downloader { self.stream_download(response, dest_path).await?; - let text = tokio::fs::read_to_string(dest_path).await + let text = tokio::fs::read_to_string(dest_path) + .await .context("读取 HTML 文件失败")?; validate_html_content(&text)?; Ok(()) - }.await; + } + .await; match res { Ok(()) => { @@ -598,14 +698,17 @@ impl Downloader { } Err(e) => { let err_msg = e.to_string(); - if err_msg.contains("人机验证") - || err_msg.contains("挑战页面") + if err_msg.contains("人机验证") + || err_msg.contains("挑战页面") || err_msg.contains("WAF") || err_msg.contains("Cloudflare") || err_msg.contains("HTTP 403") || err_msg.contains("HTTP 503") { - warn!("[Springer] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", err_msg); + warn!( + "[Springer] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", + err_msg + ); self.download_via_obscura(&url, dest_path, false).await } else { Err(e) @@ -620,9 +723,12 @@ impl Downloader { Self::maybe_delay().await; let res = async { - let response = self.client.get(url) + let response = self + .client + .get(url) .headers(build_browser_headers()) - .send().await + .send() + .await .context(format!("[{}] PDF 请求失败", label))?; let status = response.status(); @@ -635,7 +741,8 @@ impl Downloader { let bytes = tokio::fs::read(dest_path).await?; validate_pdf_content(&bytes)?; Ok(()) - }.await; + } + .await; match res { Ok(()) => { @@ -645,14 +752,17 @@ impl Downloader { Err(e) => { let _ = tokio::fs::remove_file(dest_path).await; // 清理直连失败的残留物理文件 let err_msg = e.to_string(); - if err_msg.contains("人机验证") - || err_msg.contains("挑战页面") + if err_msg.contains("人机验证") + || err_msg.contains("挑战页面") || err_msg.contains("WAF") || err_msg.contains("Cloudflare") || err_msg.contains("HTTP 403") || err_msg.contains("HTTP 503") { - warn!("[{}] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", label, err_msg); + warn!( + "[{}] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", + label, err_msg + ); self.download_via_obscura(url, dest_path, true).await } else { Err(e) @@ -667,9 +777,12 @@ impl Downloader { Self::maybe_delay().await; let res = async { - let response = self.client.get(url) + let response = self + .client + .get(url) .headers(build_browser_headers()) - .send().await + .send() + .await .context(format!("[{}] HTML 请求失败", label))?; let status = response.status(); @@ -679,11 +792,13 @@ impl Downloader { self.stream_download(response, dest_path).await?; - let text = tokio::fs::read_to_string(dest_path).await + let text = tokio::fs::read_to_string(dest_path) + .await .context("读取 HTML 文件失败")?; validate_html_content(&text)?; Ok(()) - }.await; + } + .await; match res { Ok(()) => { @@ -692,14 +807,17 @@ impl Downloader { } Err(e) => { let err_msg = e.to_string(); - if err_msg.contains("人机验证") - || err_msg.contains("挑战页面") + if err_msg.contains("人机验证") + || err_msg.contains("挑战页面") || err_msg.contains("WAF") || err_msg.contains("Cloudflare") || err_msg.contains("HTTP 403") || err_msg.contains("HTTP 503") { - warn!("[{}] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", label, err_msg); + warn!( + "[{}] 下载触发人机验证或拦截: {}。尝试使用 Obscura 后备通道...", + label, err_msg + ); self.download_via_obscura(url, dest_path, false).await } else { Err(e) @@ -715,17 +833,23 @@ impl Downloader { let api_url = format!("https://api.crossref.org/works/{}", doi); info!("[CrossRef] 查询 PDF 链接: {}", api_url); - let data: serde_json::Value = self.client.get(&api_url) + let data: serde_json::Value = self + .client + .get(&api_url) .header("Accept", "application/json") - .send().await + .send() + .await .context("CrossRef API 请求失败")? - .json().await + .json() + .await .context("CrossRef API 响应解析失败")?; - let links = data["message"]["link"].as_array() + let links = data["message"]["link"] + .as_array() .context("CrossRef 未返回 link 数组")?; - let pdf_url = links.iter() + let pdf_url = links + .iter() .find(|l| { let ct = l["content-type"].as_str().unwrap_or(""); ct.contains("pdf") || ct == "unspecified" @@ -734,7 +858,8 @@ impl Downloader { .context("CrossRef 未找到 PDF 链接")?; info!("[CrossRef] PDF 链接: {}", pdf_url); - self.download_pdf_direct(pdf_url, dest_path, "CrossRef").await + self.download_pdf_direct(pdf_url, dest_path, "CrossRef") + .await } // ─── 公共入口 ────────────────────────────────────────────── @@ -744,7 +869,11 @@ impl Downloader { /// HTML 下载优先级： /// 1. 官方 `arxiv.org/html/{id}`（2023-12 起支持，质量与 ar5iv 相同，更稳定） /// 2. ar5iv `ar5iv.labs.arxiv.org/html/{id}`（约 3% 论文转换失败时跳过） - pub async fn download_arxiv_direct(&self, arxiv_id: &str, library_dir: &Path) -> (Result, Result) { + pub async fn download_arxiv_direct( + &self, + arxiv_id: &str, + library_dir: &Path, + ) -> (Result, Result) { // 去除版本号（v1/v2/v3），arxiv.org/html/ 和 ar5iv 均只提供最新渲染版 let clean_id = strip_arxiv_version(arxiv_id); @@ -764,13 +893,19 @@ impl Downloader { // HTML 下载：官方 arxiv.org/html/ 优先 let official_html_url = format!("https://arxiv.org/html/{}", clean_id); - let html_res = match self.download_html_direct(&official_html_url, &html_dest, "arXiv-HTML").await { + let html_res = match self + .download_html_direct(&official_html_url, &html_dest, "arXiv-HTML") + .await + { Ok(_) => Ok(html_dest.clone()), Err(e) => { warn!("[arXiv-HTML] 官方 HTML 下载失败，回退 ar5iv: {:?}", e); // ar5iv 兜底：约 97% 成功率，可能有延迟 let ar5iv_url = format!("https://ar5iv.labs.arxiv.org/html/{}", clean_id); - match self.download_html_direct(&ar5iv_url, &html_dest, "ar5iv").await { + match self + .download_html_direct(&ar5iv_url, &html_dest, "ar5iv") + .await + { Ok(_) => Ok(html_dest), Err(e2) => { let err_msg = format!("arXiv HTML 下载失败 (官方: {}, ar5iv: {})", e, e2); @@ -786,14 +921,22 @@ impl Downloader { /// 下载 arXiv HTML：官方 arxiv.org/html/ 优先，ar5iv 兜底 /// arxiv_id 应已去除版本号 - async fn download_arxiv_html_with_fallback(&self, arxiv_id: &str, dest_path: &Path) -> Result<()> { + async fn download_arxiv_html_with_fallback( + &self, + arxiv_id: &str, + dest_path: &Path, + ) -> Result<()> { let official_url = format!("https://arxiv.org/html/{}", arxiv_id); - match self.download_html_direct(&official_url, dest_path, "arXiv-HTML").await { + match self + .download_html_direct(&official_url, dest_path, "arXiv-HTML") + .await + { Ok(()) => Ok(()), Err(e) => { warn!("[arXiv-HTML] 官方 HTML 失败，回退 ar5iv: {:?}", e); let ar5iv_url = format!("https://ar5iv.labs.arxiv.org/html/{}", arxiv_id); - self.download_html_direct(&ar5iv_url, dest_path, "ar5iv").await + self.download_html_direct(&ar5iv_url, dest_path, "ar5iv") + .await } } } @@ -808,7 +951,12 @@ impl Downloader { /// HTML 回退顺序： /// 1. ADS PUB_HTML 网关（IOP→ 直联 iopscience，arxiv abs → ar5iv） /// 2. ADS EPRINT_HTML 网关（arxiv abs → ar5iv） - pub async fn download_paper(&self, bibcode: &str, doi: Option<&str>, library_dir: &Path) -> (Result, Result) { + pub async fn download_paper( + &self, + bibcode: &str, + doi: Option<&str>, + library_dir: &Path, + ) -> (Result, Result) { let base = "https://ui.adsabs.harvard.edu/link_gateway"; let pdf_dest = library_dir.join("PDF").join(format!("{}.pdf", bibcode)); let html_dest = library_dir.join("HTML").join(format!("{}.html", bibcode)); @@ -834,14 +982,21 @@ impl Downloader { .trim_end_matches("/pdf") .trim_end_matches('/'); self.download_iop_pdf(doi, &pdf_dest).await - } else if resolved.contains("link.springer.com") || resolved.contains("nature.com") { + } else if resolved.contains("link.springer.com") + || resolved.contains("nature.com") + { // Springer/Nature：HTML 更可靠，PDF 用通用策略 - self.download_pdf_direct(&resolved, &pdf_dest, "Springer").await + self.download_pdf_direct(&resolved, &pdf_dest, "Springer") + .await } else { - self.download_pdf_direct(&resolved, &pdf_dest, "PUB_PDF").await + self.download_pdf_direct(&resolved, &pdf_dest, "PUB_PDF") + .await }; match result { - Ok(_) => { pdf_res = Ok(pdf_dest.clone()); break 'pdf; } + Ok(_) => { + pdf_res = Ok(pdf_dest.clone()); + break 'pdf; + } Err(e) => { let msg = format!("PUB_PDF下载失败: {}", e); warn!("{}", msg); @@ -860,8 +1015,14 @@ impl Downloader { let gw = format!("{}/{}/ADS_PDF", base, bibcode); match self.resolve_ads_gateway(&gw).await { Ok(resolved) => { - match self.download_pdf_direct(&resolved, &pdf_dest, "ADS_PDF").await { - Ok(_) => { pdf_res = Ok(pdf_dest.clone()); break 'pdf; } + match self + .download_pdf_direct(&resolved, &pdf_dest, "ADS_PDF") + .await + { + Ok(_) => { + pdf_res = Ok(pdf_dest.clone()); + break 'pdf; + } Err(e) => { let msg = format!("ADS_PDF下载失败: {}", e); warn!("{}", msg); @@ -880,8 +1041,14 @@ impl Downloader { let gw = format!("{}/{}/EPRINT_PDF", base, bibcode); match self.resolve_ads_gateway(&gw).await { Ok(resolved) => { - match self.download_pdf_direct(&resolved, &pdf_dest, "EPRINT_PDF").await { - Ok(_) => { pdf_res = Ok(pdf_dest.clone()); break 'pdf; } + match self + .download_pdf_direct(&resolved, &pdf_dest, "EPRINT_PDF") + .await + { + Ok(_) => { + pdf_res = Ok(pdf_dest.clone()); + break 'pdf; + } Err(e) => { let msg = format!("EPRINT_PDF下载失败: {}", e); warn!("{}", msg); @@ -899,7 +1066,10 @@ impl Downloader { // 1c. CrossRef API 回退（需要 DOI） if let Some(doi_str) = doi { match self.download_crossref_pdf(doi_str, &pdf_dest).await { - Ok(_) => { pdf_res = Ok(pdf_dest.clone()); break 'pdf; } + Ok(_) => { + pdf_res = Ok(pdf_dest.clone()); + break 'pdf; + } Err(e) => { let msg = format!("CrossRef下载失败: {}", e); warn!("{}", msg); @@ -910,8 +1080,13 @@ impl Downloader { // 1d. ADS SCAN 扫描版文献直接合并下载 PDF（主要针对早期/不可下载直接 PDF 的文献） let scan_url = format!("https://articles.adsabs.harvard.edu/cgi-bin/nph-iarticle_query?bibcode={}&db_key=AST&data_type=PDF_HIGH", bibcode); - match self.download_pdf_direct(&scan_url, &pdf_dest, "ADS_SCAN").await { - Ok(_) => { pdf_res = Ok(pdf_dest.clone()); } + match self + .download_pdf_direct(&scan_url, &pdf_dest, "ADS_SCAN") + .await + { + Ok(_) => { + pdf_res = Ok(pdf_dest.clone()); + } Err(e) => { let msg = format!("ADS_SCAN下载失败: {}", e); warn!("{}", msg); @@ -932,7 +1107,9 @@ impl Downloader { let gw = format!("{}/{}/PUB_HTML", base, bibcode); match self.resolve_ads_gateway(&gw).await { Ok(resolved) => { - let result = if resolved.contains("link.springer.com") || resolved.contains("nature.com") { + let result = if resolved.contains("link.springer.com") + || resolved.contains("nature.com") + { // Springer/Nature 专属 HTML 策略 let doi_part = resolved .trim_start_matches("https://link.springer.com/article/") @@ -941,12 +1118,17 @@ impl Downloader { self.download_springer_html(doi_part, &html_dest).await } else if let Some(arxiv_id) = extract_arxiv_id_from_url(&resolved) { // ADS 网关指向 arXiv abs 页面 → 优先官方 HTML，ar5iv 兜底 - self.download_arxiv_html_with_fallback(&arxiv_id, &html_dest).await + self.download_arxiv_html_with_fallback(&arxiv_id, &html_dest) + .await } else { - self.download_html_direct(&resolved, &html_dest, "PUB_HTML").await + self.download_html_direct(&resolved, &html_dest, "PUB_HTML") + .await }; match result { - Ok(_) => { html_res = Ok(html_dest.clone()); break 'html; } + Ok(_) => { + html_res = Ok(html_dest.clone()); + break 'html; + } Err(e) => { let msg = format!("PUB_HTML下载失败: {}", e); warn!("{}", msg); @@ -966,12 +1148,16 @@ impl Downloader { match self.resolve_ads_gateway(&gw).await { Ok(resolved) => { let result = if let Some(arxiv_id) = extract_arxiv_id_from_url(&resolved) { - self.download_arxiv_html_with_fallback(&arxiv_id, &html_dest).await + self.download_arxiv_html_with_fallback(&arxiv_id, &html_dest) + .await } else { - self.download_html_direct(&resolved, &html_dest, "EPRINT_HTML").await + self.download_html_direct(&resolved, &html_dest, "EPRINT_HTML") + .await }; match result { - Ok(_) => { html_res = Ok(html_dest.clone()); } + Ok(_) => { + html_res = Ok(html_dest.clone()); + } Err(e) => { let msg = format!("EPRINT_HTML下载失败: {}", e); warn!("{}", msg); @@ -1002,8 +1188,7 @@ impl Downloader { bibcode: &str, force: bool, ) -> anyhow::Result { - let paper = get_paper_from_db(db, library_dir, bibcode) - .await?; + let paper = get_paper_from_db(db, library_dir, bibcode).await?; if force { sqlx::query("UPDATE papers SET pdf_path = NULL, html_path = NULL WHERE bibcode = ?") @@ -1027,24 +1212,38 @@ impl Downloader { // 调用底层的 download pipeline let (pdf_res, html_res) = if !paper.arxiv_id.is_empty() { info!("[下载] 优先使用 arXiv 通道: {}", paper.arxiv_id); - let res = self.download_arxiv_direct(&paper.arxiv_id, library_dir).await; + let res = self + .download_arxiv_direct(&paper.arxiv_id, library_dir) + .await; if res.0.is_ok() || res.1.is_ok() { res } else { - warn!("[下载] arXiv 通道下载失败，开始回退至 ADS/出版商通道: {}", bibcode); - let doi_opt = if !paper.doi.is_empty() { Some(paper.doi.as_str()) } else { None }; + warn!( + "[下载] arXiv 通道下载失败，开始回退至 ADS/出版商通道: {}", + bibcode + ); + let doi_opt = if !paper.doi.is_empty() { + Some(paper.doi.as_str()) + } else { + None + }; self.download_paper(bibcode, doi_opt, library_dir).await } } else { - let doi_opt = if !paper.doi.is_empty() { Some(paper.doi.as_str()) } else { None }; + let doi_opt = if !paper.doi.is_empty() { + Some(paper.doi.as_str()) + } else { + None + }; self.download_paper(bibcode, doi_opt, library_dir).await }; - if pdf_res.is_err() && html_res.is_err() { - let pdf_err = pdf_res.as_ref().err().unwrap(); - let html_err = html_res.as_ref().err().unwrap(); - error!("文献 {} PDF 和 HTML 均下载失败，无可用物理文件格式", bibcode); - + if let (Err(ref pdf_err), Err(ref html_err)) = (&pdf_res, &html_res) { + error!( + "文献 {} PDF 和 HTML 均下载失败，无可用物理文件格式", + bibcode + ); + let pdf_db_err = format!("error: {}", pdf_err); let html_db_err = format!("error: {}", html_err); let _ = sqlx::query("UPDATE papers SET pdf_path = ?, html_path = ? WHERE bibcode = ?") @@ -1054,15 +1253,29 @@ impl Downloader { .execute(db) .await; - return Err(anyhow::anyhow!("下载失败。PDF: {}, HTML: {}", pdf_err, html_err)); + return Err(anyhow::anyhow!( + "下载失败。PDF: {}, HTML: {}", + pdf_err, + html_err + )); } let pdf_rel = match pdf_res { - Ok(p) => Some(p.strip_prefix(library_dir).unwrap_or(&p).to_string_lossy().to_string()), + Ok(p) => Some( + p.strip_prefix(library_dir) + .unwrap_or(&p) + .to_string_lossy() + .to_string(), + ), Err(_) => None, }; let html_rel = match html_res { - Ok(p) => Some(p.strip_prefix(library_dir).unwrap_or(&p).to_string_lossy().to_string()), + Ok(p) => Some( + p.strip_prefix(library_dir) + .unwrap_or(&p) + .to_string_lossy() + .to_string(), + ), Err(_) => None, }; @@ -1101,8 +1314,13 @@ fn extract_arxiv_id_from_url(url: &str) -> Option { for pat in &patterns { if let Some(pos) = url.find(pat) { let id_raw = &url[pos + pat.len()..]; - let mut id_clean = id_raw.split('?').next().unwrap_or(id_raw) - .split('#').next().unwrap_or(id_raw) + let mut id_clean = id_raw + .split('?') + .next() + .unwrap_or(id_raw) + .split('#') + .next() + .unwrap_or(id_raw) .trim_end_matches('/') .to_string(); if id_clean.to_lowercase().ends_with(".pdf") { @@ -1120,14 +1338,10 @@ fn extract_arxiv_id_from_url(url: &str) -> Option { None } - - - - #[cfg(test)] mod tests { use super::*; - use axum::{Router, routing::get, response::Redirect}; + use axum::{response::Redirect, routing::get, Router}; #[tokio::test] async fn test_resolve_ads_gateway_perfdrive() { @@ -1138,18 +1352,20 @@ mod tests { let target_ssc = "https%3A%2F%2Fexample.com%2Ftarget.pdf"; let redirect_to = format!("https://validate.perfdrive.com/?ssc={}", target_ssc); - let app = Router::new().route("/gate", get(move || { - let r = redirect_to.clone(); - async move { Redirect::to(&r) } - })); - - let server = axum::serve( - tokio::net::TcpListener::from_std(listener).unwrap(), - app, + let app = Router::new().route( + "/gate", + get(move || { + let r = redirect_to.clone(); + async move { Redirect::to(&r) } + }), ); - tokio::spawn(async move { let _ = server.await; }); - let downloader = Downloader::new(); + let server = axum::serve(tokio::net::TcpListener::from_std(listener).unwrap(), app); + tokio::spawn(async move { + let _ = server.await; + }); + + let downloader = Downloader::new().expect("Failed to create downloader in test"); let gateway_url = format!("http://127.0.0.1:{}/gate", port); let result = downloader.resolve_ads_gateway(&gateway_url).await; assert_eq!(result.unwrap(), "https://example.com/target.pdf"); @@ -1178,7 +1394,10 @@ mod tests { #[test] fn test_detect_anti_bot_clean() { - let result = detect_anti_bot("

      Abstract

      We study...

      ", None); + let result = detect_anti_bot( + "

      Abstract

      We study...

      ", + None, + ); assert!(result.is_ok()); } @@ -1208,19 +1427,19 @@ mod tests { #[tokio::test] #[ignore] async fn test_download_scan_pdf() -> anyhow::Result<()> { - let downloader = Downloader::new(); + let downloader = Downloader::new().expect("Failed to create downloader in test"); let bibcode = "2005MNRAS.359..315E"; let temp_dir = std::env::temp_dir(); - + let (pdf_path, _html_path) = downloader.download_paper(bibcode, None, &temp_dir).await; assert!(pdf_path.is_ok()); - + let path = pdf_path.unwrap(); assert!(path.exists()); - + let bytes = std::fs::read(&path)?; assert!(bytes.starts_with(b"%PDF")); - + let _ = std::fs::remove_file(&path); Ok(()) } @@ -1228,9 +1447,9 @@ mod tests { #[tokio::test] #[ignore] async fn test_download_via_obscura_integration() -> anyhow::Result<()> { - use axum::{Router, routing::get, response::Response as AxumResponse}; - use axum::http::{HeaderValue, header::CONTENT_TYPE}; - + use axum::http::{header::CONTENT_TYPE, HeaderValue}; + use axum::{response::Response as AxumResponse, routing::get, Router}; + // Bind to a random port let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); listener.set_nonblocking(true).unwrap(); @@ -1259,22 +1478,23 @@ mod tests { .unwrap() })); - let server = axum::serve( - tokio::net::TcpListener::from_std(listener).unwrap(), - app, - ); - tokio::spawn(async move { let _ = server.await; }); + let server = axum::serve(tokio::net::TcpListener::from_std(listener).unwrap(), app); + tokio::spawn(async move { + let _ = server.await; + }); // Temporarily set OBSCURA_ALLOW_PRIVATE_NETWORK=1 to allow loopback fetches in Obscura std::env::set_var("OBSCURA_ALLOW_PRIVATE_NETWORK", "1"); - let downloader = Downloader::new(); + let downloader = Downloader::new().expect("Failed to create downloader in test"); let temp_dir = std::env::temp_dir(); - + // 1. Test HTML download via obscura let html_dest = temp_dir.join("test_obscura_mock.html"); let html_url = format!("http://127.0.0.1:{}/mock.html", port); - downloader.download_via_obscura(&html_url, &html_dest, false).await?; + downloader + .download_via_obscura(&html_url, &html_dest, false) + .await?; assert!(html_dest.exists()); let html_content = std::fs::read_to_string(&html_dest)?; assert!(html_content.contains("introduction")); @@ -1283,7 +1503,9 @@ mod tests { // 2. Test PDF download via obscura let pdf_dest = temp_dir.join("test_obscura_mock.pdf"); let pdf_url = format!("http://127.0.0.1:{}/mock.pdf", port); - downloader.download_via_obscura(&pdf_url, &pdf_dest, true).await?; + downloader + .download_via_obscura(&pdf_url, &pdf_dest, true) + .await?; assert!(pdf_dest.exists()); let pdf_content = std::fs::read(&pdf_dest)?; assert_eq!(pdf_content, pdf_data); @@ -1293,4 +1515,3 @@ mod tests { Ok(()) } } - diff --git a/src/services/logging.rs b/src/services/logging.rs index 41252c3..319144f 100644 --- a/src/services/logging.rs +++ b/src/services/logging.rs @@ -6,10 +6,7 @@ use tracing_appender::non_blocking::WorkerGuard; use tracing_appender::rolling; use tracing_subscriber::fmt::format::Writer; use tracing_subscriber::fmt::time::FormatTime; -use tracing_subscriber::{ - fmt, layer::SubscriberExt, util::SubscriberInitExt, - EnvFilter, Layer, -}; +use tracing_subscriber::{fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; pub struct ShanghaiTime; @@ -27,12 +24,13 @@ pub fn init_logging() -> anyhow::Result> { let mut guards = Vec::new(); // 从环境变量中读取配置 - let log_level = env::var("LOG_LEVEL").unwrap_or_else(|_| "info,astroresearch=debug".to_string()); + let log_level = + env::var("LOG_LEVEL").unwrap_or_else(|_| "info,astroresearch=debug".to_string()); let log_format = env::var("LOG_FORMAT").unwrap_or_else(|_| "pretty".to_string()); let log_outputs = env::var("LOG_OUTPUTS").unwrap_or_else(|_| "stdout,file".to_string()); - let env_filter = EnvFilter::try_from_default_env() - .unwrap_or_else(|_| EnvFilter::new(&log_level)); + let env_filter = + EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new(&log_level)); let is_json = log_format.to_lowercase() == "json"; let mut layers: Vec + Send + Sync>> = Vec::new(); diff --git a/src/services/mod.rs b/src/services/mod.rs index 579ebbf..6d1514b 100644 --- a/src/services/mod.rs +++ b/src/services/mod.rs @@ -1,13 +1,13 @@ -pub mod download; -pub mod parser; -pub mod translation; -pub mod query_parser; pub mod batch; -pub mod logging; pub mod chunker; +pub mod download; +pub mod logging; +pub mod parser; +pub mod query_parser; pub mod rag; -pub mod target; pub mod search; +pub mod target; +pub mod translation; pub mod batch_sync { pub use super::batch::*; diff --git a/src/services/parser/aanda.rs b/src/services/parser/aanda.rs index 94024dc..ae24399 100644 --- a/src/services/parser/aanda.rs +++ b/src/services/parser/aanda.rs @@ -14,7 +14,9 @@ impl AandaParser { } impl JournalParser for AandaParser { - fn name(&self) -> &str { "A&A" } + fn name(&self) -> &str { + "A&A" + } fn extract_body<'a>(&self, html: &'a str) -> &'a str { if let Some(start) = html.find("]*>.*?"#).unwrap().replace_all(&h, "").to_string(); + h = Regex::new(r#"(?s)]*>.*?"#) + .unwrap() + .replace_all(&h, "") + .to_string(); // JavaScript 邮件混淆代码 - h = Regex::new(r#"(?s)]*>.*?"#).unwrap().replace_all(&h, "").to_string(); + h = Regex::new(r#"(?s)]*>.*?"#) + .unwrap() + .replace_all(&h, "") + .to_string(); // A&A 特有的页面框架 for (pattern, replacement) in &[ - (r#"(?s)]*class="[^"]*breadcrumbs[^"]*"[^>]*>.*?"#, ""), // 面包屑 - (r#"(?s)]*class="[^"]*menu[^"]*"[^>]*id="bloc"[^>]*>.*?"#, ""), // 内部 TOC - (r#"(?s)]*class="[^"]*Z3988[^"]*"[^>]*>.*?"#, ""), // COinS 元数据 - (r#"(?s)]*class="[^"]*nav-(?:article|buttons)[^"]*"[^>]*>.*?"#, ""), // 期刊导航 - (r#"(?s)]*class="[^"]*special_article[^"]*"[^>]*>.*?"#, ""), // Open Access 徽章 + ( + r#"(?s)]*class="[^"]*breadcrumbs[^"]*"[^>]*>.*?"#, + "", + ), // 面包屑 + ( + r#"(?s)]*class="[^"]*menu[^"]*"[^>]*id="bloc"[^>]*>.*?"#, + "", + ), // 内部 TOC + ( + r#"(?s)]*class="[^"]*Z3988[^"]*"[^>]*>.*?"#, + "", + ), // COinS 元数据 + ( + r#"(?s)]*class="[^"]*nav-(?:article|buttons)[^"]*"[^>]*>.*?"#, + "", + ), // 期刊导航 + ( + r#"(?s)]*class="[^"]*special_article[^"]*"[^>]*>.*?"#, + "", + ), // Open Access 徽章 // 元数据表格：从 summary.full 删除到
      - (r#"(?s)]*class="[^"]*summary\s+full[^"]*"[^>]*>.*?]*class="[^"]*summary\s+full[^"]*"[^>]*>.*?]*class="[^"]*sec\b[^"]*"[^>]*>(.*?)

    "#) - .unwrap().replace_all(&h, |caps: ®ex::Captures| { + .unwrap() + .replace_all(&h, |caps: ®ex::Captures| { let inner = super::common::strip_html_tags(&caps[1]); format!("\n\n## {}\n\n", inner.trim()) - }).to_string(); + }) + .to_string(); // h3.sec2 → ### h = Regex::new(r#"(?s)]*class="[^"]*sec2\b[^"]*"[^>]*>(.*?)

"#) - .unwrap().replace_all(&h, |caps: ®ex::Captures| { + .unwrap() + .replace_all(&h, |caps: ®ex::Captures| { let inner = super::common::strip_html_tags(&caps[1]); format!("\n\n### {}\n\n", inner.trim()) - }).to_string(); + }) + .to_string(); h } @@ -90,18 +122,20 @@ impl JournalParser for AandaParser { ).unwrap(); let bo = base_origin; - inset_re.replace_all(html, |caps: ®ex::Captures| { - let img_src = &caps[1]; - let caption_html = &caps[2]; - let full_img_src = img_src.replace("_small.", "."); - let absolute_url = if full_img_src.starts_with('/') { - format!("{}{}", bo, full_img_src) - } else { - full_img_src.to_string() - }; - let caption = super::common::strip_html_preserve_links(caption_html); - let clean_caption = caption.trim().replace('\n', " "); - format!("\n\n![{}]({})\n\n", clean_caption, absolute_url) - }).to_string() + inset_re + .replace_all(html, |caps: ®ex::Captures| { + let img_src = &caps[1]; + let caption_html = &caps[2]; + let full_img_src = img_src.replace("_small.", "."); + let absolute_url = if full_img_src.starts_with('/') { + format!("{}{}", bo, full_img_src) + } else { + full_img_src.to_string() + }; + let caption = super::common::strip_html_preserve_links(caption_html); + let clean_caption = caption.trim().replace('\n', " "); + format!("\n\n![{}]({})\n\n", clean_caption, absolute_url) + }) + .to_string() } } diff --git a/src/services/parser/ar5iv.rs b/src/services/parser/ar5iv.rs index 31782e3..3f95bfe 100644 --- a/src/services/parser/ar5iv.rs +++ b/src/services/parser/ar5iv.rs @@ -12,7 +12,9 @@ impl Ar5ivParser { } impl JournalParser for Ar5ivParser { - fn name(&self) -> &str { "ar5iv" } + fn name(&self) -> &str { + "ar5iv" + } fn extract_body<'a>(&self, html: &'a str) -> &'a str { if let Some(start) = html.find("
") { @@ -48,11 +50,15 @@ impl JournalParser for Ar5ivParser { // N → ^{N} h = Regex::new(r#"]*class="[^"]*ltx_note_mark[^"]*"[^>]*>([^<]*)"#) - .unwrap().replace_all(&h, "^{$1}").to_string(); + .unwrap() + .replace_all(&h, "^{$1}") + .to_string(); // ... → ^{...} h = Regex::new(r#"]*class="[^"]*ltx_sup[^"]*"[^>]*>(.*?)"#) - .unwrap().replace_all(&h, "^{$1}").to_string(); + .unwrap() + .replace_all(&h, "^{$1}") + .to_string(); // text → [text](url)，供后续内部链接剥离 // 匹配 class 在 href 之前或之后两种顺序 @@ -60,7 +66,10 @@ impl JournalParser for Ar5ivParser { r#"]*class="[^"]*ltx_ref[^"]*"[^>]*href="([^"]*)"[^>]*>(.*?)"#, r#"]*href="([^"]*)"[^>]*class="[^"]*ltx_ref[^"]*"[^>]*>(.*?)"#, ] { - h = Regex::new(pat).unwrap().replace_all(&h, "[$2]($1)").to_string(); + h = Regex::new(pat) + .unwrap() + .replace_all(&h, "[$2]($1)") + .to_string(); } h @@ -70,25 +79,37 @@ impl JournalParser for Ar5ivParser { let mut h = html.to_string(); // ltx_title_document → # 一级标题 - h = Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_document[^"]*"[^>]*>(.*?)"#) - .unwrap().replace_all(&h, |caps: ®ex::Captures| { - let inner = super::common::strip_html_tags(&caps[1]); - format!("\n# {}\n\n", inner.trim()) - }).to_string(); + h = Regex::new( + r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_document[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() + .replace_all(&h, |caps: ®ex::Captures| { + let inner = super::common::strip_html_tags(&caps[1]); + format!("\n# {}\n\n", inner.trim()) + }) + .to_string(); // ltx_title_section → ## 二级标题 - h = Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_section[^"]*"[^>]*>(.*?)"#) - .unwrap().replace_all(&h, |caps: ®ex::Captures| { - let inner = super::common::strip_html_tags(&caps[1]); - format!("\n\n## {}\n\n", inner.trim()) - }).to_string(); + h = Regex::new( + r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_section[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() + .replace_all(&h, |caps: ®ex::Captures| { + let inner = super::common::strip_html_tags(&caps[1]); + format!("\n\n## {}\n\n", inner.trim()) + }) + .to_string(); // ltx_title_subsection → ### 三级标题 - h = Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsection[^"]*"[^>]*>(.*?)"#) - .unwrap().replace_all(&h, |caps: ®ex::Captures| { - let inner = super::common::strip_html_tags(&caps[1]); - format!("\n\n### {}\n\n", inner.trim()) - }).to_string(); + h = Regex::new( + r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsection[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() + .replace_all(&h, |caps: ®ex::Captures| { + let inner = super::common::strip_html_tags(&caps[1]); + format!("\n\n### {}\n\n", inner.trim()) + }) + .to_string(); // ltx_title_subsubsection → #### 四级标题 h = Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsubsection[^"]*"[^>]*>(.*?)"#) @@ -104,9 +125,11 @@ impl JournalParser for Ar5ivParser { // ar5iv 的
+
由 html2md 原生处理。 // 这里将
预处理为块引用标记以获得更好的格式。 let figcaption_re = Regex::new(r#"(?s)]*>(.*?)
"#).unwrap(); - figcaption_re.replace_all(html, |caps: ®ex::Captures| { - let inner = super::common::strip_html_tags(&caps[1]); - format!("\n\n> **Figure:** {}\n", inner.trim()) - }).to_string() + figcaption_re + .replace_all(html, |caps: ®ex::Captures| { + let inner = super::common::strip_html_tags(&caps[1]); + format!("\n\n> **Figure:** {}\n", inner.trim()) + }) + .to_string() } } diff --git a/src/services/parser/common.rs b/src/services/parser/common.rs index 9e96a4f..e5dddbb 100644 --- a/src/services/parser/common.rs +++ b/src/services/parser/common.rs @@ -1,41 +1,117 @@ // parser/common.rs — Shared utilities for HTML→Markdown. All regexes compiled once via LazyLock. -use std::sync::LazyLock; use regex::Regex; +use std::sync::LazyLock; // ── Static regexes ────────────────────────────────────────────────────────── // extract_html_base_origin -static BASE_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?i)]*href="([^"]*)""#).unwrap()); -static CANONICAL_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?i)]*rel="canonical"[^>]*href="([^"]*)"|]*href="([^"]*)"[^>]*rel="canonical""#).unwrap()); -static OG_URL_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?i)]*property="og:url"[^>]*content="([^"]*)"|]*content="([^"]*)"[^>]*property="og:url""#).unwrap()); -static PRISM_URL_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?i)]*name="prism\.url"[^>]*content="([^"]*)"|]*content="([^"]*)"[^>]*name="prism\.url""#).unwrap()); -static CITATION_URL_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?i)]*name="citation_pdf_url"[^>]*content="([^"]*)"|]*content="([^"]*)"[^>]*name="citation_pdf_url""#).unwrap()); +static BASE_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?i)]*href="([^"]*)""#).unwrap()); +static CANONICAL_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?i)]*rel="canonical"[^>]*href="([^"]*)"|]*href="([^"]*)"[^>]*rel="canonical""#).unwrap() +}); +static OG_URL_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?i)]*property="og:url"[^>]*content="([^"]*)"|]*content="([^"]*)"[^>]*property="og:url""#).unwrap() +}); +static PRISM_URL_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?i)]*name="prism\.url"[^>]*content="([^"]*)"|]*content="([^"]*)"[^>]*name="prism\.url""#).unwrap() +}); +static CITATION_URL_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?i)]*name="citation_pdf_url"[^>]*content="([^"]*)"|]*content="([^"]*)"[^>]*name="citation_pdf_url""#).unwrap() +}); // strip_html_tags static TAG_STRIP_RE: LazyLock = LazyLock::new(|| Regex::new(r"<[^>]+>").unwrap()); // strip_html_preserve_links -static A_LINK_RE: LazyLock = LazyLock::new(|| Regex::new(r#"]*href="([^"]*)"[^>]*>(.*?)"#).unwrap()); +static A_LINK_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"]*href="([^"]*)"[^>]*>(.*?)"#).unwrap()); // entity decoding — static HashMap + regex -static ENTITY_RE: LazyLock = LazyLock::new(|| Regex::new(r"&#[0-9]+;|&#x[0-9a-fA-F]+;|&[a-zA-Z]+;").unwrap()); -static ENTITY_MAP: LazyLock> = LazyLock::new(|| { - let mut m = std::collections::HashMap::new(); - m.insert("α","α");m.insert("β","β");m.insert("γ","γ");m.insert("δ","δ");m.insert("ε","ε");m.insert("ζ","ζ"); - m.insert("η","η");m.insert("θ","θ");m.insert("ι","ι");m.insert("κ","κ");m.insert("λ","λ");m.insert("μ","μ"); - m.insert("ν","ν");m.insert("ξ","ξ");m.insert("ο","ο");m.insert("π","π");m.insert("ρ","ρ");m.insert("ς","ς"); - m.insert("σ","σ");m.insert("τ","τ");m.insert("υ","υ");m.insert("φ","φ");m.insert("χ","χ");m.insert("ψ","ψ");m.insert("ω","ω"); - m.insert("Α","Α");m.insert("Β","Β");m.insert("Γ","Γ");m.insert("Δ","Δ");m.insert("Ε","Ε");m.insert("Ζ","Ζ"); - m.insert("Η","Η");m.insert("Θ","Θ");m.insert("Ι","Ι");m.insert("Κ","Κ");m.insert("Λ","Λ");m.insert("Μ","Μ"); - m.insert("Ν","Ν");m.insert("Ξ","Ξ");m.insert("Ο","Ο");m.insert("Π","Π");m.insert("Ρ","Ρ");m.insert("Σ","Σ"); - m.insert("Τ","Τ");m.insert("Υ","Υ");m.insert("Φ","Φ");m.insert("Χ","Χ");m.insert("Ψ","Ψ");m.insert("Ω","Ω"); - m.insert("≤","≤");m.insert("≥","≥");m.insert("−","−");m.insert("≈","≈");m.insert("⊙","⊙");m.insert("☉","⊙"); - m.insert(" "," ");m.insert(" "," ");m.insert("∑","∑");m.insert("×","×");m.insert("√","√");m.insert("∞","∞"); - m.insert("∫","∫");m.insert("–","–");m.insert("—","—");m.insert("’","’");m.insert("′","′");m.insert("″","″"); - m.insert("⁄","⁄");m.insert("€","€");m.insert("≠","≠");m.insert("⊂","⊂");m.insert("⊃","⊃"); - m.insert("⊆","⊆");m.insert("⊇","⊇");m.insert("<","<");m.insert(">",">");m.insert("&","&"); - m -}); +static ENTITY_RE: LazyLock = + LazyLock::new(|| Regex::new(r"&#[0-9]+;|&#x[0-9a-fA-F]+;|&[a-zA-Z]+;").unwrap()); +static ENTITY_MAP: LazyLock> = + LazyLock::new(|| { + let mut m = std::collections::HashMap::new(); + m.insert("α", "α"); + m.insert("β", "β"); + m.insert("γ", "γ"); + m.insert("δ", "δ"); + m.insert("ε", "ε"); + m.insert("ζ", "ζ"); + m.insert("η", "η"); + m.insert("θ", "θ"); + m.insert("ι", "ι"); + m.insert("κ", "κ"); + m.insert("λ", "λ"); + m.insert("μ", "μ"); + m.insert("ν", "ν"); + m.insert("ξ", "ξ"); + m.insert("ο", "ο"); + m.insert("π", "π"); + m.insert("ρ", "ρ"); + m.insert("ς", "ς"); + m.insert("σ", "σ"); + m.insert("τ", "τ"); + m.insert("υ", "υ"); + m.insert("φ", "φ"); + m.insert("χ", "χ"); + m.insert("ψ", "ψ"); + m.insert("ω", "ω"); + m.insert("Α", "Α"); + m.insert("Β", "Β"); + m.insert("Γ", "Γ"); + m.insert("Δ", "Δ"); + m.insert("Ε", "Ε"); + m.insert("Ζ", "Ζ"); + m.insert("Η", "Η"); + m.insert("Θ", "Θ"); + m.insert("Ι", "Ι"); + m.insert("Κ", "Κ"); + m.insert("Λ", "Λ"); + m.insert("Μ", "Μ"); + m.insert("Ν", "Ν"); + m.insert("Ξ", "Ξ"); + m.insert("Ο", "Ο"); + m.insert("Π", "Π"); + m.insert("Ρ", "Ρ"); + m.insert("Σ", "Σ"); + m.insert("Τ", "Τ"); + m.insert("Υ", "Υ"); + m.insert("Φ", "Φ"); + m.insert("Χ", "Χ"); + m.insert("Ψ", "Ψ"); + m.insert("Ω", "Ω"); + m.insert("≤", "≤"); + m.insert("≥", "≥"); + m.insert("−", "−"); + m.insert("≈", "≈"); + m.insert("⊙", "⊙"); + m.insert("☉", "⊙"); + m.insert(" ", " "); + m.insert(" ", " "); + m.insert("∑", "∑"); + m.insert("×", "×"); + m.insert("√", "√"); + m.insert("∞", "∞"); + m.insert("∫", "∫"); + m.insert("–", "–"); + m.insert("—", "—"); + m.insert("’", "’"); + m.insert("′", "′"); + m.insert("″", "″"); + m.insert("⁄", "⁄"); + m.insert("€", "€"); + m.insert("≠", "≠"); + m.insert("⊂", "⊂"); + m.insert("⊃", "⊃"); + m.insert("⊆", "⊆"); + m.insert("⊇", "⊇"); + m.insert("<", "<"); + m.insert(">", ">"); + m.insert("&", "&"); + m + }); // html_math_to_latex sub/sup inner regexes static ITALIC_INNER_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(.*?)"#).unwrap()); @@ -45,29 +121,45 @@ static SUP_INNER_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(.*?) // convert_html_math_to_latex — dynamic regex cached via OnceLock use std::sync::OnceLock; static MATH_REGEX: OnceLock = OnceLock::new(); -static ITALIC_CHECK_RE: LazyLock = LazyLock::new(|| Regex::new(r#"([^<]*)"#).unwrap()); +static ITALIC_CHECK_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"([^<]*)"#).unwrap()); // LaTeXML table conversion -static LATEXML_TAG_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?i)<(span|div)\b([^>]*?)>|"#).unwrap()); -static LATEXML_CLASS_RE: LazyLock = LazyLock::new(|| Regex::new(r#"class="([^"]*)""#).unwrap()); +static LATEXML_TAG_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?i)<(span|div)\b([^>]*?)>|"#).unwrap()); +static LATEXML_CLASS_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"class="([^"]*)""#).unwrap()); // Table conversion -static TABLE_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*>(.*?)"#).unwrap()); +static TABLE_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?s)]*>(.*?)"#).unwrap()); static TR_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*>(.*?)"#).unwrap()); -static TD_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)<(?:td|th)[^>]*>(.*?)"#).unwrap()); +static TD_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?s)<(?:td|th)[^>]*>(.*?)"#).unwrap()); // postprocess_markdown static DIV_RE: LazyLock = LazyLock::new(|| Regex::new(r"]*>").unwrap()); static SPAN_RE: LazyLock = LazyLock::new(|| Regex::new(r"]*>").unwrap()); -static IFRAME_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?s)]*>.*?").unwrap()); -static ASTROBJ_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?s)]*>(.*?)").unwrap()); -static FIG_MERGE_RE: LazyLock = LazyLock::new(|| Regex::new(r"(!\[[^\]]*\]\([^)]*\))\s*\\>\s*\\\*\\\*(Figure|Caption):\\\*\\\*").unwrap()); -static PAGERANGE_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?m)^\^\{†\}\^\{†\}(?:pagerange|pubyear|offprints):.*\n?").unwrap()); +static IFRAME_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?s)]*>.*?").unwrap()); +static ASTROBJ_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?s)]*>(.*?)").unwrap()); +static FIG_MERGE_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(!\[[^\]]*\]\([^)]*\))\s*\\>\s*\\\*\\\*(Figure|Caption):\\\*\\\*").unwrap() +}); +static PAGERANGE_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?m)^\^\{†\}\^\{†\}(?:pagerange|pubyear|offprints):.*\n?").unwrap() +}); static EMPTY_BRACKETS_RE: LazyLock = LazyLock::new(|| Regex::new(r"\[\]").unwrap()); -static INTERNAL_LINK_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(!?)\[([^\]]*?)\]\(([^)]*(?:/articles/aa/[^)]*|#[^)]*))\)"#).unwrap()); +static INTERNAL_LINK_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(!?)\[([^\]]*?)\]\(([^)]*(?:/articles/aa/[^)]*|#[^)]*))\)"#).unwrap() +}); static EXCESSIVE_NL_RE: LazyLock = LazyLock::new(|| Regex::new(r"\n{4,}").unwrap()); -static LINK_FIX_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(!?\[[^\]]*?\])\(([^)]*?)\)"#).unwrap()); -static LATEXML_MACRO_RE: LazyLock = LazyLock::new(|| Regex::new(r"\\{1,2}(?:orgname|orgdiv|orgaddress|articletag|term|savesymbol|restoresymbol|volnopage|SInits)\b").unwrap()); +static LINK_FIX_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(!?\[[^\]]*?\])\(([^)]*?)\)"#).unwrap()); +static LATEXML_MACRO_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"\\{1,2}(?:orgname|orgdiv|orgaddress|articletag|term|savesymbol|restoresymbol|volnopage|SInits)\b").unwrap() +}); /// 清理文首的 LaTeXML 模板垃圾行（如 `second\savesymboldegree\restoresymbol...`） fn clean_latexml_preamble(md: &str) -> String { @@ -75,7 +167,10 @@ fn clean_latexml_preamble(md: &str) -> String { // 只清理明确是 LaTeXML 命令残渣的行（含 \ 且不含 $ 或 #） while let Some(first) = lines.first() { let trimmed = first.trim(); - if trimmed.is_empty() { lines.remove(0); continue; } + if trimmed.is_empty() { + lines.remove(0); + continue; + } let has_slash = trimmed.contains('\\'); let has_math = trimmed.contains('$'); let is_heading = trimmed.starts_with('#'); @@ -88,13 +183,26 @@ fn clean_latexml_preamble(md: &str) -> String { } lines.join("\n") } -static HEADING_TRAIL_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?m)^(#{1,6})\s+(.*?)\s+#+$").unwrap()); -static SECTION_PROMOTE_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?mi)^(#{3,6})[ \t]*(Abstract|Keywords|Glossary|Nomenclature|Acknowledgments|References)(:?)[ \t]*$").unwrap()); -static ABSTRACT_CLEAN_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?mi)^##\s+Abstract\s*\n\s*\n\s*\[Abstract\]\s*\n").unwrap()); -static BRACKET_INLINE_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?mi)^\[(Abstract|Keywords|Glossary|Nomenclature|Acknowledgments|References)\][ \t]+(.+)$").unwrap()); -static BRACKET_HEADER_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?mi)^\[(Abstract|Keywords|Glossary|Nomenclature|Acknowledgments|References)\][ \t]*$").unwrap()); -static BULLET_CLEAN_RE: LazyLock = LazyLock::new(|| Regex::new(r"(?m)^(\s*[\*\-+])\s*•\s*").unwrap()); -static BRACKET_NEWLINE_RE: LazyLock = LazyLock::new(|| Regex::new(r"\[\s*\n+\s*([^\]\n]+)\]").unwrap()); +static HEADING_TRAIL_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?m)^(#{1,6})\s+(.*?)\s+#+$").unwrap()); +static SECTION_PROMOTE_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?mi)^(#{3,6})[ \t]*(Abstract|Keywords|Glossary|Nomenclature|Acknowledgments|References)(:?)[ \t]*$").unwrap() +}); +static ABSTRACT_CLEAN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?mi)^##\s+Abstract\s*\n\s*\n\s*\[Abstract\]\s*\n").unwrap()); +static BRACKET_INLINE_RE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?mi)^\[(Abstract|Keywords|Glossary|Nomenclature|Acknowledgments|References)\][ \t]+(.+)$").unwrap() +}); +static BRACKET_HEADER_RE: LazyLock = LazyLock::new(|| { + Regex::new( + r"(?mi)^\[(Abstract|Keywords|Glossary|Nomenclature|Acknowledgments|References)\][ \t]*$", + ) + .unwrap() +}); +static BULLET_CLEAN_RE: LazyLock = + LazyLock::new(|| Regex::new(r"(?m)^(\s*[\*\-+])\s*•\s*").unwrap()); +static BRACKET_NEWLINE_RE: LazyLock = + LazyLock::new(|| Regex::new(r"\[\s*\n+\s*([^\]\n]+)\]").unwrap()); // Unescape patterns static UNESC_RE: LazyLock> = LazyLock::new(|| { @@ -112,28 +220,60 @@ static UNESC_RE: LazyLock> = LazyLock::new(|| { static TRUNCATE_RE: OnceLock = OnceLock::new(); // extract_img_equations -static IMG_EQ_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*class="[^"]*img-equation[^"]*"[^>]*data-latex="([^"]*)"[^>]*>.*?\s*\s*(?:([^<]*)\s*)?"#).unwrap()); +static IMG_EQ_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?s)]*class="[^"]*img-equation[^"]*"[^>]*data-latex="([^"]*)"[^>]*>.*?\s*\s*(?:([^<]*)\s*)?"#).unwrap() +}); // extract_math_blocks -static MATH_BLOCK_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*?)>(.*?)"#).unwrap()); +static MATH_BLOCK_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?s)]*?)>(.*?)"#).unwrap()); static ALTTEXT_RE: LazyLock = LazyLock::new(|| Regex::new(r#"alttext="([^"]*)""#).unwrap()); -static ANNOTATION_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*encoding="application/x-tex"[^>]*>(.*?)"#).unwrap()); +static ANNOTATION_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?s)]*encoding="application/x-tex"[^>]*>(.*?)"#) + .unwrap() +}); // convert_img_tags -static IMG_TAG_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*?)>"#).unwrap()); +static IMG_TAG_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?s)]*?)>"#).unwrap()); static IMG_SRC_RE: LazyLock = LazyLock::new(|| Regex::new(r#"src="([^"]*)""#).unwrap()); static IMG_ALT_RE: LazyLock = LazyLock::new(|| Regex::new(r#"alt="([^"]*)""#).unwrap()); // convert_common_markup -static FIGCAPTION_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*>(.*?)
"#).unwrap()); -static LTX_SEC_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_section[^"]*"[^>]*>(.*?)"#).unwrap()); -static LTX_SUBSEC_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsection[^"]*"[^>]*>(.*?)"#).unwrap()); -static LTX_SUBSUBSEC_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsubsection[^"]*"[^>]*>(.*?)"#).unwrap()); -static LTX_CAPTION_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)<(?:span|div|p)[^>]*class="[^"]*ltx_caption[^"]*"[^>]*>(.*?)"#).unwrap()); -static LTX_TITLE_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_document[^"]*"[^>]*>(.*?)"#).unwrap()); +static FIGCAPTION_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?s)]*>(.*?)"#).unwrap()); +static LTX_SEC_RE: LazyLock = LazyLock::new(|| { + Regex::new(r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_section[^"]*"[^>]*>(.*?)"#) + .unwrap() +}); +static LTX_SUBSEC_RE: LazyLock = LazyLock::new(|| { + Regex::new( + r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsection[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() +}); +static LTX_SUBSUBSEC_RE: LazyLock = LazyLock::new(|| { + Regex::new( + r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_subsubsection[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() +}); +static LTX_CAPTION_RE: LazyLock = LazyLock::new(|| { + Regex::new( + r#"(?s)<(?:span|div|p)[^>]*class="[^"]*ltx_caption[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() +}); +static LTX_TITLE_RE: LazyLock = LazyLock::new(|| { + Regex::new( + r#"(?s)<(?:h[1-6])[^>]*class="[^"]*ltx_title_document[^"]*"[^>]*>(.*?)"#, + ) + .unwrap() +}); static CITE_START_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*>"#).unwrap()); static CITE_END_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)"#).unwrap()); -static EMPTY_A_RE: LazyLock = LazyLock::new(|| Regex::new(r#"(?s)]*>\s*"#).unwrap()); +static EMPTY_A_RE: LazyLock = + LazyLock::new(|| Regex::new(r#"(?s)]*>\s*"#).unwrap()); // ── Functions ─────────────────────────────────────────────────────────────── @@ -142,51 +282,98 @@ pub fn extract_html_base_origin(html: &str) -> String { if let Some(caps) = BASE_RE.captures(html) { url_str = Some(caps[1].trim().to_string()); } else if let Some(caps) = CANONICAL_RE.captures(html) { - url_str = Some(caps.get(1).or_else(|| caps.get(2)).map(|m| m.as_str().trim().to_string()).unwrap_or_default()); + url_str = Some( + caps.get(1) + .or_else(|| caps.get(2)) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_default(), + ); } else if let Some(caps) = OG_URL_RE.captures(html) { - url_str = Some(caps.get(1).or_else(|| caps.get(2)).map(|m| m.as_str().trim().to_string()).unwrap_or_default()); + url_str = Some( + caps.get(1) + .or_else(|| caps.get(2)) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_default(), + ); } else if let Some(caps) = PRISM_URL_RE.captures(html) { - url_str = Some(caps.get(1).or_else(|| caps.get(2)).map(|m| m.as_str().trim().to_string()).unwrap_or_default()); + url_str = Some( + caps.get(1) + .or_else(|| caps.get(2)) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_default(), + ); } else if let Some(caps) = CITATION_URL_RE.captures(html) { - url_str = Some(caps.get(1).or_else(|| caps.get(2)).map(|m| m.as_str().trim().to_string()).unwrap_or_default()); + url_str = Some( + caps.get(1) + .or_else(|| caps.get(2)) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_default(), + ); } if let Some(mut u) = url_str { - if u.starts_with('/') { u = format!("https://ar5iv.labs.arxiv.org{}", u); } + if u.starts_with('/') { + u = format!("https://ar5iv.labs.arxiv.org{}", u); + } if let Ok(parsed) = url::Url::parse(&u) { let scheme = parsed.scheme(); let host = parsed.host_str().unwrap_or(""); - if !host.is_empty() { return format!("{}://{}", scheme, host); } + if !host.is_empty() { + return format!("{}://{}", scheme, host); + } } } "https://ar5iv.labs.arxiv.org".to_string() } pub fn truncate_footer(html: &str) -> &str { - if let Some(end) = html.find("
") { &html[..end] } - else if let Some(end) = html.find("") { + &html[..end] + } else if let Some(end) = html.find(" String { let text = TAG_STRIP_RE.replace_all(html, "").to_string(); - text.replace("&","&").replace("<","<").replace(">",">").replace(""","\"").replace(" "," ").replace("'","'") + text.replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace(""", "\"") + .replace(" ", " ") + .replace("'", "'") } pub fn strip_html_preserve_links(html: &str) -> String { let text = A_LINK_RE.replace_all(html, "$2").to_string(); let text = TAG_STRIP_RE.replace_all(&text, "").to_string(); - text.replace("&","&").replace("<","<").replace(">",">").replace(""","\"").replace(" "," ").replace("'","'") + text.replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace(""", "\"") + .replace(" ", " ") + .replace("'", "'") } pub fn decode_html_numeric_entities(html: &str) -> String { let entities = &*ENTITY_MAP; - ENTITY_RE.replace_all(html, |caps: ®ex::Captures| -> String { - let entity = caps[0].to_owned(); - entities.get(entity.as_str()).copied().unwrap_or(&entity).to_owned() - }).to_string() + ENTITY_RE + .replace_all(html, |caps: ®ex::Captures| -> String { + let entity = caps[0].to_owned(); + entities + .get(entity.as_str()) + .copied() + .unwrap_or(&entity) + .to_owned() + }) + .to_string() } -pub fn convert_html_math_to_latex(html: String, mut formulas: Vec<(String, bool)>) -> (String, Vec<(String, bool)>) { +pub fn convert_html_math_to_latex( + html: String, + mut formulas: Vec<(String, bool)>, +) -> (String, Vec<(String, bool)>) { let math_letter = r#"(?:[a-zA-Z]\b|dex\b|log\b|ln\b|sin\b|cos\b|tan\b|eff\b|env\b|obs\b|sys\b|sta\b|WD\b|MS\b|He\b|H\b)"#; let sub_content = r#"(?:[^<]*|[a-zA-Z0-9α-ωΑ-Ω⊙+\-λθχσϖ])+"#; let sup_content = r#"(?:[^<]*|[a-zA-Z0-9α-ωΑ-Ω⊙+\-λθχσϖ])+"#; @@ -198,76 +385,154 @@ pub fn convert_html_math_to_latex(html: String, mut formulas: Vec<(String, bool) }); let mut placeholder_counter = formulas.len(); - let result = math_regex.replace_all(&html, |caps: ®ex::Captures| { - let matched = &caps[0]; - let has_text_italic = ITALIC_CHECK_RE.captures_iter(matched).any(|c| { - let content = &c[1]; - content.len() > 2 && !content.chars().any(|ch| - ('α'..='ω').contains(&ch) || ('Α'..='Ω').contains(&ch) || - ch == '⊙' || ch == '≤' || ch == '≥' || ch == '−' || ch == '≈' - ) - }); - if !has_text_italic && ( - matched.contains(" open_count { - trailing_part = format!("){}", trailing_part); - math_part = &math_part[..math_part.len()-1]; continue; + let result = math_regex + .replace_all(&html, |caps: ®ex::Captures| { + let matched = &caps[0]; + let has_text_italic = ITALIC_CHECK_RE.captures_iter(matched).any(|c| { + let content = &c[1]; + content.len() > 2 + && !content.chars().any(|ch| { + ('α'..='ω').contains(&ch) + || ('Α'..='Ω').contains(&ch) + || ch == '⊙' + || ch == '≤' + || ch == '≥' + || ch == '−' + || ch == '≈' + }) + }); + if !has_text_italic + && (matched.contains(" open_count { + trailing_part = format!("){}", trailing_part); + math_part = &math_part[..math_part.len() - 1]; + continue; + } + } + break; } - break; + if math_part.trim().is_empty() { + return matched.to_string(); + } + let latex = html_math_to_latex(math_part); + formulas.push((latex, false)); + let p = format!(" MATHPLACEHOLDER{} ", placeholder_counter); + placeholder_counter += 1; + format!("{}{}", p, trailing_part) + } else { + matched.to_string() } - if math_part.trim().is_empty() { return matched.to_string(); } - let latex = html_math_to_latex(math_part); - formulas.push((latex, false)); - let p = format!(" MATHPLACEHOLDER{} ", placeholder_counter); - placeholder_counter += 1; - format!("{}{}", p, trailing_part) - } else { matched.to_string() } - }).to_string(); + }) + .to_string(); (result, formulas) } fn html_math_to_latex(s: &str) -> String { - let s = s.replace(' '," ").replace(' ',"\\,").replace('\u{2006}',"\\,"); - let s = s.replace('≤'," \\le ").replace('≥'," \\ge ").replace('−'," - ") - .replace('≈'," \\approx ").replace('×'," \\times ").replace('≡'," \\equiv ") - .replace('≠'," \\neq ").replace('∑'," \\sum ").replace('√'," \\sqrt ") - .replace('∞'," \\infty ").replace('∫'," \\int ").replace('–'," -- ").replace('—'," --- "); - let s = s.replace('⊙',"{\\odot}"); - let s = s.replace('α',"\\alpha").replace('β',"\\beta").replace('γ',"\\gamma").replace('δ',"\\delta").replace('ε',"\\epsilon") - .replace('ζ',"\\zeta").replace('η',"\\eta").replace('θ',"\\theta").replace('ι',"\\iota").replace('κ',"\\kappa") - .replace('λ',"\\lambda").replace('μ',"\\mu").replace('ν',"\\nu").replace('ξ',"\\xi").replace('π',"\\pi") - .replace('ρ',"\\rho").replace('σ',"\\sigma").replace('τ',"\\tau").replace('υ',"\\upsilon").replace('φ',"\\phi") - .replace('χ',"\\chi").replace('ψ',"\\psi").replace('ω',"\\omega").replace('ς',"\\varsigma") - .replace('Γ',"\\Gamma").replace('Δ',"\\Delta").replace('Θ',"\\Theta").replace('Λ',"\\Lambda") - .replace('Ξ',"\\Xi").replace('Π',"\\Pi").replace('Σ',"\\Sigma").replace('Υ',"\\Upsilon").replace('Φ',"\\Phi") - .replace('Ψ',"\\Psi").replace('Ω',"\\Omega"); - let s = s.replace("log","\\log ").replace("\\log \\,","\\log ").replace("\\log ","\\log "); + let s = s + .replace(' ', " ") + .replace([' ', '\u{2006}'], "\\,"); + let s = s + .replace('≤', " \\le ") + .replace('≥', " \\ge ") + .replace('−', " - ") + .replace('≈', " \\approx ") + .replace('×', " \\times ") + .replace('≡', " \\equiv ") + .replace('≠', " \\neq ") + .replace('∑', " \\sum ") + .replace('√', " \\sqrt ") + .replace('∞', " \\infty ") + .replace('∫', " \\int ") + .replace('–', " -- ") + .replace('—', " --- "); + let s = s.replace('⊙', "{\\odot}"); + let s = s + .replace('α', "\\alpha") + .replace('β', "\\beta") + .replace('γ', "\\gamma") + .replace('δ', "\\delta") + .replace('ε', "\\epsilon") + .replace('ζ', "\\zeta") + .replace('η', "\\eta") + .replace('θ', "\\theta") + .replace('ι', "\\iota") + .replace('κ', "\\kappa") + .replace('λ', "\\lambda") + .replace('μ', "\\mu") + .replace('ν', "\\nu") + .replace('ξ', "\\xi") + .replace('π', "\\pi") + .replace('ρ', "\\rho") + .replace('σ', "\\sigma") + .replace('τ', "\\tau") + .replace('υ', "\\upsilon") + .replace('φ', "\\phi") + .replace('χ', "\\chi") + .replace('ψ', "\\psi") + .replace('ω', "\\omega") + .replace('ς', "\\varsigma") + .replace('Γ', "\\Gamma") + .replace('Δ', "\\Delta") + .replace('Θ', "\\Theta") + .replace('Λ', "\\Lambda") + .replace('Ξ', "\\Xi") + .replace('Π', "\\Pi") + .replace('Σ', "\\Sigma") + .replace('Υ', "\\Upsilon") + .replace('Φ', "\\Phi") + .replace('Ψ', "\\Psi") + .replace('Ω', "\\Omega"); + let s = s + .replace("log", "\\log ") + .replace("\\log \\,", "\\log ") + .replace("\\log ", "\\log "); let s = ITALIC_INNER_RE.replace_all(&s, "$1").to_string(); - let s = SUB_INNER_RE.replace_all(&s, |caps: ®ex::Captures| { - let inner = ITALIC_INNER_RE.replace_all(&caps[1], "$1").to_string(); - if inner.chars().all(|c| c.is_alphabetic()) { format!("_{{\\text{{{}}}}}", inner) } - else { format!("_{{{}}}", inner) } - }).to_string(); - let s = SUP_INNER_RE.replace_all(&s, |caps: ®ex::Captures| { - let inner = ITALIC_INNER_RE.replace_all(&caps[1], "$1").to_string(); - if inner.chars().all(|c| c.is_alphabetic()) { format!("^{{\\text{{{}}}}}", inner) } - else { format!("^{{{}}}", inner) } - }).to_string(); - s.replace(" K"," \\text{K}").replace("dex"," \\text{dex}").split_whitespace().collect::>().join(" ") + let s = SUB_INNER_RE + .replace_all(&s, |caps: ®ex::Captures| { + let inner = ITALIC_INNER_RE.replace_all(&caps[1], "$1").to_string(); + if inner.chars().all(|c| c.is_alphabetic()) { + format!("_{{\\text{{{}}}}}", inner) + } else { + format!("_{{{}}}", inner) + } + }) + .to_string(); + let s = SUP_INNER_RE + .replace_all(&s, |caps: ®ex::Captures| { + let inner = ITALIC_INNER_RE.replace_all(&caps[1], "$1").to_string(); + if inner.chars().all(|c| c.is_alphabetic()) { + format!("^{{\\text{{{}}}}}", inner) + } else { + format!("^{{{}}}", inner) + } + }) + .to_string(); + s.replace(" K", " \\text{K}") + .replace("dex", " \\text{dex}") + .split_whitespace() + .collect::>() + .join(" ") } pub fn replace_latexml_tables(html: &str) -> String { @@ -283,31 +548,51 @@ pub fn replace_latexml_tables(html: &str) -> String { let mut matched_type = None; if let Some(class_cap) = LATEXML_CLASS_RE.captures(attrs) { let class_str = class_cap[1].to_lowercase(); - if class_str.contains("ltx_tabular") { matched_type = Some("table"); } - else if class_str.contains("ltx_tbody") { matched_type = Some("tbody"); } - else if class_str.contains("ltx_thead") { matched_type = Some("thead"); } - else if class_str.contains("ltx_tfoot") { matched_type = Some("tfoot"); } - else if class_str.contains("ltx_tr") { matched_type = Some("tr"); } - else if class_str.contains("ltx_th") { matched_type = Some("th"); } - else if class_str.contains("ltx_td") { matched_type = Some("td"); } + if class_str.contains("ltx_tabular") { + matched_type = Some("table"); + } else if class_str.contains("ltx_tbody") { + matched_type = Some("tbody"); + } else if class_str.contains("ltx_thead") { + matched_type = Some("thead"); + } else if class_str.contains("ltx_tfoot") { + matched_type = Some("tfoot"); + } else if class_str.contains("ltx_tr") { + matched_type = Some("tr"); + } else if class_str.contains("ltx_th") { + matched_type = Some("th"); + } else if class_str.contains("ltx_td") { + matched_type = Some("td"); + } } if let Some(t) = matched_type { let t_str = t.to_string(); result.push_str(&format!("<{}>", t_str)); stack.push((tag_name, Some(t_str))); - } else { result.push_str(mat.as_str()); stack.push((tag_name, None)); } + } else { + result.push_str(mat.as_str()); + stack.push((tag_name, None)); + } } else { let tag_name = cap.get(3).unwrap().as_str().to_lowercase(); let mut replaced = false; while let Some((open_name, open_type)) = stack.pop() { if open_name == tag_name { - if let Some(t) = open_type { result.push_str(&format!("", t)); } - else { result.push_str(&format!("", tag_name)); } - replaced = true; break; - } else if let Some(t) = open_type { result.push_str(&format!("", t)); } - else { result.push_str(&format!("", open_name)); } + if let Some(t) = open_type { + result.push_str(&format!("", t)); + } else { + result.push_str(&format!("", tag_name)); + } + replaced = true; + break; + } else if let Some(t) = open_type { + result.push_str(&format!("", t)); + } else { + result.push_str(&format!("", open_name)); + } + } + if !replaced { + result.push_str(mat.as_str()); } - if !replaced { result.push_str(mat.as_str()); } } last_pos = mat.end(); } @@ -316,12 +601,18 @@ pub fn replace_latexml_tables(html: &str) -> String { } pub fn convert_html_tables_to_markdown(markdown: &str) -> String { - TABLE_RE.replace_all(markdown, |caps: ®ex::Captures| { - let full_table = &caps[0]; - if full_table.to_lowercase().contains("colspan") || full_table.to_lowercase().contains("rowspan") { - full_table.to_string() - } else { convert_html_table_to_markdown(full_table) } - }).to_string() + TABLE_RE + .replace_all(markdown, |caps: ®ex::Captures| { + let full_table = &caps[0]; + if full_table.to_lowercase().contains("colspan") + || full_table.to_lowercase().contains("rowspan") + { + full_table.to_string() + } else { + convert_html_table_to_markdown(full_table) + } + }) + .to_string() } fn convert_html_table_to_markdown(html_table: &str) -> String { @@ -333,11 +624,15 @@ fn convert_html_table_to_markdown(html_table: &str) -> String { cells.push(strip_html_tags(&td_cap[1]).trim().replace('\n', " ")); } if !cells.is_empty() { - if col_count == 0 { col_count = cells.len(); } + if col_count == 0 { + col_count = cells.len(); + } md_rows.push(format!("| {} |", cells.join(" | "))); } } - if md_rows.is_empty() { return html_table.to_string(); } + if md_rows.is_empty() { + return html_table.to_string(); + } let separator = format!("|{}|", vec!["---"; col_count].join("|")); md_rows.insert(1, separator); md_rows.join("\n") @@ -349,9 +644,14 @@ pub fn postprocess_markdown(text: &str) -> String { let mut in_code_block = false; for line in text.lines() { let trimmed = line.trim(); - if trimmed.starts_with("```") { in_code_block = !in_code_block; } - if in_code_block { clean_lines.push(line.to_string()); } - else { clean_lines.push(trimmed.to_string()); } + if trimmed.starts_with("```") { + in_code_block = !in_code_block; + } + if in_code_block { + clean_lines.push(line.to_string()); + } else { + clean_lines.push(trimmed.to_string()); + } } let mut md = clean_lines.join("\n"); @@ -363,9 +663,15 @@ pub fn postprocess_markdown(text: &str) -> String { md = ASTROBJ_RE.replace_all(&md, "$1").to_string(); md = EMPTY_BRACKETS_RE.replace_all(&md, "").to_string(); - md = INTERNAL_LINK_RE.replace_all(&md, |caps: ®ex::Captures| { - if caps[1].starts_with('!') { caps[0].to_string() } else { caps[2].to_string() } - }).to_string(); + md = INTERNAL_LINK_RE + .replace_all(&md, |caps: ®ex::Captures| { + if caps[1].starts_with('!') { + caps[0].to_string() + } else { + caps[2].to_string() + } + }) + .to_string(); md = EXCESSIVE_NL_RE.replace_all(&md, "\n\n\n").to_string(); @@ -373,17 +679,32 @@ pub fn postprocess_markdown(text: &str) -> String { md = re.replace_all(&md, *repl).to_string(); } - md = md.replace("<","<").replace(">",">").replace("&","&").replace(""","\"").replace("'","'"); + md = md + .replace("<", "<") + .replace(">", ">") + .replace("&", "&") + .replace(""", "\"") + .replace("'", "'"); - md = LINK_FIX_RE.replace_all(&md, |caps: ®ex::Captures| { - format!("{}({})", &caps[1], caps[2].replace(r"\_","_").replace(r"\%","%")) - }).to_string(); + md = LINK_FIX_RE + .replace_all(&md, |caps: ®ex::Captures| { + format!( + "{}({})", + &caps[1], + caps[2].replace(r"\_", "_").replace(r"\%", "%") + ) + }) + .to_string(); md = LATEXML_MACRO_RE.replace_all(&md, " ").to_string(); md = HEADING_TRAIL_RE.replace_all(&md, "$1 $2").to_string(); md = SECTION_PROMOTE_RE.replace_all(&md, "## $2$3").to_string(); - md = ABSTRACT_CLEAN_RE.replace_all(&md, "## Abstract\n\n").to_string(); - md = BRACKET_INLINE_RE.replace_all(&md, "## $1\n\n$2").to_string(); + md = ABSTRACT_CLEAN_RE + .replace_all(&md, "## Abstract\n\n") + .to_string(); + md = BRACKET_INLINE_RE + .replace_all(&md, "## $1\n\n$2") + .to_string(); md = BRACKET_HEADER_RE.replace_all(&md, "## $1").to_string(); md = BULLET_CLEAN_RE.replace_all(&md, "$1 ").to_string(); md = BRACKET_NEWLINE_RE.replace_all(&md, "[$1]").to_string(); @@ -397,49 +718,80 @@ pub fn cut_acknowledgments_and_references(text: &str) -> String { let re = TRUNCATE_RE.get_or_init(|| { Regex::new(r"(?mi)^(?:(?:#+\s*)?(?:[\d\.]+\s+)?(?:acknowledgements?|acknowledgments?|acknowledgment|references?(?:\s+&\s+citations)?|literature\s+cited|bibliography|bibliographie)\s*[:\.]?\s*#*\s*$|\[(?:acknowledgements?|acknowledgments?|acknowledgment|references?|literature\s+cited|bibliography|bibliographie)\]\s*$|\{thebibliography\*?\}\s*$|\{ack(?:nowledgements?)?\}\s*$|^(?:acknowledgements?|acknowledgments?|acknowledgment|references?|literature\s+cited|bibliography|bibliographie)\b\s*[:\.]\s+.*$)").unwrap() }); - if let Some(mat) = re.find(text) { text[..mat.start()].trim_end().to_string() } - else { text.to_string() } + if let Some(mat) = re.find(text) { + text[..mat.start()].trim_end().to_string() + } else { + text.to_string() + } } pub fn extract_img_equations(html: &str) -> (String, Vec<(String, String)>) { let mut img_eqs = Vec::new(); let mut counter = 0u32; - let result = IMG_EQ_RE.replace_all(html, |caps: ®ex::Captures| { - let latex = caps[1].trim().to_string(); - let label = caps.get(2).map(|m| m.as_str().trim().to_string()).unwrap_or_default(); - img_eqs.push((latex, label)); - let p = format!(" IMGEQPLACEHOLDER{} ", counter); - counter += 1; p - }).to_string(); + let result = IMG_EQ_RE + .replace_all(html, |caps: ®ex::Captures| { + let latex = caps[1].trim().to_string(); + let label = caps + .get(2) + .map(|m| m.as_str().trim().to_string()) + .unwrap_or_default(); + img_eqs.push((latex, label)); + let p = format!(" IMGEQPLACEHOLDER{} ", counter); + counter += 1; + p + }) + .to_string(); (result, img_eqs) } pub fn extract_math_blocks(html: &str) -> (String, Vec<(String, bool)>) { let mut formulas = Vec::new(); let mut counter = 0u32; - let result = MATH_BLOCK_RE.replace_all(html, |caps: ®ex::Captures| { - let attrs = &caps[1]; - let mut alttext = ALTTEXT_RE.captures(attrs).map(|c| c[1].to_string()).unwrap_or_default(); - if alttext.is_empty() { - if let Some(ann_caps) = ANNOTATION_RE.captures(&caps[2]) { alttext = ann_caps[1].trim().to_string(); } - } - if alttext.is_empty() { alttext = strip_html_tags(&caps[2]).trim().to_string(); } - let is_block = attrs.contains("display=\"block\"") || attrs.contains("display='block'"); - formulas.push((alttext, is_block)); - let p = format!(" MATHPLACEHOLDER{} ", counter); - counter += 1; p - }).to_string(); + let result = MATH_BLOCK_RE + .replace_all(html, |caps: ®ex::Captures| { + let attrs = &caps[1]; + let mut alttext = ALTTEXT_RE + .captures(attrs) + .map(|c| c[1].to_string()) + .unwrap_or_default(); + if alttext.is_empty() { + if let Some(ann_caps) = ANNOTATION_RE.captures(&caps[2]) { + alttext = ann_caps[1].trim().to_string(); + } + } + if alttext.is_empty() { + alttext = strip_html_tags(&caps[2]).trim().to_string(); + } + let is_block = attrs.contains("display=\"block\"") || attrs.contains("display='block'"); + formulas.push((alttext, is_block)); + let p = format!(" MATHPLACEHOLDER{} ", counter); + counter += 1; + p + }) + .to_string(); (result, formulas) } pub fn convert_img_tags(html: &str, base_origin: &str) -> String { - IMG_TAG_RE.replace_all(html, |caps: ®ex::Captures| { - let attrs = &caps[1]; - let src = IMG_SRC_RE.captures(attrs).map(|c| c[1].to_string()).unwrap_or_default(); - let alt = IMG_ALT_RE.captures(attrs).map(|c| c[1].to_string()).unwrap_or_else(|| "image".to_string()); - let absolute_src = if src.starts_with('/') { format!("{}{}", base_origin, src) } else { src }; - format!("\n\n![{}]({})\n\n", alt, absolute_src) - }).to_string() + IMG_TAG_RE + .replace_all(html, |caps: ®ex::Captures| { + let attrs = &caps[1]; + let src = IMG_SRC_RE + .captures(attrs) + .map(|c| c[1].to_string()) + .unwrap_or_default(); + let alt = IMG_ALT_RE + .captures(attrs) + .map(|c| c[1].to_string()) + .unwrap_or_else(|| "image".to_string()); + let absolute_src = if src.starts_with('/') { + format!("{}{}", base_origin, src) + } else { + src + }; + format!("\n\n![{}]({})\n\n", alt, absolute_src) + }) + .to_string() } pub fn restore_math_placeholders(markdown: &str, formulas: &[(String, bool)]) -> String { @@ -447,8 +799,11 @@ pub fn restore_math_placeholders(markdown: &str, formulas: &[(String, bool)]) -> for i in (0..formulas.len()).rev() { let (ref alttext, is_block) = formulas[i]; let placeholder = format!("MATHPLACEHOLDER{}", i); - let replacement = if is_block { format!("\n\n$$\n{}\n$$\n\n", alttext) } - else { format!(" ${}$ ", alttext) }; + let replacement = if is_block { + format!("\n\n$$\n{}\n$$\n\n", alttext) + } else { + format!(" ${}$ ", alttext) + }; md = md.replace(&placeholder, &replacement); } md @@ -459,14 +814,28 @@ pub fn restore_img_equations(markdown: &str, img_eqs: &[(String, String)]) -> St for i in (0..img_eqs.len()).rev() { let (ref latex, ref label) = img_eqs[i]; let placeholder = format!("IMGEQPLACEHOLDER{}", i); - let decoded = latex.replace("&","&").replace("<","<").replace(">",">").replace(""","\"").replace("'","'"); + let decoded = latex + .replace("&", "&") + .replace("<", "<") + .replace(">", ">") + .replace(""", "\"") + .replace("'", "'"); let trimmed = decoded.trim(); let mut raw = trimmed.to_string(); - if raw.starts_with("$$") && raw.ends_with("$$") { raw = raw[2..raw.len()-2].trim().to_string(); } - else if raw.starts_with('$') && raw.ends_with('$') { raw = raw[1..raw.len()-1].trim().to_string(); } + if raw.starts_with("$$") && raw.ends_with("$$") { + raw = raw[2..raw.len() - 2].trim().to_string(); + } else if raw.starts_with('$') && raw.ends_with('$') { + raw = raw[1..raw.len() - 1].trim().to_string(); + } let replacement = if !label.is_empty() { - format!("\n\n$$\n{} \\tag{{{}}}\n$$\n\n", raw, label.trim_matches(|c| c=='('||c==')').trim()) - } else { format!("\n\n$$\n{}\n$$\n\n", raw) }; + format!( + "\n\n$$\n{} \\tag{{{}}}\n$$\n\n", + raw, + label.trim_matches(|c| c == '(' || c == ')').trim() + ) + } else { + format!("\n\n$$\n{}\n$$\n\n", raw) + }; md = md.replace(&placeholder, &replacement); } md @@ -474,14 +843,36 @@ pub fn restore_img_equations(markdown: &str, img_eqs: &[(String, String)]) -> St pub fn convert_common_markup(html: &str) -> String { let mut h = html.to_string(); - h = FIGCAPTION_RE.replace_all(&h, |caps: ®ex::Captures| { - format!("\n\n> **Figure:** {}\n", strip_html_tags(&caps[1]).trim()) - }).to_string(); - h = LTX_SEC_RE.replace_all(&h, |caps: ®ex::Captures| format!("\n\n## {}\n\n", strip_html_tags(&caps[1]).trim())).to_string(); - h = LTX_SUBSEC_RE.replace_all(&h, |caps: ®ex::Captures| format!("\n\n### {}\n\n", strip_html_tags(&caps[1]).trim())).to_string(); - h = LTX_SUBSUBSEC_RE.replace_all(&h, |caps: ®ex::Captures| format!("\n\n#### {}\n\n", strip_html_tags(&caps[1]).trim())).to_string(); - h = LTX_CAPTION_RE.replace_all(&h, |caps: ®ex::Captures| format!("\n> **Caption:** {}\n", strip_html_tags(&caps[1]).trim())).to_string(); - h = LTX_TITLE_RE.replace_all(&h, |caps: ®ex::Captures| format!("\n# {}\n\n", strip_html_tags(&caps[1]).trim())).to_string(); + h = FIGCAPTION_RE + .replace_all(&h, |caps: ®ex::Captures| { + format!("\n\n> **Figure:** {}\n", strip_html_tags(&caps[1]).trim()) + }) + .to_string(); + h = LTX_SEC_RE + .replace_all(&h, |caps: ®ex::Captures| { + format!("\n\n## {}\n\n", strip_html_tags(&caps[1]).trim()) + }) + .to_string(); + h = LTX_SUBSEC_RE + .replace_all(&h, |caps: ®ex::Captures| { + format!("\n\n### {}\n\n", strip_html_tags(&caps[1]).trim()) + }) + .to_string(); + h = LTX_SUBSUBSEC_RE + .replace_all(&h, |caps: ®ex::Captures| { + format!("\n\n#### {}\n\n", strip_html_tags(&caps[1]).trim()) + }) + .to_string(); + h = LTX_CAPTION_RE + .replace_all(&h, |caps: ®ex::Captures| { + format!("\n> **Caption:** {}\n", strip_html_tags(&caps[1]).trim()) + }) + .to_string(); + h = LTX_TITLE_RE + .replace_all(&h, |caps: ®ex::Captures| { + format!("\n# {}\n\n", strip_html_tags(&caps[1]).trim()) + }) + .to_string(); h = CITE_START_RE.replace_all(&h, "").to_string(); h = CITE_END_RE.replace_all(&h, "").to_string(); h = EMPTY_A_RE.replace_all(&h, "").to_string(); @@ -503,14 +894,25 @@ mod tests { fn test_decode_html_numeric_entities() { let input = "λ θ χ α"; let output = decode_html_numeric_entities(input); - assert!(output.contains('λ') && output.contains('θ') && output.contains('χ') && output.contains('α')); + assert!( + output.contains('λ') + && output.contains('θ') + && output.contains('χ') + && output.contains('α') + ); assert!(!output.contains("λ")); } #[test] fn test_cut_acknowledgments() { - assert_eq!(cut_acknowledgments_and_references("text\n## Acknowledgments\nthanks"), "text"); - assert_eq!(cut_acknowledgments_and_references("text\n## References\nbib"), "text"); + assert_eq!( + cut_acknowledgments_and_references("text\n## Acknowledgments\nthanks"), + "text" + ); + assert_eq!( + cut_acknowledgments_and_references("text\n## References\nbib"), + "text" + ); assert_eq!(cut_acknowledgments_and_references("text only"), "text only"); } @@ -547,7 +949,8 @@ mod tests { #[test] fn test_math_block_extraction() { - let html = r#"α"#; + let html = + r#"α"#; let (_, formulas) = extract_math_blocks(html); assert_eq!(formulas[0].0, "\\alpha"); assert!(!formulas[0].1); diff --git a/src/services/parser/generic.rs b/src/services/parser/generic.rs index d46efb2..dcde75b 100644 --- a/src/services/parser/generic.rs +++ b/src/services/parser/generic.rs @@ -4,13 +4,16 @@ use super::JournalParser; pub struct GenericParser; impl GenericParser { + #[allow(dead_code)] pub fn detect(_html: &str) -> bool { true // 永远匹配，放在检测链的最末尾 } } impl JournalParser for GenericParser { - fn name(&self) -> &str { "通用" } + fn name(&self) -> &str { + "通用" + } fn extract_body<'a>(&self, html: &'a str) -> &'a str { if let Some(start) = html.find("]*>.*?"#, r#"(?s)]*>.*?"#, ] { - h = regex::Regex::new(pat).unwrap().replace_all(&h, "").to_string(); + h = regex::Regex::new(pat) + .unwrap() + .replace_all(&h, "") + .to_string(); } h } // 以下方法使用默认空操作 —— html2md 原生处理标准 HTML 标记 - fn clean_markers(&self, html: &str) -> String { html.to_string() } - fn convert_headings(&self, html: &str) -> String { html.to_string() } - fn convert_figures(&self, html: &str, _: &str) -> String { html.to_string() } + fn clean_markers(&self, html: &str) -> String { + html.to_string() + } + fn convert_headings(&self, html: &str) -> String { + html.to_string() + } + fn convert_figures(&self, html: &str, _: &str) -> String { + html.to_string() + } } diff --git a/src/services/parser/iop.rs b/src/services/parser/iop.rs index b794d33..a579e9f 100644 --- a/src/services/parser/iop.rs +++ b/src/services/parser/iop.rs @@ -13,7 +13,9 @@ impl IopParser { } impl JournalParser for IopParser { - fn name(&self) -> &str { "IOP/AAS" } + fn name(&self) -> &str { + "IOP/AAS" + } fn extract_body<'a>(&self, html: &'a str) -> &'a str { // IOP 将正文放在 itemprop="articleBody" 的 div 中 @@ -33,14 +35,14 @@ impl JournalParser for IopParser { // IOP 页面框架垃圾 for pattern in &[ - r#"(?s)]*class="[^"]*wd-main-nav[^"]*"[^>]*>.*?"#, // 主导航 - r#"(?s)]*class="[^"]*content-nav[^"]*"[^>]*>.*?"#, // 侧边栏 + r#"(?s)]*class="[^"]*wd-main-nav[^"]*"[^>]*>.*?"#, // 主导航 + r#"(?s)]*class="[^"]*content-nav[^"]*"[^>]*>.*?"#, // 侧边栏 r#"(?s)]*class="[^"]*content-grid__full-width[^"]*"[^>]*>.*?"#, // 站点标题 - r#"(?s)]*class="[^"]*secondary-header[^"]*"[^>]*>.*?
"#, // 期刊品牌 + r#"(?s)]*class="[^"]*secondary-header[^"]*"[^>]*>.*?
"#, // 期刊品牌 r#"(?s)]*class="[^"]*article-head[^"]*"[^>]*>.*?\s*"#, // 文章元数据 - r#"(?s)]*class="[^"]*footer[^"]*"[^>]*>.*?"#, // 页脚 - r#"(?s)]*>.*?"#, // GTM noscript - r#"(?s)]*googletagmanager[^>]*>.*?"#, // GTM iframe + r#"(?s)]*class="[^"]*footer[^"]*"[^>]*>.*?"#, // 页脚 + r#"(?s)]*>.*?"#, // GTM noscript + r#"(?s)]*googletagmanager[^>]*>.*?"#, // GTM iframe r#"(?s)]*class="[^"]*partner-logo-overlay[^"]*"[^>]*>.*?"#, // AAS logo 弹窗 ] { h = Regex::new(pattern).unwrap().replace_all(&h, "").to_string(); @@ -59,11 +61,15 @@ impl JournalParser for IopParser { }).to_string(); // IOP 子节标题：

- h = Regex::new(r#"(?s)]*>\s*(?:]*>]*>)?\s*(.*?)\s*

"#) - .unwrap().replace_all(&h, |caps: ®ex::Captures| { - let inner = super::common::strip_html_tags(&caps[1]); - format!("\n\n### {}\n\n", inner.trim()) - }).to_string(); + h = Regex::new( + r#"(?s)]*>\s*(?:]*>]*>)?\s*(.*?)\s*"#, + ) + .unwrap() + .replace_all(&h, |caps: ®ex::Captures| { + let inner = super::common::strip_html_tags(&caps[1]); + format!("\n\n### {}\n\n", inner.trim()) + }) + .to_string(); h } @@ -74,42 +80,57 @@ impl JournalParser for IopParser { r#"(?s)]*class="[^"]*boxout[^"]*"[^>]*>\s*]*>\s*(.*?)\s*]*>\s*(.*?)\s*\s*\s*"# ).unwrap(); - figure_re.replace_all(html, |caps: ®ex::Captures| { - let inner_fig = &caps[1]; - let caption_raw = &caps[2]; + figure_re + .replace_all(html, |caps: ®ex::Captures| { + let inner_fig = &caps[1]; + let caption_raw = &caps[2]; - // 从 data-src 或 src 提取图片 URL - let src = Regex::new(r#"data-src="([^"]*)""#).unwrap() - .captures(inner_fig) - .map(|c| c[1].to_string()) - .or_else(|| { - Regex::new(r#"src="([^"]*)""#).unwrap() - .captures(inner_fig) - .map(|c| c[1].to_string()) - }) - .unwrap_or_default(); + // 从 data-src 或 src 提取图片 URL + let src = Regex::new(r#"data-src="([^"]*)""#) + .unwrap() + .captures(inner_fig) + .map(|c| c[1].to_string()) + .or_else(|| { + Regex::new(r#"src="([^"]*)""#) + .unwrap() + .captures(inner_fig) + .map(|c| c[1].to_string()) + }) + .unwrap_or_default(); - // 跳过 base64 占位图 - if src.starts_with("data:") { - return caps[0].to_string(); - } + // 跳过 base64 占位图 + if src.starts_with("data:") { + return caps[0].to_string(); + } - // 清理图注：去掉 Figure N. 标签和下载链接 - let caption_text = Regex::new(r#"(?s)]*>.*?"#).unwrap() - .replace_all(caption_raw, "").to_string(); - let caption_text = Regex::new(r#"(?s)]*class="[^"]*btn-multi-block[^"]*"[^>]*>.*?"#).unwrap() - .replace_all(&caption_text, "").to_string(); - let caption_text = Regex::new(r#"(?s)]*class="[^"]*print-hide[^"]*"[^>]*>.*?

"#).unwrap() - .replace_all(&caption_text, "").to_string(); - let caption = super::common::strip_html_tags(&caption_text); + // 清理图注：去掉 Figure N. 标签和下载链接 + let caption_text = Regex::new(r#"(?s)]*>.*?"#) + .unwrap() + .replace_all(caption_raw, "") + .to_string(); + let caption_text = Regex::new( + r#"(?s)]*class="[^"]*btn-multi-block[^"]*"[^>]*>.*?"#, + ) + .unwrap() + .replace_all(&caption_text, "") + .to_string(); + let caption_text = + Regex::new(r#"(?s)]*class="[^"]*print-hide[^"]*"[^>]*>.*?

"#) + .unwrap() + .replace_all(&caption_text, "") + .to_string(); + let caption = super::common::strip_html_tags(&caption_text); - format!("\n\n![{}]({})\n\n", caption.trim(), src) - }).to_string() + format!("\n\n![{}]({})\n\n", caption.trim(), src) + }) + .to_string() } fn clean_markers(&self, html: &str) -> String { // 删除残留的 JavaScript - Regex::new(r#"(?s)]*>.*?"#).unwrap() - .replace_all(html, "").to_string() + Regex::new(r#"(?s)]*>.*?"#) + .unwrap() + .replace_all(html, "") + .to_string() } } diff --git a/src/services/parser/mod.rs b/src/services/parser/mod.rs index 6848565..d3f18ac 100644 --- a/src/services/parser/mod.rs +++ b/src/services/parser/mod.rs @@ -1,24 +1,26 @@ // JournalParser trait + 期刊自动检测 + HTML→Markdown 解析管线 -mod ar5iv; mod aanda; -mod iop; -mod generic; +mod ar5iv; pub mod common; +mod generic; +mod iop; pub mod pdf; -use std::path::Path; -use tracing::{info, warn, error}; use crate::api::helpers::{check_paper_paths_in_db, get_paper_from_db}; +use std::path::Path; +use tracing::{error, info, warn}; -use ar5iv::Ar5ivParser; use aanda::AandaParser; -use iop::IopParser; +use ar5iv::Ar5ivParser; use generic::GenericParser; +use iop::IopParser; /// 期刊解析器 trait —— 每种期刊实现自己特有的预处理规则 trait JournalParser { /// 解析器名称（用于日志） - fn name(&self) -> &str { "未知" } + fn name(&self) -> &str { + "未知" + } /// 从完整 HTML 页面中提取正文区域 fn extract_body<'a>(&self, html: &'a str) -> &'a str; /// 删除页面框架垃圾（导航、面包屑、元数据表格、JS、广告） @@ -34,14 +36,20 @@ trait JournalParser { /// 自动检测期刊类型并返回对应的解析器 /// 优先级：ar5iv > IOP > A&A > 通用回退 fn detect_parser(html: &str) -> Box { - if Ar5ivParser::detect(html) { return Box::new(Ar5ivParser); } - if IopParser::detect(html) { return Box::new(IopParser); } - if AandaParser::detect(html) { return Box::new(AandaParser); } + if Ar5ivParser::detect(html) { + return Box::new(Ar5ivParser); + } + if IopParser::detect(html) { + return Box::new(IopParser); + } + if AandaParser::detect(html) { + return Box::new(AandaParser); + } Box::new(GenericParser) } // 将 PDF/MinerU 函数重新导出到 parser 模块层级 -pub use pdf::{submit_pdf_to_mineru, poll_and_extract_mineru, parse_pdf_via_mineru}; +pub use pdf::{parse_pdf_via_mineru, poll_and_extract_mineru, submit_pdf_to_mineru}; /// HTML → Markdown 核心解析管线 pub fn html_to_markdown(html_path: &Path) -> anyhow::Result { @@ -112,11 +120,15 @@ pub async fn parse_paper_service( bibcode: &str, force: bool, ) -> anyhow::Result { - info!("接收到文献结构化解析服务调用: {} (强制重新解析: {:?})", bibcode, force); + info!( + "接收到文献结构化解析服务调用: {} (强制重新解析: {:?})", + bibcode, force + ); - let (pdf_opt, html_opt, md_opt, _) = check_paper_paths_in_db(db, library_dir, bibcode) - .await? - .ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?; + let (pdf_opt, html_opt, md_opt, _) = + check_paper_paths_in_db(db, library_dir, bibcode) + .await? + .ok_or_else(|| anyhow::anyhow!("该文献未注册在数据库中"))?; // 如果先前已经解析成功过且非强制重新解析，直读 Markdown 文件返回 if !force { @@ -156,13 +168,18 @@ pub async fn parse_paper_service( parsed_markdown = format!("{}{}", front_matter, md); let md_filename = format!("{}.md", bibcode); let md_dest = library_dir.join("Markdown").join(&md_filename); - std::fs::create_dir_all(md_dest.parent().unwrap()).unwrap_or_default(); + if let Some(parent) = md_dest.parent() { + std::fs::create_dir_all(parent).unwrap_or_default(); + } if std::fs::write(&md_dest, &parsed_markdown).is_ok() { relative_md_path = format!("Markdown/{}", md_filename); } } Err(e) => { - warn!("HTML 转换为 Markdown 失败 {}: {}。将自动降级使用 PDF 结构化解析。", bibcode, e); + warn!( + "HTML 转换为 Markdown 失败 {}: {}。将自动降级使用 PDF 结构化解析。", + bibcode, e + ); } } } @@ -187,7 +204,9 @@ pub async fn parse_paper_service( parsed_markdown = format!("{}{}", front_matter, md); let md_filename = format!("{}.md", bibcode); let md_dest = library_dir.join("Markdown").join(&md_filename); - std::fs::create_dir_all(md_dest.parent().unwrap()).unwrap_or_default(); + if let Some(parent) = md_dest.parent() { + std::fs::create_dir_all(parent).unwrap_or_default(); + } if std::fs::write(&md_dest, &parsed_markdown).is_ok() { relative_md_path = format!("Markdown/{}", md_filename); } @@ -198,11 +217,17 @@ pub async fn parse_paper_service( } } } else { - error!("文献 {} 解析失败：本地 PDF 文件 {:?} 丢失", bibcode, pdf_abs); + error!( + "文献 {} 解析失败：本地 PDF 文件 {:?} 丢失", + bibcode, pdf_abs + ); return Err(anyhow::anyhow!("本地 PDF 文件未找到")); } } else { - error!("文献 {} 解析失败：请先下载该文献的 HTML 或 PDF 文件", bibcode); + error!( + "文献 {} 解析失败：请先下载该文献的 HTML 或 PDF 文件", + bibcode + ); return Err(anyhow::anyhow!("请先下载该文献的 HTML 或 PDF 文件")); } } @@ -222,8 +247,8 @@ pub async fn parse_paper_service( // ── 测试 ────────────────────────────────────────────────────────── #[cfg(test)] mod tests { + use super::common::{convert_html_tables_to_markdown, postprocess_markdown}; use super::html_to_markdown; - use super::common::{postprocess_markdown, convert_html_tables_to_markdown}; use std::io::Write; #[test] @@ -249,7 +274,8 @@ mod tests { let cleaned = postprocess_markdown(dirty); assert_eq!(cleaned, "Hello World V391 Peg \n\n\nNew Paragraph"); - let dirty_abstract = "###### Abstract\n\n[Abstract]\n\nHot subdwarfs are core helium burning stars."; + let dirty_abstract = + "###### Abstract\n\n[Abstract]\n\nHot subdwarfs are core helium burning stars."; let cleaned_abstract = postprocess_markdown(dirty_abstract); assert!(cleaned_abstract.contains("## Abstract\n\nHot subdwarfs are core")); assert!(!cleaned_abstract.contains("[Abstract]")); @@ -267,9 +293,13 @@ mod tests { assert_eq!(postprocess_markdown(with_inline_ack), "Main body."); // 图注与图片连行修复 - let merged_fig = r"![alt](https://example.com/img.png) \> \*\*Figure:\*\* This is a caption."; + let merged_fig = + r"![alt](https://example.com/img.png) \> \*\*Figure:\*\* This is a caption."; let fixed_fig = postprocess_markdown(merged_fig); - assert!(fixed_fig.contains("img.png)\n\n> **Figure:**"), "got: '{fixed_fig}'"); + assert!( + fixed_fig.contains("img.png)\n\n> **Figure:**"), + "got: '{fixed_fig}'" + ); // #bib 链接剥离 let bib_txt = "([1976](#bib.bib16)) and [1984](#bib.bib30)."; @@ -278,7 +308,8 @@ mod tests { assert!(bib_fixed.contains("(1976)")); // 误报检测：正文中 "Reference objects" 标题不应被截断 - let fp = "We compare with reference objects in Sect. 2.\n\n## 2.4 Reference objects\n\nText."; + let fp = + "We compare with reference objects in Sect. 2.\n\n## 2.4 Reference objects\n\nText."; let fixed = postprocess_markdown(fp); assert!(fixed.contains("## 2.4 Reference objects")); } @@ -326,7 +357,10 @@ mod tests { #[test] fn test_aa_full_html_conversion() -> anyhow::Result<()> { let html_path = std::path::Path::new("library/HTML/2026A&A...709A..52F.html"); - if !html_path.exists() { eprintln!("Skip: file not found"); return Ok(()); } + if !html_path.exists() { + eprintln!("Skip: file not found"); + return Ok(()); + } let md = html_to_markdown(html_path)?; assert!(!md.contains("var prefix")); assert!(!md.contains("[Home](/")); @@ -345,7 +379,10 @@ mod tests { #[test] fn test_ar5iv_full_html_conversion() -> anyhow::Result<()> { let html_path = std::path::Path::new("library/HTML/1103.1435.html"); - if !html_path.exists() { eprintln!("Skip: file not found"); return Ok(()); } + if !html_path.exists() { + eprintln!("Skip: file not found"); + return Ok(()); + } let md = html_to_markdown(html_path)?; assert!(md.contains("x1.png)\n\n> **Figure:**")); assert!(!md.contains("[1]")); @@ -363,7 +400,9 @@ mod tests { ("small", "library/HTML/0804.1287.html"), ] { let path = std::path::Path::new(file); - if !path.exists() { continue; } + if !path.exists() { + continue; + } let start = std::time::Instant::now(); let _ = html_to_markdown(path).unwrap(); let ms = start.elapsed().as_secs_f64() * 1000.0; diff --git a/src/services/parser/pdf.rs b/src/services/parser/pdf.rs index c42ce61..74d4e3f 100644 --- a/src/services/parser/pdf.rs +++ b/src/services/parser/pdf.rs @@ -1,12 +1,12 @@ // PDF 解析 —— 通过 MinerU 远程 API 提取 Markdown +use regex::Regex; +use serde::{Deserialize, Serialize}; use std::fs; use std::path::Path; -use serde::{Deserialize, Serialize}; use tracing::{info, warn}; -use regex::Regex; -use crate::Config; use crate::clients::qiniu::QiniuClient; +use crate::Config; // 复用公共后处理函数 use super::common::{convert_html_tables_to_markdown, cut_acknowledgments_and_references}; @@ -63,29 +63,29 @@ struct ExtractResult { } // 调用 MinerU 远程接口解析 PDF，并在提取出图片后自动上传至七牛云进行外链替换 -pub async fn submit_pdf_to_mineru( - pdf_path: &Path, - config: &Config -) -> anyhow::Result { +pub async fn submit_pdf_to_mineru(pdf_path: &Path, config: &Config) -> anyhow::Result { info!("正在请求 MinerU 解析本地 PDF 文献: {:?}", pdf_path); - + if config.mineru_api_url.is_empty() { return Err(anyhow::anyhow!("未在环境变量 .env 中配置 MINERU_API_URL")); } let pdf_bytes = fs::read(pdf_path)?; - let filename = pdf_path.file_name() + let filename = pdf_path + .file_name() .and_then(|f| f.to_str()) .unwrap_or("paper.pdf") .to_string(); - let bibcode = pdf_path.file_stem() + let bibcode = pdf_path + .file_stem() .and_then(|f| f.to_str()) .unwrap_or("paper") .to_string(); // 提取 base_url - let base_url = config.mineru_api_url + let base_url = config + .mineru_api_url .replace("/extract/task", "") .replace("/extract", "") .trim_end_matches('/') @@ -106,7 +106,8 @@ pub async fn submit_pdf_to_mineru( model_version: "vlm".to_string(), }; - let mut request = client.post(format!("{}/file-urls/batch/", base_url)) + let mut request = client + .post(format!("{}/file-urls/batch/", base_url)) .json(&upload_req); if !config.mineru_api_key.is_empty() { @@ -117,7 +118,11 @@ pub async fn submit_pdf_to_mineru( let status = response.status(); let res_text = response.text().await?; if !status.is_success() { - return Err(anyhow::anyhow!("请求 MinerU 批量上传 URL 失败 (状态码: {}): {}", status, res_text)); + return Err(anyhow::anyhow!( + "请求 MinerU 批量上传 URL 失败 (状态码: {}): {}", + status, + res_text + )); } let upload_res: BatchUploadResponse = serde_json::from_str(&res_text)?; @@ -125,17 +130,20 @@ pub async fn submit_pdf_to_mineru( return Err(anyhow::anyhow!("MinerU API 错误: {}", upload_res.msg)); } - let upload_url = upload_res.data.file_urls.first() + let upload_url = upload_res + .data + .file_urls + .first() .ok_or_else(|| anyhow::anyhow!("MinerU 未返回上传 URL"))?; // 2. 上传文件 (PUT) info!("MinerU: 正在直接上传 PDF 字节流至对象存储..."); - let put_res = client.put(upload_url) - .body(pdf_bytes) - .send() - .await?; + let put_res = client.put(upload_url).body(pdf_bytes).send().await?; if !put_res.status().is_success() { - return Err(anyhow::anyhow!("上传 PDF 至 MinerU 对象存储直传 URL 失败: {}", put_res.status())); + return Err(anyhow::anyhow!( + "上传 PDF 至 MinerU 对象存储直传 URL 失败: {}", + put_res.status() + )); } let batch_id = upload_res.data.batch_id; @@ -146,10 +154,11 @@ pub async fn poll_and_extract_mineru( batch_id: &str, bibcode: &str, qiniu_client: &QiniuClient, - config: &Config + config: &Config, ) -> anyhow::Result { let client = reqwest::Client::new(); - let base_url = config.mineru_api_url + let base_url = config + .mineru_api_url .replace("/extract/task", "") .replace("/extract", "") .trim_end_matches('/') @@ -159,18 +168,22 @@ pub async fn poll_and_extract_mineru( let max_polls = 45; // 45 * 10s = 7.5 min info!("MinerU: 开始轮询任务结果 (Batch ID: {})...", batch_id); - let mut full_zip_url = String::new(); + let full_zip_url; loop { poll_count += 1; if poll_count > max_polls { - return Err(anyhow::anyhow!("MinerU 结构化解析超时 (Bibcode: {})", bibcode)); + return Err(anyhow::anyhow!( + "MinerU 结构化解析超时 (Bibcode: {})", + bibcode + )); } tokio::time::sleep(std::time::Duration::from_secs(10)).await; let mut status_req = client.get(format!("{}/extract-results/batch/{}", base_url, batch_id)); if !config.mineru_api_key.is_empty() { - status_req = status_req.header("Authorization", format!("Bearer {}", config.mineru_api_key)); + status_req = + status_req.header("Authorization", format!("Bearer {}", config.mineru_api_key)); } let status_res = status_req.send().await?; @@ -225,7 +238,12 @@ pub async fn poll_and_extract_mineru( markdown = md_content; } else if file.is_file() { let lower = name.to_lowercase(); - if lower.ends_with(".png") || lower.ends_with(".jpg") || lower.ends_with(".jpeg") || lower.ends_with(".gif") || lower.ends_with(".svg") { + if lower.ends_with(".png") + || lower.ends_with(".jpg") + || lower.ends_with(".jpeg") + || lower.ends_with(".gif") + || lower.ends_with(".svg") + { let mut buf = Vec::new(); std::io::copy(&mut file, &mut buf)?; let file_basename = Path::new(&name) @@ -248,7 +266,11 @@ pub async fn poll_and_extract_mineru( let _ = fs::create_dir_all(&mineru_dir); let raw_path = mineru_dir.join(format!("{}.md", bibcode)); if let Err(e) = fs::write(&raw_path, &markdown) { - warn!("保存 MinerU 原始 Markdown 失败: {} ({})", raw_path.display(), e); + warn!( + "保存 MinerU 原始 Markdown 失败: {} ({})", + raw_path.display(), + e + ); } } @@ -258,7 +280,10 @@ pub async fn poll_and_extract_mineru( let _ = fs::create_dir_all(&local_img_dir); if qiniu_client.is_configured() { - info!("MinerU 批量模式解析出 {} 张本地插图。准备上传至七牛云...", image_buffers.len()); + info!( + "MinerU 批量模式解析出 {} 张本地插图。准备上传至七牛云...", + image_buffers.len() + ); for (img_name, img_bytes) in image_buffers { let local_path = local_img_dir.join(&img_name); let _ = fs::write(&local_path, &img_bytes); @@ -267,10 +292,13 @@ pub async fn poll_and_extract_mineru( match qiniu_client.upload_buffer(img_bytes, &qiniu_name).await { Ok(qiniu_url) => { let escaped_img_name = regex::escape(&img_name); - let link_re = Regex::new(&format!(r"\((?:[^)]*/)?{}\)", escaped_img_name)).unwrap(); - markdown = link_re.replace_all(&markdown, |_: ®ex::Captures| { - format!("({})", qiniu_url) - }).to_string(); + let link_re = + Regex::new(&format!(r"\((?:[^)]*/)?{}\)", escaped_img_name)).unwrap(); + markdown = link_re + .replace_all(&markdown, |_: ®ex::Captures| { + format!("({})", qiniu_url) + }) + .to_string(); } Err(e) => warn!("上传图片至七牛云失败 {}: {}", img_name, e), } @@ -284,7 +312,9 @@ pub async fn poll_and_extract_mineru( let escaped_img_name = regex::escape(&img_name); let link_re = Regex::new(&format!(r"\((?:[^)]*/)?{}\)", escaped_img_name)).unwrap(); let replacement_link = format!("(images/{}/{})", bibcode, img_name); - markdown = link_re.replace_all(&markdown, replacement_link.as_str()).to_string(); + markdown = link_re + .replace_all(&markdown, replacement_link.as_str()) + .to_string(); } } } @@ -295,15 +325,15 @@ pub async fn poll_and_extract_mineru( } pub async fn parse_pdf_via_mineru( - pdf_path: &Path, - qiniu_client: &QiniuClient, - config: &Config + pdf_path: &Path, + qiniu_client: &QiniuClient, + config: &Config, ) -> anyhow::Result { - let bibcode = pdf_path.file_stem() + let bibcode = pdf_path + .file_stem() .and_then(|f| f.to_str()) .unwrap_or("paper") .to_string(); let batch_id = submit_pdf_to_mineru(pdf_path, config).await?; poll_and_extract_mineru(&batch_id, &bibcode, qiniu_client, config).await } - diff --git a/src/services/query_parser.rs b/src/services/query_parser.rs index 1aa4077..4ed4a30 100644 --- a/src/services/query_parser.rs +++ b/src/services/query_parser.rs @@ -4,7 +4,7 @@ use regex::Regex; /// 清洗用户输入的检索词，转换全角字符和中文标点 pub fn clean_query(query: &str) -> String { let mut cleaned = query.to_string(); - + // 全角双引号 -> 半角双引号 cleaned = cleaned.replace("“", "\"").replace("”", "\""); // 全角单引号 -> 半角单引号 @@ -13,7 +13,7 @@ pub fn clean_query(query: &str) -> String { cleaned = cleaned.replace("（", "(").replace("）", ")"); // 全角逗号/分号 cleaned = cleaned.replace("，", ",").replace("；", ";"); - + cleaned.trim().to_string() } @@ -21,26 +21,32 @@ pub fn clean_query(query: &str) -> String { /// 例如: `hot subdwarf year:2020-2023` -> (Some(2020), Some(2023), "hot subdwarf") pub fn extract_year_filter(query: &str) -> (Option, Option, String) { let cleaned = clean_query(query); - + // 匹配 year:2020-2023 或 year:2020 let year_re = Regex::new(r"(?i)\byear:\s*(\d{4})(?:\s*-\s*(\d{4}))?\b").unwrap(); - + if let Some(caps) = year_re.captures(&cleaned) { let start_year = caps.get(1).and_then(|m| m.as_str().parse::().ok()); - let end_year = caps.get(2) + let end_year = caps + .get(2) .and_then(|m| m.as_str().parse::().ok()) .or(start_year); // 如果是单一年份 year:2020，结束年份也是 2020 - + // 将 year 过滤子句从原始检索式中移除，避免污染基础文本匹配 let without_year = year_re.replace_all(&cleaned, "").to_string(); - + // 清理可能由于移除子句导致的多余 AND/OR 逻辑符或空格 - let cleanup_re = Regex::new(r"\s+(AND|OR|NOT)\s*$|^\s*(AND|OR|NOT)\s+|\s+(AND|OR)\s+(AND|OR)\s+").unwrap(); - let final_query = cleanup_re.replace_all(&without_year, " ").trim().to_string(); - + let cleanup_re = + Regex::new(r"\s+(AND|OR|NOT)\s*$|^\s*(AND|OR|NOT)\s+|\s+(AND|OR)\s+(AND|OR)\s+") + .unwrap(); + let final_query = cleanup_re + .replace_all(&without_year, " ") + .trim() + .to_string(); + return (start_year, end_year, final_query); } - + (None, None, cleaned) } @@ -48,28 +54,28 @@ pub fn extract_year_filter(query: &str) -> (Option, Option, String) { pub fn to_ads_query(query: &str) -> String { let (start, end, rest_query) = extract_year_filter(query); let mut parts = Vec::new(); - + // 处理剩余检索词项的字段映射 (如 abs: -> abstract:) let ads_rest = rest_query .replace("abs:", "abstract:") .replace("ti:", "title:") .replace("au:", "author:"); - + if !ads_rest.trim().is_empty() { parts.push(ads_rest); } - + // 如果有时间范围，添加 Solr 范围语法 if let Some(s) = start { if let Some(e) = end { parts.push(format!("year:[{} TO {}]", s, e)); } } - + if parts.is_empty() { return "*:*".to_string(); } - + if parts.len() == 1 { parts[0].clone() } else { @@ -82,33 +88,35 @@ pub fn to_ads_query(query: &str) -> String { pub fn to_arxiv_query(query: &str) -> (String, Option<(i32, i32)>) { let (start, end, rest_query) = extract_year_filter(query); let cleaned_rest = rest_query; - + // 年份范围元组 let year_range = start.map(|s| (s, end.unwrap_or(s))); - + if cleaned_rest.trim().is_empty() { return ("all:\"\"".to_string(), year_range); } - + // 自动为未限定前缀的检索短语/单词补全前缀 // 逻辑：以空格、括号、运算符分割，为不带前缀的独立词/短语添加 "all:"。 // 用正则简单分词翻译： // 我们找出所有的双引号短语，或者无空格单词，如果它们不是运算符（AND, OR, NOT, ANDNOT）且不带冒号前缀，则加上 all: - let token_re = Regex::new(r#"(?s)(\b(?:title|author|abs|ti|au):)?("[^"]+"|\b[a-zA-Z0-9_\-\.\*]+)"#).unwrap(); - + let token_re = + Regex::new(r#"(?s)(\b(?:title|author|abs|ti|au):)?("[^"]+"|\b[a-zA-Z0-9_\-\.\*]+)"#) + .unwrap(); + let mut translated = String::new(); let mut last_pos = 0; - + for cap in token_re.captures_iter(&cleaned_rest) { let entire_match = cap.get(0).unwrap(); let prefix = cap.get(1).map(|m| m.as_str()).unwrap_or(""); let val = cap.get(2).map(|m| m.as_str()).unwrap_or(""); - + // 拼装匹配项之间的非单词字符（如空格、括号、逻辑运算符） let between = &cleaned_rest[last_pos..entire_match.start()]; translated.push_str(between); last_pos = entire_match.end(); - + let val_upper = val.to_uppercase(); if val_upper == "AND" || val_upper == "OR" || val_upper == "NOT" { // NOT 翻译为 ANDNOT，因为 arXiv 不支持单独的 NOT @@ -133,16 +141,16 @@ pub fn to_arxiv_query(query: &str) -> (String, Option<(i32, i32)>) { translated.push_str(&format!("{}{}", standard_prefix, val)); } } - + if last_pos < cleaned_rest.len() { translated.push_str(&cleaned_rest[last_pos..]); } - + // 全局清理和修饰：如果翻译后的语句中依然有单独的 NOT，将其转换为 ANDNOT let translated_clean = translated .replace(" NOT ", " ANDNOT ") .replace("(NOT ", "(ANDNOT "); - + (translated_clean.trim().to_string(), year_range) } @@ -172,16 +180,24 @@ mod tests { #[test] fn test_to_ads_query() { let ads = to_ads_query("author:\"Althaus\" AND ti:\"hot subdwarf\" year:2020-2023"); - assert_eq!(ads, "(author:\"Althaus\" AND title:\"hot subdwarf\") AND year:[2020 TO 2023]"); + assert_eq!( + ads, + "(author:\"Althaus\" AND title:\"hot subdwarf\") AND year:[2020 TO 2023]" + ); } #[test] fn test_to_arxiv_query() { - let (arxiv, year) = to_arxiv_query("author:\"Althaus\" AND ti:\"hot subdwarf\" year:2020-2023"); + let (arxiv, year) = + to_arxiv_query("author:\"Althaus\" AND ti:\"hot subdwarf\" year:2020-2023"); assert_eq!(arxiv, "au:\"Althaus\" AND ti:\"hot subdwarf\""); assert_eq!(year, Some((2020, 2023))); - let (arxiv2, _) = to_arxiv_query("(\"hot subdwarf\" OR sdOB) AND Gaia NOT \"neutron star\""); - assert_eq!(arxiv2, "(all:\"hot subdwarf\" OR all:sdOB) AND all:Gaia ANDNOT all:\"neutron star\""); + let (arxiv2, _) = + to_arxiv_query("(\"hot subdwarf\" OR sdOB) AND Gaia NOT \"neutron star\""); + assert_eq!( + arxiv2, + "(all:\"hot subdwarf\" OR all:sdOB) AND all:Gaia ANDNOT all:\"neutron star\"" + ); } } diff --git a/src/services/rag.rs b/src/services/rag.rs index 53f10a7..ae99daa 100644 --- a/src/services/rag.rs +++ b/src/services/rag.rs @@ -7,7 +7,7 @@ // 3. Complete: 组装检索上下文 -> 调用 LLM 生成最终回答 use sqlx::SqlitePool; -use tracing::{info, warn, error}; +use tracing::{error, info, warn}; use crate::clients::llm::EmbeddingClient; use crate::services::chunker::{chunk_markdown, TextChunk}; @@ -64,12 +64,11 @@ pub async fn ingest_paper( // 先清除该文献的旧切片（重新解析场景） // 因为 vec0 虚拟表没有 FK，需要先查出旧的 rowid 再手动删除 - let old_rowids: Vec<(i64,)> = sqlx::query_as( - "SELECT rowid FROM paper_chunks_content WHERE bibcode = ?" - ) - .bind(bibcode) - .fetch_all(pool) - .await?; + let old_rowids: Vec<(i64,)> = + sqlx::query_as("SELECT rowid FROM paper_chunks_content WHERE bibcode = ?") + .bind(bibcode) + .fetch_all(pool) + .await?; if !old_rowids.is_empty() { info!("清除文献 {} 的 {} 条旧切片", bibcode, old_rowids.len()); @@ -102,7 +101,7 @@ pub async fn ingest_paper( // 插入 content 表获得 rowid let result = sqlx::query( - "INSERT INTO paper_chunks_content (bibcode, paragraph_index, content) VALUES (?, ?, ?)" + "INSERT INTO paper_chunks_content (bibcode, paragraph_index, content) VALUES (?, ?, ?)", ) .bind(bibcode) .bind(chunk.paragraph_index as i64) @@ -115,20 +114,20 @@ pub async fn ingest_paper( // 将 f32 向量转换为字节切片，插入 vec0 虚拟表 let embedding_bytes = embedding_to_bytes(&embedding); - sqlx::query( - "INSERT INTO vec_paper_chunks (rowid, embedding) VALUES (?, ?)" - ) - .bind(rowid) - .bind(&embedding_bytes) - .execute(pool) - .await?; + sqlx::query("INSERT INTO vec_paper_chunks (rowid, embedding) VALUES (?, ?)") + .bind(rowid) + .bind(&embedding_bytes) + .execute(pool) + .await?; ingested_count += 1; } info!( "文献 {} 向量化完成，成功写入 {}/{} 块", - bibcode, ingested_count, chunks.len() + bibcode, + ingested_count, + chunks.len() ); Ok(ingested_count) @@ -153,7 +152,7 @@ pub async fn retrieve( FROM vec_paper_chunks v \ INNER JOIN paper_chunks_content c ON v.rowid = c.rowid \ WHERE v.embedding MATCH ? AND k = ? \ - ORDER BY v.distance" + ORDER BY v.distance", ) .bind(&query_bytes) .bind(top_k as i64) @@ -162,12 +161,14 @@ pub async fn retrieve( let results: Vec = rows .into_iter() - .map(|(bibcode, paragraph_index, content, distance)| RetrievalResult { - bibcode, - paragraph_index, - content, - distance, - }) + .map( + |(bibcode, paragraph_index, content, distance)| RetrievalResult { + bibcode, + paragraph_index, + content, + distance, + }, + ) .collect(); info!("RAG 检索完成，返回 {} 条结果", results.len()); @@ -247,7 +248,7 @@ mod tests { #[test] fn test_embedding_to_bytes_roundtrip() { - let original = vec![1.0f32, 2.5, -3.14, 0.0]; + let original = vec![1.0f32, 2.5, -std::f32::consts::PI, 0.0]; let bytes = embedding_to_bytes(&original); assert_eq!(bytes.len(), 16); // 4 floats × 4 bytes diff --git a/src/services/search.rs b/src/services/search.rs index 514b0ca..a95f035 100644 --- a/src/services/search.rs +++ b/src/services/search.rs @@ -1,11 +1,10 @@ // src/services/search.rs -use tracing::{warn, error}; +use crate::api::helpers::{ + convert_ads_doc_to_standard, convert_arxiv_to_standard, get_paper_from_db, save_paper_to_db, +}; use crate::api::AppState; use crate::api::StandardPaper; -use crate::api::helpers::{ - convert_ads_doc_to_standard, convert_arxiv_to_standard, save_paper_to_db, - get_paper_from_db, -}; +use tracing::{error, warn}; /// 统一检索逻辑，合并去重 ADS 和 arXiv 数据，并自动记录引用拓扑与同步本地状态 pub async fn search_papers( @@ -25,7 +24,7 @@ pub async fn search_papers( Ok(docs) => { for doc in docs { let paper = convert_ads_doc_to_standard(&doc); - + // 入库 SQLite if let Err(e) = save_paper_to_db(&state.db, &paper).await { warn!("保存 ADS 文献至数据库失败: {}", e); @@ -69,7 +68,7 @@ pub async fn search_papers( Ok(papers) => { for p in papers { let paper = convert_arxiv_to_standard(&p); - + // 入库 SQLite (使用 arXiv ID 暂作主键以作记录) if let Err(e) = save_paper_to_db(&state.db, &paper).await { warn!("保存 arXiv 文献至数据库失败: {}", e); @@ -87,10 +86,16 @@ pub async fn search_papers( // 对两端获取的数据进行去重合并，增加对相同 arxiv_id 的判断 let mut unique_results: Vec = Vec::new(); for r in results { - if !unique_results.iter().any(|u| u.bibcode == r.bibcode || (!u.doi.is_empty() && u.doi == r.doi) || (!u.arxiv_id.is_empty() && u.arxiv_id == r.arxiv_id)) { + if !unique_results.iter().any(|u| { + u.bibcode == r.bibcode + || (!u.doi.is_empty() && u.doi == r.doi) + || (!u.arxiv_id.is_empty() && u.arxiv_id == r.arxiv_id) + }) { let mut final_paper = r.clone(); // 如果本地数据库存在该文献，直接从数据库读取标准元数据（包括 is_downloaded, has_markdown, pdf_error, html_error 等） - if let Ok(db_paper) = get_paper_from_db(&state.db, &state.config.library_dir, &r.bibcode).await { + if let Ok(db_paper) = + get_paper_from_db(&state.db, &state.config.library_dir, &r.bibcode).await + { final_paper = db_paper; } unique_results.push(final_paper); diff --git a/src/services/target.rs b/src/services/target.rs index 8d6c1f3..244c109 100644 --- a/src/services/target.rs +++ b/src/services/target.rs @@ -9,8 +9,8 @@ use regex::Regex; use reqwest::Client; use sqlx::SqlitePool; -use tracing::{info, warn, error}; use std::sync::OnceLock; +use tracing::{error, info, warn}; /// 天体目标的标准化属性信息 #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] @@ -157,12 +157,9 @@ fn parse_sesame_xml(xml: &str, original_name: &str) -> anyhow::Result().ok()); - let parallax = extract_xml_value(xml, "plx") - .and_then(|v| v.parse::().ok()); + let spectral_type = extract_xml_value(xml, "spType").or_else(|| extract_xml_value(xml, "sp")); + let v_mag = extract_xml_value(xml, "Vmag").and_then(|v| v.parse::().ok()); + let parallax = extract_xml_value(xml, "plx").and_then(|v| v.parse::().ok()); // 提取别名列表 let mut aliases = Vec::new(); @@ -245,7 +242,8 @@ pub async fn query_target_cached( // 3. 写入缓存 if let Some(bib) = bibcode { - let aliases_json = serde_json::to_string(&info.aliases).unwrap_or_else(|_| "[]".to_string()); + let aliases_json = + serde_json::to_string(&info.aliases).unwrap_or_else(|_| "[]".to_string()); if let Err(e) = sqlx::query( "INSERT OR IGNORE INTO paper_targets (bibcode, target_name, ra, dec, parallax, spectral_type, v_magnitude, aliases) VALUES (?, ?, ?, ?, ?, ?, ?, ?)" ) diff --git a/src/services/translation.rs b/src/services/translation.rs index 535cb04..c9cc54c 100644 --- a/src/services/translation.rs +++ b/src/services/translation.rs @@ -12,6 +12,12 @@ pub struct Dictionary { terms: HashMap, } +impl Default for Dictionary { + fn default() -> Self { + Self::new() + } +} + impl Dictionary { pub fn new() -> Self { Dictionary { @@ -59,9 +65,15 @@ impl Dictionary { // 基础分词清理：保留字母数字及连接符，其余视为空格以进行精确段落划分 let clean_text = text .chars() - .map(|c| if c.is_alphanumeric() || c == '-' || c == '\'' || c == ' ' { c } else { ' ' }) + .map(|c| { + if c.is_alphanumeric() || c == '-' || c == '\'' || c == ' ' { + c + } else { + ' ' + } + }) .collect::(); - + let words: Vec<&str> = clean_text.split_whitespace().collect(); let mut matched = HashSet::new(); let mut results = Vec::new(); @@ -79,13 +91,13 @@ impl Dictionary { if i + len <= n { let phrase_slice = &words[i..i + len]; let phrase = phrase_slice.join(" ").to_lowercase(); - + if self.terms.contains_key(&phrase) { // 避免重复匹配更长名词的子词 (如已匹配 'active galactic nucleus' 就不重复提取其中的 'nucleus') if !matched.contains(&phrase) { let chinese = self.terms.get(&phrase).unwrap().clone(); let original_phrase = &words[i..i + len].join(" "); - + results.push((original_phrase.clone(), chinese.clone())); matched.insert(phrase.clone()); } @@ -107,9 +119,12 @@ impl Dictionary { // 辅助清理中文译名中的“全称/缩写”等说明性前缀，只保留实际译名 fn clean_chinese_translation(raw: &str) -> String { let raw_trimmed = raw.trim(); - if raw_trimmed.starts_with("1、") || raw_trimmed.starts_with("1. ") || raw_trimmed.starts_with("1.") { + if raw_trimmed.starts_with("1、") + || raw_trimmed.starts_with("1. ") + || raw_trimmed.starts_with("1.") + { let mut parts = Vec::new(); - for part in raw_trimmed.split(|c| c == ';' || c == '；') { + for part in raw_trimmed.split([';', '；']) { let part_trimmed = part.trim(); if let Some(pos) = part_trimmed.rfind('。') { let clean = part_trimmed[pos + '。'.len_utf8()..].trim().to_string(); @@ -121,7 +136,7 @@ fn clean_chinese_translation(raw: &str) -> String { if let Some(pos) = clean.find('、') { clean = clean[pos + '、'.len_utf8()..].trim().to_string(); } else if let Some(pos) = clean.find('.') { - if clean[..pos].chars().all(|c| c.is_digit(10)) { + if clean[..pos].chars().all(|c| c.is_ascii_digit()) { clean = clean[pos + 1..].trim().to_string(); } } @@ -164,10 +179,12 @@ pub async fn translate_markdown( // 在英文文献中扫描天文词典匹配专业词汇 let matched_terms = dict.match_text(markdown_content); let mut terms_instruction = String::new(); - + if !matched_terms.is_empty() { - terms_instruction.push_str("\n\n在翻译时，请遵循以下天文学名词对照表（严格使用对应的中文译名）：\n"); - for (en, zh) in matched_terms.iter().take(50) { // 最多注入前 50 条防止超量 + terms_instruction + .push_str("\n\n在翻译时，请遵循以下天文学名词对照表（严格使用对应的中文译名）：\n"); + for (en, zh) in matched_terms.iter().take(50) { + // 最多注入前 50 条防止超量 terms_instruction.push_str(&format!("- \"{}\" 必须翻译为 \"{}\"\n", en, zh)); } } @@ -183,13 +200,24 @@ pub async fn translate_markdown( terms_instruction ); - info!("正在请求大模型开展中英翻译。所选大模型: {}", llm_client.model()); + info!( + "正在请求大模型开展中英翻译。所选大模型: {}", + llm_client.model() + ); let start_time = std::time::Instant::now(); - match llm_client.chat_completion(&system_prompt, markdown_content).await { + match llm_client + .chat_completion(&system_prompt, markdown_content) + .await + { Ok(translated) => { let duration = start_time.elapsed(); - info!("LLM 翻译成功。所选大模型: {}, 耗时: {:?}, 译文字符数: {}", llm_client.model(), duration, translated.len()); + info!( + "LLM 翻译成功。所选大模型: {}, 耗时: {:?}, 译文字符数: {}", + llm_client.model(), + duration, + translated.len() + ); Ok(translated) } Err(e) => Err(e), @@ -205,10 +233,15 @@ mod tests { fn test_dictionary_match() { let mut dict = Dictionary::new(); // 模拟词典数据 - dict.terms.insert("active galactic nucleus".to_string(), "活动星系核".to_string()); - dict.terms.insert("galactic nucleus".to_string(), "星系核".to_string()); + dict.terms.insert( + "active galactic nucleus".to_string(), + "活动星系核".to_string(), + ); + dict.terms + .insert("galactic nucleus".to_string(), "星系核".to_string()); dict.terms.insert("nucleus".to_string(), "核心".to_string()); - dict.terms.insert("black hole".to_string(), "黑洞".to_string()); + dict.terms + .insert("black hole".to_string(), "黑洞".to_string()); let text = "We study the active galactic nucleus and its central black hole."; let matched = dict.match_text(text); @@ -216,7 +249,7 @@ mod tests { let phrases: Vec = matched.iter().map(|(en, _)| en.clone()).collect(); assert!(phrases.contains(&"active galactic nucleus".to_string())); assert!(phrases.contains(&"black hole".to_string())); - + // 验证已经匹配了更长名词的子词 (如 'galactic nucleus' 和 'nucleus') 被成功过滤去掉了，不重复提取 assert!(!phrases.contains(&"galactic nucleus".to_string())); assert!(!phrases.contains(&"nucleus".to_string())); @@ -227,9 +260,17 @@ mod tests { #[test] fn test_clean_chinese_translation_helper() { - assert_eq!(clean_chinese_translation("全称：Transiting Exoplanet Survey Satellite。凌星系外行星巡天卫星"), "凌星系外行星巡天卫星"); + assert_eq!( + clean_chinese_translation( + "全称：Transiting Exoplanet Survey Satellite。凌星系外行星巡天卫星" + ), + "凌星系外行星巡天卫星" + ); assert_eq!(clean_chinese_translation("1、全称：Anglo-Australian Observatory。英澳天文台; 2、全称：Australian Astronomical Observatory。澳大利亚天文台"), "英澳天文台; 澳大利亚天文台"); - assert_eq!(clean_chinese_translation("缩写：TESS。凌星系外行星巡天卫星"), "凌星系外行星巡天卫星"); + assert_eq!( + clean_chinese_translation("缩写：TESS。凌星系外行星巡天卫星"), + "凌星系外行星巡天卫星" + ); assert_eq!(clean_chinese_translation("黑洞"), "黑洞"); } @@ -242,13 +283,16 @@ mod tests { writeln!(file, "active galactic nucleus\t活动星系核")?; writeln!(file, "black hole\t黑洞")?; } - + let mut dict = Dictionary::new(); let res = dict.load_from_file(&path); let _ = std::fs::remove_file(&path); res?; - - assert_eq!(dict.terms.get("active galactic nucleus").unwrap(), "活动星系核"); + + assert_eq!( + dict.terms.get("active galactic nucleus").unwrap(), + "活动星系核" + ); assert_eq!(dict.terms.get("black hole").unwrap(), "黑洞"); Ok(()) }