From 077b401299f206b74da601b87677921586635cd7 Mon Sep 17 00:00:00 2001 From: zTgx <747674262@qq.com> Date: Sun, 5 Apr 2026 13:52:44 +0800 Subject: [PATCH] refactor(examples): restructure document array formatting in batch processing - Split document entries into separate lines for better readability - Format string literals with proper indentation and line breaks - Improve code formatting consistency across all document entries perf(examples): format print statements with proper line breaks - Replace inline print statements with multi-line format for better readability - Maintain character count under 100 per line as per style guidelines - Apply consistent formatting across multiple example files refactor(examples): reorder imports for better organization - Group related imports together in alphabetical order - Remove unused imports and reorder remaining ones - Maintain consistency across different example files refactor(storage): improve codec implementation logging - Format log messages with proper multiline structure - Maintain character limit per line for better readability - Enhance logging output formatting in storage implementations --- examples/batch_processing.rs | 431 ++++++++++++++++------- examples/content_aggregation.rs | 66 ++-- examples/events.rs | 105 +++--- examples/session.rs | 5 +- examples/storage_async.rs | 9 +- examples/storage_backend.rs | 13 +- examples/storage_compression.rs | 33 +- examples/storage_migration.rs | 5 +- examples/storage_workspace.rs | 5 +- src/client/builder.rs | 13 +- src/client/context.rs | 7 +- src/client/engine.rs | 57 ++- src/client/events.rs | 2 +- src/client/indexer.rs | 15 +- src/client/mod.rs | 27 +- src/client/retriever.rs | 57 +-- src/client/session.rs | 33 +- src/client/types.rs | 3 +- src/client/workspace.rs | 58 ++- src/config/docs.rs | 398 ++++++++++++++++++--- src/config/loader.rs | 31 +- src/config/merge.rs | 18 +- src/config/mod.rs | 36 +- src/config/types/fallback.rs | 4 +- src/config/types/mod.rs | 33 +- src/config/validator.rs | 157 +++++---- src/document/tree.rs | 4 +- src/error.rs | 17 +- src/index/config.rs | 2 +- src/index/mod.rs | 4 +- src/index/stages/enhance.rs | 3 +- src/index/stages/optimize.rs | 3 +- src/index/stages/parse.rs | 5 +- src/lib.rs | 8 +- src/parser/docx/parser.rs | 2 +- src/parser/markdown/parser.rs | 2 +- src/parser/pdf/parser.rs | 2 +- src/parser/registry.rs | 2 +- src/retrieval/content/aggregator.rs | 15 +- src/retrieval/content/budget.rs | 26 +- src/retrieval/content/builder.rs | 21 +- src/retrieval/content/config.rs | 6 +- src/retrieval/content/mod.rs | 8 +- src/retrieval/content/scorer.rs | 164 +++++++-- src/retrieval/mod.rs | 2 +- src/retrieval/pilot/budget.rs | 29 +- src/retrieval/pilot/builder.rs | 44 ++- src/retrieval/pilot/config.rs | 2 +- src/retrieval/pilot/decision.rs | 13 +- src/retrieval/pilot/fallback.rs | 10 +- src/retrieval/pilot/llm_pilot.rs | 47 ++- src/retrieval/pilot/metrics.rs | 32 +- src/retrieval/pilot/mod.rs | 8 +- src/retrieval/pilot/noop.rs | 11 +- src/retrieval/pilot/parser.rs | 68 +++- src/retrieval/pilot/prompts/builder.rs | 2 +- src/retrieval/pilot/prompts/mod.rs | 4 +- src/retrieval/pilot/prompts/templates.rs | 6 +- src/retrieval/pilot/trait.rs | 13 +- src/retrieval/pipeline/orchestrator.rs | 33 +- src/retrieval/search/beam.rs | 34 +- src/retrieval/search/greedy.rs | 19 +- src/retrieval/stages/judge.rs | 15 +- src/retrieval/stages/search.rs | 20 +- src/storage/async_workspace.rs | 39 +- src/storage/backend/file.rs | 6 +- src/storage/backend/memory.rs | 56 +-- src/storage/cache.rs | 57 +-- src/storage/codec.rs | 12 +- src/storage/lock.rs | 18 +- src/storage/migration.rs | 34 +- src/storage/mod.rs | 12 +- src/storage/persistence.rs | 38 +- src/storage/workspace.rs | 46 ++- src/util/mod.rs | 6 +- 75 files changed, 1768 insertions(+), 883 deletions(-) diff --git a/examples/batch_processing.rs b/examples/batch_processing.rs index 6906189f..16a29896 100644 --- a/examples/batch_processing.rs +++ b/examples/batch_processing.rs @@ -33,7 +33,9 @@ async fn main() -> Result<(), Box> { let temp_dir = tempfile::tempdir()?; let documents = vec![ - ("intro.md", r#"# Introduction + ( + "intro.md", + r#"# Introduction Welcome to the vectorless library. This is a document intelligence engine. @@ -42,8 +44,11 @@ Welcome to the vectorless library. This is a document intelligence engine. - Tree-based navigation - Multi-format support - Session management -"#), - ("api.md", r#"# API Reference +"#, + ), + ( + "api.md", + r#"# API Reference ## Engine @@ -62,8 +67,11 @@ Multi-document operations with caching. - `index(path)`: Index into session - `query_all(question)`: Query across all documents -"#), - ("guide.md", r#"# User Guide +"#, + ), + ( + "guide.md", + r#"# User Guide ## Getting Started @@ -74,8 +82,11 @@ First, create a client with workspace configuration. - Use sessions for multi-document operations - Enable caching for better performance - Monitor events for debugging -"#), - ("advanced.md", r#"# Advanced Topics +"#, + ), + ( + "advanced.md", + r#"# Advanced Topics ## Performance Tuning @@ -89,8 +100,11 @@ Configure retrieval parameters for optimal performance. ## Custom Pilots Implement custom navigation logic. -"#), - ("reference.md", r#"# Reference +"#, + ), + ( + "reference.md", + r#"# Reference ## Configuration @@ -103,8 +117,11 @@ All configuration is done via TOML files. top_k = 5 max_tokens = 4000 ``` -"#), - ("examples.md", r#"# Examples +"#, + ), + ( + "examples.md", + r#"# Examples ## Basic Usage @@ -117,8 +134,11 @@ Process multiple documents concurrently. ## Session Usage Multi-document operations with caching. -"#), - ("faq.md", r#"# FAQ +"#, + ), + ( + "faq.md", + r#"# FAQ ## Common Questions @@ -130,8 +150,11 @@ A: Use `engine.query(doc_id, question)` method. **Q: What formats are supported?** A: Markdown, PDF, DOCX, HTML. -"#), - ("changelog.md", r#"# Changelog +"#, + ), + ( + "changelog.md", + r#"# Changelog ## Version 0.1.0 @@ -144,8 +167,11 @@ A: Markdown, PDF, DOCX, HTML. - Session support - Event system - Content aggregator -"#), - ("contributing.md", r#"# Contributing +"#, + ), + ( + "contributing.md", + r#"# Contributing ## How to Contribute @@ -160,14 +186,20 @@ We welcome contributions! Please follow these steps: - Run `cargo fmt` - Run `cargo clippy` - Add tests -"#), - ("license.md", r#"# License +"#, + ), + ( + "license.md", + r#"# License Apache License, Version 2.0 Copyright 2026 vectorless developers -"#), - ("architecture.md", r#"# Architecture +"#, + ), + ( + "architecture.md", + r#"# Architecture ## Overview @@ -179,8 +211,11 @@ Vectorless uses a tree-based architecture. - Indexer: Tree building - Retriever: Content search - Storage: Persistence -"#), - ("security.md", r#"# Security +"#, + ), + ( + "security.md", + r#"# Security ## Security Considerations @@ -192,8 +227,11 @@ Vectorless uses a tree-based architecture. - Use environment variables - Rotate keys periodically -"#), - ("performance.md", r#"# Performance +"#, + ), + ( + "performance.md", + r#"# Performance ## Optimization Tips @@ -204,8 +242,11 @@ Vectorless uses a tree-based architecture. ## Benchmarks Run `cargo bench` for performance metrics. -"#), - ("testing.md", r#"# Testing +"#, + ), + ( + "testing.md", + r#"# Testing ## Running Tests @@ -218,8 +259,11 @@ cargo test - Unit tests - Integration tests - Example tests -"#), - ("deployment.md", r#"# Deployment +"#, + ), + ( + "deployment.md", + r#"# Deployment ## Production Setup @@ -230,8 +274,11 @@ cargo test ## Configuration Use TOML configuration files. -"#), - ("troubleshooting.md", r#"# Troubleshooting +"#, + ), + ( + "troubleshooting.md", + r#"# Troubleshooting ## Common Issues @@ -246,8 +293,11 @@ Ensure document is indexed. ### Performance Issues Reduce batch size or enable caching. -"#), - ("integrations.md", r#"# Integrations +"#, + ), + ( + "integrations.md", + r#"# Integrations ## LLM Providers @@ -259,8 +309,11 @@ Reduce batch size or enable caching. - File system (default) - S3 (planned) -"#), - ("migrations.md", r#"# Migrations +"#, + ), + ( + "migrations.md", + r#"# Migrations ## Version Migrations @@ -268,8 +321,11 @@ Reduce batch size or enable caching. - Update configuration format - Re-index documents -"#), - ("roadmap.md", r#"# Roadmap +"#, + ), + ( + "roadmap.md", + r#"# Roadmap ## Future Plans @@ -282,8 +338,11 @@ Reduce batch size or enable caching. - Distributed indexing - Real-time updates -"#), - ("credits.md", r#"# Credits +"#, + ), + ( + "credits.md", + r#"# Credits ## Contributors @@ -292,8 +351,11 @@ Thanks to all contributors! ## Libraries Built with Rust and many open-source libraries. -"#), - ("index.md", r#"# Index +"#, + ), + ( + "index.md", + r#"# Index ## Quick Links @@ -304,8 +366,11 @@ Built with Rust and many open-source libraries. ## Search Use the search functionality to find specific content. -"#), - ("search.md", r#"# Search +"#, + ), + ( + "search.md", + r#"# Search ## Search Functionality @@ -318,8 +383,11 @@ let results = engine.query(&doc_id, "search term").await?; ### Advanced Search Use sessions for cross-document search. -"#), - ("export.md", r#"# Export +"#, + ), + ( + "export.md", + r#"# Export ## Exporting Data @@ -332,8 +400,11 @@ let json = tree.to_structure_json(); ### Custom Formats Implement custom exporters as needed. -"#), - ("import.md", r#"# Import +"#, + ), + ( + "import.md", + r#"# Import ## Importing Data @@ -346,8 +417,11 @@ let doc_id = engine.index("./document.md").await?; ### From Memory Use the content directly with parsers. -"#), - ("validation.md", r#"# Validation +"#, + ), + ( + "validation.md", + r#"# Validation ## Input Validation @@ -362,8 +436,11 @@ Validated on load with helpful errors. ### Queries Sanitized before processing. -"#), - ("formatting.md", r#"# Formatting +"#, + ), + ( + "formatting.md", + r#"# Formatting ## Content Formatting @@ -378,8 +455,11 @@ Syntax highlighting support. ### Tables Basic table parsing. -"#), - ("localization.md", r#"# Localization +"#, + ), + ( + "localization.md", + r#"# Localization ## Internationalization @@ -391,8 +471,11 @@ Planned i18n support for: - Error messages - UI strings - Documentation -"#), - ("accessibility.md", r#"# Accessibility +"#, + ), + ( + "accessibility.md", + r#"# Accessibility ## Accessibility @@ -407,8 +490,11 @@ Consistent and intuitive naming. ### Error Messages Helpful and actionable. -"#), - ("glossary.md", r#"# Glossary +"#, + ), + ( + "glossary.md", + r#"# Glossary ## Terms @@ -416,16 +502,22 @@ Helpful and actionable. - **Session**: Multi-document context - **Workspace**: Document storage - **Retrieval**: Content search -"#), - ("appendix.md", r#"# Appendix +"#, + ), + ( + "appendix.md", + r#"# Appendix ## Additional Resources - [GitHub Repository](https://github.com) - [Documentation Site](https://docs.vectorless.dev) - [Community Discord](https://discord.gg) -"#), - ("summary.md", r#"# Summary +"#, + ), + ( + "summary.md", + r#"# Summary ## Overview @@ -436,8 +528,11 @@ This documentation covers all aspects of vectorless. - Try the examples - Join the community - Contribute! -"#), - ("conclusion.md", r#"# Conclusion +"#, + ), + ( + "conclusion.md", + r#"# Conclusion ## Thank You @@ -446,8 +541,11 @@ Thanks for using vectorless! ## Feedback We'd love to hear from you. Open an issue on GitHub. -"#), - ("revision.md", r#"# Revision History +"#, + ), + ( + "revision.md", + r#"# Revision History ## Document Versions @@ -455,8 +553,11 @@ We'd love to hear from you. Open an issue on GitHub. |---------|------------|---------------------------| | 1.0 | 2026-01-01 | Initial version | | 1.1 | 2026-02-01 | Session support | -"#), - ("feedback.md", r#"# Feedback +"#, + ), + ( + "feedback.md", + r#"# Feedback ## Providing Feedback @@ -473,8 +574,11 @@ We value your input! - Bug reports - Feature requests - Documentation improvements -"#), - ("support.md", r#"# Support +"#, + ), + ( + "support.md", + r#"# Support ## Getting Help @@ -489,8 +593,11 @@ Join our Discord for discussions. ### Enterprise Contact us for enterprise support. -"#), - ("updates.md", r#"# Updates +"#, + ), + ( + "updates.md", + r#"# Updates ## Staying Updated @@ -505,8 +612,11 @@ Apply security patches promptly. ### Deprecations Watch for deprecation notices. -"#), - ("resources.md", r#"# Resources +"#, + ), + ( + "resources.md", + r#"# Resources ## External Resources @@ -521,8 +631,11 @@ Watch for deprecation notices. - Blog posts - Tutorial videos - Example projects -"#), - ("contact.md", r#"# Contact +"#, + ), + ( + "contact.md", + r#"# Contact ## Contact Information @@ -537,8 +650,11 @@ Email: security@vectorless.dev ### Enterprise Sales Email: enterprise@vectorless.dev -"#), - ("privacy.md", r#"# Privacy Policy +"#, + ), + ( + "privacy.md", + r#"# Privacy Policy ## Data Handling @@ -551,8 +667,11 @@ We don't track usage or content. ## API Keys Stored securely in configuration files. -"#), - ("terms.md", r#"# Terms of Service +"#, + ), + ( + "terms.md", + r#"# Terms of Service ## Usage Terms @@ -565,8 +684,11 @@ By using vectorless, you agree to: ## Changes Terms may be updated. Check for revisions. -"#), - ("legal.md", r#"# Legal +"#, + ), + ( + "legal.md", + r#"# Legal ## Licensing @@ -579,8 +701,11 @@ Copyright 2026 vectorless developers ## Trademarks Vectorless is a trademark. -"#), - ("versioning.md", r#"# Versioning +"#, + ), + ( + "versioning.md", + r#"# Versioning ## Semantic Versioning @@ -593,8 +718,11 @@ We follow semver: ## Current Version 0.1.10 -"#), - ("compatibility.md", r#"# Compatibility +"#, + ), + ( + "compatibility.md", + r#"# Compatibility ## Supported Versions @@ -610,8 +738,11 @@ We follow semver: ## Breaking Changes Documented in changelog. -"#), - ("installation.md", r#"# Installation +"#, + ), + ( + "installation.md", + r#"# Installation ## Requirements @@ -629,8 +760,11 @@ cargo install vectorless ```bash vectorless --version ``` -"#), - ("quickstart.md", r#"# Quick Start +"#, + ), + ( + "quickstart.md", + r#"# Quick Start ## 5-Minute Setup @@ -644,8 +778,11 @@ let client = Engine::builder().build()?; let doc_id = client.index("./doc.md").await?; let result = client.query(&doc_id, "What is this?").await?; ``` -"#), - ("tutorial.md", r#"# Tutorial +"#, + ), + ( + "tutorial.md", + r#"# Tutorial ## Introduction @@ -666,8 +803,11 @@ Ask questions about your document. ## Step 4: Next Explore advanced features. -"#), - ("examples_overview.md", r#"# Examples Overview +"#, + ), + ( + "examples_overview.md", + r#"# Examples Overview ## Available Examples @@ -683,8 +823,11 @@ Explore advanced features. ```bash cargo run --example ``` -"#), - ("configuration.md", r#"# Configuration +"#, + ), + ( + "configuration.md", + r#"# Configuration ## Configuration File @@ -702,8 +845,11 @@ max_tokens = 4000 ## Environment Variables - `OPENAI_API_KEY`: LLM API key -"#), - ("optimization.md", r#"# Optimization +"#, + ), + ( + "optimization.md", + r#"# Optimization ## Performance Tips @@ -718,8 +864,11 @@ Documents are cached in sessions. ## Concurrency Use `buffer_unordered` for parallel indexing. -"#), - ("errors.md", r#"# Error Handling +"#, + ), + ( + "errors.md", + r#"# Error Handling ## Error Types @@ -736,8 +885,11 @@ match result { Err(e) => { /* other error */ }, } ``` -"#), - ("logging.md", r#"# Logging +"#, + ), + ( + "logging.md", + r#"# Logging ## Log Levels @@ -752,8 +904,11 @@ match result { ```bash RUST_LOG=debug cargo run ``` -"#), - ("metrics.md", r#"# Metrics +"#, + ), + ( + "metrics.md", + r#"# Metrics ## Available Metrics @@ -767,8 +922,11 @@ RUST_LOG=debug cargo run let stats = session.stats(); println!("Cache hit rate: {:.1}%", stats.cache_hit_rate() * 100.0); ``` -"#), - ("health.md", r#"# Health Checks +"#, + ), + ( + "health.md", + r#"# Health Checks ## Workspace Health @@ -782,8 +940,11 @@ println!("{} documents indexed", docs.len()); ## Session Health Monitor session statistics regularly. -"#), - ("backup.md", r#"# Backup +"#, + ), + ( + "backup.md", + r#"# Backup ## Backing Up @@ -800,8 +961,11 @@ Restore by copying back: ```bash cp -r ./workspace_backup ./workspace ``` -"#), - ("recovery.md", r#"# Recovery +"#, + ), + ( + "recovery.md", + r#"# Recovery ## Corrupted Documents @@ -815,8 +979,11 @@ engine.index(&path).await?; ## Session Recovery Create a new session if issues occur. -"#), - ("monitoring.md", r#"# Monitoring +"#, + ), + ( + "monitoring.md", + r#"# Monitoring ## Production Monitoring @@ -832,8 +999,11 @@ let events = EventEmitter::new() ## Alerts Set up alerts for error rates. -"#), - ("scaling.md", r#"# Scaling +"#, + ), + ( + "scaling.md", + r#"# Scaling ## Horizontal Scaling @@ -848,8 +1018,11 @@ Increase resources for single instance. - Storage backend - Cache coordination - Rate limiting -"#), - ("security_config.md", r#"# Security Configuration +"#, + ), + ( + "security_config.md", + r#"# Security Configuration ## API Keys @@ -867,7 +1040,8 @@ Use HTTPS for all API calls. ## Access Control Implement authentication for production. -"#), +"#, + ), ]; for (name, content) in &documents { @@ -896,13 +1070,19 @@ Implement authentication for production. let elapsed = start.elapsed(); println!(" ✓ Indexed {} documents in {:?}", doc_ids.len(), elapsed); - println!(" - Rate: {:.1} docs/sec", doc_ids.len() as f64 / elapsed.as_secs_f64()); + println!( + " - Rate: {:.1} docs/sec", + doc_ids.len() as f64 / elapsed.as_secs_f64() + ); println!(); // 4. Show session stats println!("Step 4: Session statistics:"); let stats = session.stats(); - println!(" - Documents in session: {}", session.list_documents().len()); + println!( + " - Documents in session: {}", + session.list_documents().len() + ); println!(" - Queries: {}", stats.query_count.get()); println!(); @@ -939,8 +1119,14 @@ Implement authentication for production. let elapsed = start.elapsed(); println!(" ✓ Completed {} queries in {:?}", queries.len(), elapsed); - println!(" - Success rate: {:.0}%", (success_count as f64 / queries.len() as f64) * 100.0); - println!(" - Rate: {:.1} queries/sec", queries.len() as f64 / elapsed.as_secs_f64()); + println!( + " - Success rate: {:.0}%", + (success_count as f64 / queries.len() as f64) * 100.0 + ); + println!( + " - Rate: {:.1} queries/sec", + queries.len() as f64 / elapsed.as_secs_f64() + ); println!(); // 6. Final statistics @@ -950,10 +1136,7 @@ Implement authentication for production. println!(" - Total queries: {}", stats.query_count.get()); println!(" - Cache hits: {}", stats.cache_hits.get()); println!(" - Cache misses: {}", stats.cache_misses.get()); - println!( - " - Cache hit rate: {:.1}%", - stats.cache_hit_rate() * 100.0 - ); + println!(" - Cache hit rate: {:.1}%", stats.cache_hit_rate() * 100.0); if let Some(avg_time) = stats.avg_query_time() { println!(" - Avg query time: {:?}", avg_time); } diff --git a/examples/content_aggregation.rs b/examples/content_aggregation.rs index 9ead2aeb..1bfdcffc 100644 --- a/examples/content_aggregation.rs +++ b/examples/content_aggregation.rs @@ -14,13 +14,12 @@ //! cargo run --example content_aggregation //! ``` +use indextree::Arena; +use vectorless::document::NodeId; use vectorless::retrieval::content::{ - ContentAggregator, ContentAggregatorConfig, BudgetAllocator, AllocationStrategy, - StructureBuilder, OutputFormat, RelevanceScorer, ScoringStrategyConfig, - ContentChunk, ScoringContext, + AllocationStrategy, BudgetAllocator, ContentAggregator, ContentAggregatorConfig, ContentChunk, + OutputFormat, RelevanceScorer, ScoringContext, ScoringStrategyConfig, StructureBuilder, }; -use vectorless::document::NodeId; -use indextree::Arena; fn make_node_id() -> NodeId { let mut arena = Arena::new(); @@ -78,9 +77,12 @@ fn main() { println!("\nScored chunks:"); for chunk in &chunks { let relevance = scorer.score_chunk(chunk, &ctx); - println!(" - '{}' (depth {}): score {:.3}", - chunk.title, chunk.depth, relevance.score); - println!(" Components: keyword={:.2}, bm25={:.2}, depth_penalty={:.2}, density={:.2}", + println!( + " - '{}' (depth {}): score {:.3}", + chunk.title, chunk.depth, relevance.score + ); + println!( + " Components: keyword={:.2}, bm25={:.2}, depth_penalty={:.2}, density={:.2}", relevance.components.keyword_score, relevance.components.bm25_score, relevance.components.depth_penalty, @@ -99,12 +101,14 @@ fn main() { let strategies = vec![ ("Greedy", AllocationStrategy::Greedy), - ("Hierarchical (20%/level)", AllocationStrategy::Hierarchical { min_per_level: 0.2 }), + ( + "Hierarchical (20%/level)", + AllocationStrategy::Hierarchical { min_per_level: 0.2 }, + ), ]; for (name, strategy) in strategies { - let allocator = BudgetAllocator::new(200) - .with_strategy(strategy); + let allocator = BudgetAllocator::new(200).with_strategy(strategy); let result = allocator.allocate(scored.clone(), 2); @@ -114,9 +118,15 @@ fn main() { println!(" Avg score: {:.3}", result.stats.avg_score); for content in &result.selected { - let trunc = if content.is_truncated() { " [truncated]" } else { "" }; - println!(" - '{}' ({} tokens, score {:.2}){}", - content.title, content.tokens, content.score, trunc); + let trunc = if content.is_truncated() { + " [truncated]" + } else { + "" + }; + println!( + " - '{}' ({} tokens, score {:.2}){}", + content.title, content.tokens, content.score, trunc + ); } } @@ -129,8 +139,7 @@ fn main() { ("Flat", OutputFormat::Flat), ]; - let allocator = BudgetAllocator::new(500) - .with_strategy(AllocationStrategy::Greedy); + let allocator = BudgetAllocator::new(500).with_strategy(AllocationStrategy::Greedy); let result = allocator.allocate(scored.clone(), 2); for (name, format) in formats { @@ -138,7 +147,12 @@ fn main() { let tree = vectorless::document::DocumentTree::new("Test", ""); let structured = builder.build(result.selected.clone(), &tree); - println!("\n{} Output ({} chars, {} tokens):", name, structured.content.len(), structured.metadata.total_tokens); + println!( + "\n{} Output ({} chars, {} tokens):", + name, + structured.content.len(), + structured.metadata.total_tokens + ); let preview = if structured.content.len() > 300 { format!("{}...", &structured.content[..300]) } else { @@ -153,12 +167,18 @@ fn main() { let configs = vec![ ("Default (4000 tokens)", ContentAggregatorConfig::default()), - ("Conservative (1000 tokens)", ContentAggregatorConfig::new() - .with_token_budget(1000) - .with_min_relevance(0.3)), - ("High Precision (2000 tokens, 0.5 threshold)", ContentAggregatorConfig::new() - .with_token_budget(2000) - .with_min_relevance(0.5)), + ( + "Conservative (1000 tokens)", + ContentAggregatorConfig::new() + .with_token_budget(1000) + .with_min_relevance(0.3), + ), + ( + "High Precision (2000 tokens, 0.5 threshold)", + ContentAggregatorConfig::new() + .with_token_budget(2000) + .with_min_relevance(0.5), + ), ]; for (name, config) in configs { diff --git a/examples/events.rs b/examples/events.rs index eab7b68a..8f736d88 100644 --- a/examples/events.rs +++ b/examples/events.rs @@ -14,8 +14,8 @@ //! cargo run --example events //! ``` -use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use vectorless::client::{EngineBuilder, EventEmitter, IndexEvent, QueryEvent}; @@ -36,49 +36,55 @@ async fn main() -> Result<(), Box> { let events = EventEmitter::new() // Index events - .on_index(move |e| { - match e { - IndexEvent::Started { path } => { - println!(" [INDEX] Started: {}", path); - } - IndexEvent::FormatDetected { format } => { - println!(" [INDEX] Format: {:?}", format); - } - IndexEvent::TreeBuilt { node_count } => { - println!(" [INDEX] Tree built: {} nodes", node_count); - } - IndexEvent::Complete { doc_id } => { - println!(" [INDEX] Complete: {}", &doc_id[..8]); - index_count_clone.fetch_add(1, Ordering::SeqCst); - } - IndexEvent::Error { message } => { - println!(" [INDEX] Error: {}", message); - } - _ => {} + .on_index(move |e| match e { + IndexEvent::Started { path } => { + println!(" [INDEX] Started: {}", path); + } + IndexEvent::FormatDetected { format } => { + println!(" [INDEX] Format: {:?}", format); + } + IndexEvent::TreeBuilt { node_count } => { + println!(" [INDEX] Tree built: {} nodes", node_count); + } + IndexEvent::Complete { doc_id } => { + println!(" [INDEX] Complete: {}", &doc_id[..8]); + index_count_clone.fetch_add(1, Ordering::SeqCst); } + IndexEvent::Error { message } => { + println!(" [INDEX] Error: {}", message); + } + _ => {} }) // Query events - .on_query(move |e| { - match e { - QueryEvent::Started { query } => { - println!(" [QUERY] Started: \"{}\"", query); - query_count_clone.fetch_add(1, Ordering::SeqCst); - } - QueryEvent::NodeVisited { title, score, .. } => { - println!(" [QUERY] Visited: \"{}\" (score: {:.2})", title, score); - nodes_visited_clone.fetch_add(1, Ordering::SeqCst); - } - QueryEvent::CandidateFound { node_id, score } => { - println!(" [QUERY] Candidate: {} (score: {:.2})", &node_id[..8], score); - } - QueryEvent::Complete { total_results, confidence } => { - println!(" [QUERY] Complete: {} results, confidence: {:.2}", total_results, confidence); - } - QueryEvent::Error { message } => { - println!(" [QUERY] Error: {}", message); - } - _ => {} + .on_query(move |e| match e { + QueryEvent::Started { query } => { + println!(" [QUERY] Started: \"{}\"", query); + query_count_clone.fetch_add(1, Ordering::SeqCst); + } + QueryEvent::NodeVisited { title, score, .. } => { + println!(" [QUERY] Visited: \"{}\" (score: {:.2})", title, score); + nodes_visited_clone.fetch_add(1, Ordering::SeqCst); + } + QueryEvent::CandidateFound { node_id, score } => { + println!( + " [QUERY] Candidate: {} (score: {:.2})", + &node_id[..8], + score + ); + } + QueryEvent::Complete { + total_results, + confidence, + } => { + println!( + " [QUERY] Complete: {} results, confidence: {:.2}", + total_results, confidence + ); + } + QueryEvent::Error { message } => { + println!(" [QUERY] Error: {}", message); } + _ => {} }); println!(" ✓ Event handlers configured\n"); @@ -122,7 +128,9 @@ The event system uses handlers that can be attached to the engine builder. // 4. Query the document (events will fire) println!("Step 4: Querying document (watch events)...\n"); - let result = engine.query(&doc_id, "What features are available?").await?; + let result = engine + .query(&doc_id, "What features are available?") + .await?; println!(); // 5. Show results @@ -137,9 +145,18 @@ The event system uses handlers that can be attached to the engine builder. // 6. Show statistics println!("Step 6: Event statistics:"); - println!(" - Index events fired: {}", index_count.load(Ordering::SeqCst)); - println!(" - Query events fired: {}", query_count.load(Ordering::SeqCst)); - println!(" - Nodes visited: {}", nodes_visited.load(Ordering::SeqCst)); + println!( + " - Index events fired: {}", + index_count.load(Ordering::SeqCst) + ); + println!( + " - Query events fired: {}", + query_count.load(Ordering::SeqCst) + ); + println!( + " - Nodes visited: {}", + nodes_visited.load(Ordering::SeqCst) + ); println!(); // 7. Cleanup diff --git a/examples/session.rs b/examples/session.rs index 25aaf3ab..30a2f97a 100644 --- a/examples/session.rs +++ b/examples/session.rs @@ -185,10 +185,7 @@ token_budget = 4000 println!(" - Queries: {}", stats.query_count.get()); println!(" - Cache hits: {}", stats.cache_hits.get()); println!(" - Cache misses: {}", stats.cache_misses.get()); - println!( - " - Cache hit rate: {:.1}%", - stats.cache_hit_rate() * 100.0 - ); + println!(" - Cache hit rate: {:.1}%", stats.cache_hit_rate() * 100.0); if let Some(avg_time) = stats.avg_query_time() { println!(" - Avg query time: {:?}", avg_time); } diff --git a/examples/storage_async.rs b/examples/storage_async.rs index f7ecfaec..606755b4 100644 --- a/examples/storage_async.rs +++ b/examples/storage_async.rs @@ -41,7 +41,9 @@ async fn main() -> vectorless::Result<()> { println!("2. Adding documents..."); workspace.add(&create_doc("doc-1", "Document One")).await?; workspace.add(&create_doc("doc-2", "Document Two")).await?; - workspace.add(&create_doc("doc-3", "Document Three")).await?; + workspace + .add(&create_doc("doc-3", "Document Three")) + .await?; println!(" ✓ Added 3 documents\n"); // 3. Concurrent access example @@ -83,7 +85,10 @@ async fn main() -> vectorless::Result<()> { let len2 = ws2.len().await; let len3 = ws3.len().await; - println!(" ws1.len() = {}, ws2.len() = {}, ws3.len() = {}", len1, len2, len3); + println!( + " ws1.len() = {}, ws2.len() = {}, ws3.len() = {}", + len1, len2, len3 + ); println!(" ✓ All clones share the same state\n"); // Cleanup diff --git a/examples/storage_backend.rs b/examples/storage_backend.rs index 3b9a5fd9..7ad6f6f5 100644 --- a/examples/storage_backend.rs +++ b/examples/storage_backend.rs @@ -15,9 +15,9 @@ use std::collections::HashMap; use std::sync::{Arc, RwLock}; +use vectorless::Result; use vectorless::document::DocumentTree; use vectorless::storage::{DocumentMeta, PersistedDocument, StorageBackend, Workspace}; -use vectorless::Result; /// A simple in-memory backend with logging. /// @@ -42,7 +42,16 @@ impl StorageBackend for LoggingMemoryBackend { fn get(&self, key: &str) -> Result>> { let data = self.data.read().unwrap(); let result = data.get(key).cloned(); - println!(" [{}] GET '{}' -> {}", self.name, key, if result.is_some() { "found" } else { "not found" }); + println!( + " [{}] GET '{}' -> {}", + self.name, + key, + if result.is_some() { + "found" + } else { + "not found" + } + ); Ok(result) } diff --git a/examples/storage_compression.rs b/examples/storage_compression.rs index 303f582a..80a6e10a 100644 --- a/examples/storage_compression.rs +++ b/examples/storage_compression.rs @@ -14,8 +14,8 @@ //! cargo run --example storage_compression //! ``` -use vectorless::storage::{GzipCodec, IdentityCodec, Codec}; use vectorless::Result; +use vectorless::storage::{Codec, GzipCodec, IdentityCodec}; fn main() -> Result<()> { println!("=== Compression Example ===\n"); @@ -35,8 +35,10 @@ fn main() -> Result<()> { let identity_decoded = identity.decode(&identity_encoded)?; println!(" Encoded size: {} bytes", identity_encoded.len()); - println!(" Compression ratio: {:.1}%", - (identity_encoded.len() as f64 / original.len() as f64) * 100.0); + println!( + " Compression ratio: {:.1}%", + (identity_encoded.len() as f64 / original.len() as f64) * 100.0 + ); assert_eq!(original.to_vec(), identity_decoded); println!(" ✓ Roundtrip verified\n"); @@ -47,10 +49,12 @@ fn main() -> Result<()> { let gzip = GzipCodec::new(level); let compressed = gzip.encode(original)?; - println!(" Level {}: {} bytes ({:.1}% of original)", - level, - compressed.len(), - (compressed.len() as f64 / original.len() as f64) * 100.0); + println!( + " Level {}: {} bytes ({:.1}% of original)", + level, + compressed.len(), + (compressed.len() as f64 / original.len() as f64) * 100.0 + ); } println!(); @@ -62,8 +66,11 @@ fn main() -> Result<()> { let decoded = gzip.decode(&encoded)?; assert_eq!(original.to_vec(), decoded); - println!(" ✓ Encoded {} bytes -> {} bytes", - original.len(), encoded.len()); + println!( + " ✓ Encoded {} bytes -> {} bytes", + original.len(), + encoded.len() + ); println!(" ✓ Decoded back to {} bytes", decoded.len()); println!(" ✓ Data integrity verified\n"); @@ -80,9 +87,11 @@ fn main() -> Result<()> { println!("5. Summary:"); println!(" Original: {} bytes", original.len()); println!(" Identity: {} bytes (100.0%)", identity_encoded.len()); - println!(" Gzip (lvl6): {} bytes ({:.1}%)", - encoded.len(), - (encoded.len() as f64 / original.len() as f64) * 100.0); + println!( + " Gzip (lvl6): {} bytes ({:.1}%)", + encoded.len(), + (encoded.len() as f64 / original.len() as f64) * 100.0 + ); println!(); println!("✓ Compression example complete!"); diff --git a/examples/storage_migration.rs b/examples/storage_migration.rs index 5874046c..af4350bf 100644 --- a/examples/storage_migration.rs +++ b/examples/storage_migration.rs @@ -96,7 +96,10 @@ fn main() -> vectorless::Result<()> { // 3. Migrate from v1 to v3 (multi-step) println!("3. Migrating data from v1 to v3 (via v2):"); let original_data = b"Hello, World!"; - println!(" Original (v1): {:?}", String::from_utf8_lossy(original_data)); + println!( + " Original (v1): {:?}", + String::from_utf8_lossy(original_data) + ); let migrated = migrator.migrate(original_data, 1, 3)?; println!(" Migrated (v3): {:?}", String::from_utf8_lossy(&migrated)); diff --git a/examples/storage_workspace.rs b/examples/storage_workspace.rs index 9f93310c..c0067653 100644 --- a/examples/storage_workspace.rs +++ b/examples/storage_workspace.rs @@ -73,7 +73,10 @@ fn main() -> vectorless::Result<()> { println!(" - Hits: {}", stats.hits); println!(" - Misses: {}", stats.misses); println!(" - Evictions: {}", stats.evictions); - println!(" - Utilization: {:.1}%", workspace.cache_utilization() * 100.0); + println!( + " - Utilization: {:.1}%", + workspace.cache_utilization() * 100.0 + ); println!(); // 7. Load again (should hit cache) diff --git a/src/client/builder.rs b/src/client/builder.rs index 76a335cf..3243bcfa 100644 --- a/src/client/builder.rs +++ b/src/client/builder.rs @@ -194,7 +194,9 @@ impl EngineBuilder { // LLM API key is REQUIRED for retrieval (Pilot needs it for semantic navigation) // Try retrieval config first, then fall back to summary config - let retrieval_api_key = retrieval_config.api_key.clone() + let retrieval_api_key = retrieval_config + .api_key + .clone() .or_else(|| config.summary.api_key.clone()) .ok_or(BuildError::MissingApiKey)?; @@ -207,9 +209,8 @@ impl EngineBuilder { // Configure content aggregator if enabled if retrieval_config.content.enabled { - retriever = retriever.with_content_config( - retrieval_config.content.to_aggregator_config() - ); + retriever = + retriever.with_content_config(retrieval_config.content.to_aggregator_config()); } // Build engine @@ -236,7 +237,9 @@ pub enum BuildError { Workspace(String), /// Missing API key for retrieval. - #[error("Missing API key: LLM API key is required for retrieval. Set OPENAI_API_KEY environment variable or configure retrieval.api_key")] + #[error( + "Missing API key: LLM API key is required for retrieval. Set OPENAI_API_KEY environment variable or configure retrieval.api_key" + )] MissingApiKey, /// Other error. diff --git a/src/client/context.rs b/src/client/context.rs index 344c05cb..3b8a7acd 100644 --- a/src/client/context.rs +++ b/src/client/context.rs @@ -140,9 +140,7 @@ impl ClientContext { /// Check if the request has timed out. pub fn is_timed_out(&self) -> bool { - self.deadline - .map(|d| Instant::now() > d) - .unwrap_or(false) + self.deadline.map(|d| Instant::now() > d).unwrap_or(false) } /// Get remaining time until deadline. @@ -290,8 +288,7 @@ mod tests { #[test] fn test_context_timeout() { - let ctx = ClientContext::new() - .with_timeout(Duration::from_millis(100)); + let ctx = ClientContext::new().with_timeout(Duration::from_millis(100)); assert!(!ctx.is_timed_out()); assert!(ctx.remaining_time().is_some()); diff --git a/src/client/engine.rs b/src/client/engine.rs index 0c0785c4..94745222 100644 --- a/src/client/engine.rs +++ b/src/client/engine.rs @@ -46,10 +46,10 @@ use tracing::info; use crate::config::Config; use crate::error::Result; -use crate::{DocumentTree, Error}; use crate::index::PipelineExecutor; use crate::retrieval::{PipelineRetriever, RetrieveOptions}; use crate::storage::Workspace; +use crate::{DocumentTree, Error}; use super::context::ClientContext; use super::events::EventEmitter; @@ -124,17 +124,15 @@ impl Engine { let events = EventEmitter::new(); // Create indexer client - let indexer = IndexerClient::new(executor) - .with_events(events.clone()); + let indexer = IndexerClient::new(executor).with_events(events.clone()); // Create retriever client - let retriever = RetrieverClient::new(retriever, Arc::clone(&config)) - .with_events(events.clone()); + let retriever = + RetrieverClient::new(retriever, Arc::clone(&config)).with_events(events.clone()); // Create workspace client (if workspace provided) - let workspace_client = workspace.map(|ws| { - WorkspaceClient::new(ws).with_events(events.clone()) - }); + let workspace_client = + workspace.map(|ws| WorkspaceClient::new(ws).with_events(events.clone())); Ok(Self { config, @@ -238,7 +236,10 @@ impl Engine { options.max_tokens = token_budget; } - let mut result = self.retriever.query_with_context(&tree, question, &options, ctx).await?; + let mut result = self + .retriever + .query_with_context(&tree, question, &options, ctx) + .await?; result.doc_id = doc_id.to_string(); Ok(result) @@ -291,10 +292,13 @@ 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()))?; - let doc = workspace.load(doc_id)? + let doc = workspace + .load(doc_id)? .ok_or_else(|| Error::DocumentNotFound(format!("Document not found: {}", doc_id)))?; Ok(doc.tree) @@ -309,10 +313,13 @@ 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()))?; - let doc = workspace.load(doc_id)? + let doc = workspace + .load(doc_id)? .ok_or_else(|| Error::DocumentNotFound(format!("Document not found: {}", doc_id)))?; if doc.pages.is_empty() { @@ -375,7 +382,9 @@ 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()))?; if !workspace.exists(doc_id)? { @@ -392,7 +401,9 @@ 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()))?; workspace.remove(doc_id) @@ -404,7 +415,9 @@ 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()))?; workspace.exists(doc_id) @@ -416,7 +429,9 @@ 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()))?; workspace.get_document_info(doc_id) @@ -430,7 +445,9 @@ 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()))?; workspace.batch_remove(doc_ids) @@ -444,7 +461,9 @@ 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()))?; workspace.clear() diff --git a/src/client/events.rs b/src/client/events.rs index a1d797c4..1195ce70 100644 --- a/src/client/events.rs +++ b/src/client/events.rs @@ -316,8 +316,8 @@ impl Clone for EventEmitter { #[cfg(test)] mod tests { use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; #[test] fn test_event_emitter_index() { diff --git a/src/client/indexer.rs b/src/client/indexer.rs index 8ecb25d4..4462992b 100644 --- a/src/client/indexer.rs +++ b/src/client/indexer.rs @@ -32,7 +32,7 @@ use crate::storage::{DocumentMeta, PersistedDocument}; use super::context::ClientContext; use super::events::{EventEmitter, IndexEvent}; -use super::types::{IndexOptions, IndexMode as ClientIndexMode, IndexedDocument}; +use super::types::{IndexMode as ClientIndexMode, IndexOptions, IndexedDocument}; /// Document indexing client. /// @@ -145,7 +145,8 @@ impl IndexerClient { // Detect format let format = self.detect_format(&path, &options)?; - self.events.emit_index(IndexEvent::FormatDetected { format }); + self.events + .emit_index(IndexEvent::FormatDetected { format }); info!("Indexing {:?} document: {}", format, path.display()); @@ -171,7 +172,9 @@ impl IndexerClient { // Create pipeline input and execute let input = IndexInput::file(&path); let result = { - let mut executor = self.executor.lock() + let mut executor = self + .executor + .lock() .map_err(|_| Error::Other("Pipeline executor lock poisoned".to_string()))?; executor.execute(input, pipeline_options).await? }; @@ -281,10 +284,8 @@ impl IndexerClient { ) .with_description(doc.description.clone().unwrap_or_default()); - let mut persisted = PersistedDocument::new( - meta, - doc.tree.expect("IndexedDocument must have a tree"), - ); + let mut persisted = + PersistedDocument::new(meta, doc.tree.expect("IndexedDocument must have a tree")); for page in doc.pages { persisted.add_page(page.page, &page.content); diff --git a/src/client/mod.rs b/src/client/mod.rs index 51abecd0..33093d8c 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -123,8 +123,8 @@ mod workspace; // Main Types // ============================================================ -pub use engine::Engine; pub use builder::{BuildError, EngineBuilder}; +pub use engine::Engine; // ============================================================ // Sub-Clients @@ -132,8 +132,8 @@ pub use builder::{BuildError, EngineBuilder}; pub use indexer::IndexerClient; pub use retriever::RetrieverClient; -pub use workspace::WorkspaceClient; pub use session::Session; +pub use workspace::WorkspaceClient; // ============================================================ // Context and Events @@ -141,8 +141,7 @@ pub use session::Session; pub use context::{ClientContext, FeatureFlags, RequestContextConfig}; pub use events::{ - EventEmitter, Event, EventHandler, AsyncEventHandler, - IndexEvent, QueryEvent, WorkspaceEvent, + AsyncEventHandler, Event, EventEmitter, EventHandler, IndexEvent, QueryEvent, WorkspaceEvent, }; // ============================================================ @@ -150,16 +149,18 @@ pub use events::{ // ============================================================ pub use types::{ - // Document types - IndexedDocument, PageContent, + // Error types + ClientError, + // Document info + DocumentInfo, // Index types - IndexMode, IndexOptions, + IndexMode, + IndexOptions, + // Document types + IndexedDocument, + PageContent, // Query types QueryResult, - // Document info - DocumentInfo, - // Error types - ClientError, }; // ============================================================ @@ -167,6 +168,6 @@ pub use types::{ // ============================================================ pub use indexer::{IndexerConfig, ValidationResult}; -pub use retriever::{RetrieverClientConfig, NodeContext}; +pub use retriever::{NodeContext, RetrieverClientConfig}; +pub use session::{EvictionPolicy, PreloadStrategy, SessionConfig, SessionStats}; pub use workspace::{WorkspaceClientConfig, WorkspaceStats}; -pub use session::{SessionConfig, SessionStats, EvictionPolicy, PreloadStrategy}; diff --git a/src/client/retriever.rs b/src/client/retriever.rs index ee7a0cbd..f99903f7 100644 --- a/src/client/retriever.rs +++ b/src/client/retriever.rs @@ -26,7 +26,8 @@ use crate::document::{DocumentTree, NodeId}; use crate::error::{Error, Result}; use crate::retrieval::content::ContentAggregatorConfig; use crate::retrieval::{ - QueryComplexity, RetrieveOptions, RetrieveResponse, RetrievalResult, Retriever, SufficiencyLevel, + QueryComplexity, RetrievalResult, RetrieveOptions, RetrieveResponse, Retriever, + SufficiencyLevel, }; use super::context::ClientContext; @@ -129,7 +130,8 @@ impl RetrieverClient { question: &str, options: &RetrieveOptions, ) -> Result { - self.query_with_context(tree, question, options, &ClientContext::new()).await + self.query_with_context(tree, question, options, &ClientContext::new()) + .await } /// Query with request context. @@ -167,7 +169,8 @@ impl RetrieverClient { } // Execute retrieval - let response = self.retriever + let response = self + .retriever .retrieve(tree, question, &options) .await .map_err(|e| Error::Retrieval(e.to_string()))?; @@ -259,11 +262,13 @@ impl RetrieverClient { let similarity = self.calculate_similarity(&target_keywords, &node_keywords); if similarity > 0.3 { - results.push(RetrievalResult::new(&node.title) - .with_node_id(format!("{:?}", current_id)) - .with_content(node.content.clone()) - .with_score(similarity) - .with_depth(tree.depth(current_id))); + results.push( + RetrievalResult::new(&node.title) + .with_node_id(format!("{:?}", current_id)) + .with_content(node.content.clone()) + .with_score(similarity) + .with_depth(tree.depth(current_id)), + ); } } @@ -271,7 +276,11 @@ impl RetrieverClient { } // Sort by score and take top_k - results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); + results.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); results.truncate(top_k); Ok(results) @@ -326,18 +335,22 @@ impl RetrieverClient { } if let Some(node) = tree.get(id) { - ancestors.push(RetrievalResult::new(&node.title) - .with_node_id(format!("{:?}", id)) - .with_depth(tree.depth(id))); + ancestors.push( + RetrievalResult::new(&node.title) + .with_node_id(format!("{:?}", id)) + .with_depth(tree.depth(id)), + ); // Get siblings at this level if let Some(parent_id) = tree.parent(id) { for child_id in tree.children(parent_id) { if child_id != id { if let Some(sibling) = tree.get(child_id) { - siblings.push(RetrievalResult::new(&sibling.title) - .with_node_id(format!("{:?}", child_id)) - .with_depth(tree.depth(child_id))); + siblings.push( + RetrievalResult::new(&sibling.title) + .with_node_id(format!("{:?}", child_id)) + .with_depth(tree.depth(child_id)), + ); } } } @@ -349,14 +362,12 @@ impl RetrieverClient { } // Get the target node - let target = tree - .get(node_id) - .map(|n| { - RetrievalResult::new(&n.title) - .with_node_id(format!("{:?}", node_id)) - .with_content(n.content.clone()) - .with_depth(tree.depth(node_id)) - }); + let target = tree.get(node_id).map(|n| { + RetrievalResult::new(&n.title) + .with_node_id(format!("{:?}", node_id)) + .with_content(n.content.clone()) + .with_depth(tree.depth(node_id)) + }); Ok(NodeContext { target, diff --git a/src/client/session.rs b/src/client/session.rs index f659ac75..0868caaf 100644 --- a/src/client/session.rs +++ b/src/client/session.rs @@ -31,10 +31,10 @@ use std::time::{Duration, Instant}; use tracing::info; use uuid::Uuid; -use crate::{DocumentTree, Error}; use crate::error::Result; use crate::retrieval::RetrieveOptions; use crate::storage::PersistedDocument; +use crate::{DocumentTree, Error}; use super::context::ClientContext; use super::events::EventEmitter; @@ -189,9 +189,8 @@ impl SessionStats { /// Add query time. fn add_query_time(&self, duration: Duration) { - self.total_query_time_us.set( - self.total_query_time_us.get() + duration.as_micros() as u64 - ); + self.total_query_time_us + .set(self.total_query_time_us.get() + duration.as_micros() as u64); } /// Increment cache hits. @@ -294,7 +293,8 @@ impl Session { /// /// Uses the cached tree if available, otherwise loads from workspace. pub async fn query(&self, doc_id: &str, question: &str) -> Result { - self.query_with_options(doc_id, question, RetrieveOptions::default()).await + self.query_with_options(doc_id, question, RetrieveOptions::default()) + .await } /// Query a document with options. @@ -324,7 +324,8 @@ impl Session { /// /// Searches each document and merges results. pub async fn query_all(&self, question: &str) -> Result> { - self.query_all_with_options(question, RetrieveOptions::default()).await + self.query_all_with_options(question, RetrieveOptions::default()) + .await } /// Query across all documents with options. @@ -342,7 +343,10 @@ impl Session { let mut results = Vec::new(); for doc_id in &doc_ids { - match self.query_with_options(doc_id, question, options.clone()).await { + match self + .query_with_options(doc_id, question, options.clone()) + .await + { Ok(result) => { if !result.node_ids.is_empty() { results.push(result); @@ -355,7 +359,11 @@ impl Session { } // Sort by score descending - results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal)); + results.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + }); Ok(results) } @@ -366,7 +374,10 @@ impl Session { /// Get list of documents in this session. pub fn list_documents(&self) -> Vec { - self.documents.values().map(|ctx| ctx.meta.clone()).collect() + self.documents + .values() + .map(|ctx| ctx.meta.clone()) + .collect() } /// Get a document tree (from cache or workspace). @@ -380,7 +391,9 @@ impl Session { self.stats.increment_cache_misses(); // Load from workspace - let doc = self.workspace.load(doc_id)? + let doc = self + .workspace + .load(doc_id)? .ok_or_else(|| Error::DocumentNotFound(format!("Document not found: {}", doc_id)))?; let tree = doc.tree; diff --git a/src/client/types.rs b/src/client/types.rs index 861d52e6..4e3087b9 100644 --- a/src/client/types.rs +++ b/src/client/types.rs @@ -353,8 +353,7 @@ mod tests { #[test] fn test_document_info() { - let info = DocumentInfo::new("doc-1", "Test") - .with_format("markdown"); + let info = DocumentInfo::new("doc-1", "Test").with_format("markdown"); assert_eq!(info.id, "doc-1"); assert_eq!(info.format, "markdown"); diff --git a/src/client/workspace.rs b/src/client/workspace.rs index c5525bfa..1d9575d2 100644 --- a/src/client/workspace.rs +++ b/src/client/workspace.rs @@ -27,7 +27,7 @@ use std::sync::{Arc, RwLock}; use tracing::{debug, info, warn}; -use crate::{Error}; +use crate::Error; use crate::error::Result; use crate::storage::{DocumentMetaEntry, PersistedDocument, Workspace}; @@ -107,7 +107,9 @@ impl WorkspaceClient { let doc_id = doc.meta.id.clone(); { - let mut ws = self.workspace.write() + let mut ws = self + .workspace + .write() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; ws.add(doc)?; } @@ -126,7 +128,9 @@ impl WorkspaceClient { /// /// Returns an error if the workspace read fails. pub fn load(&self, doc_id: &str) -> Result> { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; if !ws.contains(doc_id) { @@ -157,7 +161,9 @@ impl WorkspaceClient { /// Returns an error if the workspace write fails. pub fn remove(&self, doc_id: &str) -> Result { let removed = { - let mut ws = self.workspace.write() + let mut ws = self + .workspace + .write() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; ws.remove(doc_id)? }; @@ -178,7 +184,9 @@ impl WorkspaceClient { /// /// Returns an error if the workspace read fails. pub fn exists(&self, doc_id: &str) -> Result { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; Ok(ws.contains(doc_id)) } @@ -189,10 +197,13 @@ impl WorkspaceClient { /// /// Returns an error if the workspace read fails. pub fn list(&self) -> Result> { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; - Ok(ws.list_documents() + Ok(ws + .list_documents() .iter() .filter_map(|id| ws.get_meta(id)) .map(|meta| DocumentInfo { @@ -212,7 +223,9 @@ impl WorkspaceClient { /// /// Returns an error if the workspace read fails. pub fn get_meta(&self, doc_id: &str) -> Result> { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; Ok(ws.get_meta(doc_id).cloned()) } @@ -223,7 +236,9 @@ impl WorkspaceClient { /// /// Returns an error if the workspace read fails. pub fn get_document_info(&self, doc_id: &str) -> Result> { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; Ok(ws.get_meta(doc_id).map(|meta| DocumentInfo { @@ -247,7 +262,9 @@ impl WorkspaceClient { let mut removed = 0; { - let mut ws = self.workspace.write() + let mut ws = self + .workspace + .write() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; for doc_id in doc_ids { @@ -278,7 +295,9 @@ impl WorkspaceClient { let doc_ids: Vec; { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; doc_ids = ws.list_documents().iter().map(|s| s.to_string()).collect(); } @@ -286,7 +305,9 @@ impl WorkspaceClient { let count = doc_ids.len(); { - let mut ws = self.workspace.write() + let mut ws = self + .workspace + .write() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; for doc_id in &doc_ids { @@ -296,7 +317,8 @@ impl WorkspaceClient { if count > 0 { info!("Cleared workspace: {} documents removed", count); - self.events.emit_workspace(WorkspaceEvent::Cleared { count }); + self.events + .emit_workspace(WorkspaceEvent::Cleared { count }); } Ok(count) @@ -308,7 +330,9 @@ impl WorkspaceClient { /// /// Returns an error if the workspace read fails. pub fn stats(&self) -> Result { - let ws = self.workspace.read() + let ws = self + .workspace + .read() .map_err(|_| Error::Other("Workspace lock poisoned".to_string()))?; Ok(WorkspaceStats { @@ -318,9 +342,7 @@ impl WorkspaceClient { /// Get the number of documents in the workspace. pub fn len(&self) -> usize { - self.workspace.read() - .map(|ws| ws.len()) - .unwrap_or(0) + self.workspace.read().map(|ws| ws.len()).unwrap_or(0) } /// Check if the workspace is empty. @@ -354,8 +376,8 @@ pub struct WorkspaceStats { #[cfg(test)] mod tests { use super::*; - use tempfile::TempDir; use crate::storage::WorkspaceOptions; + use tempfile::TempDir; #[test] fn test_workspace_client_creation() { diff --git a/src/config/docs.rs b/src/config/docs.rs index 7e2330b9..3a0fe062 100644 --- a/src/config/docs.rs +++ b/src/config/docs.rs @@ -43,14 +43,34 @@ impl ConfigDocs { md.push_str("Controls document indexing behavior.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "subsection_threshold", "usize", "300", - "Word count threshold for splitting sections into subsections"); - self.add_row(&mut md, "max_segment_tokens", "usize", "3000", - "Maximum tokens to send in a single segmentation request"); - self.add_row(&mut md, "max_summary_tokens", "usize", "200", - "Maximum tokens for each summary"); - self.add_row(&mut md, "min_summary_tokens", "usize", "20", - "Minimum content tokens required to generate a summary"); + self.add_row( + &mut md, + "subsection_threshold", + "usize", + "300", + "Word count threshold for splitting sections into subsections", + ); + self.add_row( + &mut md, + "max_segment_tokens", + "usize", + "3000", + "Maximum tokens to send in a single segmentation request", + ); + self.add_row( + &mut md, + "max_summary_tokens", + "usize", + "200", + "Maximum tokens for each summary", + ); + self.add_row( + &mut md, + "min_summary_tokens", + "usize", + "20", + "Minimum content tokens required to generate a summary", + ); md.push_str("\n"); // Summary section @@ -58,11 +78,41 @@ impl ConfigDocs { md.push_str("LLM configuration for summary generation.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "model", "string", "gpt-4o-mini", "Model for summarization"); - self.add_row(&mut md, "endpoint", "string", "https://api.openai.com/v1", "API endpoint"); - self.add_row(&mut md, "api_key", "string?", "null", "API key (optional, can use env var)"); - self.add_row(&mut md, "max_tokens", "usize", "200", "Maximum tokens for summary generation"); - self.add_row(&mut md, "temperature", "f32", "0.0", "Temperature for summary generation"); + self.add_row( + &mut md, + "model", + "string", + "gpt-4o-mini", + "Model for summarization", + ); + self.add_row( + &mut md, + "endpoint", + "string", + "https://api.openai.com/v1", + "API endpoint", + ); + self.add_row( + &mut md, + "api_key", + "string?", + "null", + "API key (optional, can use env var)", + ); + self.add_row( + &mut md, + "max_tokens", + "usize", + "200", + "Maximum tokens for summary generation", + ); + self.add_row( + &mut md, + "temperature", + "f32", + "0.0", + "Temperature for summary generation", + ); md.push_str("\n"); // Retrieval section @@ -70,12 +120,48 @@ impl ConfigDocs { md.push_str("Retrieval model and behavior configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "model", "string", "gpt-4o", "Model for retrieval navigation"); - self.add_row(&mut md, "endpoint", "string", "https://api.openai.com/v1", "API endpoint"); - self.add_row(&mut md, "api_key", "string?", "null", "API key (defaults to summary.api_key)"); - self.add_row(&mut md, "top_k", "usize", "3", "Number of top results to return"); - self.add_row(&mut md, "max_tokens", "usize", "1000", "Maximum tokens for retrieval context"); - self.add_row(&mut md, "temperature", "f32", "0.0", "Temperature for retrieval"); + self.add_row( + &mut md, + "model", + "string", + "gpt-4o", + "Model for retrieval navigation", + ); + self.add_row( + &mut md, + "endpoint", + "string", + "https://api.openai.com/v1", + "API endpoint", + ); + self.add_row( + &mut md, + "api_key", + "string?", + "null", + "API key (defaults to summary.api_key)", + ); + self.add_row( + &mut md, + "top_k", + "usize", + "3", + "Number of top results to return", + ); + self.add_row( + &mut md, + "max_tokens", + "usize", + "1000", + "Maximum tokens for retrieval context", + ); + self.add_row( + &mut md, + "temperature", + "f32", + "0.0", + "Temperature for retrieval", + ); md.push_str("\n"); // Retrieval.search section @@ -83,10 +169,34 @@ impl ConfigDocs { md.push_str("Search algorithm configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "top_k", "usize", "5", "Number of top-k results to return"); - self.add_row(&mut md, "beam_width", "usize", "3", "Beam width for multi-path search"); - self.add_row(&mut md, "max_iterations", "usize", "10", "Maximum iterations for search algorithms"); - self.add_row(&mut md, "min_score", "f32", "0.1", "Minimum score to include a path"); + self.add_row( + &mut md, + "top_k", + "usize", + "5", + "Number of top-k results to return", + ); + self.add_row( + &mut md, + "beam_width", + "usize", + "3", + "Beam width for multi-path search", + ); + self.add_row( + &mut md, + "max_iterations", + "usize", + "10", + "Maximum iterations for search algorithms", + ); + self.add_row( + &mut md, + "min_score", + "f32", + "0.1", + "Minimum score to include a path", + ); md.push_str("\n"); // Retrieval.sufficiency section @@ -94,11 +204,41 @@ impl ConfigDocs { md.push_str("Sufficiency checker configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "min_tokens", "usize", "500", "Minimum tokens for sufficiency"); - self.add_row(&mut md, "target_tokens", "usize", "2000", "Target tokens for full sufficiency"); - self.add_row(&mut md, "max_tokens", "usize", "4000", "Maximum tokens before stopping"); - self.add_row(&mut md, "min_content_length", "usize", "200", "Minimum content length (characters)"); - self.add_row(&mut md, "confidence_threshold", "f32", "0.7", "Confidence threshold for LLM judge"); + self.add_row( + &mut md, + "min_tokens", + "usize", + "500", + "Minimum tokens for sufficiency", + ); + self.add_row( + &mut md, + "target_tokens", + "usize", + "2000", + "Target tokens for full sufficiency", + ); + self.add_row( + &mut md, + "max_tokens", + "usize", + "4000", + "Maximum tokens before stopping", + ); + self.add_row( + &mut md, + "min_content_length", + "usize", + "200", + "Minimum content length (characters)", + ); + self.add_row( + &mut md, + "confidence_threshold", + "f32", + "0.7", + "Confidence threshold for LLM judge", + ); md.push_str("\n"); // Retrieval.content section @@ -106,15 +246,69 @@ impl ConfigDocs { md.push_str("Content aggregator configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "enabled", "bool", "true", "Enable content aggregator"); - self.add_row(&mut md, "token_budget", "usize", "4000", "Maximum tokens for aggregated content"); - self.add_row(&mut md, "min_relevance_score", "f32", "0.2", "Minimum relevance score threshold (0.0-1.0)"); - self.add_row(&mut md, "scoring_strategy", "string", "keyword_bm25", "Scoring strategy (keyword_only, keyword_bm25, hybrid)"); - self.add_row(&mut md, "output_format", "string", "markdown", "Output format (markdown, json, tree, flat)"); - self.add_row(&mut md, "include_scores", "bool", "false", "Include relevance scores in output"); - self.add_row(&mut md, "hierarchical_min_per_level", "f32", "0.1", "Minimum budget allocation per depth level"); - self.add_row(&mut md, "deduplicate", "bool", "true", "Enable content deduplication"); - self.add_row(&mut md, "dedup_threshold", "f32", "0.9", "Similarity threshold for deduplication"); + self.add_row( + &mut md, + "enabled", + "bool", + "true", + "Enable content aggregator", + ); + self.add_row( + &mut md, + "token_budget", + "usize", + "4000", + "Maximum tokens for aggregated content", + ); + self.add_row( + &mut md, + "min_relevance_score", + "f32", + "0.2", + "Minimum relevance score threshold (0.0-1.0)", + ); + self.add_row( + &mut md, + "scoring_strategy", + "string", + "keyword_bm25", + "Scoring strategy (keyword_only, keyword_bm25, hybrid)", + ); + self.add_row( + &mut md, + "output_format", + "string", + "markdown", + "Output format (markdown, json, tree, flat)", + ); + self.add_row( + &mut md, + "include_scores", + "bool", + "false", + "Include relevance scores in output", + ); + self.add_row( + &mut md, + "hierarchical_min_per_level", + "f32", + "0.1", + "Minimum budget allocation per depth level", + ); + self.add_row( + &mut md, + "deduplicate", + "bool", + "true", + "Enable content deduplication", + ); + self.add_row( + &mut md, + "dedup_threshold", + "f32", + "0.9", + "Similarity threshold for deduplication", + ); md.push_str("\n"); // Retrieval.strategy section @@ -122,10 +316,34 @@ impl ConfigDocs { md.push_str("Strategy-specific configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "exploration_weight", "f32", "1.414", "MCTS exploration weight (√2)"); - self.add_row(&mut md, "similarity_threshold", "f32", "0.5", "Semantic similarity threshold"); - self.add_row(&mut md, "high_similarity_threshold", "f32", "0.8", "High similarity for 'answer' decision"); - self.add_row(&mut md, "low_similarity_threshold", "f32", "0.3", "Low similarity for 'explore' decision"); + self.add_row( + &mut md, + "exploration_weight", + "f32", + "1.414", + "MCTS exploration weight (√2)", + ); + self.add_row( + &mut md, + "similarity_threshold", + "f32", + "0.5", + "Semantic similarity threshold", + ); + self.add_row( + &mut md, + "high_similarity_threshold", + "f32", + "0.8", + "High similarity for 'answer' decision", + ); + self.add_row( + &mut md, + "low_similarity_threshold", + "f32", + "0.3", + "Low similarity for 'explore' decision", + ); md.push_str("\n"); // Storage section @@ -133,7 +351,13 @@ impl ConfigDocs { md.push_str("Storage configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "workspace_dir", "string", "./workspace", "Workspace directory for persisted documents"); + self.add_row( + &mut md, + "workspace_dir", + "string", + "./workspace", + "Workspace directory for persisted documents", + ); md.push_str("\n"); // Concurrency section @@ -141,10 +365,28 @@ impl ConfigDocs { md.push_str("Concurrency control configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "max_concurrent_requests", "usize", "10", "Maximum concurrent LLM API calls"); - self.add_row(&mut md, "requests_per_minute", "usize", "500", "Rate limit: requests per minute"); + self.add_row( + &mut md, + "max_concurrent_requests", + "usize", + "10", + "Maximum concurrent LLM API calls", + ); + self.add_row( + &mut md, + "requests_per_minute", + "usize", + "500", + "Rate limit: requests per minute", + ); self.add_row(&mut md, "enabled", "bool", "true", "Enable rate limiting"); - self.add_row(&mut md, "semaphore_enabled", "bool", "true", "Enable semaphore-based concurrency"); + self.add_row( + &mut md, + "semaphore_enabled", + "bool", + "true", + "Enable semaphore-based concurrency", + ); md.push_str("\n"); // Fallback section @@ -152,26 +394,68 @@ impl ConfigDocs { md.push_str("Fallback/error recovery configuration.\n\n"); md.push_str("| Option | Type | Default | Description |\n"); md.push_str("|--------|------|---------|-------------|\n"); - self.add_row(&mut md, "enabled", "bool", "true", "Enable graceful degradation"); - self.add_row(&mut md, "models", "[string]", "[\"gpt-4o-mini\", \"glm-4-flash\"]", "Fallback models in priority order"); - self.add_row(&mut md, "endpoints", "[string]", "[]", "Fallback endpoints in priority order"); - self.add_row(&mut md, "on_rate_limit", "string", "retry_then_fallback", "Behavior on rate limit (retry, fallback, retry_then_fallback, fail)"); - self.add_row(&mut md, "on_timeout", "string", "retry_then_fallback", "Behavior on timeout"); - self.add_row(&mut md, "on_all_failed", "string", "return_error", "Behavior when all attempts fail (return_error, return_cache)"); + self.add_row( + &mut md, + "enabled", + "bool", + "true", + "Enable graceful degradation", + ); + self.add_row( + &mut md, + "models", + "[string]", + "[\"gpt-4o-mini\", \"glm-4-flash\"]", + "Fallback models in priority order", + ); + self.add_row( + &mut md, + "endpoints", + "[string]", + "[]", + "Fallback endpoints in priority order", + ); + self.add_row( + &mut md, + "on_rate_limit", + "string", + "retry_then_fallback", + "Behavior on rate limit (retry, fallback, retry_then_fallback, fail)", + ); + self.add_row( + &mut md, + "on_timeout", + "string", + "retry_then_fallback", + "Behavior on timeout", + ); + self.add_row( + &mut md, + "on_all_failed", + "string", + "return_error", + "Behavior when all attempts fail (return_error, return_cache)", + ); md.push_str("\n"); md } fn add_row(&self, md: &mut String, name: &str, ty: &str, default: &str, desc: &str) { - md.push_str(&format!("| `{}` | {} | {} | {} |\n", name, ty, default, desc)); + md.push_str(&format!( + "| `{}` | {} | {} | {} |\n", + name, ty, default, desc + )); } /// Generate an example TOML file with all options. pub fn to_example_toml(&self) -> String { toml::to_string_pretty(&self.config).unwrap_or_else(|e| { - format!("# Error generating TOML: {}\n\n# Using default config\n{}", - e, Self::fallback_toml()) + format!( + "# Error generating TOML: {}\n\n# Using default config\n{}", + e, + Self::fallback_toml() + ) }) } @@ -249,7 +533,8 @@ models = ["gpt-4o-mini", "glm-4-flash"] on_rate_limit = "retry_then_fallback" on_timeout = "retry_then_fallback" on_all_failed = "return_error" -"#.to_string() +"# + .to_string() } /// Generate a minimal example TOML file. @@ -262,7 +547,8 @@ api_key = "your-api-key-here" [retrieval] top_k = 5 -"#.to_string() +"# + .to_string() } } diff --git a/src/config/loader.rs b/src/config/loader.rs index fe2c6736..99aef12f 100644 --- a/src/config/loader.rs +++ b/src/config/loader.rs @@ -211,7 +211,12 @@ impl ConfigLoader { } /// Set a configuration value by path. - fn set_by_path(&self, config: &mut Config, path: &[&str], value: &str) -> Result<(), ConfigError> { + fn set_by_path( + &self, + config: &mut Config, + path: &[&str], + value: &str, + ) -> Result<(), ConfigError> { match path { ["summary", "api_key"] => { config.summary.api_key = Some(value.to_string()); @@ -223,9 +228,9 @@ impl ConfigLoader { config.summary.endpoint = value.to_string(); } ["summary", "max_tokens"] => { - config.summary.max_tokens = value.parse().map_err(|e| { - ConfigError::Env(format!("Invalid max_tokens: {}", e)) - })?; + config.summary.max_tokens = value + .parse() + .map_err(|e| ConfigError::Env(format!("Invalid max_tokens: {}", e)))?; } ["retrieval", "api_key"] => { config.retrieval.api_key = Some(value.to_string()); @@ -237,9 +242,9 @@ impl ConfigLoader { config.retrieval.endpoint = value.to_string(); } ["retrieval", "top_k"] => { - config.retrieval.top_k = value.parse().map_err(|e| { - ConfigError::Env(format!("Invalid top_k: {}", e)) - })?; + config.retrieval.top_k = value + .parse() + .map_err(|e| ConfigError::Env(format!("Invalid top_k: {}", e)))?; } ["storage", "workspace_dir"] => { config.storage.workspace_dir = PathBuf::from(value); @@ -259,8 +264,7 @@ impl ConfigLoader { } /// Default configuration file names to search for. -pub const CONFIG_FILE_NAMES: &[&str] = - &["vectorless.toml", "config.toml", ".vectorless.toml"]; +pub const CONFIG_FILE_NAMES: &[&str] = &["vectorless.toml", "config.toml", ".vectorless.toml"]; /// Find a configuration file in current or parent directories. pub fn find_config_file() -> Option { @@ -317,9 +321,7 @@ mod tests { #[test] fn test_config_loader_not_found() { - let result = ConfigLoader::new() - .file("nonexistent_config.toml") - .load(); + let result = ConfigLoader::new().file("nonexistent_config.toml").load(); assert!(result.is_err()); assert!(matches!(result.unwrap_err(), ConfigError::NotFound(_))); @@ -327,10 +329,7 @@ mod tests { #[test] fn test_config_loader_with_validation() { - let config = ConfigLoader::new() - .with_validation(true) - .load() - .unwrap(); + let config = ConfigLoader::new().with_validation(true).load().unwrap(); assert_eq!(config.retrieval.model, "gpt-4o"); } diff --git a/src/config/merge.rs b/src/config/merge.rs index 438872b5..c6d995a0 100644 --- a/src/config/merge.rs +++ b/src/config/merge.rs @@ -7,9 +7,8 @@ //! enabling layered configuration from multiple sources. use super::types::{ - CacheConfig, Config, ConcurrencyConfig, ContentAggregatorConfig, FallbackConfig, - IndexerConfig, RetrievalConfig, SearchConfig, StorageConfig, StrategyConfig, SufficiencyConfig, - SummaryConfig, + CacheConfig, ConcurrencyConfig, Config, ContentAggregatorConfig, FallbackConfig, IndexerConfig, + RetrievalConfig, SearchConfig, StorageConfig, StrategyConfig, SufficiencyConfig, SummaryConfig, }; /// Configuration merge strategy. @@ -138,7 +137,8 @@ impl Merge for SufficiencyConfig { if strategy == MergeStrategy::Replace || self.min_content_length == 200 { self.min_content_length = other.min_content_length; } - if strategy == MergeStrategy::Replace || (self.confidence_threshold - 0.7).abs() < f32::EPSILON + if strategy == MergeStrategy::Replace + || (self.confidence_threshold - 0.7).abs() < f32::EPSILON { self.confidence_threshold = other.confidence_threshold; } @@ -158,12 +158,11 @@ impl Merge for CacheConfig { impl Merge for StrategyConfig { fn merge(&mut self, other: &Self, strategy: MergeStrategy) { - if strategy == MergeStrategy::Replace - || (self.exploration_weight - 1.414).abs() < 0.001 - { + if strategy == MergeStrategy::Replace || (self.exploration_weight - 1.414).abs() < 0.001 { self.exploration_weight = other.exploration_weight; } - if strategy == MergeStrategy::Replace || (self.similarity_threshold - 0.5).abs() < f32::EPSILON + if strategy == MergeStrategy::Replace + || (self.similarity_threshold - 0.5).abs() < f32::EPSILON { self.similarity_threshold = other.similarity_threshold; } @@ -188,7 +187,8 @@ impl Merge for ContentAggregatorConfig { if strategy == MergeStrategy::Replace || self.token_budget == 4000 { self.token_budget = other.token_budget; } - if strategy == MergeStrategy::Replace || (self.min_relevance_score - 0.2).abs() < f32::EPSILON + if strategy == MergeStrategy::Replace + || (self.min_relevance_score - 0.2).abs() < f32::EPSILON { self.min_relevance_score = other.min_relevance_score; } diff --git a/src/config/mod.rs b/src/config/mod.rs index d821332a..4de3984b 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -75,27 +75,37 @@ mod validator; // Re-export main types pub use docs::ConfigDocs; -pub use loader::{find_config_file, ConfigError, ConfigLoader, CONFIG_FILE_NAMES}; +pub use loader::{CONFIG_FILE_NAMES, ConfigError, ConfigLoader, find_config_file}; pub use merge::{ConfigOverlay, Merge, MergeStrategy}; pub use types::{ + CacheConfig, + CompressionAlgorithm, + CompressionConfig, + // Concurrency + ConcurrencyConfig, // Main config Config, + // Validation + ConfigValidationError, + // Content aggregator + ContentAggregatorConfig, + // Fallback + FallbackBehavior, + FallbackConfig, // Indexer IndexerConfig, // LLM configs - LlmConfig, SummaryConfig, + LlmConfig, + OnAllFailedBehavior, // Retrieval configs - RetrievalConfig, SearchConfig, + RetrievalConfig, + SearchConfig, + Severity, // Storage and sufficiency - StorageConfig, CompressionAlgorithm, CompressionConfig, - CacheConfig, StrategyConfig, SufficiencyConfig, - // Content aggregator - ContentAggregatorConfig, - // Concurrency - ConcurrencyConfig, - // Fallback - FallbackBehavior, FallbackConfig, OnAllFailedBehavior, - // Validation - ConfigValidationError, ValidationError, Severity, + StorageConfig, + StrategyConfig, + SufficiencyConfig, + SummaryConfig, + ValidationError, }; pub use validator::{ConfigValidator, ValidationRule}; diff --git a/src/config/types/fallback.rs b/src/config/types/fallback.rs index fa199b30..96c2fd89 100644 --- a/src/config/types/fallback.rs +++ b/src/config/types/fallback.rs @@ -181,8 +181,8 @@ impl FallbackConfig { let delay_ms = if attempt == 0 { self.initial_retry_delay_ms } else { - let delay = self.initial_retry_delay_ms as f32 - * self.retry_multiplier.powi(attempt as i32); + let delay = + self.initial_retry_delay_ms as f32 * self.retry_multiplier.powi(attempt as i32); delay.min(self.max_retry_delay_ms as f32) as u64 }; std::time::Duration::from_millis(delay_ms) diff --git a/src/config/types/mod.rs b/src/config/types/mod.rs index 39536156..78b6c37e 100644 --- a/src/config/types/mod.rs +++ b/src/config/types/mod.rs @@ -6,8 +6,8 @@ //! All configuration values are defined inline in `Default` trait implementations. //! Configuration is loaded from TOML files only - no environment variable magic. -mod content; mod concurrency; +mod content; mod fallback; mod indexer; mod llm; @@ -16,15 +16,15 @@ mod storage; use serde::{Deserialize, Serialize}; -pub use content::ContentAggregatorConfig; pub use concurrency::ConcurrencyConfig; +pub use content::ContentAggregatorConfig; pub use fallback::{FallbackBehavior, FallbackConfig, OnAllFailedBehavior}; pub use indexer::IndexerConfig; pub use llm::{LlmConfig, SummaryConfig}; pub use retrieval::{RetrievalConfig, SearchConfig}; pub use storage::{ - CacheConfig, CompressionAlgorithm, CompressionConfig, - StorageConfig, StrategyConfig, SufficiencyConfig, + CacheConfig, CompressionAlgorithm, CompressionConfig, StorageConfig, StrategyConfig, + SufficiencyConfig, }; /// Main configuration for vectorless. @@ -139,10 +139,13 @@ impl Config { } if self.retrieval.temperature < 0.0 || self.retrieval.temperature > 2.0 { - errors.push(ValidationError::warning( - "retrieval.temperature", - "Temperature outside typical range [0.0, 2.0]", - ).with_actual(self.retrieval.temperature.to_string())); + errors.push( + ValidationError::warning( + "retrieval.temperature", + "Temperature outside typical range [0.0, 2.0]", + ) + .with_actual(self.retrieval.temperature.to_string()), + ); } // Validate content aggregator @@ -156,12 +159,14 @@ impl Config { if self.retrieval.content.min_relevance_score < 0.0 || self.retrieval.content.min_relevance_score > 1.0 { - errors.push(ValidationError::error( - "retrieval.content.min_relevance_score", - "Min relevance score must be between 0.0 and 1.0", - ) - .with_expected("0.0 - 1.0") - .with_actual(self.retrieval.content.min_relevance_score.to_string())); + errors.push( + ValidationError::error( + "retrieval.content.min_relevance_score", + "Min relevance score must be between 0.0 and 1.0", + ) + .with_expected("0.0 - 1.0") + .with_actual(self.retrieval.content.min_relevance_score.to_string()), + ); } // Validate concurrency diff --git a/src/config/validator.rs b/src/config/validator.rs index 8a3596fd..c4000764 100644 --- a/src/config/validator.rs +++ b/src/config/validator.rs @@ -73,10 +73,13 @@ impl ValidationRule for RangeValidator { } if config.indexer.subsection_threshold > 10000 { - errors.push(ValidationError::warning( - "indexer.subsection_threshold", - "Subsection threshold is very high, may impact performance", - ).with_actual(config.indexer.subsection_threshold.to_string())); + errors.push( + ValidationError::warning( + "indexer.subsection_threshold", + "Subsection threshold is very high, may impact performance", + ) + .with_actual(config.indexer.subsection_threshold.to_string()), + ); } // Summary ranges @@ -88,10 +91,13 @@ impl ValidationRule for RangeValidator { } if config.summary.temperature < 0.0 || config.summary.temperature > 2.0 { - errors.push(ValidationError::warning( - "summary.temperature", - "Temperature outside typical range [0.0, 2.0]", - ).with_actual(config.summary.temperature.to_string())); + errors.push( + ValidationError::warning( + "summary.temperature", + "Temperature outside typical range [0.0, 2.0]", + ) + .with_actual(config.summary.temperature.to_string()), + ); } // Retrieval ranges @@ -120,12 +126,14 @@ impl ValidationRule for RangeValidator { if config.retrieval.content.min_relevance_score < 0.0 || config.retrieval.content.min_relevance_score > 1.0 { - errors.push(ValidationError::error( - "retrieval.content.min_relevance_score", - "Min relevance score must be between 0.0 and 1.0", - ) - .with_expected("0.0 - 1.0") - .with_actual(config.retrieval.content.min_relevance_score.to_string())); + errors.push( + ValidationError::error( + "retrieval.content.min_relevance_score", + "Min relevance score must be between 0.0 and 1.0", + ) + .with_expected("0.0 - 1.0") + .with_actual(config.retrieval.content.min_relevance_score.to_string()), + ); } if config.retrieval.content.hierarchical_min_per_level < 0.0 @@ -170,61 +178,71 @@ impl ValidationRule for ConsistencyValidator { fn validate(&self, config: &Config, errors: &mut Vec) { // Check if summary tokens are reasonable if config.summary.max_tokens > config.indexer.max_segment_tokens { - errors.push(ValidationError::warning( - "summary.max_tokens", - "Summary max tokens exceeds max segment tokens", - ) - .with_expected(format!("<= {}", config.indexer.max_segment_tokens)) - .with_actual(config.summary.max_tokens.to_string())); + errors.push( + ValidationError::warning( + "summary.max_tokens", + "Summary max tokens exceeds max segment tokens", + ) + .with_expected(format!("<= {}", config.indexer.max_segment_tokens)) + .with_actual(config.summary.max_tokens.to_string()), + ); } // Check if content token budget is reasonable if config.retrieval.content.token_budget > 100000 { - errors.push(ValidationError::warning( - "retrieval.content.token_budget", - "Token budget is very high, may cause performance issues", - ).with_actual(config.retrieval.content.token_budget.to_string())); + errors.push( + ValidationError::warning( + "retrieval.content.token_budget", + "Token budget is very high, may cause performance issues", + ) + .with_actual(config.retrieval.content.token_budget.to_string()), + ); } // Check if sufficiency thresholds are consistent if config.retrieval.sufficiency.min_tokens > config.retrieval.sufficiency.target_tokens { - errors.push(ValidationError::error( - "retrieval.sufficiency.min_tokens", - "Min tokens cannot exceed target tokens", - ) - .with_expected(format!("<= {}", config.retrieval.sufficiency.target_tokens)) - .with_actual(config.retrieval.sufficiency.min_tokens.to_string())); + errors.push( + ValidationError::error( + "retrieval.sufficiency.min_tokens", + "Min tokens cannot exceed target tokens", + ) + .with_expected(format!("<= {}", config.retrieval.sufficiency.target_tokens)) + .with_actual(config.retrieval.sufficiency.min_tokens.to_string()), + ); } if config.retrieval.sufficiency.target_tokens > config.retrieval.sufficiency.max_tokens { - errors.push(ValidationError::error( - "retrieval.sufficiency.target_tokens", - "Target tokens cannot exceed max tokens", - ) - .with_expected(format!("<= {}", config.retrieval.sufficiency.max_tokens)) - .with_actual(config.retrieval.sufficiency.target_tokens.to_string())); + errors.push( + ValidationError::error( + "retrieval.sufficiency.target_tokens", + "Target tokens cannot exceed max tokens", + ) + .with_expected(format!("<= {}", config.retrieval.sufficiency.max_tokens)) + .with_actual(config.retrieval.sufficiency.target_tokens.to_string()), + ); } // Check scoring strategy validity let valid_strategies = ["keyword_only", "keyword_bm25", "hybrid"]; if !valid_strategies.contains(&config.retrieval.content.scoring_strategy.as_str()) { - errors.push(ValidationError::error( - "retrieval.content.scoring_strategy", - "Invalid scoring strategy", - ) - .with_expected(format!("one of: {:?}", valid_strategies)) - .with_actual(config.retrieval.content.scoring_strategy.clone())); + errors.push( + ValidationError::error( + "retrieval.content.scoring_strategy", + "Invalid scoring strategy", + ) + .with_expected(format!("one of: {:?}", valid_strategies)) + .with_actual(config.retrieval.content.scoring_strategy.clone()), + ); } // Check output format validity let valid_formats = ["markdown", "json", "tree", "flat"]; if !valid_formats.contains(&config.retrieval.content.output_format.as_str()) { - errors.push(ValidationError::error( - "retrieval.content.output_format", - "Invalid output format", - ) - .with_expected(format!("one of: {:?}", valid_formats)) - .with_actual(config.retrieval.content.output_format.clone())); + errors.push( + ValidationError::error("retrieval.content.output_format", "Invalid output format") + .with_expected(format!("one of: {:?}", valid_formats)) + .with_actual(config.retrieval.content.output_format.clone()), + ); } } } @@ -278,25 +296,36 @@ impl ValidationRule for DependencyValidator { // Check strategy configuration if config.retrieval.strategy.exploration_weight <= 0.0 { - errors.push(ValidationError::error( - "retrieval.strategy.exploration_weight", - "Exploration weight must be positive", - ).with_actual(config.retrieval.strategy.exploration_weight.to_string())); + errors.push( + ValidationError::error( + "retrieval.strategy.exploration_weight", + "Exploration weight must be positive", + ) + .with_actual(config.retrieval.strategy.exploration_weight.to_string()), + ); } // Check similarity thresholds are ordered correctly if config.retrieval.strategy.low_similarity_threshold >= config.retrieval.strategy.high_similarity_threshold { - errors.push(ValidationError::error( - "retrieval.strategy.low_similarity_threshold", - "Low similarity threshold must be less than high similarity threshold", - ) - .with_expected(format!( - "< {}", - config.retrieval.strategy.high_similarity_threshold - )) - .with_actual(config.retrieval.strategy.low_similarity_threshold.to_string())); + errors.push( + ValidationError::error( + "retrieval.strategy.low_similarity_threshold", + "Low similarity threshold must be less than high similarity threshold", + ) + .with_expected(format!( + "< {}", + config.retrieval.strategy.high_similarity_threshold + )) + .with_actual( + config + .retrieval + .strategy + .low_similarity_threshold + .to_string(), + ), + ); } } } @@ -353,7 +382,11 @@ mod tests { // Should succeed but with warnings if let Err(err) = result { - assert!(err.errors.iter().any(|e| e.path.contains("fallback.models"))); + assert!( + err.errors + .iter() + .any(|e| e.path.contains("fallback.models")) + ); } } } diff --git a/src/document/tree.rs b/src/document/tree.rs index 090dacae..88e5ffd4 100644 --- a/src/document/tree.rs +++ b/src/document/tree.rs @@ -119,9 +119,7 @@ impl RetrievalIndex { } } // Sort by start page for consistent ordering - result.sort_by_key(|&id| { - self.node_page_range.get(&id).map(|(s, _)| *s).unwrap_or(0) - }); + result.sort_by_key(|&id| self.node_page_range.get(&id).map(|(s, _)| *s).unwrap_or(0)); result } diff --git a/src/error.rs b/src/error.rs index 615dd671..a5caacad 100644 --- a/src/error.rs +++ b/src/error.rs @@ -14,7 +14,6 @@ pub enum Error { // ========================================================================= // Document & Parsing Errors // ========================================================================= - /// An error occurred while parsing a document. #[error("Document parsing error: {0}")] Parse(String), @@ -30,7 +29,6 @@ pub enum Error { // ========================================================================= // Index Errors // ========================================================================= - /// An error occurred while building the index. #[error("Index building error: {0}")] IndexBuild(String), @@ -46,7 +44,6 @@ pub enum Error { // ========================================================================= // Retrieval Errors // ========================================================================= - /// An error occurred during retrieval. #[error("Retrieval error: {0}")] Retrieval(String), @@ -62,7 +59,6 @@ pub enum Error { // ========================================================================= // LLM Errors // ========================================================================= - /// An error occurred during LLM call. #[error("LLM error: {0}")] Llm(String), @@ -78,7 +74,6 @@ pub enum Error { // ========================================================================= // Summary Errors // ========================================================================= - /// An error occurred during summarization. #[error("Summarization error: {0}")] Summarization(String), @@ -90,7 +85,6 @@ pub enum Error { // ========================================================================= // Storage Errors // ========================================================================= - /// An error occurred during I/O operations. #[error("IO error: {0}")] Io(#[from] std::io::Error), @@ -126,7 +120,6 @@ pub enum Error { // ========================================================================= // Configuration Errors // ========================================================================= - /// TOML parsing error. #[error("TOML parsing error: {0}")] Toml(String), @@ -142,7 +135,6 @@ pub enum Error { // ========================================================================= // Node Errors // ========================================================================= - /// The requested node was not found. #[error("Node not found: {0}")] NodeNotFound(String), @@ -150,7 +142,6 @@ pub enum Error { // ========================================================================= // Input Validation Errors // ========================================================================= - /// Invalid input. #[error("Invalid input: {0}")] InvalidInput(String), @@ -178,7 +169,6 @@ pub enum Error { // ========================================================================= // Throttle Errors // ========================================================================= - /// Throttle error. #[error("Throttle error: {0}")] Throttle(String), @@ -190,7 +180,6 @@ pub enum Error { // ========================================================================= // Timeout Errors // ========================================================================= - /// Operation timeout. #[error("Operation timeout: {0}")] Timeout(String), @@ -198,7 +187,6 @@ pub enum Error { // ========================================================================= // Generic Errors // ========================================================================= - /// A generic error with a message. #[error("{0}")] Other(String), @@ -229,10 +217,7 @@ impl Error { pub fn is_retryable(&self) -> bool { matches!( self, - Self::RateLimitExceeded(_) - | Self::SearchTimeout(_) - | Self::Timeout(_) - | Self::Llm(_) + Self::RateLimitExceeded(_) | Self::SearchTimeout(_) | Self::Timeout(_) | Self::Llm(_) ) } diff --git a/src/index/config.rs b/src/index/config.rs index 55128822..f08b5968 100644 --- a/src/index/config.rs +++ b/src/index/config.rs @@ -9,8 +9,8 @@ //! - [`OptimizationConfig`] - Tree optimization settings //! - [`ThinningConfig`] - Node merging settings -use crate::config::{ConcurrencyConfig, IndexerConfig}; use super::summary::SummaryStrategy; +use crate::config::{ConcurrencyConfig, IndexerConfig}; /// Index mode for document processing. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/src/index/mod.rs b/src/index/mod.rs index 96de34a5..51a18ec5 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -48,9 +48,7 @@ pub use pipeline::{ }; // Re-export config types -pub use config::{ - IndexMode, OptimizationConfig, PipelineOptions, ThinningConfig, -}; +pub use config::{IndexMode, OptimizationConfig, PipelineOptions, ThinningConfig}; // Re-export stages pub use stages::IndexStage; diff --git a/src/index/stages/enhance.rs b/src/index/stages/enhance.rs index d1d0f6fd..278b572c 100644 --- a/src/index/stages/enhance.rs +++ b/src/index/stages/enhance.rs @@ -8,9 +8,8 @@ use std::sync::Arc; use std::time::Instant; use tracing::{info, warn}; - -use crate::error::Result; use crate::document::{DocumentTree, NodeId}; +use crate::error::Result; use crate::llm::LlmClient; use super::{IndexStage, StageResult}; diff --git a/src/index/stages/optimize.rs b/src/index/stages/optimize.rs index 571e947d..8ae8b44f 100644 --- a/src/index/stages/optimize.rs +++ b/src/index/stages/optimize.rs @@ -7,9 +7,8 @@ use super::async_trait; use std::time::Instant; use tracing::info; - +use crate::document::NodeId; use crate::error::Result; -use crate::document::{NodeId}; use crate::index::pipeline::IndexContext; use super::{IndexStage, StageResult}; diff --git a/src/index/stages/parse.rs b/src/index/stages/parse.rs index 150d1803..72657b9b 100644 --- a/src/index/stages/parse.rs +++ b/src/index/stages/parse.rs @@ -34,9 +34,8 @@ impl ParseStage { 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::Error::Parse(format!("Unknown format: {}", ext)) - }) + DocumentFormat::from_extension(ext) + .ok_or_else(|| crate::Error::Parse(format!("Unknown format: {}", ext))) } IndexInput::Content { format, .. } => Ok(*format), }, diff --git a/src/lib.rs b/src/lib.rs index cd3dc7a9..00818c67 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -128,8 +128,8 @@ pub use error::{Error, Result}; // Document types pub use document::{ - DocumentStructure, DocumentTree, NodeId, StructureNode, TocConfig, TocEntry, - TocNode, TocView, TreeNode, + DocumentStructure, DocumentTree, NodeId, StructureNode, TocConfig, TocEntry, TocNode, TocView, + TreeNode, }; // Utility functions @@ -163,7 +163,9 @@ pub use retrieval::{ }; // Storage -pub use storage::{AsyncWorkspace, DocumentMeta as StorageDocumentMeta, PersistedDocument, Workspace}; +pub use storage::{ + AsyncWorkspace, DocumentMeta as StorageDocumentMeta, PersistedDocument, Workspace, +}; // Throttle pub use throttle::{ConcurrencyConfig, ConcurrencyController, RateLimiter}; diff --git a/src/parser/docx/parser.rs b/src/parser/docx/parser.rs index 15d593c8..aa3885f6 100644 --- a/src/parser/docx/parser.rs +++ b/src/parser/docx/parser.rs @@ -32,7 +32,7 @@ use std::path::Path; use async_trait::async_trait; use zip::ZipArchive; -use crate::{Error}; +use crate::Error; use crate::error::Result; use crate::parser::{DocumentFormat, DocumentMeta, DocumentParser, ParseResult, RawNode}; diff --git a/src/parser/markdown/parser.rs b/src/parser/markdown/parser.rs index cc6df8a1..09fc8888 100644 --- a/src/parser/markdown/parser.rs +++ b/src/parser/markdown/parser.rs @@ -8,8 +8,8 @@ use pulldown_cmark::Options; use std::path::Path; use crate::error::Result; -use crate::util::estimate_tokens; use crate::parser::{DocumentFormat, DocumentMeta, DocumentParser, ParseResult, RawNode}; +use crate::util::estimate_tokens; use super::config::MarkdownConfig; use super::frontmatter; diff --git a/src/parser/pdf/parser.rs b/src/parser/pdf/parser.rs index a96bf0c2..10fe053e 100644 --- a/src/parser/pdf/parser.rs +++ b/src/parser/pdf/parser.rs @@ -8,7 +8,7 @@ use std::path::Path; use lopdf::Document as LopdfDocument; use tracing::{info, warn}; -use crate::{Error}; +use crate::Error; use crate::error::Result; use crate::parser::DocumentParser; use crate::parser::toc::TocProcessor; diff --git a/src/parser/registry.rs b/src/parser/registry.rs index ae632e4c..9667c176 100644 --- a/src/parser/registry.rs +++ b/src/parser/registry.rs @@ -11,7 +11,7 @@ use std::collections::HashMap; use std::path::Path; use std::sync::{Arc, RwLock}; -use crate::{Error}; +use crate::Error; use crate::error::Result; use crate::parser::{DocumentFormat, DocumentParser, MarkdownParser, ParseResult, PdfParser}; diff --git a/src/retrieval/content/aggregator.rs b/src/retrieval/content/aggregator.rs index 87a8f20e..0fbcbf37 100644 --- a/src/retrieval/content/aggregator.rs +++ b/src/retrieval/content/aggregator.rs @@ -35,7 +35,11 @@ impl CandidateNode { /// Create a new candidate. #[must_use] pub fn new(node_id: NodeId, score: f32, depth: usize) -> Self { - Self { node_id, score, depth } + Self { + node_id, + score, + depth, + } } } @@ -149,8 +153,7 @@ impl ContentAggregator { // Step 3: Allocate token budget let max_depth = filtered.iter().map(|r| r.chunk.depth).max().unwrap_or(0); let strategy = self.get_allocation_strategy(); - let allocator = BudgetAllocator::new(self.config.token_budget) - .with_strategy(strategy); + let allocator = BudgetAllocator::new(self.config.token_budget).with_strategy(strategy); let allocation = allocator.allocate(filtered, max_depth); @@ -162,10 +165,8 @@ impl ContentAggregator { ); // Step 4: Build structured output - let builder = StructureBuilder::from_config( - self.config.output_format, - self.config.include_scores, - ); + let builder = + StructureBuilder::from_config(self.config.output_format, self.config.include_scores); let structured = builder.build(allocation.selected.clone(), tree); diff --git a/src/retrieval/content/budget.rs b/src/retrieval/content/budget.rs index 1b4ed279..4c867e4f 100644 --- a/src/retrieval/content/budget.rs +++ b/src/retrieval/content/budget.rs @@ -223,7 +223,9 @@ impl BudgetAllocator { let remaining = self.total_budget - tokens_used; if remaining >= 50 { // Minimum useful content - if let Some(truncated) = self.truncate_content(&relevance.chunk.content, remaining) { + if let Some(truncated) = + self.truncate_content(&relevance.chunk.content, remaining) + { let truncated_tokens = estimate_tokens(&truncated); selected.push(SelectedContent { node_id: relevance.chunk.node_id, @@ -305,7 +307,9 @@ impl BudgetAllocator { // Truncate to allocated budget let remaining = self.total_budget - tokens_used; if remaining >= 50 && remaining >= allocated_budget / 2 { - if let Some(truncated) = self.truncate_content(&relevance.chunk.content, remaining.min(allocated_budget)) { + if let Some(truncated) = self + .truncate_content(&relevance.chunk.content, remaining.min(allocated_budget)) + { let truncated_tokens = estimate_tokens(&truncated); let truncated_len = truncated.len(); selected.push(SelectedContent { @@ -354,10 +358,7 @@ impl BudgetAllocator { // Group content by depth let mut by_depth: HashMap> = HashMap::new(); for c in content { - by_depth - .entry(c.chunk.depth) - .or_default() - .push(c); + by_depth.entry(c.chunk.depth).or_default().push(c); } // Sort each level by score @@ -407,9 +408,12 @@ impl BudgetAllocator { level_used += tokens; } else if level_used < per_level_budget { // Try truncated version - let remaining = (self.total_budget - tokens_used).min(per_level_budget - level_used); + let remaining = + (self.total_budget - tokens_used).min(per_level_budget - level_used); if remaining >= 50 { - if let Some(truncated) = self.truncate_content(&relevance.chunk.content, remaining) { + if let Some(truncated) = + self.truncate_content(&relevance.chunk.content, remaining) + { let truncated_tokens = estimate_tokens(&truncated); selected.push(SelectedContent { node_id: relevance.chunk.node_id, @@ -562,8 +566,7 @@ mod tests { #[test] fn test_greedy_allocation() { - let allocator = BudgetAllocator::new(100) - .with_strategy(AllocationStrategy::Greedy); + let allocator = BudgetAllocator::new(100).with_strategy(AllocationStrategy::Greedy); let content = vec![ make_relevance("High score content with enough text", 0.9, 0), @@ -577,8 +580,7 @@ mod tests { #[test] fn test_min_score_filter() { - let allocator = BudgetAllocator::new(1000) - .with_min_score(0.5); + let allocator = BudgetAllocator::new(1000).with_min_score(0.5); let content = vec![ make_relevance("Good content", 0.8, 0), diff --git a/src/retrieval/content/builder.rs b/src/retrieval/content/builder.rs index e0248e7b..93b85219 100644 --- a/src/retrieval/content/builder.rs +++ b/src/retrieval/content/builder.rs @@ -182,11 +182,7 @@ impl StructureBuilder { /// Build structured content from selected items. #[must_use] - pub fn build( - &self, - selected: Vec, - tree: &DocumentTree, - ) -> StructuredContent { + pub fn build(&self, selected: Vec, tree: &DocumentTree) -> StructuredContent { if selected.is_empty() { return StructuredContent { content: String::new(), @@ -470,9 +466,7 @@ mod tests { #[test] fn test_flat_builder() { let builder = StructureBuilder::new(OutputFormat::Flat); - let selected = vec![ - make_selected("Section 1", "Content 1", 0.9, 0), - ]; + let selected = vec![make_selected("Section 1", "Content 1", 0.9, 0)]; let tree = DocumentTree::new("Test", ""); let result = builder.build(selected, &tree); @@ -483,12 +477,9 @@ mod tests { #[test] fn test_builder_with_scores() { - let builder = StructureBuilder::new(OutputFormat::Markdown) - .with_scores(); + let builder = StructureBuilder::new(OutputFormat::Markdown).with_scores(); - let selected = vec![ - make_selected("Section 1", "Content 1", 0.95, 0), - ]; + let selected = vec![make_selected("Section 1", "Content 1", 0.95, 0)]; let tree = DocumentTree::new("Test", ""); let result = builder.build(selected, &tree); @@ -508,8 +499,8 @@ mod tests { #[test] fn test_content_tree_node() { - let mut root = ContentTreeNode::new("Root".to_string()) - .with_content("Root content".to_string(), 0.9); + let mut root = + ContentTreeNode::new("Root".to_string()).with_content("Root content".to_string(), 0.9); let child = ContentTreeNode::new("Child".to_string()) .with_content("Child content".to_string(), 0.8); diff --git a/src/retrieval/content/config.rs b/src/retrieval/content/config.rs index f9bc38b6..aa40bc8a 100644 --- a/src/retrieval/content/config.rs +++ b/src/retrieval/content/config.rs @@ -147,12 +147,10 @@ mod tests { #[test] fn test_min_relevance_clamped() { - let config = ContentAggregatorConfig::new() - .with_min_relevance(1.5); + let config = ContentAggregatorConfig::new().with_min_relevance(1.5); assert_eq!(config.min_relevance_score, 1.0); - let config = ContentAggregatorConfig::new() - .with_min_relevance(-0.5); + let config = ContentAggregatorConfig::new().with_min_relevance(-0.5); assert_eq!(config.min_relevance_score, 0.0); } } diff --git a/src/retrieval/content/mod.rs b/src/retrieval/content/mod.rs index 2a78f801..5b02588b 100644 --- a/src/retrieval/content/mod.rs +++ b/src/retrieval/content/mod.rs @@ -37,10 +37,10 @@ mod builder; mod config; mod scorer; -pub use aggregator::{ContentAggregator, AggregationResult, CandidateNode}; -pub use budget::{BudgetAllocator, AllocationStrategy, AllocationResult, SelectedContent}; -pub use builder::{StructureBuilder, OutputFormat, StructuredContent, ContentTree}; +pub use aggregator::{AggregationResult, CandidateNode, ContentAggregator}; +pub use budget::{AllocationResult, AllocationStrategy, BudgetAllocator, SelectedContent}; +pub use builder::{ContentTree, OutputFormat, StructureBuilder, StructuredContent}; pub use config::{ContentAggregatorConfig, OutputFormatConfig, ScoringStrategyConfig}; pub use scorer::{ - RelevanceScorer, ContentRelevance, ScoreComponents, ContentChunk, ScoringContext, + ContentChunk, ContentRelevance, RelevanceScorer, ScoreComponents, ScoringContext, }; diff --git a/src/retrieval/content/scorer.rs b/src/retrieval/content/scorer.rs index daf49550..78219819 100644 --- a/src/retrieval/content/scorer.rs +++ b/src/retrieval/content/scorer.rs @@ -167,7 +167,10 @@ impl RelevanceScorer { components.keyword_score = self.compute_keyword_score(&chunk.content); // 2. BM25 score (if enabled) - if matches!(self.strategy, ScoringStrategyConfig::KeywordWithBM25 | ScoringStrategyConfig::Hybrid) { + if matches!( + self.strategy, + ScoringStrategyConfig::KeywordWithBM25 | ScoringStrategyConfig::Hybrid + ) { components.bm25_score = self.compute_bm25_score(&chunk.content, ctx); } @@ -204,9 +207,8 @@ impl RelevanceScorer { } let content_lower = content.to_lowercase(); - let content_words: std::collections::HashSet<&str> = content_lower - .split_whitespace() - .collect(); + let content_words: std::collections::HashSet<&str> = + content_lower.split_whitespace().collect(); let matches = self .query_keywords @@ -232,10 +234,7 @@ impl RelevanceScorer { for term in &self.query_keywords { let term_lower = term.to_lowercase(); - let tf = content - .to_lowercase() - .matches(&term_lower) - .count() as f32; + let tf = content.to_lowercase().matches(&term_lower).count() as f32; if tf == 0.0 { continue; @@ -268,21 +267,128 @@ impl RelevanceScorer { fn extract_keywords(query: &str) -> Vec { // Common English stop words const STOPWORDS: &[&str] = &[ - "a", "an", "the", "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", "about", - "what", "which", "who", "whom", "this", "that", "these", "those", - "i", "me", "my", "myself", "we", "our", "ours", "ourselves", "you", - "your", "yours", "yourself", "yourselves", "he", "him", "his", - "himself", "she", "her", "hers", "herself", "it", "its", "itself", - "they", "them", "their", "theirs", "themselves", + "a", + "an", + "the", + "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", + "about", + "what", + "which", + "who", + "whom", + "this", + "that", + "these", + "those", + "i", + "me", + "my", + "myself", + "we", + "our", + "ours", + "ourselves", + "you", + "your", + "yours", + "yourself", + "yourselves", + "he", + "him", + "his", + "himself", + "she", + "her", + "hers", + "herself", + "it", + "its", + "itself", + "they", + "them", + "their", + "theirs", + "themselves", ]; query @@ -305,10 +411,9 @@ fn compute_density(content: &str) -> f32 { // Stopword ratio (lower is better) const STOPWORDS: &[&str] = &[ - "a", "an", "the", "is", "are", "was", "were", "be", "been", "being", - "have", "has", "had", "do", "does", "did", "will", "would", "could", - "should", "may", "might", "must", "shall", "can", "to", "of", "in", - "for", "on", "with", "at", "by", "from", "and", "but", "or", "as", + "a", "an", "the", "is", "are", "was", "were", "be", "been", "being", "have", "has", "had", + "do", "does", "did", "will", "would", "could", "should", "may", "might", "must", "shall", + "can", "to", "of", "in", "for", "on", "with", "at", "by", "from", "and", "but", "or", "as", ]; let stopword_count = words @@ -321,10 +426,7 @@ fn compute_density(content: &str) -> f32 { // Entity-like ratio (capitalized, numbers, special terms) let entity_count = words .iter() - .filter(|w| { - w.chars() - .any(|c| c.is_numeric() || c.is_uppercase()) - }) + .filter(|w| w.chars().any(|c| c.is_numeric() || c.is_uppercase())) .count(); let entity_ratio = entity_count as f32 / words.len() as f32; diff --git a/src/retrieval/mod.rs b/src/retrieval/mod.rs index 565d0fa8..4124e9a6 100644 --- a/src/retrieval/mod.rs +++ b/src/retrieval/mod.rs @@ -107,8 +107,8 @@ pub use content::{ }; // Pilot exports +pub use pilot::NoopPilot; pub use pilot::{ BudgetConfig, InterventionConfig, InterventionPoint, Pilot, PilotConfig, PilotDecision, PilotMode, RankedCandidate, SearchDirection, SearchState, }; -pub use pilot::NoopPilot; diff --git a/src/retrieval/pilot/budget.rs b/src/retrieval/pilot/budget.rs index defa1847..4776d931 100644 --- a/src/retrieval/pilot/budget.rs +++ b/src/retrieval/pilot/budget.rs @@ -7,8 +7,8 @@ //! and control costs during retrieval. use std::collections::HashMap; -use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::RwLock; +use std::sync::atomic::{AtomicUsize, Ordering}; use super::config::BudgetConfig; @@ -118,8 +118,7 @@ impl BudgetController { let tokens = self.total_tokens(); let calls = self.calls_made.load(Ordering::Relaxed); - tokens < self.config.max_tokens_per_query - && calls < self.config.max_calls_per_query + tokens < self.config.max_tokens_per_query && calls < self.config.max_calls_per_query } /// Check if a call is allowed at a specific tree level. @@ -151,7 +150,8 @@ impl BudgetController { let english_count = char_count - chinese_count; // Estimate tokens - let input_tokens = (chinese_count as f32 / 1.5 + english_count as f32 / 4.0).ceil() as usize; + let input_tokens = + (chinese_count as f32 / 1.5 + english_count as f32 / 4.0).ceil() as usize; // Add output reserve input_tokens + 100 @@ -161,8 +161,7 @@ impl BudgetController { pub fn can_afford(&self, estimated_cost: usize) -> bool { let remaining = self.remaining_tokens(); - estimated_cost <= remaining - && estimated_cost <= self.config.max_tokens_per_call + estimated_cost <= remaining && estimated_cost <= self.config.max_tokens_per_call } /// Get remaining token budget. @@ -188,7 +187,8 @@ impl BudgetController { /// * `level` - Tree level where call was made pub fn record_usage(&self, input_tokens: usize, output_tokens: usize, level: usize) { self.input_tokens.fetch_add(input_tokens, Ordering::Relaxed); - self.output_tokens.fetch_add(output_tokens, Ordering::Relaxed); + self.output_tokens + .fetch_add(output_tokens, Ordering::Relaxed); self.calls_made.fetch_add(1, Ordering::Relaxed); // Track level calls @@ -200,8 +200,7 @@ impl BudgetController { /// Get total tokens used. pub fn total_tokens(&self) -> usize { - self.input_tokens.load(Ordering::Relaxed) - + self.output_tokens.load(Ordering::Relaxed) + self.input_tokens.load(Ordering::Relaxed) + self.output_tokens.load(Ordering::Relaxed) } /// Get current usage statistics. @@ -296,13 +295,21 @@ mod tests { // English text - 26 chars ≈ 7 tokens + 100 output reserve = ~107 let english = "Hello world this is a test"; let cost = budget.estimate_cost(english); - assert!(cost > 100 && cost < 150, "Expected cost between 100-150, got {}", cost); + assert!( + cost > 100 && cost < 150, + "Expected cost between 100-150, got {}", + cost + ); // Chinese text - 6 chars ≈ 4 tokens + 100 output reserve = ~104 let chinese = "这是一个测试"; let cost_chinese = budget.estimate_cost(chinese); // Both have ~100 token base from output reserve, so just check it's reasonable - assert!(cost_chinese > 100, "Expected Chinese cost > 100, got {}", cost_chinese); + assert!( + cost_chinese > 100, + "Expected Chinese cost > 100, got {}", + cost_chinese + ); } #[test] diff --git a/src/retrieval/pilot/builder.rs b/src/retrieval/pilot/builder.rs index 931c19b0..fd9d4581 100644 --- a/src/retrieval/pilot/builder.rs +++ b/src/retrieval/pilot/builder.rs @@ -16,8 +16,8 @@ use std::collections::HashSet; -use crate::document::{DocumentTree, NodeId}; use super::SearchState; +use crate::document::{DocumentTree, NodeId}; /// Token budget distribution for context building. #[derive(Debug, Clone)] @@ -47,7 +47,13 @@ impl TokenBudget { } /// Create budget with custom distribution. - pub fn with_distribution(total: usize, query_pct: f32, path_pct: f32, candidates_pct: f32, siblings_pct: f32) -> Self { + pub fn with_distribution( + total: usize, + query_pct: f32, + path_pct: f32, + candidates_pct: f32, + siblings_pct: f32, + ) -> Self { let sum = query_pct + path_pct + candidates_pct + siblings_pct; Self { total, @@ -85,10 +91,7 @@ impl PilotContext { pub fn to_string(&self) -> String { format!( "{}\n{}\n{}\n{}", - self.query_section, - self.path_section, - self.candidates_section, - self.toc_section + self.query_section, self.path_section, self.candidates_section, self.toc_section ) } @@ -222,7 +225,10 @@ impl ContextBuilder { ctx.estimated_tokens += self.estimate_tokens(&ctx.query_section); // Show failed path - ctx.path_section = format!("Failed path:\n{}", self.build_path_section(state.tree, failed_path)); + ctx.path_section = format!( + "Failed path:\n{}", + self.build_path_section(state.tree, failed_path) + ); ctx.estimated_tokens += self.estimate_tokens(&ctx.path_section); // Show unvisited alternatives @@ -330,7 +336,11 @@ impl ContextBuilder { for sibling_id in siblings.iter().take(8) { if let Some(node) = tree.get(*sibling_id) { - let marker = if *sibling_id == current_id { "⭐ " } else { "" }; + let marker = if *sibling_id == current_id { + "⭐ " + } else { + "" + }; result.push_str(&format!(" {}{}\n", marker, node.title)); } } @@ -366,7 +376,15 @@ impl ContextBuilder { // Only show children for first few levels if depth < max_depth { for child_id in tree.children(node_id) { - build_toc_recursive(tree, child_id, depth + 1, result, tokens_used, max_tokens, max_depth); + build_toc_recursive( + tree, + child_id, + depth + 1, + result, + tokens_used, + max_tokens, + max_depth, + ); } } } @@ -470,7 +488,7 @@ mod tests { fn test_token_budget_distribution() { let budget = TokenBudget::new(500); assert_eq!(budget.query, 150); // 30% - assert_eq!(budget.path, 100); // 20% + assert_eq!(budget.path, 100); // 20% assert_eq!(budget.candidates, 200); // 40% assert_eq!(budget.siblings, 50); // 10% } @@ -496,7 +514,11 @@ mod tests { let builder = ContextBuilder::new(20); // Very small budget - 20 * 0.30 = 6 tokens for query = ~24 chars let long_query = "This is a very long query that should be truncated because it exceeds the token budget"; let result = builder.build_query_section(long_query); - assert!(result.contains("..."), "Expected truncation, got: {}", result); + assert!( + result.contains("..."), + "Expected truncation, got: {}", + result + ); } #[test] diff --git a/src/retrieval/pilot/config.rs b/src/retrieval/pilot/config.rs index f381393e..14eb01c4 100644 --- a/src/retrieval/pilot/config.rs +++ b/src/retrieval/pilot/config.rs @@ -131,7 +131,7 @@ impl PilotMode { /// Get the fork threshold multiplier for this mode. pub fn fork_threshold_multiplier(&self) -> f32 { match self { - PilotMode::Aggressive => 0.5, // Lower threshold = more interventions + PilotMode::Aggressive => 0.5, // Lower threshold = more interventions PilotMode::Balanced => 1.0, PilotMode::Conservative => 2.0, // Higher threshold = fewer interventions PilotMode::AlgorithmOnly => f32::MAX, diff --git a/src/retrieval/pilot/decision.rs b/src/retrieval/pilot/decision.rs index 084582c2..4ecaf901 100644 --- a/src/retrieval/pilot/decision.rs +++ b/src/retrieval/pilot/decision.rs @@ -288,10 +288,12 @@ mod tests { assert_eq!(candidate.score, 0.8); assert!(candidate.reason.is_none()); - let candidate_with_reason = - RankedCandidate::with_reason(node_ids[0], 0.9, "test reason"); + let candidate_with_reason = RankedCandidate::with_reason(node_ids[0], 0.9, "test reason"); assert_eq!(candidate_with_reason.score, 0.9); - assert_eq!(candidate_with_reason.reason, Some("test reason".to_string())); + assert_eq!( + candidate_with_reason.reason, + Some("test reason".to_string()) + ); } #[test] @@ -300,7 +302,10 @@ mod tests { assert!(matches!(deeper, SearchDirection::GoDeeper { .. })); let found = SearchDirection::found_answer(0.9); - assert!(matches!(found, SearchDirection::FoundAnswer { confidence: 0.9 })); + assert!(matches!( + found, + SearchDirection::FoundAnswer { confidence: 0.9 } + )); } #[test] diff --git a/src/retrieval/pilot/fallback.rs b/src/retrieval/pilot/fallback.rs index 874cdd96..da93354e 100644 --- a/src/retrieval/pilot/fallback.rs +++ b/src/retrieval/pilot/fallback.rs @@ -170,7 +170,10 @@ impl std::fmt::Debug for FallbackManager { f.debug_struct("FallbackManager") .field("config", &self.config) .field("current_level", &self.current_level()) - .field("consecutive_failures", &self.consecutive_failures.load(Ordering::Relaxed)) + .field( + "consecutive_failures", + &self.consecutive_failures.load(Ordering::Relaxed), + ) .finish() } } @@ -393,7 +396,10 @@ mod tests { // Trigger failures to escalate for _ in 0..manager.config.failures_before_escalate { let action = manager.record_failure(&FallbackError::Network("test".to_string())); - assert!(matches!(action, FallbackAction::Retry | FallbackAction::Escalate)); + assert!(matches!( + action, + FallbackAction::Retry | FallbackAction::Escalate + )); } assert_eq!(manager.current_level(), FallbackLevel::Retry); diff --git a/src/retrieval/pilot/llm_pilot.rs b/src/retrieval/pilot/llm_pilot.rs index 10118ff0..414bbb3e 100644 --- a/src/retrieval/pilot/llm_pilot.rs +++ b/src/retrieval/pilot/llm_pilot.rs @@ -13,8 +13,8 @@ use tracing::{debug, info, warn}; use crate::document::DocumentTree; use crate::llm::LlmClient; -use super::builder::ContextBuilder; use super::budget::BudgetController; +use super::builder::ContextBuilder; use super::config::PilotConfig; use super::decision::{InterventionPoint, PilotDecision}; use super::parser::ResponseParser; @@ -126,7 +126,8 @@ impl LlmPilot { fn scores_are_close(&self, state: &SearchState<'_>) -> bool { // Use the config's score_gap_threshold with the state's best_score // If best_score is low, consider scores as close - state.candidates.len() >= 2 && state.best_score < self.config.intervention.score_gap_threshold + state.candidates.len() >= 2 + && state.best_score < self.config.intervention.score_gap_threshold } /// Determine the intervention point type. @@ -154,7 +155,10 @@ impl LlmPilot { // Check if we can afford this call if !self.budget.can_afford(prompt.estimated_tokens) { - warn!("Budget cannot afford LLM call (estimated: {} tokens)", prompt.estimated_tokens); + warn!( + "Budget cannot afford LLM call (estimated: {} tokens)", + prompt.estimated_tokens + ); return self.default_decision(candidates, point); } @@ -168,7 +172,8 @@ impl LlmPilot { Ok(response) => { // Record usage (estimate output tokens) let output_tokens = self.estimate_tokens(&response); - self.budget.record_usage(prompt.estimated_tokens, output_tokens, 0); + self.budget + .record_usage(prompt.estimated_tokens, output_tokens, 0); // Parse response let decision = self.response_parser.parse(&response, candidates, point); @@ -251,7 +256,10 @@ impl Pilot for LlmPilot { // Condition 1: Fork point with enough candidates if state.candidates.len() > intervention.fork_threshold { - debug!("Intervening: fork point with {} candidates", state.candidates.len()); + debug!( + "Intervening: fork point with {} candidates", + state.candidates.len() + ); return true; } @@ -263,7 +271,10 @@ impl Pilot for LlmPilot { // Condition 3: Low confidence (best score too low) if intervention.is_low_confidence(state.best_score) { - debug!("Intervening: low confidence (best_score={:.2})", state.best_score); + debug!( + "Intervening: low confidence (best_score={:.2})", + state.best_score + ); return true; } @@ -286,11 +297,7 @@ impl Pilot for LlmPilot { self.call_llm(point, &context, state.candidates).await } - async fn guide_start( - &self, - tree: &DocumentTree, - query: &str, - ) -> Option { + async fn guide_start(&self, tree: &DocumentTree, query: &str) -> Option { // Check if guide_at_start is enabled if !self.config.guide_at_start { return None; @@ -309,7 +316,9 @@ impl Pilot for LlmPilot { let candidates = tree.children(tree.root()); // Make LLM call - let decision = self.call_llm(InterventionPoint::Start, &context, &candidates).await; + let decision = self + .call_llm(InterventionPoint::Start, &context, &candidates) + .await; info!( "Pilot start guidance: confidence={}, candidates={}", decision.confidence, @@ -319,10 +328,7 @@ impl Pilot for LlmPilot { Some(decision) } - async fn guide_backtrack( - &self, - state: &SearchState<'_>, - ) -> Option { + async fn guide_backtrack(&self, state: &SearchState<'_>) -> Option { // Check if guide_at_backtrack is enabled if !self.config.guide_at_backtrack { return None; @@ -334,10 +340,15 @@ impl Pilot for LlmPilot { } // Build backtrack context - let context = self.context_builder.build_backtrack_context(state, state.path); + let context = self + .context_builder + .build_backtrack_context(state, state.path); // Make LLM call - Some(self.call_llm(InterventionPoint::Backtrack, &context, state.candidates).await) + Some( + self.call_llm(InterventionPoint::Backtrack, &context, state.candidates) + .await, + ) } fn config(&self) -> &PilotConfig { diff --git a/src/retrieval/pilot/metrics.rs b/src/retrieval/pilot/metrics.rs index d9e7b656..dba2d9b5 100644 --- a/src/retrieval/pilot/metrics.rs +++ b/src/retrieval/pilot/metrics.rs @@ -208,12 +208,15 @@ impl MetricsCollector { } // Update token counters - self.total_input_tokens.fetch_add(input_tokens, Ordering::Relaxed); - self.total_output_tokens.fetch_add(output_tokens, Ordering::Relaxed); + self.total_input_tokens + .fetch_add(input_tokens, Ordering::Relaxed); + self.total_output_tokens + .fetch_add(output_tokens, Ordering::Relaxed); // Update latency let latency_ms = latency.as_millis() as u64; - self.total_latency_ms.fetch_add(latency_ms, Ordering::Relaxed); + self.total_latency_ms + .fetch_add(latency_ms, Ordering::Relaxed); // Store latency sample if let Ok(mut samples) = self.latency_samples.write() { @@ -328,7 +331,10 @@ impl MetricsCollector { let p99_idx = (latencies.len() as f64 * 0.99) as usize; let p50 = latencies.get(p50_idx).copied().unwrap_or(0); - let p99 = latencies.get(p99_idx.min(latencies.len() - 1)).copied().unwrap_or(0); + let p99 = latencies + .get(p99_idx.min(latencies.len() - 1)) + .copied() + .unwrap_or(0); (p50, p99) } else { @@ -445,7 +451,14 @@ mod tests { fn test_token_utilization() { let metrics = MetricsCollector::new(); - metrics.record_call(InterventionPoint::Fork, 500, 200, Duration::ZERO, true, false); + metrics.record_call( + InterventionPoint::Fork, + 500, + 200, + Duration::ZERO, + true, + false, + ); let utilization = metrics.snapshot().token_utilization(1000); assert!((utilization - 0.7).abs() < 0.01); @@ -480,7 +493,14 @@ mod tests { fn test_reset() { let metrics = MetricsCollector::new(); - metrics.record_call(InterventionPoint::Fork, 100, 50, Duration::from_millis(200), true, false); + metrics.record_call( + InterventionPoint::Fork, + 100, + 50, + Duration::from_millis(200), + true, + false, + ); assert!(metrics.total_calls() > 0); metrics.reset(); diff --git a/src/retrieval/pilot/mod.rs b/src/retrieval/pilot/mod.rs index 6aaa9eb5..0a3a9a61 100644 --- a/src/retrieval/pilot/mod.rs +++ b/src/retrieval/pilot/mod.rs @@ -65,12 +65,8 @@ mod r#trait; pub use budget::{BudgetController, BudgetUsage}; pub use builder::{ContextBuilder, PilotContext, TokenBudget}; -pub use config::{ - BudgetConfig, InterventionConfig, PilotConfig, PilotMode, -}; -pub use decision::{ - InterventionPoint, PilotDecision, RankedCandidate, SearchDirection, -}; +pub use config::{BudgetConfig, InterventionConfig, PilotConfig, PilotMode}; +pub use decision::{InterventionPoint, PilotDecision, RankedCandidate, SearchDirection}; pub use fallback::{FallbackAction, FallbackConfig, FallbackError, FallbackLevel, FallbackManager}; pub use llm_pilot::LlmPilot; pub use metrics::{CallRecord, MetricsCollector, PilotMetrics}; diff --git a/src/retrieval/pilot/noop.rs b/src/retrieval/pilot/noop.rs index b79156a5..ffedf5b8 100644 --- a/src/retrieval/pilot/noop.rs +++ b/src/retrieval/pilot/noop.rs @@ -69,19 +69,12 @@ impl Pilot for NoopPilot { } } - async fn guide_start( - &self, - _tree: &DocumentTree, - _query: &str, - ) -> Option { + async fn guide_start(&self, _tree: &DocumentTree, _query: &str) -> Option { // No guidance at start None } - async fn guide_backtrack( - &self, - _state: &SearchState<'_>, - ) -> Option { + async fn guide_backtrack(&self, _state: &SearchState<'_>) -> Option { // No guidance during backtrack None } diff --git a/src/retrieval/pilot/parser.rs b/src/retrieval/pilot/parser.rs index ca88ff26..1d47d9ae 100644 --- a/src/retrieval/pilot/parser.rs +++ b/src/retrieval/pilot/parser.rs @@ -13,8 +13,8 @@ use regex::Regex; use serde::{Deserialize, Serialize}; use tracing::warn; +use super::decision::{InterventionPoint, PilotDecision, RankedCandidate, SearchDirection}; use crate::document::NodeId; -use super::decision::{PilotDecision, RankedCandidate, SearchDirection, InterventionPoint}; /// Parsed response from LLM. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -100,8 +100,9 @@ impl ResponseParser { confidence_regex: Regex::new(r"(?i)confidence[:\s]+([0-9.]+)").unwrap(), // Match direction keywords direction_regex: Regex::new( - r"(?i)(go.?deeper|explore.?siblings|backtrack|found.?answer)" - ).unwrap(), + r"(?i)(go.?deeper|explore.?siblings|backtrack|found.?answer)", + ) + .unwrap(), } } @@ -170,30 +171,40 @@ impl ResponseParser { point: InterventionPoint, ) -> Option { // Extract confidence - let confidence = self.confidence_regex + let confidence = self + .confidence_regex .captures(response) .and_then(|caps| caps.get(1)?.as_str().parse::().ok()) .unwrap_or(0.5) .clamp(0.0, 1.0); // Extract direction - let direction = self.direction_regex + let direction = self + .direction_regex .captures(response) .map(|caps| { let dir = caps.get(1)?.as_str().to_lowercase(); match dir.as_str() { - d if d.contains("deeper") => Some(SearchDirection::GoDeeper { reason: String::new() }), - d if d.contains("sibling") => Some(SearchDirection::ExploreSiblings { recommended: vec![] }), + d if d.contains("deeper") => Some(SearchDirection::GoDeeper { + reason: String::new(), + }), + d if d.contains("sibling") => Some(SearchDirection::ExploreSiblings { + recommended: vec![], + }), d if d.contains("backtrack") => Some(SearchDirection::Backtrack { reason: String::new(), alternative_branches: vec![], }), - d if d.contains("found") || d.contains("answer") => Some(SearchDirection::FoundAnswer { confidence }), + d if d.contains("found") || d.contains("answer") => { + Some(SearchDirection::FoundAnswer { confidence }) + } _ => None, } }) .flatten() - .unwrap_or_else(|| SearchDirection::GoDeeper { reason: String::new() }); + .unwrap_or_else(|| SearchDirection::GoDeeper { + reason: String::new(), + }); // Try to extract candidate rankings from numbered list let ranked = self.extract_ranked_candidates(response, candidates); @@ -212,16 +223,23 @@ impl ResponseParser { } /// Extract ranked candidates from text using patterns. - fn extract_ranked_candidates(&self, response: &str, candidates: &[NodeId]) -> Vec { + fn extract_ranked_candidates( + &self, + response: &str, + candidates: &[NodeId], + ) -> Vec { let mut ranked = Vec::new(); // Pattern: "1. Candidate Name (score: 0.8)" - let ranking_pattern = Regex::new(r"(\d+)[.\)]\s*(?:Candidate\s*)?(\d+)[\s:]+(?:score[:\s]*)?([0-9.]+)?").unwrap(); + let ranking_pattern = + Regex::new(r"(\d+)[.\)]\s*(?:Candidate\s*)?(\d+)[\s:]+(?:score[:\s]*)?([0-9.]+)?") + .unwrap(); for caps in ranking_pattern.captures_iter(response) { if let Some(index_match) = caps.get(2) { if let Ok(index) = index_match.as_str().parse::() { - let score: f32 = caps.get(3) + let score: f32 = caps + .get(3) .and_then(|m| m.as_str().parse().ok()) .unwrap_or(0.5); @@ -296,11 +314,19 @@ impl ResponseParser { reason: llm_response.reasoning.clone(), }, DirectionResponse::ExploreSiblings => SearchDirection::ExploreSiblings { - recommended: ranked_candidates.iter().take(3).map(|c| c.node_id).collect(), + recommended: ranked_candidates + .iter() + .take(3) + .map(|c| c.node_id) + .collect(), }, DirectionResponse::Backtrack => SearchDirection::Backtrack { reason: llm_response.reasoning.clone(), - alternative_branches: ranked_candidates.iter().take(3).map(|c| c.node_id).collect(), + alternative_branches: ranked_candidates + .iter() + .take(3) + .map(|c| c.node_id) + .collect(), }, DirectionResponse::FoundAnswer => SearchDirection::FoundAnswer { confidence: llm_response.confidence, @@ -331,7 +357,9 @@ impl ResponseParser { PilotDecision { ranked_candidates: ranked, - direction: SearchDirection::GoDeeper { reason: String::new() }, + direction: SearchDirection::GoDeeper { + reason: String::new(), + }, confidence: 0.0, reasoning: "Default decision (parsing failed)".to_string(), intervention_point: point, @@ -387,7 +415,10 @@ mod tests { assert_eq!(decision.ranked_candidates.len(), 2); assert_eq!(decision.ranked_candidates[0].node_id, candidates[1]); assert!((decision.confidence - 0.85).abs() < 0.01); - assert!(matches!(decision.direction, SearchDirection::GoDeeper { .. })); + assert!(matches!( + decision.direction, + SearchDirection::GoDeeper { .. } + )); } #[test] @@ -477,7 +508,10 @@ Direction: go_deeper ]; for (dir_json, should_parse) in test_cases { - let response = format!(r#"{{"ranked_candidates": [], "confidence": 0.5, {}}}"#, dir_json); + let response = format!( + r#"{{"ranked_candidates": [], "confidence": 0.5, {}}}"#, + dir_json + ); let decision = parser.parse(&response, &candidates, InterventionPoint::Fork); assert!(should_parse, "Direction should parse correctly"); } diff --git a/src/retrieval/pilot/prompts/builder.rs b/src/retrieval/pilot/prompts/builder.rs index 47e90d47..0d03f09d 100644 --- a/src/retrieval/pilot/prompts/builder.rs +++ b/src/retrieval/pilot/prompts/builder.rs @@ -7,7 +7,7 @@ use super::super::builder::PilotContext; use super::super::decision::InterventionPoint; -use super::templates::{ForkPrompt, PromptTemplate, StartPrompt, BacktrackPrompt, EvaluatePrompt}; +use super::templates::{BacktrackPrompt, EvaluatePrompt, ForkPrompt, PromptTemplate, StartPrompt}; /// Built prompt ready for LLM call. #[derive(Debug, Clone)] diff --git a/src/retrieval/pilot/prompts/mod.rs b/src/retrieval/pilot/prompts/mod.rs index 52684dcb..0bb336b2 100644 --- a/src/retrieval/pilot/prompts/mod.rs +++ b/src/retrieval/pilot/prompts/mod.rs @@ -13,6 +13,4 @@ mod builder; mod templates; pub use builder::PromptBuilder; -pub use templates::{ - ForkPrompt, PromptTemplate, StartPrompt, BacktrackPrompt, EvaluatePrompt, -}; +pub use templates::{BacktrackPrompt, EvaluatePrompt, ForkPrompt, PromptTemplate, StartPrompt}; diff --git a/src/retrieval/pilot/prompts/templates.rs b/src/retrieval/pilot/prompts/templates.rs index d9106f6d..5f9f75ff 100644 --- a/src/retrieval/pilot/prompts/templates.rs +++ b/src/retrieval/pilot/prompts/templates.rs @@ -268,7 +268,8 @@ Respond in JSON format with your analysis."#.to_string() {context} -Respond in JSON format with ranked candidates."#.to_string() +Respond in JSON format with ranked candidates."# + .to_string() } pub fn system_backtrack() -> String { @@ -292,7 +293,8 @@ Respond in JSON format with alternative branches."#.to_string() {context} -Respond in JSON format with your evaluation."#.to_string() +Respond in JSON format with your evaluation."# + .to_string() } } diff --git a/src/retrieval/pilot/trait.rs b/src/retrieval/pilot/trait.rs index 94e7fac7..5873d2b0 100644 --- a/src/retrieval/pilot/trait.rs +++ b/src/retrieval/pilot/trait.rs @@ -13,7 +13,7 @@ use std::sync::LazyLock; use crate::document::{DocumentTree, NodeId}; -use super::{PilotConfig, PilotDecision, InterventionPoint}; +use super::{InterventionPoint, PilotConfig, PilotDecision}; /// Empty HashSet for use in SearchState::for_start static EMPTY_VISITED: LazyLock> = LazyLock::new(HashSet::new); @@ -168,11 +168,7 @@ pub trait Pilot: Send + Sync { /// the starting point and initial direction. /// /// Returns `None` if no guidance is available or needed. - async fn guide_start( - &self, - tree: &DocumentTree, - query: &str, - ) -> Option; + async fn guide_start(&self, tree: &DocumentTree, query: &str) -> Option; /// Provide guidance during backtracking. /// @@ -181,10 +177,7 @@ pub trait Pilot: Send + Sync { /// alternative paths. /// /// Returns `None` if no guidance is available. - async fn guide_backtrack( - &self, - state: &SearchState<'_>, - ) -> Option; + async fn guide_backtrack(&self, state: &SearchState<'_>) -> Option; /// Get the current configuration. fn config(&self) -> &PilotConfig; diff --git a/src/retrieval/pipeline/orchestrator.rs b/src/retrieval/pipeline/orchestrator.rs index fc013014..0dce81f1 100644 --- a/src/retrieval/pipeline/orchestrator.rs +++ b/src/retrieval/pipeline/orchestrator.rs @@ -15,7 +15,7 @@ use std::sync::Arc; use std::time::Instant; use tracing::{debug, error, info, warn}; -use crate::document::{DocumentTree}; +use crate::document::DocumentTree; use crate::error::Result; use crate::retrieval::pilot::{Pilot, SearchState}; // FailurePolicy is re-exported for stages @@ -381,11 +381,11 @@ impl RetrievalOrchestrator { if let Some(ref pilot) = self.pilot { if pilot.config().guide_at_backtrack { // Build search state for Pilot - let visited: std::collections::HashSet<_> = - ctx.search_paths - .iter() - .flat_map(|p| p.nodes.iter().copied()) - .collect(); + let visited: std::collections::HashSet<_> = ctx + .search_paths + .iter() + .flat_map(|p| p.nodes.iter().copied()) + .collect(); let candidates: Vec<_> = ctx.candidates.iter().map(|c| c.node_id).collect(); @@ -460,13 +460,16 @@ impl RetrievalOrchestrator { if target_stage == "search" { if let Some(ref pilot) = self.pilot { if pilot.config().guide_at_backtrack { - let visited: std::collections::HashSet<_> = - ctx.search_paths - .iter() - .flat_map(|p| p.nodes.iter().copied()) - .collect(); - let candidates: Vec<_> = - ctx.candidates.iter().map(|c| c.node_id).collect(); + let visited: std::collections::HashSet<_> = ctx + .search_paths + .iter() + .flat_map(|p| p.nodes.iter().copied()) + .collect(); + let candidates: Vec<_> = ctx + .candidates + .iter() + .map(|c| c.node_id) + .collect(); let state = SearchState::new( &ctx.tree, @@ -476,7 +479,9 @@ impl RetrievalOrchestrator { &visited, ); - if let Some(guidance) = pilot.guide_backtrack(&state).await { + if let Some(guidance) = + pilot.guide_backtrack(&state).await + { debug!( "Pilot backtrack guidance for explicit backtrack: confidence={}", guidance.confidence diff --git a/src/retrieval/search/beam.rs b/src/retrieval/search/beam.rs index ea73051c..9fba59e9 100644 --- a/src/retrieval/search/beam.rs +++ b/src/retrieval/search/beam.rs @@ -76,7 +76,8 @@ impl BeamSearch { let beta = 0.6 * pilot_decision.confidence; // Build a map from node_id to pilot score - let mut pilot_scores: std::collections::HashMap = std::collections::HashMap::new(); + let mut pilot_scores: std::collections::HashMap = + std::collections::HashMap::new(); for ranked in &pilot_decision.ranked_candidates { pilot_scores.insert(ranked.node_id, ranked.score); } @@ -100,9 +101,7 @@ impl BeamSearch { .collect(); // Sort by merged score - merged.sort_by(|a, b| { - b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal) - }); + merged.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); merged } @@ -137,12 +136,20 @@ impl SearchTree for BeamSearch { let initial_candidates = if let Some(p) = pilot { if p.config().guide_at_start { if let Some(guidance) = p.guide_start(tree, &context.query).await { - debug!("Pilot provided start guidance with confidence {}", guidance.confidence); + debug!( + "Pilot provided start guidance with confidence {}", + guidance.confidence + ); pilot_interventions += 1; // Use Pilot's ranked order if available if guidance.has_candidates() { - self.merge_with_pilot_decision(tree, &root_children, &guidance, &context.query) + self.merge_with_pilot_decision( + tree, + &root_children, + &guidance, + &context.query, + ) } else { self.score_candidates_with_query(tree, &root_children, &context.query) } @@ -203,7 +210,10 @@ impl SearchTree for BeamSearch { // Check if Pilot wants to intervene if p.should_intervene(&state) { - trace!("Pilot intervening at fork with {} candidates", children.len()); + trace!( + "Pilot intervening at fork with {} candidates", + children.len() + ); match p.decide(&state).await { decision => { @@ -215,7 +225,12 @@ impl SearchTree for BeamSearch { ); // Merge algorithm scores with Pilot decision - self.merge_with_pilot_decision(tree, &children, &decision, &context.query) + self.merge_with_pilot_decision( + tree, + &children, + &decision, + &context.query, + ) } } } else { @@ -276,7 +291,8 @@ impl SearchTree for BeamSearch { if result.paths.is_empty() && config.min_score > 0.0 { debug!("No results above min_score, adding best candidates as fallback"); // Re-score initial candidates and take top-k - let all_candidates = self.score_candidates_with_query(tree, &tree.children(tree.root()), &context.query); + let all_candidates = + self.score_candidates_with_query(tree, &tree.children(tree.root()), &context.query); for (node_id, score) in all_candidates.into_iter().take(config.top_k) { result.paths.push(SearchPath::from_node(node_id, score)); } diff --git a/src/retrieval/search/greedy.rs b/src/retrieval/search/greedy.rs index 89357225..b539cd23 100644 --- a/src/retrieval/search/greedy.rs +++ b/src/retrieval/search/greedy.rs @@ -57,7 +57,8 @@ impl GreedySearch { let beta = 0.6 * pilot_decision.confidence; // Build a map from node_id to pilot score - let mut pilot_scores: std::collections::HashMap = std::collections::HashMap::new(); + let mut pilot_scores: std::collections::HashMap = + std::collections::HashMap::new(); for ranked in &pilot_decision.ranked_candidates { pilot_scores.insert(ranked.node_id, ranked.score); } @@ -81,9 +82,7 @@ impl GreedySearch { .collect(); // Sort by merged score - merged.sort_by(|a, b| { - b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal) - }); + merged.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); merged } @@ -140,7 +139,10 @@ impl SearchTree for GreedySearch { // Check if Pilot wants to intervene if p.should_intervene(&state) { - trace!("Pilot intervening at greedy decision point with {} candidates", children.len()); + trace!( + "Pilot intervening at greedy decision point with {} candidates", + children.len() + ); match p.decide(&state).await { decision => { @@ -152,7 +154,12 @@ impl SearchTree for GreedySearch { ); // Merge algorithm scores with Pilot decision - self.merge_with_pilot_decision(tree, &children, &decision, &context.query) + self.merge_with_pilot_decision( + tree, + &children, + &decision, + &context.query, + ) } } } else { diff --git a/src/retrieval/stages/judge.rs b/src/retrieval/stages/judge.rs index 1178f402..8378371a 100644 --- a/src/retrieval/stages/judge.rs +++ b/src/retrieval/stages/judge.rs @@ -10,12 +10,12 @@ use async_trait::async_trait; // Arc is used for async sharing use tracing::{info, warn}; -use crate::util::estimate_tokens; use crate::llm::LlmClient; use crate::retrieval::content::{ContentAggregator, ContentAggregatorConfig}; use crate::retrieval::pipeline::{FailurePolicy, PipelineContext, RetrievalStage, StageOutcome}; use crate::retrieval::sufficiency::{LlmJudge, SufficiencyChecker, ThresholdChecker}; use crate::retrieval::types::{RetrievalResult, RetrieveResponse, SufficiencyLevel}; +use crate::util::estimate_tokens; /// Judge Stage - evaluates retrieval sufficiency. /// @@ -108,7 +108,8 @@ impl JudgeStage { if let Some(ref aggregator) = self.content_aggregator { use crate::retrieval::content::CandidateNode; - let candidates: Vec = ctx.candidates + let candidates: Vec = ctx + .candidates .iter() .map(|c| CandidateNode::new(c.node_id, c.score, c.depth)) .collect(); @@ -116,9 +117,7 @@ impl JudgeStage { let result = aggregator.aggregate(&candidates, &ctx.tree, &ctx.query); info!( "ContentAggregator: {} nodes, {} tokens, avg score {:.2}", - result.nodes_included, - result.tokens_used, - result.avg_score + result.nodes_included, result.tokens_used, result.avg_score ); return (result.content, result.tokens_used); } @@ -167,7 +166,11 @@ impl JudgeStage { } /// Collect content from leaf descendants of a node (excluding the node itself). - fn collect_leaf_content(&self, tree: &crate::document::DocumentTree, node_id: crate::document::NodeId) -> String { + fn collect_leaf_content( + &self, + tree: &crate::document::DocumentTree, + node_id: crate::document::NodeId, + ) -> String { let mut content_parts = Vec::new(); // Start with children, not the node itself diff --git a/src/retrieval/stages/search.rs b/src/retrieval/stages/search.rs index 121378f5..58a6e681 100644 --- a/src/retrieval/stages/search.rs +++ b/src/retrieval/stages/search.rs @@ -13,8 +13,8 @@ use tracing::{info, warn}; use crate::document::DocumentTree; // LlmClient is used via strategy -use crate::retrieval::pilot::Pilot; use crate::retrieval::RetrievalContext; // Legacy context +use crate::retrieval::pilot::Pilot; use crate::retrieval::pipeline::{ CandidateNode, FailurePolicy, PipelineContext, RetrievalStage, SearchAlgorithm, StageOutcome, }; @@ -204,7 +204,11 @@ impl RetrievalStage for SearchStage { "Executing search: algorithm={:?}, beam_width={}, pilot={}", algorithm, config.beam_width, - if self.has_pilot() { "enabled" } else { "disabled" } + if self.has_pilot() { + "enabled" + } else { + "disabled" + } ); // Increment search iteration @@ -233,16 +237,22 @@ impl RetrievalStage for SearchStage { let result = match algorithm { SearchAlgorithm::Greedy => { let search = GreedySearch::new(); - search.search(&ctx.tree, &legacy_ctx, &search_config, pilot_ref).await + search + .search(&ctx.tree, &legacy_ctx, &search_config, pilot_ref) + .await } SearchAlgorithm::Beam => { let search = BeamSearch::new(); - search.search(&ctx.tree, &legacy_ctx, &search_config, pilot_ref).await + search + .search(&ctx.tree, &legacy_ctx, &search_config, pilot_ref) + .await } SearchAlgorithm::Mcts => { // Use beam search as fallback for now let search = BeamSearch::new(); - search.search(&ctx.tree, &legacy_ctx, &search_config, pilot_ref).await + search + .search(&ctx.tree, &legacy_ctx, &search_config, pilot_ref) + .await } }; diff --git a/src/storage/async_workspace.rs b/src/storage/async_workspace.rs index 56c43373..0f72ddc3 100644 --- a/src/storage/async_workspace.rs +++ b/src/storage/async_workspace.rs @@ -36,8 +36,8 @@ use tracing::{debug, info, warn}; use super::backend::{FileBackend, StorageBackend}; use super::cache::DocumentCache; use super::persistence::{PersistedDocument, load_document_from_bytes, save_document_to_bytes}; -use crate::error::Result; use crate::Error; +use crate::error::Result; const META_KEY: &str = "_meta"; const DEFAULT_CACHE_SIZE: usize = 100; @@ -121,8 +121,7 @@ pub struct AsyncWorkspace { impl std::fmt::Debug for AsyncWorkspace { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("AsyncWorkspace") - .finish() + f.debug_struct("AsyncWorkspace").finish() } } @@ -158,14 +157,21 @@ impl AsyncWorkspace { /// Create a new async workspace with custom cache size. pub async fn with_cache_size(path: impl Into, cache_size: usize) -> Result { - Self::with_options(path, AsyncWorkspaceOptions { - cache_size, - ..Default::default() - }).await + Self::with_options( + path, + AsyncWorkspaceOptions { + cache_size, + ..Default::default() + }, + ) + .await } /// Create a new async workspace with custom options. - pub async fn with_options(path: impl Into, options: AsyncWorkspaceOptions) -> Result { + pub async fn with_options( + path: impl Into, + options: AsyncWorkspaceOptions, + ) -> Result { let root = path.into(); let backend = Arc::new(FileBackend::new(&root)?); @@ -229,7 +235,11 @@ impl AsyncWorkspace { .source_path .as_ref() .map(|p| p.to_string_lossy().to_string()), - page_count: if doc.pages.is_empty() { None } else { Some(doc.pages.len()) }, + page_count: if doc.pages.is_empty() { + None + } else { + Some(doc.pages.len()) + }, line_count: doc.meta.line_count, }; @@ -417,10 +427,7 @@ impl AsyncWorkspace { /// Rebuild the meta index from existing documents. fn rebuild_meta_index(inner: &mut AsyncWorkspaceInner) -> Result<()> { let keys = inner.backend.keys()?; - let doc_keys: Vec<_> = keys - .iter() - .filter(|k| k.starts_with("doc:")) - .collect(); + let doc_keys: Vec<_> = keys.iter().filter(|k| k.starts_with("doc:")).collect(); for key in doc_keys { if let Some(bytes) = inner.backend.get(key)? { @@ -436,7 +443,11 @@ impl AsyncWorkspace { .source_path .as_ref() .map(|p| p.to_string_lossy().to_string()), - page_count: if doc.pages.is_empty() { None } else { Some(doc.pages.len()) }, + page_count: if doc.pages.is_empty() { + None + } else { + Some(doc.pages.len()) + }, line_count: doc.meta.line_count, }; inner.meta_index.insert(doc_id, meta_entry); diff --git a/src/storage/backend/file.rs b/src/storage/backend/file.rs index 915d0b4c..0eabc5c8 100644 --- a/src/storage/backend/file.rs +++ b/src/storage/backend/file.rs @@ -10,8 +10,8 @@ use std::sync::RwLock; use tracing::{debug, warn}; use super::StorageBackend; -use crate::error::Result; use crate::Error; +use crate::error::Result; /// File system storage backend. /// @@ -68,9 +68,7 @@ impl FileBackend { /// Convert a key to a file path. fn key_to_path(&self, key: &str) -> PathBuf { // Sanitize key to prevent path traversal - let sanitized = key - .replace("..", "_") - .replace(['/', '\\', ':'], "_"); + let sanitized = key.replace("..", "_").replace(['/', '\\', ':'], "_"); self.root.join(format!("{}.bin", sanitized)) } diff --git a/src/storage/backend/memory.rs b/src/storage/backend/memory.rs index 013c87f9..4844f8e2 100644 --- a/src/storage/backend/memory.rs +++ b/src/storage/backend/memory.rs @@ -39,60 +39,68 @@ impl MemoryBackend { impl StorageBackend for MemoryBackend { fn get(&self, key: &str) -> Result>> { - let data = self.data.read().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let data = self + .data + .read() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; Ok(data.get(key).cloned()) } fn put(&self, key: &str, value: &[u8]) -> Result<()> { - let mut data = self.data.write().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let mut data = self + .data + .write() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; data.insert(key.to_string(), value.to_vec()); Ok(()) } fn delete(&self, key: &str) -> Result { - let mut data = self.data.write().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let mut data = self + .data + .write() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; Ok(data.remove(key).is_some()) } fn exists(&self, key: &str) -> Result { - let data = self.data.read().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let data = self + .data + .read() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; Ok(data.contains_key(key)) } fn keys(&self) -> Result> { - let data = self.data.read().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let data = self + .data + .read() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; Ok(data.keys().cloned().collect()) } fn len(&self) -> Result { - let data = self.data.read().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let data = self + .data + .read() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; Ok(data.len()) } fn clear(&self) -> Result<()> { - let mut data = self.data.write().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let mut data = self + .data + .write() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; data.clear(); Ok(()) } fn batch_put(&self, items: &[(&str, &[u8])]) -> Result<()> { - let mut data = self.data.write().map_err(|_| { - crate::Error::Cache("Memory backend lock poisoned".to_string()) - })?; + let mut data = self + .data + .write() + .map_err(|_| crate::Error::Cache("Memory backend lock poisoned".to_string()))?; for (key, value) in items { data.insert(key.to_string(), value.to_vec()); } diff --git a/src/storage/cache.rs b/src/storage/cache.rs index 4e7e6a57..70f7d4a9 100644 --- a/src/storage/cache.rs +++ b/src/storage/cache.rs @@ -15,14 +15,14 @@ //! - Utilization: Current usage as percentage of capacity use std::num::NonZeroUsize; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; use lru::LruCache; use super::persistence::PersistedDocument; -use crate::error::Result; use crate::Error; +use crate::error::Result; /// Default cache size (number of documents). const DEFAULT_CACHE_SIZE: usize = 100; @@ -67,9 +67,8 @@ impl DocumentCache { #[must_use] pub fn with_capacity(capacity: usize) -> Self { let capacity = capacity.max(1); - let non_zero = NonZeroUsize::new(capacity).unwrap_or_else(|| { - NonZeroUsize::new(DEFAULT_CACHE_SIZE).expect("default is non-zero") - }); + let non_zero = NonZeroUsize::new(capacity) + .unwrap_or_else(|| NonZeroUsize::new(DEFAULT_CACHE_SIZE).expect("default is non-zero")); Self { inner: Mutex::new(LruCache::new(non_zero)), @@ -104,9 +103,7 @@ impl DocumentCache { /// Check if a document is in the cache. pub fn contains(&self, id: &str) -> bool { - self.lock() - .map(|cache| cache.contains(id)) - .unwrap_or(false) + self.lock().map(|cache| cache.contains(id)).unwrap_or(false) } /// Put a document into the cache. @@ -158,9 +155,7 @@ impl DocumentCache { /// Get the number of entries currently in the cache. pub fn len(&self) -> usize { - self.lock() - .map(|cache| cache.len()) - .unwrap_or(0) + self.lock().map(|cache| cache.len()).unwrap_or(0) } /// Check if the cache is empty. @@ -240,9 +235,9 @@ impl DocumentCache { /// Lock the inner cache. fn lock(&self) -> Result>> { - self.inner.lock().map_err(|_| { - Error::Cache("Cache lock poisoned".to_string()) - }) + self.inner + .lock() + .map_err(|_| Error::Cache("Cache lock poisoned".to_string())) } } @@ -272,8 +267,8 @@ pub struct CacheStats { #[cfg(test)] mod tests { use super::*; - use crate::storage::{DocumentMeta, PersistedDocument}; use crate::document::DocumentTree; + use crate::storage::{DocumentMeta, PersistedDocument}; fn create_test_doc(id: &str) -> PersistedDocument { let meta = DocumentMeta::new(id, "Test Doc", "md"); @@ -316,9 +311,15 @@ mod tests { fn test_cache_eviction() { let cache = DocumentCache::with_capacity(2); - cache.put("doc1".to_string(), create_test_doc("doc1")).unwrap(); - cache.put("doc2".to_string(), create_test_doc("doc2")).unwrap(); - cache.put("doc3".to_string(), create_test_doc("doc3")).unwrap(); + cache + .put("doc1".to_string(), create_test_doc("doc1")) + .unwrap(); + cache + .put("doc2".to_string(), create_test_doc("doc2")) + .unwrap(); + cache + .put("doc3".to_string(), create_test_doc("doc3")) + .unwrap(); // doc1 should be evicted (least recently used) assert!(!cache.contains("doc1")); @@ -330,7 +331,9 @@ mod tests { fn test_cache_remove() { let cache = DocumentCache::new(); - cache.put("doc1".to_string(), create_test_doc("doc1")).unwrap(); + cache + .put("doc1".to_string(), create_test_doc("doc1")) + .unwrap(); assert!(cache.contains("doc1")); let removed = cache.remove("doc1").unwrap(); @@ -345,8 +348,12 @@ mod tests { fn test_cache_clear() { let cache = DocumentCache::new(); - cache.put("doc1".to_string(), create_test_doc("doc1")).unwrap(); - cache.put("doc2".to_string(), create_test_doc("doc2")).unwrap(); + cache + .put("doc1".to_string(), create_test_doc("doc1")) + .unwrap(); + cache + .put("doc2".to_string(), create_test_doc("doc2")) + .unwrap(); assert_eq!(cache.len(), 2); @@ -361,10 +368,14 @@ mod tests { assert_eq!(cache.utilization(), 0.0); - cache.put("doc1".to_string(), create_test_doc("doc1")).unwrap(); + cache + .put("doc1".to_string(), create_test_doc("doc1")) + .unwrap(); assert!((cache.utilization() - 0.1).abs() < 0.01); - cache.put("doc2".to_string(), create_test_doc("doc2")).unwrap(); + cache + .put("doc2".to_string(), create_test_doc("doc2")) + .unwrap(); assert!((cache.utilization() - 0.2).abs() < 0.01); } } diff --git a/src/storage/codec.rs b/src/storage/codec.rs index 3fcfd055..4c7e864e 100644 --- a/src/storage/codec.rs +++ b/src/storage/codec.rs @@ -26,12 +26,12 @@ use std::fmt::Debug; use std::io::{Read, Write}; +use flate2::Compression; use flate2::read::GzDecoder; use flate2::write::GzEncoder; -use flate2::Compression; -use crate::error::Result; use crate::Error; +use crate::error::Result; /// Codec trait for compression/decompression. pub trait Codec: Debug + Send + Sync { @@ -120,8 +120,12 @@ impl Default for GzipCodec { impl Codec for GzipCodec { fn encode(&self, data: &[u8]) -> Result> { let mut encoder = GzEncoder::new(Vec::new(), Compression::new(self.level)); - encoder.write_all(data).map_err(|e| Error::Parse(format!("Gzip encode error: {}", e)))?; - encoder.finish().map_err(|e| Error::Parse(format!("Gzip finish error: {}", e))) + encoder + .write_all(data) + .map_err(|e| Error::Parse(format!("Gzip encode error: {}", e)))?; + encoder + .finish() + .map_err(|e| Error::Parse(format!("Gzip finish error: {}", e))) } fn decode(&self, data: &[u8]) -> Result> { diff --git a/src/storage/lock.rs b/src/storage/lock.rs index 66a65d46..6318e8a5 100644 --- a/src/storage/lock.rs +++ b/src/storage/lock.rs @@ -15,8 +15,8 @@ use std::fs::{File, OpenOptions}; use std::path::Path; -use crate::error::Result; use crate::Error; +use crate::error::Result; /// A file lock that is automatically released when dropped. /// @@ -88,7 +88,7 @@ impl FileLock { { use std::os::windows::fs::OpenOptionsExt; use windows_sys::Win32::Storage::FileSystem::{ - LockFileEx, LOCKFILE_EXCLUSIVE_LOCK, LOCKFILE_FAIL_IMMEDIATELY, + LOCKFILE_EXCLUSIVE_LOCK, LOCKFILE_FAIL_IMMEDIATELY, LockFileEx, }; let handle = std::os::windows::io::AsRawHandle::as_raw_handle(&file); @@ -97,7 +97,11 @@ impl FileLock { let result = unsafe { LockFileEx( handle, - if exclusive { LOCKFILE_EXCLUSIVE_LOCK } else { 0 } | LOCKFILE_FAIL_IMMEDIATELY, + if exclusive { + LOCKFILE_EXCLUSIVE_LOCK + } else { + 0 + } | LOCKFILE_FAIL_IMMEDIATELY, 0, 0xFFFFFFFF, 0xFFFFFFFF, @@ -136,11 +140,11 @@ impl FileLock { exclusive: bool, ) -> Result> { match Self::try_lock(&path.into(), exclusive) { - Ok(lock) => Ok(Some(lock)), - Err(Error::WorkspaceLocked) => Ok(None), - Err(e) => Err(e), - } + Ok(lock) => Ok(Some(lock)), + Err(Error::WorkspaceLocked) => Ok(None), + Err(e) => Err(e), } + } /// Check if the lock file is locked by another process. /// diff --git a/src/storage/migration.rs b/src/storage/migration.rs index b73c0f6e..1711169b 100644 --- a/src/storage/migration.rs +++ b/src/storage/migration.rs @@ -35,8 +35,8 @@ use std::collections::HashMap; use tracing::{debug, info, warn}; -use crate::error::Result; use crate::Error; +use crate::error::Result; /// Current data format version. pub const CURRENT_VERSION: u32 = 1; @@ -133,10 +133,7 @@ impl Migrator { /// Register a migration. pub fn register(&mut self, migration: Box) { let key = (migration.from_version(), migration.to_version()); - debug!( - "Registering migration: v{} -> v{}", - key.0, key.1 - ); + debug!("Registering migration: v{} -> v{}", key.0, key.1); self.migrations.insert(key, migration); } @@ -213,11 +210,14 @@ impl Migrator { } // Find migration path - let path = self.find_migration_path(from_version, to_version) - .ok_or_else(|| Error::VersionMismatch(format!( - "No migration path from v{} to v{}", - from_version, to_version - )))?; + let path = self + .find_migration_path(from_version, to_version) + .ok_or_else(|| { + Error::VersionMismatch(format!( + "No migration path from v{} to v{}", + from_version, to_version + )) + })?; if path.len() < 2 { return Ok(data.to_vec()); @@ -233,17 +233,20 @@ impl Migrator { for next_version in path.iter().skip(1) { let key = (current_version, *next_version); - let migration = self.migrations.get(&key) - .ok_or_else(|| Error::VersionMismatch(format!( + let migration = self.migrations.get(&key).ok_or_else(|| { + Error::VersionMismatch(format!( "Missing migration from v{} to v{}", current_version, next_version - )))?; + )) + })?; let ctx = MigrationContext::new(current_version, *next_version); debug!( "Applying migration: v{} -> v{} ({})", - current_version, next_version, migration.description() + current_version, + next_version, + migration.description() ); current_data = migration.migrate(¤t_data, &ctx)?; @@ -316,8 +319,7 @@ mod tests { #[test] fn test_migration_context() { - let ctx = MigrationContext::new(1, 2) - .with_metadata("key", "value"); + let ctx = MigrationContext::new(1, 2).with_metadata("key", "value"); assert_eq!(ctx.from_version, 1); assert_eq!(ctx.to_version, 2); diff --git a/src/storage/mod.rs b/src/storage/mod.rs index f8d97b1f..b2aee4a9 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -39,18 +39,16 @@ mod persistence; mod workspace; // Re-export main types +pub use async_workspace::{AsyncDocumentMetaEntry, AsyncWorkspace, AsyncWorkspaceOptions}; pub use backend::{FileBackend, MemoryBackend, StorageBackend}; pub use cache::DocumentCache; pub use codec::{Codec, GzipCodec, IdentityCodec, codec_from_config}; -pub use migration::{Migration, MigrationContext, Migrator, CURRENT_VERSION}; pub use lock::{FileLock, ScopedLock}; +pub use migration::{CURRENT_VERSION, Migration, MigrationContext, Migrator}; pub use persistence::{ - DocumentMeta, PageContent, PersistedDocument, - load_document, load_document_from_bytes, load_document_with_options, - load_index, load_index_from_bytes, load_index_with_options, - save_document, save_document_to_bytes, save_document_with_options, + DocumentMeta, PageContent, PersistedDocument, PersistenceOptions, load_document, + load_document_from_bytes, load_document_with_options, load_index, load_index_from_bytes, + load_index_with_options, save_document, save_document_to_bytes, save_document_with_options, save_index, save_index_to_bytes, save_index_with_options, - PersistenceOptions, }; -pub use async_workspace::{AsyncDocumentMetaEntry, AsyncWorkspace, AsyncWorkspaceOptions}; pub use workspace::{DocumentMetaEntry, Workspace, WorkspaceOptions}; diff --git a/src/storage/persistence.rs b/src/storage/persistence.rs index 245f33a6..1d63cd47 100644 --- a/src/storage/persistence.rs +++ b/src/storage/persistence.rs @@ -9,15 +9,15 @@ //! - **Checksum verification**: SHA-256 checksums for data integrity //! - **Version header**: Format version for future migrations -use sha2::{Digest, Sha256}; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use std::fs::File; use std::io::{BufReader, BufWriter, Read, Write}; use std::path::{Path, PathBuf}; +use crate::Error; use crate::document::DocumentTree; use crate::error::Result; -use crate::Error; /// Current format version for persisted documents. const FORMAT_VERSION: u32 = 1; @@ -209,8 +209,7 @@ pub fn save_document_with_options( options: &PersistenceOptions, ) -> Result<()> { // Serialize the payload first - let payload_bytes = serde_json::to_vec(doc) - .map_err(|e| Error::Serialization(e.to_string()))?; + let payload_bytes = serde_json::to_vec(doc).map_err(|e| Error::Serialization(e.to_string()))?; // Calculate checksum let checksum = calculate_checksum(&payload_bytes); @@ -223,8 +222,8 @@ pub fn save_document_with_options( }; // Serialize wrapper - let json = serde_json::to_string_pretty(&wrapper) - .map_err(|e| Error::Serialization(e.to_string()))?; + let json = + serde_json::to_string_pretty(&wrapper).map_err(|e| Error::Serialization(e.to_string()))?; if options.atomic_writes { // Atomic write: write to temp file, then rename @@ -280,9 +279,7 @@ pub fn load_document_with_options( options: &PersistenceOptions, ) -> Result { if !path.exists() { - return Err(Error::DocumentNotFound( - path.display().to_string() - )); + return Err(Error::DocumentNotFound(path.display().to_string())); } let file = File::open(path).map_err(Error::Io)?; @@ -330,8 +327,8 @@ pub fn save_index_with_options( options: &PersistenceOptions, ) -> Result<()> { // Serialize payload - let payload_bytes = serde_json::to_vec(entries) - .map_err(|e| Error::Serialization(e.to_string()))?; + let payload_bytes = + serde_json::to_vec(entries).map_err(|e| Error::Serialization(e.to_string()))?; let checksum = calculate_checksum(&payload_bytes); @@ -341,8 +338,8 @@ pub fn save_index_with_options( payload: entries.to_vec(), }; - let json = serde_json::to_string_pretty(&wrapper) - .map_err(|e| Error::Serialization(e.to_string()))?; + let json = + serde_json::to_string_pretty(&wrapper).map_err(|e| Error::Serialization(e.to_string()))?; if options.atomic_writes { let temp_path = path.with_extension("tmp"); @@ -424,8 +421,7 @@ pub fn load_index_with_options( /// This is useful for storage backends that work with byte arrays. pub fn save_document_to_bytes(doc: &PersistedDocument) -> Result> { // Serialize the payload first - let payload_bytes = serde_json::to_vec(doc) - .map_err(|e| Error::Serialization(e.to_string()))?; + let payload_bytes = serde_json::to_vec(doc).map_err(|e| Error::Serialization(e.to_string()))?; // Calculate checksum let checksum = calculate_checksum(&payload_bytes); @@ -438,8 +434,7 @@ pub fn save_document_to_bytes(doc: &PersistedDocument) -> Result> { }; // Serialize wrapper - serde_json::to_vec(&wrapper) - .map_err(|e| Error::Serialization(e.to_string())) + serde_json::to_vec(&wrapper).map_err(|e| Error::Serialization(e.to_string())) } /// Deserialize a document from bytes. @@ -486,8 +481,8 @@ pub fn load_document_from_bytes_with_options( /// Serialize an index to bytes. pub fn save_index_to_bytes(entries: &[DocumentMeta]) -> Result> { - let payload_bytes = serde_json::to_vec(entries) - .map_err(|e| Error::Serialization(e.to_string()))?; + let payload_bytes = + serde_json::to_vec(entries).map_err(|e| Error::Serialization(e.to_string()))?; let checksum = calculate_checksum(&payload_bytes); @@ -497,8 +492,7 @@ pub fn save_index_to_bytes(entries: &[DocumentMeta]) -> Result> { payload: entries.to_vec(), }; - serde_json::to_vec(&wrapper) - .map_err(|e| Error::Serialization(e.to_string())) + serde_json::to_vec(&wrapper).map_err(|e| Error::Serialization(e.to_string())) } /// Deserialize an index from bytes. @@ -620,7 +614,7 @@ mod tests { // Change the checksum value but keep the payload intact let corrupted = content.replace( &calculate_checksum(&serde_json::to_vec(&doc).unwrap()), - "0000000000000000000000000000000000000000000000000000000000000000" + "0000000000000000000000000000000000000000000000000000000000000000", ); std::fs::write(&path, corrupted).unwrap(); diff --git a/src/storage/workspace.rs b/src/storage/workspace.rs index c13e99ec..610592e1 100644 --- a/src/storage/workspace.rs +++ b/src/storage/workspace.rs @@ -38,8 +38,8 @@ use super::backend::{FileBackend, StorageBackend}; use super::cache::DocumentCache; use super::lock::FileLock; use super::persistence::{PersistedDocument, load_document_from_bytes, save_document_to_bytes}; -use crate::error::Result; use crate::Error; +use crate::error::Result; const META_KEY: &str = "_meta"; const LOCK_FILE: &str = ".workspace.lock"; @@ -168,10 +168,13 @@ impl Workspace { /// Create a new workspace with custom LRU cache size. pub fn with_cache_size(path: impl Into, cache_size: usize) -> Result { - Self::with_options(path, WorkspaceOptions { - cache_size, - ..Default::default() - }) + Self::with_options( + path, + WorkspaceOptions { + cache_size, + ..Default::default() + }, + ) } /// Create a new workspace with custom options. @@ -210,10 +213,13 @@ impl Workspace { path: impl Into + Clone, cache_size: usize, ) -> Result { - Self::open_with_options(path, WorkspaceOptions { - cache_size, - ..Default::default() - }) + Self::open_with_options( + path, + WorkspaceOptions { + cache_size, + ..Default::default() + }, + ) } /// Open with custom options. @@ -290,7 +296,11 @@ impl Workspace { .source_path .as_ref() .map(|p| p.to_string_lossy().to_string()), - page_count: if doc.pages.is_empty() { None } else { Some(doc.pages.len()) }, + page_count: if doc.pages.is_empty() { + None + } else { + Some(doc.pages.len()) + }, line_count: doc.meta.line_count, }; @@ -426,10 +436,7 @@ impl Workspace { /// Rebuild the meta index from existing documents. fn rebuild_meta_index(&mut self) -> Result<()> { let keys = self.backend.keys()?; - let doc_keys: Vec<_> = keys - .iter() - .filter(|k| k.starts_with("doc:")) - .collect(); + let doc_keys: Vec<_> = keys.iter().filter(|k| k.starts_with("doc:")).collect(); for key in doc_keys { if let Some(bytes) = self.backend.get(key)? { @@ -445,7 +452,11 @@ impl Workspace { .source_path .as_ref() .map(|p| p.to_string_lossy().to_string()), - page_count: if doc.pages.is_empty() { None } else { Some(doc.pages.len()) }, + page_count: if doc.pages.is_empty() { + None + } else { + Some(doc.pages.len()) + }, line_count: doc.meta.line_count, }; self.meta_index.insert(doc_id, meta_entry); @@ -455,10 +466,7 @@ impl Workspace { if !self.meta_index.is_empty() { self.save_meta_index()?; - info!( - "Rebuilt index from {} document(s)", - self.meta_index.len() - ); + info!("Rebuilt index from {} document(s)", self.meta_index.len()); } Ok(()) diff --git a/src/util/mod.rs b/src/util/mod.rs index 9ec7184e..6af20f59 100644 --- a/src/util/mod.rs +++ b/src/util/mod.rs @@ -14,8 +14,8 @@ mod timing; mod token; pub use format::{ - clean_whitespace, format_bytes, format_number, format_percent, indent, line_count, - truncate, truncate_words, word_count, + clean_whitespace, format_bytes, format_number, format_percent, indent, line_count, truncate, + truncate_words, word_count, }; -pub use timing::{format_duration, format_duration_compact, Timer}; +pub use timing::{Timer, format_duration, format_duration_compact}; pub use token::{estimate_tokens, estimate_tokens_batch, estimate_tokens_fast};