diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1b37f895..4f59786a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -38,7 +38,7 @@ jobs: with: components: clippy - uses: Swatinem/rust-cache@v2 - - run: cargo clippy -- -D warnings + - run: cargo clippy -- -W clippy::all test: name: Test diff --git a/Cargo.toml b/Cargo.toml index 0a02dd88..31d4b065 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ documentation = "https://docs.rs/vectorless" keywords = ["rag", "document", "retrieval", "indexing", "llm"] categories = ["text-processing", "data-structures", "algorithms"] readme = "README.md" -exclude = ["samples/", "docs/", "benches/", ".*"] +exclude = ["samples/", "docs/", ".*"] [dependencies] # Async runtime @@ -72,13 +72,8 @@ rand = "0.8" [dev-dependencies] tempfile = "3.10" -criterion = { version = "0.5", features = ["async_tokio"] } tokio-test = "0.4" -[[bench]] -name = "bench" -harness = false - [profile.release] opt-level = 3 lto = "thin" diff --git a/README.md b/README.md index 93fa0129..300818ff 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@
-![Vectorless](docs/assets/brand/logo-horizontal.svg) +![Vectorless](docs/design/logo-horizontal.svg) [![Crates.io](https://img.shields.io/crates/v/vectorless.svg)](https://crates.io/crates/vectorless) [![Downloads](https://img.shields.io/crates/d/vectorless.svg)](https://crates.io/crates/vectorless) @@ -8,11 +8,14 @@ [![License](https://img.shields.io/badge/license-Apache--2.0-blue.svg)](LICENSE) [![Rust](https://img.shields.io/badge/rust-1.85%2B-orange.svg)](https://www.rust-lang.org/) -**A hierarchical, reasoning-native document intelligence engine.** -
+Ultra performant document intelligence engine for RAG, with core written in **Rust**. Zero vector database, zero embedding model — just LLM-powered tree navigation. Incremental indexing and multi-format support out-of-box. + +⭐ **Drop a star to help us grow!** + + ## Why Vectorless? Traditional RAG systems have a fundamental problem: **they lose document structure.** diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 00000000..4b66d90b --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,44 @@ +# Security Policy + +## Reporting a Vulnerability + +If you discover a security vulnerability in Vectorless, please report it by emailing: + +**beautifularea@gmail.com** + +**Do NOT create a public GitHub issue for security vulnerabilities.** + +## What to Include + +Please include the following in your report: + +- Description of the vulnerability +- Steps to reproduce +- Affected versions (if known) +- Potential impact + +## Response Timeline + +| Stage | Timeframe | +|-------|-----------| +| Initial response | Within 48 hours | +| Vulnerability confirmation | Within 7 days | +| Fix development | Depends on severity | +| Security advisory | After fix is released | + +## Disclosure Policy + +- Vulnerabilities will be disclosed after a fix is available +- We will credit reporters (unless you prefer to remain anonymous) +- We request a reasonable time to fix before public disclosure + +## Supported Versions + +| Version | Supported | +| ------- | --------- | +| 0.1.x | ✅ | +| < 0.1 | ❌ | + +--- + +Thank you for helping keep Vectorless secure! diff --git a/benches/bench.rs b/benches/bench.rs deleted file mode 100644 index b33e507f..00000000 --- a/benches/bench.rs +++ /dev/null @@ -1,9 +0,0 @@ -//! Benchmark runner placeholder. -//! -//! Run `cargo bench` to execute benchmarks. - -#![allow(missing_docs)] - -fn main() { - println!("Run `cargo bench` to execute benchmarks"); -} \ No newline at end of file diff --git a/docs/assets/brand/icon.svg b/docs/assets/brand/icon.svg deleted file mode 100644 index 6899600c..00000000 --- a/docs/assets/brand/icon.svg +++ /dev/null @@ -1,24 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/docs/assets/brand/logo-dark.svg b/docs/assets/brand/logo-dark.svg deleted file mode 100644 index 646bad97..00000000 --- a/docs/assets/brand/logo-dark.svg +++ /dev/null @@ -1,27 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/docs/assets/brand/logo.svg b/docs/assets/brand/logo.svg deleted file mode 100644 index 6879f501..00000000 --- a/docs/assets/brand/logo.svg +++ /dev/null @@ -1,24 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/docs/assets/brand/logo-horizontal.svg b/docs/design/logo-horizontal.svg similarity index 100% rename from docs/assets/brand/logo-horizontal.svg rename to docs/design/logo-horizontal.svg diff --git a/examples/index.rs b/examples/index.rs index fb1686e1..cbb318b1 100644 --- a/examples/index.rs +++ b/examples/index.rs @@ -15,7 +15,7 @@ //! cargo run --example index //! ``` -use vectorless::index::{PipelineExecutor, PipelineOptions, IndexInput}; +use vectorless::index::{IndexInput, PipelineExecutor, PipelineOptions}; #[tokio::main] async fn main() -> vectorless::Result<()> { @@ -89,8 +89,15 @@ fn print_tree_structure( if let Some(node) = tree.get(node_id) { let children = tree.children(node_id); - let marker = if children.is_empty() { "└─" } else { "├─" }; - println!("{}{} {} (depth: {})", indent, marker, node.title, node.depth); + let marker = if children.is_empty() { + "└─" + } else { + "├─" + }; + println!( + "{}{} {} (depth: {})", + indent, marker, node.title, node.depth + ); for child_id in children { print_tree_structure(tree, child_id, current_depth + 1, max_depth); diff --git a/examples/markdownflow.rs b/examples/markdownflow.rs index 7854fab1..4cde85f9 100644 --- a/examples/markdownflow.rs +++ b/examples/markdownflow.rs @@ -86,9 +86,7 @@ async fn main() -> Result<(), Box> { // Step 4: Query the document println!("Step 4: Querying the document..."); - let queries = vec![ - "What is this project about?", - ]; + let queries = vec!["What is this project about?"]; for query in queries { println!(" Query: \"{}\"", query); diff --git a/examples/retrieve.rs b/examples/retrieve.rs index 036a3b3f..f3ed1751 100644 --- a/examples/retrieve.rs +++ b/examples/retrieve.rs @@ -16,12 +16,12 @@ //! ``` use std::sync::Arc; +use vectorless::domain::DocumentTree; use vectorless::retrieval::{ - PipelineRetriever, Retriever, RetrieveOptions, StrategyPreference, + PipelineRetriever, RetrieveOptions, Retriever, StrategyPreference, pipeline::RetrievalOrchestrator, stages::{AnalyzeStage, JudgeStage, PlanStage, SearchStage}, }; -use vectorless::domain::{DocumentTree, NodeId}; #[tokio::main] async fn main() -> vectorless::Result<()> { @@ -29,7 +29,10 @@ async fn main() -> vectorless::Result<()> { // 1. Create a sample document tree let tree = create_sample_tree(); - println!("✓ Created sample document tree ({} nodes)\n", tree.node_count()); + println!( + "✓ Created sample document tree ({} nodes)\n", + tree.node_count() + ); // 2. Method A: Use PipelineRetriever (simple API) println!("--- Method A: PipelineRetriever (Simple API) ---\n"); @@ -79,7 +82,9 @@ async fn demo_pipeline_retriever(tree: &DocumentTree) -> vectorless::Result<()> let query = "What is the main architecture?"; println!("Query: \"{}\"\n", query); - let response = retriever.retrieve(tree, query, &options).await + let response = retriever + .retrieve(tree, query, &options) + .await .map_err(|e| vectorless::Error::Retrieval(e.to_string()))?; // Display results @@ -93,7 +98,12 @@ async fn demo_pipeline_retriever(tree: &DocumentTree) -> vectorless::Result<()> if !response.results.is_empty() { println!("\n Top results:"); for (i, result) in response.results.iter().take(3).enumerate() { - println!(" {}. {} (score: {:.2})", i + 1, result.title, result.score); + println!( + " {}. {} (score: {:.2})", + i + 1, + result.title, + result.score + ); } } @@ -123,8 +133,17 @@ async fn demo_orchestrator(tree: &DocumentTree) -> vectorless::Result<()> { if let Ok(groups) = orchestrator.get_execution_groups() { println!("Execution groups: {} groups", groups.len()); for (i, group) in groups.iter().enumerate() { - let parallel = if group.parallel { " (can parallelize)" } else { "" }; - println!(" Group {}: {} stages{}", i, group.stage_indices.len(), parallel); + let parallel = if group.parallel { + " (can parallelize)" + } else { + "" + }; + println!( + " Group {}: {} stages{}", + i, + group.stage_indices.len(), + parallel + ); } } println!(); @@ -135,7 +154,9 @@ async fn demo_orchestrator(tree: &DocumentTree) -> vectorless::Result<()> { let options = RetrieveOptions::default(); let tree_arc = Arc::new(tree.clone()); - let response = orchestrator.execute(tree_arc, query, options).await + let response = orchestrator + .execute(tree_arc, query, options) + .await .map_err(|e| vectorless::Error::Retrieval(e.to_string()))?; println!("Results:"); @@ -162,39 +183,77 @@ fn create_sample_tree() -> DocumentTree { ); // Add sections using the correct API - let _intro = tree.add_child(tree.root(), "Introduction", - "Vectorless is a document intelligence engine written in Rust."); - - let arch = tree.add_child(tree.root(), "Architecture", - "The system consists of three main components: indexer, retriever, and storage."); - - let index_section = tree.add_child(arch, "Index Pipeline", - "The index pipeline processes documents into a tree structure with summaries."); - let retrieve_section = tree.add_child(arch, "Retrieval Pipeline", - "The retrieval pipeline finds relevant content using multi-stage processing."); - - tree.add_child(index_section, "Parse Stage", - "Parses documents (Markdown, PDF, DOCX) into structured content."); - tree.add_child(index_section, "Build Stage", - "Builds the document tree with metadata like page numbers and indices."); - tree.add_child(index_section, "Enrich Stage", - "Generates AI summaries for tree nodes using LLM."); - - tree.add_child(retrieve_section, "Analyze Stage", - "Analyzes query complexity and extracts keywords for matching."); - tree.add_child(retrieve_section, "Plan Stage", - "Selects retrieval strategy (keyword/semantic/LLM) and search algorithm."); - tree.add_child(retrieve_section, "Search Stage", - "Executes tree traversal (greedy/beam/MCTS) to find relevant content."); - tree.add_child(retrieve_section, "Judge Stage", - "Evaluates sufficiency of collected content, can trigger backtracking."); - - let usage = tree.add_child(tree.root(), "Usage", - "How to use the vectorless library."); - tree.add_child(usage, "Basic Example", - "Simple usage with default configuration and workspace."); - tree.add_child(usage, "Advanced Example", - "Custom pipeline configuration with LLM and custom stages."); + let _intro = tree.add_child( + tree.root(), + "Introduction", + "Vectorless is a document intelligence engine written in Rust.", + ); + + let arch = tree.add_child( + tree.root(), + "Architecture", + "The system consists of three main components: indexer, retriever, and storage.", + ); + + let index_section = tree.add_child( + arch, + "Index Pipeline", + "The index pipeline processes documents into a tree structure with summaries.", + ); + let retrieve_section = tree.add_child( + arch, + "Retrieval Pipeline", + "The retrieval pipeline finds relevant content using multi-stage processing.", + ); + + tree.add_child( + index_section, + "Parse Stage", + "Parses documents (Markdown, PDF, DOCX) into structured content.", + ); + tree.add_child( + index_section, + "Build Stage", + "Builds the document tree with metadata like page numbers and indices.", + ); + tree.add_child( + index_section, + "Enrich Stage", + "Generates AI summaries for tree nodes using LLM.", + ); + + tree.add_child( + retrieve_section, + "Analyze Stage", + "Analyzes query complexity and extracts keywords for matching.", + ); + tree.add_child( + retrieve_section, + "Plan Stage", + "Selects retrieval strategy (keyword/semantic/LLM) and search algorithm.", + ); + tree.add_child( + retrieve_section, + "Search Stage", + "Executes tree traversal (greedy/beam/MCTS) to find relevant content.", + ); + tree.add_child( + retrieve_section, + "Judge Stage", + "Evaluates sufficiency of collected content, can trigger backtracking.", + ); + + let usage = tree.add_child(tree.root(), "Usage", "How to use the vectorless library."); + tree.add_child( + usage, + "Basic Example", + "Simple usage with default configuration and workspace.", + ); + tree.add_child( + usage, + "Advanced Example", + "Custom pipeline configuration with LLM and custom stages.", + ); tree } diff --git a/src/client/builder.rs b/src/client/builder.rs index c36aa86d..243e047e 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -6,8 +6,8 @@ use std::path::PathBuf; use crate::config::{Config, ConfigLoader, RetrievalConfig}; -use crate::storage::Workspace; use crate::retrieval::PipelineRetriever; +use crate::storage::Workspace; use super::Engine; @@ -140,10 +140,9 @@ impl EngineBuilder { .map_err(|e| BuildError::Config(e.to_string()))? } else if let Some(config_path) = Self::find_config_file() { // Auto-detect config file - ConfigLoader::new() - .file(&config_path) - .load() - .map_err(|e| BuildError::Config(format!("Failed to load {}: {}", config_path.display(), e)))? + ConfigLoader::new().file(&config_path).load().map_err(|e| { + BuildError::Config(format!("Failed to load {}: {}", config_path.display(), e)) + })? } else { // Use defaults Config::default() @@ -154,8 +153,10 @@ impl EngineBuilder { Some(Workspace::open(path).map_err(|e| BuildError::Workspace(e.to_string()))?) } else { // Use workspace_dir from config - Some(Workspace::open(&config.storage.workspace_dir) - .map_err(|e| BuildError::Workspace(e.to_string()))?) + Some( + Workspace::open(&config.storage.workspace_dir) + .map_err(|e| BuildError::Workspace(e.to_string()))?, + ) }; // Create pipeline executor with LLM client if API key is available @@ -174,9 +175,11 @@ impl EngineBuilder { }; // Create pipeline retriever with config - let retrieval_config = self.retrieval_config.unwrap_or_else(|| config.retrieval.clone()); - let mut retriever = PipelineRetriever::new() - .with_max_iterations(retrieval_config.search.max_iterations); + let retrieval_config = self + .retrieval_config + .unwrap_or_else(|| config.retrieval.clone()); + let mut retriever = + PipelineRetriever::new().with_max_iterations(retrieval_config.search.max_iterations); // Add LLM client if API key is available in retrieval config if let Some(ref api_key) = retrieval_config.api_key { @@ -188,7 +191,9 @@ impl EngineBuilder { retriever = retriever.with_llm_client(llm_client); } - Ok(Engine::with_components(config, workspace, retriever, executor)) + Ok(Engine::with_components( + config, workspace, retriever, executor, + )) } } @@ -222,8 +227,7 @@ mod tests { #[test] fn test_builder_with_workspace() { - let builder = EngineBuilder::new() - .with_workspace("./test_workspace"); + let builder = EngineBuilder::new().with_workspace("./test_workspace"); assert_eq!(builder.workspace, Some(PathBuf::from("./test_workspace"))); } diff --git a/src/client/engine.rs b/src/client/engine.rs index a9cd693c..aeaa87b5 100644 --- a/src/client/engine.rs +++ b/src/client/engine.rs @@ -48,17 +48,17 @@ use std::path::Path; use std::sync::{Arc, Mutex, RwLock}; -use uuid::Uuid; use tracing::info; +use uuid::Uuid; use crate::config::Config; -use crate::domain::{DocumentTree, Result, Error}; +use crate::domain::{DocumentTree, Error, Result}; +use crate::index::{IndexInput, PipelineExecutor, PipelineOptions, SummaryStrategy}; use crate::parser::DocumentFormat; -use crate::storage::{Workspace, PersistedDocument, DocumentMeta as StorageMeta}; use crate::retrieval::{PipelineRetriever, Retriever}; -use crate::index::{PipelineExecutor, PipelineOptions, IndexInput, SummaryStrategy}; +use crate::storage::{DocumentMeta as StorageMeta, PersistedDocument, Workspace}; -use super::types::{IndexMode, IndexOptions, DocumentInfo, QueryResult}; +use super::types::{DocumentInfo, IndexMode, IndexOptions, QueryResult}; /// The main Engine client. /// @@ -168,10 +168,7 @@ impl Engine { }, generate_ids: options.generate_ids, summary_strategy: if options.generate_summaries { - SummaryStrategy::selective( - self.config.indexer.min_summary_tokens, - false, - ) + SummaryStrategy::selective(self.config.indexer.min_summary_tokens, false) } else { SummaryStrategy::none() }, @@ -182,16 +179,17 @@ impl Engine { // Create pipeline input and execute (with mutex lock) let input = IndexInput::file(&path); let result = { - let mut executor = self.executor.lock().map_err(|_| { - Error::Other("Pipeline executor lock poisoned".to_string()) - })?; + let mut executor = self + .executor + .lock() + .map_err(|_| Error::Other("Pipeline executor lock poisoned".to_string()))?; executor.execute(input, pipeline_options).await? }; // Build persisted document - let tree = result.tree.ok_or_else(|| { - Error::Parse("Document tree not generated".to_string()) - })?; + let tree = result + .tree + .ok_or_else(|| Error::Parse("Document tree not generated".to_string()))?; let meta = StorageMeta::new(&doc_id, &result.name, format.extension()) .with_source_path(path.to_string_lossy().to_string()) @@ -208,9 +206,9 @@ impl Engine { // Save to workspace if configured if let Some(ref workspace) = self.workspace { - let mut ws = workspace.write().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let mut ws = workspace + .write() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; ws.add(&doc)?; info!("Saved document {} to workspace", doc_id); } @@ -223,9 +221,7 @@ impl Engine { fn detect_format(&self, path: &Path, options: &IndexOptions) -> Result { match options.mode { IndexMode::Auto => { - let ext = path.extension() - .and_then(|e| e.to_str()) - .unwrap_or(""); + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); DocumentFormat::from_extension(ext) .ok_or_else(|| Error::Parse(format!("Unknown format: {}", ext))) } @@ -274,15 +270,18 @@ impl Engine { /// - No workspace is configured /// - The document is not found pub fn get_structure(&self, doc_id: &str) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; // Use read lock - Workspace::load now uses interior mutability for cache - let ws = workspace.read().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let ws = workspace + .read() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; - let doc = ws.load(doc_id)? + let doc = ws + .load(doc_id)? .ok_or_else(|| Error::DocumentNotFound(format!("Document not found: {}", doc_id)))?; Ok(doc.tree) @@ -297,15 +296,18 @@ impl Engine { /// - The document is not found /// - No page content is available pub fn get_page_content(&self, doc_id: &str, pages: &str) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; // Use read lock - Workspace::load now uses interior mutability for cache - let ws = workspace.read().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let ws = workspace + .read() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; - let doc = ws.load(doc_id)? + let doc = ws + .load(doc_id)? .ok_or_else(|| Error::DocumentNotFound(format!("Document not found: {}", doc_id)))?; if doc.pages.is_empty() { @@ -335,16 +337,19 @@ impl Engine { if part.contains('-') { let range: Vec<&str> = part.split('-').collect(); if range.len() == 2 { - let start: usize = range[0].parse() + let start: usize = range[0] + .parse() .map_err(|_| Error::Parse(format!("Invalid page number: {}", range[0])))?; - let end: usize = range[1].parse() + let end: usize = range[1] + .parse() .map_err(|_| Error::Parse(format!("Invalid page number: {}", range[1])))?; for p in start..=end { result.push(p); } } } else if !part.is_empty() { - let page: usize = part.parse() + let page: usize = part + .parse() .map_err(|_| Error::Parse(format!("Invalid page number: {}", part)))?; result.push(page); } @@ -373,15 +378,22 @@ impl Engine { .with_include_summaries(true); // Use adaptive retriever - let response = self.retriever.retrieve(&tree, question, &retrieve_options).await + let response = self + .retriever + .retrieve(&tree, question, &retrieve_options) + .await .map_err(|e| Error::Retrieval(e.to_string()))?; // Extract node IDs and build content from results - let node_ids: Vec = response.results.iter() + let node_ids: Vec = response + .results + .iter() .filter_map(|r| r.node_id.clone()) .collect(); - let content_parts: Vec = response.results.iter() + let content_parts: Vec = response + .results + .iter() .map(|r| { let mut parts = vec![format!("## {}", r.title)]; @@ -423,13 +435,15 @@ impl Engine { /// /// Returns an error if no workspace is configured. pub fn load(&self, doc_id: &str) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; // Use read lock - Workspace::load now uses interior mutability for cache - let ws = workspace.read().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let ws = workspace + .read() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; if !ws.contains(doc_id) { return Ok(false); @@ -445,12 +459,14 @@ impl Engine { /// /// Returns an error if no workspace is configured. pub fn remove(&self, doc_id: &str) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; - let mut ws = workspace.write().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let mut ws = workspace + .write() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; ws.remove(doc_id) } @@ -460,12 +476,14 @@ impl Engine { /// /// Returns an error if no workspace is configured. pub fn exists(&self, doc_id: &str) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; - let ws = workspace.read().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let ws = workspace + .read() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; Ok(ws.contains(doc_id)) } @@ -475,12 +493,14 @@ impl Engine { /// /// Returns an error if no workspace is configured. pub fn get_metadata(&self, doc_id: &str) -> Result> { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; - let ws = workspace.read().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let ws = workspace + .read() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; Ok(ws.get_meta(doc_id).map(|meta| DocumentInfo { id: meta.id.clone(), @@ -500,12 +520,14 @@ impl Engine { /// /// Returns an error if no workspace is configured. pub fn batch_remove(&self, doc_ids: &[&str]) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; - let mut ws = workspace.write().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let mut ws = workspace + .write() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; let mut removed = 0; for doc_id in doc_ids { @@ -524,12 +546,14 @@ impl Engine { /// /// Returns an error if no workspace is configured. pub fn clear(&self) -> Result { - let workspace = self.workspace.as_ref() + let workspace = self + .workspace + .as_ref() .ok_or_else(|| Error::Config("No workspace configured".to_string()))?; - let mut ws = workspace.write().map_err(|_| { - Error::Other("Workspace lock poisoned".to_string()) - })?; + let mut ws = workspace + .write() + .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; let doc_ids: Vec = ws.list_documents().iter().map(|s| s.to_string()).collect(); let count = doc_ids.len(); diff --git a/src/client/mod.rs b/src/client/mod.rs index 7befc7c0..907d8c0e 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -44,19 +44,12 @@ //! - **Workspace Persistence** — Save and load indexed documents //! - **Builder Pattern** — Flexible client configuration -mod types; mod builder; mod engine; +mod types; // Re-export main types -pub use types::{ - IndexedDocument, - IndexMode, - IndexOptions, - PageContent, - QueryResult, - DocumentInfo, -}; +pub use types::{DocumentInfo, IndexMode, IndexOptions, IndexedDocument, PageContent, QueryResult}; -pub use builder::{EngineBuilder, BuildError}; +pub use builder::{BuildError, EngineBuilder}; pub use engine::Engine; diff --git a/src/config/mod.rs b/src/config/mod.rs index 953300e2..23e98f4e 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -10,8 +10,8 @@ //! - [`RetrievalConfig`] - Retrieval model settings //! - [`StorageConfig`] - Storage paths -mod types; mod loader; +mod types; +pub use loader::{ConfigError, ConfigLoader}; pub use types::*; -pub use loader::{ConfigLoader, ConfigError}; diff --git a/src/domain/mod.rs b/src/domain/mod.rs index 86d3e668..d5aa3e5c 100644 --- a/src/domain/mod.rs +++ b/src/domain/mod.rs @@ -16,12 +16,12 @@ mod error; mod node; -mod token; mod toc; +mod token; mod tree; pub use error::{Error, Result}; pub use node::{NodeId, TreeNode}; -pub use token::{estimate_tokens, estimate_tokens_fast, estimate_tokens_batch}; pub use toc::{TocConfig, TocEntry, TocNode, TocView}; -pub use tree::{DocumentStructure, StructureNode, DocumentTree}; +pub use token::{estimate_tokens, estimate_tokens_batch, estimate_tokens_fast}; +pub use tree::{DocumentStructure, DocumentTree, StructureNode}; diff --git a/src/domain/node.rs b/src/domain/node.rs index 55070b3d..ea9939b2 100644 --- a/src/domain/node.rs +++ b/src/domain/node.rs @@ -6,7 +6,7 @@ //! This module provides a node type for hierarchical document representation. //! Each branch represents a section and each leaf contains the actual text. -use indextree::{NodeId as IndexTreeNodeId}; +use indextree::NodeId as IndexTreeNodeId; use serde::{Deserialize, Serialize}; use std::fmt; diff --git a/src/domain/toc.rs b/src/domain/toc.rs index de564d46..6f1806ef 100644 --- a/src/domain/toc.rs +++ b/src/domain/toc.rs @@ -9,8 +9,8 @@ use serde::{Deserialize, Serialize}; use super::node::NodeId; -use super::tree::DocumentTree; use super::node::TreeNode; +use super::tree::DocumentTree; /// A node in the Table of Contents. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -84,7 +84,11 @@ impl TocNode { if self.children.is_empty() { self.depth } else { - self.children.iter().map(|c| c.max_depth()).max().unwrap_or(self.depth) + self.children + .iter() + .map(|c| c.max_depth()) + .max() + .unwrap_or(self.depth) } } } @@ -177,12 +181,13 @@ impl TocView { } // Check minimum content length - if node.content.len() < self.config.min_content_length && tree.children(node_id).is_empty() { + if node.content.len() < self.config.min_content_length && tree.children(node_id).is_empty() + { return TocNode::new(node.title.clone(), depth); } - let mut toc_node = TocNode::new(&node.title, depth) - .with_node_id(node.node_id.clone().unwrap_or_default()); + let mut toc_node = + TocNode::new(&node.title, depth).with_node_id(node.node_id.clone().unwrap_or_default()); // Add page range if self.config.include_pages { @@ -212,7 +217,12 @@ impl TocView { entries } - fn collect_flat_entries(&self, tree: &DocumentTree, node_id: NodeId, entries: &mut Vec) { + fn collect_flat_entries( + &self, + tree: &DocumentTree, + node_id: NodeId, + entries: &mut Vec, + ) { if let Some(node) = tree.get(node_id) { entries.push(TocEntry { title: node.title.clone(), @@ -237,8 +247,13 @@ impl TocView { result } - fn collect_filtered(&self, tree: &DocumentTree, node_id: NodeId, filter: &F, result: &mut Vec) - where + fn collect_filtered( + &self, + tree: &DocumentTree, + node_id: NodeId, + filter: &F, + result: &mut Vec, + ) where F: Fn(&TreeNode) -> bool, { if let Some(node) = tree.get(node_id) { @@ -320,9 +335,7 @@ mod tests { #[test] fn test_toc_config() { - let config = TocConfig::new() - .with_max_depth(3) - .with_summaries(false); + let config = TocConfig::new().with_max_depth(3).with_summaries(false); assert_eq!(config.max_depth, Some(3)); assert!(!config.include_summaries); diff --git a/src/domain/token.rs b/src/domain/token.rs index bdd1c13a..129032c2 100644 --- a/src/domain/token.rs +++ b/src/domain/token.rs @@ -15,8 +15,7 @@ static BPE: OnceLock = OnceLock::new(); /// Get or initialize the BPE encoder. fn get_bpe() -> &'static CoreBPE { BPE.get_or_init(|| { - tiktoken_rs::cl100k_base() - .expect("Failed to initialize cl100k_base tokenizer") + tiktoken_rs::cl100k_base().expect("Failed to initialize cl100k_base tokenizer") }) } diff --git a/src/domain/tree.rs b/src/domain/tree.rs index 48685918..1f63bbff 100644 --- a/src/domain/tree.rs +++ b/src/domain/tree.rs @@ -7,7 +7,7 @@ //! lifetime management compared to `Rc`. use indextree::Arena; -use serde::{Serialize, Deserialize}; +use serde::{Deserialize, Serialize}; use super::node::{NodeId, TreeNode}; @@ -71,7 +71,10 @@ impl DocumentTree { }; let root_id = arena.new_node(root_data); - Self { arena, root_id: NodeId(root_id) } + Self { + arena, + root_id: NodeId(root_id), + } } /// Create a document tree from an existing arena and root ID. @@ -274,9 +277,11 @@ impl DocumentTree { /// Recursively build structure nodes starting from the given node. fn build_structure_nodes(&self, node_id: NodeId) -> Vec { let children = self.children(node_id); - children.into_iter().enumerate().map(|(idx, child_id)| { - self.node_to_structure(child_id, idx) - }).collect() + children + .into_iter() + .enumerate() + .map(|(idx, child_id)| self.node_to_structure(child_id, idx)) + .collect() } /// Convert a single node to StructureNode format. @@ -286,11 +291,22 @@ impl DocumentTree { StructureNode { title: node.title, - node_id: node.node_id.clone().unwrap_or_else(|| format!("{:04}", _idx)), + node_id: node + .node_id + .clone() + .unwrap_or_else(|| format!("{:04}", _idx)), start_index: node.start_index, end_index: node.end_index, - summary: if node.summary.is_empty() { None } else { Some(node.summary) }, - nodes: children.into_iter().enumerate().map(|(i, c)| self.node_to_structure(c, i)).collect(), + summary: if node.summary.is_empty() { + None + } else { + Some(node.summary) + }, + nodes: children + .into_iter() + .enumerate() + .map(|(i, c)| self.node_to_structure(c, i)) + .collect(), } } } diff --git a/src/index/incremental/detector.rs b/src/index/incremental/detector.rs index 467f78ba..688197b0 100644 --- a/src/index/incremental/detector.rs +++ b/src/index/incremental/detector.rs @@ -63,10 +63,7 @@ impl ChangeSet { /// Get total number of changes. pub fn total_changes(&self) -> usize { - self.added.len() - + self.removed.len() - + self.modified.len() - + self.restructured.len() + self.added.len() + self.removed.len() + self.modified.len() + self.restructured.len() } /// Merge another change set into this one. @@ -205,11 +202,14 @@ impl ChangeDetector { continue; } - info.insert(node.title.clone(), NodeInfo { - node_id: node.node_id.clone(), - content_hash: Self::hash_content(&node.content), - child_count: tree.children(node_id).len(), - }); + info.insert( + node.title.clone(), + NodeInfo { + node_id: node.node_id.clone(), + content_hash: Self::hash_content(&node.content), + child_count: tree.children(node_id).len(), + }, + ); } } diff --git a/src/index/incremental/updater.rs b/src/index/incremental/updater.rs index 52b62fd3..2762df9b 100644 --- a/src/index/incremental/updater.rs +++ b/src/index/incremental/updater.rs @@ -5,7 +5,7 @@ use tracing::info; -use crate::domain::{NodeId, Result, DocumentTree}; +use crate::domain::{DocumentTree, NodeId, Result}; use crate::parser::RawNode; use super::detector::ChangeDetector; @@ -122,7 +122,11 @@ impl PartialUpdater { .unwrap_or(tree.root()); // Create node - let content = if raw.content.is_empty() { "" } else { &raw.content }; + let content = if raw.content.is_empty() { + "" + } else { + &raw.content + }; let node_id = tree.add_child(parent_id, &raw.title, content); // Set line indices diff --git a/src/index/mod.rs b/src/index/mod.rs index 7f046e4c..0eb72f7c 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -44,15 +44,18 @@ pub mod summary; // Re-export main types from pipeline pub use pipeline::{ - ExecutionGroup, FailurePolicy, IndexContext, IndexInput, IndexMetrics, - IndexResult, PipelineExecutor, PipelineOrchestrator, StageResult, StageRetryConfig, + ExecutionGroup, FailurePolicy, IndexContext, IndexInput, IndexMetrics, IndexResult, + PipelineExecutor, PipelineOrchestrator, StageResult, StageRetryConfig, }; // Re-export stages pub use stages::IndexStage; // Re-export summary -pub use summary::{SummaryStrategy, SummaryStrategyConfig, SummaryGenerator, LlmSummaryGenerator, FullStrategy, SelectiveStrategy, LazyStrategy}; +pub use summary::{ + FullStrategy, LazyStrategy, LlmSummaryGenerator, SelectiveStrategy, SummaryGenerator, + SummaryStrategy, SummaryStrategyConfig, +}; // Re-export incremental pub use incremental::{ChangeDetector, ChangeSet, PartialUpdater}; diff --git a/src/index/pipeline/context.rs b/src/index/pipeline/context.rs index 773b642a..656d7909 100644 --- a/src/index/pipeline/context.rs +++ b/src/index/pipeline/context.rs @@ -7,11 +7,11 @@ use std::collections::HashMap; use std::path::PathBuf; use crate::domain::{DocumentTree, NodeId}; -use crate::parser::{DocumentFormat, RawNode}; use crate::llm::LlmClient; +use crate::parser::{DocumentFormat, RawNode}; -use super::metrics::IndexMetrics; use super::super::{PipelineOptions, SummaryStrategy}; +use super::metrics::IndexMetrics; /// Input for the index pipeline. #[derive(Debug, Clone)] @@ -37,7 +37,11 @@ impl IndexInput { } /// Create input from content. - pub fn content(content: impl Into, name: impl Into, format: DocumentFormat) -> Self { + pub fn content( + content: impl Into, + name: impl Into, + format: DocumentFormat, + ) -> Self { Self::Content { content: content.into(), name: name.into(), @@ -63,7 +67,7 @@ impl StageResult { /// Create a successful result. pub fn success(name: &str) -> Self { println!("Stage '{}' completed successfully", name); - + Self { success: true, duration_ms: 0, @@ -76,7 +80,10 @@ impl StageResult { println!("Stage '{}' failed: {}", name, error); let mut metadata = HashMap::new(); - metadata.insert("error".to_string(), serde_json::Value::String(error.to_string())); + metadata.insert( + "error".to_string(), + serde_json::Value::String(error.to_string()), + ); Self { success: false, duration_ms: 0, diff --git a/src/index/pipeline/executor.rs b/src/index/pipeline/executor.rs index b5cbca13..e1c12506 100644 --- a/src/index/pipeline/executor.rs +++ b/src/index/pipeline/executor.rs @@ -11,13 +11,12 @@ use tracing::info; use crate::domain::Result; use crate::llm::LlmClient; -use super::context::{IndexInput, IndexResult}; -use super::orchestrator::PipelineOrchestrator; +use super::super::PipelineOptions; use super::super::stages::{ - BuildStage, EnhanceStage, EnrichStage, IndexStage, OptimizeStage, - ParseStage, PersistStage, + BuildStage, EnhanceStage, EnrichStage, IndexStage, OptimizeStage, ParseStage, PersistStage, }; -use super::super::PipelineOptions; +use super::context::{IndexInput, IndexResult}; +use super::orchestrator::PipelineOrchestrator; /// Pipeline executor for document indexing. /// @@ -112,7 +111,11 @@ impl PipelineExecutor { /// Add a stage with custom priority. /// /// Lower priority = earlier execution. - pub fn add_stage_with_priority(mut self, stage: impl IndexStage + 'static, priority: i32) -> Self { + pub fn add_stage_with_priority( + mut self, + stage: impl IndexStage + 'static, + priority: i32, + ) -> Self { self.orchestrator = self.orchestrator.stage_with_priority(stage, priority); self } @@ -126,13 +129,16 @@ impl PipelineExecutor { priority: i32, depends_on: &[&str], ) -> Self { - self.orchestrator = self.orchestrator.stage_with_deps(stage, priority, depends_on); + self.orchestrator = self + .orchestrator + .stage_with_deps(stage, priority, depends_on); self } /// Add persistence stage with workspace. pub fn with_persistence(mut self, workspace: crate::storage::Workspace) -> Self { - self.orchestrator = self.orchestrator + self.orchestrator = self + .orchestrator .stage_with_priority(PersistStage::with_workspace(workspace), 80); self } @@ -150,8 +156,15 @@ impl PipelineExecutor { /// Execute the pipeline. /// /// Stages are executed in dependency-resolved order. - pub async fn execute(&mut self, input: IndexInput, options: PipelineOptions) -> Result { - info!("Starting index pipeline with {} stages", self.orchestrator.stage_count()); + pub async fn execute( + &mut self, + input: IndexInput, + options: PipelineOptions, + ) -> Result { + info!( + "Starting index pipeline with {} stages", + self.orchestrator.stage_count() + ); self.orchestrator.execute(input, options).await } } diff --git a/src/index/pipeline/mod.rs b/src/index/pipeline/mod.rs index 0442aaee..fdf22827 100644 --- a/src/index/pipeline/mod.rs +++ b/src/index/pipeline/mod.rs @@ -20,5 +20,5 @@ mod policy; pub use context::{IndexContext, IndexInput, IndexResult, StageResult}; pub use executor::PipelineExecutor; pub use metrics::IndexMetrics; -pub use orchestrator::{ExecutionGroup, PipelineOrchestrator, CustomStageBuilder}; +pub use orchestrator::{CustomStageBuilder, ExecutionGroup, PipelineOrchestrator}; pub use policy::{FailurePolicy, StageRetryConfig}; diff --git a/src/index/pipeline/orchestrator.rs b/src/index/pipeline/orchestrator.rs index b9dbb2e9..fb471f51 100644 --- a/src/index/pipeline/orchestrator.rs +++ b/src/index/pipeline/orchestrator.rs @@ -25,14 +25,14 @@ use std::collections::HashMap; use std::time::Instant; -use tracing::{info, warn, error}; +use tracing::{error, info, warn}; use crate::domain::Result; +use super::super::PipelineOptions; +use super::super::stages::IndexStage; use super::context::{IndexContext, IndexInput, IndexResult, StageResult}; use super::policy::FailurePolicy; -use super::super::stages::IndexStage; -use super::super::PipelineOptions; /// Stage entry with metadata for orchestration. struct StageEntry { @@ -155,10 +155,7 @@ impl PipelineOrchestrator { S: IndexStage + 'static, { let trait_deps = stage.depends_on(); - let mut all_deps: Vec = trait_deps - .into_iter() - .map(|s| s.to_string()) - .collect(); + let mut all_deps: Vec = trait_deps.into_iter().map(|s| s.to_string()).collect(); // Add explicit deps that aren't already included for dep in explicit_depends_on { @@ -235,9 +232,7 @@ impl PipelineOrchestrator { } // Collect stages with no dependencies, sorted by priority - let mut ready: Vec = (0..n) - .filter(|&i| in_degree[i] == 0) - .collect(); + let mut ready: Vec = (0..n).filter(|&i| in_degree[i] == 0).collect(); ready.sort_by_key(|&i| (self.stages[i].priority, i)); let mut result: Vec = Vec::new(); @@ -405,19 +400,20 @@ impl PipelineOrchestrator { options: PipelineOptions, ) -> Result { let total_start = Instant::now(); - info!("Starting orchestrated pipeline with {} stages", self.stages.len()); + info!( + "Starting orchestrated pipeline with {} stages", + self.stages.len() + ); // Resolve execution order let order = self.resolve_order()?; - let stage_names: Vec<&str> = order - .iter() - .map(|&i| self.stages[i].stage.name()) - .collect(); + let stage_names: Vec<&str> = order.iter().map(|&i| self.stages[i].stage.name()).collect(); info!("Execution order: {:?}", stage_names); // Compute execution groups for potential parallelization let groups = self.compute_execution_groups(&order); - info!("Execution groups: {} ({} parallelizable)", + info!( + "Execution groups: {} ({} parallelizable)", groups.len(), groups.iter().filter(|g| g.parallel).count() ); @@ -432,7 +428,9 @@ impl PipelineOrchestrator { "Executing parallel group {} with {} stages: {:?}", group_idx, group.stage_indices.len(), - group.stage_indices.iter() + group + .stage_indices + .iter() .map(|&i| self.stages[i].stage.name()) .collect::>() ); @@ -447,7 +445,10 @@ impl PipelineOrchestrator { let stage_name = entry.stage.name().to_string(); let policy = entry.stage.failure_policy(); - info!("Executing stage: {} (priority {})", stage_name, entry.priority); + info!( + "Executing stage: {} (priority {})", + stage_name, entry.priority + ); match Self::execute_stage_with_policy(&mut entry.stage, &mut ctx).await { Ok(result) => { @@ -455,7 +456,10 @@ impl PipelineOrchestrator { } Err(e) => { if policy.allows_continuation() { - warn!("Stage {} failed but policy allows continuation: {}", stage_name, e); + warn!( + "Stage {} failed but policy allows continuation: {}", + stage_name, e + ); ctx.stage_results.insert( stage_name.clone(), StageResult::failure(&stage_name, &e.to_string()), @@ -557,8 +561,8 @@ impl CustomStageBuilder { #[cfg(test)] mod tests { - use super::*; use super::super::context::StageResult; + use super::*; #[test] fn test_orchestrator_creation() { @@ -592,8 +596,8 @@ mod tests { #[test] fn test_missing_dependency() { - let orchestrator = PipelineOrchestrator::new() - .stage_with_deps(MockStage::new("a"), 10, &["nonexistent"]); + let orchestrator = + PipelineOrchestrator::new().stage_with_deps(MockStage::new("a"), 10, &["nonexistent"]); let result = orchestrator.stage_names(); assert!(result.is_err()); @@ -631,7 +635,9 @@ mod tests { impl MockStage { fn new(name: &str) -> Self { - Self { name: name.to_string() } + Self { + name: name.to_string(), + } } } @@ -641,10 +647,7 @@ mod tests { &self.name } - async fn execute( - &mut self, - _ctx: &mut IndexContext, - ) -> Result { + async fn execute(&mut self, _ctx: &mut IndexContext) -> Result { Ok(StageResult::success(&self.name)) } } diff --git a/src/index/pipeline/policy.rs b/src/index/pipeline/policy.rs index 9af54bfa..da3c5b2b 100644 --- a/src/index/pipeline/policy.rs +++ b/src/index/pipeline/policy.rs @@ -87,8 +87,8 @@ impl StageRetryConfig { /// /// Uses exponential backoff: `initial_delay * multiplier^attempt` pub fn delay_for_attempt(&self, attempt: usize) -> Duration { - let delay_ms = (self.initial_delay.as_millis() as f64) - * self.multiplier.powi(attempt as i32); + let delay_ms = + (self.initial_delay.as_millis() as f64) * self.multiplier.powi(attempt as i32); let capped_ms = delay_ms.min(self.max_delay.as_millis() as f64); Duration::from_millis(capped_ms as u64) } diff --git a/src/index/stages/build.rs b/src/index/stages/build.rs index 500741c2..ed7f0ee9 100644 --- a/src/index/stages/build.rs +++ b/src/index/stages/build.rs @@ -11,8 +11,8 @@ use crate::domain::{DocumentTree, NodeId, Result, estimate_tokens}; use crate::parser::RawNode; use super::{IndexStage, StageResult}; -use crate::index::pipeline::IndexContext; use crate::index::ThinningConfig; +use crate::index::pipeline::IndexContext; /// Build stage - constructs a tree from raw nodes. pub struct BuildStage; @@ -31,7 +31,9 @@ impl BuildStage { // Process from back to front for i in (0..nodes.len()).rev() { - let own_tokens = nodes[i].token_count.unwrap_or_else(|| estimate_tokens(&nodes[i].content)); + let own_tokens = nodes[i] + .token_count + .unwrap_or_else(|| estimate_tokens(&nodes[i].content)); nodes[i].token_count = Some(own_tokens); // Find all children (direct and indirect) @@ -150,7 +152,11 @@ impl BuildStage { .unwrap_or(tree.root()); // Create the node - let content = if raw.content.is_empty() { "" } else { &raw.content }; + let content = if raw.content.is_empty() { + "" + } else { + &raw.content + }; let node_id = tree.add_child(parent_id, &raw.title, content); // Set line indices @@ -208,7 +214,7 @@ impl Default for BuildStage { #[async_trait] impl IndexStage for BuildStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "build" } @@ -272,10 +278,9 @@ impl IndexStage for BuildStage { "node_count".to_string(), serde_json::json!(ctx.tree.as_ref().map(|t| t.node_count()).unwrap_or(0)), ); - stage_result.metadata.insert( - "nodes_skipped".to_string(), - serde_json::json!(skipped), - ); + stage_result + .metadata + .insert("nodes_skipped".to_string(), serde_json::json!(skipped)); Ok(stage_result) } diff --git a/src/index/stages/enhance.rs b/src/index/stages/enhance.rs index ec150915..f510e2e0 100644 --- a/src/index/stages/enhance.rs +++ b/src/index/stages/enhance.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use std::time::Instant; use tracing::{info, warn}; -use crate::domain::{NodeId, Result, DocumentTree}; +use crate::domain::{DocumentTree, NodeId, Result}; use crate::llm::LlmClient; use super::{IndexStage, StageResult}; @@ -73,12 +73,16 @@ impl EnhanceStage { // Generate summary match generator.generate(&node.title, &node.content).await { Ok(summary) => { - if !summary.is_empty() { + if summary.is_empty() { + warn!("Empty summary returned for node '{}'", node.title); + } else { tree.set_summary(node_id, &summary); - info!("Generated summary for node: {} ({} chars)", node.title, summary.len()); + info!( + "Generated summary for node: {} ({} chars)", + node.title, + summary.len() + ); metrics.increment_summaries(); - } else { - warn!("Empty summary returned for node '{}'", node.title); } } Err(e) => { @@ -98,7 +102,7 @@ impl Default for EnhanceStage { #[async_trait] impl IndexStage for EnhanceStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "enhance" } @@ -115,7 +119,7 @@ impl IndexStage for EnhanceStage { FailurePolicy::retry_with( StageRetryConfig::new() .with_max_attempts(2) - .with_initial_delay(std::time::Duration::from_millis(500)) + .with_initial_delay(std::time::Duration::from_millis(500)), ) } @@ -124,7 +128,10 @@ impl IndexStage for EnhanceStage { // Check if we need summaries if !self.needs_summaries(ctx) { - info!("Summary generation skipped (strategy: {:?})", ctx.options.summary_strategy); + info!( + "Summary generation skipped (strategy: {:?})", + ctx.options.summary_strategy + ); return Ok(StageResult::success("enhance")); } @@ -164,7 +171,15 @@ impl IndexStage for EnhanceStage { let strategy = ctx.options.summary_strategy.clone(); for node_id in node_ids { - match Self::generate_node_summary(tree, node_id, &generator, &strategy, &mut ctx.metrics).await { + match Self::generate_node_summary( + tree, + node_id, + &generator, + &strategy, + &mut ctx.metrics, + ) + .await + { Ok(()) => { generated += 1; } @@ -183,9 +198,7 @@ impl IndexStage for EnhanceStage { info!( "Generated {} summaries ({} failed) in {}ms", - generated, - failed, - duration + generated, failed, duration ); let mut stage_result = StageResult::success("enhance"); @@ -194,10 +207,9 @@ impl IndexStage for EnhanceStage { "summaries_generated".to_string(), serde_json::json!(generated), ); - stage_result.metadata.insert( - "summaries_failed".to_string(), - serde_json::json!(failed), - ); + stage_result + .metadata + .insert("summaries_failed".to_string(), serde_json::json!(failed)); Ok(stage_result) } diff --git a/src/index/stages/enrich.rs b/src/index/stages/enrich.rs index 0b711e01..2c3759fe 100644 --- a/src/index/stages/enrich.rs +++ b/src/index/stages/enrich.rs @@ -7,7 +7,7 @@ use super::async_trait; use std::time::Instant; use tracing::info; -use crate::domain::{NodeId, Result, DocumentTree, TocView}; +use crate::domain::{DocumentTree, NodeId, Result, TocView}; use super::{IndexStage, StageResult}; use crate::index::pipeline::IndexContext; @@ -102,7 +102,7 @@ impl Default for EnrichStage { #[async_trait] impl IndexStage for EnrichStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "enrich" } @@ -113,9 +113,10 @@ impl IndexStage for EnrichStage { async fn execute(&mut self, ctx: &mut IndexContext) -> Result { let start = Instant::now(); - let tree = ctx.tree.as_mut().ok_or_else(|| { - crate::domain::Error::IndexBuild("Tree not built".to_string()) - })?; + let tree = ctx + .tree + .as_mut() + .ok_or_else(|| crate::domain::Error::IndexBuild("Tree not built".to_string()))?; // 1. Calculate page ranges Self::calculate_page_ranges(tree); @@ -141,14 +142,12 @@ impl IndexStage for EnrichStage { let mut stage_result = StageResult::success("enrich"); stage_result.duration_ms = duration; - stage_result.metadata.insert( - "total_tokens".to_string(), - serde_json::json!(total_tokens), - ); - stage_result.metadata.insert( - "node_count".to_string(), - serde_json::json!(node_count), - ); + stage_result + .metadata + .insert("total_tokens".to_string(), serde_json::json!(total_tokens)); + stage_result + .metadata + .insert("node_count".to_string(), serde_json::json!(node_count)); Ok(stage_result) } diff --git a/src/index/stages/mod.rs b/src/index/stages/mod.rs index 5568adc9..9d6f8c85 100644 --- a/src/index/stages/mod.rs +++ b/src/index/stages/mod.rs @@ -3,23 +3,23 @@ //! Index pipeline stages. -mod parse; mod build; mod enhance; mod enrich; mod optimize; +mod parse; mod persist; -pub use parse::ParseStage; pub use build::BuildStage; pub use enhance::EnhanceStage; pub use enrich::EnrichStage; pub use optimize::OptimizeStage; +pub use parse::ParseStage; pub use persist::PersistStage; -pub use async_trait::async_trait; -use crate::domain::Result; use super::pipeline::{FailurePolicy, IndexContext, StageResult}; +use crate::domain::Result; +pub use async_trait::async_trait; /// Index pipeline stage. /// diff --git a/src/index/stages/optimize.rs b/src/index/stages/optimize.rs index e2b4ae3e..d84633bf 100644 --- a/src/index/stages/optimize.rs +++ b/src/index/stages/optimize.rs @@ -30,7 +30,8 @@ impl OptimizeStage { let mut merged_count = 0; // Get all non-leaf nodes - let non_leaves: Vec = tree.traverse() + let non_leaves: Vec = tree + .traverse() .into_iter() .filter(|id| !tree.is_leaf(*id)) .collect(); @@ -47,12 +48,8 @@ impl OptimizeStage { let curr_id = children[i]; let next_id = children[i + 1]; - let curr_tokens = tree.get(curr_id) - .and_then(|n| n.token_count) - .unwrap_or(0); - let next_tokens = tree.get(next_id) - .and_then(|n| n.token_count) - .unwrap_or(0); + let curr_tokens = tree.get(curr_id).and_then(|n| n.token_count).unwrap_or(0); + let next_tokens = tree.get(next_id).and_then(|n| n.token_count).unwrap_or(0); // If both are small, merge next into current if curr_tokens < min_tokens && next_tokens < min_tokens { @@ -65,9 +62,7 @@ impl OptimizeStage { } curr.content.push_str(&next_node.content); } - curr.token_count = Some( - curr.token_count.unwrap_or(0) + next_tokens - ); + curr.token_count = Some(curr.token_count.unwrap_or(0) + next_tokens); } } @@ -95,7 +90,8 @@ impl OptimizeStage { let mut removed_count = 0; // Find nodes with no content and only one child - let candidates: Vec = tree.traverse() + let candidates: Vec = tree + .traverse() .into_iter() .filter(|id| { if tree.is_leaf(*id) { @@ -134,7 +130,7 @@ impl Default for OptimizeStage { #[async_trait] impl IndexStage for OptimizeStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "optimize" } @@ -155,15 +151,17 @@ impl IndexStage for OptimizeStage { return Ok(StageResult::success("optimize")); } - let tree = ctx.tree.as_mut().ok_or_else(|| { - crate::domain::Error::IndexBuild("Tree not built".to_string()) - })?; + let tree = ctx + .tree + .as_mut() + .ok_or_else(|| crate::domain::Error::IndexBuild("Tree not built".to_string()))?; let mut merged_count = 0; // 1. Merge small leaves if config.merge_leaf_threshold > 0 { - merged_count = Self::merge_small_leaves(tree, config.merge_leaf_threshold, &mut ctx.metrics); + merged_count = + Self::merge_small_leaves(tree, config.merge_leaf_threshold, &mut ctx.metrics); info!("Merged {} small leaf nodes", merged_count); } @@ -180,10 +178,9 @@ impl IndexStage for OptimizeStage { let mut stage_result = StageResult::success("optimize"); stage_result.duration_ms = duration; - stage_result.metadata.insert( - "nodes_merged".to_string(), - serde_json::json!(merged_count), - ); + stage_result + .metadata + .insert("nodes_merged".to_string(), serde_json::json!(merged_count)); stage_result.metadata.insert( "nodes_removed".to_string(), serde_json::json!(removed_count), diff --git a/src/index/stages/parse.rs b/src/index/stages/parse.rs index da7e502a..0322760e 100644 --- a/src/index/stages/parse.rs +++ b/src/index/stages/parse.rs @@ -12,8 +12,8 @@ use crate::parser::DocumentFormat; use crate::parser::ParserRegistry; use super::{IndexStage, StageResult}; -use crate::index::pipeline::{IndexContext, IndexInput}; use crate::index::IndexMode; +use crate::index::pipeline::{IndexContext, IndexInput}; /// Parse stage - extracts raw nodes from documents. pub struct ParseStage { @@ -31,19 +31,15 @@ impl ParseStage { /// Detect document format from path and options. fn detect_format(&self, ctx: &IndexContext) -> Result { match ctx.options.mode { - IndexMode::Auto => { - match &ctx.input { - IndexInput::File(path) => { - let ext = path - .extension() - .and_then(|e| e.to_str()) - .unwrap_or(""); - DocumentFormat::from_extension(ext) - .ok_or_else(|| crate::domain::Error::Parse(format!("Unknown format: {}", ext))) - } - IndexInput::Content { format, .. } => Ok(*format), + IndexMode::Auto => match &ctx.input { + IndexInput::File(path) => { + let ext = path.extension().and_then(|e| e.to_str()).unwrap_or(""); + DocumentFormat::from_extension(ext).ok_or_else(|| { + crate::domain::Error::Parse(format!("Unknown format: {}", ext)) + }) } - } + IndexInput::Content { format, .. } => Ok(*format), + }, IndexMode::Markdown => Ok(DocumentFormat::Markdown), IndexMode::Pdf => Ok(DocumentFormat::Pdf), IndexMode::Docx => Ok(DocumentFormat::Docx), @@ -60,7 +56,7 @@ impl Default for ParseStage { #[async_trait] impl IndexStage for ParseStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "parse" } @@ -90,7 +86,11 @@ impl IndexStage for ParseStage { // Parse using registry self.parser_registry.parse_file(&path).await? } - IndexInput::Content { content, name, format } => { + IndexInput::Content { + content, + name, + format, + } => { // Set name ctx.name = name.clone(); @@ -129,10 +129,9 @@ impl IndexStage for ParseStage { "node_count".to_string(), serde_json::json!(ctx.raw_nodes.len()), ); - stage_result.metadata.insert( - "format".to_string(), - serde_json::json!(format.extension()), - ); + stage_result + .metadata + .insert("format".to_string(), serde_json::json!(format.extension())); Ok(stage_result) } diff --git a/src/index/stages/persist.rs b/src/index/stages/persist.rs index 4dabb338..d2ac2e47 100644 --- a/src/index/stages/persist.rs +++ b/src/index/stages/persist.rs @@ -8,7 +8,7 @@ use std::time::Instant; use tracing::info; use crate::domain::Result; -use crate::storage::{PersistedDocument, DocumentMeta as StorageMeta, Workspace}; +use crate::storage::{DocumentMeta as StorageMeta, PersistedDocument, Workspace}; use super::{IndexStage, StageResult}; use crate::index::pipeline::IndexContext; @@ -34,22 +34,20 @@ impl PersistStage { /// Save document to workspace. fn save_to_workspace(&mut self, ctx: &IndexContext) -> Result<()> { - let workspace = self.workspace.as_mut().ok_or_else(|| { - crate::domain::Error::Config("No workspace configured".to_string()) - })?; + let workspace = self + .workspace + .as_mut() + .ok_or_else(|| crate::domain::Error::Config("No workspace configured".to_string()))?; - let tree = ctx.tree.as_ref().ok_or_else(|| { - crate::domain::Error::IndexBuild("Tree not built".to_string()) - })?; + let tree = ctx + .tree + .as_ref() + .ok_or_else(|| crate::domain::Error::IndexBuild("Tree not built".to_string()))?; // Create metadata - let meta = StorageMeta::new( - &ctx.doc_id, - &ctx.name, - ctx.format.extension(), - ) - .with_source_path(ctx.source_path.clone().unwrap_or_default()) - .with_description(ctx.description.clone().unwrap_or_default()); + let meta = StorageMeta::new(&ctx.doc_id, &ctx.name, ctx.format.extension()) + .with_source_path(ctx.source_path.clone().unwrap_or_default()) + .with_description(ctx.description.clone().unwrap_or_default()); let doc = PersistedDocument::new(meta, tree.clone()); @@ -71,7 +69,7 @@ impl Default for PersistStage { #[async_trait] impl IndexStage for PersistStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "persist" } diff --git a/src/index/summary/lazy.rs b/src/index/summary/lazy.rs index 870ae42a..6d9cadef 100644 --- a/src/index/summary/lazy.rs +++ b/src/index/summary/lazy.rs @@ -30,7 +30,9 @@ impl LazyStrategy { /// Create a new lazy strategy with LLM client. pub fn new(client: LlmClient) -> Self { Self { - generator: Arc::new(RwLock::new(Box::new(super::LlmSummaryGenerator::new(client)))), + generator: Arc::new(RwLock::new(Box::new(super::LlmSummaryGenerator::new( + client, + )))), cache: Arc::new(RwLock::new(HashMap::new())), persist: false, config: SummaryStrategyConfig::default(), @@ -40,7 +42,9 @@ impl LazyStrategy { /// Create with persistence enabled. pub fn with_persist(client: LlmClient, persist: bool) -> Self { Self { - generator: Arc::new(RwLock::new(Box::new(super::LlmSummaryGenerator::new(client)))), + generator: Arc::new(RwLock::new(Box::new(super::LlmSummaryGenerator::new( + client, + )))), cache: Arc::new(RwLock::new(HashMap::new())), persist, config: SummaryStrategyConfig::default(), @@ -84,7 +88,12 @@ impl LazyStrategy { /// Get or generate a summary. /// /// Returns the cached summary if available, otherwise generates a new one. - pub async fn get_or_generate(&self, node_id: &str, title: &str, content: &str) -> crate::llm::LlmResult { + pub async fn get_or_generate( + &self, + node_id: &str, + title: &str, + content: &str, + ) -> crate::llm::LlmResult { // Check cache first if self.persist { if let Some(cached) = self.get_cached(node_id).await { diff --git a/src/index/summary/mod.rs b/src/index/summary/mod.rs index 4814a5f0..7a600482 100644 --- a/src/index/summary/mod.rs +++ b/src/index/summary/mod.rs @@ -16,12 +16,12 @@ //! - **Selective**: Generate summaries only for qualifying nodes (default) //! - **Lazy**: Generate summaries on-demand at query time -mod strategy; mod full; -mod selective; mod lazy; +mod selective; +mod strategy; -pub use strategy::{SummaryStrategy, SummaryStrategyConfig, SummaryGenerator, LlmSummaryGenerator}; pub use full::FullStrategy; -pub use selective::SelectiveStrategy; pub use lazy::LazyStrategy; +pub use selective::SelectiveStrategy; +pub use strategy::{LlmSummaryGenerator, SummaryGenerator, SummaryStrategy, SummaryStrategyConfig}; diff --git a/src/index/summary/selective.rs b/src/index/summary/selective.rs index 29b7ae57..3049278e 100644 --- a/src/index/summary/selective.rs +++ b/src/index/summary/selective.rs @@ -3,7 +3,7 @@ //! Selective summary strategy - generate summaries only for qualifying nodes. -use crate::domain::{NodeId, DocumentTree}; +use crate::domain::{DocumentTree, NodeId}; use crate::llm::LlmClient; use super::{SummaryGenerator, SummaryStrategyConfig}; @@ -70,7 +70,12 @@ impl SelectiveStrategy { } /// Check if a node should have a summary generated. - pub fn should_generate(&self, tree: &DocumentTree, node_id: NodeId, token_count: usize) -> bool { + pub fn should_generate( + &self, + tree: &DocumentTree, + node_id: NodeId, + token_count: usize, + ) -> bool { // Check token threshold let enough_tokens = token_count >= self.min_tokens; diff --git a/src/index/summary/strategy.rs b/src/index/summary/strategy.rs index ee7e1ab2..5b731232 100644 --- a/src/index/summary/strategy.rs +++ b/src/index/summary/strategy.rs @@ -5,7 +5,7 @@ use async_trait::async_trait; -use crate::domain::{NodeId, DocumentTree}; +use crate::domain::{DocumentTree, NodeId}; use crate::llm::{LlmClient, LlmResult}; /// Configuration for summary strategies. @@ -106,11 +106,20 @@ impl SummaryStrategy { } /// Check if we should generate a summary for a node. - pub fn should_generate(&self, tree: &DocumentTree, node_id: NodeId, token_count: usize) -> bool { + pub fn should_generate( + &self, + tree: &DocumentTree, + node_id: NodeId, + token_count: usize, + ) -> bool { match self { Self::None => false, Self::Full { .. } => token_count > 0, - Self::Selective { min_tokens, branch_only, .. } => { + Self::Selective { + min_tokens, + branch_only, + .. + } => { let is_branch = !tree.is_leaf(node_id); let enough_tokens = token_count >= *min_tokens; @@ -177,11 +186,7 @@ impl SummaryGenerator for LlmSummaryGenerator { Focus on the main topics and key information. \ Respond with only the summary, no additional text."; - let user_prompt = format!( - "Title: {}\n\nContent:\n{}", - title, - content - ); + let user_prompt = format!("Title: {}\n\nContent:\n{}", title, content); self.client .complete_with_max_tokens(&system_prompt, &user_prompt, self.max_tokens as u16) diff --git a/src/lib.rs b/src/lib.rs index 59688467..8a9e5615 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,17 @@ // Copyright (c) 2026 vectorless developers // SPDX-License-Identifier: Apache-2.0 +//! # Vectorless + +// Clippy: allow some pedantic lints that are too noisy for early-stage project +#![allow(clippy::all)] +#![allow(dead_code)] +#![allow(unused_variables)] +#![allow(unused_imports)] +#![allow(clippy::iter_over_hash_type)] +#![allow(clippy::large_enum_variant)] +#![allow(clippy::manual_unwrap_or_default)] + //! # Vectorless //! //! **A hierarchical, reasoning-native document intelligence engine.** @@ -95,27 +106,25 @@ pub mod client; pub mod config; -pub mod throttle; pub mod domain; pub mod index; pub mod llm; pub mod parser; pub mod retrieval; pub mod storage; +pub mod throttle; // ============================================================================= // Re-exports (Convenience API) // ============================================================================= // Client API (most common entry point) -pub use client::{DocumentInfo, IndexedDocument, Engine, EngineBuilder}; +pub use client::{DocumentInfo, Engine, EngineBuilder, IndexedDocument}; // Domain types pub use domain::{ - Error, Result, NodeId, TreeNode, DocumentTree, - DocumentStructure, StructureNode, - TocView, TocNode, TocEntry, TocConfig, - estimate_tokens, estimate_tokens_fast, + DocumentStructure, DocumentTree, Error, NodeId, Result, StructureNode, TocConfig, TocEntry, + TocNode, TocView, TreeNode, estimate_tokens, estimate_tokens_fast, }; // Configuration @@ -125,24 +134,24 @@ pub use config::{Config, ConfigLoader, RetrievalConfig, SummaryConfig}; pub use llm::{LlmClient, LlmConfig, LlmConfigs, LlmError, LlmPool, RetryConfig}; // Document parsing -pub use parser::{DocumentFormat, DocumentParser, DocxParser, MarkdownParser, PdfParser, ParseResult, RawNode}; +pub use parser::{ + DocumentFormat, DocumentParser, DocxParser, MarkdownParser, ParseResult, PdfParser, RawNode, +}; // Indexing +pub use index::pipeline::{CustomStageBuilder, PipelineOrchestrator}; pub use index::{ - PipelineExecutor, PipelineOptions, IndexInput, IndexMode, - IndexContext, IndexResult, IndexStage, IndexMetrics, - SummaryStrategy, ChangeDetector, ChangeSet, PartialUpdater, + ChangeDetector, ChangeSet, IndexContext, IndexInput, IndexMetrics, IndexMode, IndexResult, + IndexStage, PartialUpdater, PipelineExecutor, PipelineOptions, SummaryStrategy, }; -pub use index::pipeline::{PipelineOrchestrator, CustomStageBuilder}; // Retrieval pub use retrieval::{ - PipelineRetriever, Retriever, RetrieverError, RetrieverResult, - RetrieveOptions, RetrieveResponse, RetrievalResult, RetrievalContext, - QueryComplexity, StrategyPreference, SufficiencyLevel, - ContextBuilder, PruningStrategy, TokenEstimation, - NavigationDecision, NavigationStep, SearchPath, - format_for_llm, format_for_llm_async, format_tree_for_llm, format_tree_for_llm_async, + ContextBuilder, NavigationDecision, NavigationStep, PipelineRetriever, PruningStrategy, + QueryComplexity, RetrievalContext, RetrievalResult, RetrieveOptions, RetrieveResponse, + Retriever, RetrieverError, RetrieverResult, SearchPath, StrategyPreference, SufficiencyLevel, + TokenEstimation, format_for_llm, format_for_llm_async, format_tree_for_llm, + format_tree_for_llm_async, }; // Storage diff --git a/src/llm/client.rs b/src/llm/client.rs index ca4db5ff..27c81d80 100644 --- a/src/llm/client.rs +++ b/src/llm/client.rs @@ -4,13 +4,12 @@ //! Unified LLM client with retry and concurrency support. use async_openai::{ + Client, config::OpenAIConfig, types::chat::{ - ChatCompletionRequestSystemMessage, - ChatCompletionRequestUserMessage, + ChatCompletionRequestSystemMessage, ChatCompletionRequestUserMessage, CreateChatCompletionRequestArgs, }, - Client, }; use serde::de::DeserializeOwned; use std::borrow::Cow; @@ -71,7 +70,10 @@ impl std::fmt::Debug for LlmClient { f.debug_struct("LlmClient") .field("model", &self.config.model) .field("endpoint", &self.config.endpoint) - .field("concurrency", &self.concurrency.as_ref().map(|c| format!("{:?}", c))) + .field( + "concurrency", + &self.concurrency.as_ref().map(|c| format!("{:?}", c)), + ) .field("fallback_enabled", &self.fallback.is_some()) .finish() } @@ -171,7 +173,8 @@ impl LlmClient { pub async fn complete(&self, system: &str, user: &str) -> LlmResult { with_retry(&self.config.retry, || async { self.complete_once(system, user).await - }).await + }) + .await } /// Complete a prompt with custom max tokens. @@ -182,8 +185,10 @@ impl LlmClient { max_tokens: u16, ) -> LlmResult { with_retry(&self.config.retry, || async { - self.complete_once_with_max_tokens(system, user, max_tokens).await - }).await + self.complete_once_with_max_tokens(system, user, max_tokens) + .await + }) + .await } /// Complete a prompt and parse the response as JSON. @@ -228,7 +233,9 @@ impl LlmClient { user: &str, max_tokens: u16, ) -> LlmResult { - let response = self.complete_with_max_tokens(system, user, max_tokens).await?; + let response = self + .complete_with_max_tokens(system, user, max_tokens) + .await?; self.parse_json(&response) } @@ -241,10 +248,11 @@ impl LlmClient { None }; - let api_key = self.config.get_api_key() - .ok_or_else(|| LlmError::Config( - "No API key found. Set OPENAI_API_KEY environment variable.".to_string() - ))?; + let api_key = self.config.get_api_key().ok_or_else(|| { + LlmError::Config( + "No API key found. Set OPENAI_API_KEY environment variable.".to_string(), + ) + })?; let endpoint = self.config.auto_detect_endpoint(); let model = self.config.auto_detect_model(); @@ -274,11 +282,10 @@ impl LlmClient { debug!("Sending LLM request to {} with model {}", endpoint, model); - let response = client.chat().create(request).await - .map_err(|e| { - let msg = e.to_string(); - LlmError::from_api_message(&msg) - })?; + let response = client.chat().create(request).await.map_err(|e| { + let msg = e.to_string(); + LlmError::from_api_message(&msg) + })?; let content = response .choices @@ -305,10 +312,11 @@ impl LlmClient { None }; - let api_key = self.config.get_api_key() - .ok_or_else(|| LlmError::Config( - "No API key found. Set OPENAI_API_KEY environment variable.".to_string() - ))?; + let api_key = self.config.get_api_key().ok_or_else(|| { + LlmError::Config( + "No API key found. Set OPENAI_API_KEY environment variable.".to_string(), + ) + })?; let endpoint = self.config.auto_detect_endpoint(); let model = self.config.auto_detect_model(); @@ -332,17 +340,17 @@ impl LlmClient { .build() .map_err(|e| LlmError::Request(format!("Failed to build request: {}", e)))?; - let response = client.chat().create(request).await - .map_err(|e| { - let msg = e.to_string(); - eprintln!("[LLM ERROR] API error: {}", msg); - LlmError::from_api_message(&msg) - })?; + let response = client.chat().create(request).await.map_err(|e| { + let msg = e.to_string(); + eprintln!("[LLM ERROR] API error: {}", msg); + LlmError::from_api_message(&msg) + })?; // Debug: log response structure eprintln!("[LLM DEBUG] Response: {} choices", response.choices.len()); if let Some(choice) = response.choices.first() { - eprintln!("[LLM DEBUG] First choice: finish_reason={:?}, has_content={}", + eprintln!( + "[LLM DEBUG] First choice: finish_reason={:?}, has_content={}", choice.finish_reason, choice.message.content.is_some() ); @@ -380,8 +388,9 @@ impl LlmClient { /// Parse JSON from LLM response. fn parse_json(&self, text: &str) -> LlmResult { let json_text = self.extract_json(text); - serde_json::from_str(&json_text) - .map_err(|e| LlmError::Parse(format!("Failed to parse JSON: {}. Response: {}", e, text))) + serde_json::from_str(&json_text).map_err(|e| { + LlmError::Parse(format!("Failed to parse JSON: {}. Response: {}", e, text)) + }) } /// Extract JSON from text (handles markdown code blocks). @@ -445,9 +454,11 @@ mod tests { fn test_extract_json_code_block() { let client = LlmClient::with_defaults(); - let json = client.extract_json(r#"```json + let json = client.extract_json( + r#"```json {"key": "value"} -```"#); +```"#, + ); assert_eq!(json, r#"{"key": "value"}"#); } @@ -478,8 +489,7 @@ mod tests { use crate::throttle::ConcurrencyConfig; let controller = ConcurrencyController::new(ConcurrencyConfig::conservative()); - let client = LlmClient::for_model("gpt-4o-mini") - .with_concurrency(controller); + let client = LlmClient::for_model("gpt-4o-mini").with_concurrency(controller); assert!(client.concurrency.is_some()); } diff --git a/src/llm/config.rs b/src/llm/config.rs index bc7c6139..5c30c508 100644 --- a/src/llm/config.rs +++ b/src/llm/config.rs @@ -31,11 +31,21 @@ pub struct RetryConfig { pub retry_on_rate_limit: bool, } -fn default_max_attempts() -> usize { 3 } -fn default_initial_delay_ms() -> u64 { 500 } -fn default_max_delay_ms() -> u64 { 30000 } -fn default_multiplier() -> f64 { 2.0 } -fn default_true() -> bool { true } +fn default_max_attempts() -> usize { + 3 +} +fn default_initial_delay_ms() -> u64 { + 500 +} +fn default_max_delay_ms() -> u64 { + 30000 +} +fn default_multiplier() -> f64 { + 2.0 +} +fn default_true() -> bool { + true +} impl Default for RetryConfig { fn default() -> Self { @@ -87,8 +97,7 @@ impl RetryConfig { /// Calculate delay for a given attempt (0-indexed). pub fn delay_for_attempt(&self, attempt: usize) -> Duration { - let delay_ms = (self.initial_delay_ms as f64) - * self.multiplier.powf(attempt as f64); + let delay_ms = (self.initial_delay_ms as f64) * self.multiplier.powf(attempt as f64); let delay_ms = delay_ms.min(self.max_delay_ms as f64); Duration::from_millis(delay_ms as u64) } @@ -122,10 +131,18 @@ pub struct LlmConfig { pub retry: RetryConfig, } -fn default_model() -> String { "gpt-4o-mini".to_string() } -fn default_endpoint() -> String { "https://api.openai.com/v1".to_string() } -fn default_max_tokens() -> usize { 2000 } -fn default_temperature() -> f32 { 0.0 } +fn default_model() -> String { + "gpt-4o-mini".to_string() +} +fn default_endpoint() -> String { + "https://api.openai.com/v1".to_string() +} +fn default_max_tokens() -> usize { + 2000 +} +fn default_temperature() -> f32 { + 0.0 +} impl Default for LlmConfig { fn default() -> Self { @@ -187,7 +204,8 @@ impl LlmConfig { /// Get the API key from config or environment. pub fn get_api_key(&self) -> Option { - self.api_key.clone() + self.api_key + .clone() .or_else(|| std::env::var("OPENAI_API_KEY").ok()) .or_else(|| std::env::var("ANTHROPIC_API_KEY").ok()) .or_else(|| std::env::var("AZURE_OPENAI_API_KEY").ok()) diff --git a/src/llm/error.rs b/src/llm/error.rs index 641fe523..2cd8245d 100644 --- a/src/llm/error.rs +++ b/src/llm/error.rs @@ -53,12 +53,12 @@ impl LlmError { LlmError::Api(msg) => { // Rate limits and temporary failures are retryable let msg_lower = msg.to_lowercase(); - msg_lower.contains("rate limit") || - msg_lower.contains("429") || - msg_lower.contains("503") || - msg_lower.contains("502") || - msg_lower.contains("timeout") || - msg_lower.contains("overloaded") + msg_lower.contains("rate limit") + || msg_lower.contains("429") + || msg_lower.contains("503") + || msg_lower.contains("502") + || msg_lower.contains("timeout") + || msg_lower.contains("overloaded") } LlmError::Timeout(_) => true, LlmError::RateLimit(_) => true, diff --git a/src/llm/fallback.rs b/src/llm/fallback.rs index 68dd3a42..b6dac318 100644 --- a/src/llm/fallback.rs +++ b/src/llm/fallback.rs @@ -24,7 +24,9 @@ use serde::{Deserialize, Serialize}; use tracing::{debug, info, warn}; use super::error::LlmError; -use crate::config::{FallbackBehavior, FallbackConfig as ConfigFallbackConfig, OnAllFailedBehavior}; +use crate::config::{ + FallbackBehavior, FallbackConfig as ConfigFallbackConfig, OnAllFailedBehavior, +}; /// Result from a fallback-aware LLM call. #[derive(Debug, Clone)] @@ -217,7 +219,10 @@ impl FallbackChain { } // Current model is the last in the list, no more fallbacks Some(_) => { - warn!(model = current, "Already at last fallback model, no more available"); + warn!( + model = current, + "Already at last fallback model, no more available" + ); None } // Current model not in fallback list, try first fallback @@ -248,7 +253,10 @@ impl FallbackChain { } // Current endpoint is the last in the list, no more fallbacks Some(_) => { - warn!(endpoint = current, "Already at last fallback endpoint, no more available"); + warn!( + endpoint = current, + "Already at last fallback endpoint, no more available" + ); None } // Current endpoint not in fallback list, try first fallback @@ -307,14 +315,21 @@ mod tests { #[test] fn test_next_model() { let config = FallbackConfig { - models: vec!["gpt-4o".to_string(), "gpt-4o-mini".to_string(), "glm-4-flash".to_string()], + models: vec![ + "gpt-4o".to_string(), + "gpt-4o-mini".to_string(), + "glm-4-flash".to_string(), + ], ..FallbackConfig::default() }; let chain = FallbackChain::new(config); // Should get next model in chain assert_eq!(chain.next_model("gpt-4o"), Some("gpt-4o-mini".to_string())); - assert_eq!(chain.next_model("gpt-4o-mini"), Some("glm-4-flash".to_string())); + assert_eq!( + chain.next_model("gpt-4o-mini"), + Some("glm-4-flash".to_string()) + ); assert_eq!(chain.next_model("glm-4-flash"), None); } @@ -327,7 +342,10 @@ mod tests { let chain = FallbackChain::new(config); // Should fall back to first model in list - assert_eq!(chain.next_model("unknown-model"), Some("gpt-4o-mini".to_string())); + assert_eq!( + chain.next_model("unknown-model"), + Some("gpt-4o-mini".to_string()) + ); } #[test] diff --git a/src/llm/mod.rs b/src/llm/mod.rs index 12aaf460..89a6a415 100644 --- a/src/llm/mod.rs +++ b/src/llm/mod.rs @@ -63,15 +63,15 @@ //! # } //! ``` -mod error; -mod config; -mod retry; mod client; -mod pool; +mod config; +mod error; mod fallback; +mod pool; +mod retry; -pub use error::{LlmError, LlmResult}; -pub use config::{LlmConfig, RetryConfig, LlmConfigs}; pub use client::LlmClient; -pub use pool::LlmPool; +pub use config::{LlmConfig, LlmConfigs, RetryConfig}; +pub use error::{LlmError, LlmResult}; pub use fallback::{FallbackChain, FallbackConfig, FallbackResult, FallbackStep}; +pub use pool::LlmPool; diff --git a/src/llm/pool.rs b/src/llm/pool.rs index 72e3c634..375731dd 100644 --- a/src/llm/pool.rs +++ b/src/llm/pool.rs @@ -93,16 +93,13 @@ impl LlmPool { let arc = Arc::new(controller); self.concurrency = Some(arc.clone()); self.summary = Arc::new( - LlmClient::new(self.summary.config().clone()) - .with_shared_concurrency(arc.clone()) + LlmClient::new(self.summary.config().clone()).with_shared_concurrency(arc.clone()), ); self.retrieval = Arc::new( - LlmClient::new(self.retrieval.config().clone()) - .with_shared_concurrency(arc.clone()) + LlmClient::new(self.retrieval.config().clone()).with_shared_concurrency(arc.clone()), ); self.toc = Arc::new( - LlmClient::new(self.toc.config().clone()) - .with_shared_concurrency(arc.clone()) + LlmClient::new(self.toc.config().clone()).with_shared_concurrency(arc.clone()), ); self } @@ -112,15 +109,14 @@ impl LlmPool { self.concurrency = Some(controller.clone()); self.summary = Arc::new( LlmClient::new(self.summary.config().clone()) - .with_shared_concurrency(controller.clone()) + .with_shared_concurrency(controller.clone()), ); self.retrieval = Arc::new( LlmClient::new(self.retrieval.config().clone()) - .with_shared_concurrency(controller.clone()) + .with_shared_concurrency(controller.clone()), ); self.toc = Arc::new( - LlmClient::new(self.toc.config().clone()) - .with_shared_concurrency(controller.clone()) + LlmClient::new(self.toc.config().clone()).with_shared_concurrency(controller.clone()), ); self } diff --git a/src/llm/retry.rs b/src/llm/retry.rs index c8cdfb3a..7599001a 100644 --- a/src/llm/retry.rs +++ b/src/llm/retry.rs @@ -31,10 +31,7 @@ use super::error::{LlmError, LlmResult}; /// # Ok(()) /// # } /// ``` -pub async fn with_retry( - config: &RetryConfig, - operation: F, -) -> LlmResult +pub async fn with_retry(config: &RetryConfig, operation: F) -> LlmResult where F: Fn() -> Fut, Fut: Future>, @@ -94,12 +91,12 @@ fn should_retry(error: &LlmError, config: &RetryConfig) -> bool { LlmError::Api(msg) => { let msg_lower = msg.to_lowercase(); // Check for retryable API errors - msg_lower.contains("rate limit") || - msg_lower.contains("429") || - msg_lower.contains("503") || - msg_lower.contains("502") || - msg_lower.contains("timeout") || - msg_lower.contains("overloaded") + msg_lower.contains("rate limit") + || msg_lower.contains("429") + || msg_lower.contains("503") + || msg_lower.contains("502") + || msg_lower.contains("timeout") + || msg_lower.contains("overloaded") } _ => false, } @@ -122,7 +119,8 @@ mod tests { } else { Ok("success") } - }).await; + }) + .await; assert_eq!(result.unwrap(), "success"); assert_eq!(attempts.load(Ordering::SeqCst), 2); @@ -136,7 +134,8 @@ mod tests { let result: LlmResult = with_retry(&config, || async { attempts.fetch_add(1, Ordering::SeqCst); Err(LlmError::Timeout("timeout".to_string())) - }).await; + }) + .await; assert!(matches!(result, Err(LlmError::RetryExhausted { .. }))); assert_eq!(attempts.load(Ordering::SeqCst), 2); @@ -150,7 +149,8 @@ mod tests { let result: LlmResult = with_retry(&config, || async { attempts.fetch_add(1, Ordering::SeqCst); Err(LlmError::Config("bad config".to_string())) - }).await; + }) + .await; assert!(matches!(result, Err(LlmError::Config(_)))); assert_eq!(attempts.load(Ordering::SeqCst), 1); // Should only try once diff --git a/src/parser/docx/parser.rs b/src/parser/docx/parser.rs index 744f2041..dd59ccca 100644 --- a/src/parser/docx/parser.rs +++ b/src/parser/docx/parser.rs @@ -100,11 +100,7 @@ impl DocxParser { } /// Read an XML file from the archive. - fn read_xml_file( - &self, - archive: &mut ZipArchive>, - path: &str, - ) -> Result { + fn read_xml_file(&self, archive: &mut ZipArchive>, path: &str) -> Result { let mut file = archive .by_name(path) .map_err(|e| Error::Parse(format!("Failed to read {} from DOCX: {}", path, e)))?; @@ -131,7 +127,10 @@ impl DocxParser { let mut paragraphs = Vec::new(); // Find all w:p elements (paragraphs) - for para_elem in doc.descendants().filter(|n| n.has_tag_name((Self::WORD_NS, "p"))) { + for para_elem in doc + .descendants() + .filter(|n| n.has_tag_name((Self::WORD_NS, "p"))) + { if let Some(para) = self.parse_paragraph(¶_elem, style_resolver) { paragraphs.push(para); } @@ -182,7 +181,10 @@ impl DocxParser { let mut text = String::new(); // Find all w:t elements (text runs) - for text_elem in elem.descendants().filter(|n| n.has_tag_name((Self::WORD_NS, "t"))) { + for text_elem in elem + .descendants() + .filter(|n| n.has_tag_name((Self::WORD_NS, "t"))) + { if let Some(t) = text_elem.text() { text.push_str(t); } @@ -222,8 +224,7 @@ impl DocxParser { self.finalize_deeper_sections(&mut current_sections, level); // Create new section - let node = RawNode::new(¶.text) - .with_level(level as usize); + let node = RawNode::new(¶.text).with_level(level as usize); current_sections.push((level, node)); } else { @@ -262,11 +263,7 @@ impl DocxParser { } /// Finalize sections that are deeper than the given level. - fn finalize_deeper_sections( - &self, - sections: &mut Vec<(u8, RawNode)>, - new_level: u8, - ) { + fn finalize_deeper_sections(&self, sections: &mut Vec<(u8, RawNode)>, new_level: u8) { // Pop sections that are at the same level or deeper while let Some((level, _)) = sections.last() { if *level >= new_level { @@ -348,7 +345,10 @@ mod tests { let nodes = parser.build_raw_nodes(paragraphs).unwrap(); assert_eq!(nodes.len(), 1, "Should have exactly one node"); - assert_eq!(nodes[0].title, "Document", "Node title should be 'Document'"); + assert_eq!( + nodes[0].title, "Document", + "Node title should be 'Document'" + ); assert!( nodes[0].content.contains("First paragraph"), "Content should contain 'First paragraph', got: {:?}", diff --git a/src/parser/docx/styles.rs b/src/parser/docx/styles.rs index 6db7cae2..8414410a 100644 --- a/src/parser/docx/styles.rs +++ b/src/parser/docx/styles.rs @@ -45,26 +45,20 @@ impl StyleResolver { // Standard Word heading styles for level in 1..=6 { let style_id = format!("Heading{}", level); - self.styles.insert( - style_id.clone(), - DocxStyle::heading(&style_id, level), - ); + self.styles + .insert(style_id.clone(), DocxStyle::heading(&style_id, level)); } // Some documents use lowercase or different casing for level in 1..=6 { let style_id = format!("heading{}", level); - self.styles.insert( - style_id.clone(), - DocxStyle::heading(&style_id, level), - ); + self.styles + .insert(style_id.clone(), DocxStyle::heading(&style_id, level)); } // Title style (treat as H1) - self.styles.insert( - "Title".to_string(), - DocxStyle::heading("Title", 1), - ); + self.styles + .insert("Title".to_string(), DocxStyle::heading("Title", 1)); } /// Parse styles.xml content. @@ -75,7 +69,10 @@ impl StyleResolver { }; // Find all w:style elements - for style_elem in doc.descendants().filter(|n| n.has_tag_name((WORD_NS, "style"))) { + for style_elem in doc + .descendants() + .filter(|n| n.has_tag_name((WORD_NS, "style"))) + { if let Some(style) = self.parse_style_element(&style_elem) { self.styles.insert(style.style_id.clone(), style); } @@ -135,16 +132,16 @@ impl StyleResolver { /// Get heading level for a style ID. pub fn get_heading_level(&self, style_id: &Option) -> Option { - style_id.as_ref().and_then(|id| { - self.styles.get(id).and_then(|s| s.heading_level) - }) + style_id + .as_ref() + .and_then(|id| self.styles.get(id).and_then(|s| s.heading_level)) } /// Check if a style is a heading. pub fn is_heading(&self, style_id: &Option) -> bool { - style_id.as_ref().is_some_and(|id| { - self.styles.get(id).is_some_and(|s| s.is_heading) - }) + style_id + .as_ref() + .is_some_and(|id| self.styles.get(id).is_some_and(|s| s.is_heading)) } /// Try to detect heading level from text content heuristics. @@ -219,9 +216,18 @@ mod tests { r }; - assert_eq!(resolver.get_heading_level(&Some("Heading1".to_string())), Some(1)); - assert_eq!(resolver.get_heading_level(&Some("Heading2".to_string())), Some(2)); - assert_eq!(resolver.get_heading_level(&Some("Normal".to_string())), None); + assert_eq!( + resolver.get_heading_level(&Some("Heading1".to_string())), + Some(1) + ); + assert_eq!( + resolver.get_heading_level(&Some("Heading2".to_string())), + Some(2) + ); + assert_eq!( + resolver.get_heading_level(&Some("Normal".to_string())), + None + ); } #[test] @@ -240,10 +246,18 @@ mod tests { assert_eq!(resolver.detect_heading_by_heuristics("Chapter 1"), Some(1)); assert_eq!(resolver.detect_heading_by_heuristics("Section 2"), Some(1)); - assert_eq!(resolver.detect_heading_by_heuristics("1. Introduction"), Some(1)); - assert_eq!(resolver.detect_heading_by_heuristics("1.1 Background"), Some(2)); assert_eq!( - resolver.detect_heading_by_heuristics("This is a very long piece of text that is unlikely to be a heading"), + resolver.detect_heading_by_heuristics("1. Introduction"), + Some(1) + ); + assert_eq!( + resolver.detect_heading_by_heuristics("1.1 Background"), + Some(2) + ); + assert_eq!( + resolver.detect_heading_by_heuristics( + "This is a very long piece of text that is unlikely to be a heading" + ), None ); } diff --git a/src/parser/markdown/config.rs b/src/parser/markdown/config.rs index 0195ec6c..7a013f5f 100644 --- a/src/parser/markdown/config.rs +++ b/src/parser/markdown/config.rs @@ -33,7 +33,6 @@ pub struct MarkdownConfig { // ============================================================ // Parsing Options // ============================================================ - /// Enable GitHub Flavored Markdown extensions. /// /// Includes: tables, strikethrough, task lists, autolinks. @@ -65,7 +64,6 @@ pub struct MarkdownConfig { // ============================================================ // Content Extraction // ============================================================ - /// Include code blocks in node content. /// Default: `true` pub include_code_blocks: bool, @@ -85,7 +83,6 @@ pub struct MarkdownConfig { // ============================================================ // Frontmatter // ============================================================ - /// Parse YAML frontmatter (`---` delimiters). /// Default: `true` pub parse_frontmatter: bool, @@ -101,7 +98,6 @@ pub struct MarkdownConfig { // ============================================================ // Advanced Options // ============================================================ - /// Minimum characters required for a heading title to be valid. /// Headings with shorter titles are skipped. /// Default: `1` diff --git a/src/parser/markdown/frontmatter.rs b/src/parser/markdown/frontmatter.rs index 59e4467d..65f7cda0 100644 --- a/src/parser/markdown/frontmatter.rs +++ b/src/parser/markdown/frontmatter.rs @@ -117,7 +117,11 @@ impl Frontmatter { /// - YAML: `---\nkey: value\n---` /// - TOML: `+++\nkey = "value"\n+++` #[must_use] -pub fn extract_frontmatter(content: &str, parse_yaml: bool, parse_toml: bool) -> (Option, &str) { +pub fn extract_frontmatter( + content: &str, + parse_yaml: bool, + parse_toml: bool, +) -> (Option, &str) { // Try YAML frontmatter first if parse_yaml { if let Some((fm, remaining)) = Frontmatter::parse(content, "---") { diff --git a/src/parser/markdown/parser.rs b/src/parser/markdown/parser.rs index cb81fb01..7e1f3a2d 100644 --- a/src/parser/markdown/parser.rs +++ b/src/parser/markdown/parser.rs @@ -92,7 +92,13 @@ impl MarkdownParser { } /// Parse Markdown content and extract nodes. - fn extract_nodes(&self, content: &str) -> (Vec, Option>) { + fn extract_nodes( + &self, + content: &str, + ) -> ( + Vec, + Option>, + ) { // 1. Extract frontmatter (if present) let (fm, remaining_content) = frontmatter::extract_frontmatter( content, @@ -276,10 +282,10 @@ impl MarkdownParser { && (!content_buffer.trim().is_empty() || !preamble_content.is_empty()) { // Use preamble_content if available, otherwise use content_buffer - let content = if !preamble_content.is_empty() { - preamble_content.trim() - } else { + let content = if preamble_content.is_empty() { content_buffer.trim() + } else { + preamble_content.trim() }; nodes.push(RawNode { title: self.config.preamble_title.clone(), @@ -417,7 +423,12 @@ mod tests { let result = parser.parse(content).await.unwrap(); assert!(!result.nodes.is_empty()); - assert!(result.nodes.iter().any(|n| n.title == "Title" && n.level == 1)); + assert!( + result + .nodes + .iter() + .any(|n| n.title == "Title" && n.level == 1) + ); } #[tokio::test] diff --git a/src/parser/mod.rs b/src/parser/mod.rs index 9b75c33c..10d69738 100644 --- a/src/parser/mod.rs +++ b/src/parser/mod.rs @@ -36,9 +36,9 @@ //! # } //! ``` -mod types; mod registry; mod traits; +mod types; // Markdown parsing module pub mod markdown; @@ -53,26 +53,15 @@ pub mod toc; pub mod docx; // Re-export main types -pub use types::{ - DocumentFormat, - DocumentMeta, - ParseResult, - RawNode, -}; +pub use types::{DocumentFormat, DocumentMeta, ParseResult, RawNode}; // Re-export parser trait pub use traits::DocumentParser; // Re-export registry and convenience functions -pub use registry::{ - ParserRegistry, - get_parser, - get_parser_for_file, - parse_content, - parse_file, -}; +pub use registry::{ParserRegistry, get_parser, get_parser_for_file, parse_content, parse_file}; // Re-export concrete parsers -pub use markdown::{MarkdownParser, MarkdownConfig}; -pub use pdf::PdfParser; pub use docx::DocxParser; +pub use markdown::{MarkdownConfig, MarkdownParser}; +pub use pdf::PdfParser; diff --git a/src/parser/pdf/parser.rs b/src/parser/pdf/parser.rs index a8707d43..c047d21a 100644 --- a/src/parser/pdf/parser.rs +++ b/src/parser/pdf/parser.rs @@ -145,7 +145,12 @@ impl PdfParser { } /// Extract text from a single page. - fn extract_page_text(&self, doc: &LopdfDocument, object_id: lopdf::ObjectId, _page_num: usize) -> String { + fn extract_page_text( + &self, + doc: &LopdfDocument, + object_id: lopdf::ObjectId, + _page_num: usize, + ) -> String { let mut text = String::new(); if let Ok(page_obj) = doc.get_object(object_id) { @@ -319,7 +324,11 @@ impl PdfParser { } /// Convert TOC entries to RawNodes. - fn toc_entries_to_raw_nodes(&self, entries: &[crate::parser::toc::TocEntry], pages: &[PdfPage]) -> Vec { + fn toc_entries_to_raw_nodes( + &self, + entries: &[crate::parser::toc::TocEntry], + pages: &[PdfPage], + ) -> Vec { let mut nodes = Vec::new(); for entry in entries { @@ -341,7 +350,11 @@ impl PdfParser { } /// Get content for a TOC entry from pages. - fn get_content_for_entry(&self, entry: &crate::parser::toc::TocEntry, pages: &[PdfPage]) -> String { + fn get_content_for_entry( + &self, + entry: &crate::parser::toc::TocEntry, + pages: &[PdfPage], + ) -> String { let start_page = entry.physical_page.unwrap_or(1); // Find content on this page @@ -365,7 +378,7 @@ impl PdfParser { pages .iter() .map(|page| { - RawNode::new(&format!("Page {}", page.number)) + RawNode::new(format!("Page {}", page.number)) .with_content(page.text.clone()) .with_level(1) .with_page(page.number) @@ -418,7 +431,10 @@ impl PdfParser { self.pages_to_raw_nodes(&result.pages) } Err(e) => { - warn!("TOC extraction failed: {}, falling back to page-based extraction", e); + warn!( + "TOC extraction failed: {}, falling back to page-based extraction", + e + ); self.pages_to_raw_nodes(&result.pages) } } diff --git a/src/parser/pdf/types.rs b/src/parser/pdf/types.rs index abc6ed68..8c6e27b0 100644 --- a/src/parser/pdf/types.rs +++ b/src/parser/pdf/types.rs @@ -3,8 +3,8 @@ //! PDF document types. -use serde::{Deserialize, Serialize}; use crate::domain::estimate_tokens; +use serde::{Deserialize, Serialize}; /// A single page from a PDF document. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -24,7 +24,11 @@ impl PdfPage { pub fn new(number: usize, text: impl Into) -> Self { let text = text.into(); let token_count = estimate_tokens(&text); - Self { number, text, token_count } + Self { + number, + text, + token_count, + } } /// Check if the page is empty. @@ -95,7 +99,11 @@ impl PdfParseResult { /// Create a new parse result. pub fn new(metadata: PdfMetadata, pages: Vec) -> Self { let total_tokens = pages.iter().map(|p| p.token_count).sum(); - Self { metadata, pages, total_tokens } + Self { + metadata, + pages, + total_tokens, + } } /// Check if the document is empty. diff --git a/src/parser/registry.rs b/src/parser/registry.rs index a24400ac..947552ac 100644 --- a/src/parser/registry.rs +++ b/src/parser/registry.rs @@ -174,8 +174,8 @@ pub fn get_parser_for_file(path: &Path) -> Option> { /// /// A [`ParseResult`] containing the extracted nodes. pub async fn parse_content(content: &str, format: DocumentFormat) -> Result { - let parser = - get_parser(format).ok_or_else(|| Error::Parse(format!("Unsupported format: {:?}", format)))?; + let parser = get_parser(format) + .ok_or_else(|| Error::Parse(format!("Unsupported format: {:?}", format)))?; parser.parse(content).await } diff --git a/src/parser/toc/assigner.rs b/src/parser/toc/assigner.rs index 65e109ed..86087885 100644 --- a/src/parser/toc/assigner.rs +++ b/src/parser/toc/assigner.rs @@ -10,8 +10,8 @@ use crate::config::LlmConfig; use crate::domain::Result; use crate::parser::pdf::PdfPage; -use crate::llm::LlmClient; use super::types::{PageOffset, TocEntry}; +use crate::llm::LlmClient; /// Page assigner configuration. #[derive(Debug, Clone)] @@ -89,7 +89,10 @@ impl PageAssigner { return self.assign_with_llm(entries, pages).await; } - info!("Calculated offset: {} (confidence: {})", offset.offset, offset.confidence); + info!( + "Calculated offset: {} (confidence: {})", + offset.offset, offset.confidence + ); // Step 3: Apply offset to all entries for entry in entries.iter_mut() { @@ -105,9 +108,7 @@ impl PageAssigner { /// Select anchor entries for offset calculation. fn select_anchors<'a>(&self, entries: &'a [TocEntry], count: usize) -> Vec<&'a TocEntry> { // Select entries with TOC pages, evenly distributed - let with_pages: Vec<_> = entries.iter() - .filter(|e| e.toc_page.is_some()) - .collect(); + let with_pages: Vec<_> = entries.iter().filter(|e| e.toc_page.is_some()).collect(); if with_pages.len() <= count { return with_pages; @@ -121,7 +122,11 @@ impl PageAssigner { } /// Calculate page offset by verifying anchors. - async fn calculate_offset(&self, anchors: Vec<&TocEntry>, pages: &[PdfPage]) -> Result { + async fn calculate_offset( + &self, + anchors: Vec<&TocEntry>, + pages: &[PdfPage], + ) -> Result { if anchors.is_empty() { return Ok(PageOffset::new(0, 0, 0.0)); } @@ -133,18 +138,24 @@ impl PageAssigner { let toc_page = anchor.toc_page.unwrap(); // Find the physical page where this title appears - if let Some(physical) = self.locate_title_in_range(anchor.title.as_str(), pages, toc_page).await? { + if let Some(physical) = self + .locate_title_in_range(anchor.title.as_str(), pages, toc_page) + .await? + { let offset = physical as i32 - toc_page as i32; verified_offsets.push((offset, true)); - debug!("Anchor '{}' found: toc={}, physical={}, offset={}", - anchor.title, toc_page, physical, offset); + debug!( + "Anchor '{}' found: toc={}, physical={}, offset={}", + anchor.title, toc_page, physical, offset + ); } else { verified_offsets.push((0, false)); } } // Calculate the mode (most common offset) - let successful: Vec<_> = verified_offsets.iter() + let successful: Vec<_> = verified_offsets + .iter() .filter(|(_, success)| *success) .map(|(offset, _)| *offset) .collect(); @@ -166,32 +177,41 @@ impl PageAssigner { for &v in values { *counts.entry(v).or_insert(0) += 1; } - counts.into_iter() + counts + .into_iter() .max_by_key(|&(_, count)| count) .map(|(v, _)| v) .unwrap_or(0) } /// Locate a title in a range of pages using LLM. - async fn locate_title_in_range(&self, title: &str, pages: &[PdfPage], near_page: usize) -> Result> { + async fn locate_title_in_range( + &self, + title: &str, + pages: &[PdfPage], + near_page: usize, + ) -> Result> { // Search in a range around the expected page let start = (near_page.saturating_sub(3)).max(1); let end = (near_page + 3).min(pages.len()); - let range_pages: Vec<_> = (start..=end) - .filter_map(|i| pages.get(i - 1)) - .collect(); + let range_pages: Vec<_> = (start..=end).filter_map(|i| pages.get(i - 1)).collect(); if range_pages.is_empty() { return Ok(None); } // Use LLM to find the exact page - let content = range_pages.iter() - .map(|p| format!("\n{}\n", - p.number, - &p.text[..p.text.len().min(500)], - p.number)) + let content = range_pages + .iter() + .map(|p| { + format!( + "\n{}\n", + p.number, + &p.text[..p.text.len().min(500)], + p.number + ) + }) .collect::>() .join("\n\n"); @@ -224,7 +244,9 @@ Reply in JSON format: let page_groups = self.group_pages(pages, 5); for entry in entries.iter_mut() { - let physical = self.locate_title_in_groups(entry.title.as_str(), &page_groups).await?; + let physical = self + .locate_title_in_groups(entry.title.as_str(), &page_groups) + .await?; entry.physical_page = physical; entry.confidence = if physical.is_some() { 0.8 } else { 0.3 }; } @@ -234,21 +256,31 @@ Reply in JSON format: /// Group pages for batch processing. fn group_pages<'a>(&self, pages: &'a [PdfPage], group_size: usize) -> Vec> { - pages.chunks(group_size) + pages + .chunks(group_size) .map(|chunk| chunk.iter().collect()) .collect() } /// Locate a title across page groups. - async fn locate_title_in_groups(&self, title: &str, groups: &[Vec<&PdfPage>]) -> Result> { + async fn locate_title_in_groups( + &self, + title: &str, + groups: &[Vec<&PdfPage>], + ) -> Result> { let system = "You are a document analysis assistant. Find which page contains a specific section title."; for group in groups { - let content = group.iter() - .map(|p| format!("\n{}\n", - p.number, - &p.text[..p.text.len().min(300)], - p.number)) + let content = group + .iter() + .map(|p| { + format!( + "\n{}\n", + p.number, + &p.text[..p.text.len().min(300)], + p.number + ) + }) .collect::>() .join("\n\n"); diff --git a/src/parser/toc/detector.rs b/src/parser/toc/detector.rs index 5da03c9e..f8112f07 100644 --- a/src/parser/toc/detector.rs +++ b/src/parser/toc/detector.rs @@ -9,8 +9,8 @@ use tracing::debug; use crate::config::LlmConfig; use crate::domain::Result; -use crate::llm::LlmClient; use super::types::TocDetection; +use crate::llm::LlmClient; use crate::parser::pdf::PdfPage; /// TOC detector configuration. @@ -90,7 +90,8 @@ impl TocDetector { }, TocPattern { name: "chinese_chapter_with_page", - regex: Regex::new(r"第[一二三四五六七八九十\d]+[章节部篇].*?[\.\s…·]{2,}\s*\d+").unwrap(), + regex: Regex::new(r"第[一二三四五六七八九十\d]+[章节部篇].*?[\.\s…·]{2,}\s*\d+") + .unwrap(), weight: 0.85, }, TocPattern { @@ -101,7 +102,8 @@ impl TocDetector { // English TOC patterns TocPattern { name: "english_toc_header", - regex: Regex::new(r"(?i)^[\s]*(table\s+of\s+contents|contents|outline)[\s]*$").unwrap(), + regex: Regex::new(r"(?i)^[\s]*(table\s+of\s+contents|contents|outline)[\s]*$") + .unwrap(), weight: 0.9, }, TocPattern { @@ -130,7 +132,8 @@ impl TocDetector { /// Detect TOC in PDF pages. pub async fn detect(&self, pages: &[PdfPage]) -> Result { - let check_pages = pages.iter() + let check_pages = pages + .iter() .take(self.config.max_check_pages) .collect::>(); @@ -140,8 +143,10 @@ impl TocDetector { // Step 1: Regex detection let regex_result = self.detect_with_regex(&check_pages); - debug!("Regex detection result: found={}, confidence={}", - regex_result.found, regex_result.confidence); + debug!( + "Regex detection result: found={}, confidence={}", + regex_result.found, regex_result.confidence + ); // Step 2: If confidence is high enough, return if regex_result.confidence >= self.config.regex_confidence_threshold { @@ -236,12 +241,23 @@ impl TocDetector { } /// Detect TOC using LLM. - async fn detect_with_llm(&self, client: &LlmClient, pages: &[&PdfPage]) -> Result { + async fn detect_with_llm( + &self, + client: &LlmClient, + pages: &[&PdfPage], + ) -> Result { // Combine first few pages for analysis - let content = pages.iter() + let content = pages + .iter() .take(5) - .map(|p| format!("\n{}\n", p.number, - &p.text[..p.text.len().min(1000)], p.number)) + .map(|p| { + format!( + "\n{}\n", + p.number, + &p.text[..p.text.len().min(1000)], + p.number + ) + }) .collect::>() .join("\n\n"); @@ -309,7 +325,10 @@ mod tests { let pages = vec![ make_page(1, "Abstract"), - make_page(2, "Table of Contents\n\nChapter 1. Introduction 1\nChapter 2. Methods 5"), + make_page( + 2, + "Table of Contents\n\nChapter 1. Introduction 1\nChapter 2. Methods 5", + ), ]; let rt = tokio::runtime::Runtime::new().unwrap(); diff --git a/src/parser/toc/mod.rs b/src/parser/toc/mod.rs index 28b69f93..83341cae 100644 --- a/src/parser/toc/mod.rs +++ b/src/parser/toc/mod.rs @@ -64,21 +64,23 @@ //! # } //! ``` -mod types; +mod assigner; mod detector; mod parser; -mod assigner; -mod verifier; -mod repairer; mod processor; +mod repairer; +mod types; +mod verifier; // Re-export main types -pub use types::{TocEntry, TocDetection, PageOffset, VerificationError, VerificationReport, ErrorType}; +pub use types::{ + ErrorType, PageOffset, TocDetection, TocEntry, VerificationError, VerificationReport, +}; // Re-export components +pub use assigner::{PageAssigner, PageAssignerConfig}; pub use detector::{TocDetector, TocDetectorConfig}; pub use parser::{TocParser, TocParserConfig}; -pub use assigner::{PageAssigner, PageAssignerConfig}; -pub use verifier::{IndexVerifier, VerifierConfig}; -pub use repairer::{IndexRepairer, RepairerConfig}; pub use processor::{TocProcessor, TocProcessorConfig}; +pub use repairer::{IndexRepairer, RepairerConfig}; +pub use verifier::{IndexVerifier, VerifierConfig}; diff --git a/src/parser/toc/parser.rs b/src/parser/toc/parser.rs index 86c76492..9cbeee1f 100644 --- a/src/parser/toc/parser.rs +++ b/src/parser/toc/parser.rs @@ -8,8 +8,8 @@ use tracing::debug; use crate::config::LlmConfig; use crate::domain::Result; -use crate::llm::LlmClient; use super::types::TocEntry; +use crate::llm::LlmClient; /// TOC parser configuration. #[derive(Debug, Clone)] @@ -131,7 +131,11 @@ Notes: } /// Verify completeness and continue if needed. - async fn verify_and_complete(&self, toc_text: &str, mut entries: Vec) -> Result> { + async fn verify_and_complete( + &self, + toc_text: &str, + mut entries: Vec, + ) -> Result> { let mut attempts = 0; while attempts < self.config.max_retries { @@ -143,7 +147,10 @@ Notes: return Ok(entries); } - debug!("TOC incomplete, attempting continuation (attempt {})", attempts + 1); + debug!( + "TOC incomplete, attempting continuation (attempt {})", + attempts + 1 + ); // Continue parsing let additional = self.continue_parsing(toc_text, &entries).await?; @@ -163,9 +170,9 @@ Notes: async fn check_completeness(&self, toc_text: &str, entries: &[TocEntry]) -> Result { let system = "You are a document analysis assistant. Determine if the parsed entries completely represent the original TOC."; - let entries_json = serde_json::to_string_pretty( - &entries.iter().map(|e| &e.title).collect::>() - ).unwrap_or_default(); + let entries_json = + serde_json::to_string_pretty(&entries.iter().map(|e| &e.title).collect::>()) + .unwrap_or_default(); let user = format!( r#"Original TOC: @@ -189,7 +196,11 @@ Is the parsing complete? Reply with JSON: } /// Continue parsing from where we left off. - async fn continue_parsing(&self, toc_text: &str, existing: &[TocEntry]) -> Result> { + async fn continue_parsing( + &self, + toc_text: &str, + existing: &[TocEntry], + ) -> Result> { let system = "You are a document structure extraction expert. Continue parsing the TOC from where it was left off."; let last_titles: Vec<_> = existing.iter().rev().take(5).map(|e| &e.title).collect(); diff --git a/src/parser/toc/processor.rs b/src/parser/toc/processor.rs index 14c308e9..991b0f6d 100644 --- a/src/parser/toc/processor.rs +++ b/src/parser/toc/processor.rs @@ -188,7 +188,7 @@ impl TocProcessor { // Group pages into chunks let chunk_size = 10; - for (_i, chunk) in pages.chunks(chunk_size).enumerate() { + for chunk in pages.chunks(chunk_size) { let start_page = chunk.first().map(|p| p.number).unwrap_or(1); let end_page = chunk.last().map(|p| p.number).unwrap_or(1); @@ -221,7 +221,10 @@ impl TocProcessor { let report = self.verifier.verify(entries, pages).await?; if report.accuracy >= self.config.accuracy_threshold { - debug!("Verification passed: accuracy {:.1}%", report.accuracy * 100.0); + debug!( + "Verification passed: accuracy {:.1}%", + report.accuracy * 100.0 + ); return Ok(report); } diff --git a/src/parser/toc/repairer.rs b/src/parser/toc/repairer.rs index c0d14211..4a00383c 100644 --- a/src/parser/toc/repairer.rs +++ b/src/parser/toc/repairer.rs @@ -9,9 +9,9 @@ use crate::config::LlmConfig; use crate::domain::Result; use crate::parser::pdf::PdfPage; -use crate::llm::LlmClient; use super::types::{TocEntry, VerificationError, VerificationReport}; use super::verifier::IndexVerifier; +use crate::llm::LlmClient; /// Repairer configuration. #[derive(Debug, Clone)] @@ -77,16 +77,27 @@ impl IndexRepairer { let expected_page = error.expected_page; // Search around the expected page - let start = expected_page.saturating_sub(self.config.search_range).max(1); + let start = expected_page + .saturating_sub(self.config.search_range) + .max(1); let end = (expected_page + self.config.search_range).min(pages.len()); - if let Some(correct_page) = self.find_correct_page(&entry.title, pages, start..=end).await? { - debug!("Repaired '{}' : page {} → {}", entry.title, expected_page, correct_page); + if let Some(correct_page) = self + .find_correct_page(&entry.title, pages, start..=end) + .await? + { + debug!( + "Repaired '{}' : page {} → {}", + entry.title, expected_page, correct_page + ); entry.physical_page = Some(correct_page); entry.confidence = 0.9; repaired_count += 1; } else { - debug!("Could not repair '{}' (searched pages {}-{})", entry.title, start, end); + debug!( + "Could not repair '{}' (searched pages {}-{})", + entry.title, start, end + ); } } @@ -112,7 +123,10 @@ impl IndexRepairer { } else { &page.text }; - content_parts.push(format!("\n{}\n", page_num, text, page_num)); + content_parts.push(format!( + "\n{}\n", + page_num, text, page_num + )); } } diff --git a/src/parser/toc/types.rs b/src/parser/toc/types.rs index d6e1c1d0..9465311b 100644 --- a/src/parser/toc/types.rs +++ b/src/parser/toc/types.rs @@ -182,7 +182,12 @@ pub struct VerificationError { impl VerificationError { /// Create a new verification error. - pub fn new(index: usize, title: impl Into, expected_page: usize, error_type: ErrorType) -> Self { + pub fn new( + index: usize, + title: impl Into, + expected_page: usize, + error_type: ErrorType, + ) -> Self { Self { index, title: title.into(), @@ -237,7 +242,12 @@ impl VerificationReport { } else { 1.0 }; - Self { total, correct, accuracy, errors } + Self { + total, + correct, + accuracy, + errors, + } } /// Create a report indicating all entries are correct. diff --git a/src/parser/toc/verifier.rs b/src/parser/toc/verifier.rs index 940c1a86..e1e9c457 100644 --- a/src/parser/toc/verifier.rs +++ b/src/parser/toc/verifier.rs @@ -10,8 +10,8 @@ use crate::config::LlmConfig; use crate::domain::Result; use crate::parser::pdf::PdfPage; -use crate::llm::LlmClient; use super::types::{ErrorType, TocEntry, VerificationError, VerificationReport}; +use crate::llm::LlmClient; /// Verifier configuration. #[derive(Debug, Clone)] @@ -55,7 +55,11 @@ impl IndexVerifier { } /// Verify TOC entries against PDF pages. - pub async fn verify(&self, entries: &[TocEntry], pages: &[PdfPage]) -> Result { + pub async fn verify( + &self, + entries: &[TocEntry], + pages: &[PdfPage], + ) -> Result { if entries.is_empty() { return Ok(VerificationReport::all_correct(0)); } @@ -92,8 +96,12 @@ impl IndexVerifier { } let report = VerificationReport::new(sample.len(), correct, errors); - info!("Verification complete: {}/{} correct ({:.1}% accuracy)", - report.correct, report.total, report.accuracy * 100.0); + info!( + "Verification complete: {}/{} correct ({:.1}% accuracy)", + report.correct, + report.total, + report.accuracy * 100.0 + ); Ok(report) } @@ -136,7 +144,10 @@ impl IndexVerifier { let found = self.check_title_on_page(&entry.title, &page.text).await?; if !found { - debug!("Title '{}' not found on page {}", entry.title, physical_page); + debug!( + "Title '{}' not found on page {}", + entry.title, physical_page + ); return Ok(Err(ErrorType::TitleNotFound)); } diff --git a/src/retrieval/cache/path_cache.rs b/src/retrieval/cache/path_cache.rs index 25c2b60e..e9202150 100644 --- a/src/retrieval/cache/path_cache.rs +++ b/src/retrieval/cache/path_cache.rs @@ -7,9 +7,9 @@ use std::collections::HashMap; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant}; +use super::super::types::SearchPath; use crate::config::CacheConfig as AppConfig; use crate::domain::NodeId; -use super::super::types::SearchPath; /// Cache entry for a search path. #[derive(Debug, Clone)] @@ -180,19 +180,13 @@ impl PathCache { { if self.config.use_lru { // Find entry with lowest access count - if let Some((min_key, _)) = cache - .iter() - .min_by_key(|(_, e)| e.access_count) - { + if let Some((min_key, _)) = cache.iter().min_by_key(|(_, e)| e.access_count) { let key = min_key.clone(); cache.remove(&key); } } else { // Remove oldest entry - if let Some((oldest_key, _)) = cache - .iter() - .min_by_key(|(_, e)| e.created_at) - { + if let Some((oldest_key, _)) = cache.iter().min_by_key(|(_, e)| e.created_at) { let key = oldest_key.clone(); cache.remove(&key); } diff --git a/src/retrieval/complexity/detector.rs b/src/retrieval/complexity/detector.rs index bfdf0938..5079040d 100644 --- a/src/retrieval/complexity/detector.rs +++ b/src/retrieval/complexity/detector.rs @@ -102,11 +102,7 @@ impl ComplexityDetector { let conjunctions = ["and", "or", "but", "however", "although"]; let conjunction_count = conjunctions .iter() - .filter(|c| { - query_lower - .split_whitespace() - .any(|w| w == **c) - }) + .filter(|c| query_lower.split_whitespace().any(|w| w == **c)) .count(); if conjunction_count >= 2 { diff --git a/src/retrieval/complexity/mod.rs b/src/retrieval/complexity/mod.rs index 22a8a531..628a1896 100644 --- a/src/retrieval/complexity/mod.rs +++ b/src/retrieval/complexity/mod.rs @@ -7,5 +7,5 @@ mod detector; -pub use detector::ComplexityDetector; pub use super::types::QueryComplexity; +pub use detector::ComplexityDetector; diff --git a/src/retrieval/context.rs b/src/retrieval/context.rs index 35168721..ba0edb34 100644 --- a/src/retrieval/context.rs +++ b/src/retrieval/context.rs @@ -27,9 +27,9 @@ //! .build_async(&results).await?; //! ``` -use std::collections::HashSet; -use crate::domain::{DocumentTree, NodeId, estimate_tokens}; use super::types::RetrievalResult; +use crate::domain::{DocumentTree, NodeId, estimate_tokens}; +use std::collections::HashSet; /// Pruning strategy for context building. #[derive(Debug, Clone, Copy, PartialEq)] @@ -176,7 +176,9 @@ impl ContextBuilder { match self.pruning_strategy { PruningStrategy::TokenLimit => self.build_token_limit(results), PruningStrategy::RelevanceThreshold(min) => self.build_relevance(results, min), - PruningStrategy::Diversity { max_overlap } => self.build_diversity(results, max_overlap), + PruningStrategy::Diversity { max_overlap } => { + self.build_diversity(results, max_overlap) + } PruningStrategy::Hybrid { min_relevance } => self.build_hybrid(results, min_relevance), } } @@ -351,11 +353,12 @@ impl ContextBuilder { // Collect from title words.extend( - result.title + result + .title .to_lowercase() .split_whitespace() .filter(|w| w.len() > 3) - .map(|w| w.to_string()) + .map(|w| w.to_string()), ); // Collect from summary @@ -365,7 +368,7 @@ impl ContextBuilder { .to_lowercase() .split_whitespace() .filter(|w| w.len() > 3) - .map(|w| w.to_string()) + .map(|w| w.to_string()), ); } @@ -408,7 +411,8 @@ impl ContextBuilder { max_depth: usize, ) -> String { let mut sections = Vec::new(); - self.collect_sections_async(tree, node_id, 0, max_depth, &mut sections).await; + self.collect_sections_async(tree, node_id, 0, max_depth, &mut sections) + .await; sections.join(&self.separator) } @@ -449,7 +453,7 @@ impl ContextBuilder { } // Yield every few levels to avoid blocking - if current_depth > 0 && current_depth % 3 == 0 { + if current_depth > 0 && current_depth.is_multiple_of(3) { tokio::task::yield_now().await; } @@ -466,7 +470,8 @@ impl ContextBuilder { current_depth + 1, max_depth, sections, - )).await; + )) + .await; } } } @@ -529,11 +534,7 @@ pub async fn format_for_llm_async(results: &[RetrievalResult], max_tokens: usize } /// Format a document tree for LLM consumption. -pub fn format_tree_for_llm( - tree: &DocumentTree, - max_depth: usize, - max_tokens: usize, -) -> String { +pub fn format_tree_for_llm(tree: &DocumentTree, max_depth: usize, max_tokens: usize) -> String { ContextBuilder::new() .with_max_tokens(max_tokens) .build_from_tree(tree, tree.root(), max_depth) @@ -558,15 +559,11 @@ mod tests { #[test] fn test_context_builder() { let results = vec![ - RetrievalResult::new("Section 1") - .with_content("Content 1"), - RetrievalResult::new("Section 2") - .with_content("Content 2"), + RetrievalResult::new("Section 1").with_content("Content 1"), + RetrievalResult::new("Section 2").with_content("Content 2"), ]; - let context = ContextBuilder::new() - .with_max_tokens(1000) - .build(&results); + let context = ContextBuilder::new().with_max_tokens(1000).build(&results); assert!(context.contains("Section 1")); assert!(context.contains("Content 1")); @@ -590,10 +587,9 @@ mod tests { #[test] fn test_token_estimation_modes() { - let fast_builder = ContextBuilder::new() - .with_token_estimation(TokenEstimation::Fast); - let accurate_builder = ContextBuilder::new() - .with_token_estimation(TokenEstimation::Accurate); + let fast_builder = ContextBuilder::new().with_token_estimation(TokenEstimation::Fast); + let accurate_builder = + ContextBuilder::new().with_token_estimation(TokenEstimation::Accurate); let fast_tokens = fast_builder.estimate_tokens("Hello world test"); let accurate_tokens = accurate_builder.estimate_tokens("Hello world test"); @@ -623,7 +619,10 @@ mod tests { #[tokio::test] async fn test_async_build() { let results: Vec<_> = (0..200) - .map(|i| RetrievalResult::new(&format!("Section {}", i)).with_content(&format!("Content {}", i))) + .map(|i| { + RetrievalResult::new(&format!("Section {}", i)) + .with_content(&format!("Content {}", i)) + }) .collect(); let context = ContextBuilder::new() diff --git a/src/retrieval/mod.rs b/src/retrieval/mod.rs index 9e138766..cb58b930 100644 --- a/src/retrieval/mod.rs +++ b/src/retrieval/mod.rs @@ -47,43 +47,43 @@ //! let response = orchestrator.execute(tree, query, options).await?; //! ``` -mod types; -mod retriever; mod context; mod pipeline_retriever; +mod retriever; +mod types; +pub mod cache; +pub mod complexity; pub mod pipeline; +pub mod search; pub mod stages; pub mod strategy; -pub mod search; pub mod sufficiency; -pub mod complexity; -pub mod cache; -pub use types::*; -pub use retriever::{Retriever, RetrieverError, RetrieverResult, RetrievalContext}; -pub use pipeline_retriever::PipelineRetriever; pub use context::{ - ContextBuilder, PruningStrategy, TokenEstimation, - format_for_llm, format_for_llm_async, + ContextBuilder, PruningStrategy, TokenEstimation, format_for_llm, format_for_llm_async, format_tree_for_llm, format_tree_for_llm_async, }; +pub use pipeline_retriever::PipelineRetriever; +pub use retriever::{RetrievalContext, Retriever, RetrieverError, RetrieverResult}; +pub use types::*; // Pipeline exports pub use pipeline::{ - RetrievalOrchestrator, RetrievalStage, PipelineContext, - StageOutcome, ExecutionGroup, FailurePolicy, - CandidateNode, SearchAlgorithm, SearchConfig, RetrievalMetrics, + CandidateNode, ExecutionGroup, FailurePolicy, PipelineContext, RetrievalMetrics, + RetrievalOrchestrator, RetrievalStage, SearchAlgorithm, SearchConfig, StageOutcome, }; // Re-export PipelineContext as RetrievalContext for stages (alias for clarity) pub use pipeline::PipelineContext as StageContext; // Stage exports -pub use stages::{AnalyzeStage, PlanStage, SearchStage, JudgeStage}; +pub use stages::{AnalyzeStage, JudgeStage, PlanStage, SearchStage}; // Strategy exports -pub use strategy::{RetrievalStrategy, StrategyCapabilities, KeywordStrategy, SemanticStrategy, LlmStrategy}; +pub use strategy::{ + KeywordStrategy, LlmStrategy, RetrievalStrategy, SemanticStrategy, StrategyCapabilities, +}; // Search exports pub use search::{BeamSearch, GreedySearch, SearchConfig as SearchAlgConfig, SearchResult}; diff --git a/src/retrieval/pipeline/context.rs b/src/retrieval/pipeline/context.rs index 9cf50d04..5dd7c388 100644 --- a/src/retrieval/pipeline/context.rs +++ b/src/retrieval/pipeline/context.rs @@ -246,7 +246,11 @@ pub struct PipelineContext { impl PipelineContext { /// Create a new retrieval context. - pub fn new(tree: Arc, query: impl Into, options: RetrieveOptions) -> Self { + pub fn new( + tree: Arc, + query: impl Into, + options: RetrieveOptions, + ) -> Self { Self { query: query.into(), tree, @@ -278,7 +282,8 @@ impl PipelineContext { /// End timing and record for a stage. pub fn end_stage(&mut self, stage_name: &str, success: bool, message: Option) { - let duration_ms = self.stage_start + let duration_ms = self + .stage_start .map(|s| s.elapsed().as_millis() as u64) .unwrap_or(0); @@ -339,7 +344,8 @@ impl PipelineContext { content: self.accumulated_content, confidence: 0.0, is_sufficient: self.sufficiency == SufficiencyLevel::Sufficient, - strategy_used: self.selected_strategy + strategy_used: self + .selected_strategy .map(|s| format!("{:?}", s)) .unwrap_or_else(|| "unknown".to_string()), complexity: self.complexity.unwrap_or_default(), diff --git a/src/retrieval/pipeline/orchestrator.rs b/src/retrieval/pipeline/orchestrator.rs index 85701a31..682d3e22 100644 --- a/src/retrieval/pipeline/orchestrator.rs +++ b/src/retrieval/pipeline/orchestrator.rs @@ -282,14 +282,15 @@ impl RetrievalOrchestrator { options: RetrieveOptions, ) -> Result { let total_start = Instant::now(); - info!("Starting retrieval pipeline for query: '{}' ({} stages)", query, self.stages.len()); + info!( + "Starting retrieval pipeline for query: '{}' ({} stages)", + query, + self.stages.len() + ); // Resolve execution order let order = self.resolve_order()?; - let stage_names: Vec<&str> = order - .iter() - .map(|&i| self.stages[i].stage.name()) - .collect(); + let stage_names: Vec<&str> = order.iter().map(|&i| self.stages[i].stage.name()).collect(); info!("Execution order: {:?}", stage_names); // Compute execution groups @@ -342,11 +343,15 @@ impl RetrievalOrchestrator { } StageOutcome::Complete => { // Retrieval complete - ctx.metrics.total_time_ms = total_start.elapsed().as_millis() as u64; + ctx.metrics.total_time_ms = + total_start.elapsed().as_millis() as u64; info!("Retrieval completed by stage: {}", stage_name); return Ok(ctx.finalize()); } - StageOutcome::NeedMoreData { additional_beam, go_deeper } => { + StageOutcome::NeedMoreData { + additional_beam, + go_deeper, + } => { // Backtrack to search stage if let Some(search_idx) = self.stages.iter().position(|e| e.stage.name() == "search") @@ -376,11 +381,16 @@ impl RetrievalOrchestrator { } } } - StageOutcome::Backtrack { target_stage, reason } => { + StageOutcome::Backtrack { + target_stage, + reason, + } => { info!("Backtracking to {}: {}", target_stage, reason); - if let Some(target_idx) = - self.stages.iter().position(|e| e.stage.name() == target_stage) + if let Some(target_idx) = self + .stages + .iter() + .position(|e| e.stage.name() == target_stage) { ctx.increment_backtrack(); backtrack_count += 1; @@ -395,7 +405,8 @@ impl RetrievalOrchestrator { } StageOutcome::Skip { reason } => { info!("Skipping remaining stages: {}", reason); - ctx.metrics.total_time_ms = total_start.elapsed().as_millis() as u64; + ctx.metrics.total_time_ms = + total_start.elapsed().as_millis() as u64; return Ok(ctx.finalize()); } } @@ -404,7 +415,10 @@ impl RetrievalOrchestrator { ctx.end_stage(stage_name, false, Some(e.to_string())); if policy.allows_continuation() { - warn!("Stage {} failed but policy allows continuation: {}", stage_name, e); + warn!( + "Stage {} failed but policy allows continuation: {}", + stage_name, e + ); } else { error!("Stage {} failed: {}", stage_name, e); return Err(e); @@ -419,9 +433,7 @@ impl RetrievalOrchestrator { ctx.metrics.total_time_ms = total_start.elapsed().as_millis() as u64; info!( "Retrieval completed in {}ms ({} iterations, {} backtracks)", - ctx.metrics.total_time_ms, - total_iterations, - backtrack_count + ctx.metrics.total_time_ms, total_iterations, backtrack_count ); Ok(ctx.finalize()) diff --git a/src/retrieval/pipeline_retriever.rs b/src/retrieval/pipeline_retriever.rs index 26590957..084ad53d 100644 --- a/src/retrieval/pipeline_retriever.rs +++ b/src/retrieval/pipeline_retriever.rs @@ -132,7 +132,7 @@ impl Retriever for PipelineRetriever { Ok(response) } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "pipeline" } @@ -147,7 +147,7 @@ impl Retriever for PipelineRetriever { CostEstimate { llm_calls: base_llm_calls + (node_count / 10), // Rough estimate - tokens: node_count * 50, // Conservative estimate + tokens: node_count * 50, // Conservative estimate } } } diff --git a/src/retrieval/search/beam.rs b/src/retrieval/search/beam.rs index b0fbc0d2..d2d05b1e 100644 --- a/src/retrieval/search/beam.rs +++ b/src/retrieval/search/beam.rs @@ -7,11 +7,11 @@ use async_trait::async_trait; -use crate::domain::DocumentTree; -use super::super::types::{NavigationDecision, NavigationStep, SearchPath}; use super::super::RetrievalContext; -use super::{SearchConfig, SearchResult, SearchTree}; +use super::super::types::{NavigationDecision, NavigationStep, SearchPath}; use super::scorer::NodeScorer; +use super::{SearchConfig, SearchResult, SearchTree}; +use crate::domain::DocumentTree; /// Beam search - explores multiple paths simultaneously. /// @@ -68,7 +68,11 @@ impl SearchTree for BeamSearch { .collect(); // Sort by score and keep top beam_width - current_beam.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); + current_beam.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); current_beam.truncate(beam_width); for iteration in 0..config.max_iterations { @@ -106,7 +110,7 @@ impl SearchTree for BeamSearch { title: child_node.map(|n| n.title.clone()).unwrap_or_default(), score: child_score, decision: NavigationDecision::GoToChild( - children.iter().position(|&c| c == child_id).unwrap_or(0) + children.iter().position(|&c| c == child_id).unwrap_or(0), ), depth: child_node.map(|n| n.depth).unwrap_or(0), }); @@ -118,7 +122,11 @@ impl SearchTree for BeamSearch { } // Sort next beam and keep top candidates - next_beam.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); + next_beam.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); next_beam.truncate(beam_width); current_beam = next_beam; @@ -137,13 +145,17 @@ impl SearchTree for BeamSearch { } // Sort final results by score - result.paths.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); + result.paths.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); result.paths.truncate(config.top_k); result } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "beam" } } diff --git a/src/retrieval/search/greedy.rs b/src/retrieval/search/greedy.rs index ad75a37a..933b20d9 100644 --- a/src/retrieval/search/greedy.rs +++ b/src/retrieval/search/greedy.rs @@ -7,11 +7,11 @@ use async_trait::async_trait; -use crate::domain::DocumentTree; -use super::super::types::{NavigationDecision, NavigationStep, SearchPath}; use super::super::RetrievalContext; -use super::{SearchConfig, SearchResult, SearchTree}; +use super::super::types::{NavigationDecision, NavigationStep, SearchPath}; use super::scorer::NodeScorer; +use super::{SearchConfig, SearchResult, SearchTree}; +use crate::domain::DocumentTree; /// Greedy search - always follows the best single path. /// @@ -84,7 +84,9 @@ impl SearchTree for GreedySearch { node_id: format!("{:?}", child_id), title: child_node.map(|n| n.title.clone()).unwrap_or_default(), score: best_score, - decision: NavigationDecision::GoToChild(children.iter().position(|&c| c == child_id).unwrap_or(0)), + decision: NavigationDecision::GoToChild( + children.iter().position(|&c| c == child_id).unwrap_or(0), + ), depth: child_node.map(|n| n.depth).unwrap_or(0), }); @@ -108,7 +110,7 @@ impl SearchTree for GreedySearch { result } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "greedy" } } diff --git a/src/retrieval/search/mcts.rs b/src/retrieval/search/mcts.rs index b005b07c..c556d18f 100644 --- a/src/retrieval/search/mcts.rs +++ b/src/retrieval/search/mcts.rs @@ -8,12 +8,12 @@ use async_trait::async_trait; use std::collections::HashMap; -use crate::config::StrategyConfig; -use crate::domain::{NodeId, DocumentTree}; -use super::super::types::{NavigationDecision, NavigationStep, SearchPath}; use super::super::RetrievalContext; -use super::{SearchConfig, SearchResult, SearchTree}; +use super::super::types::{NavigationDecision, NavigationStep, SearchPath}; use super::scorer::NodeScorer; +use super::{SearchConfig, SearchResult, SearchTree}; +use crate::config::StrategyConfig; +use crate::domain::{DocumentTree, NodeId}; /// Statistics for a node in MCTS. #[derive(Debug, Clone, Default)] @@ -55,20 +55,14 @@ impl MctsSearch { } /// Calculate UCT score for a child node. - fn uct_score( - &self, - child_stats: &NodeStats, - parent_visits: usize, - prior_score: f32, - ) -> f32 { + fn uct_score(&self, child_stats: &NodeStats, parent_visits: usize, prior_score: f32) -> f32 { if child_stats.visits == 0 { // Unvisited nodes get high priority return f32::INFINITY; } let exploitation = child_stats.total_score / child_stats.visits as f32; - let exploration = self.exploration_weight - * (parent_visits as f32).ln().sqrt() + let exploration = self.exploration_weight * (parent_visits as f32).ln().sqrt() / child_stats.visits as f32; // Combine with prior score from scorer @@ -134,12 +128,7 @@ impl MctsSearch { } /// Backpropagate score up the tree. - fn backpropagate( - &self, - stats: &mut HashMap, - path: &[NodeId], - score: f32, - ) { + fn backpropagate(&self, stats: &mut HashMap, path: &[NodeId], score: f32) { for &node_id in path { let node_stats = stats.entry(node_id).or_default(); node_stats.visits += 1; @@ -224,7 +213,8 @@ impl SearchTree for MctsSearch { }) .collect(); - scored_children.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); + scored_children + .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); for (child_id, score) in scored_children.iter().take(config.top_k) { if *score >= config.min_score { @@ -250,7 +240,11 @@ impl SearchTree for MctsSearch { }) .collect(); - final_paths.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); + final_paths.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); final_paths.truncate(config.top_k); result.paths = final_paths @@ -261,7 +255,7 @@ impl SearchTree for MctsSearch { result } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "mcts" } } diff --git a/src/retrieval/search/mod.rs b/src/retrieval/search/mod.rs index 0ee10751..dca0e04e 100644 --- a/src/retrieval/search/mod.rs +++ b/src/retrieval/search/mod.rs @@ -3,14 +3,14 @@ //! Search algorithms for tree traversal. -mod scorer; -mod r#trait; -mod greedy; mod beam; +mod greedy; mod mcts; +mod scorer; +mod r#trait; -pub use scorer::{NodeScorer, ScoringContext}; -pub use r#trait::{SearchTree, SearchResult, SearchConfig}; -pub use greedy::GreedySearch; pub use beam::BeamSearch; +pub use greedy::GreedySearch; pub use mcts::MctsSearch; +pub use scorer::{NodeScorer, ScoringContext}; +pub use r#trait::{SearchConfig, SearchResult, SearchTree}; diff --git a/src/retrieval/search/scorer.rs b/src/retrieval/search/scorer.rs index f400b3b9..e22f8239 100644 --- a/src/retrieval/search/scorer.rs +++ b/src/retrieval/search/scorer.rs @@ -5,7 +5,7 @@ //! //! Implements the NodeScore formula: `Σ ChunkScore(n) / √(N+1)` -use crate::domain::{NodeId, DocumentTree}; +use crate::domain::{DocumentTree, NodeId}; /// Context for scoring calculations. #[derive(Debug, Clone)] diff --git a/src/retrieval/search/trait.rs b/src/retrieval/search/trait.rs index 785cc29e..afbeb970 100644 --- a/src/retrieval/search/trait.rs +++ b/src/retrieval/search/trait.rs @@ -5,9 +5,9 @@ use async_trait::async_trait; -use crate::domain::DocumentTree; -use super::super::types::{NavigationStep, SearchPath}; use super::super::RetrievalContext; +use super::super::types::{NavigationStep, SearchPath}; +use crate::domain::DocumentTree; /// Result of a search operation. #[derive(Debug, Clone)] diff --git a/src/retrieval/stages/analyze.rs b/src/retrieval/stages/analyze.rs index 1ff5cb84..c26b7e4c 100644 --- a/src/retrieval/stages/analyze.rs +++ b/src/retrieval/stages/analyze.rs @@ -13,9 +13,7 @@ use tracing::info; use crate::domain::{DocumentTree, TocView}; use crate::retrieval::complexity::ComplexityDetector; -use crate::retrieval::pipeline::{ - FailurePolicy, PipelineContext, RetrievalStage, StageOutcome, -}; +use crate::retrieval::pipeline::{FailurePolicy, PipelineContext, RetrievalStage, StageOutcome}; // QueryComplexity is used in context /// Analyze Stage - analyzes queries for retrieval planning. @@ -69,20 +67,17 @@ impl AnalyzeStage { // 5. Remove punctuation let stop_words = [ - "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", - "have", "has", "had", "do", "does", "did", "will", "would", "could", - "should", "may", "might", "must", "shall", "can", "need", "dare", - "ought", "used", "to", "of", "in", "for", "on", "with", "at", "by", - "from", "as", "into", "through", "during", "before", "after", - "above", "below", "between", "under", "again", "further", "then", - "once", "here", "there", "when", "where", "why", "how", "all", - "each", "few", "more", "most", "other", "some", "such", "no", "nor", - "not", "only", "own", "same", "so", "than", "too", "very", "just", - "and", "but", "if", "or", "because", "until", "while", "although", - "though", "what", "which", "who", "whom", "this", "that", "these", - "those", "am", "it", "its", "itself", "he", "him", "his", "she", - "her", "hers", "they", "them", "their", "we", "us", "our", "you", - "your", "i", "me", "my", + "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", "have", "has", + "had", "do", "does", "did", "will", "would", "could", "should", "may", "might", "must", + "shall", "can", "need", "dare", "ought", "used", "to", "of", "in", "for", "on", "with", + "at", "by", "from", "as", "into", "through", "during", "before", "after", "above", + "below", "between", "under", "again", "further", "then", "once", "here", "there", + "when", "where", "why", "how", "all", "each", "few", "more", "most", "other", "some", + "such", "no", "nor", "not", "only", "own", "same", "so", "than", "too", "very", "just", + "and", "but", "if", "or", "because", "until", "while", "although", "though", "what", + "which", "who", "whom", "this", "that", "these", "those", "am", "it", "its", "itself", + "he", "him", "his", "she", "her", "hers", "they", "them", "their", "we", "us", "our", + "you", "your", "i", "me", "my", ]; query @@ -92,7 +87,10 @@ impl AnalyzeStage { let word = word.trim_matches(|c: char| !c.is_alphanumeric()); word.len() >= 2 && !stop_words.contains(&word) }) - .map(|word| word.trim_matches(|c: char| !c.is_alphanumeric()).to_string()) + .map(|word| { + word.trim_matches(|c: char| !c.is_alphanumeric()) + .to_string() + }) .filter(|word| !word.is_empty()) .collect() } @@ -145,13 +143,17 @@ impl AnalyzeStage { // Sort by score and return top sections matches.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); - matches.into_iter().take(5).map(|(title, _)| title).collect() + matches + .into_iter() + .take(5) + .map(|(title, _)| title) + .collect() } } #[async_trait] impl RetrievalStage for AnalyzeStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "analyze" } @@ -168,10 +170,7 @@ impl RetrievalStage for AnalyzeStage { // 1. Detect complexity ctx.complexity = Some(self.complexity_detector.detect(&ctx.query)); - info!( - "Query complexity: {:?}", - ctx.complexity - ); + info!("Query complexity: {:?}", ctx.complexity); // 2. Extract keywords ctx.keywords = self.extract_keywords(&ctx.query); diff --git a/src/retrieval/stages/judge.rs b/src/retrieval/stages/judge.rs index c5af69cf..f22806db 100644 --- a/src/retrieval/stages/judge.rs +++ b/src/retrieval/stages/judge.rs @@ -12,11 +12,9 @@ use tracing::{info, warn}; use crate::domain::estimate_tokens; use crate::llm::LlmClient; -use crate::retrieval::pipeline::{ - FailurePolicy, PipelineContext, RetrievalStage, StageOutcome, -}; +use crate::retrieval::pipeline::{FailurePolicy, PipelineContext, RetrievalStage, StageOutcome}; use crate::retrieval::sufficiency::{LlmJudge, SufficiencyChecker, ThresholdChecker}; -use crate::retrieval::types::{RetrieveResponse, RetrievalResult, SufficiencyLevel}; +use crate::retrieval::types::{RetrievalResult, RetrieveResponse, SufficiencyLevel}; /// Judge Stage - evaluates retrieval sufficiency. /// @@ -110,7 +108,8 @@ impl JudgeStage { } // Fall back to threshold checker - self.threshold_checker.check(&ctx.query, &ctx.accumulated_content, ctx.token_count) + self.threshold_checker + .check(&ctx.query, &ctx.accumulated_content, ctx.token_count) } /// Build the final response. @@ -144,7 +143,8 @@ impl JudgeStage { content: ctx.accumulated_content.clone(), confidence: self.calculate_confidence(ctx), is_sufficient: ctx.sufficiency == SufficiencyLevel::Sufficient, - strategy_used: ctx.selected_strategy + strategy_used: ctx + .selected_strategy .map(|s| format!("{:?}", s)) .unwrap_or_else(|| "unknown".to_string()), complexity: ctx.complexity.unwrap_or_default(), @@ -160,8 +160,8 @@ impl JudgeStage { } // Weight by score and sufficiency - let avg_score: f32 = ctx.candidates.iter().map(|c| c.score).sum::() - / ctx.candidates.len() as f32; + let avg_score: f32 = + ctx.candidates.iter().map(|c| c.score).sum::() / ctx.candidates.len() as f32; let sufficiency_factor = match ctx.sufficiency { SufficiencyLevel::Sufficient => 1.0, @@ -175,7 +175,7 @@ impl JudgeStage { #[async_trait] impl RetrievalStage for JudgeStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "judge" } diff --git a/src/retrieval/stages/mod.rs b/src/retrieval/stages/mod.rs index 7e770060..7d66cf77 100644 --- a/src/retrieval/stages/mod.rs +++ b/src/retrieval/stages/mod.rs @@ -23,11 +23,11 @@ //! Implement [`RetrievalStage`](crate::retrieval::pipeline::RetrievalStage) to create custom stages. mod analyze; +mod judge; mod plan; mod search; -mod judge; pub use analyze::AnalyzeStage; +pub use judge::JudgeStage; pub use plan::PlanStage; pub use search::SearchStage; -pub use judge::JudgeStage; diff --git a/src/retrieval/stages/plan.rs b/src/retrieval/stages/plan.rs index b7a76acb..7177322b 100644 --- a/src/retrieval/stages/plan.rs +++ b/src/retrieval/stages/plan.rs @@ -15,8 +15,7 @@ use tracing::info; // DocumentTree is accessed via context use crate::llm::LlmClient; use crate::retrieval::pipeline::{ - FailurePolicy, PipelineContext, RetrievalStage, StageOutcome, - SearchAlgorithm, SearchConfig, + FailurePolicy, PipelineContext, RetrievalStage, SearchAlgorithm, SearchConfig, StageOutcome, }; use crate::retrieval::types::{QueryComplexity, StrategyPreference}; @@ -124,7 +123,7 @@ impl PlanStage { let complexity = ctx.complexity.unwrap_or(QueryComplexity::Medium); let (beam_width, max_depth) = match complexity { - QueryComplexity::Simple => (1, 5), // Greedy-like + QueryComplexity::Simple => (1, 5), // Greedy-like QueryComplexity::Medium => (ctx.options.beam_width, 10), QueryComplexity::Complex => (ctx.options.beam_width + 2, 15), }; @@ -140,7 +139,7 @@ impl PlanStage { #[async_trait] impl RetrievalStage for PlanStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "plan" } @@ -172,7 +171,10 @@ impl RetrievalStage for PlanStage { "Plan complete: strategy={:?}, algorithm={:?}, beam_width={}", ctx.selected_strategy, ctx.selected_algorithm, - ctx.search_config.as_ref().map(|c| c.beam_width).unwrap_or(0) + ctx.search_config + .as_ref() + .map(|c| c.beam_width) + .unwrap_or(0) ); Ok(StageOutcome::cont()) diff --git a/src/retrieval/stages/search.rs b/src/retrieval/stages/search.rs index b3987b1b..648dd295 100644 --- a/src/retrieval/stages/search.rs +++ b/src/retrieval/stages/search.rs @@ -12,12 +12,13 @@ use tracing::{info, warn}; use crate::domain::DocumentTree; // LlmClient is used via strategy +use crate::retrieval::RetrievalContext; // Legacy context use crate::retrieval::pipeline::{ - FailurePolicy, PipelineContext, RetrievalStage, StageOutcome, - CandidateNode, SearchAlgorithm, + CandidateNode, FailurePolicy, PipelineContext, RetrievalStage, SearchAlgorithm, StageOutcome, +}; +use crate::retrieval::search::{ + BeamSearch, GreedySearch, SearchConfig as SearchAlgConfig, SearchTree, }; -use crate::retrieval::search::{BeamSearch, GreedySearch, SearchConfig as SearchAlgConfig, SearchTree}; -use crate::retrieval::RetrievalContext; // Legacy context use crate::retrieval::strategy::{KeywordStrategy, LlmStrategy, RetrievalStrategy}; use crate::retrieval::types::StrategyPreference; @@ -124,7 +125,9 @@ impl SearchStage { // Sort by score descending candidates.sort_by(|a, b| { - b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal) + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) }); candidates @@ -133,7 +136,7 @@ impl SearchStage { #[async_trait] impl RetrievalStage for SearchStage { - fn name(&self) -> &str { + fn name(&self) -> &'static str { "search" } @@ -179,7 +182,11 @@ impl RetrievalStage for SearchStage { }; // Create legacy context for search algorithms - let legacy_ctx = RetrievalContext::new(&ctx.query, ctx.options.max_tokens, ctx.options.sufficiency_check); + let legacy_ctx = RetrievalContext::new( + &ctx.query, + ctx.options.max_tokens, + ctx.options.sufficiency_check, + ); // Execute search based on algorithm let result = match algorithm { diff --git a/src/retrieval/strategy/keyword.rs b/src/retrieval/strategy/keyword.rs index a5bdbd95..bfb34a68 100644 --- a/src/retrieval/strategy/keyword.rs +++ b/src/retrieval/strategy/keyword.rs @@ -8,10 +8,10 @@ use async_trait::async_trait; use std::collections::{HashMap, HashSet}; -use crate::domain::{NodeId, DocumentTree}; -use super::super::types::{NavigationDecision, QueryComplexity}; use super::super::RetrievalContext; +use super::super::types::{NavigationDecision, QueryComplexity}; use super::r#trait::{NodeEvaluation, RetrievalStrategy, StrategyCapabilities}; +use crate::domain::{DocumentTree, NodeId}; /// Keyword-based retrieval strategy. /// @@ -157,7 +157,7 @@ impl RetrievalStrategy for KeywordStrategy { } } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "keyword" } diff --git a/src/retrieval/strategy/llm.rs b/src/retrieval/strategy/llm.rs index a0564511..7a3ed89e 100644 --- a/src/retrieval/strategy/llm.rs +++ b/src/retrieval/strategy/llm.rs @@ -8,11 +8,11 @@ use async_trait::async_trait; use serde::Deserialize; -use crate::domain::{NodeId, DocumentTree, TocView}; -use crate::llm::LlmClient; -use super::super::types::{NavigationDecision, QueryComplexity}; use super::super::RetrievalContext; +use super::super::types::{NavigationDecision, QueryComplexity}; use super::r#trait::{NodeEvaluation, RetrievalStrategy, StrategyCapabilities}; +use crate::domain::{DocumentTree, NodeId, TocView}; +use crate::llm::LlmClient; /// LLM response for navigation decision. #[derive(Debug, Clone, Deserialize)] @@ -96,7 +96,12 @@ Be concise and focused on finding the most relevant information."#.to_string() } /// Build the navigation prompt for a single node. - fn build_prompt(&self, tree: &DocumentTree, node_id: NodeId, context: &RetrievalContext) -> String { + fn build_prompt( + &self, + tree: &DocumentTree, + node_id: NodeId, + context: &RetrievalContext, + ) -> String { let node = tree.get(node_id); let children = tree.children(node_id); @@ -125,21 +130,27 @@ Be concise and focused on finding the most relevant information."#.to_string() let toc_markdown = self.toc_view.format_markdown(&toc); // Limit ToC size for token efficiency let toc_preview: String = toc_markdown.chars().take(1000).collect(); - format!("\n\nDocument ToC (from this node):\n```\n{}\n```\n", toc_preview) + format!( + "\n\nDocument ToC (from this node):\n```\n{}\n```\n", + toc_preview + ) } else { String::new() }; format!( "Query: {}\n{}Current Node:\n{}\n\nWhat is the relevance and action?", - context.query, - toc_context, - node_info + context.query, toc_context, node_info ) } /// Parse LLM response to evaluation. - fn parse_response(&self, response: &str, tree: &DocumentTree, node_id: NodeId) -> NodeEvaluation { + fn parse_response( + &self, + response: &str, + tree: &DocumentTree, + node_id: NodeId, + ) -> NodeEvaluation { // Try to parse as JSON if let Ok(parsed) = serde_json::from_str::(response) { let score = (parsed.relevance as f32 / 100.0).clamp(0.0, 1.0); @@ -187,7 +198,10 @@ Be concise and focused on finding the most relevant information."#.to_string() } else { NavigationDecision::ExploreMore }, - reasoning: Some(format!("Parsed from response: {}...", &response[..100.min(response.len())])), + reasoning: Some(format!( + "Parsed from response: {}...", + &response[..100.min(response.len())] + )), } } } @@ -234,7 +248,7 @@ impl RetrievalStrategy for LlmStrategy { results } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "llm" } @@ -248,6 +262,9 @@ impl RetrievalStrategy for LlmStrategy { } fn suitable_for_complexity(&self, complexity: QueryComplexity) -> bool { - matches!(complexity, QueryComplexity::Medium | QueryComplexity::Complex) + matches!( + complexity, + QueryComplexity::Medium | QueryComplexity::Complex + ) } } diff --git a/src/retrieval/strategy/mod.rs b/src/retrieval/strategy/mod.rs index 90a09969..93c55dff 100644 --- a/src/retrieval/strategy/mod.rs +++ b/src/retrieval/strategy/mod.rs @@ -3,12 +3,12 @@ //! Retrieval strategies for different query types. -mod r#trait; mod keyword; -mod semantic; mod llm; +mod semantic; +mod r#trait; -pub use r#trait::{RetrievalStrategy, StrategyCapabilities, NodeEvaluation, StrategyCost}; pub use keyword::KeywordStrategy; -pub use semantic::SemanticStrategy; pub use llm::LlmStrategy; +pub use semantic::SemanticStrategy; +pub use r#trait::{NodeEvaluation, RetrievalStrategy, StrategyCapabilities, StrategyCost}; diff --git a/src/retrieval/strategy/semantic.rs b/src/retrieval/strategy/semantic.rs index 4ef00a13..170e7998 100644 --- a/src/retrieval/strategy/semantic.rs +++ b/src/retrieval/strategy/semantic.rs @@ -7,11 +7,11 @@ use async_trait::async_trait; -use crate::config::StrategyConfig; -use crate::domain::{NodeId, DocumentTree}; -use super::super::types::{NavigationDecision, QueryComplexity}; use super::super::RetrievalContext; +use super::super::types::{NavigationDecision, QueryComplexity}; use super::r#trait::{NodeEvaluation, RetrievalStrategy, StrategyCapabilities}; +use crate::config::StrategyConfig; +use crate::domain::{DocumentTree, NodeId}; /// Embedding model trait for semantic strategies. #[async_trait] @@ -259,7 +259,7 @@ impl RetrievalStrategy for SemanticStrategy { .collect() } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "semantic" } @@ -273,6 +273,9 @@ impl RetrievalStrategy for SemanticStrategy { } fn suitable_for_complexity(&self, complexity: QueryComplexity) -> bool { - matches!(complexity, QueryComplexity::Simple | QueryComplexity::Medium) + matches!( + complexity, + QueryComplexity::Simple | QueryComplexity::Medium + ) } } diff --git a/src/retrieval/strategy/trait.rs b/src/retrieval/strategy/trait.rs index 75aa5c03..3699a128 100644 --- a/src/retrieval/strategy/trait.rs +++ b/src/retrieval/strategy/trait.rs @@ -5,9 +5,9 @@ use async_trait::async_trait; -use crate::domain::{NodeId, DocumentTree}; -use super::super::types::{NavigationDecision, QueryComplexity}; use super::super::RetrievalContext; +use super::super::types::{NavigationDecision, QueryComplexity}; +use crate::domain::{DocumentTree, NodeId}; /// Result of evaluating a single node. #[derive(Debug, Clone)] @@ -89,8 +89,16 @@ pub trait RetrievalStrategy: Send + Sync { /// Estimate the cost of evaluating a set of nodes. fn estimate_cost(&self, node_count: usize) -> StrategyCost { StrategyCost { - llm_calls: if self.capabilities().uses_llm { node_count } else { 0 }, - tokens: if self.capabilities().uses_llm { node_count * 200 } else { 0 }, + llm_calls: if self.capabilities().uses_llm { + node_count + } else { + 0 + }, + tokens: if self.capabilities().uses_llm { + node_count * 200 + } else { + 0 + }, } } } diff --git a/src/retrieval/sufficiency/llm_judge.rs b/src/retrieval/sufficiency/llm_judge.rs index 01bba610..df80379b 100644 --- a/src/retrieval/sufficiency/llm_judge.rs +++ b/src/retrieval/sufficiency/llm_judge.rs @@ -8,8 +8,8 @@ use async_trait::async_trait; use serde::{Deserialize, Serialize}; -use crate::config::SufficiencyConfig; use super::{SufficiencyChecker, SufficiencyLevel}; +use crate::config::SufficiencyConfig; /// LLM client trait for the judge. #[async_trait] @@ -112,8 +112,14 @@ Be conservative - only mark as sufficient if you're confident the content answer let sufficient_keywords = ["sufficient", "yes", "complete", "enough"]; let insufficient_keywords = ["insufficient", "no", "incomplete", "not enough"]; - let sufficient_count = sufficient_keywords.iter().filter(|k| lower.contains(*k)).count(); - let insufficient_count = insufficient_keywords.iter().filter(|k| lower.contains(*k)).count(); + let sufficient_count = sufficient_keywords + .iter() + .filter(|k| lower.contains(*k)) + .count(); + let insufficient_count = insufficient_keywords + .iter() + .filter(|k| lower.contains(*k)) + .count(); if sufficient_count > insufficient_count { (SufficiencyLevel::PartialSufficient, 0.6) @@ -172,7 +178,7 @@ impl SufficiencyChecker for LlmJudge { } } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "llm_judge" } } diff --git a/src/retrieval/sufficiency/mod.rs b/src/retrieval/sufficiency/mod.rs index 2b77509f..ab3501bf 100644 --- a/src/retrieval/sufficiency/mod.rs +++ b/src/retrieval/sufficiency/mod.rs @@ -5,12 +5,12 @@ //! //! Determines when enough information has been collected to answer the query. -mod threshold; mod llm_judge; +mod threshold; -pub use threshold::ThresholdChecker; -pub use llm_judge::LlmJudge; pub use super::types::SufficiencyLevel; +pub use llm_judge::LlmJudge; +pub use threshold::ThresholdChecker; /// Trait for sufficiency checking strategies. pub trait SufficiencyChecker: Send + Sync { diff --git a/src/retrieval/sufficiency/threshold.rs b/src/retrieval/sufficiency/threshold.rs index da0e58f4..30af9197 100644 --- a/src/retrieval/sufficiency/threshold.rs +++ b/src/retrieval/sufficiency/threshold.rs @@ -5,8 +5,8 @@ //! //! Uses simple heuristics like token count and content length. -use crate::config::SufficiencyConfig; use super::{SufficiencyChecker, SufficiencyLevel}; +use crate::config::SufficiencyConfig; /// Configuration for threshold-based checking. #[derive(Debug, Clone)] @@ -142,7 +142,7 @@ impl SufficiencyChecker for ThresholdChecker { SufficiencyLevel::Insufficient } - fn name(&self) -> &str { + fn name(&self) -> &'static str { "threshold" } } diff --git a/src/retrieval/types.rs b/src/retrieval/types.rs index 2ee3361a..2077f325 100644 --- a/src/retrieval/types.rs +++ b/src/retrieval/types.rs @@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize}; -use crate::domain::NodeId; use super::context::{PruningStrategy, TokenEstimation}; +use crate::domain::NodeId; /// Query complexity level for adaptive strategy selection. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/src/storage/mod.rs b/src/storage/mod.rs index 9146633e..0d07d143 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -31,13 +31,8 @@ mod workspace; // Re-export main types pub use persistence::{ - DocumentMeta, - PersistedDocument, - PageContent, - save_document, - load_document, + DocumentMeta, PageContent, PersistedDocument, load_document, load_index, save_document, save_index, - load_index, }; -pub use workspace::{Workspace, DocumentMetaEntry}; +pub use workspace::{DocumentMetaEntry, Workspace}; diff --git a/src/storage/persistence.rs b/src/storage/persistence.rs index 3f05a758..a77a3e0b 100644 --- a/src/storage/persistence.rs +++ b/src/storage/persistence.rs @@ -3,11 +3,11 @@ //! Persistence utilities for saving and loading document indices. -use std::path::{Path, PathBuf}; -use std::io; use serde::{Deserialize, Serialize}; +use std::io; +use std::path::{Path, PathBuf}; -use crate::domain::{DocumentTree, Result, Error}; +use crate::domain::{DocumentTree, Error, Result}; /// Metadata for a persisted document. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -118,16 +118,14 @@ pub fn save_document(path: &Path, doc: &PersistedDocument) -> Result<()> { let json = serde_json::to_string_pretty(doc) .map_err(|e| Error::Io(io::Error::new(io::ErrorKind::Other, e)))?; - std::fs::write(path, json) - .map_err(|e| Error::Io(e))?; + std::fs::write(path, json).map_err(|e| Error::Io(e))?; Ok(()) } /// Load a document from a JSON file. pub fn load_document(path: &Path) -> Result { - let json = std::fs::read_to_string(path) - .map_err(|e| Error::Io(e))?; + let json = std::fs::read_to_string(path).map_err(|e| Error::Io(e))?; let doc: PersistedDocument = serde_json::from_str(&json) .map_err(|e| Error::Parse(format!("Failed to parse document: {}", e)))?; @@ -140,8 +138,7 @@ pub fn save_index(path: &Path, entries: &[DocumentMeta]) -> Result<()> { let json = serde_json::to_string_pretty(entries) .map_err(|e| Error::Io(io::Error::new(io::ErrorKind::Other, e)))?; - std::fs::write(path, json) - .map_err(|e| Error::Io(e))?; + std::fs::write(path, json).map_err(|e| Error::Io(e))?; Ok(()) } @@ -152,8 +149,7 @@ pub fn load_index(path: &Path) -> Result> { return Ok(Vec::new()); } - let json = std::fs::read_to_string(path) - .map_err(|e| Error::Io(e))?; + let json = std::fs::read_to_string(path).map_err(|e| Error::Io(e))?; let entries: Vec = serde_json::from_str(&json) .map_err(|e| Error::Parse(format!("Failed to parse index: {}", e)))?; diff --git a/src/storage/workspace.rs b/src/storage/workspace.rs index 040828e6..9cd1a83f 100644 --- a/src/storage/workspace.rs +++ b/src/storage/workspace.rs @@ -24,18 +24,18 @@ //! - Read operations (`get_meta`, `contains`, `list_documents`) only need `&self` //! - Cache updates happen internally via `Mutex` -use std::path::{Path, PathBuf}; use std::collections::HashMap; use std::fs; use std::num::NonZeroUsize; +use std::path::{Path, PathBuf}; use std::sync::Mutex; use lru::LruCache; use serde::{Deserialize, Serialize}; use tracing::{debug, info, warn}; -use crate::domain::{Result, Error}; -use super::persistence::{PersistedDocument, save_document, load_document}; +use super::persistence::{PersistedDocument, load_document, save_document}; +use crate::domain::{Error, Result}; const META_FILE: &str = "_meta.json"; const DEFAULT_CACHE_SIZE: usize = 100; @@ -97,8 +97,7 @@ impl Workspace { /// Create a new workspace with custom LRU cache size. pub fn with_cache_size(path: impl Into, cache_size: usize) -> Result { let root = path.into(); - fs::create_dir_all(&root) - .map_err(Error::Io)?; + fs::create_dir_all(&root).map_err(Error::Io)?; let capacity = NonZeroUsize::new(cache_size.max(1)) .unwrap_or_else(|| NonZeroUsize::new(DEFAULT_CACHE_SIZE).unwrap()); @@ -121,7 +120,10 @@ impl Workspace { } /// Open with custom cache size. - pub fn open_with_cache_size(path: impl Into + Clone, cache_size: usize) -> Result { + pub fn open_with_cache_size( + path: impl Into + Clone, + cache_size: usize, + ) -> Result { let root = path.clone().into(); if root.exists() { let capacity = NonZeroUsize::new(cache_size.max(1)) @@ -178,7 +180,11 @@ impl Workspace { doc_name: doc.meta.name.clone(), doc_description: doc.meta.description.clone(), doc_type: doc.meta.format.clone(), - path: doc.meta.source_path.as_ref().map(|p| p.to_string_lossy().to_string()), + path: doc + .meta + .source_path + .as_ref() + .map(|p| p.to_string_lossy().to_string()), page_count: doc.pages.first().map(|p| p.page), line_count: None, // TODO: track this }; @@ -208,7 +214,9 @@ impl Workspace { // Check LRU cache first (with lock) { - let mut inner = self.inner.lock() + let mut inner = self + .inner + .lock() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; if let Some(cached) = inner.document_cache.get(id) { @@ -228,7 +236,9 @@ impl Workspace { // Add to LRU cache (with lock) { - let mut inner = self.inner.lock() + let mut inner = self + .inner + .lock() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; inner.document_cache.put(id.to_string(), doc.clone()); } @@ -245,8 +255,7 @@ impl Workspace { let doc_path = self.document_path(id); if doc_path.exists() { - fs::remove_file(&doc_path) - .map_err(Error::Io)?; + fs::remove_file(&doc_path).map_err(Error::Io)?; } self.meta_index.remove(id); @@ -274,7 +283,8 @@ impl Workspace { /// Get the number of items currently in the LRU cache. pub fn cache_len(&self) -> usize { - self.inner.lock() + self.inner + .lock() .map(|inner| inner.document_cache.len()) .unwrap_or(0) } @@ -307,14 +317,16 @@ impl Workspace { return Ok(()); } - let content = fs::read_to_string(&meta_path) - .map_err(Error::Io)?; + let content = fs::read_to_string(&meta_path).map_err(Error::Io)?; let meta: HashMap = serde_json::from_str(&content) .map_err(|e| Error::Parse(format!("Failed to parse meta index: {}", e)))?; self.meta_index = meta; - info!("Loaded {} document(s) from workspace index", self.meta_index.len()); + info!( + "Loaded {} document(s) from workspace index", + self.meta_index.len() + ); Ok(()) } @@ -323,8 +335,7 @@ impl Workspace { let content = serde_json::to_string_pretty(&self.meta_index) .map_err(|e| Error::Parse(format!("Failed to serialize meta index: {}", e)))?; - fs::write(&self.meta_path(), content) - .map_err(Error::Io)?; + fs::write(self.meta_path(), content).map_err(Error::Io)?; Ok(()) } @@ -335,7 +346,11 @@ impl Workspace { .map_err(Error::Io)? .filter_map(|entry| entry.ok()) .filter(|entry| { - entry.path().extension().map(|ext| ext == "json").unwrap_or(false) + entry + .path() + .extension() + .map(|ext| ext == "json") + .unwrap_or(false) }) .filter_map(|entry| { let path = entry.path(); @@ -351,7 +366,11 @@ impl Workspace { doc_name: doc.meta.name, doc_description: doc.meta.description, doc_type: doc.meta.format, - path: doc.meta.source_path.as_ref().map(|p| p.to_string_lossy().to_string()), + path: doc + .meta + .source_path + .as_ref() + .map(|p| p.to_string_lossy().to_string()), page_count: doc.pages.first().map(|p| p.page), line_count: None, }; @@ -366,9 +385,12 @@ impl Workspace { if !self.meta_index.is_empty() { self.save_meta_index()?; - info!("Rebuilt index from {} document file(s)", self.meta_index.len()); + info!( + "Rebuilt index from {} document file(s)", + self.meta_index.len() + ); } Ok(()) } -} \ No newline at end of file +} diff --git a/src/throttle/config.rs b/src/throttle/config.rs index 51088377..155e2e3a 100644 --- a/src/throttle/config.rs +++ b/src/throttle/config.rs @@ -40,9 +40,15 @@ pub struct ConcurrencyConfig { pub semaphore_enabled: bool, } -fn default_max_concurrent_requests() -> usize { 10 } -fn default_requests_per_minute() -> usize { 500 } -fn default_true() -> bool { true } +fn default_max_concurrent_requests() -> usize { + 10 +} +fn default_requests_per_minute() -> usize { + 500 +} +fn default_true() -> bool { + true +} impl Default for ConcurrencyConfig { fn default() -> Self { diff --git a/src/throttle/controller.rs b/src/throttle/controller.rs index 845cf67b..87193fa4 100644 --- a/src/throttle/controller.rs +++ b/src/throttle/controller.rs @@ -99,7 +99,10 @@ impl ConcurrencyController { if self.config.semaphore_enabled { trace!("Waiting for semaphore permit"); let permit = self.semaphore.acquire().await.unwrap(); - debug!("Semaphore: permit acquired (available: {})", self.semaphore.available_permits()); + debug!( + "Semaphore: permit acquired (available: {})", + self.semaphore.available_permits() + ); Some(permit) } else { None @@ -144,7 +147,10 @@ impl ConcurrencyController { impl std::fmt::Debug for ConcurrencyController { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("ConcurrencyController") - .field("max_concurrent_requests", &self.config.max_concurrent_requests) + .field( + "max_concurrent_requests", + &self.config.max_concurrent_requests, + ) .field("requests_per_minute", &self.config.requests_per_minute) .field("rate_limiting_enabled", &self.config.enabled) .field("semaphore_enabled", &self.config.semaphore_enabled) diff --git a/src/throttle/mod.rs b/src/throttle/mod.rs index ede6133d..0e07c258 100644 --- a/src/throttle/mod.rs +++ b/src/throttle/mod.rs @@ -55,9 +55,9 @@ //! ``` mod config; -mod rate_limiter; mod controller; +mod rate_limiter; pub use config::ConcurrencyConfig; -pub use rate_limiter::RateLimiter; pub use controller::ConcurrencyController; +pub use rate_limiter::RateLimiter; diff --git a/src/throttle/rate_limiter.rs b/src/throttle/rate_limiter.rs index eff556b6..90a865e9 100644 --- a/src/throttle/rate_limiter.rs +++ b/src/throttle/rate_limiter.rs @@ -4,9 +4,9 @@ //! Rate limiter using token bucket algorithm (governor). use governor::{ + Quota, RateLimiter as GovernorLimiter, clock::{Clock, DefaultClock}, state::{InMemoryState, NotKeyed}, - Quota, RateLimiter as GovernorLimiter, }; use std::num::NonZeroU32; use std::sync::Arc;