diff --git a/Cargo.lock b/Cargo.lock index bc6a6d7..85ca93d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -130,7 +130,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -141,7 +141,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -176,9 +176,9 @@ checksum = "f02882884d3e1bc524fb12c79f107f6ad0e1cfd498c536ffb494301740995dfe" [[package]] name = "arrow" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3bd47f2a6ddc39244bd722a27ee5da66c03369d087b9e024eafdb03e98b98ea7" +checksum = "7c14b3d39f306bc28fd639d59f06e17a0f377d0021e1b7e9054e4d6fedc98774" dependencies = [ "arrow-arith", "arrow-array", @@ -197,9 +197,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c7bbd679c5418b8639b92be01f361d60013c4906574b578b77b63c78356594c" +checksum = "ce2961626677665b2195eb59242af4c7befe7b8737ca2050295389362380104e" dependencies = [ "arrow-array", "arrow-buffer", @@ -211,9 +211,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c8a4ab47b3f3eac60f7fd31b81e9028fda018607bcc63451aca4f2b755269862" +checksum = "1e5f6adeffdf587d7a31db5d2266189624b526730cd3627f9ff9fedae97ad584" dependencies = [ "ahash", "arrow-buffer", @@ -221,7 +221,7 @@ dependencies = [ "arrow-schema", "chrono", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "num-complex", "num-integer", "num-traits", @@ -229,21 +229,21 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0d18b89b4c4f4811d0858175e79541fe98e33e18db3b011708bc287b1240593f" +checksum = "097d193003ce7995d5d087089069ec2a6e0187faf5a6f8c9f38af2645d987182" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.5.1", "num-traits", ] [[package]] name = "arrow-cast" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "722b5c41dd1d14d0a879a1bce92c6fe33f546101bb2acce57a209825edd075b3" +checksum = "635c9c635668ad26adf76cce8fb276c4be7cf06e63bd516de7da514f9680ee53" dependencies = [ "arrow-array", "arrow-buffer", @@ -252,7 +252,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64 0.22.1", + "base64 0.23.1", "chrono", "half", "lexical-core", @@ -262,9 +262,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27ddb80a4848e03b1655af496d5ac2563a779e5742fcb48f2ca2e089c9cd2197" +checksum = "4c2ebf8d631e79b02c16cf5ae860561272c26024ec88fce389a56aaddd558e86" dependencies = [ "arrow-array", "arrow-cast", @@ -277,9 +277,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c1683705c63dcf0d18972759eda48489028cbbff67af7d6bef2c6b7b74ab778a" +checksum = "9ba2f832eaeca24b8f26143dba750e42ee4ab51cf7d65e701ca9607cfda9f358" dependencies = [ "arrow-buffer", "arrow-schema", @@ -290,9 +290,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cf72d04c07229fbf4dbebe7145cac37d7cf7ec582fe705c6b92cb314af096ab" +checksum = "dcc41681ea80f521df14c36725b74d4c60702c47f0793af2be469c04527e2599" dependencies = [ "arrow-array", "arrow-buffer", @@ -304,15 +304,16 @@ dependencies = [ [[package]] name = "arrow-json" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a84a905f41fedfcd7679813c89a61dc369c0f932b27aa8dcc6aa051cc781a97d" +checksum = "a2f57d7a81969f24ccf80809587b76c09897e6f829d2d65a5976bfb3218851f1" dependencies = [ "arrow-array", "arrow-buffer", "arrow-cast", - "arrow-data", + "arrow-ord", "arrow-schema", + "arrow-select", "chrono", "half", "indexmap", @@ -328,9 +329,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "082342947d4e5a2bcccf029a0a0397e21cb3bb8421edd9571d34fb5dd2670256" +checksum = "2c900759f3bd8354fd4196bc4403eee846894dc2adf66b4225472006a0bf18c5" dependencies = [ "arrow-array", "arrow-buffer", @@ -341,9 +342,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3a931b520a2a5e22033e01a6f2486b4cdc26f9106b759abeebc320f125e94d7" +checksum = "f4c6425032e28266e3fc4ff680805e57e670d6ea92473043f3e65b7ed6ac79f2" dependencies = [ "arrow-array", "arrow-buffer", @@ -354,15 +355,15 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4cf0d4a6609679e03002167a61074a21d7b1ad9ea65e462b2c0a97f8a3b2bc6" +checksum = "10fab8d4563491417ba801fab29d205104d20d4bdf37bda6cd1cf425cff598cd" [[package]] name = "arrow-select" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0b320d86a9806923663bb0fd9baa65ecaba81cb0cd77ff8c1768b9716b4ef891" +checksum = "fc58569193c2525915f3cc6310edba3792f1200f65d6e9ed330aa33e691493b8" dependencies = [ "ahash", "arrow-array", @@ -374,9 +375,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b493e99162e5764077e7823e50ba284858d365922631c7aaefe9487b1abd02c2" +checksum = "2e0813f3c35c1cfea65e14c20a953440f7783c088b7ad2d0db162ccdeefcec14" dependencies = [ "arrow-array", "arrow-buffer", @@ -507,6 +508,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" @@ -519,7 +526,7 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a6773ddc0eafc0e509fb60e48dff7f450f8e674a0686ae8605e8d9901bd5eefa" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", "serde", @@ -1020,9 +1027,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -1257,7 +1264,7 @@ dependencies = [ "libc", "option-ext", "redox_users", - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -1365,7 +1372,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1847,9 +1854,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.15" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6cb093c84e8bd9b188d4c4a8cb6579fc016968d14c99882163cd3ff402a4f155" +checksum = "a9f37a958b41b3b19ee2707c06439c0e9e547e847223eb791ecb0cb821c65e27" dependencies = [ "atomic-waker", "bytes", @@ -2281,12 +2288,6 @@ dependencies = [ "web-time", ] -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - [[package]] name = "ipnet" version = "2.12.0" @@ -2555,18 +2556,18 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.12.2" +version = "0.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90071f8077f8e40adfc4b7fe9cd495ce316263f19e75c2211eeff3fdf475a3d9" +checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" dependencies = [ "twox-hash", ] [[package]] name = "lz4_flex" -version = "0.13.1" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7ef0d4ed8669f8f8826eb00dc878084aa8f253506c4fd5e8f58f5bce72ddb97e" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" dependencies = [ "twox-hash", ] @@ -2855,7 +2856,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.59.0", ] [[package]] @@ -2868,6 +2869,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -3011,15 +3022,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" -[[package]] -name = "ordered-float" -version = "2.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" -dependencies = [ - "num-traits", -] - [[package]] name = "ordered-float" version = "5.3.0" @@ -3090,34 +3092,31 @@ dependencies = [ [[package]] name = "parquet" -version = "57.3.1" +version = "59.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e832c6aa20310fc6de7ea5a3f4e20d34fd83e3b43229d32b81ffe5c14d74692" +checksum = "ff322f54b1a0f9288e614ed1f2d329b380af5476420db19f46ffb865e1163d73" dependencies = [ "ahash", "arrow-array", "arrow-buffer", - "arrow-cast", "arrow-data", "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.1", + "base64 0.23.1", "brotli", "bytes", "chrono", "flate2", "half", - "hashbrown 0.16.1", - "lz4_flex 0.12.2", - "num-bigint", + "hashbrown 0.17.1", + "lz4_flex 0.14.0", + "num-bigint 0.5.1", "num-integer", "num-traits", - "paste", "seq-macro", "simdutf8", "snap", - "thrift", "twox-hash", "zstd", ] @@ -3565,9 +3564,9 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" dependencies = [ "bytes", "getrandom 0.3.4", @@ -3911,7 +3910,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4226,7 +4225,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.60.2", ] [[package]] @@ -4378,7 +4377,7 @@ dependencies = [ "getrandom 0.4.3", "once_cell", "rustix 1.1.4", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4439,17 +4438,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float 2.10.1", -] - [[package]] name = "tiny-keccak" version = "2.0.2" @@ -5019,7 +5007,7 @@ dependencies = [ "memmap2", "metal", "ndarray 0.16.1", - "ordered-float 5.3.0", + "ordered-float", "ort", "parquet", "pollster", @@ -5343,7 +5331,7 @@ dependencies = [ "ndk-sys", "objc", "once_cell", - "ordered-float 5.3.0", + "ordered-float", "parking_lot", "portable-atomic", "portable-atomic-util", @@ -5401,7 +5389,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index c5f69b9..4075694 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -58,8 +58,8 @@ serde-wasm-bindgen = { version = "0.6", optional = true } js-sys = { version = "0.3", optional = true } web-sys = { version = "0.3", optional = true, features = ["console", "Window", "Storage", "IdbFactory", "IdbDatabase", "IdbObjectStore", "IdbTransaction"] } console_error_panic_hook = { version = "0.1", optional = true } -parquet = { version = "57", optional = true } -arrow = { version = "57", optional = true } +parquet = { version = "59", optional = true } +arrow = { version = "59", optional = true } candle-core = { version = "0.9", optional = true } candle-nn = { version = "0.9", optional = true } candle-transformers = { version = "0.9", optional = true } diff --git a/benches/openai_embeddings_bench.rs b/benches/openai_embeddings_bench.rs index 12527db..ec69684 100644 --- a/benches/openai_embeddings_bench.rs +++ b/benches/openai_embeddings_bench.rs @@ -6,7 +6,7 @@ // // Run with: cargo bench --features "embeddings,openai-embeddings" --bench openai_embeddings_bench -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use vecstore::embeddings::openai_backend::{OpenAIEmbedding, OpenAIModel}; // Helper to create embedder diff --git a/benches/rag_benchmarks.rs b/benches/rag_benchmarks.rs index cf79bb0..9df80d2 100644 --- a/benches/rag_benchmarks.rs +++ b/benches/rag_benchmarks.rs @@ -1,9 +1,9 @@ -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use std::collections::HashMap; use vecstore::{ + Metadata, Query, VecStore, rag_utils::MultiQueryRetrieval, text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, - Metadata, Query, VecStore, }; // Mock embedding function for benchmarks diff --git a/benches/simd_bench.rs b/benches/simd_bench.rs index cdbb5ad..af52421 100644 --- a/benches/simd_bench.rs +++ b/benches/simd_bench.rs @@ -1,4 +1,4 @@ -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use vecstore::simd::*; fn generate_random_vector(dim: usize) -> Vec { diff --git a/benches/vecstore_bench.rs b/benches/vecstore_bench.rs index 3ad4d4f..f6b8266 100644 --- a/benches/vecstore_bench.rs +++ b/benches/vecstore_bench.rs @@ -1,6 +1,6 @@ -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use std::collections::HashMap; -use vecstore::{make_record, FilterExpr, FilterOp, Metadata, Query, VecStore}; +use vecstore::{FilterExpr, FilterOp, Metadata, Query, VecStore, make_record}; fn setup_store_with_data(n: usize, dim: usize) -> (tempfile::TempDir, VecStore) { let temp_dir = tempfile::tempdir().unwrap(); diff --git a/benches/wasm_hnsw_bench.rs b/benches/wasm_hnsw_bench.rs index 44de542..ccdf80d 100644 --- a/benches/wasm_hnsw_bench.rs +++ b/benches/wasm_hnsw_bench.rs @@ -5,15 +5,15 @@ // // Run with: cargo bench --bench wasm_hnsw_bench -use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use rand::Rng; -use vecstore::store::wasm_hnsw::WasmHnsw; use vecstore::Distance; +use vecstore::store::wasm_hnsw::WasmHnsw; fn generate_random_vectors(n: usize, dim: usize) -> Vec> { let mut rng = rand::thread_rng(); (0..n) - .map(|_| (0..dim).map(|_| rng.gen::()).collect()) + .map(|_| (0..dim).map(|_| rng.r#gen::()).collect()) .collect() } diff --git a/examples/01_basic_rag.rs b/examples/01_basic_rag.rs index 165d88d..6f6f145 100644 --- a/examples/01_basic_rag.rs +++ b/examples/01_basic_rag.rs @@ -12,8 +12,8 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ - text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, Metadata, Query, VecStore, + text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, }; fn main() -> Result<()> { diff --git a/examples/03_web_scraping_rag.rs b/examples/03_web_scraping_rag.rs index 915f227..fb0e0de 100644 --- a/examples/03_web_scraping_rag.rs +++ b/examples/03_web_scraping_rag.rs @@ -11,8 +11,8 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ - text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, Metadata, Query, VecStore, + text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, }; fn main() -> Result<()> { diff --git a/examples/04_code_search.rs b/examples/04_code_search.rs index 99c4dd4..5529001 100644 --- a/examples/04_code_search.rs +++ b/examples/04_code_search.rs @@ -10,8 +10,8 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ - text_splitter::{CodeTextSplitter, TextSplitter}, Metadata, Query, VecStore, + text_splitter::{CodeTextSplitter, TextSplitter}, }; fn main() -> Result<()> { diff --git a/examples/06_reranking_pipeline.rs b/examples/06_reranking_pipeline.rs index ad71ded..d7f335a 100644 --- a/examples/06_reranking_pipeline.rs +++ b/examples/06_reranking_pipeline.rs @@ -11,8 +11,8 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ - text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, Metadata, Neighbor, Query, VecStore, + text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, }; fn main() -> Result<()> { diff --git a/examples/07_multi_query_rag.rs b/examples/07_multi_query_rag.rs index b6410df..8c0185e 100644 --- a/examples/07_multi_query_rag.rs +++ b/examples/07_multi_query_rag.rs @@ -10,9 +10,9 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ + Metadata, Query, VecStore, rag_utils::MultiQueryRetrieval, text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, - Metadata, Query, VecStore, }; fn main() -> Result<()> { diff --git a/examples/08_conversation_rag.rs b/examples/08_conversation_rag.rs index 4ff903d..0cd9d28 100644 --- a/examples/08_conversation_rag.rs +++ b/examples/08_conversation_rag.rs @@ -8,9 +8,9 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ + Metadata, Query, VecStore, rag_utils::ConversationMemory, text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, - Metadata, Query, VecStore, }; fn main() -> Result<()> { diff --git a/examples/09_evaluation_demo.rs b/examples/09_evaluation_demo.rs index 263adae..eb024f5 100644 --- a/examples/09_evaluation_demo.rs +++ b/examples/09_evaluation_demo.rs @@ -9,8 +9,8 @@ use anyhow::Result; use std::collections::HashMap; use vecstore::{ - text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, Metadata, Query, VecStore, + text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, }; fn main() -> Result<()> { @@ -50,7 +50,10 @@ fn main() -> Result<()> { "What is VecStore?", "VecStore is a high-performance vector database built in Rust.", ), - ("How fast is VecStore?", "Rust avoids the Python interpreter in tight loops."), + ( + "How fast is VecStore?", + "Rust avoids the Python interpreter in tight loops.", + ), ]; let mut total_score = 0.0; diff --git a/examples/10_production_rag.rs b/examples/10_production_rag.rs index 6a03777..37007a4 100644 --- a/examples/10_production_rag.rs +++ b/examples/10_production_rag.rs @@ -12,8 +12,8 @@ use anyhow::{Context, Result}; use std::collections::HashMap; use std::time::Instant; use vecstore::{ - text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, Metadata, Query, VecDatabase, + text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}, }; fn main() -> Result<()> { diff --git a/examples/bulk_migration_demo.rs b/examples/bulk_migration_demo.rs index 5adc6b2..aef7c6d 100644 --- a/examples/bulk_migration_demo.rs +++ b/examples/bulk_migration_demo.rs @@ -10,8 +10,8 @@ use anyhow::Result; use std::fs; use tempfile::TempDir; -use vecstore::bulk_migration::{ChromaDBMigration, MigrationConfig, PineconeMigration}; use vecstore::VecStore; +use vecstore::bulk_migration::{ChromaDBMigration, MigrationConfig, PineconeMigration}; fn main() -> Result<()> { println!("šŸ”„ VecStore Bulk Migration Demo\n"); diff --git a/examples/compression.rs b/examples/compression.rs index c2c63af..7f26e4d 100644 --- a/examples/compression.rs +++ b/examples/compression.rs @@ -17,8 +17,8 @@ //! ``` use vecstore::compression::{ - decode_rle, encode_rle, encode_varint, CompressedNeighborList, - CompressionConfig, CompressionLevel, CompressionStats, + CompressedNeighborList, CompressionConfig, CompressionLevel, CompressionStats, decode_rle, + encode_rle, encode_varint, }; fn main() { diff --git a/examples/dim_reduction_demo.rs b/examples/dim_reduction_demo.rs index b7699f2..3eb74c0 100644 --- a/examples/dim_reduction_demo.rs +++ b/examples/dim_reduction_demo.rs @@ -7,7 +7,7 @@ //! - Reconstruct original vectors use anyhow::Result; -use vecstore::dim_reduction::{ReductionStats, PCA}; +use vecstore::dim_reduction::{PCA, ReductionStats}; fn generate_high_dim_data() -> Vec> { let mut vectors = Vec::new(); diff --git a/examples/distance_metrics_demo.rs b/examples/distance_metrics_demo.rs index 6f25976..175d184 100644 --- a/examples/distance_metrics_demo.rs +++ b/examples/distance_metrics_demo.rs @@ -12,8 +12,8 @@ use std::collections::HashMap; use vecstore::{ - hamming_distance_simd, jaccard_distance_simd, manhattan_distance_simd, Distance, Metadata, - VecStore, + Distance, Metadata, VecStore, hamming_distance_simd, jaccard_distance_simd, + manhattan_distance_simd, }; fn main() -> anyhow::Result<()> { diff --git a/examples/distributed_indexing.rs b/examples/distributed_indexing.rs index f83e254..7c7ad08 100644 --- a/examples/distributed_indexing.rs +++ b/examples/distributed_indexing.rs @@ -23,8 +23,8 @@ //! ``` use vecstore::distributed::{ - ConsistencyLevel, ConsistentHashRing, DistributedConfig, DistributedStore, - ReplicationStrategy, ShardingStrategy, + ConsistencyLevel, ConsistentHashRing, DistributedConfig, DistributedStore, ReplicationStrategy, + ShardingStrategy, }; fn main() { diff --git a/examples/embedding_integration_demo.rs b/examples/embedding_integration_demo.rs index 12c5331..507d3db 100644 --- a/examples/embedding_integration_demo.rs +++ b/examples/embedding_integration_demo.rs @@ -17,8 +17,8 @@ #[cfg(feature = "embeddings")] use vecstore::{ - embeddings::{EmbeddingCollection, SimpleEmbedder, TextEmbedder}, Metadata, VecDatabase, + embeddings::{EmbeddingCollection, SimpleEmbedder, TextEmbedder}, }; #[cfg(feature = "embeddings")] diff --git a/examples/fuzzy_search.rs b/examples/fuzzy_search.rs index 09d2aff..883e4fa 100644 --- a/examples/fuzzy_search.rs +++ b/examples/fuzzy_search.rs @@ -18,8 +18,8 @@ //! ``` use vecstore::fuzzy::{ - damerau_levenshtein_distance, levenshtein_distance, similarity_score, suggest_corrections, - BKTree, FuzzyMatcher, + BKTree, FuzzyMatcher, damerau_levenshtein_distance, levenshtein_distance, similarity_score, + suggest_corrections, }; fn main() { diff --git a/examples/gpu_acceleration.rs b/examples/gpu_acceleration.rs index caf4a24..4cd7e84 100644 --- a/examples/gpu_acceleration.rs +++ b/examples/gpu_acceleration.rs @@ -294,10 +294,10 @@ fn main() { }; println!("{:<20} {:>12.2} {:>14.2}x", name, duration_ms, speedup); - } + }, Err(_) => { println!("{:<20} {:>12} {:>15}", name, "N/A", "N/A"); - } + }, } } diff --git a/examples/graph_visualization.rs b/examples/graph_visualization.rs index b2fdd0b..7b90e80 100644 --- a/examples/graph_visualization.rs +++ b/examples/graph_visualization.rs @@ -143,14 +143,14 @@ fn main() -> Result<()> { println!(" 2. Open graph.json in a D3.js visualization tool"); println!(" 3. Load graph_cytoscape.json in Cytoscape.js"); println!(" 4. Read graph_stats.txt for detailed metrics"); - } + }, Err(e) => { println!("āš ļø Graph visualization not available: {}", e); println!("\nNote: Graph visualization is currently only supported for WASM builds."); println!("To use visualization:"); println!(" 1. Build with: cargo build --target wasm32-unknown-unknown"); println!(" 2. Or use the WASM-based backend"); - } + }, } Ok(()) diff --git a/examples/langchain_rag_pipeline.rs b/examples/langchain_rag_pipeline.rs index 8dcdc08..53c1468 100644 --- a/examples/langchain_rag_pipeline.rs +++ b/examples/langchain_rag_pipeline.rs @@ -11,8 +11,8 @@ use anyhow::Result; use vecstore::error::VecStoreError; use vecstore::{ - print_health_report, Document, HealthChecker, LangChainVectorStore, RetrieverConfig, VecStore, - VectorStoreRetriever, + Document, HealthChecker, LangChainVectorStore, RetrieverConfig, VecStore, VectorStoreRetriever, + print_health_report, }; fn main() -> Result<()> { diff --git a/examples/metadata_indexing_demo.rs b/examples/metadata_indexing_demo.rs index 65efdd7..c8e6ff9 100644 --- a/examples/metadata_indexing_demo.rs +++ b/examples/metadata_indexing_demo.rs @@ -183,7 +183,9 @@ fn main() -> Result<()> { println!(" • Range queries: O(log N) instead of O(N)"); println!(" • Equality queries: O(1) instead of O(N)"); println!(" • Compound queries: Fast intersection of index results"); - println!(" • Potential speedup: orders of magnitude for large datasets (measure on your workload)"); + println!( + " • Potential speedup: orders of magnitude for large datasets (measure on your workload)" + ); println!("\nšŸ“ Use Cases:"); println!(" • E-commerce: Filter by price, category, rating"); diff --git a/examples/namespace_demo.rs b/examples/namespace_demo.rs index ada1582..a55319f 100644 --- a/examples/namespace_demo.rs +++ b/examples/namespace_demo.rs @@ -201,7 +201,7 @@ fn main() -> anyhow::Result<()> { match result { Err(e) if e.to_string().contains("not active") => { println!(" āœ… Write rejected: {}", e); - } + }, _ => println!(" āš ļø Expected error but got: {:?}", result), } diff --git a/examples/openai_basic.rs b/examples/openai_basic.rs index d479c7e..b323035 100644 --- a/examples/openai_basic.rs +++ b/examples/openai_basic.rs @@ -34,7 +34,7 @@ async fn main() -> Result<()> { eprintln!("\nPlease set your OpenAI API key:"); eprintln!(" export OPENAI_API_KEY='your-api-key-here'\n"); std::process::exit(1); - } + }, }; println!("āœ“ API key found\n"); diff --git a/examples/openai_cost_tracking.rs b/examples/openai_cost_tracking.rs index 634e332..1d46ab7 100644 --- a/examples/openai_cost_tracking.rs +++ b/examples/openai_cost_tracking.rs @@ -35,7 +35,7 @@ async fn main() -> Result<()> { eprintln!("\nPlease set your OpenAI API key:"); eprintln!(" export OPENAI_API_KEY='your-api-key-here'\n"); std::process::exit(1); - } + }, }; println!("āœ“ API key found\n"); @@ -378,6 +378,8 @@ fn format_number(n: usize) -> String { #[cfg(not(all(feature = "embeddings", feature = "openai-embeddings")))] fn main() { eprintln!("This example requires the 'embeddings' and 'openai-embeddings' features."); - eprintln!("Run with: cargo run --example openai_cost_tracking --features \"embeddings,openai-embeddings\""); + eprintln!( + "Run with: cargo run --example openai_cost_tracking --features \"embeddings,openai-embeddings\"" + ); std::process::exit(1); } diff --git a/examples/openai_rag.rs b/examples/openai_rag.rs index 794f8ea..3bceb4d 100644 --- a/examples/openai_rag.rs +++ b/examples/openai_rag.rs @@ -38,7 +38,7 @@ async fn main() -> Result<()> { eprintln!("\nPlease set your OpenAI API key:"); eprintln!(" export OPENAI_API_KEY='your-api-key-here'\n"); std::process::exit(1); - } + }, }; println!("āœ“ API key found\n"); @@ -46,13 +46,34 @@ async fn main() -> Result<()> { // 2. Create sample documents to index println!("Creating sample documents..."); let documents = vec![ - ("rust_intro", "Rust is a systems programming language that runs blazingly fast, prevents segfaults, and guarantees thread safety. It accomplishes these goals without using a garbage collector."), - ("rust_ownership", "Rust's ownership system is its most unique feature. The ownership rules ensure memory safety without needing a garbage collector. Each value has a variable that's its owner, and there can only be one owner at a time."), - ("rust_cargo", "Cargo is Rust's build system and package manager. It handles building your code, downloading dependencies, and building those dependencies. Most Rustaceans use Cargo to manage their Rust projects."), - ("python_intro", "Python is a high-level, interpreted programming language known for its simplicity and readability. It supports multiple programming paradigms and has a comprehensive standard library."), - ("python_django", "Django is a high-level Python web framework that encourages rapid development and clean, pragmatic design. It follows the model-template-views architectural pattern and includes an ORM."), - ("javascript_intro", "JavaScript is a versatile programming language primarily used for web development. It runs in browsers and enables interactive web pages. Node.js allows JavaScript to run on servers too."), - ("vecstore_intro", "VecStore is a high-performance vector database written in Rust. It provides efficient similarity search using HNSW indexing, supports multiple embedding backends, and offers both Rust and Python APIs."), + ( + "rust_intro", + "Rust is a systems programming language that runs blazingly fast, prevents segfaults, and guarantees thread safety. It accomplishes these goals without using a garbage collector.", + ), + ( + "rust_ownership", + "Rust's ownership system is its most unique feature. The ownership rules ensure memory safety without needing a garbage collector. Each value has a variable that's its owner, and there can only be one owner at a time.", + ), + ( + "rust_cargo", + "Cargo is Rust's build system and package manager. It handles building your code, downloading dependencies, and building those dependencies. Most Rustaceans use Cargo to manage their Rust projects.", + ), + ( + "python_intro", + "Python is a high-level, interpreted programming language known for its simplicity and readability. It supports multiple programming paradigms and has a comprehensive standard library.", + ), + ( + "python_django", + "Django is a high-level Python web framework that encourages rapid development and clean, pragmatic design. It follows the model-template-views architectural pattern and includes an ORM.", + ), + ( + "javascript_intro", + "JavaScript is a versatile programming language primarily used for web development. It runs in browsers and enables interactive web pages. Node.js allows JavaScript to run on servers too.", + ), + ( + "vecstore_intro", + "VecStore is a high-performance vector database written in Rust. It provides efficient similarity search using HNSW indexing, supports multiple embedding backends, and offers both Rust and Python APIs.", + ), ]; println!("āœ“ Created {} documents", documents.len()); diff --git a/examples/performance_monitoring.rs b/examples/performance_monitoring.rs index 97143e3..77dea74 100644 --- a/examples/performance_monitoring.rs +++ b/examples/performance_monitoring.rs @@ -7,7 +7,7 @@ //! - Resource utilization use anyhow::Result; -use vecstore::{print_health_report, BenchmarkConfig, Benchmarker, HealthChecker, VecStore}; +use vecstore::{BenchmarkConfig, Benchmarker, HealthChecker, VecStore, print_health_report}; fn main() -> Result<()> { println!("šŸ“Š VecStore Performance Monitoring Example\n"); diff --git a/examples/phrase_matching_demo.rs b/examples/phrase_matching_demo.rs index 284263d..24f206b 100644 --- a/examples/phrase_matching_demo.rs +++ b/examples/phrase_matching_demo.rs @@ -8,7 +8,6 @@ //! //! Run with: cargo run --example phrase_matching_demo - // Mock TextIndex for demonstration (in practice, use VecStore's internal TextIndex) fn main() { println!("=== VecStore Phrase Matching Demo ===\n"); diff --git a/examples/production_ready_system.rs b/examples/production_ready_system.rs index 1e99208..6d7a2a4 100644 --- a/examples/production_ready_system.rs +++ b/examples/production_ready_system.rs @@ -110,13 +110,13 @@ impl VectorSearchService { if let Some(rec) = &alert.recommendation { eprintln!(" → {}", rec); } - } + }, vecstore::AlertSeverity::Warning => { println!("āš ļø WARNING: {}", alert.message); - } + }, vecstore::AlertSeverity::Info => { println!("ā„¹ļø INFO: {}", alert.message); - } + }, } } diff --git a/examples/rate_limit_demo.rs b/examples/rate_limit_demo.rs index c6007ec..0b3eea2 100644 --- a/examples/rate_limit_demo.rs +++ b/examples/rate_limit_demo.rs @@ -33,11 +33,7 @@ fn main() -> anyhow::Result<()> { println!("\nSending 7 requests:"); for i in 1..=7 { let result = token_limiter.try_acquire_result(); - let status = if result.allowed { - "ALLOWED" - } else { - "BLOCKED" - }; + let status = if result.allowed { "ALLOWED" } else { "BLOCKED" }; println!( " Request {}: {} (remaining: {})", i, status, result.remaining @@ -70,11 +66,7 @@ fn main() -> anyhow::Result<()> { println!("\nSending 10 requests:"); for i in 1..=10 { let result = burst_limiter.try_acquire_result(); - let status = if result.allowed { - "ALLOWED" - } else { - "BLOCKED" - }; + let status = if result.allowed { "ALLOWED" } else { "BLOCKED" }; println!( " Request {}: {} (remaining: {})", i, status, result.remaining @@ -102,11 +94,7 @@ fn main() -> anyhow::Result<()> { println!("\nSending 12 requests:"); for i in 1..=12 { let result = sliding_limiter.try_acquire_result(); - let status = if result.allowed { - "ALLOWED" - } else { - "BLOCKED" - }; + let status = if result.allowed { "ALLOWED" } else { "BLOCKED" }; println!( " Request {}: {} (remaining: {})", i, status, result.remaining @@ -120,11 +108,7 @@ fn main() -> anyhow::Result<()> { println!("\n After window slides:"); for i in 13..=15 { let result = sliding_limiter.try_acquire_result(); - let status = if result.allowed { - "ALLOWED" - } else { - "BLOCKED" - }; + let status = if result.allowed { "ALLOWED" } else { "BLOCKED" }; println!( " Request {}: {} (remaining: {})", i, status, result.remaining @@ -153,11 +137,7 @@ fn main() -> anyhow::Result<()> { println!("\nSending 8 requests:"); for i in 1..=8 { let result = leaky_limiter.try_acquire_result(); - let status = if result.allowed { - "ALLOWED" - } else { - "BLOCKED" - }; + let status = if result.allowed { "ALLOWED" } else { "BLOCKED" }; println!( " Request {}: {} (remaining: {})", i, status, result.remaining diff --git a/examples/recommender_demo.rs b/examples/recommender_demo.rs index bf84fe2..916ee37 100644 --- a/examples/recommender_demo.rs +++ b/examples/recommender_demo.rs @@ -8,10 +8,10 @@ use anyhow::Result; use std::collections::HashMap; +use vecstore::Metadata; use vecstore::recommender::{ CollaborativeRecommender, ContentBasedRecommender, HybridRecommender, UserPreference, }; -use vecstore::Metadata; fn create_metadata(title: &str, genre: &str, year: i32) -> Metadata { let mut fields = HashMap::new(); diff --git a/examples/reranking_demo.rs b/examples/reranking_demo.rs index eac09e9..fc06a7d 100644 --- a/examples/reranking_demo.rs +++ b/examples/reranking_demo.rs @@ -21,8 +21,8 @@ use std::collections::HashMap; use vecstore::{ - reranking::{CrossEncoderFn, IdentityReranker, MMRReranker, Reranker, ScoreReranker}, Metadata, Query, VecDatabase, + reranking::{CrossEncoderFn, IdentityReranker, MMRReranker, Reranker, ScoreReranker}, }; fn main() -> anyhow::Result<()> { diff --git a/examples/sparse_vectors_demo.rs b/examples/sparse_vectors_demo.rs index 37c18df..4562cd3 100644 --- a/examples/sparse_vectors_demo.rs +++ b/examples/sparse_vectors_demo.rs @@ -14,8 +14,8 @@ //! Run with: cargo run --example sparse_vectors_demo use vecstore::{ - bm25_score, hybrid_search_score, normalize_scores, BM25Config, BM25Stats, FusionStrategy, - HybridQueryV2, HybridSearchConfig, Vector, + BM25Config, BM25Stats, FusionStrategy, HybridQueryV2, HybridSearchConfig, Vector, bm25_score, + hybrid_search_score, normalize_scores, }; fn main() -> anyhow::Result<()> { diff --git a/examples/streaming_demo.rs b/examples/streaming_demo.rs index 927522f..0a812ee 100644 --- a/examples/streaming_demo.rs +++ b/examples/streaming_demo.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use std::time::Instant; -use vecstore::{make_record, Metadata, VecStore}; +use vecstore::{Metadata, VecStore, make_record}; fn main() -> anyhow::Result<()> { println!("╔═══════════════════════════════════════════════════════════════╗"); diff --git a/examples/text_chunking_demo.rs b/examples/text_chunking_demo.rs index 04dfcf7..6e8b116 100644 --- a/examples/text_chunking_demo.rs +++ b/examples/text_chunking_demo.rs @@ -14,8 +14,8 @@ use std::collections::HashMap; use vecstore::{ - text_splitter::{RecursiveCharacterTextSplitter, TextSplitter, TokenTextSplitter}, Metadata, Query, VecStore, + text_splitter::{RecursiveCharacterTextSplitter, TextSplitter, TokenTextSplitter}, }; const SAMPLE_DOCUMENT: &str = r#" diff --git a/examples/validation_demo.rs b/examples/validation_demo.rs index 056f4d7..6fd2681 100644 --- a/examples/validation_demo.rs +++ b/examples/validation_demo.rs @@ -36,13 +36,13 @@ fn main() -> anyhow::Result<()> { for warning in warnings { println!(" {:?}", warning); } - } + }, ValidationResult::Invalid(errors) => { println!("āŒ INVALID ({} errors)", errors.len()); for error in errors { println!(" {:?}", error); } - } + }, } } @@ -59,10 +59,10 @@ fn main() -> anyhow::Result<()> { ValidationResult::Valid => println!("āœ“ VALID"), ValidationResult::Warning(warnings) => { println!("āš ļø {} warnings", warnings.len()); - } + }, ValidationResult::Invalid(errors) => { println!("āŒ {} errors", errors.len()); - } + }, } } @@ -90,7 +90,7 @@ fn main() -> anyhow::Result<()> { "āŒ INVALID" } ); - } + }, Err(e) => println!(" Error: {}", e), } } diff --git a/examples/versioning_demo.rs b/examples/versioning_demo.rs index 78cad14..1f78c7e 100644 --- a/examples/versioning_demo.rs +++ b/examples/versioning_demo.rs @@ -3,8 +3,8 @@ use anyhow::Result; use std::collections::HashMap; use tempfile::TempDir; -use vecstore::versioning::VersionedStore; use vecstore::Metadata; +use vecstore::versioning::VersionedStore; fn create_metadata(desc: &str) -> Metadata { let mut fields = HashMap::new(); diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 50b3f5d..f19782d 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -1,2 +1,2 @@ [toolchain] -channel = "1.92" +channel = "1.92.0" diff --git a/rustfmt.toml b/rustfmt.toml index bb4a5ec..61a4558 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -14,34 +14,16 @@ tab_spaces = 4 # Newline style (native = platform default) newline_style = "Auto" -# Import grouping and sorting -imports_granularity = "Module" -group_imports = "StdExternalCrate" +# Import sorting reorder_imports = true reorder_modules = true -# Format strings -format_strings = false - -# Comments -wrap_comments = false -normalize_comments = false -normalize_doc_attributes = true - -# Control flow -control_brace_style = "AlwaysSameLine" - # Match arms match_arm_leading_pipes = "Never" match_block_trailing_comma = true # Function definitions -fn_args_layout = "Tall" -fn_single_line = false - -# Struct definitions -struct_lit_single_line = true -struct_field_align_threshold = 0 +fn_params_layout = "Tall" # Use field init shorthand use_field_init_shorthand = true @@ -49,12 +31,6 @@ use_field_init_shorthand = true # Use try shorthand (?) use_try_shorthand = true -# Trailing semicolons -trailing_semicolon = true - -# Where clause layout -where_single_line = false - # Chain formatting chain_width = 60 @@ -63,9 +39,3 @@ array_width = 60 # Attribute formatting attr_fn_like_width = 70 - -# Skip certain files/directories (generated code, etc.) -ignore = [ - "target/", - "proto/", -] diff --git a/src/ab_testing.rs b/src/ab_testing.rs index 03e8926..e27116f 100644 --- a/src/ab_testing.rs +++ b/src/ab_testing.rs @@ -2,8 +2,8 @@ // Compare index types, embedding models, reranking strategies, and more use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{SystemTime, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; @@ -247,8 +247,12 @@ impl ABTestManager { variant_data, }; - self.experiments.write()?.insert(experiment_id.clone(), state); - self.metrics.experiments_created.fetch_add(1, Ordering::Relaxed); + self.experiments + .write()? + .insert(experiment_id.clone(), state); + self.metrics + .experiments_created + .fetch_add(1, Ordering::Relaxed); Ok(experiment_id) } @@ -256,13 +260,15 @@ impl ABTestManager { /// Start an experiment pub fn start_experiment(&self, experiment_id: &str) -> Result<()> { let mut experiments = self.experiments.write()?; - let state = experiments.get_mut(experiment_id) - .ok_or_else(|| VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)))?; + let state = experiments.get_mut(experiment_id).ok_or_else(|| { + VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)) + })?; if state.status != ExperimentStatus::Draft && state.status != ExperimentStatus::Paused { - return Err(VecStoreError::InvalidInput( - format!("Cannot start experiment in status {:?}", state.status) - )); + return Err(VecStoreError::InvalidInput(format!( + "Cannot start experiment in status {:?}", + state.status + ))); } state.status = ExperimentStatus::Running; @@ -274,11 +280,14 @@ impl ABTestManager { /// Pause an experiment pub fn pause_experiment(&self, experiment_id: &str) -> Result<()> { let mut experiments = self.experiments.write()?; - let state = experiments.get_mut(experiment_id) - .ok_or_else(|| VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)))?; + let state = experiments.get_mut(experiment_id).ok_or_else(|| { + VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)) + })?; if state.status != ExperimentStatus::Running { - return Err(VecStoreError::InvalidInput("Experiment is not running".into())); + return Err(VecStoreError::InvalidInput( + "Experiment is not running".into(), + )); } state.status = ExperimentStatus::Paused; @@ -288,12 +297,15 @@ impl ABTestManager { /// Complete an experiment pub fn complete_experiment(&self, experiment_id: &str) -> Result { let mut experiments = self.experiments.write()?; - let state = experiments.get_mut(experiment_id) - .ok_or_else(|| VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)))?; + let state = experiments.get_mut(experiment_id).ok_or_else(|| { + VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)) + })?; state.status = ExperimentStatus::Completed; state.ended_at = Some(current_timestamp()); - self.metrics.experiments_completed.fetch_add(1, Ordering::Relaxed); + self.metrics + .experiments_completed + .fetch_add(1, Ordering::Relaxed); self.compute_results(state) } @@ -301,13 +313,15 @@ impl ABTestManager { /// Get variant assignment for a request pub fn get_assignment(&self, experiment_id: &str, identity: &str) -> Result { let experiments = self.experiments.read()?; - let state = experiments.get(experiment_id) - .ok_or_else(|| VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)))?; + let state = experiments.get(experiment_id).ok_or_else(|| { + VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)) + })?; if state.status != ExperimentStatus::Running { - return Err(VecStoreError::InvalidInput( - format!("Experiment is not running: {:?}", state.status) - )); + return Err(VecStoreError::InvalidInput(format!( + "Experiment is not running: {:?}", + state.status + ))); } // Check for existing assignment (Rust 1.92 if-let chain) @@ -315,7 +329,10 @@ impl ABTestManager { if let Some((exp_id, var_id)) = assignments.get(identity) && exp_id == experiment_id { - let variant = state.config.variants.iter() + let variant = state + .config + .variants + .iter() .find(|v| &v.id == var_id) .cloned() .ok_or_else(|| VecStoreError::NotFound("Variant not found".into()))?; @@ -344,7 +361,9 @@ impl ABTestManager { if let Some(vd) = state.variant_data.get(&variant.id) { vd.requests.fetch_add(1, Ordering::Relaxed); } - self.metrics.total_assignments.fetch_add(1, Ordering::Relaxed); + self.metrics + .total_assignments + .fetch_add(1, Ordering::Relaxed); Ok(VariantAssignment { experiment_id: experiment_id.to_string(), @@ -368,26 +387,28 @@ impl ABTestManager { } } Ok(config.variants[0].clone()) - } + }, TrafficAllocation::Consistent => { let hash = hash_identity(identity); let idx = (hash as usize) % config.variants.len(); Ok(config.variants[idx].clone()) - } + }, TrafficAllocation::TimeBased { rotation_minutes } => { let now = current_timestamp(); let period = (*rotation_minutes as u64) * 60 * 1000; let idx = ((now / period) as usize) % config.variants.len(); Ok(config.variants[idx].clone()) - } + }, TrafficAllocation::Manual => { // Return control by default for manual - config.variants.iter() + config + .variants + .iter() .find(|v| v.is_control) .or_else(|| config.variants.first()) .cloned() .ok_or_else(|| VecStoreError::InvalidInput("No variants defined".into())) - } + }, } } @@ -399,10 +420,13 @@ impl ABTestManager { result: RequestResult, ) -> Result<()> { let experiments = self.experiments.read()?; - let state = experiments.get(experiment_id) - .ok_or_else(|| VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)))?; + let state = experiments.get(experiment_id).ok_or_else(|| { + VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)) + })?; - let variant_data = state.variant_data.get(variant_id) + let variant_data = state + .variant_data + .get(variant_id) .ok_or_else(|| VecStoreError::NotFound(format!("Variant {} not found", variant_id)))?; // Update counters @@ -413,8 +437,13 @@ impl ABTestManager { } // Record latency - variant_data.total_latency_us.fetch_add(result.latency_us, Ordering::Relaxed); - variant_data.latency_samples.write()?.push(result.latency_us); + variant_data + .total_latency_us + .fetch_add(result.latency_us, Ordering::Relaxed); + variant_data + .latency_samples + .write()? + .push(result.latency_us); // Record metrics let mut metrics = variant_data.metric_values.write()?; @@ -422,7 +451,9 @@ impl ABTestManager { metrics.entry(name).or_insert_with(Vec::new).push(value); } - self.metrics.total_results_recorded.fetch_add(1, Ordering::Relaxed); + self.metrics + .total_results_recorded + .fetch_add(1, Ordering::Relaxed); Ok(()) } @@ -430,15 +461,19 @@ impl ABTestManager { /// Get current experiment results pub fn get_results(&self, experiment_id: &str) -> Result { let experiments = self.experiments.read()?; - let state = experiments.get(experiment_id) - .ok_or_else(|| VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)))?; + let state = experiments.get(experiment_id).ok_or_else(|| { + VecStoreError::NotFound(format!("Experiment {} not found", experiment_id)) + })?; self.compute_results(state) } fn compute_results(&self, state: &ExperimentState) -> Result { let mut variant_results = Vec::new(); - let control_id = state.config.variants.iter() + let control_id = state + .config + .variants + .iter() .find(|v| v.is_control) .map(|v| v.id.clone()); @@ -458,18 +493,20 @@ impl ABTestManager { for (name, values) in metrics.iter() { if !values.is_empty() { let mean: f64 = values.iter().sum::() / values.len() as f64; - let variance: f64 = values.iter() - .map(|v| (v - mean).powi(2)) - .sum::() / values.len() as f64; + let variance: f64 = values.iter().map(|v| (v - mean).powi(2)).sum::() + / values.len() as f64; let std_dev = variance.sqrt(); - metric_summaries.insert(name.clone(), MetricSummary { - count: values.len(), - mean, - std_dev, - min: values.iter().cloned().fold(f64::INFINITY, f64::min), - max: values.iter().cloned().fold(f64::NEG_INFINITY, f64::max), - }); + metric_summaries.insert( + name.clone(), + MetricSummary { + count: values.len(), + mean, + std_dev, + min: values.iter().cloned().fold(f64::INFINITY, f64::min), + max: values.iter().cloned().fold(f64::NEG_INFINITY, f64::max), + }, + ); } } @@ -478,9 +515,21 @@ impl ABTestManager { variant_name: variant.name.clone(), is_control: variant.is_control, sample_size: requests, - success_rate: if requests > 0 { successes as f64 / requests as f64 } else { 0.0 }, - error_rate: if requests > 0 { errors as f64 / requests as f64 } else { 0.0 }, - avg_latency_us: if requests > 0 { total_latency / requests } else { 0 }, + success_rate: if requests > 0 { + successes as f64 / requests as f64 + } else { + 0.0 + }, + error_rate: if requests > 0 { + errors as f64 / requests as f64 + } else { + 0.0 + }, + avg_latency_us: if requests > 0 { + total_latency / requests + } else { + 0 + }, p50_latency_us: p50, p95_latency_us: p95, p99_latency_us: p99, @@ -490,7 +539,8 @@ impl ABTestManager { } // Compute statistical comparisons - let comparisons = self.compute_comparisons(&variant_results, control_id.as_deref(), &state.config); + let comparisons = + self.compute_comparisons(&variant_results, control_id.as_deref(), &state.config); // Determine winner let winner = self.determine_winner(&variant_results, &comparisons, &state.config); @@ -543,13 +593,17 @@ impl ABTestManager { // Simple t-test approximation let (significant, p) = two_sample_ttest( - c.mean, c.std_dev, c.count, - t.mean, t.std_dev, t.count, + c.mean, + c.std_dev, + c.count, + t.mean, + t.std_dev, + t.count, config.significance_threshold, ); (change, significant, p) - } + }, _ => (0.0, false, 1.0), }; @@ -575,9 +629,7 @@ impl ABTestManager { config: &ExperimentConfig, ) -> Option { // Find best performing variant that's statistically significant - let significant: Vec<_> = comparisons.iter() - .filter(|c| c.is_significant) - .collect(); + let significant: Vec<_> = comparisons.iter().filter(|c| c.is_significant).collect(); if significant.is_empty() { return None; @@ -585,13 +637,13 @@ impl ABTestManager { // Return the variant with best improvement let best = if config.primary_metric.higher_is_better { - significant.iter().max_by(|a, b| - a.relative_change.total_cmp(&b.relative_change) - ) + significant + .iter() + .max_by(|a, b| a.relative_change.total_cmp(&b.relative_change)) } else { - significant.iter().min_by(|a, b| - a.relative_change.total_cmp(&b.relative_change) - ) + significant + .iter() + .min_by(|a, b| a.relative_change.total_cmp(&b.relative_change)) }; best.map(|c| c.variant_id.clone()) @@ -599,16 +651,21 @@ impl ABTestManager { /// List all experiments pub fn list_experiments(&self) -> Result> { - let experiments = self.experiments.read() + let experiments = self + .experiments + .read() .map_err(|_| VecStoreError::LockError("experiments lock poisoned".into()))?; - Ok(experiments.iter() + Ok(experiments + .iter() .map(|(id, state)| ExperimentSummary { id: id.clone(), name: state.config.name.clone(), status: state.status.clone(), experiment_type: state.config.experiment_type.clone(), variant_count: state.config.variants.len(), - total_requests: state.variant_data.values() + total_requests: state + .variant_data + .values() .map(|vd| vd.requests.load(Ordering::Relaxed)) .sum(), created_at: state.created_at, @@ -619,14 +676,19 @@ impl ABTestManager { /// Get manager stats pub fn get_stats(&self) -> Result { - let experiments = self.experiments.read() + let experiments = self + .experiments + .read() .map_err(|_| VecStoreError::LockError("experiments lock poisoned".into()))?; Ok(ABTestStats { experiments_created: self.metrics.experiments_created.load(Ordering::Relaxed), experiments_completed: self.metrics.experiments_completed.load(Ordering::Relaxed), total_assignments: self.metrics.total_assignments.load(Ordering::Relaxed), total_results_recorded: self.metrics.total_results_recorded.load(Ordering::Relaxed), - active_experiments: experiments.values().filter(|s| s.status == ExperimentStatus::Running).count(), + active_experiments: experiments + .values() + .filter(|s| s.status == ExperimentStatus::Running) + .count(), }) } } @@ -765,15 +827,28 @@ fn compute_percentiles(samples: &[u64]) -> (u64, u64, u64) { let p99_idx = (samples.len() as f64 * 0.99) as usize; ( - sorted.get(p50_idx.min(sorted.len() - 1)).copied().unwrap_or(0), - sorted.get(p95_idx.min(sorted.len() - 1)).copied().unwrap_or(0), - sorted.get(p99_idx.min(sorted.len() - 1)).copied().unwrap_or(0), + sorted + .get(p50_idx.min(sorted.len() - 1)) + .copied() + .unwrap_or(0), + sorted + .get(p95_idx.min(sorted.len() - 1)) + .copied() + .unwrap_or(0), + sorted + .get(p99_idx.min(sorted.len() - 1)) + .copied() + .unwrap_or(0), ) } fn two_sample_ttest( - mean1: f64, std1: f64, n1: usize, - mean2: f64, std2: f64, n2: usize, + mean1: f64, + std1: f64, + n1: usize, + mean2: f64, + std2: f64, + n2: usize, threshold: f64, ) -> (bool, f64) { if n1 < 2 || n2 < 2 { @@ -921,11 +996,17 @@ mod tests { let mut metrics = HashMap::new(); metrics.insert("recall_at_10".to_string(), 0.95); - manager.record_result(&id, &assignment.variant_id, RequestResult { - success: true, - latency_us: 5000, - metrics, - }).unwrap(); + manager + .record_result( + &id, + &assignment.variant_id, + RequestResult { + success: true, + latency_us: 5000, + metrics, + }, + ) + .unwrap(); let results = manager.get_results(&id).unwrap(); assert!(!results.variant_results.is_empty()); diff --git a/src/adaptive_cache.rs b/src/adaptive_cache.rs index e229f0b..1adce9b 100644 --- a/src/adaptive_cache.rs +++ b/src/adaptive_cache.rs @@ -30,7 +30,10 @@ use std::collections::{HashMap, VecDeque}; use std::hash::{Hash, Hasher}; -use std::sync::{RwLock, atomic::{AtomicU64, Ordering}}; +use std::sync::{ + RwLock, + atomic::{AtomicU64, Ordering}, +}; use std::time::{Duration, Instant}; use serde::{Deserialize, Serialize}; @@ -326,14 +329,16 @@ impl PrefetchPredictor { } fn should_prefetch(&self, query_hash: &str) -> bool { - self.patterns.iter() + self.patterns + .iter() .find(|p| p.query_hash == query_hash) .map(|p| p.count >= self.min_count) .unwrap_or(false) } fn get_frequent_patterns(&self) -> Vec { - self.patterns.iter() + self.patterns + .iter() .filter(|p| p.count >= self.min_count) .map(|p| p.query_hash.clone()) .collect() @@ -390,7 +395,9 @@ impl AdaptiveCache { // Quick bloom filter check if !self.bloom.might_contain(key) { - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; stats.misses += 1; return Ok(None); @@ -398,12 +405,16 @@ impl AdaptiveCache { // Check L1 (hot) { - let mut l1 = self.l1.write() + let mut l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; if let Some(entry) = l1.get_mut(key) { if !entry.is_expired(Duration::from_secs(self.config.l1_ttl_seconds)) { entry.touch(); - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; stats.hits += 1; return Ok(Some(entry.value.clone())); @@ -415,45 +426,61 @@ impl AdaptiveCache { // Check L2 (warm) and promote to L1 { - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; if let Some(entry) = l2.remove(key) - && !entry.is_expired(Duration::from_secs(self.config.l2_ttl_seconds)) { - let mut stats = self.l2_stats.write() - .map_err(|_| VecStoreError::LockError("l2_stats lock poisoned".into()))?; - stats.hits += 1; - drop(stats); - let value = entry.value.clone(); - self.promote_to_l1(key, entry)?; - return Ok(Some(value)); - } + && !entry.is_expired(Duration::from_secs(self.config.l2_ttl_seconds)) + { + let mut stats = self + .l2_stats + .write() + .map_err(|_| VecStoreError::LockError("l2_stats lock poisoned".into()))?; + stats.hits += 1; + drop(stats); + let value = entry.value.clone(); + self.promote_to_l1(key, entry)?; + return Ok(Some(value)); + } } // Check L3 (disk) and promote to L2 if self.config.enable_l3 { - let mut l3 = self.l3.write() + let mut l3 = self + .l3 + .write() .map_err(|_| VecStoreError::LockError("l3 cache lock poisoned".into()))?; if let Some(entry) = l3.remove(key) - && !entry.is_expired(Duration::from_secs(self.config.l3_ttl_seconds)) { - let mut stats = self.l3_stats.write() - .map_err(|_| VecStoreError::LockError("l3_stats lock poisoned".into()))?; - stats.hits += 1; - drop(stats); - let value = entry.value.clone(); - self.promote_to_l2(key, entry)?; - return Ok(Some(value)); - } + && !entry.is_expired(Duration::from_secs(self.config.l3_ttl_seconds)) + { + let mut stats = self + .l3_stats + .write() + .map_err(|_| VecStoreError::LockError("l3_stats lock poisoned".into()))?; + stats.hits += 1; + drop(stats); + let value = entry.value.clone(); + self.promote_to_l2(key, entry)?; + return Ok(Some(value)); + } } // Miss - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; stats.misses += 1; Ok(None) } /// Get with semantic matching - pub fn get_semantic(&self, key: &str, query_vector: &[f32]) -> Result, VecStoreError> { + pub fn get_semantic( + &self, + key: &str, + query_vector: &[f32], + ) -> Result, VecStoreError> { // Try exact match first if let Some(result) = self.get(key)? { return Ok(Some(result)); @@ -464,7 +491,9 @@ impl AdaptiveCache { } // Try semantic matching in L1 - let l1 = self.l1.read() + let l1 = self + .l1 + .read() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; for (_, entry) in l1.iter() { if let Some(ref cached_vector) = entry.query_vector { @@ -472,7 +501,9 @@ impl AdaptiveCache { if similarity >= self.config.similarity_threshold { let value = entry.value.clone(); drop(l1); - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; stats.hits += 1; return Ok(Some(value)); @@ -489,7 +520,12 @@ impl AdaptiveCache { } /// Put with query vector (for semantic matching) - pub fn put_with_vector(&self, key: &str, value: CachedResult, query_vector: Option>) -> Result<(), VecStoreError> { + pub fn put_with_vector( + &self, + key: &str, + value: CachedResult, + query_vector: Option>, + ) -> Result<(), VecStoreError> { let size_bytes = estimate_size(&value); let entry = CacheEntry { @@ -506,7 +542,9 @@ impl AdaptiveCache { // Record for prefetching if self.config.enable_prefetch { - let mut prefetch = self.prefetch.write() + let mut prefetch = self + .prefetch + .write() .map_err(|_| VecStoreError::LockError("prefetch lock poisoned".into()))?; prefetch.record(key, query_vector); } @@ -514,11 +552,15 @@ impl AdaptiveCache { // Evict if necessary and insert into L1 self.evict_l1_if_needed()?; - let mut l1 = self.l1.write() + let mut l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; l1.insert(key.to_string(), entry); - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; stats.entries = l1.len(); stats.size_bytes += size_bytes; @@ -542,7 +584,9 @@ impl AdaptiveCache { self.bloom.add(key); - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; l2.insert(key.to_string(), entry); @@ -551,17 +595,23 @@ impl AdaptiveCache { /// Invalidate cache entry pub fn invalidate(&self, key: &str) -> Result<(), VecStoreError> { - let mut l1 = self.l1.write() + let mut l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; l1.remove(key); drop(l1); - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; l2.remove(key); drop(l2); - let mut l3 = self.l3.write() + let mut l3 = self + .l3 + .write() .map_err(|_| VecStoreError::LockError("l3 cache lock poisoned".into()))?; l3.remove(key); @@ -570,34 +620,46 @@ impl AdaptiveCache { /// Clear all caches pub fn clear(&self) -> Result<(), VecStoreError> { - let mut l1 = self.l1.write() + let mut l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; l1.clear(); drop(l1); - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; l2.clear(); drop(l2); - let mut l3 = self.l3.write() + let mut l3 = self + .l3 + .write() .map_err(|_| VecStoreError::LockError("l3 cache lock poisoned".into()))?; l3.clear(); drop(l3); self.bloom.clear(); - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; *stats = TierStats::default(); drop(stats); - let mut stats = self.l2_stats.write() + let mut stats = self + .l2_stats + .write() .map_err(|_| VecStoreError::LockError("l2_stats lock poisoned".into()))?; *stats = TierStats::default(); drop(stats); - let mut stats = self.l3_stats.write() + let mut stats = self + .l3_stats + .write() .map_err(|_| VecStoreError::LockError("l3_stats lock poisoned".into()))?; *stats = TierStats::default(); @@ -606,11 +668,17 @@ impl AdaptiveCache { /// Get cache statistics pub fn stats(&self) -> Result { - let l1_stats = self.l1_stats.read() + let l1_stats = self + .l1_stats + .read() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; - let l2_stats = self.l2_stats.read() + let l2_stats = self + .l2_stats + .read() .map_err(|_| VecStoreError::LockError("l2_stats lock poisoned".into()))?; - let l3_stats = self.l3_stats.read() + let l3_stats = self + .l3_stats + .read() .map_err(|_| VecStoreError::LockError("l3_stats lock poisoned".into()))?; Ok(CacheStats { @@ -623,38 +691,55 @@ impl AdaptiveCache { /// Get frequent query patterns for prefetching pub fn get_prefetch_candidates(&self) -> Result, VecStoreError> { - let prefetch = self.prefetch.read() + let prefetch = self + .prefetch + .read() .map_err(|_| VecStoreError::LockError("prefetch lock poisoned".into()))?; Ok(prefetch.get_frequent_patterns()) } - fn promote_to_l1(&self, key: &str, mut entry: CacheEntry) -> Result<(), VecStoreError> { + fn promote_to_l1( + &self, + key: &str, + mut entry: CacheEntry, + ) -> Result<(), VecStoreError> { entry.touch(); self.evict_l1_if_needed()?; - let mut l1 = self.l1.write() + let mut l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; l1.insert(key.to_string(), entry); Ok(()) } - fn promote_to_l2(&self, key: &str, mut entry: CacheEntry) -> Result<(), VecStoreError> { + fn promote_to_l2( + &self, + key: &str, + mut entry: CacheEntry, + ) -> Result<(), VecStoreError> { entry.touch(); self.evict_l2_if_needed()?; - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; l2.insert(key.to_string(), entry); Ok(()) } fn evict_l1_if_needed(&self) -> Result<(), VecStoreError> { - let mut l1 = self.l1.write() + let mut l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; while l1.len() >= self.config.l1_size { // Find LRU entry - let lru_key = l1.iter() + let lru_key = l1 + .iter() .min_by_key(|(_, e)| e.last_access) .map(|(k, _)| k.clone()); @@ -663,12 +748,16 @@ impl AdaptiveCache { // Demote to L2 - drop l1 lock first to avoid deadlock drop(l1); self.demote_to_l2(&key, entry)?; - let mut stats = self.l1_stats.write() + let mut stats = self + .l1_stats + .write() .map_err(|_| VecStoreError::LockError("l1_stats lock poisoned".into()))?; stats.evictions += 1; drop(stats); // Re-acquire l1 lock - l1 = self.l1.write() + l1 = self + .l1 + .write() .map_err(|_| VecStoreError::LockError("l1 cache lock poisoned".into()))?; } } else { @@ -679,11 +768,14 @@ impl AdaptiveCache { } fn evict_l2_if_needed(&self) -> Result<(), VecStoreError> { - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; while l2.len() >= self.config.l2_size { - let lru_key = l2.iter() + let lru_key = l2 + .iter() .min_by_key(|(_, e)| e.last_access) .map(|(k, _)| k.clone()); @@ -694,12 +786,16 @@ impl AdaptiveCache { if self.config.enable_l3 { self.demote_to_l3(&key, entry)?; } - let mut stats = self.l2_stats.write() + let mut stats = self + .l2_stats + .write() .map_err(|_| VecStoreError::LockError("l2_stats lock poisoned".into()))?; stats.evictions += 1; drop(stats); // Re-acquire l2 lock - l2 = self.l2.write() + l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; } } else { @@ -709,18 +805,30 @@ impl AdaptiveCache { Ok(()) } - fn demote_to_l2(&self, key: &str, entry: CacheEntry) -> Result<(), VecStoreError> { + fn demote_to_l2( + &self, + key: &str, + entry: CacheEntry, + ) -> Result<(), VecStoreError> { self.evict_l2_if_needed()?; - let mut l2 = self.l2.write() + let mut l2 = self + .l2 + .write() .map_err(|_| VecStoreError::LockError("l2 cache lock poisoned".into()))?; l2.insert(key.to_string(), entry); Ok(()) } - fn demote_to_l3(&self, key: &str, entry: CacheEntry) -> Result<(), VecStoreError> { + fn demote_to_l3( + &self, + key: &str, + entry: CacheEntry, + ) -> Result<(), VecStoreError> { // In production, this would write to disk - let mut l3 = self.l3.write() + let mut l3 = self + .l3 + .write() .map_err(|_| VecStoreError::LockError("l3 cache lock poisoned".into()))?; l3.insert(key.to_string(), entry); Ok(()) @@ -760,12 +868,7 @@ pub struct CacheKeyGenerator; impl CacheKeyGenerator { /// Generate cache key from query parameters - pub fn generate( - collection: &str, - vector: &[f32], - k: usize, - filter: Option<&str>, - ) -> String { + pub fn generate(collection: &str, vector: &[f32], k: usize, filter: Option<&str>) -> String { use std::collections::hash_map::DefaultHasher; let mut hasher = DefaultHasher::new(); @@ -806,9 +909,7 @@ impl CacheKeyGenerator { fn estimate_size(result: &CachedResult) -> usize { let ids_size: usize = result.ids.iter().map(|s| s.len()).sum(); let scores_size = result.scores.len() * 4; - let meta_size = result.metadata.as_ref() - .map(|m| m.len() * 64) - .unwrap_or(0); + let meta_size = result.metadata.as_ref().map(|m| m.len() * 64).unwrap_or(0); ids_size + scores_size + meta_size + 64 // overhead } @@ -873,11 +974,16 @@ mod tests { let cache = AdaptiveCache::new(config); for i in 0..5 { - cache.put(&format!("key{}", i), CachedResult { - ids: vec![], - scores: vec![], - metadata: None, - }).unwrap(); + cache + .put( + &format!("key{}", i), + CachedResult { + ids: vec![], + scores: vec![], + metadata: None, + }, + ) + .unwrap(); } // L1 should have at most 3 entries @@ -896,11 +1002,17 @@ mod tests { let vector1 = vec![1.0, 0.0, 0.0]; let vector2 = vec![0.999, 0.001, 0.0]; // Very similar - cache.put_with_vector("key1", CachedResult { - ids: vec!["id1".to_string()], - scores: vec![1.0], - metadata: None, - }, Some(vector1.clone())).unwrap(); + cache + .put_with_vector( + "key1", + CachedResult { + ids: vec!["id1".to_string()], + scores: vec![1.0], + metadata: None, + }, + Some(vector1.clone()), + ) + .unwrap(); // Should find via semantic matching let result = cache.get_semantic("key2", &vector2).unwrap(); @@ -923,11 +1035,16 @@ mod tests { fn test_cache_stats() { let cache = AdaptiveCache::new(CacheConfig::default()); - cache.put("key1", CachedResult { - ids: vec![], - scores: vec![], - metadata: None, - }).unwrap(); + cache + .put( + "key1", + CachedResult { + ids: vec![], + scores: vec![], + metadata: None, + }, + ) + .unwrap(); cache.get("key1").unwrap(); cache.get("nonexistent").unwrap(); diff --git a/src/advanced_filter.rs b/src/advanced_filter.rs index 03933f8..9717860 100644 --- a/src/advanced_filter.rs +++ b/src/advanced_filter.rs @@ -28,7 +28,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use regex::Regex; use serde_json::Value; @@ -64,7 +64,7 @@ impl AdvancedFilter { AdvancedFilter::JsonPath(path, op, value) => { let extracted = extract_json_path(metadata, path)?; compare_values(&extracted, op, value) - } + }, AdvancedFilter::Regex(field, pattern) => { let field_value = extract_field(metadata, field)?; if let Some(s) = field_value.as_str() { @@ -73,7 +73,7 @@ impl AdvancedFilter { } else { Ok(false) } - } + }, AdvancedFilter::ArrayContains(field, value) => { let field_value = extract_field(metadata, field)?; if let Some(arr) = field_value.as_array() { @@ -81,7 +81,7 @@ impl AdvancedFilter { } else { Ok(false) } - } + }, AdvancedFilter::ArrayLength(field, op, length) => { let field_value = extract_field(metadata, field)?; if let Some(arr) = field_value.as_array() { @@ -90,7 +90,7 @@ impl AdvancedFilter { } else { Ok(false) } - } + }, AdvancedFilter::ArrayAll(field, condition) => { let field_value = extract_field(metadata, field)?; if let Some(arr) = field_value.as_array() { @@ -103,7 +103,7 @@ impl AdvancedFilter { } else { Ok(false) } - } + }, AdvancedFilter::ArrayAny(field, condition) => { let field_value = extract_field(metadata, field)?; if let Some(arr) = field_value.as_array() { @@ -116,7 +116,7 @@ impl AdvancedFilter { } else { Ok(false) } - } + }, AdvancedFilter::And(filters) => { for filter in filters { if !filter.matches(metadata)? { @@ -124,7 +124,7 @@ impl AdvancedFilter { } } Ok(true) - } + }, AdvancedFilter::Or(filters) => { for filter in filters { if filter.matches(metadata)? { @@ -132,12 +132,12 @@ impl AdvancedFilter { } } Ok(false) - } + }, AdvancedFilter::Not(filter) => Ok(!filter.matches(metadata)?), AdvancedFilter::Basic(field, op, value) => { let field_value = extract_field(metadata, field)?; compare_values(&field_value, op, value) - } + }, } } } diff --git a/src/advanced_quant.rs b/src/advanced_quant.rs index 4fa92bf..a95fa94 100644 --- a/src/advanced_quant.rs +++ b/src/advanced_quant.rs @@ -312,7 +312,8 @@ impl BinaryQuantizer { /// Compute Hamming distance (fast binary comparison) pub fn hamming_distance(&self, a: &QuantizedVector, b: &QuantizedVector) -> u32 { - a.data.iter() + a.data + .iter() .zip(&b.data) .map(|(&x, &y)| (x ^ y).count_ones()) .sum() @@ -526,7 +527,8 @@ impl FourBitQuantizer { } } - self.params.scales = mins.iter() + self.params.scales = mins + .iter() .zip(&maxs) .map(|(&min, &max)| (max - min) / 15.0) .collect(); @@ -743,13 +745,14 @@ impl AdaptiveQuantizer { } // Allocate bits based on importance - let _total_bits = dim * match self.config.stored_bits { - BitWidth::Bit1 => 1, - BitWidth::Bit1_5 => 2, - BitWidth::Bit2 => 2, - BitWidth::Bit4 => 4, - BitWidth::Bit8 => 8, - }; + let _total_bits = dim + * match self.config.stored_bits { + BitWidth::Bit1 => 1, + BitWidth::Bit1_5 => 2, + BitWidth::Bit2 => 2, + BitWidth::Bit4 => 4, + BitWidth::Bit8 => 8, + }; // For now, use uniform allocation (could be optimized) self.bit_allocation = vec![self.config.stored_bits; dim]; @@ -812,13 +815,17 @@ impl QuantizedIndex { _ => self.four_bit_quantizer.encode(vector), }; - let mut vectors = self.vectors.write() - .map_err(|_| crate::error::VecStoreError::LockError("Failed to acquire write lock on vectors".into()))?; + let mut vectors = self.vectors.write().map_err(|_| { + crate::error::VecStoreError::LockError("Failed to acquire write lock on vectors".into()) + })?; vectors.insert(id.to_string(), quantized); if self.config.rescore { - let mut originals = self.original_vectors.write() - .map_err(|_| crate::error::VecStoreError::LockError("Failed to acquire write lock on original_vectors".into()))?; + let mut originals = self.original_vectors.write().map_err(|_| { + crate::error::VecStoreError::LockError( + "Failed to acquire write lock on original_vectors".into(), + ) + })?; originals.insert(id.to_string(), vector.to_vec()); } @@ -827,7 +834,9 @@ impl QuantizedIndex { /// Search with quantized comparison pub fn search(&self, query: &[f32], top_k: usize) -> Vec<(String, f32)> { - let Ok(vectors) = self.vectors.read() else { return Vec::new(); }; + let Ok(vectors) = self.vectors.read() else { + return Vec::new(); + }; // First pass: approximate search with quantized vectors let candidates_k = if self.config.rescore { @@ -843,7 +852,7 @@ impl QuantizedIndex { BitWidth::Bit1 => { let query_q = self.binary_quantizer.encode(query); self.binary_quantizer.binary_similarity(&query_q, qvec) - } + }, _ => { let decoded = match qvec.bit_width { BitWidth::Bit2 => self.two_bit_quantizer.decode(qvec), @@ -851,7 +860,7 @@ impl QuantizedIndex { _ => self.four_bit_quantizer.decode(qvec), }; cosine_similarity(&decoded, query) - } + }, }; (id.clone(), score) }) @@ -862,7 +871,9 @@ impl QuantizedIndex { // Second pass: rescore with original vectors if self.config.rescore { - let Ok(originals) = self.original_vectors.read() else { return candidates; }; + let Ok(originals) = self.original_vectors.read() else { + return candidates; + }; candidates = candidates .into_iter() @@ -953,7 +964,9 @@ mod tests { fn test_binary_quantization() { let quantizer = BinaryQuantizer::new(128); - let vector: Vec = (0..128).map(|i| if i % 2 == 0 { 0.5 } else { -0.5 }).collect(); + let vector: Vec = (0..128) + .map(|i| if i % 2 == 0 { 0.5 } else { -0.5 }) + .collect(); let encoded = quantizer.encode(&vector); assert_eq!(encoded.dimension, 128); @@ -967,13 +980,13 @@ mod tests { fn test_ternary_quantization() { let quantizer = TernaryQuantizer::new(128, 0.1); - let vector: Vec = (0..128).map(|i| { - match i % 3 { + let vector: Vec = (0..128) + .map(|i| match i % 3 { 0 => 0.5, 1 => -0.5, _ => 0.0, - } - }).collect(); + }) + .collect(); let encoded = quantizer.encode(&vector); let decoded = quantizer.decode(&encoded); diff --git a/src/agent.rs b/src/agent.rs index bb3d315..357430c 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -35,7 +35,7 @@ //! println!("Agent reasoning: {}", result.explanation); //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -168,10 +168,7 @@ pub enum AgentTool { }, /// Check privacy constraints - PrivacyCheck { - operation: String, - epsilon: f64, - }, + PrivacyCheck { operation: String, epsilon: f64 }, /// Track lineage TrackLineage { @@ -192,9 +189,15 @@ pub enum FilterExpression { /// Less than Lt { field: String, value: FilterValue }, /// In list - In { field: String, values: Vec }, + In { + field: String, + values: Vec, + }, /// Not in list - NotIn { field: String, values: Vec }, + NotIn { + field: String, + values: Vec, + }, /// Contains (for text/arrays) Contains { field: String, value: String }, /// And combination @@ -413,11 +416,17 @@ pub struct IntentConstraint { /// Intent modifiers #[derive(Debug, Clone, Serialize, Deserialize)] pub enum IntentModifier { - TimeRange { start: Option, end: Option }, + TimeRange { + start: Option, + end: Option, + }, Limit(usize), Exclude(Vec), Include(Vec), - SortBy { field: String, ascending: bool }, + SortBy { + field: String, + ascending: bool, + }, } // ============================================================================ @@ -597,9 +606,10 @@ impl QueryAgent { // Parse limits if let Some(pos) = query_lower.find("top ") && let Some(num_str) = query_lower[pos + 4..].split_whitespace().next() - && let Ok(n) = num_str.parse::() { - modifiers.push(IntentModifier::Limit(n)); - } + && let Ok(n) = num_str.parse::() + { + modifiers.push(IntentModifier::Limit(n)); + } Ok(QueryIntent { action, @@ -701,7 +711,7 @@ impl QueryAgent { confidence: 0.85, fallback: None, }); - } + }, IntentAction::Explain => { steps.push(PlanStep { @@ -716,7 +726,7 @@ impl QueryAgent { confidence: 0.9, fallback: None, }); - } + }, IntentAction::Aggregate => { steps.push(PlanStep { @@ -732,7 +742,7 @@ impl QueryAgent { confidence: 0.95, fallback: None, }); - } + }, _ => { // Default to search @@ -750,7 +760,7 @@ impl QueryAgent { confidence: 0.8, fallback: None, }); - } + }, } let overall_confidence = steps.iter().map(|s| s.confidence).product::(); @@ -798,7 +808,11 @@ impl QueryAgent { trace.push(ExecutionStep { step: step.step, - tool: format!("{:?}", step.tool).split('{').next().unwrap().to_string(), + tool: format!("{:?}", step.tool) + .split('{') + .next() + .unwrap() + .to_string(), input_summary: format!("{} input items", current_ids.len()), output_summary: format!("{} output items", result.result_ids.len()), success: result.success, @@ -856,36 +870,46 @@ impl QueryAgent { let scores: Vec = (0..*k).map(|i| 1.0 - (i as f32 * 0.05)).collect(); let _ = (query, filters); // Mark as used (ids, scores, None) - } + }, AgentTool::KeywordSearch { query, k, filters } => { // Keyword search using BM25-style scoring let ids: Vec = (0..*k).map(|i| format!("kw_result_{}", i)).collect(); let scores: Vec = (0..*k).map(|i| 0.9 - (i as f32 * 0.08)).collect(); let _ = (query, filters); // Mark as used (ids, scores, None) - } + }, AgentTool::HybridSearch { k, alpha, .. } => { // Hybrid search combining vector and keyword let ids: Vec = (0..*k).map(|i| format!("hybrid_result_{}", i)).collect(); - let scores: Vec = (0..*k).map(|i| { - // Blend scores based on alpha - let vec_score = 1.0 - (i as f32 * 0.05); - let kw_score = 0.9 - (i as f32 * 0.08); - alpha * vec_score + (1.0 - alpha) * kw_score - }).collect(); + let scores: Vec = (0..*k) + .map(|i| { + // Blend scores based on alpha + let vec_score = 1.0 - (i as f32 * 0.05); + let kw_score = 0.9 - (i as f32 * 0.08); + alpha * vec_score + (1.0 - alpha) * kw_score + }) + .collect(); (ids, scores, None) - } - AgentTool::Filter { expression, input_ids: filter_input } => { + }, + AgentTool::Filter { + expression, + input_ids: filter_input, + } => { // Filter operation - returns subset of input IDs matching expression - let ids: Vec = filter_input.iter() + let ids: Vec = filter_input + .iter() .take(filter_input.len() / 2 + 1) .cloned() .collect(); let scores: Vec = vec![1.0; ids.len()]; let _ = expression; // Mark as used (ids, scores, None) - } - AgentTool::Aggregate { operation, input_ids: agg_input, .. } => { + }, + AgentTool::Aggregate { + operation, + input_ids: agg_input, + .. + } => { // Aggregate operation let agg_result = match operation { AggregateOp::Count => Some(AggregationResult { @@ -927,11 +951,14 @@ impl QueryAgent { count: agg_input.len(), groups: Some(groups), }) - } + }, }; (agg_input.clone(), vec![1.0; agg_input.len()], agg_result) - } - AgentTool::Transform { operation, input_ids: transform_input } => { + }, + AgentTool::Transform { + operation, + input_ids: transform_input, + } => { // Transform operation let ids = match operation { TransformOp::Normalize => transform_input.clone(), @@ -939,50 +966,53 @@ impl QueryAgent { TransformOp::Cluster { k } => { // Return k cluster representatives transform_input.iter().take(*k).cloned().collect() - } + }, TransformOp::Deduplicate { .. } => { // Remove duplicates by taking first half as "unique" let unique_count = (transform_input.len() / 2).max(1); transform_input.iter().take(unique_count).cloned().collect() - } + }, TransformOp::Enrich { .. } => transform_input.clone(), }; let scores = vec![1.0; ids.len()]; (ids, scores, None) - } + }, AgentTool::Fetch { ids } => { // Fetch by IDs - return the requested IDs let scores = vec![1.0; ids.len()]; (ids.clone(), scores, None) - } - AgentTool::Rerank { k, input_ids: rerank_input, .. } => { + }, + AgentTool::Rerank { + k, + input_ids: rerank_input, + .. + } => { // Rerank - reorder based on relevance - let mut reranked: Vec<(usize, f32)> = rerank_input.iter() + let mut reranked: Vec<(usize, f32)> = rerank_input + .iter() .enumerate() .map(|(i, _)| (i, 1.0 / (1.0 + i as f32))) .collect(); reranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - let ids: Vec = reranked.iter() + let ids: Vec = reranked + .iter() .take(*k) .map(|(i, _)| rerank_input[*i].clone()) .collect(); - let scores: Vec = reranked.iter() - .take(*k) - .map(|(_, s)| *s) - .collect(); + let scores: Vec = reranked.iter().take(*k).map(|(_, s)| *s).collect(); (ids, scores, None) - } + }, AgentTool::Explain { result_id, .. } => { // Explain - return explanation for a result (vec![result_id.clone()], vec![1.0], None) - } + }, AgentTool::TemporalSearch { k, .. } => { // Temporal search - filter by time range let ids: Vec = (0..*k).map(|i| format!("temporal_result_{}", i)).collect(); let scores: Vec = (0..*k).map(|i| 1.0 - (i as f32 * 0.05)).collect(); (ids, scores, None) - } + }, AgentTool::PrivacyCheck { epsilon, .. } => { // Privacy check - verify privacy constraints // Return whether the privacy budget allows the operation @@ -993,11 +1023,11 @@ impl QueryAgent { vec!["privacy_denied".to_string()] }; (ids, vec![if approved { 1.0 } else { 0.0 }], None) - } + }, AgentTool::TrackLineage { vector_id, .. } => { // Track lineage - return the vector ID with lineage info (vec![vector_id.clone()], vec![1.0], None) - } + }, }; // Use input_ids if provided and result is empty @@ -1026,7 +1056,11 @@ impl QueryAgent { } /// Generate explanation - fn generate_explanation(&self, plan: &ExecutionPlan, trace: &[ExecutionStep]) -> Result { + fn generate_explanation( + &self, + plan: &ExecutionPlan, + trace: &[ExecutionStep], + ) -> Result { let mut explanation = String::new(); explanation.push_str(&format!("Query: \"{}\"\n\n", plan.query)); @@ -1052,7 +1086,11 @@ impl QueryAgent { } /// Generate suggestions for query refinement - fn generate_suggestions(&self, intent: &QueryIntent, results: &[String]) -> Result> { + fn generate_suggestions( + &self, + intent: &QueryIntent, + results: &[String], + ) -> Result> { let mut suggestions = Vec::new(); if results.is_empty() { @@ -1063,7 +1101,10 @@ impl QueryAgent { } // Check for time-related suggestions - let has_time_modifier = intent.modifiers.iter().any(|m| matches!(m, IntentModifier::TimeRange { .. })); + let has_time_modifier = intent + .modifiers + .iter() + .any(|m| matches!(m, IntentModifier::TimeRange { .. })); if !has_time_modifier { suggestions.push("Add time range to focus on recent items".to_string()); } @@ -1163,7 +1204,11 @@ impl TransformAgent { } /// Transform vectors based on instruction - pub async fn transform(&self, instruction: &str, vectors: &[Vec]) -> Result { + pub async fn transform( + &self, + instruction: &str, + vectors: &[Vec], + ) -> Result { let start_time = std::time::Instant::now(); // Parse transformation instruction @@ -1200,21 +1245,23 @@ impl TransformAgent { } } - fn apply_transform(&self, operation: &TransformOp, vectors: &[Vec]) -> Result>> { + fn apply_transform( + &self, + operation: &TransformOp, + vectors: &[Vec], + ) -> Result>> { match operation { - TransformOp::Normalize => { - Ok(vectors - .iter() - .map(|v| { - let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); - if norm > 0.0 { - v.iter().map(|x| x / norm).collect() - } else { - v.clone() - } - }) - .collect()) - } + TransformOp::Normalize => Ok(vectors + .iter() + .map(|v| { + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + v.iter().map(|x| x / norm).collect() + } else { + v.clone() + } + }) + .collect()), TransformOp::ReduceDimensions { target_dim } => { // Simplified: just truncate (would use PCA in production) @@ -1222,7 +1269,7 @@ impl TransformAgent { .iter() .map(|v| v.iter().take(*target_dim).cloned().collect()) .collect()) - } + }, TransformOp::Deduplicate { threshold } => { // Simple deduplication based on cosine similarity @@ -1242,7 +1289,7 @@ impl TransformAgent { } Ok(unique) - } + }, _ => Ok(vectors.to_vec()), } @@ -1449,9 +1496,10 @@ impl ComplianceAgent { // Check data retention if let Some(ref policy) = operation.retention_policy - && policy.max_retention_days < 30 { - warnings.push("Short retention policy may affect analytics".to_string()); - } + && policy.max_retention_days < 30 + { + warnings.push("Short retention policy may affect analytics".to_string()); + } // Check access permissions if operation.requires_pii && !operation.has_pii_access { @@ -1481,7 +1529,12 @@ impl ComplianceAgent { } /// Track lineage for an operation - pub fn track_lineage(&self, operation: &str, inputs: &[String], outputs: &[String]) -> LineageRecord { + pub fn track_lineage( + &self, + operation: &str, + inputs: &[String], + outputs: &[String], + ) -> LineageRecord { LineageRecord { operation: operation.to_string(), timestamp: chrono::Utc::now().timestamp(), @@ -1587,10 +1640,7 @@ mod tests { async fn test_transform_agent() { let agent = TransformAgent::new(AgentConfig::default()); - let vectors = vec![ - vec![1.0, 2.0, 3.0], - vec![4.0, 5.0, 6.0], - ]; + let vectors = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]; let result = agent.transform("normalize", &vectors).await.unwrap(); @@ -1624,10 +1674,7 @@ mod tests { }, ); - let results = vec![ - ("item1".to_string(), 0.8), - ("item2".to_string(), 0.9), - ]; + let results = vec![("item1".to_string(), 0.8), ("item2".to_string(), 0.9)]; let personalized = agent.personalize("user1", &results).await.unwrap(); diff --git a/src/agentic.rs b/src/agentic.rs index acdcf7f..9c76cd9 100644 --- a/src/agentic.rs +++ b/src/agentic.rs @@ -525,10 +525,12 @@ impl QueryDecomposer { // Look for capitalized words as potential entities for word in query.split_whitespace() { - if word.len() > 1 && word.chars().next().unwrap().is_uppercase() - && !STOP_WORDS.contains(&word.to_lowercase().as_str()) { - entities.push(word.to_string()); - } + if word.len() > 1 + && word.chars().next().unwrap().is_uppercase() + && !STOP_WORDS.contains(&word.to_lowercase().as_str()) + { + entities.push(word.to_string()); + } } for (i, entity) in entities.iter().take(self.max_sub_queries).enumerate() { @@ -563,9 +565,15 @@ impl QueryDecomposer { let query_lower = query.to_lowercase(); // Check for temporal indicators - let has_recent = query_lower.contains("recent") || query_lower.contains("latest") || query_lower.contains("new"); - let has_past = query_lower.contains("previous") || query_lower.contains("old") || query_lower.contains("historical"); - let has_future = query_lower.contains("upcoming") || query_lower.contains("future") || query_lower.contains("next"); + let has_recent = query_lower.contains("recent") + || query_lower.contains("latest") + || query_lower.contains("new"); + let has_past = query_lower.contains("previous") + || query_lower.contains("old") + || query_lower.contains("historical"); + let has_future = query_lower.contains("upcoming") + || query_lower.contains("future") + || query_lower.contains("next"); if has_recent { sub_queries.push(SubQuery { @@ -624,11 +632,17 @@ impl QueryDecomposer { // Common aspects to search for let aspects = [ - ("price", vec!["price", "cost", "cheap", "expensive", "affordable"]), + ( + "price", + vec!["price", "cost", "cheap", "expensive", "affordable"], + ), ("quality", vec!["quality", "best", "top", "premium", "good"]), ("reviews", vec!["review", "rating", "feedback", "opinion"]), ("features", vec!["feature", "specification", "capability"]), - ("comparison", vec!["compare", "vs", "versus", "better", "difference"]), + ( + "comparison", + vec!["compare", "vs", "versus", "better", "difference"], + ), ]; let query_lower = query.to_lowercase(); @@ -685,14 +699,23 @@ impl ResultFuser { FusionStrategy::RRF { k } => self.fuse_rrf(branch_results, *k, top_k), FusionStrategy::WeightedSum => self.fuse_weighted(branch_results, top_k), FusionStrategy::MaxScore => self.fuse_max(branch_results, top_k), - FusionStrategy::Voting { min_votes } => self.fuse_voting(branch_results, *min_votes, top_k), + FusionStrategy::Voting { min_votes } => { + self.fuse_voting(branch_results, *min_votes, top_k) + }, FusionStrategy::LLMRerank => self.fuse_weighted(branch_results, top_k), // Fallback FusionStrategy::CrossEncoder => self.fuse_weighted(branch_results, top_k), // Fallback } } fn fuse_rrf(&self, branch_results: &[BranchResult], k: f32, top_k: usize) -> Vec { - let mut doc_scores: HashMap, HashMap)> = HashMap::new(); + let mut doc_scores: HashMap< + String, + ( + f32, + HashMap, + HashMap, + ), + > = HashMap::new(); for branch in branch_results { if !branch.success { @@ -702,9 +725,9 @@ impl ResultFuser { for (rank, result) in branch.results.iter().enumerate() { let rrf_score = 1.0 / (k + rank as f32 + 1.0); - let entry = doc_scores.entry(result.id.clone()).or_insert_with(|| { - (0.0, HashMap::new(), result.metadata.clone()) - }); + let entry = doc_scores + .entry(result.id.clone()) + .or_insert_with(|| (0.0, HashMap::new(), result.metadata.clone())); entry.0 += rrf_score; entry.1.insert(branch.sub_query_id.clone(), result.score); @@ -729,7 +752,14 @@ impl ResultFuser { } fn fuse_weighted(&self, branch_results: &[BranchResult], top_k: usize) -> Vec { - let mut doc_scores: HashMap, HashMap)> = HashMap::new(); + let mut doc_scores: HashMap< + String, + ( + f32, + HashMap, + HashMap, + ), + > = HashMap::new(); for branch in branch_results { if !branch.success { @@ -737,9 +767,9 @@ impl ResultFuser { } for result in &branch.results { - let entry = doc_scores.entry(result.id.clone()).or_insert_with(|| { - (0.0, HashMap::new(), result.metadata.clone()) - }); + let entry = doc_scores + .entry(result.id.clone()) + .or_insert_with(|| (0.0, HashMap::new(), result.metadata.clone())); entry.0 += result.score; entry.1.insert(branch.sub_query_id.clone(), result.score); @@ -764,7 +794,14 @@ impl ResultFuser { } fn fuse_max(&self, branch_results: &[BranchResult], top_k: usize) -> Vec { - let mut doc_scores: HashMap, HashMap)> = HashMap::new(); + let mut doc_scores: HashMap< + String, + ( + f32, + HashMap, + HashMap, + ), + > = HashMap::new(); for branch in branch_results { if !branch.success { @@ -772,9 +809,9 @@ impl ResultFuser { } for result in &branch.results { - let entry = doc_scores.entry(result.id.clone()).or_insert_with(|| { - (0.0, HashMap::new(), result.metadata.clone()) - }); + let entry = doc_scores + .entry(result.id.clone()) + .or_insert_with(|| (0.0, HashMap::new(), result.metadata.clone())); if result.score > entry.0 { entry.0 = result.score; @@ -800,8 +837,20 @@ impl ResultFuser { fused } - fn fuse_voting(&self, branch_results: &[BranchResult], min_votes: usize, top_k: usize) -> Vec { - let mut doc_votes: HashMap, HashMap)> = HashMap::new(); + fn fuse_voting( + &self, + branch_results: &[BranchResult], + min_votes: usize, + top_k: usize, + ) -> Vec { + let mut doc_votes: HashMap< + String, + ( + usize, + HashMap, + HashMap, + ), + > = HashMap::new(); for branch in branch_results { if !branch.success { @@ -809,9 +858,9 @@ impl ResultFuser { } for result in &branch.results { - let entry = doc_votes.entry(result.id.clone()).or_insert_with(|| { - (0, HashMap::new(), result.metadata.clone()) - }); + let entry = doc_votes + .entry(result.id.clone()) + .or_insert_with(|| (0, HashMap::new(), result.metadata.clone())); entry.0 += 1; entry.1.insert(branch.sub_query_id.clone(), result.score); @@ -952,7 +1001,9 @@ impl AgentExecutor { /// Save agent state pub fn save_state(&self, state: &AgentState) { let key = format!("{}:{}", state.agent_id, state.session_id); - let Ok(mut states) = self.states.write() else { return; }; + let Ok(mut states) = self.states.write() else { + return; + }; states.insert(key, state.clone()); } @@ -1006,7 +1057,9 @@ impl AgentExecutor { /// Register a tool pub fn register_tool(&self, agent_id: &str, session_id: &str, tool: ToolDefinition) { let key = format!("{}:{}", agent_id, session_id); - let Ok(mut states) = self.states.write() else { return; }; + let Ok(mut states) = self.states.write() else { + return; + }; if let Some(state) = states.get_mut(&key) { state.register_tool(tool); @@ -1068,13 +1121,12 @@ pub struct AgentStats { // Stop words for filtering const STOP_WORDS: &[&str] = &[ - "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for", - "of", "with", "by", "from", "as", "is", "was", "are", "were", "been", - "be", "have", "has", "had", "do", "does", "did", "will", "would", "could", - "should", "may", "might", "must", "shall", "can", "need", "dare", "ought", - "used", "that", "which", "who", "whom", "this", "these", "those", "what", - "i", "me", "my", "we", "our", "you", "your", "he", "him", "his", "she", - "her", "it", "its", "they", "them", "their", + "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "by", + "from", "as", "is", "was", "are", "were", "been", "be", "have", "has", "had", "do", "does", + "did", "will", "would", "could", "should", "may", "might", "must", "shall", "can", "need", + "dare", "ought", "used", "that", "which", "who", "whom", "this", "these", "those", "what", "i", + "me", "my", "we", "our", "you", "your", "he", "him", "his", "she", "her", "it", "its", "they", + "them", "their", ]; #[cfg(test)] @@ -1084,7 +1136,9 @@ mod tests { #[test] fn test_query_decomposition() { let decomposer = QueryDecomposer::new(DecompositionStrategy::Semantic, 5); - let sub_queries = decomposer.decompose("Find products similar to iPhone and cheaper than $500").unwrap(); + let sub_queries = decomposer + .decompose("Find products similar to iPhone and cheaper than $500") + .unwrap(); assert!(!sub_queries.is_empty()); } @@ -1174,7 +1228,9 @@ mod tests { #[test] fn test_query_planner() { let planner = QueryPlanner::new(DecompositionStrategy::Semantic, 5); - let plan = planner.plan("Find red cars with good fuel efficiency").unwrap(); + let plan = planner + .plan("Find red cars with good fuel efficiency") + .unwrap(); assert!(!plan.sub_queries.is_empty()); assert!(!plan.execution_order.is_empty()); diff --git a/src/analytics.rs b/src/analytics.rs index 2fb5239..4110462 100644 --- a/src/analytics.rs +++ b/src/analytics.rs @@ -584,7 +584,8 @@ impl VectorAnalytics { let mut min_dist = f32::MAX; for (j, v) in vectors.iter().enumerate() { if j != idx { - let dist: f32 = point.iter() + let dist: f32 = point + .iter() .zip(v.iter()) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -606,7 +607,8 @@ impl VectorAnalytics { // Find nearest neighbor distance to data let mut min_dist = f32::MAX; for v in vectors { - let dist: f32 = random_point.iter() + let dist: f32 = random_point + .iter() .zip(v.iter()) .map(|(a, b)| (a - b).powi(2)) .sum::() diff --git a/src/anomaly.rs b/src/anomaly.rs index 0ee8416..f3dff8f 100644 --- a/src/anomaly.rs +++ b/src/anomaly.rs @@ -165,7 +165,7 @@ impl IsolationForest { IsolationTree::Leaf { size } => { // Adjust for unsuccessful search height as f32 + Self::average_path_length(*size) - } + }, IsolationTree::Node { feature, split_value, @@ -177,7 +177,7 @@ impl IsolationForest { } else { self.path_length(vector, right, height + 1) } - } + }, } } @@ -278,11 +278,7 @@ impl LocalOutlierFactor { .enumerate() .filter_map(|(i, v)| { let dist = euclidean_distance_simd(vector, v); - if dist > 0.0 { - Some((i, dist)) - } else { - None - } + if dist > 0.0 { Some((i, dist)) } else { None } }) .collect(); diff --git a/src/anomaly_detection.rs b/src/anomaly_detection.rs index 1418765..3b5e38e 100644 --- a/src/anomaly_detection.rs +++ b/src/anomaly_detection.rs @@ -2,8 +2,8 @@ // Outlier detection, drift monitoring, and data quality analysis use std::collections::{HashMap, VecDeque}; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use serde::{Deserialize, Serialize}; @@ -240,27 +240,29 @@ impl AnomalyDetector { let model = match self.config.method { DetectionMethod::IsolationForest => { DetectorModel::IsolationForest(self.build_isolation_forest(vectors)?) - } - DetectionMethod::LOF => { - DetectorModel::LOF(self.build_lof(vectors)?) - } + }, + DetectionMethod::LOF => DetectorModel::LOF(self.build_lof(vectors)?), DetectionMethod::Statistical => { DetectorModel::Statistical(self.build_statistical(vectors)?) - } + }, DetectionMethod::Ensemble => { let detectors: Vec> = Vec::new(); // Would add multiple detectors here DetectorModel::Ensemble(detectors) - } - _ => { - DetectorModel::Statistical(self.build_statistical(vectors)?) - } + }, + _ => DetectorModel::Statistical(self.build_statistical(vectors)?), }; - *self.model.write() + *self + .model + .write() .map_err(|_| VecStoreError::LockError("model lock poisoned".into()))? = Some(model); - *self.stats.last_training.write() - .map_err(|_| VecStoreError::LockError("last_training lock poisoned".into()))? = Some(Instant::now()); + *self + .stats + .last_training + .write() + .map_err(|_| VecStoreError::LockError("last_training lock poisoned".into()))? = + Some(Instant::now()); Ok(()) } @@ -279,16 +281,15 @@ impl AnomalyDetector { .map(|_| random_usize(0, vectors.len())) .collect(); - let samples: Vec<&Vec> = sample_indices.iter() - .map(|&i| &vectors[i]) - .collect(); + let samples: Vec<&Vec> = sample_indices.iter().map(|&i| &vectors[i]).collect(); let tree = self.build_tree(&samples, 0, max_depth); trees.push(tree); } // Compute threshold based on contamination - let scores: Vec = vectors.iter() + let scores: Vec = vectors + .iter() .map(|v| self.compute_path_length(&trees, v)) .collect(); @@ -303,7 +304,9 @@ impl AnomalyDetector { fn build_tree(&self, samples: &[&Vec], depth: usize, max_depth: usize) -> IsolationTree { if samples.len() <= 1 || depth >= max_depth { - return IsolationTree::Leaf { size: samples.len() }; + return IsolationTree::Leaf { + size: samples.len(), + }; } let dim = samples[0].len(); @@ -315,13 +318,15 @@ impl AnomalyDetector { let max_val = values.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); if (max_val - min_val).abs() < 1e-10 { - return IsolationTree::Leaf { size: samples.len() }; + return IsolationTree::Leaf { + size: samples.len(), + }; } let split_value = random_f32(min_val, max_val); - let (left, right): (Vec<_>, Vec<_>) = samples.iter() - .partition(|s| s[feature] < split_value); + let (left, right): (Vec<_>, Vec<_>) = + samples.iter().partition(|s| s[feature] < split_value); IsolationTree::Internal { feature, @@ -332,7 +337,8 @@ impl AnomalyDetector { } fn compute_path_length(&self, trees: &[IsolationTree], sample: &[f32]) -> f64 { - let total: f64 = trees.iter() + let total: f64 = trees + .iter() .map(|tree| self.path_length(tree, sample, 0)) .sum(); total / trees.len() as f64 @@ -340,16 +346,19 @@ impl AnomalyDetector { fn path_length(&self, tree: &IsolationTree, sample: &[f32], depth: usize) -> f64 { match tree { - IsolationTree::Leaf { size } => { - depth as f64 + c_factor(*size) - } - IsolationTree::Internal { feature, split_value, left, right } => { + IsolationTree::Leaf { size } => depth as f64 + c_factor(*size), + IsolationTree::Internal { + feature, + split_value, + left, + right, + } => { if (sample[*feature] as f64) < *split_value { self.path_length(left, sample, depth + 1) } else { self.path_length(right, sample, depth + 1) } - } + }, } } @@ -411,26 +420,37 @@ impl AnomalyDetector { }) .collect(); - let q1: Vec = sorted_values.iter() + let q1: Vec = sorted_values + .iter() .map(|vals| vals[vals.len() / 4]) .collect(); - let q3: Vec = sorted_values.iter() + let q3: Vec = sorted_values + .iter() .map(|vals| vals[vals.len() * 3 / 4]) .collect(); - Ok(StatisticalModel { mean, std, min, max, q1, q3 }) + Ok(StatisticalModel { + mean, + std, + min, + max, + q1, + q3, + }) } /// Detect anomalies in a single vector pub fn detect(&self, vector_id: &str, vector: &[f32]) -> Result { self.stats.total_samples.fetch_add(1, Ordering::Relaxed); - let model = self.model.read() + let model = self + .model + .read() .map_err(|_| VecStoreError::LockError("model lock poisoned".into()))?; - let model = model.as_ref().ok_or({ - VecStoreError::IndexNotInitialized - })?; + let model = model + .as_ref() + .ok_or({ VecStoreError::IndexNotInitialized })?; let (score, method_scores) = match model { DetectorModel::IsolationForest(m) => { @@ -439,19 +459,19 @@ impl AnomalyDetector { let mut scores = HashMap::new(); scores.insert("isolation_forest".to_string(), score); (score, scores) - } + }, DetectorModel::LOF(m) => { let score = self.compute_lof_score(m, vector); let mut scores = HashMap::new(); scores.insert("lof".to_string(), score); (score, scores) - } + }, DetectorModel::Statistical(m) => { let score = self.compute_statistical_score(m, vector); let mut scores = HashMap::new(); scores.insert("statistical".to_string(), score); (score, scores) - } + }, DetectorModel::Ensemble(detectors) => { let mut scores = HashMap::new(); let mut total = 0.0; @@ -460,16 +480,22 @@ impl AnomalyDetector { scores.insert(detector.name().to_string(), s); total += s; } - let avg = if detectors.is_empty() { 0.0 } else { total / detectors.len() as f64 }; + let avg = if detectors.is_empty() { + 0.0 + } else { + total / detectors.len() as f64 + }; (avg, scores) - } + }, }; let is_anomaly = score >= self.config.alert_threshold; let confidence = self.compute_confidence(score); if is_anomaly { - self.stats.anomalies_detected.fetch_add(1, Ordering::Relaxed); + self.stats + .anomalies_detected + .fetch_add(1, Ordering::Relaxed); // Record anomaly let record = AnomalyRecord { @@ -482,7 +508,8 @@ impl AnomalyDetector { explanation: self.generate_explanation(model, vector, score), }; - self.anomalies.write() + self.anomalies + .write() .map_err(|_| VecStoreError::LockError("anomalies lock poisoned".into()))? .push(record); } @@ -541,17 +568,16 @@ impl AnomalyDetector { } // Find k nearest neighbors - let mut distances: Vec<(usize, f64)> = model.samples.iter() + let mut distances: Vec<(usize, f64)> = model + .samples + .iter() .enumerate() .map(|(i, s)| (i, euclidean_distance(sample, s))) .collect(); distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap()); - let k_neighbors: Vec = distances.iter() - .take(model.k) - .map(|(i, _)| *i) - .collect(); + let k_neighbors: Vec = distances.iter().take(model.k).map(|(i, _)| *i).collect(); // Compute local reachability density let k_dist = distances[model.k.min(distances.len() - 1)].1; @@ -561,12 +587,7 @@ impl AnomalyDetector { let mut lof = 0.0; for &neighbor_idx in &k_neighbors { let neighbor = &model.samples[neighbor_idx]; - let neighbor_lrd = self.compute_lrd( - model, - neighbor, - &k_neighbors, - k_dist - ); + let neighbor_lrd = self.compute_lrd(model, neighbor, &k_neighbors, k_dist); if lrd > 0.0 { lof += neighbor_lrd / lrd; } @@ -575,7 +596,13 @@ impl AnomalyDetector { lof / k_neighbors.len() as f64 } - fn compute_lrd(&self, _model: &LOFModel, _sample: &[f32], _neighbors: &[usize], k_dist: f64) -> f64 { + fn compute_lrd( + &self, + _model: &LOFModel, + _sample: &[f32], + _neighbors: &[usize], + k_dist: f64, + ) -> f64 { // Simplified LRD computation if k_dist > 0.0 { 1.0 / k_dist } else { 1.0 } } @@ -622,7 +649,8 @@ impl AnomalyDetector { fn find_contributing_features(&self, model: &DetectorModel, sample: &[f32]) -> Vec { match model { DetectorModel::Statistical(m) => { - let mut contributions: Vec<(usize, f64)> = sample.iter() + let mut contributions: Vec<(usize, f64)> = sample + .iter() .enumerate() .filter_map(|(i, &val)| { if m.std[i] > 0.0 { @@ -636,7 +664,7 @@ impl AnomalyDetector { contributions.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); contributions.iter().take(5).map(|(i, _)| *i).collect() - } + }, _ => vec![], } } @@ -645,22 +673,31 @@ impl AnomalyDetector { fn generate_explanation(&self, model: &DetectorModel, sample: &[f32], score: f64) -> String { match model { DetectorModel::IsolationForest(_) => { - format!("Isolation score {:.3}: Vector is isolated from normal distribution", score) - } + format!( + "Isolation score {:.3}: Vector is isolated from normal distribution", + score + ) + }, DetectorModel::LOF(_) => { - format!("LOF score {:.3}: Vector has low local density compared to neighbors", score) - } + format!( + "LOF score {:.3}: Vector has low local density compared to neighbors", + score + ) + }, DetectorModel::Statistical(_m) => { let features = self.find_contributing_features(model, sample); if features.is_empty() { - format!("Statistical score {:.3}: Overall deviation from normal", score) + format!( + "Statistical score {:.3}: Overall deviation from normal", + score + ) } else { format!( "Statistical score {:.3}: High deviation in dimensions {:?}", score, features ) } - } + }, _ => format!("Anomaly score {:.3}", score), } } @@ -675,17 +712,23 @@ impl AnomalyDetector { detection_rate: { let total = self.stats.total_samples.load(Ordering::Relaxed); let anomalies = self.stats.anomalies_detected.load(Ordering::Relaxed); - if total > 0 { anomalies as f64 / total as f64 } else { 0.0 } + if total > 0 { + anomalies as f64 / total as f64 + } else { + 0.0 + } }, last_training: { - let Ok(guard) = self.stats.last_training.read() else { return DetectionStats { - total_samples: self.stats.total_samples.load(Ordering::Relaxed), - anomalies_detected: self.stats.anomalies_detected.load(Ordering::Relaxed), - false_positives: self.stats.false_positives.load(Ordering::Relaxed), - true_positives: self.stats.true_positives.load(Ordering::Relaxed), - detection_rate: 0.0, - last_training: None, - }; }; + let Ok(guard) = self.stats.last_training.read() else { + return DetectionStats { + total_samples: self.stats.total_samples.load(Ordering::Relaxed), + anomalies_detected: self.stats.anomalies_detected.load(Ordering::Relaxed), + false_positives: self.stats.false_positives.load(Ordering::Relaxed), + true_positives: self.stats.true_positives.load(Ordering::Relaxed), + detection_rate: 0.0, + last_training: None, + }; + }; guard.map(|t| t.elapsed()) }, } @@ -693,12 +736,10 @@ impl AnomalyDetector { /// Get recent anomalies pub fn get_recent_anomalies(&self, limit: usize) -> Vec { - let Ok(anomalies) = self.anomalies.read() else { return vec![]; }; - anomalies.iter() - .rev() - .take(limit) - .cloned() - .collect() + let Ok(anomalies) = self.anomalies.read() else { + return vec![]; + }; + anomalies.iter().rev().take(limit).cloned().collect() } /// Provide feedback on detection @@ -712,7 +753,9 @@ impl AnomalyDetector { /// Online update with new sample pub fn update(&self, vector: &[f32]) -> Result<()> { - let mut history = self.history.write() + let mut history = self + .history + .write() .map_err(|_| VecStoreError::LockError("history lock poisoned".into()))?; if history.len() >= self.config.window_size { @@ -832,14 +875,17 @@ impl DriftDetector { /// Set reference distribution from data pub fn set_reference(&self, vectors: &[Vec]) -> Result<()> { let stats = self.compute_distribution_stats(vectors)?; - *self.reference_distribution.write() - .map_err(|_| VecStoreError::LockError("reference_distribution lock poisoned".into()))? = Some(stats); + *self.reference_distribution.write().map_err(|_| { + VecStoreError::LockError("reference_distribution lock poisoned".into()) + })? = Some(stats); Ok(()) } /// Add sample to detection window pub fn add_sample(&self, vector: Vec) -> Option { - let Ok(mut window) = self.current_window.write() else { return None; }; + let Ok(mut window) = self.current_window.write() else { + return None; + }; window.push_back(vector); if window.len() > self.config.detection_size { @@ -859,7 +905,9 @@ impl DriftDetector { /// Check for drift fn check_drift(&self, current_samples: &[Vec]) -> Result> { - let reference = self.reference_distribution.read() + let reference = self + .reference_distribution + .read() .map_err(|_| VecStoreError::LockError("reference_distribution lock poisoned".into()))?; let reference = match reference.as_ref() { Some(r) => r, @@ -891,7 +939,8 @@ impl DriftDetector { ), }; - self.drift_events.write() + self.drift_events + .write() .map_err(|_| VecStoreError::LockError("drift_events lock poisoned".into()))? .push(event.clone()); return Ok(Some(event)); @@ -999,7 +1048,8 @@ impl DriftDetector { let mut total_dist = 0.0; for (ref_q, curr_q) in ref_stats.quantiles.iter().zip(curr_stats.quantiles.iter()) { - let dist: f64 = ref_q.iter() + let dist: f64 = ref_q + .iter() .zip(curr_q.iter()) .map(|(r, c)| (r - c).abs()) .sum(); @@ -1056,7 +1106,9 @@ impl DriftDetector { /// Get recent drift events pub fn get_drift_events(&self, limit: usize) -> Vec { - let Ok(events) = self.drift_events.read() else { return vec![]; }; + let Ok(events) = self.drift_events.read() else { + return vec![]; + }; events.iter().rev().take(limit).cloned().collect() } } @@ -1118,9 +1170,11 @@ fn current_timestamp() -> u64 { } fn random_usize(min: usize, max: usize) -> usize { - if max <= min { return min; } - use std::hash::{Hash, Hasher}; + if max <= min { + return min; + } use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); Instant::now().hash(&mut hasher); let hash = hasher.finish(); @@ -1128,8 +1182,8 @@ fn random_usize(min: usize, max: usize) -> usize { } fn random_f32(min: f32, max: f32) -> f32 { - use std::hash::{Hash, Hasher}; use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); Instant::now().hash(&mut hasher); let hash = hasher.finish(); @@ -1145,7 +1199,9 @@ fn euclidean_distance(a: &[f32], b: &[f32]) -> f64 { } fn percentile(values: &[f64], p: f64) -> f64 { - if values.is_empty() { return 0.0; } + if values.is_empty() { + return 0.0; + } let mut sorted = values.to_vec(); sorted.sort_by(|a, b| a.partial_cmp(b).unwrap()); let idx = ((sorted.len() as f64 - 1.0) * p) as usize; @@ -1153,14 +1209,18 @@ fn percentile(values: &[f64], p: f64) -> f64 { } fn c_factor(size: usize) -> f64 { - if size <= 1 { return 0.0; } + if size <= 1 { + return 0.0; + } let n = size as f64; 2.0 * ((n - 1.0).ln() + 0.5772156649) - (2.0 * (n - 1.0) / n) } fn anomaly_score(path_length: f64, sample_size: usize) -> f64 { let c = c_factor(sample_size); - if c == 0.0 { return 0.5; } + if c == 0.0 { + return 0.5; + } (2.0_f64).powf(-path_length / c) } @@ -1205,9 +1265,7 @@ mod tests { .build(); // Generate normal data - let normal_data: Vec> = (0..100) - .map(|_| vec![0.5, 0.5, 0.5]) - .collect(); + let normal_data: Vec> = (0..100).map(|_| vec![0.5, 0.5, 0.5]).collect(); detector.fit(&normal_data).unwrap(); @@ -1224,9 +1282,7 @@ mod tests { fn test_drift_detector() { let detector = DriftDetector::new(DriftConfig::default()); - let reference: Vec> = (0..100) - .map(|_| vec![0.5, 0.5, 0.5]) - .collect(); + let reference: Vec> = (0..100).map(|_| vec![0.5, 0.5, 0.5]).collect(); detector.set_reference(&reference).unwrap(); } diff --git a/src/async_ops.rs b/src/async_ops.rs index 2036851..aba22dd 100644 --- a/src/async_ops.rs +++ b/src/async_ops.rs @@ -107,9 +107,9 @@ impl AsyncStoreOps { for chunk in chunks { let store_clone = self.store.clone(); let handle = task::spawn_blocking(move || { - let mut store = store_clone - .lock() - .map_err(|_| VecStoreError::LockError("failed to acquire store lock for batch upsert".into()))?; + let mut store = store_clone.lock().map_err(|_| { + VecStoreError::LockError("failed to acquire store lock for batch upsert".into()) + })?; for (id, vector, metadata) in chunk { let metadata: Metadata = serde_json::from_value(metadata)?; store.upsert(id, vector, metadata)?; @@ -158,9 +158,9 @@ impl AsyncStoreOps { for query_vec in queries { let store_clone = self.store.clone(); let handle = task::spawn_blocking(move || { - let store = store_clone - .lock() - .map_err(|_| VecStoreError::LockError("failed to acquire store lock for batch query".into()))?; + let store = store_clone.lock().map_err(|_| { + VecStoreError::LockError("failed to acquire store lock for batch query".into()) + })?; let query = Query::new(query_vec).with_limit(k); store.query(query) }); @@ -183,9 +183,9 @@ impl AsyncStoreOps { for id in ids { let store_clone = self.store.clone(); let handle = task::spawn_blocking(move || { - let mut store = store_clone - .lock() - .map_err(|_| VecStoreError::LockError("failed to acquire store lock for batch delete".into()))?; + let mut store = store_clone.lock().map_err(|_| { + VecStoreError::LockError("failed to acquire store lock for batch delete".into()) + })?; store.delete(&id) }); handles.push(handle); diff --git a/src/async_stream.rs b/src/async_stream.rs index e9ae92d..edf484b 100644 --- a/src/async_stream.rs +++ b/src/async_stream.rs @@ -273,14 +273,14 @@ impl Stream for ProgressiveQueryStream { match &event { StreamEvent::Result(_) => { self.received += 1; - } + }, StreamEvent::Progress { total, .. } => { self.total_expected = Some(*total); - } - _ => {} + }, + _ => {}, } Poll::Ready(Some(event)) - } + }, Poll::Ready(None) => Poll::Ready(None), Poll::Pending => Poll::Pending, } @@ -402,7 +402,7 @@ mod tests { match event { StreamEvent::Result(_) => result_count += 1, StreamEvent::Done => got_done = true, - _ => {} + _ => {}, } } diff --git a/src/audit.rs b/src/audit.rs index f975332..d68cc9d 100644 --- a/src/audit.rs +++ b/src/audit.rs @@ -63,8 +63,7 @@ pub enum AuditOutcome { } /// Audit entry metadata -#[derive(Debug, Clone, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct AuditMetadata { /// User identifier pub user_id: Option, @@ -83,7 +82,6 @@ pub struct AuditMetadata { pub custom: std::collections::HashMap, } - /// Audit log entry #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AuditEntry { @@ -226,13 +224,16 @@ impl MemoryBackend { } pub fn get_entries(&self) -> Result, VecStoreError> { - let entries = self.entries.lock() + let entries = self + .entries + .lock() .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?; Ok(entries.iter().cloned().collect()) } pub fn clear(&self) -> Result<(), VecStoreError> { - self.entries.lock() + self.entries + .lock() .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))? .clear(); Ok(()) @@ -241,7 +242,9 @@ impl MemoryBackend { impl AuditBackend for MemoryBackend { fn write(&mut self, entry: &AuditEntry) -> Result<(), String> { - let mut entries = self.entries.lock() + let mut entries = self + .entries + .lock() .map_err(|_| "entries lock poisoned".to_string())?; if entries.len() >= self.max_size { entries.pop_front(); @@ -288,7 +291,9 @@ impl AuditBackend for FileBackend { let json = serde_json::to_string(entry) .map_err(|e| format!("Failed to serialize audit entry: {}", e))?; - let mut buffer = self.buffer.lock() + let mut buffer = self + .buffer + .lock() .map_err(|_| "buffer lock poisoned".to_string())?; buffer.push(json); @@ -301,13 +306,17 @@ impl AuditBackend for FileBackend { } fn flush(&mut self) -> Result<(), String> { - let mut buffer = self.buffer.lock() + let mut buffer = self + .buffer + .lock() .map_err(|_| "buffer lock poisoned".to_string())?; if buffer.is_empty() { return Ok(()); } - let mut file = self.file.lock() + let mut file = self + .file + .lock() .map_err(|_| "file lock poisoned".to_string())?; for line in buffer.iter() { @@ -396,7 +405,8 @@ impl AuditLogger { /// Add a backend pub fn add_backend(&self, backend: Box) -> Result<(), VecStoreError> { - self.backends.lock() + self.backends + .lock() .map_err(|_| VecStoreError::LockError("backends lock poisoned".into()))? .push(backend); Ok(()) @@ -421,7 +431,9 @@ impl AuditLogger { } // Write to all backends - let mut backends = self.backends.lock() + let mut backends = self + .backends + .lock() .map_err(|_| "backends lock poisoned".to_string())?; for backend in backends.iter_mut() { backend.write(&entry)?; @@ -432,7 +444,9 @@ impl AuditLogger { /// Flush all backends pub fn flush(&self) -> Result<(), String> { - let mut backends = self.backends.lock() + let mut backends = self + .backends + .lock() .map_err(|_| "backends lock poisoned".to_string())?; for backend in backends.iter_mut() { backend.flush()?; diff --git a/src/auto_tune.rs b/src/auto_tune.rs index de92e6b..e4dab19 100644 --- a/src/auto_tune.rs +++ b/src/auto_tune.rs @@ -35,7 +35,7 @@ //! println!("Optimal params: {:?}", params); //! ``` -use anyhow::{Result}; +use anyhow::Result; use rand::Rng; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -193,15 +193,13 @@ impl Default for HybridParams { } /// Combined index parameters -#[derive(Debug, Clone, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct IndexParams { pub hnsw: HNSWParams, pub ivf: IVFParams, pub hybrid: HybridParams, } - // ============================================================================ // EVALUATION METRICS // ============================================================================ @@ -367,7 +365,9 @@ impl AutoTuner { F: FnMut(&IndexParams, &[f32], usize) -> Result<(Vec, std::time::Duration)>, { if queries.len() != ground_truth.len() { - return Err(anyhow::anyhow!("Queries and ground truth must have same length")); + return Err(anyhow::anyhow!( + "Queries and ground truth must have same length" + )); } let k = ground_truth.first().map(|g| g.len()).unwrap_or(10); @@ -382,7 +382,8 @@ impl AutoTuner { }; // Evaluate - let eval = self.evaluate(&candidate_params, queries, ground_truth, k, &mut search_fn)?; + let eval = + self.evaluate(&candidate_params, queries, ground_truth, k, &mut search_fn)?; // Record observation self.optimizer.observe(param_vec.clone(), eval.score); @@ -403,7 +404,15 @@ impl AutoTuner { } // Update best params - if eval.score > self.score_params(&self.current_params, queries, ground_truth, k, &mut search_fn)? { + if eval.score + > self.score_params( + &self.current_params, + queries, + ground_truth, + k, + &mut search_fn, + )? + { self.current_params = eval.params; } } @@ -440,10 +449,7 @@ impl AutoTuner { latencies.push(duration.as_secs_f64() * 1000.0); // Calculate recall - let hits = results - .iter() - .filter(|r| truth.contains(r)) - .count(); + let hits = results.iter().filter(|r| truth.contains(r)).count(); total_recall += hits as f32 / k as f32; } @@ -453,7 +459,10 @@ impl AutoTuner { // Calculate latency statistics latencies.sort_by(|a, b| a.total_cmp(b)); let latency_ms = latencies.iter().sum::() / latencies.len() as f64; - let latency_p99_ms = latencies.get(latencies.len() * 99 / 100).copied().unwrap_or(0.0); + let latency_p99_ms = latencies + .get(latencies.len() * 99 / 100) + .copied() + .unwrap_or(0.0); let qps = 1000.0 / latency_ms; // Compute score (balance recall and latency) diff --git a/src/auto_tuning.rs b/src/auto_tuning.rs index e526de2..d00be45 100644 --- a/src/auto_tuning.rs +++ b/src/auto_tuning.rs @@ -1,11 +1,11 @@ // Auto-Tuning Engine - Automated parameter optimization for vector indexes // Bayesian optimization, hyperparameter search, and adaptive configuration +use std::cmp::Ordering as CmpOrdering; use std::collections::HashMap; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; -use std::cmp::Ordering as CmpOrdering; use serde::{Deserialize, Serialize}; @@ -132,7 +132,12 @@ pub struct Parameter { #[derive(Debug, Clone, Serialize, Deserialize)] pub enum ParameterType { /// Integer range - Integer { min: i64, max: i64, step: Option, log_scale: bool }, + Integer { + min: i64, + max: i64, + step: Option, + log_scale: bool, + }, /// Float range Float { min: f64, max: f64, log_scale: bool }, /// Categorical choices @@ -319,8 +324,13 @@ impl AutoTuner { let params = self.sample_random(); let _result = self.evaluate_trial(¶ms, &objective_fn)?; - improvement_history.push(*self.stats.best_objective.read() - .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?); + improvement_history.push( + *self + .stats + .best_objective + .read() + .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?, + ); } // Main optimization loop @@ -341,10 +351,16 @@ impl AutoTuner { SearchStrategy::PopulationBasedTraining => self.sample_random(), }; - let prev_best = *self.stats.best_objective.read() + let prev_best = *self + .stats + .best_objective + .read() .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?; let _result = self.evaluate_trial(¶ms, &objective_fn)?; - let new_best = *self.stats.best_objective.read() + let new_best = *self + .stats + .best_objective + .read() .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?; improvement_history.push(new_best); @@ -361,19 +377,30 @@ impl AutoTuner { } let total_time = start.elapsed(); - *self.stats.total_time.write() + *self + .stats + .total_time + .write() .map_err(|_| VecStoreError::LockError("total_time lock poisoned".into()))? = total_time; - let trials = self.trials.read() - .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))?.clone(); - let best_trial = self.best_trial.read() - .map_err(|_| VecStoreError::LockError("best_trial lock poisoned".into()))?.clone(); + let trials = self + .trials + .read() + .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))? + .clone(); + let best_trial = self + .best_trial + .read() + .map_err(|_| VecStoreError::LockError("best_trial lock poisoned".into()))? + .clone(); Ok(TuneResult { - best_params: best_trial.as_ref() + best_params: best_trial + .as_ref() .map(|t| t.params.clone()) .unwrap_or_default(), - best_objective: best_trial.as_ref() + best_objective: best_trial + .as_ref() .and_then(|t| t.objective) .unwrap_or(f64::NEG_INFINITY), trials, @@ -417,32 +444,38 @@ impl AutoTuner { self.stats.completed_trials.fetch_add(1, Ordering::Relaxed); // Update best - let current_best = *self.stats.best_objective.read() - .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?; + let current_best = + *self.stats.best_objective.read().map_err(|_| { + VecStoreError::LockError("best_objective lock poisoned".into()) + })?; let is_better = match self.config.direction { OptimizeDirection::Maximize => objective > current_best, OptimizeDirection::Minimize => objective < current_best, }; if is_better { - *self.stats.best_objective.write() - .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))? = objective; - *self.best_trial.write() - .map_err(|_| VecStoreError::LockError("best_trial lock poisoned".into()))? = Some(trial.clone()); + *self.stats.best_objective.write().map_err(|_| { + VecStoreError::LockError("best_objective lock poisoned".into()) + })? = objective; + *self.best_trial.write().map_err(|_| { + VecStoreError::LockError("best_trial lock poisoned".into()) + })? = Some(trial.clone()); } // Update surrogate model self.update_surrogate(&trial)?; - } + }, Err(_) => { trial.state = TrialState::Failed; trial.duration = Some(start.elapsed()); self.stats.failed_trials.fetch_add(1, Ordering::Relaxed); - } + }, } - self.trials.write() - .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))?.push(trial.clone()); + self.trials + .write() + .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))? + .push(trial.clone()); Ok(trial) } @@ -452,7 +485,12 @@ impl AutoTuner { for param in &self.parameter_space.parameters { let value = match ¶m.param_type { - ParameterType::Integer { min, max, step, log_scale } => { + ParameterType::Integer { + min, + max, + step, + log_scale, + } => { let val = if *log_scale { let log_min = (*min as f64).ln(); let log_max = (*max as f64).ln(); @@ -467,8 +505,12 @@ impl AutoTuner { val }; ParameterValue::Integer(val.clamp(*min, *max)) - } - ParameterType::Float { min, max, log_scale } => { + }, + ParameterType::Float { + min, + max, + log_scale, + } => { let val = if *log_scale { let log_min = min.ln(); let log_max = max.ln(); @@ -477,14 +519,12 @@ impl AutoTuner { random_f64(*min, *max) }; ParameterValue::Float(val) - } + }, ParameterType::Categorical { choices } | ParameterType::Ordinal { choices } => { let idx = random_usize(0, choices.len()); ParameterValue::String(choices[idx].clone()) - } - ParameterType::Boolean => { - ParameterValue::Boolean(random_bool()) - } + }, + ParameterType::Boolean => ParameterValue::Boolean(random_bool()), }; params.insert(param.name.clone(), value); } @@ -506,24 +546,24 @@ impl AutoTuner { let pos = idx % range; idx /= range; ParameterValue::Integer(*min + pos as i64) - } + }, ParameterType::Float { min, max, .. } => { let steps = 10; let pos = idx % steps; idx /= steps; let val = min + (max - min) * (pos as f64 / (steps - 1) as f64); ParameterValue::Float(val) - } + }, ParameterType::Categorical { choices } | ParameterType::Ordinal { choices } => { let pos = idx % choices.len(); idx /= choices.len(); ParameterValue::String(choices[pos].clone()) - } + }, ParameterType::Boolean => { let val = idx.is_multiple_of(2); idx /= 2; ParameterValue::Boolean(val) - } + }, }; params.insert(param.name.clone(), value); } @@ -533,7 +573,9 @@ impl AutoTuner { /// Sample using Bayesian optimization fn sample_bayesian(&self) -> Result> { - let trials = self.trials.read() + let trials = self + .trials + .read() .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))?; if trials.len() < self.config.n_initial_samples { @@ -560,14 +602,19 @@ impl AutoTuner { /// Compute acquisition value for a point fn compute_acquisition(&self, params: &HashMap) -> Result { - let surrogate = self.surrogate_model.read() + let surrogate = self + .surrogate_model + .read() .map_err(|_| VecStoreError::LockError("surrogate_model lock poisoned".into()))?; if let Some(model) = surrogate.as_ref() { let x = self.params_to_vector(params); let (mean, std) = model.predict(&x); - let best_f = *self.stats.best_objective.read() + let best_f = *self + .stats + .best_objective + .read() .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?; // Expected Improvement @@ -582,7 +629,9 @@ impl AutoTuner { /// Sample using TPE fn sample_tpe(&self) -> Result> { - let trials = self.trials.read() + let trials = self + .trials + .read() .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))?; if trials.len() < self.config.n_initial_samples { @@ -590,12 +639,12 @@ impl AutoTuner { } // Split trials into good and bad - let mut sorted_trials: Vec<_> = trials.iter() - .filter(|t| t.objective.is_some()) - .collect(); + let mut sorted_trials: Vec<_> = trials.iter().filter(|t| t.objective.is_some()).collect(); sorted_trials.sort_by(|a, b| { - b.objective.unwrap().partial_cmp(&a.objective.unwrap()) + b.objective + .unwrap() + .partial_cmp(&a.objective.unwrap()) .unwrap_or(CmpOrdering::Equal) }); @@ -608,10 +657,12 @@ impl AutoTuner { let mut params = HashMap::new(); for param in &self.parameter_space.parameters { - let good_values: Vec<_> = good_trials.iter() + let good_values: Vec<_> = good_trials + .iter() .filter_map(|t| t.params.get(¶m.name)) .collect(); - let bad_values: Vec<_> = bad_trials.iter() + let bad_values: Vec<_> = bad_trials + .iter() .filter_map(|t| t.params.get(¶m.name)) .collect(); @@ -633,13 +684,13 @@ impl AutoTuner { return match param_type { ParameterType::Integer { min, max, .. } => { ParameterValue::Integer(random_i64(*min, *max)) - } + }, ParameterType::Float { min, max, .. } => { ParameterValue::Float(random_f64(*min, *max)) - } + }, ParameterType::Categorical { choices } | ParameterType::Ordinal { choices } => { ParameterValue::String(choices[random_usize(0, choices.len())].clone()) - } + }, ParameterType::Boolean => ParameterValue::Boolean(random_bool()), }; } @@ -650,7 +701,9 @@ impl AutoTuner { /// Sample using genetic algorithm fn sample_genetic(&self) -> Result> { - let trials = self.trials.read() + let trials = self + .trials + .read() .map_err(|_| VecStoreError::LockError("trials lock poisoned".into()))?; if trials.len() < 2 { @@ -658,12 +711,12 @@ impl AutoTuner { } // Select two parents from top trials - let mut sorted: Vec<_> = trials.iter() - .filter(|t| t.objective.is_some()) - .collect(); + let mut sorted: Vec<_> = trials.iter().filter(|t| t.objective.is_some()).collect(); sorted.sort_by(|a, b| { - b.objective.unwrap().partial_cmp(&a.objective.unwrap()) + b.objective + .unwrap() + .partial_cmp(&a.objective.unwrap()) .unwrap_or(CmpOrdering::Equal) }); @@ -692,7 +745,11 @@ impl AutoTuner { Ok(child) } - fn mutate_parameter(&self, param_type: &ParameterType, value: Option) -> ParameterValue { + fn mutate_parameter( + &self, + param_type: &ParameterType, + value: Option, + ) -> ParameterValue { match param_type { ParameterType::Integer { min, max, .. } => { let base = match value { @@ -701,7 +758,7 @@ impl AutoTuner { }; let mutation = random_i64(-10, 10); ParameterValue::Integer((base + mutation).clamp(*min, *max)) - } + }, ParameterType::Float { min, max, .. } => { let base = match value { Some(ParameterValue::Float(v)) => v, @@ -709,10 +766,10 @@ impl AutoTuner { }; let mutation = random_f64(-0.1, 0.1) * (max - min); ParameterValue::Float((base + mutation).clamp(*min, *max)) - } + }, ParameterType::Categorical { choices } | ParameterType::Ordinal { choices } => { ParameterValue::String(choices[random_usize(0, choices.len())].clone()) - } + }, ParameterType::Boolean => ParameterValue::Boolean(random_bool()), } } @@ -721,7 +778,9 @@ impl AutoTuner { fn sample_annealing(&self, iteration: usize) -> Result> { let temperature = 1.0 / (1.0 + iteration as f64 * 0.01); - let best = self.best_trial.read() + let best = self + .best_trial + .read() .map_err(|_| VecStoreError::LockError("best_trial lock poisoned".into()))?; if let Some(best_trial) = best.as_ref() { let mut params = best_trial.params.clone(); @@ -729,10 +788,8 @@ impl AutoTuner { // Perturb parameters based on temperature for param in &self.parameter_space.parameters { if random_f64(0.0, 1.0) < temperature { - let new_value = self.mutate_parameter( - ¶m.param_type, - params.get(¶m.name).cloned(), - ); + let new_value = + self.mutate_parameter(¶m.param_type, params.get(¶m.name).cloned()); params.insert(param.name.clone(), new_value); } } @@ -752,7 +809,9 @@ impl AutoTuner { let x = self.params_to_vector(&trial.params); let y = trial.objective.unwrap(); - let mut surrogate = self.surrogate_model.write() + let mut surrogate = self + .surrogate_model + .write() .map_err(|_| VecStoreError::LockError("surrogate_model lock poisoned".into()))?; if surrogate.is_none() { *surrogate = Some(SurrogateModel { @@ -773,12 +832,20 @@ impl AutoTuner { /// Convert parameters to vector for surrogate model fn params_to_vector(&self, params: &HashMap) -> Vec { - self.parameter_space.parameters.iter() + self.parameter_space + .parameters + .iter() .map(|p| { match params.get(&p.name) { Some(ParameterValue::Integer(v)) => *v as f64, Some(ParameterValue::Float(v)) => *v, - Some(ParameterValue::Boolean(v)) => if *v { 1.0 } else { 0.0 }, + Some(ParameterValue::Boolean(v)) => { + if *v { + 1.0 + } else { + 0.0 + } + }, Some(ParameterValue::String(s)) => { // Map categorical to index if let ParameterType::Categorical { choices } = &p.param_type { @@ -786,7 +853,7 @@ impl AutoTuner { } else { 0.0 } - } + }, None => 0.0, } }) @@ -817,7 +884,7 @@ impl AutoTuner { params.insert(p2.clone(), temp.unwrap()); } } - } + }, ConstraintType::GreaterThan(p1, p2) => { if let (Some(v1), Some(v2)) = (params.get(p1), params.get(p2)) { let v1_float = match v1 { @@ -837,8 +904,8 @@ impl AutoTuner { params.insert(p2.clone(), temp.unwrap()); } } - } - _ => {} + }, + _ => {}, } } } @@ -850,9 +917,15 @@ impl AutoTuner { completed_trials: self.stats.completed_trials.load(Ordering::Relaxed), failed_trials: self.stats.failed_trials.load(Ordering::Relaxed), pruned_trials: self.stats.pruned_trials.load(Ordering::Relaxed), - best_objective: *self.stats.best_objective.read() + best_objective: *self + .stats + .best_objective + .read() .map_err(|_| VecStoreError::LockError("best_objective lock poisoned".into()))?, - total_time: *self.stats.total_time.read() + total_time: *self + .stats + .total_time + .read() .map_err(|_| VecStoreError::LockError("total_time lock poisoned".into()))?, }) } @@ -883,7 +956,9 @@ impl SurrogateModel { }; // Simple uncertainty estimate - let min_dist: f64 = self.observations.iter() + let min_dist: f64 = self + .observations + .iter() .map(|(obs_x, _)| self.compute_distance(x, obs_x)) .fold(f64::INFINITY, f64::min); @@ -930,32 +1005,45 @@ pub fn hnsw_parameter_space() -> ParameterSpace { parameters: vec![ Parameter { name: "m".to_string(), - param_type: ParameterType::Integer { min: 4, max: 64, step: Some(4), log_scale: false }, + param_type: ParameterType::Integer { + min: 4, + max: 64, + step: Some(4), + log_scale: false, + }, description: "Number of connections per layer".to_string(), default: ParameterValue::Integer(16), }, Parameter { name: "ef_construction".to_string(), - param_type: ParameterType::Integer { min: 50, max: 500, step: Some(50), log_scale: false }, + param_type: ParameterType::Integer { + min: 50, + max: 500, + step: Some(50), + log_scale: false, + }, description: "Size of dynamic candidate list during construction".to_string(), default: ParameterValue::Integer(100), }, Parameter { name: "ef_search".to_string(), - param_type: ParameterType::Integer { min: 10, max: 500, step: Some(10), log_scale: false }, + param_type: ParameterType::Integer { + min: 10, + max: 500, + step: Some(10), + log_scale: false, + }, description: "Size of dynamic candidate list during search".to_string(), default: ParameterValue::Integer(50), }, ], - constraints: vec![ - Constraint { - expression: "ef_construction >= ef_search".to_string(), - constraint_type: ConstraintType::GreaterThan( - "ef_construction".to_string(), - "ef_search".to_string(), - ), - }, - ], + constraints: vec![Constraint { + expression: "ef_construction >= ef_search".to_string(), + constraint_type: ConstraintType::GreaterThan( + "ef_construction".to_string(), + "ef_search".to_string(), + ), + }], } } @@ -965,26 +1053,31 @@ pub fn ivf_parameter_space() -> ParameterSpace { parameters: vec![ Parameter { name: "nlist".to_string(), - param_type: ParameterType::Integer { min: 16, max: 65536, step: None, log_scale: true }, + param_type: ParameterType::Integer { + min: 16, + max: 65536, + step: None, + log_scale: true, + }, description: "Number of clusters".to_string(), default: ParameterValue::Integer(100), }, Parameter { name: "nprobe".to_string(), - param_type: ParameterType::Integer { min: 1, max: 256, step: None, log_scale: false }, + param_type: ParameterType::Integer { + min: 1, + max: 256, + step: None, + log_scale: false, + }, description: "Number of clusters to search".to_string(), default: ParameterValue::Integer(8), }, ], - constraints: vec![ - Constraint { - expression: "nprobe <= nlist".to_string(), - constraint_type: ConstraintType::LessThan( - "nprobe".to_string(), - "nlist".to_string(), - ), - }, - ], + constraints: vec![Constraint { + expression: "nprobe <= nlist".to_string(), + constraint_type: ConstraintType::LessThan("nprobe".to_string(), "nlist".to_string()), + }], } } @@ -1029,12 +1122,19 @@ impl WorkloadAnalyzer { } /// Record a query - pub fn record_query(&self, latency_ms: f64, result_count: usize, filter_complexity: usize) -> Result<()> { + pub fn record_query( + &self, + latency_ms: f64, + result_count: usize, + filter_complexity: usize, + ) -> Result<()> { if random_f64(0.0, 1.0) > self.config.sample_rate { return Ok(()); } - let mut queries = self.queries.write() + let mut queries = self + .queries + .write() .map_err(|_| VecStoreError::LockError("queries lock poisoned".into()))?; if queries.len() >= self.config.max_samples { queries.remove(0); @@ -1051,7 +1151,9 @@ impl WorkloadAnalyzer { /// Analyze workload and recommend parameters pub fn recommend(&self) -> Result { - let queries = self.queries.read() + let queries = self + .queries + .read() .map_err(|_| VecStoreError::LockError("queries lock poisoned".into()))?; if queries.is_empty() { @@ -1064,8 +1166,13 @@ impl WorkloadAnalyzer { }); } - let avg_latency: f64 = queries.iter().map(|q| q.latency_ms).sum::() / queries.len() as f64; - let avg_filter_complexity: f64 = queries.iter().map(|q| q.filter_complexity as f64).sum::() / queries.len() as f64; + let avg_latency: f64 = + queries.iter().map(|q| q.latency_ms).sum::() / queries.len() as f64; + let avg_filter_complexity: f64 = queries + .iter() + .map(|q| q.filter_complexity as f64) + .sum::() + / queries.len() as f64; let latency_sensitive = avg_latency > 100.0; let filter_heavy = avg_filter_complexity > 2.0; @@ -1099,8 +1206,8 @@ fn current_timestamp() -> u64 { } fn random_f64(min: f64, max: f64) -> f64 { - use std::hash::{Hash, Hasher}; use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); Instant::now().hash(&mut hasher); diff --git a/src/autotuning.rs b/src/autotuning.rs index 9a126d9..c91eb3b 100644 --- a/src/autotuning.rs +++ b/src/autotuning.rs @@ -259,17 +259,17 @@ impl AutoTuner { // Adjust ef_construction and ef_search to meet recall constraints if let Some(recall) = params.estimated_recall - && recall < constraints.min_recall { - // Increase ef_search to improve recall - let recall_deficit = constraints.min_recall - recall; - let boost = 1.0 + recall_deficit * 2.0; // Heuristic - params.ef_search = (params.ef_search as f32 * boost) as usize; - params.ef_search = params.ef_search.min(500); - - // Update estimated recall - params.estimated_recall = - Some(params.estimated_recall.unwrap() + recall_deficit * 0.5); - } + && recall < constraints.min_recall + { + // Increase ef_search to improve recall + let recall_deficit = constraints.min_recall - recall; + let boost = 1.0 + recall_deficit * 2.0; // Heuristic + params.ef_search = (params.ef_search as f32 * boost) as usize; + params.ef_search = params.ef_search.min(500); + + // Update estimated recall + params.estimated_recall = Some(params.estimated_recall.unwrap() + recall_deficit * 0.5); + } Ok(params) } diff --git a/src/backpressure.rs b/src/backpressure.rs index 4e07c4a..33243ad 100644 --- a/src/backpressure.rs +++ b/src/backpressure.rs @@ -33,7 +33,10 @@ //! ``` use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, RwLock, atomic::{AtomicU64, AtomicUsize, Ordering}}; +use std::sync::{ + Arc, RwLock, + atomic::{AtomicU64, AtomicUsize, Ordering}, +}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; @@ -218,7 +221,9 @@ impl BackpressureController { /// Check if should accept more records pub fn should_accept(&self) -> bool { - let Ok(state_guard) = self.state.read() else { return false; }; + let Ok(state_guard) = self.state.read() else { + return false; + }; let state = *state_guard; match state { @@ -227,11 +232,11 @@ impl BackpressureController { // Only accept if below low watermark let fill_ratio = self.buffer_fill_ratio(); fill_ratio < self.config.low_watermark - } + }, _ => { // Accept if rate limiter allows self.rate_limiter.try_acquire() - } + }, } } @@ -292,7 +297,9 @@ impl BackpressureController { /// Commit checkpoint pub fn commit_checkpoint(&self, offsets: HashMap) { - let Ok(mut checkpoint) = self.checkpoint.write() else { return; }; + let Ok(mut checkpoint) = self.checkpoint.write() else { + return; + }; checkpoint.offsets = offsets; checkpoint.timestamp = unix_timestamp(); checkpoint.pending_txns.clear(); @@ -327,9 +334,11 @@ impl BackpressureController { f64::INFINITY }; - let partition_lag = self.partition_offsets.read() + let partition_lag = self + .partition_offsets + .read() .map(|guard| guard.clone()) - .unwrap_or_default(); // Safe: returns empty HashMap on lock failure + .unwrap_or_default(); // Safe: returns empty HashMap on lock failure LagMetrics { pending_records: pending, @@ -342,13 +351,15 @@ impl BackpressureController { /// Get current state pub fn get_state(&self) -> BackpressureState { - let Ok(state) = self.state.read() else { return BackpressureState::Normal; }; + let Ok(state) = self.state.read() else { + return BackpressureState::Normal; + }; *state } /// Get statistics pub fn get_stats(&self) -> StreamStats { - let dlq_size = self.dlq.read().map(|guard| guard.len()).unwrap_or(0); // Safe: returns 0 on lock failure + let dlq_size = self.dlq.read().map(|guard| guard.len()).unwrap_or(0); // Safe: returns 0 on lock failure StreamStats { received: self.received.load(Ordering::Relaxed), processed: self.processed.load(Ordering::Relaxed), @@ -362,7 +373,9 @@ impl BackpressureController { /// Get DLQ records pub fn get_dlq_records(&self, limit: usize) -> Vec { - let Ok(dlq) = self.dlq.read() else { return Vec::new(); }; + let Ok(dlq) = self.dlq.read() else { + return Vec::new(); + }; dlq.iter().take(limit).cloned().collect() } @@ -386,12 +399,16 @@ impl BackpressureController { } fn calculate_processing_rate(&self) -> f64 { - let Ok(starts) = self.processing_starts.read() else { return 0.0; }; + let Ok(starts) = self.processing_starts.read() else { + return 0.0; + }; if starts.len() < 2 { return 0.0; } - let Some(oldest) = starts.front() else { return 0.0; }; + let Some(oldest) = starts.front() else { + return 0.0; + }; let elapsed = oldest.elapsed().as_secs_f64(); if elapsed > 0.0 { @@ -402,7 +419,9 @@ impl BackpressureController { } fn update_state(&self) { - let Ok(mut state) = self.state.write() else { return; }; + let Ok(mut state) = self.state.write() else { + return; + }; // Check circuit breaker if self.circuit_breaker.is_open() { @@ -421,27 +440,29 @@ impl BackpressureController { if fill_ratio >= self.config.high_watermark { *state = BackpressureState::Backpressure; } - } + }, BackpressureState::Backpressure => { if fill_ratio < self.config.low_watermark { *state = BackpressureState::Normal; } - } + }, BackpressureState::CircuitOpen => { if self.circuit_breaker.should_allow_request() { *state = BackpressureState::Recovering; } - } + }, BackpressureState::Recovering => { if !self.circuit_breaker.is_open() { *state = BackpressureState::Normal; } - } + }, } } fn add_to_dlq(&self, record: &StreamRecord, error: &str) { - let Ok(mut dlq) = self.dlq.write() else { return; }; + let Ok(mut dlq) = self.dlq.write() else { + return; + }; // Evict oldest if full while dlq.len() >= self.config.dlq_max_size { @@ -507,19 +528,20 @@ impl RateLimiter { if current == 0 { return false; } - if self.tokens.compare_exchange( - current, - current - 1, - Ordering::Relaxed, - Ordering::Relaxed, - ).is_ok() { + if self + .tokens + .compare_exchange(current, current - 1, Ordering::Relaxed, Ordering::Relaxed) + .is_ok() + { return true; } } } fn refill(&self) { - let Ok(mut last) = self.last_refill.write() else { return; }; + let Ok(mut last) = self.last_refill.write() else { + return; + }; let elapsed = last.elapsed(); if elapsed >= Duration::from_millis(100) { @@ -566,7 +588,9 @@ impl CircuitBreaker { fn record_success(&self) { self.success_count.fetch_add(1, Ordering::Relaxed); - let Ok(mut state) = self.state.write() else { return; }; + let Ok(mut state) = self.state.write() else { + return; + }; if *state == CircuitState::HalfOpen { // Reset on success in half-open *state = CircuitState::Closed; @@ -581,36 +605,44 @@ impl CircuitBreaker { } if count >= self.failure_threshold - && let Ok(mut state) = self.state.write() { - *state = CircuitState::Open; - } + && let Ok(mut state) = self.state.write() + { + *state = CircuitState::Open; + } } fn is_open(&self) -> bool { - let Ok(state_guard) = self.state.read() else { return true; }; // Fail-safe: treat as open + let Ok(state_guard) = self.state.read() else { + return true; + }; // Fail-safe: treat as open let state = *state_guard; drop(state_guard); match state { CircuitState::Open => { // Check if timeout has passed - let Ok(last_failure_guard) = self.last_failure.read() else { return true; }; + let Ok(last_failure_guard) = self.last_failure.read() else { + return true; + }; if let Some(last) = *last_failure_guard - && last.elapsed() > Duration::from_secs(self.timeout_seconds) { - drop(last_failure_guard); - if let Ok(mut state) = self.state.write() { - *state = CircuitState::HalfOpen; - } - return false; + && last.elapsed() > Duration::from_secs(self.timeout_seconds) + { + drop(last_failure_guard); + if let Ok(mut state) = self.state.write() { + *state = CircuitState::HalfOpen; } + return false; + } true - } + }, _ => false, } } fn should_allow_request(&self) -> bool { - let Ok(state) = self.state.read() else { return false; }; // Fail-safe: deny requests on lock failure + let Ok(state) = self.state.read() else { + return false; + }; // Fail-safe: deny requests on lock failure *state != CircuitState::Open } } @@ -640,7 +672,9 @@ impl BatchProcessor { self.controller.record_received(); - let Ok(mut buffer) = self.buffer.write() else { return false; }; + let Ok(mut buffer) = self.buffer.write() else { + return false; + }; buffer.push(record); true @@ -648,13 +682,17 @@ impl BatchProcessor { /// Check if batch is ready pub fn is_batch_ready(&self) -> bool { - let Ok(buffer) = self.buffer.read() else { return false; }; + let Ok(buffer) = self.buffer.read() else { + return false; + }; buffer.len() >= self.batch_size } /// Get batch for processing pub fn take_batch(&self) -> Vec> { - let Ok(mut buffer) = self.buffer.write() else { return Vec::new(); }; + let Ok(mut buffer) = self.buffer.write() else { + return Vec::new(); + }; let batch: Vec<_> = buffer.drain(..).collect(); batch } diff --git a/src/backup.rs b/src/backup.rs index 374de92..145602d 100644 --- a/src/backup.rs +++ b/src/backup.rs @@ -31,7 +31,7 @@ //! # } //! ``` -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::fs::{self, File}; @@ -166,11 +166,7 @@ impl BackupManager { // Include milliseconds and a random suffix for uniqueness use rand::Rng; let random: u32 = rand::rng().random_range(0..10000); - format!( - "backup_{}_{}", - now.format("%Y%m%d_%H%M%S_%3f"), - random - ) + format!("backup_{}_{}", now.format("%Y%m%d_%H%M%S_%3f"), random) } /// Create a backup of a VecStore directory @@ -194,11 +190,7 @@ impl BackupManager { let mut total_size = 0u64; // Copy data files - let data_files = vec![ - "manifest.json", - "records.bin", - "id_mapping.bin", - ]; + let data_files = vec!["manifest.json", "records.bin", "id_mapping.bin"]; for file in &data_files { let src = source.join(file); @@ -319,7 +311,11 @@ impl BackupManager { let entry = entry?; let path = entry.path(); - if path.is_file() && path.file_name().map(|n| n != "backup_info.json").unwrap_or(false) + if path.is_file() + && path + .file_name() + .map(|n| n != "backup_info.json") + .unwrap_or(false) { let mut file = File::open(&path)?; let mut buffer = [0u8; 8192]; @@ -387,11 +383,7 @@ impl BackupManager { /// /// The target directory will be created if it doesn't exist. /// If it exists and is not empty, this will fail unless `force` is true. - pub fn restore_backup( - &self, - backup_id: &str, - target_path: impl AsRef, - ) -> Result<()> { + pub fn restore_backup(&self, backup_id: &str, target_path: impl AsRef) -> Result<()> { self.restore_backup_with_options(backup_id, target_path, false) } @@ -621,7 +613,9 @@ mod walkdir { } pub fn file_type(&self) -> FileType { - FileType { is_file: self.is_file } + FileType { + is_file: self.is_file, + } } } @@ -652,7 +646,7 @@ mod walkdir { for e in entries.flatten() { self.stack.push(e.path()); } - } + }, Err(e) => return Some(Err(e)), } } else { @@ -689,12 +683,18 @@ mod tests { // Create source with some files fs::create_dir_all(&source_dir).unwrap(); - fs::write(source_dir.join("manifest.json"), r#"{"dimension":128,"record_count":100}"#).unwrap(); + fs::write( + source_dir.join("manifest.json"), + r#"{"dimension":128,"record_count":100}"#, + ) + .unwrap(); fs::write(source_dir.join("records.bin"), vec![0u8; 1000]).unwrap(); // Create backup let mgr = BackupManager::new(&backup_dir).unwrap(); - let backup_id = mgr.create_backup(&source_dir, BackupConfig::default()).unwrap(); + let backup_id = mgr + .create_backup(&source_dir, BackupConfig::default()) + .unwrap(); // Verify backup exists let backups = mgr.list_backups().unwrap(); @@ -719,10 +719,15 @@ mod tests { fs::write(source_dir.join("records.bin"), vec![1u8; 100]).unwrap(); let mgr = BackupManager::new(&backup_dir).unwrap(); - let backup_id = mgr.create_backup(&source_dir, BackupConfig { - verify: true, - ..Default::default() - }).unwrap(); + let backup_id = mgr + .create_backup( + &source_dir, + BackupConfig { + verify: true, + ..Default::default() + }, + ) + .unwrap(); // Verification should pass assert!(mgr.verify_backup(&backup_id).unwrap()); @@ -742,17 +747,31 @@ mod tests { // Create 3 backups with delay to ensure unique timestamps for i in 0..3 { let result = mgr.create_backup(&source_dir, BackupConfig::default()); - assert!(result.is_ok(), "Failed to create backup {}: {:?}", i, result.err()); + assert!( + result.is_ok(), + "Failed to create backup {}: {:?}", + i, + result.err() + ); std::thread::sleep(std::time::Duration::from_millis(100)); } let backups = mgr.list_backups().unwrap(); - assert_eq!(backups.len(), 3, "Expected 3 backups, got {}", backups.len()); + assert_eq!( + backups.len(), + 3, + "Expected 3 backups, got {}", + backups.len() + ); // Prune to keep only 1 let deleted = mgr.prune(1).unwrap(); assert_eq!(deleted.len(), 2, "Expected to delete 2 backups"); - assert_eq!(mgr.list_backups().unwrap().len(), 1, "Expected 1 backup remaining"); + assert_eq!( + mgr.list_backups().unwrap().len(), + 1, + "Expected 1 backup remaining" + ); } #[test] @@ -775,7 +794,13 @@ mod tests { BackupArchive::extract_archive(&archive_path, &extract_dir).unwrap(); // Verify - assert_eq!(fs::read_to_string(extract_dir.join("file1.txt")).unwrap(), "hello"); - assert_eq!(fs::read(extract_dir.join("file2.bin")).unwrap(), vec![1, 2, 3, 4]); + assert_eq!( + fs::read_to_string(extract_dir.join("file1.txt")).unwrap(), + "hello" + ); + assert_eq!( + fs::read(extract_dir.join("file2.bin")).unwrap(), + vec![1, 2, 3, 4] + ); } } diff --git a/src/benchmark.rs b/src/benchmark.rs index ed717ec..10410e1 100644 --- a/src/benchmark.rs +++ b/src/benchmark.rs @@ -814,7 +814,10 @@ impl Benchmarker { // Calculate scalability factor let max_threads = multi_thread_qps.keys().max().copied().unwrap_or(1); - let max_qps = multi_thread_qps.get(&max_threads).copied().unwrap_or(single_thread_qps); + let max_qps = multi_thread_qps + .get(&max_threads) + .copied() + .unwrap_or(single_thread_qps); let scalability_factor = if single_thread_qps > 0.0 { (max_qps / single_thread_qps) / max_threads as f64 } else { diff --git a/src/bin/vecstore.rs b/src/bin/vecstore.rs index d96e403..fdf12b2 100644 --- a/src/bin/vecstore.rs +++ b/src/bin/vecstore.rs @@ -3,9 +3,9 @@ use anyhow::{Context, Result}; use clap::{Parser, Subcommand, ValueEnum}; +use flate2::Compression; use flate2::read::GzDecoder; use flate2::write::GzEncoder; -use flate2::Compression; use std::collections::HashMap; use std::fs::{self, File}; use std::io::{BufReader, BufWriter}; @@ -315,7 +315,7 @@ fn main() -> Result<()> { } else { println!(" Dimension will be auto-detected from first insert"); } - } + }, Commands::Ingest { dir, id, vec, meta } => { let mut store = VecStore::open(&dir)?; @@ -335,7 +335,7 @@ fn main() -> Result<()> { store.save()?; println!("āœ“ Ingested record: {}", id); - } + }, Commands::IngestBatch { dir, jsonl } => { let mut store = VecStore::open(&dir)?; @@ -368,7 +368,7 @@ fn main() -> Result<()> { elapsed.as_secs_f64(), count as f64 / elapsed.as_secs_f64() ); - } + }, Commands::Query { dir, @@ -417,7 +417,7 @@ fn main() -> Result<()> { } } } - } + }, Commands::Stats { dir, detailed } => { let store = VecStore::open(&dir)?; @@ -435,7 +435,7 @@ fn main() -> Result<()> { (store.count() * store.dimension() * 4) / 1_048_576 ); } - } + }, Commands::Export { dir, @@ -452,43 +452,47 @@ fn main() -> Result<()> { let _lines: Vec = Vec::new(); // Implementation would iterate and export println!("āœ“ Exported to JSONL format"); - } + }, ExportFormat::Csv => { println!("āœ“ Exported to CSV format"); - } + }, ExportFormat::Parquet => { println!("āœ“ Exported to Parquet format"); - } + }, ExportFormat::Npy => { println!("āœ“ Exported to NumPy format"); - } + }, } - } + }, - Commands::Import { dir: _, input, format } => { + Commands::Import { + dir: _, + input, + format, + } => { println!("Importing from {:?} ({:?} format)...", input, format); match format { ImportFormat::Jsonl => { println!("āœ“ Imported from JSONL"); - } + }, ImportFormat::Csv => { println!("āœ“ Imported from CSV"); - } + }, ImportFormat::Npy => { println!("āœ“ Imported from NumPy"); - } + }, ImportFormat::Pinecone => { println!("āœ“ Imported from Pinecone export"); - } + }, ImportFormat::Weaviate => { println!("āœ“ Imported from Weaviate export"); - } + }, ImportFormat::Qdrant => { println!("āœ“ Imported from Qdrant export"); - } + }, } - } + }, Commands::Migrate { source, @@ -503,25 +507,25 @@ fn main() -> Result<()> { MigrationSource::Pinecone => { println!(" Connecting to Pinecone..."); println!("āœ“ Migration complete!"); - } + }, MigrationSource::Weaviate => { println!(" Connecting to Weaviate..."); println!("āœ“ Migration complete!"); - } + }, MigrationSource::Qdrant => { println!(" Connecting to Qdrant..."); println!("āœ“ Migration complete!"); - } + }, MigrationSource::ChromaDB => { println!(" Connecting to ChromaDB..."); println!("āœ“ Migration complete!"); - } + }, MigrationSource::Milvus => { println!(" Connecting to Milvus..."); println!("āœ“ Migration complete!"); - } + }, } - } + }, Commands::Backup { dir, @@ -545,7 +549,9 @@ fn main() -> Result<()> { // Determine output path with proper extension let output_path = if compress { - if !output.to_string_lossy().ends_with(".tar.gz") && !output.to_string_lossy().ends_with(".tgz") { + if !output.to_string_lossy().ends_with(".tar.gz") + && !output.to_string_lossy().ends_with(".tgz") + { output.with_extension("tar.gz") } else { output.clone() @@ -587,15 +593,20 @@ fn main() -> Result<()> { file_count += 1; let mut f = File::open(path)?; - archive.append_file(relative_path, &mut f) - .with_context(|| format!("Failed to add file to archive: {:?}", path))?; + archive + .append_file(relative_path, &mut f) + .with_context(|| { + format!("Failed to add file to archive: {:?}", path) + })?; } else if path.is_dir() && path != dir { - archive.append_dir(relative_path, path) - .with_context(|| format!("Failed to add directory to archive: {:?}", path))?; + archive.append_dir(relative_path, path).with_context(|| { + format!("Failed to add directory to archive: {:?}", path) + })?; } } - archive.finish() + archive + .finish() .context("Failed to finalize backup archive")?; } else { println!(" Compression: disabled"); @@ -611,15 +622,20 @@ fn main() -> Result<()> { file_count += 1; let mut f = File::open(path)?; - archive.append_file(relative_path, &mut f) - .with_context(|| format!("Failed to add file to archive: {:?}", path))?; + archive + .append_file(relative_path, &mut f) + .with_context(|| { + format!("Failed to add file to archive: {:?}", path) + })?; } else if path.is_dir() && path != dir { - archive.append_dir(relative_path, path) - .with_context(|| format!("Failed to add directory to archive: {:?}", path))?; + archive.append_dir(relative_path, path).with_context(|| { + format!("Failed to add directory to archive: {:?}", path) + })?; } } - archive.finish() + archive + .finish() .context("Failed to finalize backup archive")?; } @@ -633,14 +649,17 @@ fn main() -> Result<()> { println!("āœ“ Backup complete!"); println!(" Files: {}", file_count); - println!(" Original size: {:.2} MB", total_bytes as f64 / 1_048_576.0); + println!( + " Original size: {:.2} MB", + total_bytes as f64 / 1_048_576.0 + ); println!(" Backup size: {:.2} MB", output_size as f64 / 1_048_576.0); if compress { println!(" Compression ratio: {:.1}%", compression_ratio); } println!(" Time: {:.2}s", elapsed.as_secs_f64()); println!(" Output: {:?}", output_path); - } + }, Commands::Restore { backup, dest } => { println!("šŸ“„ Restoring from backup..."); @@ -662,8 +681,9 @@ fn main() -> Result<()> { if dest.exists() { println!(" āš ļø Destination exists, files may be overwritten"); } else { - fs::create_dir_all(&dest) - .with_context(|| format!("Failed to create destination directory: {:?}", dest))?; + fs::create_dir_all(&dest).with_context(|| { + format!("Failed to create destination directory: {:?}", dest) + })?; } let file = File::open(&backup) @@ -734,7 +754,7 @@ fn main() -> Result<()> { println!(" Total size: {:.2} MB", total_bytes as f64 / 1_048_576.0); println!(" Time: {:.2}s", elapsed.as_secs_f64()); println!(" Restored to: {:?}", dest); - } + }, Commands::Optimize { dir, rebuild } => { let mut store = VecStore::open(&dir)?; @@ -751,7 +771,7 @@ fn main() -> Result<()> { let elapsed = start.elapsed(); println!("āœ“ Optimization complete in {:.2}s", elapsed.as_secs_f64()); - } + }, Commands::Benchmark { dir, queries, k } => { let store = VecStore::open(&dir)?; @@ -793,7 +813,7 @@ fn main() -> Result<()> { " Throughput: {:.0} queries/sec", queries as f64 / total_time ); - } + }, Commands::Health { dir } => { let store = VecStore::open(&dir)?; @@ -808,7 +828,7 @@ fn main() -> Result<()> { if store.count() == 0 { println!("āš ļø Warning: Store is empty"); } - } + }, Commands::Collection(cmd) => match cmd { CollectionCommands::List { dir } => { @@ -820,7 +840,7 @@ fn main() -> Result<()> { for name in collections { println!(" - {}", name); } - } + }, CollectionCommands::Create { dir, @@ -831,14 +851,14 @@ fn main() -> Result<()> { db.create_collection(&name)?; println!("āœ“ Created collection '{}'", name); - } + }, CollectionCommands::Drop { dir, name } => { let mut db = VecDatabase::open(&dir)?; db.delete_collection(&name)?; println!("āœ“ Dropped collection '{}'", name); - } + }, CollectionCommands::Info { dir, name } => { let db = VecDatabase::open(&dir)?; @@ -853,7 +873,7 @@ fn main() -> Result<()> { } else { println!("Collection '{}' not found", name); } - } + }, }, Commands::Delete { dir, id, filter } => { @@ -871,7 +891,7 @@ fn main() -> Result<()> { eprintln!("Error: Must specify either --id or --filter"); std::process::exit(1); } - } + }, Commands::Compact { dir } => { let mut store = VecStore::open(&dir)?; @@ -890,7 +910,7 @@ fn main() -> Result<()> { " Removed: {} deleted vectors", before.saturating_sub(after) ); - } + }, } Ok(()) diff --git a/src/bulk_migration.rs b/src/bulk_migration.rs index 7116d12..f963a20 100644 --- a/src/bulk_migration.rs +++ b/src/bulk_migration.rs @@ -139,18 +139,19 @@ impl PineconeMigration { for (i, vector_data) in vectors.iter().enumerate().skip(start_offset) { if i % self.config.batch_size == 0 - && let Some(ref callback) = self.progress_callback { - callback(migrated, total); - } + && let Some(ref callback) = self.progress_callback + { + callback(migrated, total); + } match self.import_vector(vector_data, store) { Ok(size) => { migrated += 1; bytes += size; - } + }, Err(_) => { errors += 1; - } + }, } } @@ -238,9 +239,10 @@ impl QdrantMigration { for (i, line) in reader.lines().enumerate() { if let Some(offset) = self.config.resume_from - && i < offset { - continue; - } + && i < offset + { + continue; + } let line = line.map_err(|e| anyhow::anyhow!("Read error: {}", e))?; @@ -248,10 +250,10 @@ impl QdrantMigration { Ok(size) => { migrated += 1; bytes += size; - } + }, Err(_) => { errors += 1; - } + }, } } @@ -371,10 +373,10 @@ impl ChromaDBMigration { Ok(_) => { migrated += 1; bytes += (id.len() + embedding.len() * 4 + 100) as u64; - } + }, Err(_) => { errors += 1; - } + }, } } diff --git a/src/cache.rs b/src/cache.rs index 3943206..c642c5e 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -13,8 +13,8 @@ //! - Default capacity: 1000 entries //! - Configurable max size -use std::collections::hash_map::DefaultHasher; use std::collections::HashMap; +use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -77,12 +77,13 @@ impl QueryCache { if let Some(entry) = self.entries.get_mut(&hash) { // Check TTL if enabled if let Some(ttl) = self.ttl - && entry.inserted_at.elapsed() > ttl { - // Entry expired, remove it - self.entries.remove(&hash); - self.misses += 1; - return None; - } + && entry.inserted_at.elapsed() > ttl + { + // Entry expired, remove it + self.entries.remove(&hash); + self.misses += 1; + return None; + } // Update access time for LRU self.access_counter += 1; @@ -244,19 +245,25 @@ impl SharedQueryCache { /// Get a cached value pub fn get(&self, key: &[f32]) -> Option { - let Ok(mut guard) = self.inner.lock() else { return None; }; + let Ok(mut guard) = self.inner.lock() else { + return None; + }; guard.get(key) } /// Insert a value pub fn insert(&self, key: &[f32], value: V) { - let Ok(mut guard) = self.inner.lock() else { return; }; + let Ok(mut guard) = self.inner.lock() else { + return; + }; guard.insert(key, value); } /// Clear the cache pub fn clear(&self) { - let Ok(mut guard) = self.inner.lock() else { return; }; + let Ok(mut guard) = self.inner.lock() else { + return; + }; guard.clear(); } @@ -277,13 +284,17 @@ impl SharedQueryCache { /// Evict expired entries pub fn evict_expired(&self) -> usize { - let Ok(mut guard) = self.inner.lock() else { return 0; }; + let Ok(mut guard) = self.inner.lock() else { + return 0; + }; guard.evict_expired() } /// Get cache hit rate pub fn hit_rate(&self) -> f32 { - let Ok(guard) = self.inner.lock() else { return 0.0; }; + let Ok(guard) = self.inner.lock() else { + return 0.0; + }; guard.hit_rate() } @@ -327,8 +338,7 @@ pub mod early_termination { impl Ord for HeapEntry { fn cmp(&self, other: &Self) -> Ordering { // Max-heap based on distance (peek gives largest distance) - self.distance - .total_cmp(&other.distance) + self.distance.total_cmp(&other.distance) } } diff --git a/src/cdc.rs b/src/cdc.rs index b1ede48..5724b3c 100644 --- a/src/cdc.rs +++ b/src/cdc.rs @@ -275,45 +275,51 @@ impl ChangeFilter { if !ns.contains(event_ns) { return false; } - } + }, None => return false, } } // Check change type if let Some(ref cts) = self.change_types - && !cts.contains(&event.change_type) { - return false; - } + && !cts.contains(&event.change_type) + { + return false; + } // Check ID prefix if let Some(ref prefix) = self.id_prefix - && !event.vector_id.starts_with(prefix) { - return false; - } + && !event.vector_id.starts_with(prefix) + { + return false; + } // Check ID pattern (simple glob) if let Some(ref pattern) = self.id_pattern - && !glob_match(pattern, &event.vector_id) { - return false; - } + && !glob_match(pattern, &event.vector_id) + { + return false; + } // Check source if let Some(ref sources) = self.sources - && !sources.contains(&event.source) { - return false; - } + && !sources.contains(&event.source) + { + return false; + } // Check sequence range if let Some(from) = self.from_sequence - && event.sequence < from { - return false; - } + && event.sequence < from + { + return false; + } if let Some(to) = self.to_sequence - && event.sequence > to { - return false; - } + && event.sequence > to + { + return false; + } true } @@ -333,7 +339,7 @@ fn glob_match_impl(pattern: &[char], text: &[char]) -> bool { // Try matching zero or more characters glob_match_impl(&pattern[1..], text) || (!text.is_empty() && glob_match_impl(pattern, &text[1..])) - } + }, (Some('?'), Some(_)) => glob_match_impl(&pattern[1..], &text[1..]), (Some(p), Some(t)) if p == t => glob_match_impl(&pattern[1..], &text[1..]), _ => false, @@ -380,7 +386,7 @@ impl ChangeStream { } else { self.stats.filtered_out.fetch_add(1, Ordering::Relaxed); } - } + }, Err(_) => return None, } } @@ -406,7 +412,7 @@ impl ChangeStream { } else { self.stats.filtered_out.fetch_add(1, Ordering::Relaxed); } - } + }, Err(mpsc::RecvTimeoutError::Timeout) => return None, Err(mpsc::RecvTimeoutError::Disconnected) => return None, } @@ -427,7 +433,7 @@ impl ChangeStream { } else { self.stats.filtered_out.fetch_add(1, Ordering::Relaxed); } - } + }, Err(TryRecvError::Empty) => return None, Err(TryRecvError::Disconnected) => return None, } @@ -556,11 +562,11 @@ impl EventBroadcaster { Ok(_) => { self.inner.stats.published.fetch_add(1, Ordering::Relaxed); true - } + }, Err(_) => { // Subscriber dropped (disconnected) false - } + }, } }); @@ -653,7 +659,7 @@ impl AsyncChangeStream { } else { self.stats.filtered_out.fetch_add(1, Ordering::Relaxed); } - } + }, None => return None, } } @@ -753,7 +759,7 @@ impl AsyncEventBroadcaster { Ok(_) => { self.inner.stats.published.fetch_add(1, Ordering::Relaxed); true - } + }, Err(_) => false, }); @@ -825,12 +831,16 @@ mod tests { assert!(!ChangeFilter::new().namespace("users").matches(&event)); // Match by change type - assert!(ChangeFilter::new() - .change_type(ChangeType::Insert) - .matches(&event)); - assert!(!ChangeFilter::new() - .change_type(ChangeType::Delete) - .matches(&event)); + assert!( + ChangeFilter::new() + .change_type(ChangeType::Insert) + .matches(&event) + ); + assert!( + !ChangeFilter::new() + .change_type(ChangeType::Delete) + .matches(&event) + ); // Match by ID prefix assert!(ChangeFilter::new().id_prefix("products:").matches(&event)); @@ -891,8 +901,9 @@ mod tests { broadcaster.publish( VectorChangeEvent::insert("p1", vec![0.1]).with_namespace("products".to_string()), ); - broadcaster - .publish(VectorChangeEvent::insert("u1", vec![0.2]).with_namespace("users".to_string())); + broadcaster.publish( + VectorChangeEvent::insert("u1", vec![0.2]).with_namespace("users".to_string()), + ); broadcaster.publish( VectorChangeEvent::insert("p2", vec![0.3]).with_namespace("products".to_string()), ); diff --git a/src/change_streams.rs b/src/change_streams.rs index e5a29a5..62992df 100644 --- a/src/change_streams.rs +++ b/src/change_streams.rs @@ -35,7 +35,10 @@ //! ``` use std::collections::{HashMap, VecDeque}; -use std::sync::{Arc, RwLock, atomic::{AtomicU64, AtomicBool, Ordering}}; +use std::sync::{ + Arc, RwLock, + atomic::{AtomicBool, AtomicU64, Ordering}, +}; use std::time::{SystemTime, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; @@ -271,15 +274,17 @@ impl ChangeFilter { fn matches(&self, event: &ChangeEvent) -> bool { // Check operation if let Some(ops) = &self.operations - && !ops.contains(&event.operation) { - return false; - } + && !ops.contains(&event.operation) + { + return false; + } // Check document ID if let Some(ids) = &self.document_ids - && !ids.contains(&event.document_id) { - return false; - } + && !ids.contains(&event.document_id) + { + return false; + } // Check metadata if let Some(filter) = &self.metadata_filter { @@ -322,12 +327,16 @@ impl Subscription { /// Get next event (non-blocking) pub fn try_next(&self) -> Result, VecStoreError> { - let mut events = self.events.write() + let mut events = self + .events + .write() .map_err(|_| VecStoreError::LockError("events lock poisoned".into()))?; let event = events.pop_front(); if let Some(ref e) = event { - let mut token = self.last_token.write() + let mut token = self + .last_token + .write() .map_err(|_| VecStoreError::LockError("last_token lock poisoned".into()))?; *token = Some(e.resume_token.clone()); } @@ -337,13 +346,17 @@ impl Subscription { /// Get batch of events pub fn try_batch(&self, max: usize) -> Result, VecStoreError> { - let mut events = self.events.write() + let mut events = self + .events + .write() .map_err(|_| VecStoreError::LockError("events lock poisoned".into()))?; let count = max.min(events.len()); let batch: Vec = events.drain(..count).collect(); if let Some(last) = batch.last() { - let mut token = self.last_token.write() + let mut token = self + .last_token + .write() .map_err(|_| VecStoreError::LockError("last_token lock poisoned".into()))?; *token = Some(last.resume_token.clone()); } @@ -353,7 +366,9 @@ impl Subscription { /// Get last resume token pub fn resume_token(&self) -> Result, VecStoreError> { - let token = self.last_token.read() + let token = self + .last_token + .read() .map_err(|_| VecStoreError::LockError("last_token lock poisoned".into()))?; Ok(token.clone()) } @@ -365,14 +380,18 @@ impl Subscription { /// Check if events are pending pub fn has_events(&self) -> Result { - let events = self.events.read() + let events = self + .events + .read() .map_err(|_| VecStoreError::LockError("events lock poisoned".into()))?; Ok(!events.is_empty()) } /// Get pending event count pub fn pending_count(&self) -> Result { - let events = self.events.read() + let events = self + .events + .read() .map_err(|_| VecStoreError::LockError("events lock poisoned".into()))?; Ok(events.len()) } @@ -433,7 +452,9 @@ impl ChangeStream { last_token: Arc::new(RwLock::new(None)), }); - let mut subs = self.subscriptions.write() + let mut subs = self + .subscriptions + .write() .map_err(|_| VecStoreError::LockError("subscriptions lock poisoned".into()))?; subs.insert(id, subscription.clone()); @@ -442,7 +463,9 @@ impl ChangeStream { /// Unwatch (close subscription) pub fn unwatch(&self, subscription_id: u64) -> Result<(), VecStoreError> { - let mut subs = self.subscriptions.write() + let mut subs = self + .subscriptions + .write() .map_err(|_| VecStoreError::LockError("subscriptions lock poisoned".into()))?; if let Some(sub) = subs.remove(&subscription_id) { sub.close(); @@ -451,7 +474,13 @@ impl ChangeStream { } /// Emit an event - pub fn emit(&self, operation: Operation, document_id: &str, before: Option, after: Option) -> Result<(), VecStoreError> { + pub fn emit( + &self, + operation: Operation, + document_id: &str, + before: Option, + after: Option, + ) -> Result<(), VecStoreError> { let event_id = self.event_counter.fetch_add(1, Ordering::Relaxed); let timestamp = unix_timestamp(); @@ -475,7 +504,9 @@ impl ChangeStream { // Add to log { - let mut log = self.event_log.write() + let mut log = self + .event_log + .write() .map_err(|_| VecStoreError::LockError("event_log lock poisoned".into()))?; log.push_back(event.clone()); while log.len() > self.max_log_size { @@ -489,7 +520,12 @@ impl ChangeStream { } /// Emit insert event - pub fn emit_insert(&self, id: &str, vector: Vec, metadata: HashMap) -> Result<(), VecStoreError> { + pub fn emit_insert( + &self, + id: &str, + vector: Vec, + metadata: HashMap, + ) -> Result<(), VecStoreError> { let snapshot = DocumentSnapshot { id: id.to_string(), vector: Some(vector), @@ -501,23 +537,35 @@ impl ChangeStream { } /// Emit update event - pub fn emit_update(&self, id: &str, before: Option, after: DocumentSnapshot) -> Result<(), VecStoreError> { + pub fn emit_update( + &self, + id: &str, + before: Option, + after: DocumentSnapshot, + ) -> Result<(), VecStoreError> { self.emit(Operation::Update, id, before, Some(after)) } /// Emit delete event - pub fn emit_delete(&self, id: &str, before: Option) -> Result<(), VecStoreError> { + pub fn emit_delete( + &self, + id: &str, + before: Option, + ) -> Result<(), VecStoreError> { self.emit(Operation::Delete, id, before, None) } fn notify_subscriptions(&self, event: &ChangeEvent) -> Result<(), VecStoreError> { - let subs = self.subscriptions.read() + let subs = self + .subscriptions + .read() .map_err(|_| VecStoreError::LockError("subscriptions lock poisoned".into()))?; for sub in subs.values() { if sub.is_active() && sub.filter.matches(event) { - let mut events = sub.events.write() - .map_err(|_| VecStoreError::LockError("subscription events lock poisoned".into()))?; + let mut events = sub.events.write().map_err(|_| { + VecStoreError::LockError("subscription events lock poisoned".into()) + })?; // Apply backpressure if events.len() < sub.filter.batch_size * 10 { @@ -529,38 +577,49 @@ impl ChangeStream { Ok(()) } - fn replay_from(&self, token: &ResumeToken, filter: &ChangeFilter) -> Result, VecStoreError> { - let log = self.event_log.read() + fn replay_from( + &self, + token: &ResumeToken, + filter: &ChangeFilter, + ) -> Result, VecStoreError> { + let log = self + .event_log + .read() .map_err(|_| VecStoreError::LockError("event_log lock poisoned".into()))?; // Find starting position - let start_id: u64 = token.token + let start_id: u64 = token + .token .split(':') .nth(1) .and_then(|s| s.parse().ok()) .unwrap_or(0); - Ok(log.iter() + Ok(log + .iter() .filter(|e| e.id > start_id && filter.matches(e)) .cloned() .collect()) } fn replay_all(&self, filter: &ChangeFilter) -> Result, VecStoreError> { - let log = self.event_log.read() + let log = self + .event_log + .read() .map_err(|_| VecStoreError::LockError("event_log lock poisoned".into()))?; - Ok(log.iter() - .filter(|e| filter.matches(e)) - .cloned() - .collect()) + Ok(log.iter().filter(|e| filter.matches(e)).cloned().collect()) } /// Get statistics pub fn stats(&self) -> Result { - let subs = self.subscriptions.read() + let subs = self + .subscriptions + .read() .map_err(|_| VecStoreError::LockError("subscriptions lock poisoned".into()))?; - let log = self.event_log.read() + let log = self + .event_log + .read() .map_err(|_| VecStoreError::LockError("event_log lock poisoned".into()))?; Ok(ChangeStreamStats { @@ -596,7 +655,9 @@ impl ChangeStreamManager { /// Get or create stream for collection pub fn stream(&self, collection: &str) -> Result, VecStoreError> { - let streams = self.streams.read() + let streams = self + .streams + .read() .map_err(|_| VecStoreError::LockError("streams lock poisoned".into()))?; if let Some(stream) = streams.get(collection) { @@ -605,7 +666,9 @@ impl ChangeStreamManager { drop(streams); - let mut streams = self.streams.write() + let mut streams = self + .streams + .write() .map_err(|_| VecStoreError::LockError("streams lock poisoned".into()))?; Ok(streams .entry(collection.to_string()) @@ -615,20 +678,20 @@ impl ChangeStreamManager { /// Watch all collections pub fn watch_all(&self, filter: ChangeFilter) -> Result>, VecStoreError> { - let streams = self.streams.read() + let streams = self + .streams + .read() .map_err(|_| VecStoreError::LockError("streams lock poisoned".into()))?; - streams.values() - .map(|s| s.watch(filter.clone())) - .collect() + streams.values().map(|s| s.watch(filter.clone())).collect() } /// Get all statistics pub fn stats(&self) -> Result, VecStoreError> { - let streams = self.streams.read() + let streams = self + .streams + .read() .map_err(|_| VecStoreError::LockError("streams lock poisoned".into()))?; - streams.values() - .map(|s| s.stats()) - .collect() + streams.values().map(|s| s.stats()).collect() } } @@ -657,8 +720,12 @@ mod tests { let sub = stream.watch(ChangeFilter::all()).unwrap(); // Emit events - stream.emit_insert("doc1", vec![0.1, 0.2], HashMap::new()).unwrap(); - stream.emit_insert("doc2", vec![0.3, 0.4], HashMap::new()).unwrap(); + stream + .emit_insert("doc1", vec![0.1, 0.2], HashMap::new()) + .unwrap(); + stream + .emit_insert("doc2", vec![0.3, 0.4], HashMap::new()) + .unwrap(); // Check events let event1 = sub.try_next().unwrap().unwrap(); @@ -677,9 +744,13 @@ mod tests { let sub = stream.watch(ChangeFilter::inserts()).unwrap(); // Emit mixed events - stream.emit_insert("doc1", vec![0.1], HashMap::new()).unwrap(); + stream + .emit_insert("doc1", vec![0.1], HashMap::new()) + .unwrap(); stream.emit_delete("doc2", None).unwrap(); - stream.emit_insert("doc3", vec![0.2], HashMap::new()).unwrap(); + stream + .emit_insert("doc3", vec![0.2], HashMap::new()) + .unwrap(); // Should only get inserts let events = sub.try_batch(10).unwrap(); @@ -692,9 +763,15 @@ mod tests { let stream = ChangeStream::new("test"); // Emit some events - stream.emit_insert("doc1", vec![0.1], HashMap::new()).unwrap(); - stream.emit_insert("doc2", vec![0.2], HashMap::new()).unwrap(); - stream.emit_insert("doc3", vec![0.3], HashMap::new()).unwrap(); + stream + .emit_insert("doc1", vec![0.1], HashMap::new()) + .unwrap(); + stream + .emit_insert("doc2", vec![0.2], HashMap::new()) + .unwrap(); + stream + .emit_insert("doc3", vec![0.3], HashMap::new()) + .unwrap(); // First subscriber gets all events let sub1 = stream.watch(ChangeFilter::all()).unwrap(); @@ -704,7 +781,9 @@ mod tests { let token = events[1].resume_token.clone(); // Second subscriber resumes after doc2 - let sub2 = stream.watch(ChangeFilter::all().resume_after(token)).unwrap(); + let sub2 = stream + .watch(ChangeFilter::all().resume_after(token)) + .unwrap(); // Should get doc3 only let events = sub2.try_batch(10).unwrap(); @@ -717,13 +796,26 @@ mod tests { let stream = ChangeStream::new("test"); // Subscribe to specific documents - let sub = stream.watch(ChangeFilter::for_documents(vec!["doc1".to_string(), "doc3".to_string()])).unwrap(); + let sub = stream + .watch(ChangeFilter::for_documents(vec![ + "doc1".to_string(), + "doc3".to_string(), + ])) + .unwrap(); // Emit events - stream.emit_insert("doc1", vec![0.1], HashMap::new()).unwrap(); - stream.emit_insert("doc2", vec![0.2], HashMap::new()).unwrap(); - stream.emit_insert("doc3", vec![0.3], HashMap::new()).unwrap(); - stream.emit_insert("doc4", vec![0.4], HashMap::new()).unwrap(); + stream + .emit_insert("doc1", vec![0.1], HashMap::new()) + .unwrap(); + stream + .emit_insert("doc2", vec![0.2], HashMap::new()) + .unwrap(); + stream + .emit_insert("doc3", vec![0.3], HashMap::new()) + .unwrap(); + stream + .emit_insert("doc4", vec![0.4], HashMap::new()) + .unwrap(); // Should only get doc1 and doc3 let events = sub.try_batch(10).unwrap(); @@ -739,8 +831,12 @@ mod tests { let stream1 = manager.stream("collection1").unwrap(); let stream2 = manager.stream("collection2").unwrap(); - stream1.emit_insert("doc1", vec![0.1], HashMap::new()).unwrap(); - stream2.emit_insert("doc2", vec![0.2], HashMap::new()).unwrap(); + stream1 + .emit_insert("doc1", vec![0.1], HashMap::new()) + .unwrap(); + stream2 + .emit_insert("doc2", vec![0.2], HashMap::new()) + .unwrap(); let stats = manager.stats().unwrap(); assert_eq!(stats.len(), 2); @@ -765,7 +861,9 @@ mod tests { timestamp: unix_timestamp(), }; - stream.emit_update("doc1", Some(before.clone()), after).unwrap(); + stream + .emit_update("doc1", Some(before.clone()), after) + .unwrap(); let event = sub.try_next().unwrap().unwrap(); assert_eq!(event.operation, Operation::Update); diff --git a/src/clustering.rs b/src/clustering.rs index 951a912..ca287c7 100644 --- a/src/clustering.rs +++ b/src/clustering.rs @@ -603,14 +603,14 @@ impl HierarchicalClustering { .iter() .flat_map(|&i| cluster2.iter().map(move |&j| distances[i][j])) .fold(f32::INFINITY, f32::min) - } + }, LinkageMethod::Complete => { // Maximum distance cluster1 .iter() .flat_map(|&i| cluster2.iter().map(move |&j| distances[i][j])) .fold(0.0, f32::max) - } + }, LinkageMethod::Average => { // Average distance let sum: f32 = cluster1 @@ -618,7 +618,7 @@ impl HierarchicalClustering { .flat_map(|&i| cluster2.iter().map(move |&j| distances[i][j])) .sum(); sum / (cluster1.len() * cluster2.len()) as f32 - } + }, } } diff --git a/src/collection.rs b/src/collection.rs index e5fce27..b94467e 100644 --- a/src/collection.rs +++ b/src/collection.rs @@ -137,8 +137,9 @@ impl VecDatabase { // Create namespace with quotas { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .create_namespace(namespace_id.clone(), description, Some(config.quotas)) .map_err(|e| crate::error::VecStoreError::Other(e.to_string()))?; @@ -168,8 +169,9 @@ impl VecDatabase { pub fn get_collection(&self, name: &str) -> Result> { let namespace_id: NamespaceId = name.to_string(); - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; // Check if namespace exists match manager.get_namespace(&namespace_id) { Ok(_) => Ok(Some(Collection { @@ -196,8 +198,9 @@ impl VecDatabase { /// # Ok::<(), anyhow::Error>(()) /// ``` pub fn list_collections(&self) -> Result> { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; let namespaces = manager.list_namespaces(); Ok(namespaces.into_iter().map(|ns| ns.id).collect()) } @@ -215,8 +218,9 @@ impl VecDatabase { /// ``` pub fn delete_collection(&mut self, name: &str) -> Result<()> { let namespace_id: NamespaceId = name.to_string(); - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .delete_namespace(&namespace_id) .map_err(|e| crate::error::VecStoreError::Other(e.to_string())) @@ -315,8 +319,9 @@ impl Collection { /// /// Returns information about vector count, storage usage, etc. pub fn stats(&self) -> Result { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .get_stats(&self.namespace_id) .map_err(|e| crate::error::VecStoreError::Other(e.to_string())) @@ -324,8 +329,9 @@ impl Collection { /// Get namespace metadata (includes resource usage) pub fn namespace(&self) -> Result { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .get_namespace(&self.namespace_id) .map_err(|e| crate::error::VecStoreError::Other(e.to_string())) @@ -350,8 +356,9 @@ impl Collection { /// # Ok::<(), anyhow::Error>(()) /// ``` pub fn upsert(&mut self, id: String, vector: Vec, metadata: Metadata) -> Result<()> { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .upsert(&self.namespace_id, id, vector, metadata) .map_err(|e| crate::error::VecStoreError::Other(e.to_string())) @@ -379,8 +386,9 @@ impl Collection { /// # Ok::<(), anyhow::Error>(()) /// ``` pub fn query(&self, query: Query) -> Result> { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .query(&self.namespace_id, query) .map_err(|e| crate::error::VecStoreError::Other(e.to_string())) @@ -388,8 +396,9 @@ impl Collection { /// Delete a vector by ID pub fn delete(&mut self, id: &str) -> Result<()> { - let manager = self.manager.read() - .map_err(|_| crate::error::VecStoreError::LockError("manager read lock poisoned".into()))?; + let manager = self.manager.read().map_err(|_| { + crate::error::VecStoreError::LockError("manager read lock poisoned".into()) + })?; manager .remove(&self.namespace_id, id) .map_err(|e| crate::error::VecStoreError::Other(e.to_string())) diff --git a/src/columnar.rs b/src/columnar.rs index 38e1904..7011bb3 100644 --- a/src/columnar.rs +++ b/src/columnar.rs @@ -38,15 +38,14 @@ //! let stats = store.column_stats(0)?; //! ``` +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::path::{Path, PathBuf}; -use serde::{Deserialize, Serialize}; use crate::error::VecStoreError; /// Compression type for columnar data -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] pub enum CompressionType { /// No compression #[default] @@ -63,7 +62,6 @@ pub enum CompressionType { Dictionary, } - /// Configuration for columnar storage #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ColumnarConfig { @@ -170,27 +168,31 @@ impl ColumnarChunk { pub fn get_vector(&self, local_idx: usize) -> Result, VecStoreError> { if local_idx >= self.count { return Err(VecStoreError::NotFound(format!( - "Vector index {} out of range", local_idx + "Vector index {} out of range", + local_idx ))); } - let vector: Vec = self.columns - .iter() - .map(|col| col[local_idx]) - .collect(); + let vector: Vec = self.columns.iter().map(|col| col[local_idx]).collect(); Ok(vector) } /// Get specific dimensions for a vector - pub fn get_dimensions(&self, local_idx: usize, dims: &[usize]) -> Result, VecStoreError> { + pub fn get_dimensions( + &self, + local_idx: usize, + dims: &[usize], + ) -> Result, VecStoreError> { if local_idx >= self.count { return Err(VecStoreError::NotFound(format!( - "Vector index {} out of range", local_idx + "Vector index {} out of range", + local_idx ))); } - let values: Vec = dims.iter() + let values: Vec = dims + .iter() .filter_map(|&d| { if d < self.columns.len() { Some(self.columns[d][local_idx]) @@ -205,7 +207,8 @@ impl ColumnarChunk { /// Compute statistics for all columns pub fn compute_stats(&mut self) { - let stats: Vec = self.columns + let stats: Vec = self + .columns .iter() .enumerate() .map(|(dim, col)| Self::compute_column_stats(dim, col)) @@ -233,12 +236,14 @@ impl ColumnarChunk { let sum: f64 = column.iter().map(|&x| x as f64).sum(); let mean = sum / column.len() as f64; - let variance: f64 = column.iter() + let variance: f64 = column + .iter() .map(|&x| { let diff = x as f64 - mean; diff * diff }) - .sum::() / column.len() as f64; + .sum::() + / column.len() as f64; let non_zero_count = column.iter().filter(|&&x| x != 0.0).count(); @@ -260,38 +265,39 @@ impl ColumnarChunk { CompressionType::None => { self.compressed_columns = None; Ok(()) - } + }, CompressionType::Delta => { - let compressed: Vec> = self.columns + let compressed: Vec> = self + .columns .iter() .map(|col| Self::delta_encode(col)) .collect(); self.compressed_columns = Some(compressed); Ok(()) - } + }, CompressionType::Rle => { - let compressed: Vec> = self.columns + let compressed: Vec> = self + .columns .iter() .map(|col| Self::rle_encode(col)) .collect(); self.compressed_columns = Some(compressed); Ok(()) - } + }, _ => { // LZ4, Zstd would require external crates // For now, store raw bytes - let compressed: Vec> = self.columns + let compressed: Vec> = self + .columns .iter() .map(|col| { - let bytes: Vec = col.iter() - .flat_map(|&f| f.to_le_bytes()) - .collect(); + let bytes: Vec = col.iter().flat_map(|&f| f.to_le_bytes()).collect(); bytes }) .collect(); self.compressed_columns = Some(compressed); Ok(()) - } + }, } } @@ -382,7 +388,11 @@ impl ColumnarStore { } /// Open or create a columnar store at path - pub fn open(path: impl AsRef, dimensions: usize, config: ColumnarConfig) -> Result { + pub fn open( + path: impl AsRef, + dimensions: usize, + config: ColumnarConfig, + ) -> Result { let path = path.as_ref().to_path_buf(); if path.exists() { @@ -398,8 +408,7 @@ impl ColumnarStore { fn load(path: &Path, config: ColumnarConfig) -> Result { let meta_path = path.join("columnar_meta.json"); - let file = std::fs::File::open(&meta_path) - .map_err(VecStoreError::Io)?; + let file = std::fs::File::open(&meta_path).map_err(VecStoreError::Io)?; let meta: ColumnarMeta = serde_json::from_reader(file) .map_err(|e| VecStoreError::Serialization(e.to_string()))?; @@ -426,8 +435,7 @@ impl ColumnarStore { } fn load_chunk(&self, path: &Path) -> Result { - let file = std::fs::File::open(path) - .map_err(VecStoreError::Io)?; + let file = std::fs::File::open(path).map_err(VecStoreError::Io)?; let chunk: ColumnarChunk = bincode::deserialize_from(file) .map_err(|e| VecStoreError::Serialization(e.to_string()))?; @@ -487,7 +495,9 @@ impl ColumnarStore { /// Get a vector by ID pub fn get(&self, id: &str) -> Result, VecStoreError> { - let (chunk_idx, local_idx) = self.id_map.get(id) + let (chunk_idx, local_idx) = self + .id_map + .get(id) .ok_or_else(|| VecStoreError::NotFound(format!("Vector {} not found", id)))?; self.chunks[*chunk_idx].get_vector(*local_idx) @@ -495,7 +505,9 @@ impl ColumnarStore { /// Read specific dimensions for a vector pub fn read_dimensions(&self, id: &str, dims: &[usize]) -> Result, VecStoreError> { - let (chunk_idx, local_idx) = self.id_map.get(id) + let (chunk_idx, local_idx) = self + .id_map + .get(id) .ok_or_else(|| VecStoreError::NotFound(format!("Vector {} not found", id)))?; self.chunks[*chunk_idx].get_dimensions(*local_idx, dims) @@ -505,7 +517,8 @@ impl ColumnarStore { pub fn read_column(&self, dimension: usize) -> Result, VecStoreError> { if dimension >= self.dimensions { return Err(VecStoreError::InvalidInput(format!( - "Dimension {} out of range (max {})", dimension, self.dimensions + "Dimension {} out of range (max {})", + dimension, self.dimensions ))); } @@ -522,7 +535,8 @@ impl ColumnarStore { pub fn column_stats(&self, dimension: usize) -> Result { if dimension >= self.dimensions { return Err(VecStoreError::InvalidInput(format!( - "Dimension {} out of range", dimension + "Dimension {} out of range", + dimension ))); } @@ -534,9 +548,8 @@ impl ColumnarStore { /// Get global statistics (across all dimensions) pub fn global_stats(&mut self) -> Result<&[ColumnStats], VecStoreError> { if self.global_stats.is_none() { - let stats: Result, _> = (0..self.dimensions) - .map(|d| self.column_stats(d)) - .collect(); + let stats: Result, _> = + (0..self.dimensions).map(|d| self.column_stats(d)).collect(); self.global_stats = Some(stats?); } @@ -545,7 +558,9 @@ impl ColumnarStore { /// Save to disk pub fn save(&self) -> Result<(), VecStoreError> { - let path = self.path.as_ref() + let path = self + .path + .as_ref() .ok_or_else(|| VecStoreError::InvalidInput("No path set for columnar store".into()))?; std::fs::create_dir_all(path).map_err(VecStoreError::Io)?; @@ -559,16 +574,14 @@ impl ColumnarStore { }; let meta_path = path.join("columnar_meta.json"); - let file = std::fs::File::create(&meta_path) - .map_err(VecStoreError::Io)?; + let file = std::fs::File::create(&meta_path).map_err(VecStoreError::Io)?; serde_json::to_writer(file, &meta) .map_err(|e| VecStoreError::Serialization(e.to_string()))?; // Save chunks for (i, chunk) in self.chunks.iter().enumerate() { let chunk_path = path.join(format!("chunk_{}.bin", i)); - let file = std::fs::File::create(&chunk_path) - .map_err(VecStoreError::Io)?; + let file = std::fs::File::create(&chunk_path).map_err(VecStoreError::Io)?; bincode::serialize_into(file, chunk) .map_err(|e| VecStoreError::Serialization(e.to_string()))?; } @@ -623,7 +636,8 @@ impl ColumnarStore { } new_chunks.push(current_chunk); let new_idx = new_chunks.len(); - current_chunk = ColumnarChunk::new(new_idx, new_count, self.dimensions, self.config.chunk_size); + current_chunk = + ColumnarChunk::new(new_idx, new_count, self.dimensions, self.config.chunk_size); } let local_idx = current_chunk.count; @@ -650,7 +664,10 @@ impl ColumnarStore { } /// Batch distance calculation (leverages columnar layout) - pub fn batch_euclidean_distance(&self, query: &[f32]) -> Result, VecStoreError> { + pub fn batch_euclidean_distance( + &self, + query: &[f32], + ) -> Result, VecStoreError> { if query.len() != self.dimensions { return Err(VecStoreError::DimensionMismatch { expected: self.dimensions, @@ -693,7 +710,8 @@ impl ColumnarStore { ) -> Result, VecStoreError> { if dimension >= self.dimensions { return Err(VecStoreError::InvalidInput(format!( - "Dimension {} out of range", dimension + "Dimension {} out of range", + dimension ))); } @@ -730,7 +748,9 @@ impl ColumnarStore { if let Some(ref compressed) = chunk.compressed_columns { size += compressed.iter().map(|c| c.len()).sum::(); } else { - size += chunk.columns.iter() + size += chunk + .columns + .iter() .map(|c| c.len() * size_of::()) .sum::(); } @@ -772,9 +792,10 @@ impl<'a> Iterator for ColumnarIterator<'a> { fn next(&mut self) -> Option { if let Some((id, (chunk_idx, local_idx))) = self.id_iter.next() - && let Ok(vec) = self.store.chunks[*chunk_idx].get_vector(*local_idx) { - return Some((id.clone(), vec)); - } + && let Ok(vec) = self.store.chunks[*chunk_idx].get_vector(*local_idx) + { + return Some((id.clone(), vec)); + } None } } diff --git a/src/compression.rs b/src/compression.rs index 46afb78..63efb93 100644 --- a/src/compression.rs +++ b/src/compression.rs @@ -41,12 +41,11 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; /// Compression level -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[derive(Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum CompressionLevel { /// No compression (fastest) None, @@ -62,7 +61,6 @@ pub enum CompressionLevel { Max, } - /// Compression configuration #[derive(Debug, Clone)] pub struct CompressionConfig { @@ -236,7 +234,7 @@ impl CompressionConfig { { Ok(data.to_vec()) } - } + }, CompressionLevel::Balanced => { #[cfg(feature = "compression")] { @@ -253,7 +251,7 @@ impl CompressionConfig { { Ok(data.to_vec()) } - } + }, CompressionLevel::Max => { #[cfg(feature = "compression")] { @@ -270,7 +268,7 @@ impl CompressionConfig { { Ok(data.to_vec()) } - } + }, } } @@ -289,17 +287,17 @@ impl CompressionConfig { let decompressed = lz4_flex::decompress_size_prepended(&data[1..]) .map_err(|e| anyhow!("LZ4 decompression failed: {}", e))?; Ok(decompressed) - } + }, Some(0x5A) => { // ZSTD format let decompressed = zstd::decode_all(std::io::Cursor::new(&data[1..])) .map_err(|e| anyhow!("ZSTD decompression failed: {}", e))?; Ok(decompressed) - } + }, _ => { // Uncompressed data (no magic byte or unknown format) Ok(data.to_vec()) - } + }, } } #[cfg(not(feature = "compression"))] @@ -329,7 +327,7 @@ impl CompressionConfig { } Ok(output) - } + }, 16 => { // 16-bit quantization let min = values.iter().copied().fold(f32::INFINITY, f32::min); @@ -346,7 +344,7 @@ impl CompressionConfig { } Ok(output) - } + }, 32 => { // No compression, just convert to bytes let mut output = Vec::with_capacity(values.len() * 4); @@ -354,7 +352,7 @@ impl CompressionConfig { output.extend_from_slice(&v.to_le_bytes()); } Ok(output) - } + }, _ => Err(anyhow!("Unsupported precision bits: {}", precision_bits)), } } @@ -385,7 +383,7 @@ impl CompressionConfig { } Ok(values) - } + }, 16 => { if data.len() < 8 { return Err(anyhow!("Insufficient data for float decompression")); @@ -406,7 +404,7 @@ impl CompressionConfig { } Ok(values) - } + }, 32 => { let mut values = Vec::with_capacity(count); for i in 0..count { @@ -417,7 +415,7 @@ impl CompressionConfig { values.push(f32::from_le_bytes(data[offset..offset + 4].try_into()?)); } Ok(values) - } + }, _ => Err(anyhow!("Unsupported precision bits: {}", precision_bits)), } } diff --git a/src/cost_optimizer.rs b/src/cost_optimizer.rs index b3d9f9c..eb1c8c3 100644 --- a/src/cost_optimizer.rs +++ b/src/cost_optimizer.rs @@ -2,8 +2,8 @@ // Automatic scaling recommendations, resource right-sizing, and cost forecasting use std::collections::VecDeque; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use serde::{Deserialize, Serialize}; @@ -179,13 +179,36 @@ pub enum EffortLevel { /// Optimization action #[derive(Debug, Clone, Serialize)] pub enum OptimizationAction { - ScaleDown { resource: String, current: f64, target: f64 }, - ScaleUp { resource: String, current: f64, target: f64 }, - EnableFeature { feature: String }, - DisableFeature { feature: String }, - ChangeConfig { key: String, current: String, target: String }, - MigrateStorage { from_tier: String, to_tier: String, size_gb: f64 }, - AdjustReplicas { current: u32, target: u32 }, + ScaleDown { + resource: String, + current: f64, + target: f64, + }, + ScaleUp { + resource: String, + current: f64, + target: f64, + }, + EnableFeature { + feature: String, + }, + DisableFeature { + feature: String, + }, + ChangeConfig { + key: String, + current: String, + target: String, + }, + MigrateStorage { + from_tier: String, + to_tier: String, + size_gb: f64, + }, + AdjustReplicas { + current: u32, + target: u32, + }, } /// Recommendation status @@ -280,7 +303,9 @@ impl CostOptimizer { latency_ms: usage.latency_ms, }; - let Ok(mut history) = self.resource_usage.write() else { return; }; + let Ok(mut history) = self.resource_usage.write() else { + return; + }; if history.samples.len() >= history.max_samples { history.samples.pop_front(); } @@ -288,8 +313,12 @@ impl CostOptimizer { // Calculate and record cost let cost = self.calculate_cost(&usage); - self.stats.total_cost.fetch_add((cost * 100.0) as u64, Ordering::Relaxed); - self.stats.total_queries.fetch_add(usage.queries, Ordering::Relaxed); + self.stats + .total_cost + .fetch_add((cost * 100.0) as u64, Ordering::Relaxed); + self.stats + .total_queries + .fetch_add(usage.queries, Ordering::Relaxed); } /// Calculate cost for resource usage @@ -309,7 +338,9 @@ impl CostOptimizer { /// Analyze current costs and generate recommendations pub fn analyze(&self) -> Result { - let history = self.resource_usage.read() + let history = self + .resource_usage + .read() .map_err(|_| VecStoreError::LockError("resource_usage lock poisoned".into()))?; if history.samples.is_empty() { @@ -344,13 +375,14 @@ impl CostOptimizer { // Generate recommendations let recommendations = self.generate_recommendations(&avg_usage, &history.samples); - let potential_savings: f64 = recommendations.iter() - .map(|r| r.estimated_savings) - .sum(); + let potential_savings: f64 = recommendations.iter().map(|r| r.estimated_savings).sum(); // Store recommendations - *self.recommendations.write() - .map_err(|_| VecStoreError::LockError("recommendations lock poisoned".into()))? = recommendations.clone(); + *self + .recommendations + .write() + .map_err(|_| VecStoreError::LockError("recommendations lock poisoned".into()))? = + recommendations.clone(); Ok(CostAnalysis { current_costs, @@ -397,7 +429,8 @@ impl CostOptimizer { return 0.0; } - let total_cost: f64 = samples.iter() + let total_cost: f64 = samples + .iter() .map(|s| { let usage = ResourceUsage { cpu_cores: s.cpu_cores, @@ -430,13 +463,17 @@ impl CostOptimizer { let first_half: Vec<_> = samples.iter().take(half).collect(); let second_half: Vec<_> = samples.iter().skip(half).collect(); - let first_avg: f64 = first_half.iter() + let first_avg: f64 = first_half + .iter() .map(|s| self.calculate_cost(&ResourceUsage::from_sample(s))) - .sum::() / first_half.len() as f64; + .sum::() + / first_half.len() as f64; - let second_avg: f64 = second_half.iter() + let second_avg: f64 = second_half + .iter() .map(|s| self.calculate_cost(&ResourceUsage::from_sample(s))) - .sum::() / second_half.len() as f64; + .sum::() + / second_half.len() as f64; let change_percent = if first_avg > 0.0 { ((second_avg - first_avg) / first_avg) * 100.0 @@ -469,7 +506,8 @@ impl CostOptimizer { } let total_queries: u64 = samples.iter().map(|s| s.queries).sum(); - let total_cost: f64 = samples.iter() + let total_cost: f64 = samples + .iter() .map(|s| self.calculate_cost(&ResourceUsage::from_sample(s))) .sum(); @@ -500,7 +538,8 @@ impl CostOptimizer { if cpu_utilization < 0.3 { let target_cores = (avg_usage.cpu_cores * 0.5).max(1.0); let savings = (avg_usage.cpu_cores - target_cores) - * self.config.pricing.cpu_per_core_hour * 720.0; + * self.config.pricing.cpu_per_core_hour + * 720.0; recommendations.push(CostRecommendation { id: format!("cpu-rightsize-{}", now), @@ -508,7 +547,9 @@ impl CostOptimizer { title: "CPU Over-provisioned".to_string(), description: format!( "CPU utilization is only {:.0}%. Consider reducing cores from {:.1} to {:.1}.", - cpu_utilization * 100.0, avg_usage.cpu_cores, target_cores + cpu_utilization * 100.0, + avg_usage.cpu_cores, + target_cores ), estimated_savings: savings, savings_percent: savings / self.project_monthly_cost(samples) * 100.0, @@ -529,7 +570,8 @@ impl CostOptimizer { if memory_utilization < 0.4 { let target_memory = (avg_usage.memory_gb * 0.6).max(1.0); let savings = (avg_usage.memory_gb - target_memory) - * self.config.pricing.memory_per_gb_hour * 720.0; + * self.config.pricing.memory_per_gb_hour + * 720.0; recommendations.push(CostRecommendation { id: format!("mem-rightsize-{}", now), @@ -537,7 +579,9 @@ impl CostOptimizer { title: "Memory Over-provisioned".to_string(), description: format!( "Memory utilization is only {:.0}%. Consider reducing from {:.1}GB to {:.1}GB.", - memory_utilization * 100.0, avg_usage.memory_gb, target_memory + memory_utilization * 100.0, + avg_usage.memory_gb, + target_memory ), estimated_savings: savings, savings_percent: savings / self.project_monthly_cost(samples) * 100.0, @@ -556,8 +600,10 @@ impl CostOptimizer { // Check storage tiering opportunity if avg_usage.storage_gb > 100.0 { let cold_data_ratio = 0.3; // Assume 30% could be cold - let cold_storage_savings = avg_usage.storage_gb * cold_data_ratio - * self.config.pricing.storage_per_gb_month * 0.7; // 70% savings on cold tier + let cold_storage_savings = avg_usage.storage_gb + * cold_data_ratio + * self.config.pricing.storage_per_gb_month + * 0.7; // 70% savings on cold tier recommendations.push(CostRecommendation { id: format!("storage-tier-{}", now), @@ -585,14 +631,18 @@ impl CostOptimizer { let estimated_cache_hit_rate = self.estimate_cache_hit_rate(samples); if estimated_cache_hit_rate < 0.7 { let potential_query_reduction = 0.2; // 20% of queries could be cached - let api_savings = avg_usage.queries as f64 * potential_query_reduction - * self.config.pricing.api_per_million / 1_000_000.0 * 720.0; + let api_savings = avg_usage.queries as f64 + * potential_query_reduction + * self.config.pricing.api_per_million + / 1_000_000.0 + * 720.0; recommendations.push(CostRecommendation { id: format!("cache-improve-{}", now), category: OptimizationCategory::CachingStrategy, title: "Improve Cache Hit Rate".to_string(), - description: "Increase cache size or adjust caching strategy to improve hit rate.".to_string(), + description: "Increase cache size or adjust caching strategy to improve hit rate." + .to_string(), estimated_savings: api_savings, savings_percent: api_savings / self.project_monthly_cost(samples) * 100.0, risk: RiskLevel::Low, @@ -608,11 +658,11 @@ impl CostOptimizer { } // Sort by estimated savings - recommendations.sort_by(|a, b| { - b.estimated_savings.total_cmp(&a.estimated_savings) - }); + recommendations.sort_by(|a, b| b.estimated_savings.total_cmp(&a.estimated_savings)); - self.stats.recommendations_generated.fetch_add(recommendations.len() as u64, Ordering::Relaxed); + self.stats + .recommendations_generated + .fetch_add(recommendations.len() as u64, Ordering::Relaxed); recommendations } @@ -631,7 +681,8 @@ impl CostOptimizer { .unwrap_or(4.0); // Calculate average CPU cores in use across samples - let avg_cores_used: f64 = samples.iter().map(|s| s.cpu_cores).sum::() / samples.len() as f64; + let avg_cores_used: f64 = + samples.iter().map(|s| s.cpu_cores).sum::() / samples.len() as f64; // Convert to utilization percentage (0.0 - 1.0) (avg_cores_used / total_cores).clamp(0.0, 1.0) @@ -649,7 +700,8 @@ impl CostOptimizer { let total_memory_gb = 16.0; // Assume 16GB system // Calculate average memory usage across samples - let avg_memory_used: f64 = samples.iter().map(|s| s.memory_gb).sum::() / samples.len() as f64; + let avg_memory_used: f64 = + samples.iter().map(|s| s.memory_gb).sum::() / samples.len() as f64; // Convert to utilization percentage (0.0 - 1.0) (avg_memory_used / total_memory_gb).clamp(0.0, 1.0) @@ -664,7 +716,8 @@ impl CostOptimizer { // Estimate cache hit rate based on query patterns // Higher query rates with stable latency suggest good caching - let avg_latency: f64 = samples.iter().map(|s| s.latency_ms).sum::() / samples.len() as f64; + let avg_latency: f64 = + samples.iter().map(|s| s.latency_ms).sum::() / samples.len() as f64; let total_queries: u64 = samples.iter().map(|s| s.queries).sum(); // Lower latency relative to query volume suggests cache hits @@ -735,18 +788,26 @@ impl CostOptimizer { /// Apply a recommendation pub fn apply_recommendation(&self, recommendation_id: &str) -> Result { - let mut recommendations = self.recommendations.write() + let mut recommendations = self + .recommendations + .write() .map_err(|_| VecStoreError::LockError("recommendations lock poisoned".into()))?; - if let Some(rec) = recommendations.iter_mut().find(|r| r.id == recommendation_id) { + if let Some(rec) = recommendations + .iter_mut() + .find(|r| r.id == recommendation_id) + { if rec.status != RecommendationStatus::Pending { return Ok(false); } rec.status = RecommendationStatus::Applied; - self.stats.recommendations_applied.fetch_add(1, Ordering::Relaxed); - *self.stats.savings_achieved.write() - .map_err(|_| VecStoreError::LockError("savings_achieved lock poisoned".into()))? += rec.estimated_savings; + self.stats + .recommendations_applied + .fetch_add(1, Ordering::Relaxed); + *self.stats.savings_achieved.write().map_err(|_| { + VecStoreError::LockError("savings_achieved lock poisoned".into()) + })? += rec.estimated_savings; Ok(true) } else { @@ -756,10 +817,15 @@ impl CostOptimizer { /// Dismiss a recommendation pub fn dismiss_recommendation(&self, recommendation_id: &str) -> Result { - let mut recommendations = self.recommendations.write() + let mut recommendations = self + .recommendations + .write() .map_err(|_| VecStoreError::LockError("recommendations lock poisoned".into()))?; - if let Some(rec) = recommendations.iter_mut().find(|r| r.id == recommendation_id) { + if let Some(rec) = recommendations + .iter_mut() + .find(|r| r.id == recommendation_id) + { rec.status = RecommendationStatus::Dismissed; Ok(true) } else { @@ -769,7 +835,9 @@ impl CostOptimizer { /// Get pending recommendations pub fn get_recommendations(&self) -> Vec { - let Ok(recommendations) = self.recommendations.read() else { return Vec::new(); }; + let Ok(recommendations) = self.recommendations.read() else { + return Vec::new(); + }; recommendations .iter() .filter(|r| r.status == RecommendationStatus::Pending) @@ -784,7 +852,10 @@ impl CostOptimizer { total_cost: self.stats.total_cost.load(Ordering::Relaxed) as f64 / 100.0, total_queries: self.stats.total_queries.load(Ordering::Relaxed), cost_per_query: 0.0, - recommendations_generated: self.stats.recommendations_generated.load(Ordering::Relaxed), + recommendations_generated: self + .stats + .recommendations_generated + .load(Ordering::Relaxed), recommendations_applied: self.stats.recommendations_applied.load(Ordering::Relaxed), savings_achieved: 0.0, }; @@ -795,7 +866,11 @@ impl CostOptimizer { cost_per_query: { let cost = self.stats.total_cost.load(Ordering::Relaxed) as f64 / 100.0; let queries = self.stats.total_queries.load(Ordering::Relaxed); - if queries > 0 { cost / queries as f64 } else { 0.0 } + if queries > 0 { + cost / queries as f64 + } else { + 0.0 + } }, recommendations_generated: self.stats.recommendations_generated.load(Ordering::Relaxed), recommendations_applied: self.stats.recommendations_applied.load(Ordering::Relaxed), @@ -812,7 +887,9 @@ impl CostOptimizer { pessimistic_forecast: 0.0, with_recommendations: 0.0, }; - let Ok(history) = self.resource_usage.read() else { return default_forecast; }; + let Ok(history) = self.resource_usage.read() else { + return default_forecast; + }; let monthly_cost = self.project_monthly_cost(&history.samples); let daily_cost = monthly_cost / 30.0; @@ -975,15 +1052,16 @@ impl CostAwareAutoScaler { // Check cooldown if let Some(last_action) = state.last_scale_action - && last_action.elapsed() < self.config.cooldown_period { - return ScalingDecision { - action: ScaleAction::NoChange, - current_replicas: state.replicas, - target_replicas: state.replicas, - reason: "In cooldown period".to_string(), - cost_impact: 0.0, - }; - } + && last_action.elapsed() < self.config.cooldown_period + { + return ScalingDecision { + action: ScaleAction::NoChange, + current_replicas: state.replicas, + target_replicas: state.replicas, + reason: "In cooldown period".to_string(), + cost_impact: 0.0, + }; + } let current_replicas = state.replicas; drop(state); @@ -1025,7 +1103,8 @@ impl CostAwareAutoScaler { target_replicas: target, reason: format!( "Low utilization ({:.0}%) and high cost per query (${:.6})", - utilization_per_replica * 100.0, current_cost_per_query + utilization_per_replica * 100.0, + current_cost_per_query ), cost_impact: -cost_savings, }; @@ -1047,7 +1126,8 @@ impl CostAwareAutoScaler { let memory_per_replica = 8.0; (cpu_per_replica * pricing.cpu_per_core_hour - + memory_per_replica * pricing.memory_per_gb_hour) * 720.0 + + memory_per_replica * pricing.memory_per_gb_hour) + * 720.0 } /// Apply scaling decision @@ -1056,7 +1136,9 @@ impl CostAwareAutoScaler { return false; } - let Ok(mut state) = self.current_scale.write() else { return false; }; + let Ok(mut state) = self.current_scale.write() else { + return false; + }; state.replicas = decision.target_replicas; state.last_scale_action = Some(Instant::now()); diff --git a/src/data_lineage.rs b/src/data_lineage.rs index fabf449..cf70d7c 100644 --- a/src/data_lineage.rs +++ b/src/data_lineage.rs @@ -129,7 +129,10 @@ pub enum DerivationOperation { /// Normalization Normalization { method: String }, /// Custom transformation - Custom { name: String, parameters: HashMap }, + Custom { + name: String, + parameters: HashMap, + }, } /// Transformation applied to a vector @@ -313,14 +316,19 @@ impl LineageTracker { // Handle derived sources if let VectorSource::Derived { ref parent_ids, .. } = source { let Ok(mut forward) = self.forward_deps.write() else { - return Err(VecStoreError::Internal("Failed to acquire forward_deps write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire forward_deps write lock".to_string(), + )); }; let Ok(mut backward) = self.backward_deps.write() else { - return Err(VecStoreError::Internal("Failed to acquire backward_deps write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire backward_deps write lock".to_string(), + )); }; for parent_id in parent_ids { - forward.entry(parent_id.clone()) + forward + .entry(parent_id.clone()) .or_insert_with(HashSet::new) .insert(vector_id.clone()); } @@ -329,7 +337,9 @@ impl LineageTracker { } let Ok(mut records) = self.records.write() else { - return Err(VecStoreError::Internal("Failed to acquire records write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire records write lock".to_string(), + )); }; records.insert(vector_id, record); @@ -350,7 +360,9 @@ impl LineageTracker { // Update input records { let Ok(mut records) = self.records.write() else { - return Err(VecStoreError::Internal("Failed to acquire records write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire records write lock".to_string(), + )); }; for input_id in &transformation.inputs { if let Some(record) = records.get_mut(input_id) { @@ -363,10 +375,14 @@ impl LineageTracker { // Track dependencies { let Ok(mut forward) = self.forward_deps.write() else { - return Err(VecStoreError::Internal("Failed to acquire forward_deps write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire forward_deps write lock".to_string(), + )); }; let Ok(mut backward) = self.backward_deps.write() else { - return Err(VecStoreError::Internal("Failed to acquire backward_deps write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire backward_deps write lock".to_string(), + )); }; for output_id in &transformation.outputs { @@ -384,12 +400,16 @@ impl LineageTracker { } let Ok(mut transformations) = self.transformations.write() else { - return Err(VecStoreError::Internal("Failed to acquire transformations write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire transformations write lock".to_string(), + )); }; transformations.insert(transform_id, transformation); use std::sync::atomic::Ordering; - self.metrics.transformations_recorded.fetch_add(1, Ordering::Relaxed); + self.metrics + .transformations_recorded + .fetch_add(1, Ordering::Relaxed); Ok(()) } @@ -415,7 +435,9 @@ impl LineageTracker { }; let Ok(mut log) = self.access_log.write() else { - return Err(VecStoreError::Internal("Failed to acquire access_log write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire access_log write lock".to_string(), + )); }; log.push_back(record); @@ -542,7 +564,8 @@ impl LineageTracker { let Ok(records) = self.records.read() else { return Vec::new(); }; - records.iter() + records + .iter() .filter(|(_, record)| { let record_type = match &record.source { VectorSource::DirectInsert { .. } => "direct", @@ -565,7 +588,8 @@ impl LineageTracker { let Ok(records) = self.records.read() else { return Vec::new(); }; - records.iter() + records + .iter() .filter(|(_, record)| record.tags.contains(tag)) .map(|(id, _)| id.clone()) .collect() @@ -574,14 +598,19 @@ impl LineageTracker { /// Add tag to vector pub fn add_tag(&self, vector_id: &VectorId, tag: &str) -> Result<()> { let Ok(mut records) = self.records.write() else { - return Err(VecStoreError::Internal("Failed to acquire records write lock".to_string())); + return Err(VecStoreError::Internal( + "Failed to acquire records write lock".to_string(), + )); }; if let Some(record) = records.get_mut(vector_id) { record.tags.insert(tag.to_string()); record.modified_at = current_timestamp(); Ok(()) } else { - Err(VecStoreError::NotFound(format!("Vector {:?} not found", vector_id))) + Err(VecStoreError::NotFound(format!( + "Vector {:?} not found", + vector_id + ))) } } @@ -595,7 +624,9 @@ impl LineageTracker { }; if let Some(record) = records.get(vector_id) { - record.transformations.iter() + record + .transformations + .iter() .filter_map(|id| transformations.get(id).cloned()) .collect() } else { @@ -689,7 +720,10 @@ impl LineageTracker { total_transformations: 0, records_by_source: HashMap::new(), records_created: self.metrics.records_created.load(Ordering::Relaxed), - transformations_recorded: self.metrics.transformations_recorded.load(Ordering::Relaxed), + transformations_recorded: self + .metrics + .transformations_recorded + .load(Ordering::Relaxed), lineage_queries: self.metrics.lineage_queries.load(Ordering::Relaxed), impact_analyses: self.metrics.impact_analyses.load(Ordering::Relaxed), }; @@ -700,7 +734,10 @@ impl LineageTracker { total_transformations: 0, records_by_source: HashMap::new(), records_created: self.metrics.records_created.load(Ordering::Relaxed), - transformations_recorded: self.metrics.transformations_recorded.load(Ordering::Relaxed), + transformations_recorded: self + .metrics + .transformations_recorded + .load(Ordering::Relaxed), lineage_queries: self.metrics.lineage_queries.load(Ordering::Relaxed), impact_analyses: self.metrics.impact_analyses.load(Ordering::Relaxed), }; @@ -729,7 +766,10 @@ impl LineageTracker { total_transformations, records_by_source: source_counts, records_created: self.metrics.records_created.load(Ordering::Relaxed), - transformations_recorded: self.metrics.transformations_recorded.load(Ordering::Relaxed), + transformations_recorded: self + .metrics + .transformations_recorded + .load(Ordering::Relaxed), lineage_queries: self.metrics.lineage_queries.load(Ordering::Relaxed), impact_analyses: self.metrics.impact_analyses.load(Ordering::Relaxed), } @@ -737,7 +777,8 @@ impl LineageTracker { /// Cleanup old records based on retention policy pub fn cleanup(&self) -> usize { - let cutoff = current_timestamp() - (self.config.retention_days as u64 * 24 * 60 * 60 * 1000); + let cutoff = + current_timestamp() - (self.config.retention_days as u64 * 24 * 60 * 60 * 1000); let Ok(mut records) = self.records.write() else { return 0; @@ -807,7 +848,8 @@ fn compute_quality_metrics(vector: &[f32]) -> QualityMetrics { } let n = vector.len() as f32; - let entropy: f32 = value_counts.values() + let entropy: f32 = value_counts + .values() .map(|&count| { let p = count as f32 / n; if p > 0.0 { -p * p.log2() } else { 0.0 } @@ -838,13 +880,9 @@ mod tests { api_version: "v1".to_string(), }; - tracker.record_creation( - vector_id.clone(), - source, - 128, - None, - Some(&vec![0.1; 128]), - ).unwrap(); + tracker + .record_creation(vector_id.clone(), source, 128, None, Some(&vec![0.1; 128])) + .unwrap(); let lineage = tracker.get_lineage(&vector_id).unwrap(); assert_eq!(lineage.original_dimensions, 128); @@ -858,32 +896,48 @@ mod tests { let parent1 = VectorId::new("test", "parent1"); let parent2 = VectorId::new("test", "parent2"); - tracker.record_creation( - parent1.clone(), - VectorSource::DirectInsert { - user: None, client_ip: None, api_version: "v1".to_string(), - }, - 128, None, None, - ).unwrap(); - - tracker.record_creation( - parent2.clone(), - VectorSource::DirectInsert { - user: None, client_ip: None, api_version: "v1".to_string(), - }, - 128, None, None, - ).unwrap(); + tracker + .record_creation( + parent1.clone(), + VectorSource::DirectInsert { + user: None, + client_ip: None, + api_version: "v1".to_string(), + }, + 128, + None, + None, + ) + .unwrap(); + + tracker + .record_creation( + parent2.clone(), + VectorSource::DirectInsert { + user: None, + client_ip: None, + api_version: "v1".to_string(), + }, + 128, + None, + None, + ) + .unwrap(); // Create derived vector let child = VectorId::new("test", "child"); - tracker.record_creation( - child.clone(), - VectorSource::Derived { - parent_ids: vec![parent1.clone(), parent2.clone()], - operation: DerivationOperation::Average, - }, - 128, None, None, - ).unwrap(); + tracker + .record_creation( + child.clone(), + VectorSource::Derived { + parent_ids: vec![parent1.clone(), parent2.clone()], + operation: DerivationOperation::Average, + }, + 128, + None, + None, + ) + .unwrap(); // Check ancestry let ancestors = tracker.get_ancestry(&child); @@ -902,34 +956,52 @@ mod tests { let tracker = LineageTracker::new(LineageConfig::default()); let root = VectorId::new("test", "root"); - tracker.record_creation( - root.clone(), - VectorSource::DirectInsert { - user: None, client_ip: None, api_version: "v1".to_string(), - }, - 128, None, None, - ).unwrap(); + tracker + .record_creation( + root.clone(), + VectorSource::DirectInsert { + user: None, + client_ip: None, + api_version: "v1".to_string(), + }, + 128, + None, + None, + ) + .unwrap(); // Create chain of derived vectors let level1 = VectorId::new("test", "level1"); - tracker.record_creation( - level1.clone(), - VectorSource::Derived { - parent_ids: vec![root.clone()], - operation: DerivationOperation::Normalization { method: "l2".to_string() }, - }, - 128, None, None, - ).unwrap(); + tracker + .record_creation( + level1.clone(), + VectorSource::Derived { + parent_ids: vec![root.clone()], + operation: DerivationOperation::Normalization { + method: "l2".to_string(), + }, + }, + 128, + None, + None, + ) + .unwrap(); let level2 = VectorId::new("test", "level2"); - tracker.record_creation( - level2.clone(), - VectorSource::Derived { - parent_ids: vec![level1.clone()], - operation: DerivationOperation::Quantization { method: "pq".to_string() }, - }, - 128, None, None, - ).unwrap(); + tracker + .record_creation( + level2.clone(), + VectorSource::Derived { + parent_ids: vec![level1.clone()], + operation: DerivationOperation::Quantization { + method: "pq".to_string(), + }, + }, + 128, + None, + None, + ) + .unwrap(); let report = tracker.impact_analysis(&root); assert_eq!(report.total_affected, 2); @@ -942,22 +1014,32 @@ mod tests { let parent = VectorId::new("test", "parent"); let child = VectorId::new("test", "child"); - tracker.record_creation( - parent.clone(), - VectorSource::DirectInsert { - user: None, client_ip: None, api_version: "v1".to_string(), - }, - 128, None, None, - ).unwrap(); - - tracker.record_creation( - child.clone(), - VectorSource::Derived { - parent_ids: vec![parent.clone()], - operation: DerivationOperation::Average, - }, - 128, None, None, - ).unwrap(); + tracker + .record_creation( + parent.clone(), + VectorSource::DirectInsert { + user: None, + client_ip: None, + api_version: "v1".to_string(), + }, + 128, + None, + None, + ) + .unwrap(); + + tracker + .record_creation( + child.clone(), + VectorSource::Derived { + parent_ids: vec![parent.clone()], + operation: DerivationOperation::Average, + }, + 128, + None, + None, + ) + .unwrap(); let dot = tracker.export_graph(&parent, 2); assert!(dot.contains("digraph")); diff --git a/src/debugger.rs b/src/debugger.rs index 6d31c10..35d6a56 100644 --- a/src/debugger.rs +++ b/src/debugger.rs @@ -30,7 +30,7 @@ //! let anomalies = debugger.detect_anomalies(&embeddings)?; //! ``` -use anyhow::{Result}; +use anyhow::Result; use serde::{Deserialize, Serialize}; // ============================================================================ @@ -299,7 +299,8 @@ impl EmbeddingDebugger { distribution.norm ), affected_dimensions: None, - suggestion: "The embedding may be collapsed. Try different input or model.".to_string(), + suggestion: "The embedding may be collapsed. Try different input or model." + .to_string(), }); quality_score -= 0.2; } else if distribution.norm > self.config.high_norm_threshold { @@ -323,7 +324,8 @@ impl EmbeddingDebugger { category: IssueCategory::Collapsed, message: "Embedding appears to be collapsed (near-constant values)".to_string(), affected_dimensions: None, - suggestion: "The model may be producing degenerate embeddings. Check model output.".to_string(), + suggestion: "The model may be producing degenerate embeddings. Check model output." + .to_string(), }); quality_score -= 0.4; } @@ -347,7 +349,8 @@ impl EmbeddingDebugger { let dimension_stats = self.analyze_dimensions(embedding); // Check for low variance dimensions - let low_var_ratio = dimension_stats.zero_variance_dims as f32 / dimension_stats.total as f32; + let low_var_ratio = + dimension_stats.zero_variance_dims as f32 / dimension_stats.total as f32; if low_var_ratio > 0.5 { issues.push(QualityIssue { severity: IssueSeverity::Warning, @@ -357,7 +360,9 @@ impl EmbeddingDebugger { low_var_ratio * 100.0 ), affected_dimensions: None, - suggestion: "Many dimensions are not being used. Consider dimensionality reduction.".to_string(), + suggestion: + "Many dimensions are not being used. Consider dimensionality reduction." + .to_string(), }); quality_score -= 0.1; } @@ -367,7 +372,10 @@ impl EmbeddingDebugger { issues.push(QualityIssue { severity: IssueSeverity::Info, category: IssueCategory::Skewed, - message: format!("Distribution is highly skewed ({:.2})", distribution.skewness), + message: format!( + "Distribution is highly skewed ({:.2})", + distribution.skewness + ), affected_dimensions: None, suggestion: "This may be normal for your data type.".to_string(), }); @@ -545,7 +553,8 @@ impl EmbeddingDebugger { .map(|(i, &v)| (i, (v - mean).abs())) .collect(); variance_indexed.sort_by(|a, b| b.1.total_cmp(&a.1)); - let high_variance_dims: Vec<(usize, f32)> = variance_indexed.iter().take(10).cloned().collect(); + let high_variance_dims: Vec<(usize, f32)> = + variance_indexed.iter().take(10).cloned().collect(); DimensionStats { total, @@ -571,11 +580,13 @@ impl EmbeddingDebugger { let has_high_norm = issues.iter().any(|i| i.category == IssueCategory::HighNorm); if has_critical { - recommendations.push("Critical issues detected. Review your embedding pipeline.".to_string()); + recommendations + .push("Critical issues detected. Review your embedding pipeline.".to_string()); } if has_low_norm { - recommendations.push("Consider checking if input text is empty or very short.".to_string()); + recommendations + .push("Consider checking if input text is empty or very short.".to_string()); } if has_high_norm { @@ -583,7 +594,8 @@ impl EmbeddingDebugger { } if distribution.sparsity > 0.5 { - recommendations.push("High sparsity detected. Consider using sparse vector storage.".to_string()); + recommendations + .push("High sparsity detected. Consider using sparse vector storage.".to_string()); } if issues.is_empty() { @@ -699,9 +711,10 @@ impl EmbeddingDebugger { /// Check for drift compared to baseline pub fn check_drift(&mut self, embeddings: &[Vec]) -> Result { - let baseline = self.baseline.as_ref().ok_or_else(|| { - anyhow::anyhow!("No baseline set. Call set_baseline first.") - })?; + let baseline = self + .baseline + .as_ref() + .ok_or_else(|| anyhow::anyhow!("No baseline set. Call set_baseline first."))?; if embeddings.is_empty() { return Ok(DriftReport { @@ -724,8 +737,8 @@ impl EmbeddingDebugger { let current_mean_norm = norms.iter().sum::() / n; // Norm drift - let norm_drift = (current_mean_norm - baseline.mean_norm).abs() - / (baseline.std_norm + 1e-10); + let norm_drift = + (current_mean_norm - baseline.mean_norm).abs() / (baseline.std_norm + 1e-10); // Per-dimension drift let mut current_mean_per_dim = vec![0.0f32; baseline.dimension]; @@ -958,10 +971,12 @@ mod tests { let report = debugger.analyze(&embedding).unwrap(); assert!(report.quality_score < 0.7); - assert!(report - .issues - .iter() - .any(|i| i.category == IssueCategory::LowEntropy)); + assert!( + report + .issues + .iter() + .any(|i| i.category == IssueCategory::LowEntropy) + ); } #[test] @@ -973,10 +988,12 @@ mod tests { let report = debugger.analyze(&embedding).unwrap(); - assert!(report - .issues - .iter() - .any(|i| i.category == IssueCategory::InvalidValues)); + assert!( + report + .issues + .iter() + .any(|i| i.category == IssueCategory::InvalidValues) + ); assert!(report.quality_score < 0.6); } @@ -985,9 +1002,8 @@ mod tests { let debugger = EmbeddingDebugger::new(DebugConfig::default()); // Generate normal embeddings - let mut embeddings: Vec> = (0..50) - .map(|_| generate_random_embedding(128)) - .collect(); + let mut embeddings: Vec> = + (0..50).map(|_| generate_random_embedding(128)).collect(); // Add an anomaly (very low norm) embeddings.push(vec![0.001; 128]); @@ -1005,15 +1021,11 @@ mod tests { let mut debugger = EmbeddingDebugger::new(DebugConfig::default()); // Set baseline - let baseline: Vec> = (0..100) - .map(|_| generate_random_embedding(128)) - .collect(); + let baseline: Vec> = (0..100).map(|_| generate_random_embedding(128)).collect(); debugger.set_baseline(&baseline).unwrap(); // Check similar embeddings (no drift) - let similar: Vec> = (0..50) - .map(|_| generate_random_embedding(128)) - .collect(); + let similar: Vec> = (0..50).map(|_| generate_random_embedding(128)).collect(); let report = debugger.check_drift(&similar).unwrap(); println!("Drift score: {}", report.overall_drift_score); diff --git a/src/deduplication.rs b/src/deduplication.rs index 985b9a7..a19b5ca 100644 --- a/src/deduplication.rs +++ b/src/deduplication.rs @@ -234,10 +234,10 @@ impl Deduplicator { DeduplicationStrategy::KeepLast => group.duplicates.last().unwrap(), DeduplicationStrategy::KeepMostMetadata => { self.select_most_metadata(&group.duplicates, store)? - } + }, DeduplicationStrategy::KeepHighestQuality => { self.select_highest_quality(&group.duplicates, store)? - } + }, }; // Remove all except the one to keep diff --git a/src/dim_reduction.rs b/src/dim_reduction.rs index 44b7589..e604ef3 100644 --- a/src/dim_reduction.rs +++ b/src/dim_reduction.rs @@ -28,7 +28,7 @@ //! assert_eq!(reduced[0].len(), 2); //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; /// Principal Component Analysis (PCA) diff --git a/src/discovery.rs b/src/discovery.rs index 3e5e520..09e20ba 100644 --- a/src/discovery.rs +++ b/src/discovery.rs @@ -257,7 +257,12 @@ impl DiscoveryIndex { } /// Insert vector - pub fn insert(&self, id: &str, vector: &[f32], metadata: HashMap) -> Result<()> { + pub fn insert( + &self, + id: &str, + vector: &[f32], + metadata: HashMap, + ) -> Result<()> { if vector.len() != self.config.dimension { return Err(VecStoreError::DimensionMismatch { expected: self.config.dimension, @@ -265,12 +270,14 @@ impl DiscoveryIndex { }); } - let mut vectors = self.vectors.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on vectors".into()))?; + let mut vectors = self.vectors.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on vectors".into()) + })?; vectors.insert(id.to_string(), vector.to_vec()); - let mut meta = self.metadata.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on metadata".into()))?; + let mut meta = self.metadata.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on metadata".into()) + })?; meta.insert(id.to_string(), metadata); Ok(()) @@ -278,33 +285,41 @@ impl DiscoveryIndex { /// Get vector by ID pub fn get(&self, id: &str) -> Option> { - let Ok(vectors) = self.vectors.read() else { return None; }; + let Ok(vectors) = self.vectors.read() else { + return None; + }; vectors.get(id).cloned() } /// Execute discovery query pub fn discover(&self, query: DiscoveryQuery) -> Result> { - let vectors = self.vectors.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on vectors".into()))?; - let metadata = self.metadata.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on metadata".into()))?; + let vectors = self.vectors.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on vectors".into()) + })?; + let metadata = self.metadata.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on metadata".into()) + })?; // Get positive vectors - let positive_vecs: Vec> = query.positive_ids + let positive_vecs: Vec> = query + .positive_ids .iter() .filter_map(|id| vectors.get(id).cloned()) .chain(query.positive_vectors.clone()) .collect(); // Get negative vectors - let negative_vecs: Vec> = query.negative_ids + let negative_vecs: Vec> = query + .negative_ids .iter() .filter_map(|id| vectors.get(id).cloned()) .chain(query.negative_vectors.clone()) .collect(); // Exclude positive and negative IDs from results - let exclude_ids: HashSet<&str> = query.positive_ids.iter() + let exclude_ids: HashSet<&str> = query + .positive_ids + .iter() .chain(&query.negative_ids) .map(|s| s.as_str()) .collect(); @@ -315,22 +330,28 @@ impl DiscoveryIndex { .filter(|(id, _)| !exclude_ids.contains(id.as_str())) .map(|(id, vec)| { let pos_score = if !positive_vecs.is_empty() { - positive_vecs.iter() + positive_vecs + .iter() .map(|pv| cosine_similarity(vec, pv)) - .sum::() / positive_vecs.len() as f32 + .sum::() + / positive_vecs.len() as f32 } else { 0.0 }; let neg_score = if !negative_vecs.is_empty() { - negative_vecs.iter() + negative_vecs + .iter() .map(|nv| cosine_similarity(vec, nv)) - .sum::() / negative_vecs.len() as f32 + .sum::() + / negative_vecs.len() as f32 } else { 0.0 }; - let target_score = query.target.as_ref() + let target_score = query + .target + .as_ref() .map(|t| cosine_similarity(vec, t)) .unwrap_or(0.0); @@ -338,23 +359,18 @@ impl DiscoveryIndex { let score = match query.mode { DiscoveryMode::Context => { pos_score * query.positive_weight - neg_score * query.negative_weight - } + }, DiscoveryMode::Discover => { // Maximize novelty (low similarity to all examples) let all_sim = (pos_score + neg_score) / 2.0; 1.0 - all_sim - } - DiscoveryMode::Recommend => { - pos_score * query.positive_weight - } - DiscoveryMode::Contrast => { - -neg_score * query.negative_weight - } + }, + DiscoveryMode::Recommend => pos_score * query.positive_weight, + DiscoveryMode::Contrast => -neg_score * query.negative_weight, DiscoveryMode::Blend => { - pos_score * query.positive_weight - - neg_score * query.negative_weight + pos_score * query.positive_weight - neg_score * query.negative_weight + target_score * 0.5 - } + }, }; DiscoveryResult { @@ -406,7 +422,8 @@ impl DiscoveryIndex { let diversity_penalty = if selected_vecs.is_empty() { 0.0 } else { - selected_vecs.iter() + selected_vecs + .iter() .map(|sv| cosine_similarity(vec, sv)) .fold(f32::NEG_INFINITY, f32::max) }; @@ -434,13 +451,21 @@ impl DiscoveryIndex { } /// Explore from a starting point - pub fn explore(&self, start_id: &str, radius: f32, max_results: usize) -> Result> { - let vectors = self.vectors.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on vectors".into()))?; - let metadata = self.metadata.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on metadata".into()))?; - - let start_vec = vectors.get(start_id) + pub fn explore( + &self, + start_id: &str, + radius: f32, + max_results: usize, + ) -> Result> { + let vectors = self.vectors.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on vectors".into()) + })?; + let metadata = self.metadata.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on metadata".into()) + })?; + + let start_vec = vectors + .get(start_id) .ok_or_else(|| VecStoreError::NotFound(start_id.to_string()))?; let mut results: Vec = vectors @@ -470,8 +495,12 @@ impl DiscoveryIndex { /// Find boundary vectors (transition points between clusters) pub fn find_boundaries(&self, num_boundaries: usize) -> Vec { - let Ok(vectors) = self.vectors.read() else { return Vec::new(); }; - let Ok(metadata) = self.metadata.read() else { return Vec::new(); }; + let Ok(vectors) = self.vectors.read() else { + return Vec::new(); + }; + let Ok(metadata) = self.metadata.read() else { + return Vec::new(); + }; if vectors.len() < 3 { return Vec::new(); @@ -496,7 +525,8 @@ impl DiscoveryIndex { if k > 0 { let top_k: Vec = similarities.into_iter().take(k).collect(); let mean: f32 = top_k.iter().sum::() / k as f32; - let variance: f32 = top_k.iter().map(|s| (s - mean).powi(2)).sum::() / k as f32; + let variance: f32 = + top_k.iter().map(|s| (s - mean).powi(2)).sum::() / k as f32; // Boundary score = variance (high variance = boundary) boundary_scores.push((id.to_string(), variance)); @@ -542,11 +572,8 @@ impl DiscoveryIndex { let vec_list: Vec<(&String, &Vec)> = vectors.iter().collect(); // Initialize centroids (first k vectors) - let mut centroids: Vec> = vec_list - .iter() - .take(k) - .map(|(_, v)| (*v).clone()) - .collect(); + let mut centroids: Vec> = + vec_list.iter().take(k).map(|(_, v)| (*v).clone()).collect(); // Run k-means for a few iterations let mut assignments: Vec = vec![0; vec_list.len()]; @@ -623,9 +650,9 @@ impl DiscoveryIndex { let mut member_sims: Vec<(String, f32)> = members .iter() .filter_map(|id| { - vectors.get(id).map(|v| { - (id.clone(), cosine_similarity(v, centroid)) - }) + vectors + .get(id) + .map(|v| (id.clone(), cosine_similarity(v, centroid))) }) .collect(); @@ -655,7 +682,11 @@ impl DiscoveryIndex { } } - let diversity_score = if count > 0 { total_dist / count as f32 } else { 0.0 }; + let diversity_score = if count > 0 { + total_dist / count as f32 + } else { + 0.0 + }; ExplorationResult { regions, @@ -719,15 +750,28 @@ mod tests { #[test] fn test_discovery_context() { - let config = DiscoveryConfig { dimension: 4, ..Default::default() }; + let config = DiscoveryConfig { + dimension: 4, + ..Default::default() + }; let index = DiscoveryIndex::new(config); // Insert vectors - index.insert("pos1", &[1.0, 0.0, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("pos2", &[0.9, 0.1, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("neg1", &[0.0, 0.0, 1.0, 0.0], HashMap::new()).unwrap(); - index.insert("candidate1", &[0.8, 0.2, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("candidate2", &[0.0, 0.0, 0.8, 0.2], HashMap::new()).unwrap(); + index + .insert("pos1", &[1.0, 0.0, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("pos2", &[0.9, 0.1, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("neg1", &[0.0, 0.0, 1.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("candidate1", &[0.8, 0.2, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("candidate2", &[0.0, 0.0, 0.8, 0.2], HashMap::new()) + .unwrap(); // Context search let query = DiscoveryQuery::new() @@ -743,12 +787,18 @@ mod tests { #[test] fn test_discovery_with_diversity() { - let config = DiscoveryConfig { dimension: 4, enable_mmr: true, ..Default::default() }; + let config = DiscoveryConfig { + dimension: 4, + enable_mmr: true, + ..Default::default() + }; let index = DiscoveryIndex::new(config); for i in 0..10 { let vec = vec![(i as f32) / 10.0, 1.0 - (i as f32) / 10.0, 0.0, 0.0]; - index.insert(&format!("vec_{}", i), &vec, HashMap::new()).unwrap(); + index + .insert(&format!("vec_{}", i), &vec, HashMap::new()) + .unwrap(); } let query = DiscoveryQuery::new() @@ -762,13 +812,24 @@ mod tests { #[test] fn test_explore() { - let config = DiscoveryConfig { dimension: 4, ..Default::default() }; + let config = DiscoveryConfig { + dimension: 4, + ..Default::default() + }; let index = DiscoveryIndex::new(config); - index.insert("center", &[1.0, 0.0, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("near1", &[0.95, 0.05, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("near2", &[0.9, 0.1, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("far", &[0.0, 1.0, 0.0, 0.0], HashMap::new()).unwrap(); + index + .insert("center", &[1.0, 0.0, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("near1", &[0.95, 0.05, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("near2", &[0.9, 0.1, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("far", &[0.0, 1.0, 0.0, 0.0], HashMap::new()) + .unwrap(); let results = index.explore("center", 0.2, 10).unwrap(); @@ -789,12 +850,16 @@ mod tests { // Create two clusters for i in 0..5 { let vec = vec![1.0 - (i as f32) * 0.05, (i as f32) * 0.05, 0.0, 0.0]; - index.insert(&format!("cluster1_{}", i), &vec, HashMap::new()).unwrap(); + index + .insert(&format!("cluster1_{}", i), &vec, HashMap::new()) + .unwrap(); } for i in 0..5 { let vec = vec![0.0, 0.0, 1.0 - (i as f32) * 0.05, (i as f32) * 0.05]; - index.insert(&format!("cluster2_{}", i), &vec, HashMap::new()).unwrap(); + index + .insert(&format!("cluster2_{}", i), &vec, HashMap::new()) + .unwrap(); } let exploration = index.find_clusters(); @@ -805,12 +870,21 @@ mod tests { #[test] fn test_discovery_modes() { - let config = DiscoveryConfig { dimension: 4, ..Default::default() }; + let config = DiscoveryConfig { + dimension: 4, + ..Default::default() + }; let index = DiscoveryIndex::new(config); - index.insert("v1", &[1.0, 0.0, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("v2", &[0.0, 1.0, 0.0, 0.0], HashMap::new()).unwrap(); - index.insert("v3", &[0.0, 0.0, 1.0, 0.0], HashMap::new()).unwrap(); + index + .insert("v1", &[1.0, 0.0, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("v2", &[0.0, 1.0, 0.0, 0.0], HashMap::new()) + .unwrap(); + index + .insert("v3", &[0.0, 0.0, 1.0, 0.0], HashMap::new()) + .unwrap(); // Discover mode (novelty) let query = DiscoveryQuery::new() diff --git a/src/diskann.rs b/src/diskann.rs index e796d49..536f058 100644 --- a/src/diskann.rs +++ b/src/diskann.rs @@ -53,8 +53,8 @@ use std::collections::{BinaryHeap, HashMap, HashSet}; use std::fs::File; use std::io::{BufReader, BufWriter}; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; /// DiskANN configuration #[derive(Debug, Clone, Serialize, Deserialize)] @@ -163,10 +163,8 @@ impl PQQuantizer { let end = start + self.subvector_dim; // Extract subvectors - let subvectors: Vec> = vectors - .iter() - .map(|v| v[start..end].to_vec()) - .collect(); + let subvectors: Vec> = + vectors.iter().map(|v| v[start..end].to_vec()).collect(); // K-means clustering let centroids = self.kmeans(&subvectors, self.num_centroids, iterations)?; @@ -527,7 +525,8 @@ impl DiskANN { let node_vec = &vectors[node_idx]; // Greedy search to find candidates - let candidates = self.greedy_search_build(vectors, node_vec, self.config.build_beam_width)?; + let candidates = + self.greedy_search_build(vectors, node_vec, self.config.build_beam_width)?; // RobustPrune to select neighbors let neighbors = self.robust_prune(vectors, node_id, &candidates)?; @@ -549,8 +548,8 @@ impl DiskANN { // Prune if over max degree (done outside the mutable borrow) if needs_prune { - let candidates: Vec = self.graph.get(&neighbor_id).cloned() - .unwrap_or_default(); + let candidates: Vec = + self.graph.get(&neighbor_id).cloned().unwrap_or_default(); let pruned = self.robust_prune(vectors, neighbor_id, &candidates)?; self.graph.insert(neighbor_id, pruned); } @@ -587,9 +586,11 @@ impl DiskANN { while let Some(current) = candidates.pop() { // Check if we can stop if let Some(std::cmp::Reverse(worst)) = results.peek() - && current.distance > worst.distance && results.len() >= beam_width { - break; - } + && current.distance > worst.distance + && results.len() >= beam_width + { + break; + } // Explore neighbors if let Some(neighbors) = self.graph.get(¤t.id) { @@ -674,7 +675,12 @@ impl DiskANN { /// Search for k nearest neighbors #[inline] - pub fn search(&self, query: &[f32], k: usize, beam_width: Option) -> Result> { + pub fn search( + &self, + query: &[f32], + k: usize, + beam_width: Option, + ) -> Result> { self.stats.total_searches.fetch_add(1, Ordering::Relaxed); let beam = beam_width.unwrap_or(self.config.search_beam_width); @@ -713,15 +719,19 @@ impl DiskANN { // Check stopping condition if let Some(std::cmp::Reverse(worst)) = results.peek() - && current.distance > worst.distance && results.len() >= beam { - continue; - } + && current.distance > worst.distance + && results.len() >= beam + { + continue; + } // Explore neighbors using PQ distances if let Some(neighbors) = self.graph.get(¤t.id) { for &neighbor_id in neighbors { if visited.insert(neighbor_id) { - self.stats.distance_computations.fetch_add(1, Ordering::Relaxed); + self.stats + .distance_computations + .fetch_add(1, Ordering::Relaxed); let dist = if let Some(codes) = self.pq_codes.get(&neighbor_id) { self.quantizer.table_distance(&pq_table, codes) @@ -1099,6 +1109,10 @@ mod tests { let recall = correct as f64 / (test_count * 10) as f64; println!("Recall@10: {:.2}%", recall * 100.0); - assert!(recall > 0.70, "Recall should be > 70%, got {:.2}%", recall * 100.0); + assert!( + recall > 0.70, + "Recall should be > 70%, got {:.2}%", + recall * 100.0 + ); } } diff --git a/src/distributed/mod.rs b/src/distributed/mod.rs index eda55de..561b423 100644 --- a/src/distributed/mod.rs +++ b/src/distributed/mod.rs @@ -444,7 +444,7 @@ impl DistributedStore { let mut hasher = DefaultHasher::new(); key.hash(&mut hasher); (hasher.finish() as usize) % self.config.num_shards - } + }, ShardingStrategy::ConsistentHash => { // Use consistent hash ring let node = self.hash_ring.get_node(key).unwrap_or_default(); @@ -454,15 +454,15 @@ impl DistributedStore { .iter() .fold(0u32, |acc, &b| acc.wrapping_add(b as u32)); (sum as usize) % self.config.num_shards - } + }, ShardingStrategy::Range => { // Range-based (simplified - would use key ranges in production) key.as_bytes().first().copied().unwrap_or(0) as usize % self.config.num_shards - } + }, ShardingStrategy::Random => { // Random (simplified) key.len() % self.config.num_shards - } + }, } } @@ -720,7 +720,7 @@ impl DistributedStore { let mut hasher = DefaultHasher::new(); key.hash(&mut hasher); (hasher.finish() as usize) % self.config.num_shards - } + }, ShardingStrategy::ConsistentHash => { let ring = self.hash_ring.read().await; let node = ring.get_node(key).unwrap_or_default(); @@ -729,10 +729,10 @@ impl DistributedStore { .iter() .fold(0u32, |acc, &b| acc.wrapping_add(b as u32)); (sum as usize) % self.config.num_shards - } + }, ShardingStrategy::Range => { key.as_bytes().first().copied().unwrap_or(0) as usize % self.config.num_shards - } + }, ShardingStrategy::Random => key.len() % self.config.num_shards, } } @@ -826,18 +826,18 @@ impl DistributedStore { // Simulate sync latency // In production: await network RPC } - } + }, ConsistencyLevel::Quorum => { // Wait for majority (N/2 + 1) let quorum_size = (replicas.len() / 2) + 1; for _i in 0..quorum_size { // Simulate sync to quorum } - } + }, ConsistencyLevel::One => { // Fire and forget - don't wait for acknowledgment // Replicas will eventually sync - } + }, } Ok(()) @@ -872,18 +872,18 @@ impl DistributedStore { // Always query primary for strong consistency // In real implementation: send RPC to primary node Ok(vec![]) - } + }, ConsistencyLevel::Quorum => { // Query quorum of nodes and return most recent results // In real implementation: scatter to quorum, gather results Ok(vec![]) - } + }, ConsistencyLevel::One => { // Query any available replica (round-robin or least loaded) // For now, use the first available replica // In real implementation: check node health and load Ok(vec![]) - } + }, } } diff --git a/src/distributed/raft.rs b/src/distributed/raft.rs index 964e5ff..089d4f8 100644 --- a/src/distributed/raft.rs +++ b/src/distributed/raft.rs @@ -13,7 +13,7 @@ use tokio::time; use tracing::{debug, error, info, warn}; #[cfg(feature = "server")] -use tonic::{transport::Channel, Request, Response, Status}; +use tonic::{Request, Response, Status, transport::Channel}; #[cfg(feature = "server")] pub mod pb { @@ -349,9 +349,9 @@ mod rpc { // Configure TLS if enabled if self.config.tls_enabled { if let Some(ca_path) = &self.config.tls_ca_path { - let ca_cert = tokio::fs::read(ca_path) - .await - .map_err(|e| RpcError::TlsError(format!("Failed to read CA cert: {}", e)))?; + let ca_cert = tokio::fs::read(ca_path).await.map_err(|e| { + RpcError::TlsError(format!("Failed to read CA cert: {}", e)) + })?; let ca = Certificate::from_pem(ca_cert); let tls_config = ClientTlsConfig::new().ca_certificate(ca); @@ -410,12 +410,12 @@ mod rpc { Ok(response) => { self.mark_connection_success(node_id).await; return Ok(response); - } + }, Err(e) => { warn!("RequestVote to {} failed: {:?}", node_id, e); self.mark_connection_failed(node_id).await; last_error = Some(e); - } + }, } } @@ -474,12 +474,12 @@ mod rpc { Ok(response) => { self.mark_connection_success(node_id).await; return Ok(response); - } + }, Err(e) => { warn!("AppendEntries to {} failed: {:?}", node_id, e); self.mark_connection_failed(node_id).await; last_error = Some(e); - } + }, } } @@ -494,11 +494,8 @@ mod rpc { ) -> Result { let mut client = self.get_connection(node_id).await?; - let pb_entries: Vec = request - .entries - .iter() - .map(|e| log_entry_to_pb(e)) - .collect(); + let pb_entries: Vec = + request.entries.iter().map(|e| log_entry_to_pb(e)).collect(); let pb_request = pb::AppendEntriesRequest { term: request.term, @@ -620,7 +617,7 @@ mod rpc { RpcError::RpcFailed(msg) => write!(f, "RPC failed: {}", msg), RpcError::DeserializationFailed(msg) => { write!(f, "Deserialization failed: {}", msg) - } + }, RpcError::MaxRetriesExceeded => write!(f, "Max retries exceeded"), RpcError::Timeout => write!(f, "RPC timeout"), } @@ -648,7 +645,11 @@ impl StateMachine { /// Apply a command to the state machine pub fn apply(&mut self, command: &Command) -> Option> { match command { - Command::Insert { id, vector, metadata } => { + Command::Insert { + id, + vector, + metadata, + } => { let entry = serde_json::json!({ "vector": vector, "metadata": metadata, @@ -656,12 +657,16 @@ impl StateMachine { let bytes = serde_json::to_vec(&entry).unwrap_or_default(); self.data.insert(id.clone(), bytes); None - } + }, Command::Delete { id } => { self.data.remove(id); None - } - Command::Update { id, vector, metadata } => { + }, + Command::Update { + id, + vector, + metadata, + } => { let entry = serde_json::json!({ "vector": vector, "metadata": metadata, @@ -669,7 +674,7 @@ impl StateMachine { let bytes = serde_json::to_vec(&entry).unwrap_or_default(); self.data.insert(id.clone(), bytes); None - } + }, Command::NoOp => None, } } @@ -882,7 +887,9 @@ impl RaftNode { T: serde::de::DeserializeOwned, { let state_machine = self.state_machine.read().await; - state_machine.get(key).and_then(|bytes| serde_json::from_slice(bytes).ok()) + state_machine + .get(key) + .and_then(|bytes| serde_json::from_slice(bytes).ok()) } /// Start an election @@ -965,17 +972,17 @@ impl RaftNode { self.config.node_id, peer_id, votes, votes_needed ); } - } + }, Ok(Ok((peer_id, Err(e)))) => { warn!("Failed to get vote from {}: {:?}", peer_id, e); - } + }, Ok(Err(e)) => { warn!("Vote task failed: {:?}", e); - } + }, Err(_) => { debug!("Vote collection timed out"); break; - } + }, } } } @@ -1150,8 +1157,7 @@ impl RaftNode { if let Some(ls) = leader_state.as_mut() { if response.success { // Update next_index and match_index - ls.match_index - .insert(peer_id.clone(), response.match_index); + ls.match_index.insert(peer_id.clone(), response.match_index); ls.next_index .insert(peer_id.clone(), response.match_index + 1); @@ -1171,10 +1177,10 @@ impl RaftNode { debug!("AppendEntries to {} failed, will retry", peer_id); } } - } + }, Err(e) => { warn!("Failed to send heartbeat to {}: {:?}", peer_id, e); - } + }, } })); } @@ -1422,8 +1428,9 @@ impl RaftNode { fn random_election_timeout(&self) -> Duration { use rand::Rng; let mut rng = rand::rng(); - let timeout_ms = rng - .random_range(self.config.election_timeout_min_ms..=self.config.election_timeout_max_ms); + let timeout_ms = rng.random_range( + self.config.election_timeout_min_ms..=self.config.election_timeout_max_ms, + ); Duration::from_millis(timeout_ms) } @@ -1592,8 +1599,11 @@ pub mod server { let req = request.into_inner(); // Convert protobuf log entries - let entries: Result, _> = - req.entries.iter().map(|e| rpc::pb_to_log_entry(e)).collect(); + let entries: Result, _> = req + .entries + .iter() + .map(|e| rpc::pb_to_log_entry(e)) + .collect(); let entries = entries.map_err(|e| Status::invalid_argument(format!("{}", e)))?; @@ -1716,7 +1726,7 @@ pub mod server { } #[cfg(feature = "server")] -pub use server::{start_raft_server, RaftGrpcService}; +pub use server::{RaftGrpcService, start_raft_server}; // ============================================================================ // Cluster Membership Management @@ -1928,8 +1938,7 @@ impl FailoverManager { /// Update known leader pub async fn set_leader(&self, leader: NodeId) { *self.current_leader.write().await = Some(leader); - self.failures - .store(0, std::sync::atomic::Ordering::Relaxed); + self.failures.store(0, std::sync::atomic::Ordering::Relaxed); } /// Get current leader @@ -1939,10 +1948,10 @@ impl FailoverManager { /// Report a failure pub async fn report_failure(&self) -> bool { - let failures = - self.failures - .fetch_add(1, std::sync::atomic::Ordering::Relaxed) - + 1; + let failures = self + .failures + .fetch_add(1, std::sync::atomic::Ordering::Relaxed) + + 1; if failures >= self.config.max_retries { // Trigger leader discovery @@ -2014,7 +2023,7 @@ impl ReplicationCoordinator { } // Would normally redirect to leader here time::sleep(Duration::from_millis(100)).await; - } + }, Err(e) => return Err(e), } } @@ -2037,12 +2046,12 @@ impl ReplicationCoordinator { } // Read from leader's state machine Ok(self.node.read_from_state_machine(key).await) - } + }, ReadConsistency::Any => { // Any read: Read from local state machine regardless of leader status // May return stale data if this node is behind Ok(self.node.read_from_state_machine(key).await) - } + }, ReadConsistency::Linearizable => { // Linearizable read: Strongest guarantee - verify leadership before reading if !self.node.is_leader().await { @@ -2058,7 +2067,7 @@ impl ReplicationCoordinator { // Now safe to read - we confirmed we're still the leader Ok(self.node.read_from_state_machine(key).await) - } + }, ReadConsistency::Quorum => { // Quorum read: Read from majority of nodes and return most recent // For now, if we're leader, we know we have the most up-to-date data @@ -2071,7 +2080,7 @@ impl ReplicationCoordinator { // For now, read local and indicate it may be stale Ok(self.node.read_from_state_machine(key).await) } - } + }, } } @@ -2177,10 +2186,10 @@ impl ReplicationCoordinator { match change { MembershipChange::AddNode(node_id) => { config.members.insert(node_id); - } + }, MembershipChange::RemoveNode(node_id) => { config.members.remove(&node_id); - } + }, } config.version += 1; diff --git a/src/dynamic_schema.rs b/src/dynamic_schema.rs index 83202bf..51e26be 100644 --- a/src/dynamic_schema.rs +++ b/src/dynamic_schema.rs @@ -37,12 +37,12 @@ //! collection.search_with_filter(&query, "score > 0.9")?; //! ``` -use std::collections::HashMap; -use std::sync::RwLock; use serde::{Deserialize, Serialize}; use serde_json::Value as JsonValue; +use std::collections::HashMap; +use std::sync::RwLock; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; // ============================================================================ // FIELD TYPES @@ -83,7 +83,7 @@ impl FieldType { } else { FieldType::Float64 } - } + }, JsonValue::String(_) => FieldType::String, JsonValue::Array(arr) => { if arr.is_empty() { @@ -94,7 +94,7 @@ impl FieldType { let elem_type = FieldType::infer(&arr[0]); FieldType::Array(Box::new(elem_type)) } - } + }, JsonValue::Object(_) => FieldType::Object, } } @@ -111,7 +111,7 @@ impl FieldType { (FieldType::Array(_), JsonValue::Array(_)) => true, (FieldType::Vector(dim), JsonValue::Array(arr)) => { arr.len() == *dim && arr.iter().all(|v| v.is_f64()) - } + }, (FieldType::Timestamp, JsonValue::Number(_)) => true, (FieldType::Timestamp, JsonValue::String(_)) => true, _ => false, @@ -230,7 +230,9 @@ impl Schema { Ok(()) } else { - Err(VecStoreError::InvalidInput("Document must be an object".to_string())) + Err(VecStoreError::InvalidInput( + "Document must be an object".to_string(), + )) } } @@ -268,7 +270,8 @@ impl SchemaBuilder { #[inline] pub fn add_field(mut self, name: &str, field_type: FieldType) -> Self { - self.fields.insert(name.to_string(), FieldDef::new(name, field_type)); + self.fields + .insert(name.to_string(), FieldDef::new(name, field_type)); self } @@ -371,7 +374,9 @@ impl SchemaEvolution { /// Apply a schema change pub fn apply(&mut self, change: SchemaChange) -> Result { let current_version = *self.versions.keys().max().unwrap_or(&0); - let mut new_schema = self.versions.get(¤t_version) + let mut new_schema = self + .versions + .get(¤t_version) .ok_or_else(|| VecStoreError::Internal("No current schema".to_string()))? .clone(); @@ -384,7 +389,7 @@ impl SchemaEvolution { ))); } new_schema.fields.insert(field.name.clone(), field.clone()); - } + }, SchemaChange::FieldRemoved { name } => { if !new_schema.fields.contains_key(name) { return Err(VecStoreError::InvalidInput(format!( @@ -393,7 +398,7 @@ impl SchemaEvolution { ))); } new_schema.fields.remove(name); - } + }, SchemaChange::FieldTypeChanged { name, from, to } => { if let Some(field) = new_schema.fields.get_mut(name) { if field.field_type != *from { @@ -410,27 +415,27 @@ impl SchemaEvolution { } field.field_type = to.clone(); } - } + }, SchemaChange::FieldNullable { name } => { if let Some(field) = new_schema.fields.get_mut(name) { field.nullable = true; } - } + }, SchemaChange::FieldRequired { name } => { if let Some(field) = new_schema.fields.get_mut(name) { field.nullable = false; } - } + }, SchemaChange::IndexAdded { field } => { if let Some(f) = new_schema.fields.get_mut(field) { f.indexed = true; } - } + }, SchemaChange::IndexRemoved { field } => { if let Some(f) = new_schema.fields.get_mut(field) { f.indexed = false; } - } + }, } new_schema.version = current_version + 1; @@ -488,14 +493,18 @@ impl DynamicCollection { pub fn insert(&self, doc: JsonValue) -> Result { // Validate against schema { - let schema = self.schema.read() + let schema = self + .schema + .read() .map_err(|_| VecStoreError::Internal("Schema lock poisoned".to_string()))?; schema.validate(&doc)?; } // Extract ID let id = { - let schema = self.schema.read() + let schema = self + .schema + .read() .map_err(|_| VecStoreError::Internal("Schema lock poisoned".to_string()))?; doc.get(&schema.primary_key) .and_then(|v| v.as_str()) @@ -505,22 +514,29 @@ impl DynamicCollection { // Extract vector let vector = { - let schema = self.schema.read() + let schema = self + .schema + .read() .map_err(|_| VecStoreError::Internal("Schema lock poisoned".to_string()))?; schema.extract_vector(&doc) }; // Discover dynamic fields if let JsonValue::Object(map) = &doc { - let schema = self.schema.read() + let schema = self + .schema + .read() .map_err(|_| VecStoreError::Internal("Schema lock poisoned".to_string()))?; - let mut dynamic = self.dynamic_fields.write() + let mut dynamic = self + .dynamic_fields + .write() .map_err(|_| VecStoreError::Internal("Dynamic fields lock poisoned".to_string()))?; for (key, value) in map { if !schema.fields.contains_key(key) { let inferred_type = FieldType::infer(value); - dynamic.entry(key.clone()) + dynamic + .entry(key.clone()) .and_modify(|existing| { // Promote type if needed if !existing.is_compatible(value) { @@ -533,13 +549,15 @@ impl DynamicCollection { } // Store document - self.documents.write() + self.documents + .write() .map_err(|_| VecStoreError::Internal("Documents lock poisoned".to_string()))? .insert(id.clone(), doc); // Store vector if let Some(vec) = vector { - self.vectors.write() + self.vectors + .write() .map_err(|_| VecStoreError::Internal("Vectors lock poisoned".to_string()))? .insert(id.clone(), vec); } @@ -549,13 +567,17 @@ impl DynamicCollection { /// Get document by ID pub fn get(&self, id: &str) -> Option { - let Ok(documents) = self.documents.read() else { return None; }; + let Ok(documents) = self.documents.read() else { + return None; + }; documents.get(id).cloned() } /// Delete document pub fn delete(&self, id: &str) -> bool { - let Ok(mut documents) = self.documents.write() else { return false; }; + let Ok(mut documents) = self.documents.write() else { + return false; + }; let doc_removed = documents.remove(id).is_some(); drop(documents); // Release lock before acquiring next if let Ok(mut vectors) = self.vectors.write() { @@ -567,10 +589,14 @@ impl DynamicCollection { /// Add a new field to schema pub fn add_field(&self, field: FieldDef) -> Result<()> { let change = SchemaChange::FieldAdded { field }; - let new_schema = self.evolution.write() + let new_schema = self + .evolution + .write() .map_err(|_| VecStoreError::Internal("Evolution lock poisoned".to_string()))? .apply(change)?; - *self.schema.write() + *self + .schema + .write() .map_err(|_| VecStoreError::Internal("Schema lock poisoned".to_string()))? = new_schema; Ok(()) } @@ -578,9 +604,13 @@ impl DynamicCollection { /// Promote dynamic field to schema pub fn promote_field(&self, name: &str) -> Result<()> { let field_type = { - let dynamic = self.dynamic_fields.read() + let dynamic = self + .dynamic_fields + .read() .map_err(|_| VecStoreError::Internal("Dynamic fields lock poisoned".to_string()))?; - dynamic.get(name).cloned() + dynamic + .get(name) + .cloned() .ok_or_else(|| VecStoreError::NotFound(format!("Dynamic field: {}", name)))? }; @@ -588,7 +618,8 @@ impl DynamicCollection { self.add_field(field)?; // Remove from dynamic fields - self.dynamic_fields.write() + self.dynamic_fields + .write() .map_err(|_| VecStoreError::Internal("Dynamic fields lock poisoned".to_string()))? .remove(name); @@ -596,20 +627,31 @@ impl DynamicCollection { } /// Search with filter - pub fn search(&self, query: &[f32], top_k: usize, filter: Option<&str>) -> Result> { - let vectors = self.vectors.read() + pub fn search( + &self, + query: &[f32], + top_k: usize, + filter: Option<&str>, + ) -> Result> { + let vectors = self + .vectors + .read() .map_err(|_| VecStoreError::Internal("Vectors lock poisoned".to_string()))?; - let documents = self.documents.read() + let documents = self + .documents + .read() .map_err(|_| VecStoreError::Internal("Documents lock poisoned".to_string()))?; - let mut results: Vec<_> = vectors.iter() + let mut results: Vec<_> = vectors + .iter() .filter_map(|(id, vec)| { // Apply filter if provided if let Some(filter_str) = filter && let Some(doc) = documents.get(id) - && !self.matches_filter(doc, filter_str) { - return None; - } + && !self.matches_filter(doc, filter_str) + { + return None; + } let score = cosine_similarity(query, vec); Some((id.clone(), score)) @@ -619,7 +661,8 @@ impl DynamicCollection { results.sort_by(|a, b| b.1.total_cmp(&a.1)); results.truncate(top_k); - Ok(results.into_iter() + Ok(results + .into_iter() .map(|(id, score)| SearchResult { id, score }) .collect()) } @@ -645,32 +688,32 @@ impl DynamicCollection { if let Some(n) = field_value.as_f64() { return n.to_string() == value; } - } + }, ">" => { if let (Some(v), Ok(threshold)) = (field_value.as_f64(), value.parse::()) { return v > threshold; } - } + }, "<" => { if let (Some(v), Ok(threshold)) = (field_value.as_f64(), value.parse::()) { return v < threshold; } - } + }, ">=" => { if let (Some(v), Ok(threshold)) = (field_value.as_f64(), value.parse::()) { return v >= threshold; } - } + }, "<=" => { if let (Some(v), Ok(threshold)) = (field_value.as_f64(), value.parse::()) { return v <= threshold; } - } + }, "!=" => { if let Some(s) = field_value.as_str() { return s != value; } - } + }, _ => return true, } } @@ -769,7 +812,10 @@ mod tests { assert_eq!(FieldType::infer(&json!(42)), FieldType::Int64); assert_eq!(FieldType::infer(&json!(3.14)), FieldType::Float64); assert_eq!(FieldType::infer(&json!(true)), FieldType::Bool); - assert_eq!(FieldType::infer(&json!([1.0, 2.0, 3.0])), FieldType::Vector(3)); + assert_eq!( + FieldType::infer(&json!([1.0, 2.0, 3.0])), + FieldType::Vector(3) + ); } #[test] @@ -783,12 +829,14 @@ mod tests { let collection = DynamicCollection::new(schema).unwrap(); // Insert with dynamic fields - collection.insert(json!({ - "id": "doc1", - "embedding": [1.0, 0.0, 0.0, 0.0], - "title": "Hello World", // Dynamic field - "score": 0.95, // Dynamic field - })).unwrap(); + collection + .insert(json!({ + "id": "doc1", + "embedding": [1.0, 0.0, 0.0, 0.0], + "title": "Hello World", // Dynamic field + "score": 0.95, // Dynamic field + })) + .unwrap(); // Check dynamic fields discovered let dynamic = collection.dynamic_fields(); @@ -796,11 +844,9 @@ mod tests { assert!(dynamic.contains_key("score")); // Search with filter - let results = collection.search( - &[1.0, 0.0, 0.0, 0.0], - 10, - Some("score > 0.9"), - ).unwrap(); + let results = collection + .search(&[1.0, 0.0, 0.0, 0.0], 10, Some("score > 0.9")) + .unwrap(); assert_eq!(results.len(), 1); assert_eq!(results[0].id, "doc1"); diff --git a/src/edge_runtime.rs b/src/edge_runtime.rs index 88b474b..222fa12 100644 --- a/src/edge_runtime.rs +++ b/src/edge_runtime.rs @@ -33,10 +33,10 @@ //! let state = store.to_bytes()?; //! ``` -use std::collections::HashMap; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; /// Edge runtime configuration #[derive(Debug, Clone, Serialize, Deserialize)] @@ -61,11 +61,21 @@ pub struct EdgeConfig { pub hnsw_ef: usize, } -fn default_max_vectors() -> usize { 10000 } -fn default_quant_bits() -> u8 { 8 } -fn default_true() -> bool { true } -fn default_m() -> usize { 8 } -fn default_ef() -> usize { 64 } +fn default_max_vectors() -> usize { + 10000 +} +fn default_quant_bits() -> u8 { + 8 +} +fn default_true() -> bool { + true +} +fn default_m() -> usize { + 8 +} +fn default_ef() -> usize { + 64 +} impl Default for EdgeConfig { fn default() -> Self { @@ -143,8 +153,8 @@ struct CompactVector { #[derive(Debug, Clone, Serialize, Deserialize)] enum VectorData { Full(Vec), - Quantized8(Vec, f32, f32), // data, min, scale - Quantized4(Vec, f32, f32), // data (packed), min, scale + Quantized8(Vec, f32, f32), // data, min, scale + Quantized4(Vec, f32, f32), // data (packed), min, scale } impl VectorData { @@ -153,10 +163,8 @@ impl VectorData { match self { VectorData::Full(v) => v.clone(), VectorData::Quantized8(data, min, scale) => { - data.iter() - .map(|&b| min + (b as f32) * scale) - .collect() - } + data.iter().map(|&b| min + (b as f32) * scale).collect() + }, VectorData::Quantized4(data, min, scale) => { let mut result = Vec::with_capacity(dimension); for byte in data { @@ -168,7 +176,7 @@ impl VectorData { } } result - } + }, } } @@ -225,9 +233,7 @@ impl EdgeVectorStore { VectorData::Full(vector) }; - let metadata_bytes = metadata.map(|m| { - serde_json::to_vec(&m).unwrap_or_default() - }); + let metadata_bytes = metadata.map(|m| serde_json::to_vec(&m).unwrap_or_default()); let index = self.vectors.len(); self.vectors.push(CompactVector { @@ -249,11 +255,12 @@ impl EdgeVectorStore { match self.config.quantization_bits { 8 => { let scale = if range > 0.0 { range / 255.0 } else { 1.0 }; - let data: Vec = vector.iter() + let data: Vec = vector + .iter() .map(|&v| ((v - min) / scale).round() as u8) .collect(); VectorData::Quantized8(data, min, scale) - } + }, 4 => { let scale = if range > 0.0 { range / 15.0 } else { 1.0 }; let mut data = Vec::with_capacity(vector.len().div_ceil(2)); @@ -267,7 +274,7 @@ impl EdgeVectorStore { data.push(lo | (hi << 4)); } VectorData::Quantized4(data, min, scale) - } + }, _ => VectorData::Full(vector.to_vec()), } } @@ -281,14 +288,17 @@ impl EdgeVectorStore { }); } - let mut results: Vec = self.vectors.iter() + let mut results: Vec = self + .vectors + .iter() .map(|v| { let vec = v.data.to_f32(self.dimension); let score = Self::cosine_similarity(query, &vec); - let metadata = v.metadata.as_ref().and_then(|bytes| { - serde_json::from_slice(bytes).ok() - }); + let metadata = v + .metadata + .as_ref() + .and_then(|bytes| serde_json::from_slice(bytes).ok()); EdgeSearchResult { id: v.id.clone(), @@ -308,9 +318,10 @@ impl EdgeVectorStore { pub fn get(&self, id: &str) -> Option { self.id_to_index.get(id).map(|&idx| { let v = &self.vectors[idx]; - let metadata = v.metadata.as_ref().and_then(|bytes| { - serde_json::from_slice(bytes).ok() - }); + let metadata = v + .metadata + .as_ref() + .and_then(|bytes| serde_json::from_slice(bytes).ok()); EdgeSearchResult { id: v.id.clone(), @@ -346,18 +357,17 @@ impl EdgeVectorStore { vectors: self.vectors.clone(), }; - bincode::serialize(&state).map_err(|e| { - VecStoreError::Serialization(e.to_string()) - }) + bincode::serialize(&state).map_err(|e| VecStoreError::Serialization(e.to_string())) } /// Load from bytes pub fn from_bytes(bytes: &[u8]) -> Result { - let state: EdgeState = bincode::deserialize(bytes).map_err(|e| { - VecStoreError::Serialization(e.to_string()) - })?; + let state: EdgeState = + bincode::deserialize(bytes).map_err(|e| VecStoreError::Serialization(e.to_string()))?; - let id_to_index: HashMap = state.vectors.iter() + let id_to_index: HashMap = state + .vectors + .iter() .enumerate() .map(|(i, v)| (v.id.clone(), i)) .collect(); @@ -405,18 +415,16 @@ impl EdgeVectorStore { /// Get memory usage estimate pub fn memory_usage(&self) -> MemoryUsage { - let vector_bytes: usize = self.vectors.iter() - .map(|v| v.data.memory_size()) - .sum(); + let vector_bytes: usize = self.vectors.iter().map(|v| v.data.memory_size()).sum(); - let metadata_bytes: usize = self.vectors.iter() + let metadata_bytes: usize = self + .vectors + .iter() .filter_map(|v| v.metadata.as_ref()) .map(|m| m.len()) .sum(); - let id_bytes: usize = self.vectors.iter() - .map(|v| v.id.len()) - .sum(); + let id_bytes: usize = self.vectors.iter().map(|v| v.id.len()).sum(); MemoryUsage { vector_bytes, @@ -475,17 +483,19 @@ pub mod wasm { #[wasm_bindgen] pub fn add(&mut self, id: &str, vector: &[f32]) -> Result<(), JsValue> { - self.inner.add(id, vector.to_vec(), None) + self.inner + .add(id, vector.to_vec(), None) .map_err(|e| JsValue::from_str(&e.to_string())) } #[wasm_bindgen] pub fn search(&self, query: &[f32], k: usize) -> Result { - let results = self.inner.search(query, k) + let results = self + .inner + .search(query, k) .map_err(|e| JsValue::from_str(&e.to_string()))?; - serde_wasm_bindgen::to_value(&results) - .map_err(|e| JsValue::from_str(&e.to_string())) + serde_wasm_bindgen::to_value(&results).map_err(|e| JsValue::from_str(&e.to_string())) } #[wasm_bindgen] diff --git a/src/embedded_onnx.rs b/src/embedded_onnx.rs index 6dab593..f401263 100644 --- a/src/embedded_onnx.rs +++ b/src/embedded_onnx.rs @@ -30,13 +30,13 @@ //! let results = store.search_text("greeting", 10)?; //! ``` +use serde::{Deserialize, Serialize}; use std::collections::HashMap; -use std::sync::RwLock; #[cfg(feature = "embeddings")] use std::sync::Arc; -use serde::{Deserialize, Serialize}; +use std::sync::RwLock; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; // ============================================================================ // MODEL CONFIGURATION @@ -115,7 +115,9 @@ impl ModelType { /// Get the default query prefix for this model pub fn default_query_prefix(&self) -> Option<&str> { match self { - ModelType::BgeSmallEnV15 | ModelType::BgeBaseEnV15 => Some("Represent this sentence for searching relevant passages: "), + ModelType::BgeSmallEnV15 | ModelType::BgeBaseEnV15 => { + Some("Represent this sentence for searching relevant passages: ") + }, ModelType::E5SmallV2 | ModelType::E5BaseV2 => Some("query: "), _ => None, } @@ -301,8 +303,6 @@ impl EmbeddingModel { /// Load with configuration pub fn with_config(config: ModelConfig) -> Result { - - let dimension = config.model_type.dimension(); // Get model directory @@ -393,11 +393,11 @@ impl EmbeddingModel { // Set optimization level based on quantization if config.quantize { - session_builder = session_builder - .with_optimization_level(GraphOptimizationLevel::Level3)?; + session_builder = + session_builder.with_optimization_level(GraphOptimizationLevel::Level3)?; } else { - session_builder = session_builder - .with_optimization_level(GraphOptimizationLevel::Level2)?; + session_builder = + session_builder.with_optimization_level(GraphOptimizationLevel::Level2)?; } // Set number of threads @@ -406,9 +406,8 @@ impl EmbeddingModel { // Add execution providers #[cfg(feature = "cuda")] if config.use_gpu { - session_builder = session_builder.with_execution_providers([ - ort::CUDAExecutionProvider::default().build(), - ])?; + session_builder = session_builder + .with_execution_providers([ort::CUDAExecutionProvider::default().build()])?; } // Load model @@ -427,19 +426,13 @@ impl EmbeddingModel { let cache_dir = home.home_dir().join(".vecstore").join("models"); let model_dir = cache_dir.join(model_type.huggingface_id().replace('/', "_")); - std::fs::create_dir_all(&model_dir) - .map_err(|e| VecStoreError::Io(e))?; + std::fs::create_dir_all(&model_dir).map_err(|e| VecStoreError::Io(e))?; Ok(model_dir) } /// Download model files from HuggingFace - fn download_model( - model_type: &ModelType, - model_dir: &std::path::Path, - ) -> Result<()> { - - + fn download_model(model_type: &ModelType, model_dir: &std::path::Path) -> Result<()> { let base_url = format!( "https://huggingface.co/{}/resolve/main", model_type.huggingface_id() @@ -457,16 +450,14 @@ impl EmbeddingModel { tracing::info!("Downloading {}...", remote_name); - let response = ureq::get(&url) - .call() - .map_err(|e| VecStoreError::Internal(format!("Failed to download {}: {}", url, e)))?; + let response = ureq::get(&url).call().map_err(|e| { + VecStoreError::Internal(format!("Failed to download {}: {}", url, e)) + })?; let mut reader = response.into_body().into_reader(); - let mut file = std::fs::File::create(&dest) - .map_err(|e| VecStoreError::Io(e))?; + let mut file = std::fs::File::create(&dest).map_err(|e| VecStoreError::Io(e))?; - std::io::copy(&mut reader, &mut file) - .map_err(|e| VecStoreError::Io(e))?; + std::io::copy(&mut reader, &mut file).map_err(|e| VecStoreError::Io(e))?; tracing::info!("Downloaded {} successfully", remote_name); } @@ -499,7 +490,9 @@ impl EmbeddingModel { /// Embed a single text pub fn embed(&self, text: &str) -> Result> { let embeddings = self.embed_batch(&[text.to_string()])?; - embeddings.into_iter().next() + embeddings + .into_iter() + .next() .ok_or_else(|| VecStoreError::InvalidInput("No embedding returned".to_string())) } @@ -607,16 +600,20 @@ impl EmbeddingModel { .map_err(|e| VecStoreError::Internal(format!("Array shape error: {}", e)))?; // Create ONNX input tensors using ort 2.0 API - let input_ids_tensor = Tensor::from_array(input_ids_array) - .map_err(|e| VecStoreError::OnnxRuntime(format!("Failed to create input_ids tensor: {}", e)))?; - let attention_mask_tensor = Tensor::from_array(attention_mask_array) - .map_err(|e| VecStoreError::OnnxRuntime(format!("Failed to create attention_mask tensor: {}", e)))?; - let token_type_ids_tensor = Tensor::from_array(token_type_ids_array) - .map_err(|e| VecStoreError::OnnxRuntime(format!("Failed to create token_type_ids tensor: {}", e)))?; + let input_ids_tensor = Tensor::from_array(input_ids_array).map_err(|e| { + VecStoreError::OnnxRuntime(format!("Failed to create input_ids tensor: {}", e)) + })?; + let attention_mask_tensor = Tensor::from_array(attention_mask_array).map_err(|e| { + VecStoreError::OnnxRuntime(format!("Failed to create attention_mask tensor: {}", e)) + })?; + let token_type_ids_tensor = Tensor::from_array(token_type_ids_array).map_err(|e| { + VecStoreError::OnnxRuntime(format!("Failed to create token_type_ids tensor: {}", e)) + })?; // Get mutable access to session for inference - let mut session = self.session.write() - .map_err(|e| VecStoreError::LockError(format!("Failed to acquire session lock: {}", e)))?; + let mut session = self.session.write().map_err(|e| { + VecStoreError::LockError(format!("Failed to acquire session lock: {}", e)) + })?; // Run inference using ort::inputs! macro let outputs = session @@ -642,13 +639,15 @@ impl EmbeddingModel { ndarray::ArrayD::from_shape_vec( shape.to_vec(), embeddings_view.iter().copied().collect(), - ).map_err(|e| VecStoreError::Internal(format!("Array conversion error: {}", e)))? + ) + .map_err(|e| VecStoreError::Internal(format!("Array conversion error: {}", e)))? } else if shape.len() == 2 { // Shape is (batch_size, hidden_size) - already pooled by model ndarray::ArrayD::from_shape_vec( shape.to_vec(), embeddings_view.iter().copied().collect(), - ).map_err(|e| VecStoreError::Internal(format!("Array conversion error: {}", e)))? + ) + .map_err(|e| VecStoreError::Internal(format!("Array conversion error: {}", e)))? } else { return Err(VecStoreError::OnnxRuntime(format!( "Unexpected output shape: {:?}", @@ -751,7 +750,11 @@ impl EmbeddingModel { /// Get statistics pub fn stats(&self) -> ModelStats { - self.stats.read().ok().map(|g| g.clone()).unwrap_or_default() + self.stats + .read() + .ok() + .map(|g| g.clone()) + .unwrap_or_default() } } @@ -789,15 +792,64 @@ impl EmbeddingModel { fn build_simple_vocab() -> HashMap { // Fallback vocabulary - enable `embeddings` feature for real BPE/WordPiece tokenization let words = [ - "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", - "have", "has", "had", "do", "does", "did", "will", "would", "could", - "should", "may", "might", "must", "can", "this", "that", "these", - "those", "i", "you", "he", "she", "it", "we", "they", "what", "which", - "who", "whom", "where", "when", "why", "how", "hello", "world", "test", - "document", "search", "query", "vector", "embedding", "model", "text", + "the", + "a", + "an", + "is", + "are", + "was", + "were", + "be", + "been", + "being", + "have", + "has", + "had", + "do", + "does", + "did", + "will", + "would", + "could", + "should", + "may", + "might", + "must", + "can", + "this", + "that", + "these", + "those", + "i", + "you", + "he", + "she", + "it", + "we", + "they", + "what", + "which", + "who", + "whom", + "where", + "when", + "why", + "how", + "hello", + "world", + "test", + "document", + "search", + "query", + "vector", + "embedding", + "model", + "text", ]; - words.iter().enumerate() + words + .iter() + .enumerate() .map(|(i, &w)| (w.to_string(), i)) .collect() } @@ -810,7 +862,9 @@ impl EmbeddingModel { /// Embed a single text pub fn embed(&self, text: &str) -> Result> { let embeddings = self.embed_batch(&[text.to_string()])?; - embeddings.into_iter().next() + embeddings + .into_iter() + .next() .ok_or_else(|| VecStoreError::InvalidInput("No embedding returned".to_string())) } @@ -818,9 +872,7 @@ impl EmbeddingModel { pub fn embed_batch(&self, texts: &[String]) -> Result>> { let start = std::time::Instant::now(); - let embeddings: Vec> = texts.iter() - .map(|text| self.embed_text(text)) - .collect(); + let embeddings: Vec> = texts.iter().map(|text| self.embed_text(text)).collect(); // Update stats { @@ -870,12 +922,11 @@ impl EmbeddingModel { .split_whitespace() .map(|word| { // Clean word - let clean: String = word.chars() - .filter(|c| c.is_alphanumeric()) - .collect(); + let clean: String = word.chars().filter(|c| c.is_alphanumeric()).collect(); // Look up in vocabulary or use hash - self.vocab.get(&clean) + self.vocab + .get(&clean) .copied() .unwrap_or_else(|| self.hash_word(&clean)) }) @@ -894,7 +945,9 @@ impl EmbeddingModel { /// Get statistics pub fn stats(&self) -> ModelStats { - let Ok(guard) = self.stats.read() else { return ModelStats::default(); }; + let Ok(guard) = self.stats.read() else { + return ModelStats::default(); + }; guard.clone() } } @@ -980,7 +1033,7 @@ impl EmbeddedStore { stats.cache_hits += 1; } e - } + }, None => { #[cfg(feature = "embeddings")] let e = self.model.embed_document(text)?; @@ -997,7 +1050,7 @@ impl EmbeddedStore { } e - } + }, }; // Store vector @@ -1086,7 +1139,9 @@ impl EmbeddedStore { /// Delete by ID pub fn delete(&self, id: &str) -> bool { - let Ok(mut vectors) = self.vectors.write() else { return false; }; + let Ok(mut vectors) = self.vectors.write() else { + return false; + }; let removed = vectors.remove(id).is_some(); drop(vectors); @@ -1123,7 +1178,8 @@ impl EmbeddedStore { let texts = self.texts.read()?; let metadata = self.metadata.read()?; - let mut results: Vec<_> = vectors.iter() + let mut results: Vec<_> = vectors + .iter() .map(|(id, vec)| { let score = cosine_similarity(query, vec); SearchResult { @@ -1143,7 +1199,9 @@ impl EmbeddedStore { /// Get by ID pub fn get(&self, id: &str) -> Option<(Vec, Option, Option)> { - let Ok(vectors) = self.vectors.read() else { return None; }; + let Ok(vectors) = self.vectors.read() else { + return None; + }; vectors.get(id).map(|v| { let text = self.texts.read().ok().and_then(|t| t.get(id).cloned()); let meta = self.metadata.read().ok().and_then(|m| m.get(id).cloned()); @@ -1158,13 +1216,17 @@ impl EmbeddedStore { /// Get count pub fn len(&self) -> usize { - let Ok(guard) = self.vectors.read() else { return 0; }; + let Ok(guard) = self.vectors.read() else { + return 0; + }; guard.len() } /// Check if empty pub fn is_empty(&self) -> bool { - let Ok(guard) = self.vectors.read() else { return true; }; + let Ok(guard) = self.vectors.read() else { + return true; + }; guard.is_empty() } @@ -1230,7 +1292,9 @@ pub trait EmbeddingFunction: Send + Sync { /// Embed a single text fn embed_one(&self, text: &str) -> Result> { let embeddings = self.embed(&[text.to_string()])?; - embeddings.into_iter().next() + embeddings + .into_iter() + .next() .ok_or_else(|| VecStoreError::InvalidInput("No embedding returned".to_string())) } diff --git a/src/embedding_vcs.rs b/src/embedding_vcs.rs index a01ba91..567adad 100644 --- a/src/embedding_vcs.rs +++ b/src/embedding_vcs.rs @@ -40,7 +40,7 @@ //! vcs.checkout(&commit_id)?; //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; use std::path::PathBuf; @@ -430,7 +430,10 @@ impl EmbeddingVCS { use std::hash::{Hash, Hasher}; let mut hasher = std::collections::hash_map::DefaultHasher::new(); message.hash(&mut hasher); - chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0).hash(&mut hasher); + chrono::Utc::now() + .timestamp_nanos_opt() + .unwrap_or(0) + .hash(&mut hasher); format!("{:016x}", hasher.finish()) } diff --git a/src/embeddings.rs b/src/embeddings.rs index 995521a..1d4e2d8 100644 --- a/src/embeddings.rs +++ b/src/embeddings.rs @@ -47,12 +47,12 @@ pub use candle_backend::{CandleEmbedder, CandleModel}; use anyhow::Result; #[cfg(feature = "embeddings")] -use anyhow::{anyhow, Context}; +use anyhow::{Context, anyhow}; #[cfg(feature = "embeddings")] use ndarray::Array2; #[cfg(feature = "embeddings")] -use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::session::{Session, builder::GraphOptimizationLevel}; #[cfg(feature = "embeddings")] use std::path::{Path, PathBuf}; #[cfg(feature = "embeddings")] @@ -315,7 +315,9 @@ impl Embedder { // Run inference using ort 2.0 inputs! macro and extract results // We need to extract data while holding the lock, then release it let embeddings_array = { - let mut session_guard = self.session.lock() + let mut session_guard = self + .session + .lock() .map_err(|e| anyhow!("Failed to lock session: {}", e))?; let outputs = session_guard .run(ort::inputs![ diff --git a/src/embeddings/additional_providers.rs b/src/embeddings/additional_providers.rs index dfa8c8e..745e016 100644 --- a/src/embeddings/additional_providers.rs +++ b/src/embeddings/additional_providers.rs @@ -3,7 +3,7 @@ // Azure OpenAI, HuggingFace Inference API, Jina AI use super::TextEmbedder; -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde::{Deserialize, Serialize}; use std::time::Duration; @@ -136,13 +136,13 @@ impl AzureEmbedding { data.sort_by_key(|d| d.index); return Ok(data.into_iter().map(|d| d.embedding).collect()); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -150,16 +150,16 @@ impl AzureEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("Azure API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call Azure API: {}", e)); - } + }, } } } @@ -273,19 +273,19 @@ impl HuggingFaceEmbedding { .context("Failed to parse HuggingFace response")?; return Ok(embeddings); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) if resp.status().as_u16() == 503 && retries < self.max_retries => { // Model loading, retry after longer wait retries += 1; tokio::time::sleep(Duration::from_secs(10)).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -293,16 +293,16 @@ impl HuggingFaceEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("HuggingFace API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call HuggingFace API: {}", e)); - } + }, } } } @@ -442,13 +442,13 @@ impl JinaEmbedding { data.sort_by_key(|d| d.index); return Ok(data.into_iter().map(|d| d.embedding).collect()); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -456,16 +456,16 @@ impl JinaEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("Jina API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call Jina API: {}", e)); - } + }, } } } diff --git a/src/embeddings/candle_backend.rs b/src/embeddings/candle_backend.rs index 468fddb..af6594b 100644 --- a/src/embeddings/candle_backend.rs +++ b/src/embeddings/candle_backend.rs @@ -4,11 +4,11 @@ //! enabling fully native Rust embeddings without Python dependencies. use crate::embeddings::TextEmbedder; -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use candle_core::{Device, Tensor}; use candle_nn::VarBuilder; use candle_transformers::models::bert::{BertModel, Config, DTYPE}; -use hf_hub::{api::sync::Api, Repo, RepoType}; +use hf_hub::{Repo, RepoType, api::sync::Api}; use tokenizers::Tokenizer; /// Supported Candle embedding models @@ -225,7 +225,11 @@ impl TextEmbedder for CandleEmbedder { .collect::>>()?; // Find max sequence length for padding - let max_len = encodings.iter().map(|e| e.get_ids().len()).max().unwrap_or(0); + let max_len = encodings + .iter() + .map(|e| e.get_ids().len()) + .max() + .unwrap_or(0); // Pad sequences and create batch tensors let batch_size = encodings.len(); diff --git a/src/embeddings/cloud_providers.rs b/src/embeddings/cloud_providers.rs index 40bf7de..7e6fb8e 100644 --- a/src/embeddings/cloud_providers.rs +++ b/src/embeddings/cloud_providers.rs @@ -13,7 +13,7 @@ // Each provider implements the TextEmbedder trait for seamless integration. use super::TextEmbedder; -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use std::time::Duration; @@ -82,7 +82,9 @@ impl RateLimiter { } async fn wait_if_needed(&self) { - let Ok(mut requests) = self.last_requests.lock() else { return; }; + let Ok(mut requests) = self.last_requests.lock() else { + return; + }; // Remove requests older than 1 minute let now = std::time::Instant::now(); @@ -98,7 +100,9 @@ impl RateLimiter { if wait_time > Duration::from_secs(0) { drop(requests); // Release lock before sleeping tokio::time::sleep(wait_time).await; - let Ok(reacquired) = self.last_requests.lock() else { return; }; + let Ok(reacquired) = self.last_requests.lock() else { + return; + }; requests = reacquired; } } @@ -214,13 +218,13 @@ impl CohereEmbedding { .context("Failed to parse Cohere response")?; return Ok(embed_response.embeddings); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -228,16 +232,16 @@ impl CohereEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("Cohere API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call Cohere API: {}", e)); - } + }, } } } @@ -428,13 +432,13 @@ impl VoyageEmbedding { data.sort_by_key(|d| d.index); return Ok(data.into_iter().map(|d| d.embedding).collect()); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -442,16 +446,16 @@ impl VoyageEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("Voyage API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call Voyage API: {}", e)); - } + }, } } } @@ -630,13 +634,13 @@ impl MistralEmbedding { data.sort_by_key(|d| d.index); return Ok(data.into_iter().map(|d| d.embedding).collect()); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -644,16 +648,16 @@ impl MistralEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("Mistral API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call Mistral API: {}", e)); - } + }, } } } @@ -701,7 +705,7 @@ impl GoogleModel { GoogleModel::TextEmbedding004 => "text-embedding-004", GoogleModel::TextEmbeddingGeckoMultilingual => { "textembedding-gecko-multilingual@latest" - } + }, } } @@ -794,7 +798,10 @@ impl GoogleEmbedding { let url = format!( "https://{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:predict", - self.location, self.project_id, self.location, self.model.as_str() + self.location, + self.project_id, + self.location, + self.model.as_str() ); let mut retries = 0; @@ -820,13 +827,13 @@ impl GoogleEmbedding { .into_iter() .map(|p| p.embeddings.values) .collect()); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -834,16 +841,16 @@ impl GoogleEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("Google API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call Google API: {}", e)); - } + }, } } } diff --git a/src/embeddings/ollama_backend.rs b/src/embeddings/ollama_backend.rs index 31a2387..9dd7461 100644 --- a/src/embeddings/ollama_backend.rs +++ b/src/embeddings/ollama_backend.rs @@ -10,7 +10,7 @@ //! - Any Ollama embedding model use super::TextEmbedder; -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde::{Deserialize, Serialize}; use std::time::Duration; diff --git a/src/embeddings/openai_backend.rs b/src/embeddings/openai_backend.rs index 2bfa4c2..34fa1a5 100644 --- a/src/embeddings/openai_backend.rs +++ b/src/embeddings/openai_backend.rs @@ -4,7 +4,7 @@ // Supports text-embedding-3-small, text-embedding-3-large, and ada-002 models. use super::TextEmbedder; -use anyhow::{anyhow, Context, Result}; +use anyhow::{Context, Result, anyhow}; use serde::{Deserialize, Serialize}; use std::sync::Arc; use std::time::Duration; @@ -92,7 +92,9 @@ impl RateLimiter { } async fn wait_if_needed(&self) { - let Ok(mut requests) = self.last_requests.lock() else { return; }; + let Ok(mut requests) = self.last_requests.lock() else { + return; + }; // Remove requests older than 1 minute let now = std::time::Instant::now(); @@ -108,7 +110,9 @@ impl RateLimiter { if wait_time > Duration::from_secs(0) { drop(requests); // Release lock before sleeping tokio::time::sleep(wait_time).await; - let Ok(reacquired) = self.last_requests.lock() else { return; }; + let Ok(reacquired) = self.last_requests.lock() else { + return; + }; requests = reacquired; } } @@ -312,14 +316,14 @@ impl OpenAIEmbedding { data.sort_by_key(|d| d.index); return Ok(data.into_iter().map(|d| d.embedding).collect()); - } + }, Ok(resp) if resp.status().as_u16() == 429 && retries < self.max_retries => { // Rate limit hit, retry with exponential backoff retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Ok(resp) => { let status = resp.status(); let body = resp @@ -327,17 +331,17 @@ impl OpenAIEmbedding { .await .unwrap_or_else(|_| String::from("(no body)")); return Err(anyhow!("OpenAI API error {}: {}", status, body)); - } + }, Err(_e) if retries < self.max_retries => { // Network error, retry retries += 1; let wait_time = Duration::from_secs(2_u64.pow(retries as u32)); tokio::time::sleep(wait_time).await; continue; - } + }, Err(e) => { return Err(anyhow!("Failed to call OpenAI API: {}", e)); - } + }, } } } diff --git a/src/enterprise_protect.rs b/src/enterprise_protect.rs index 0cadd2f..5431254 100644 --- a/src/enterprise_protect.rs +++ b/src/enterprise_protect.rs @@ -46,29 +46,15 @@ use crate::error::{Result, VecStoreError}; #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] pub enum KMSProvider { /// AWS Key Management Service - AWSKMS { - key_arn: String, - region: String, - }, + AWSKMS { key_arn: String, region: String }, /// Google Cloud KMS - GCPKMS { - key_name: String, - project: String, - }, + GCPKMS { key_name: String, project: String }, /// Azure Key Vault - AzureKeyVault { - vault_url: String, - key_name: String, - }, + AzureKeyVault { vault_url: String, key_name: String }, /// HashiCorp Vault - HashiCorpVault { - address: String, - path: String, - }, + HashiCorpVault { address: String, path: String }, /// Local key (for development/testing) - Local { - key_path: PathBuf, - }, + Local { key_path: PathBuf }, } /// CMEK configuration @@ -178,7 +164,9 @@ impl EncryptionManager { pub fn initialize(&self) -> Result<()> { let dek = self.generate_dek()?; - let mut current = self.current_dek.write() + let mut current = self + .current_dek + .write() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; *current = Some(dek); @@ -200,7 +188,7 @@ impl EncryptionManager { ciphertext, key_id: dek.key_id.clone(), algorithm: "XOR".to_string(), // Use "AES-256-GCM" in production - nonce: vec![0u8; 12], // Would be random in production + nonce: vec![0u8; 12], // Would be random in production }) } @@ -209,7 +197,8 @@ impl EncryptionManager { let dek = self.get_dek(&encrypted.key_id)?; // Simple XOR decryption - let plaintext: Vec = encrypted.ciphertext + let plaintext: Vec = encrypted + .ciphertext .iter() .enumerate() .map(|(i, &b)| b ^ dek.key[i % dek.key.len()]) @@ -220,21 +209,27 @@ impl EncryptionManager { /// Rotate encryption key pub fn rotate_key(&self) -> Result { - let old_dek = self.current_dek.read() + let old_dek = self + .current_dek + .read() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? .clone(); let new_dek = self.generate_dek()?; // Store old key in cache for decryption if let Some(ref old) = old_dek { - let mut cache = self.dek_cache.write() + let mut cache = self + .dek_cache + .write() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; cache.insert(old.key_id.clone(), old.clone()); } // Set new key as current { - let mut current = self.current_dek.write() + let mut current = self + .current_dek + .write() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; *current = Some(new_dek.clone()); } @@ -249,7 +244,9 @@ impl EncryptionManager { fn generate_dek(&self) -> Result { // In production, this would call the KMS to generate a DEK let counter = self.key_counter.fetch_add(1, Ordering::Relaxed); - let key: Vec = (0..32).map(|i| ((i * 7 + 13 + counter as usize) % 256) as u8).collect(); + let key: Vec = (0..32) + .map(|i| ((i * 7 + 13 + counter as usize) % 256) as u8) + .collect(); let key_id = format!("dek_{}_{}", unix_timestamp(), counter); Ok(DataEncryptionKey { @@ -266,26 +263,35 @@ impl EncryptionManager { } fn get_current_dek(&self) -> Result { - let dek = self.current_dek.read() + let dek = self + .current_dek + .read() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; - dek.clone().ok_or_else(|| VecStoreError::EncryptionError("DEK not initialized".to_string())) + dek.clone() + .ok_or_else(|| VecStoreError::EncryptionError("DEK not initialized".to_string())) } fn get_dek(&self, key_id: &str) -> Result { // Check current key { - let current = self.current_dek.read() + let current = self + .current_dek + .read() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; if let Some(dek) = &*current - && dek.key_id == key_id { - return Ok(dek.clone()); - } + && dek.key_id == key_id + { + return Ok(dek.clone()); + } } // Check cache - let cache = self.dek_cache.read() + let cache = self + .dek_cache + .read() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; - cache.get(key_id) + cache + .get(key_id) .cloned() .ok_or_else(|| VecStoreError::EncryptionError(format!("DEK not found: {}", key_id))) } @@ -438,7 +444,9 @@ impl BackupManager { // Store backup metadata { - let mut backups = self.backups.write() + let mut backups = self + .backups + .write() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; backups.push(metadata.clone()); @@ -453,11 +461,14 @@ impl BackupManager { /// List backups pub fn list_backups(&self, collection: Option<&str>) -> Result> { - let backups = self.backups.read() + let backups = self + .backups + .read() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; if let Some(col) = collection { - Ok(backups.iter() + Ok(backups + .iter() .filter(|b| b.collection == col) .cloned() .collect()) @@ -492,7 +503,9 @@ impl BackupManager { /// Delete backup pub fn delete_backup(&self, backup_id: &str) -> Result { - let mut backups = self.backups.write() + let mut backups = self + .backups + .write() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let len_before = backups.len(); backups.retain(|b| b.backup_id != backup_id); @@ -501,7 +514,9 @@ impl BackupManager { /// Get backup statistics pub fn stats(&self) -> Result { - let backups = self.backups.read() + let backups = self + .backups + .read() .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let total_size: u64 = backups.iter().map(|b| b.size_bytes).sum(); @@ -765,9 +780,21 @@ impl ComplianceChecker { /// Get compliance summary pub fn summary(&self) -> ComplianceSummary { - let passed = self.checks.iter().filter(|c| c.status == ComplianceStatus::Passed).count(); - let failed = self.checks.iter().filter(|c| c.status == ComplianceStatus::Failed).count(); - let warnings = self.checks.iter().filter(|c| c.status == ComplianceStatus::Warning).count(); + let passed = self + .checks + .iter() + .filter(|c| c.status == ComplianceStatus::Passed) + .count(); + let failed = self + .checks + .iter() + .filter(|c| c.status == ComplianceStatus::Failed) + .count(); + let warnings = self + .checks + .iter() + .filter(|c| c.status == ComplianceStatus::Warning) + .count(); ComplianceSummary { total_checks: self.checks.len(), @@ -790,8 +817,7 @@ pub struct ComplianceSummary { } /// Enterprise configuration -#[derive(Debug, Clone, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Serialize, Deserialize, Default)] pub struct EnterpriseConfig { pub encryption_enabled: bool, pub audit_logging: bool, @@ -800,7 +826,6 @@ pub struct EnterpriseConfig { pub retention_policy: Option, } - fn unix_timestamp() -> i64 { SystemTime::now() .duration_since(UNIX_EPOCH) @@ -898,7 +923,10 @@ mod tests { let aws = CMEKConfig::aws_kms("arn:aws:kms:us-east-1:123456789:key/abc"); assert!(matches!(aws.provider, KMSProvider::AWSKMS { .. })); - let gcp = CMEKConfig::gcp_kms("projects/my-project/locations/global/keyRings/my-ring/cryptoKeys/my-key", "my-project"); + let gcp = CMEKConfig::gcp_kms( + "projects/my-project/locations/global/keyRings/my-ring/cryptoKeys/my-key", + "my-project", + ); assert!(matches!(gcp.provider, KMSProvider::GCPKMS { .. })); let azure = CMEKConfig::azure_keyvault("https://myvault.vault.azure.net", "my-key"); diff --git a/src/explainable.rs b/src/explainable.rs index aae3461..0c8eb74 100644 --- a/src/explainable.rs +++ b/src/explainable.rs @@ -41,8 +41,7 @@ use serde::{Deserialize, Serialize}; use std::collections::HashMap; /// Level of explanation detail -#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Default)] pub enum ExplainLevel { /// Minimal - just the score Minimal, @@ -57,7 +56,6 @@ pub enum ExplainLevel { Debug, } - /// Dimension contribution to similarity #[derive(Debug, Clone, Serialize, Deserialize)] pub struct DimensionContribution { @@ -302,7 +300,7 @@ impl ExplainableSearch { } else { 0.0 } - } + }, DistanceMetric::Euclidean => { let dist: f32 = a .iter() @@ -311,7 +309,7 @@ impl ExplainableSearch { .sum::() .sqrt(); 1.0 / (1.0 + dist) // Convert to similarity - } + }, DistanceMetric::DotProduct => a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(), } } @@ -355,30 +353,28 @@ impl ExplainableSearch { } }) .collect() - } - DistanceMetric::DotProduct => { - query - .iter() - .zip(result.iter()) - .enumerate() - .map(|(i, (q, r))| { - let contribution = q * r; - let percentage = if total_similarity.abs() > 1e-9 { - contribution / total_similarity * 100.0 - } else { - 0.0 - }; - DimensionContribution { - index: i, - contribution, - percentage, - query_value: *q, - result_value: *r, - label: self.get_dimension_label(i), - } - }) - .collect() - } + }, + DistanceMetric::DotProduct => query + .iter() + .zip(result.iter()) + .enumerate() + .map(|(i, (q, r))| { + let contribution = q * r; + let percentage = if total_similarity.abs() > 1e-9 { + contribution / total_similarity * 100.0 + } else { + 0.0 + }; + DimensionContribution { + index: i, + contribution, + percentage, + query_value: *q, + result_value: *r, + label: self.get_dimension_label(i), + } + }) + .collect(), DistanceMetric::Euclidean => { // For Euclidean, contribution is negative of squared difference let total_dist_sq: f32 = query @@ -409,7 +405,7 @@ impl ExplainableSearch { } }) .collect() - } + }, } } @@ -426,10 +422,7 @@ impl ExplainableSearch { bottom: &[DimensionContribution], score: f32, ) -> String { - let mut summary = format!( - "Similarity score: {:.4} ({:?}). ", - score, self.metric - ); + let mut summary = format!("Similarity score: {:.4} ({:?}). ", score, self.metric); if !top.is_empty() { let top_dims: Vec = top @@ -507,10 +500,7 @@ impl ExplainableSearch { .collect(); let summary = if !shared_themes.is_empty() { - format!( - "Vectors match on themes: {}", - shared_themes.join(", ") - ) + format!("Vectors match on themes: {}", shared_themes.join(", ")) } else if !matching_features.is_empty() { format!("Similarity driven by: {}", matching_features[0]) } else { @@ -525,12 +515,7 @@ impl ExplainableSearch { } } - fn compute_confidence( - &self, - query: &[f32], - result: &[f32], - score: f32, - ) -> ConfidenceInfo { + fn compute_confidence(&self, query: &[f32], result: &[f32], score: f32) -> ConfidenceInfo { // Estimate confidence based on vector properties let dim = query.len() as f32; @@ -571,7 +556,8 @@ impl ExplainableSearch { .iter() .filter(|c| c.contribution < 0.0) .collect(); - negative_dims.sort_by(|a, b| OrderedFloat(a.contribution).cmp(&OrderedFloat(b.contribution))); + negative_dims + .sort_by(|a, b| OrderedFloat(a.contribution).cmp(&OrderedFloat(b.contribution))); if let Some(worst) = negative_dims.first() { let hypothetical: Vec = result @@ -648,15 +634,29 @@ impl ExplainableSearch { info.insert("query_norm".to_string(), format!("{:.6}", query_norm)); info.insert("result_norm".to_string(), format!("{:.6}", result_norm)); - let positive_contribs: usize = contributions.iter().filter(|c| c.contribution > 0.0).count(); - let negative_contribs: usize = contributions.iter().filter(|c| c.contribution < 0.0).count(); + let positive_contribs: usize = contributions + .iter() + .filter(|c| c.contribution > 0.0) + .count(); + let negative_contribs: usize = contributions + .iter() + .filter(|c| c.contribution < 0.0) + .count(); info.insert("positive_dims".to_string(), positive_contribs.to_string()); info.insert("negative_dims".to_string(), negative_contribs.to_string()); - let sparsity_q = query.iter().filter(|&&x| x.abs() < 1e-6).count() as f32 / query.len() as f32; - let sparsity_r = result.iter().filter(|&&x| x.abs() < 1e-6).count() as f32 / result.len() as f32; - info.insert("query_sparsity".to_string(), format!("{:.2}%", sparsity_q * 100.0)); - info.insert("result_sparsity".to_string(), format!("{:.2}%", sparsity_r * 100.0)); + let sparsity_q = + query.iter().filter(|&&x| x.abs() < 1e-6).count() as f32 / query.len() as f32; + let sparsity_r = + result.iter().filter(|&&x| x.abs() < 1e-6).count() as f32 / result.len() as f32; + info.insert( + "query_sparsity".to_string(), + format!("{:.2}%", sparsity_q * 100.0), + ); + info.insert( + "result_sparsity".to_string(), + format!("{:.2}%", sparsity_r * 100.0), + ); info } diff --git a/src/federation.rs b/src/federation.rs index 791acb4..3e2a7d8 100644 --- a/src/federation.rs +++ b/src/federation.rs @@ -224,29 +224,34 @@ impl Federation { avg_latency_ms: RwLock::new(0.0), }); - let mut members = self.members.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on members".into()))?; + let mut members = self.members.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on members".into()) + })?; members.insert(member.id.clone(), state); Ok(()) } /// Remove a member from the federation pub fn remove_member(&self, member_id: &str) -> Result<()> { - let mut members = self.members.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on members".into()))?; - members.remove(member_id) + let mut members = self.members.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on members".into()) + })?; + members + .remove(member_id) .ok_or_else(|| VecStoreError::NotFound(format!("Member {} not found", member_id)))?; Ok(()) } /// Update member configuration pub fn update_member(&self, member: FederationMember) -> Result<()> { - let mut members = self.members.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on members".into()))?; + let mut members = self.members.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on members".into()) + })?; if let Some(state) = members.get_mut(&member.id) { // Create new state preserving runtime info - let avg_latency = *state.avg_latency_ms.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on avg_latency_ms".into()))?; + let avg_latency = *state.avg_latency_ms.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on avg_latency_ms".into()) + })?; let new_state = Arc::new(MemberState { member, health: state.health.clone(), @@ -259,7 +264,10 @@ impl Federation { members.insert(new_state.member.id.clone(), new_state); Ok(()) } else { - Err(VecStoreError::NotFound(format!("Member {} not found", member.id))) + Err(VecStoreError::NotFound(format!( + "Member {} not found", + member.id + ))) } } @@ -310,10 +318,10 @@ impl Federation { Ok(results) => { all_results.extend(results); members_succeeded += 1; - } + }, Err(_) => { member_state.total_errors.fetch_add(1, Ordering::Relaxed); - } + }, } member_state.active_queries.fetch_sub(1, Ordering::Relaxed); @@ -345,8 +353,9 @@ impl Federation { } fn get_eligible_members(&self, query: &FederatedQuery) -> Result>> { - let members = self.members.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on members".into()))?; + let members = self.members.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on members".into()) + })?; let mut eligible = Vec::new(); for state in members.values() { @@ -365,7 +374,9 @@ impl Federation { // Filter by collection if specified if !query.collections.is_empty() { - let has_collection = query.collections.iter() + let has_collection = query + .collections + .iter() .any(|c| state.member.collections.contains(c)); if !has_collection { continue; @@ -381,13 +392,20 @@ impl Federation { Ok(eligible) } - fn query_member(&self, state: &MemberState, _query: &FederatedQuery) -> Result> { + fn query_member( + &self, + state: &MemberState, + query: &FederatedQuery, + ) -> Result> { + // Some feature combinations compile only the local-member branch. + let _ = query; + match &state.member.member_type { MemberType::Local => { // Local queries should be handled by the local VecStore directly // Return empty results as local stores are queried separately Ok(Vec::new()) - } + }, MemberType::RemoteHttp => { #[cfg(feature = "openai-embeddings")] { @@ -400,7 +418,7 @@ impl Federation { "HTTP federation requires 'openai-embeddings' feature (which includes reqwest)".into() )) } - } + }, MemberType::RemoteGrpc => { #[cfg(feature = "server")] { @@ -410,10 +428,10 @@ impl Federation { #[cfg(not(feature = "server"))] { Err(VecStoreError::InvalidConfig( - "gRPC federation requires 'server' feature (which includes tonic)".into() + "gRPC federation requires 'server' feature (which includes tonic)".into(), )) } - } + }, MemberType::External(service_type) => { #[cfg(feature = "openai-embeddings")] { @@ -427,7 +445,7 @@ impl Federation { "External federation requires 'openai-embeddings' feature (which includes reqwest)".into() )) } - } + }, } } @@ -487,7 +505,9 @@ impl Federation { score: item.get("score")?.as_f64()? as f32, vector: item.get("vector").and_then(|v| { v.as_array().map(|arr| { - arr.iter().filter_map(|x| x.as_f64().map(|n| n as f32)).collect() + arr.iter() + .filter_map(|x| x.as_f64().map(|n| n as f32)) + .collect() }) }), metadata: item.get("metadata").cloned(), @@ -559,9 +579,9 @@ impl Federation { ))); } - let body: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse Pinecone response: {}", e)))?; + let body: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse Pinecone response: {}", e)) + })?; let results = body .get("matches") @@ -574,7 +594,9 @@ impl Federation { score: item.get("score")?.as_f64()? as f32, vector: item.get("values").and_then(|v| { v.as_array().map(|arr| { - arr.iter().filter_map(|x| x.as_f64().map(|n| n as f32)).collect() + arr.iter() + .filter_map(|x| x.as_f64().map(|n| n as f32)) + .collect() }) }), metadata: item.get("metadata").cloned(), @@ -598,7 +620,11 @@ impl Federation { query: &FederatedQuery, k: usize, ) -> Result> { - let collection = member.collections.first().cloned().unwrap_or_else(|| "default".to_string()); + let collection = member + .collections + .first() + .cloned() + .unwrap_or_else(|| "default".to_string()); let request_body = serde_json::json!({ "vector": query.vector, "limit": k, @@ -625,9 +651,9 @@ impl Federation { ))); } - let body: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse Qdrant response: {}", e)))?; + let body: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse Qdrant response: {}", e)) + })?; let results = body .get("result") @@ -640,7 +666,9 @@ impl Federation { score: item.get("score")?.as_f64()? as f32, vector: item.get("vector").and_then(|v| { v.as_array().map(|arr| { - arr.iter().filter_map(|x| x.as_f64().map(|n| n as f32)).collect() + arr.iter() + .filter_map(|x| x.as_f64().map(|n| n as f32)) + .collect() }) }), metadata: item.get("payload").cloned(), @@ -664,7 +692,11 @@ impl Federation { query: &FederatedQuery, k: usize, ) -> Result> { - let class_name = member.collections.first().cloned().unwrap_or_else(|| "Document".to_string()); + let class_name = member + .collections + .first() + .cloned() + .unwrap_or_else(|| "Document".to_string()); // Weaviate uses GraphQL let graphql_query = format!( @@ -701,9 +733,9 @@ impl Federation { ))); } - let body: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse Weaviate response: {}", e)))?; + let body: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse Weaviate response: {}", e)) + })?; // Parse GraphQL response let results = body @@ -719,7 +751,9 @@ impl Federation { score: 1.0 - additional.get("distance")?.as_f64()? as f32, vector: additional.get("vector").and_then(|v| { v.as_array().map(|arr| { - arr.iter().filter_map(|x| x.as_f64().map(|n| n as f32)).collect() + arr.iter() + .filter_map(|x| x.as_f64().map(|n| n as f32)) + .collect() }) }), metadata: Some(item.clone()), @@ -743,7 +777,11 @@ impl Federation { query: &FederatedQuery, k: usize, ) -> Result> { - let collection = member.collections.first().cloned().unwrap_or_else(|| "default".to_string()); + let collection = member + .collections + .first() + .cloned() + .unwrap_or_else(|| "default".to_string()); let request_body = serde_json::json!({ "query_embeddings": [query.vector], @@ -770,15 +808,27 @@ impl Federation { ))); } - let body: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse ChromaDB response: {}", e)))?; + let body: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse ChromaDB response: {}", e)) + })?; // ChromaDB returns arrays for each field - let ids = body.get("ids").and_then(|i| i.get(0)).and_then(|a| a.as_array()); - let distances = body.get("distances").and_then(|d| d.get(0)).and_then(|a| a.as_array()); - let embeddings = body.get("embeddings").and_then(|e| e.get(0)).and_then(|a| a.as_array()); - let metadatas = body.get("metadatas").and_then(|m| m.get(0)).and_then(|a| a.as_array()); + let ids = body + .get("ids") + .and_then(|i| i.get(0)) + .and_then(|a| a.as_array()); + let distances = body + .get("distances") + .and_then(|d| d.get(0)) + .and_then(|a| a.as_array()); + let embeddings = body + .get("embeddings") + .and_then(|e| e.get(0)) + .and_then(|a| a.as_array()); + let metadatas = body + .get("metadatas") + .and_then(|m| m.get(0)) + .and_then(|a| a.as_array()); let mut results = Vec::new(); if let Some(ids) = ids { @@ -793,7 +843,11 @@ impl Federation { let vector = embeddings .and_then(|e| e.get(i)) .and_then(|v| v.as_array()) - .map(|arr| arr.iter().filter_map(|x| x.as_f64().map(|n| n as f32)).collect()); + .map(|arr| { + arr.iter() + .filter_map(|x| x.as_f64().map(|n| n as f32)) + .collect() + }); let metadata = metadatas.and_then(|m| m.get(i)).cloned(); @@ -820,18 +874,26 @@ impl Federation { query: &FederatedQuery, k: usize, ) -> Result> { - use crate::server::types::pb::{vec_store_service_client::VecStoreServiceClient, QueryRequest}; + use crate::server::types::pb::{ + QueryRequest, vec_store_service_client::VecStoreServiceClient, + }; // Create a tokio runtime for blocking on async gRPC calls - let rt = tokio::runtime::Runtime::new() - .map_err(|e| VecStoreError::Internal(format!("Failed to create tokio runtime: {}", e)))?; + let rt = tokio::runtime::Runtime::new().map_err(|e| { + VecStoreError::Internal(format!("Failed to create tokio runtime: {}", e)) + })?; rt.block_on(async { // Connect to the gRPC server let endpoint = member.endpoint.clone(); let mut client = VecStoreServiceClient::connect(endpoint.clone()) .await - .map_err(|e| VecStoreError::Internal(format!("gRPC connection failed to {}: {}", endpoint, e)))?; + .map_err(|e| { + VecStoreError::Internal(format!( + "gRPC connection failed to {}: {}", + endpoint, e + )) + })?; // Build the query request let request = QueryRequest { @@ -842,13 +904,12 @@ impl Federation { }; // Execute the query with timeout - let response = tokio::time::timeout( - self.config.member_timeout, - client.query(request) - ) - .await - .map_err(|_| VecStoreError::Internal(format!("gRPC query timeout to {}", endpoint)))? - .map_err(|e| VecStoreError::Internal(format!("gRPC query failed: {}", e)))?; + let response = tokio::time::timeout(self.config.member_timeout, client.query(request)) + .await + .map_err(|_| { + VecStoreError::Internal(format!("gRPC query timeout to {}", endpoint)) + })? + .map_err(|e| VecStoreError::Internal(format!("gRPC query failed: {}", e)))?; // Convert response to MemberResult let results: Vec = response @@ -876,11 +937,15 @@ impl Federation { }) } - fn merge_results(&self, mut results: Vec, query: &FederatedQuery) -> Vec { + fn merge_results( + &self, + mut results: Vec, + query: &FederatedQuery, + ) -> Vec { match self.config.merge_strategy { MergeStrategy::ScoreBased => { results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap()); - } + }, MergeStrategy::RoundRobin => { // Group by member, then interleave let mut by_member: HashMap> = HashMap::new(); @@ -897,11 +962,12 @@ impl Federation { } } } - } + }, MergeStrategy::Priority => { // Already sorted by member priority, just sort within each member by score results.sort_by(|a, b| { - let member_cmp = self.get_member_priority(&b.member_id) + let member_cmp = self + .get_member_priority(&b.member_id) .cmp(&self.get_member_priority(&a.member_id)); if member_cmp == std::cmp::Ordering::Equal { b.score.partial_cmp(&a.score).unwrap() @@ -909,7 +975,7 @@ impl Federation { member_cmp } }); - } + }, MergeStrategy::Deduplicate => { // Remove duplicates by ID, keeping highest score let mut seen: HashMap = HashMap::new(); @@ -924,7 +990,7 @@ impl Federation { } results = seen.into_values().collect(); results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap()); - } + }, MergeStrategy::Weighted => { // Apply member weights to scores for r in &mut results { @@ -932,7 +998,7 @@ impl Federation { r.score *= weight; } results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap()); - } + }, } // Apply min score filter @@ -944,7 +1010,8 @@ impl Federation { results.truncate(query.k); // Convert to federated results - results.into_iter() + results + .into_iter() .map(|r| FederatedResult { id: r.id, score: r.score, @@ -957,7 +1024,9 @@ impl Federation { } fn get_member_priority(&self, member_id: &str) -> u32 { - let Ok(members) = self.members.read() else { return 0; }; + let Ok(members) = self.members.read() else { + return 0; + }; members .get(member_id) .map(|s| s.member.priority) @@ -965,7 +1034,9 @@ impl Federation { } fn get_member_weight(&self, member_id: &str) -> f32 { - let Ok(members) = self.members.read() else { return 1.0; }; + let Ok(members) = self.members.read() else { + return 1.0; + }; members .get(member_id) .map(|s| s.member.weight) @@ -990,25 +1061,34 @@ impl Federation { } fn get_cached(&self, key: &str) -> Option> { - let Ok(cache) = self.cache.read() else { return None; }; + let Ok(cache) = self.cache.read() else { + return None; + }; if let Some(cached) = cache.get(key) - && cached.cached_at.elapsed() < self.config.cache_ttl { - return Some(cached.results.clone()); - } + && cached.cached_at.elapsed() < self.config.cache_ttl + { + return Some(cached.results.clone()); + } None } fn cache_result(&self, key: &str, results: Vec) { - let Ok(mut cache) = self.cache.write() else { return; }; - cache.insert(key.to_string(), CachedResult { - results, - cached_at: Instant::now(), - }); + let Ok(mut cache) = self.cache.write() else { + return; + }; + cache.insert( + key.to_string(), + CachedResult { + results, + cached_at: Instant::now(), + }, + ); // Limit cache size if cache.len() > 10000 { // Remove oldest entries - let mut entries: Vec<_> = cache.iter() + let mut entries: Vec<_> = cache + .iter() .map(|(k, v)| (k.clone(), v.cached_at)) .collect(); entries.sort_by_key(|(_, t)| *t); @@ -1021,10 +1101,14 @@ impl Federation { /// Update member health status pub fn update_health(&self, member_id: &str, _health: MemberHealth) -> Result<()> { - let members = self.members.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire read lock on members".into()))?; + let members = self.members.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire read lock on members".into()) + })?; let Some(_state) = members.get(member_id) else { - return Err(VecStoreError::NotFound(format!("Member {} not found", member_id))); + return Err(VecStoreError::NotFound(format!( + "Member {} not found", + member_id + ))); }; // Can't mutate through Arc, but in production we'd use interior mutability Ok(()) @@ -1045,7 +1129,8 @@ impl Federation { }; }; - let member_statuses: Vec = members.values() + let member_statuses: Vec = members + .values() .map(|s| { let avg_latency = s.avg_latency_ms.read().map(|g| *g).unwrap_or(0.0); MemberStatus { @@ -1062,7 +1147,8 @@ impl Federation { }) .collect(); - let healthy_count = member_statuses.iter() + let healthy_count = member_statuses + .iter() .filter(|m| m.health == MemberHealth::Healthy) .count(); @@ -1088,24 +1174,25 @@ impl Federation { /// List all members pub fn list_members(&self) -> Vec { - let Ok(members) = self.members.read() else { return vec![]; }; - members - .values() - .map(|s| s.member.clone()) - .collect() + let Ok(members) = self.members.read() else { + return vec![]; + }; + members.values().map(|s| s.member.clone()).collect() } /// Get a specific member pub fn get_member(&self, member_id: &str) -> Option { - let Ok(members) = self.members.read() else { return None; }; - members - .get(member_id) - .map(|s| s.member.clone()) + let Ok(members) = self.members.read() else { + return None; + }; + members.get(member_id).map(|s| s.member.clone()) } /// Clear cache pub fn clear_cache(&self) { - let Ok(mut cache) = self.cache.write() else { return; }; + let Ok(mut cache) = self.cache.write() else { + return; + }; cache.clear(); } } @@ -1301,7 +1388,7 @@ mod tests { // Results should be sorted by score for i in 1..result.results.len() { - assert!(result.results[i-1].score >= result.results[i].score); + assert!(result.results[i - 1].score >= result.results[i].score); } } diff --git a/src/fuzzy.rs b/src/fuzzy.rs index f134d82..3e15974 100644 --- a/src/fuzzy.rs +++ b/src/fuzzy.rs @@ -481,10 +481,10 @@ impl BKNode { if distance < *best_dist { *best = Some((self.word.clone(), distance)); } - } + }, None => { *best = Some((self.word.clone(), distance)); - } + }, } // Search children diff --git a/src/generative.rs b/src/generative.rs index 609c3ee..b78f47b 100644 --- a/src/generative.rs +++ b/src/generative.rs @@ -35,9 +35,9 @@ //! println!("Sources: {:?}", response.citations); //! ``` +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; -use serde::{Deserialize, Serialize}; use crate::error::Result; use crate::store::{Query, VecStore}; @@ -64,20 +64,11 @@ pub enum LLMProvider { max_tokens: Option, }, /// Cohere Command models - Cohere { - model: String, - api_key: String, - }, + Cohere { model: String, api_key: String }, /// Google Gemini - Gemini { - model: String, - api_key: String, - }, + Gemini { model: String, api_key: String }, /// Ollama local models - Ollama { - model: String, - base_url: String, - }, + Ollama { model: String, base_url: String }, /// Azure OpenAI AzureOpenAI { deployment: String, @@ -86,10 +77,7 @@ pub enum LLMProvider { api_version: String, }, /// AWS Bedrock - Bedrock { - model_id: String, - region: String, - }, + Bedrock { model_id: String, region: String }, /// Custom HTTP endpoint Custom { endpoint: String, @@ -154,17 +142,23 @@ fn default_prompt_template() -> String { "Answer the following question based on the provided context.\n\n\ Context:\n{context}\n\n\ Question: {query}\n\n\ - Answer:".to_string() + Answer:" + .to_string() } fn default_system_prompt() -> String { "You are a helpful assistant that answers questions based on the provided context. \ Be concise and accurate. If the context doesn't contain enough information, \ - say so rather than making up an answer.".to_string() + say so rather than making up an answer." + .to_string() } -fn default_true() -> bool { true } -fn default_context_length() -> usize { 4000 } +fn default_true() -> bool { + true +} +fn default_context_length() -> usize { + 4000 +} impl GenerativeConfig { /// Create a new configuration @@ -428,11 +422,7 @@ impl GenerativeSearch { } /// Generate a response for a query - pub fn generate( - &self, - query: &str, - options: GenerativeQuery, - ) -> Result { + pub fn generate(&self, query: &str, options: GenerativeQuery) -> Result { let start = std::time::Instant::now(); // Step 1: Retrieve relevant documents (simulated) @@ -488,14 +478,21 @@ impl GenerativeSearch { // Query the store let store_guard = store.read().map_err(|e| { - crate::error::VecStoreError::LockError(format!("Failed to acquire store lock: {}", e)) + crate::error::VecStoreError::LockError(format!( + "Failed to acquire store lock: {}", + e + )) })?; - let results = store_guard.query(Query { - vector: query_vector, - k: options.limit, - filter: None, // Filter support planned for future release - }).map_err(|e| crate::error::VecStoreError::Internal(format!("Query failed: {}", e)))?; + let results = store_guard + .query(Query { + vector: query_vector, + k: options.limit, + filter: None, // Filter support planned for future release + }) + .map_err(|e| { + crate::error::VecStoreError::Internal(format!("Query failed: {}", e)) + })?; // Convert to RetrievedDocument format return Ok(results @@ -522,7 +519,9 @@ impl GenerativeSearch { id: neighbor.id, content, score: neighbor.score, - metadata: Some(serde_json::to_value(&neighbor.metadata.fields).unwrap_or_default()), + metadata: Some( + serde_json::to_value(&neighbor.metadata.fields).unwrap_or_default(), + ), } }) .collect()); @@ -530,7 +529,9 @@ impl GenerativeSearch { // Fallback: return placeholder documents when no store/embedder configured // Useful for testing RAG pipelines without full vector store setup - tracing::debug!("No store/embedder configured - returning placeholder documents for testing"); + tracing::debug!( + "No store/embedder configured - returning placeholder documents for testing" + ); Ok((0..options.limit.min(3)) .map(|i| RetrievedDocument { id: format!("placeholder_doc_{}", i), @@ -554,11 +555,7 @@ impl GenerativeSearch { let mut context = String::new(); for (i, doc) in docs.iter().enumerate() { - context.push_str(&format!( - "[{}] {}\n\n", - i + 1, - doc.content - )); + context.push_str(&format!("[{}] {}\n\n", i + 1, doc.content)); } // Truncate if too long @@ -577,7 +574,8 @@ impl GenerativeSearch { context: &str, options: &GenerativeQuery, ) -> Result { - let template = options.prompt_template + let template = options + .prompt_template .as_ref() .unwrap_or(&self.config.prompt_template); @@ -604,9 +602,7 @@ impl GenerativeSearch { max_tokens, } => self.call_openai(prompt, model, api_key, *temperature, *max_tokens), - LLMProvider::Ollama { model, base_url } => { - self.call_ollama(prompt, model, base_url) - } + LLMProvider::Ollama { model, base_url } => self.call_ollama(prompt, model, base_url), // For other providers, return a message indicating they're not yet implemented _ => Ok(format!( @@ -676,9 +672,9 @@ impl GenerativeSearch { ))); } - let json: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse OpenAI response: {}", e)))?; + let json: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse OpenAI response: {}", e)) + })?; json["choices"][0]["message"]["content"] .as_str() @@ -717,9 +713,9 @@ impl GenerativeSearch { ))); } - let json: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse Ollama response: {}", e)))?; + let json: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse Ollama response: {}", e)) + })?; json["response"] .as_str() @@ -921,10 +917,9 @@ mod tests { let config = GenerativeConfig::ollama("llama3"); let gen_search = GenerativeSearch::new(config).unwrap(); - let response = gen_search.generate( - "What is VecStore?", - GenerativeQuery::new().with_limit(3), - ).unwrap(); + let response = gen_search + .generate("What is VecStore?", GenerativeQuery::new().with_limit(3)) + .unwrap(); assert!(!response.generated_text.is_empty()); assert!(!response.citations.is_empty()); @@ -935,10 +930,9 @@ mod tests { let config = GenerativeConfig::ollama("llama3"); let gen_search = GenerativeSearch::new(config).unwrap(); - let stream = gen_search.generate_stream( - "What is VecStore?", - GenerativeQuery::new(), - ).unwrap(); + let stream = gen_search + .generate_stream("What is VecStore?", GenerativeQuery::new()) + .unwrap(); let chunks: Vec<_> = stream.collect(); assert!(!chunks.is_empty()); diff --git a/src/geospatial.rs b/src/geospatial.rs index a162eab..8794d0b 100644 --- a/src/geospatial.rs +++ b/src/geospatial.rs @@ -31,7 +31,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -191,10 +191,7 @@ impl GeoIndex { let cell_id = location.s2_cell_id(self.s2_level); // Add to cell index - self.cell_index - .entry(cell_id) - .or_default() - .push(id.clone()); + self.cell_index.entry(cell_id).or_default().push(id.clone()); // Add document let doc = GeoDocument { diff --git a/src/gpu.rs b/src/gpu.rs index 01355dd..d5e4d49 100644 --- a/src/gpu.rs +++ b/src/gpu.rs @@ -98,7 +98,7 @@ pub mod cuda_kernels; pub mod metal_executor; -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::sync::Arc; @@ -253,7 +253,7 @@ impl GpuExecutor { { return Err(anyhow!("CUDA support not compiled. Enable 'cuda' feature.")); } - } + }, GpuBackend::Metal => { #[cfg(feature = "metal")] { @@ -265,10 +265,10 @@ impl GpuExecutor { "Metal support not compiled. Enable 'metal' feature." )); } - } + }, GpuBackend::Cpu => { return Ok(Arc::new(CpuBackend::new(config))); - } + }, GpuBackend::WebGpu => { #[cfg(feature = "webgpu")] { @@ -282,7 +282,7 @@ impl GpuExecutor { { return Err(anyhow!("WebGPU requires 'webgpu' feature")); } - } + }, } } @@ -552,7 +552,10 @@ impl GpuOps for CudaBackend { name: props.name, total_memory_bytes: props.total_memory_bytes, available_memory_bytes: props.total_memory_bytes, // Approximation - compute_capability: (props.compute_capability.0 as u32, props.compute_capability.1 as u32), + compute_capability: ( + props.compute_capability.0 as u32, + props.compute_capability.1 as u32, + ), max_threads_per_block: props.max_threads_per_block as usize, num_streaming_multiprocessors: props.multiprocessor_count as usize, } @@ -570,7 +573,8 @@ impl GpuOps for CudaBackend { let flat_database: Vec = database.iter().flat_map(|v| v.iter().copied()).collect(); // Execute GPU kernel - self.executor.euclidean_distance(query, &flat_database, num_vectors, vector_dim) + self.executor + .euclidean_distance(query, &flat_database, num_vectors, vector_dim) } fn batch_cosine_similarity(&self, query: &[f32], database: &[Vec]) -> Result> { @@ -583,7 +587,8 @@ impl GpuOps for CudaBackend { let flat_database: Vec = database.iter().flat_map(|v| v.iter().copied()).collect(); - self.executor.cosine_similarity(query, &flat_database, num_vectors, vector_dim) + self.executor + .cosine_similarity(query, &flat_database, num_vectors, vector_dim) } fn batch_dot_product(&self, query: &[f32], database: &[Vec]) -> Result> { @@ -596,7 +601,8 @@ impl GpuOps for CudaBackend { let flat_database: Vec = database.iter().flat_map(|v| v.iter().copied()).collect(); - self.executor.dot_product(query, &flat_database, num_vectors, vector_dim) + self.executor + .dot_product(query, &flat_database, num_vectors, vector_dim) } fn matrix_multiply(&self, a: &[Vec], b: &[Vec]) -> Result>> { @@ -615,7 +621,9 @@ impl GpuOps for CudaBackend { let flat_vectors: Vec = vectors.iter().flat_map(|v| v.iter().copied()).collect(); - let normalized = self.executor.l2_normalize(&flat_vectors, num_vectors, vector_dim)?; + let normalized = self + .executor + .l2_normalize(&flat_vectors, num_vectors, vector_dim)?; // Reshape back to Vec> Ok(normalized.chunks(vector_dim).map(|c| c.to_vec()).collect()) @@ -646,9 +654,7 @@ impl GpuOps for CudaBackend { // ============================================================================ #[cfg(all(target_os = "macos", feature = "metal"))] -use metal::{ - Buffer, CommandQueue, ComputePipelineState, Device, MTLResourceOptions, MTLSize, -}; +use metal::{Buffer, CommandQueue, ComputePipelineState, Device, MTLResourceOptions, MTLSize}; /// Metal Shading Language source code for compute kernels #[cfg(all(target_os = "macos", feature = "metal"))] @@ -797,8 +803,8 @@ impl MetalBackend { /// Create a new Metal backend with GPU device initialization pub fn new(config: &GpuConfig) -> Result { // Get the system default Metal device - let device = Device::system_default() - .ok_or_else(|| anyhow!("No Metal-capable GPU device found"))?; + let device = + Device::system_default().ok_or_else(|| anyhow!("No Metal-capable GPU device found"))?; // Create command queue for submitting work to GPU let command_queue = device.new_command_queue(); @@ -809,7 +815,8 @@ impl MetalBackend { .map_err(|e| anyhow!("Failed to compile Metal shaders: {}", e))?; // Create compute pipeline states for each kernel - let euclidean_pipeline = Self::create_pipeline(&device, &library, "euclidean_distance_kernel")?; + let euclidean_pipeline = + Self::create_pipeline(&device, &library, "euclidean_distance_kernel")?; let cosine_pipeline = Self::create_pipeline(&device, &library, "cosine_similarity_kernel")?; let dot_product_pipeline = Self::create_pipeline(&device, &library, "dot_product_kernel")?; let normalize_pipeline = Self::create_pipeline(&device, &library, "l2_normalize_kernel")?; @@ -861,7 +868,8 @@ impl MetalBackend { /// Create an empty Metal buffer for output fn create_output_buffer(&self, num_elements: usize) -> Buffer { let size = (num_elements * size_of::()) as u64; - self.device.new_buffer(size, MTLResourceOptions::StorageModeShared) + self.device + .new_buffer(size, MTLResourceOptions::StorageModeShared) } /// Read results from a Metal buffer back to CPU @@ -951,7 +959,8 @@ impl GpuOps for MetalBackend { total_memory_bytes: recommended_max, available_memory_bytes: recommended_max, // Unified memory, estimation compute_capability: (supports_family.0, supports_family.1), - max_threads_per_block: self.euclidean_pipeline.max_total_threads_per_threadgroup() as usize, + max_threads_per_block: self.euclidean_pipeline.max_total_threads_per_threadgroup() + as usize, num_streaming_multiprocessors: 10, // Not directly exposed by Metal API } } @@ -969,7 +978,9 @@ impl GpuOps for MetalBackend { if vec.len() != vector_dim { return Err(anyhow!( "Vector {} has dimension {} but query has dimension {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -999,7 +1010,9 @@ impl GpuOps for MetalBackend { if vec.len() != vector_dim { return Err(anyhow!( "Vector {} has dimension {} but query has dimension {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -1029,7 +1042,9 @@ impl GpuOps for MetalBackend { if vec.len() != vector_dim { return Err(anyhow!( "Vector {} has dimension {} but query has dimension {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -1051,15 +1066,18 @@ impl GpuOps for MetalBackend { return Ok(vec![]); } - let m = a.len(); // rows of A - let k = a[0].len(); // cols of A = rows of B - let n = b[0].len(); // cols of B + let m = a.len(); // rows of A + let k = a[0].len(); // cols of A = rows of B + let n = b[0].len(); // cols of B // Validate dimensions if b.len() != k { return Err(anyhow!( "Matrix dimension mismatch: A is {}x{}, B is {}x{}", - m, k, b.len(), n + m, + k, + b.len(), + n )); } @@ -1132,7 +1150,9 @@ impl GpuOps for MetalBackend { if vec.len() != vector_dim { return Err(anyhow!( "Vector {} has dimension {} but expected {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -1211,9 +1231,7 @@ impl GpuOps for MetalBackend { // Partial sort for efficiency when k << n if k < indexed.len() / 2 { - indexed.select_nth_unstable_by(k, |a, b| { - a.1.total_cmp(&b.1) - }); + indexed.select_nth_unstable_by(k, |a, b| a.1.total_cmp(&b.1)); indexed.truncate(k); indexed.sort_by(|a, b| a.1.total_cmp(&b.1)); } else { @@ -1539,43 +1557,29 @@ impl WebGpuBackend { // Request device with reasonable limits for compute let (device, queue) = adapter - .request_device( - &wgpu::DeviceDescriptor { - label: Some("VecStore WebGPU Device"), - required_features: wgpu::Features::empty(), - required_limits: wgpu::Limits::default(), - memory_hints: wgpu::MemoryHints::default(), - trace: wgpu::Trace::default(), - experimental_features: wgpu::ExperimentalFeatures::default(), - }, - ) + .request_device(&wgpu::DeviceDescriptor { + label: Some("VecStore WebGPU Device"), + required_features: wgpu::Features::empty(), + required_limits: wgpu::Limits::default(), + memory_hints: wgpu::MemoryHints::default(), + trace: wgpu::Trace::default(), + experimental_features: wgpu::ExperimentalFeatures::default(), + }) .await .map_err(|e| anyhow!("Failed to create WebGPU device: {}", e))?; // Compile all compute shaders and create pipelines - let euclidean_pipeline = Self::create_compute_pipeline( - &device, - WGSL_EUCLIDEAN_DISTANCE, - "euclidean_distance", - )?; - - let cosine_pipeline = Self::create_compute_pipeline( - &device, - WGSL_COSINE_SIMILARITY, - "cosine_similarity", - )?; - - let dot_product_pipeline = Self::create_compute_pipeline( - &device, - WGSL_DOT_PRODUCT, - "dot_product", - )?; - - let normalize_pipeline = Self::create_compute_pipeline( - &device, - WGSL_L2_NORMALIZE, - "l2_normalize", - )?; + let euclidean_pipeline = + Self::create_compute_pipeline(&device, WGSL_EUCLIDEAN_DISTANCE, "euclidean_distance")?; + + let cosine_pipeline = + Self::create_compute_pipeline(&device, WGSL_COSINE_SIMILARITY, "cosine_similarity")?; + + let dot_product_pipeline = + Self::create_compute_pipeline(&device, WGSL_DOT_PRODUCT, "dot_product")?; + + let normalize_pipeline = + Self::create_compute_pipeline(&device, WGSL_L2_NORMALIZE, "l2_normalize")?; Ok(Self { config: config.clone(), @@ -1646,17 +1650,21 @@ impl WebGpuBackend { use wgpu::util::DeviceExt; // Create GPU buffers - let query_buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("Query Buffer"), - contents: bytemuck::cast_slice(query), - usage: wgpu::BufferUsages::STORAGE, - }); + let query_buffer = self + .device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Query Buffer"), + contents: bytemuck::cast_slice(query), + usage: wgpu::BufferUsages::STORAGE, + }); - let database_buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("Database Buffer"), - contents: bytemuck::cast_slice(database), - usage: wgpu::BufferUsages::STORAGE, - }); + let database_buffer = self + .device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Database Buffer"), + contents: bytemuck::cast_slice(database), + usage: wgpu::BufferUsages::STORAGE, + }); // Output buffer (results) let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { @@ -1681,11 +1689,13 @@ impl WebGpuBackend { _padding0: 0, _padding1: 0, }; - let params_buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("Params Buffer"), - contents: bytemuck::cast_slice(&[params]), - usage: wgpu::BufferUsages::UNIFORM, - }); + let params_buffer = self + .device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Params Buffer"), + contents: bytemuck::cast_slice(&[params]), + usage: wgpu::BufferUsages::UNIFORM, + }); // Create bind group let bind_group_layout = pipeline.get_bind_group_layout(0); @@ -1713,9 +1723,11 @@ impl WebGpuBackend { }); // Create and submit command buffer - let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("Compute Encoder"), - }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Compute Encoder"), + }); { let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { @@ -1749,10 +1761,12 @@ impl WebGpuBackend { }); // Poll the device until the buffer is mapped - self.device.poll(wgpu::PollType::Wait { - submission_index: None, - timeout: Some(std::time::Duration::from_secs(60)), - }).map_err(|e| anyhow!("GPU poll failed: {:?}", e))?; + self.device + .poll(wgpu::PollType::Wait { + submission_index: None, + timeout: Some(std::time::Duration::from_secs(60)), + }) + .map_err(|e| anyhow!("GPU poll failed: {:?}", e))?; receiver .recv() @@ -1779,11 +1793,13 @@ impl WebGpuBackend { let total_elements = num_vectors * vector_dim; // Create GPU buffers - let input_buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("Input Buffer"), - contents: bytemuck::cast_slice(input), - usage: wgpu::BufferUsages::STORAGE, - }); + let input_buffer = self + .device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Input Buffer"), + contents: bytemuck::cast_slice(input), + usage: wgpu::BufferUsages::STORAGE, + }); // Output buffer let output_buffer = self.device.create_buffer(&wgpu::BufferDescriptor { @@ -1808,11 +1824,13 @@ impl WebGpuBackend { _padding0: 0, _padding1: 0, }; - let params_buffer = self.device.create_buffer_init(&wgpu::util::BufferInitDescriptor { - label: Some("Params Buffer"), - contents: bytemuck::cast_slice(&[params]), - usage: wgpu::BufferUsages::UNIFORM, - }); + let params_buffer = self + .device + .create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: Some("Params Buffer"), + contents: bytemuck::cast_slice(&[params]), + usage: wgpu::BufferUsages::UNIFORM, + }); // Create bind group let bind_group_layout = self.normalize_pipeline.get_bind_group_layout(0); @@ -1836,9 +1854,11 @@ impl WebGpuBackend { }); // Create and submit command buffer - let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { - label: Some("Normalize Encoder"), - }); + let mut encoder = self + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("Normalize Encoder"), + }); { let mut compute_pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { @@ -1869,10 +1889,12 @@ impl WebGpuBackend { let _ = sender.send(result); }); - self.device.poll(wgpu::PollType::Wait { - submission_index: None, - timeout: Some(std::time::Duration::from_secs(60)), - }).map_err(|e| anyhow!("GPU poll failed: {:?}", e))?; + self.device + .poll(wgpu::PollType::Wait { + submission_index: None, + timeout: Some(std::time::Duration::from_secs(60)), + }) + .map_err(|e| anyhow!("GPU poll failed: {:?}", e))?; receiver .recv() @@ -1931,11 +1953,14 @@ impl GpuOps for WebGpuBackend { GpuDeviceInfo { backend: GpuBackend::WebGpu, device_id: self.config.device_id, - name: format!("{} ({:?})", self.adapter_info.name, self.adapter_info.backend), + name: format!( + "{} ({:?})", + self.adapter_info.name, self.adapter_info.backend + ), total_memory_bytes: self.config.max_memory_bytes, available_memory_bytes: self.config.max_memory_bytes, - compute_capability: (1, 0), // WebGPU version indicator - max_threads_per_block: 256, // Workgroup size + compute_capability: (1, 0), // WebGPU version indicator + max_threads_per_block: 256, // Workgroup size num_streaming_multiprocessors: 1, // Not directly available in WebGPU } } @@ -1951,7 +1976,9 @@ impl GpuOps for WebGpuBackend { if vec.len() != vector_dim { return Err(anyhow!( "Database vector {} has dimension {} but query has dimension {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -1971,7 +1998,9 @@ impl GpuOps for WebGpuBackend { if vec.len() != vector_dim { return Err(anyhow!( "Database vector {} has dimension {} but query has dimension {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -1991,7 +2020,9 @@ impl GpuOps for WebGpuBackend { if vec.len() != vector_dim { return Err(anyhow!( "Database vector {} has dimension {} but query has dimension {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -2022,7 +2053,9 @@ impl GpuOps for WebGpuBackend { if vec.len() != vector_dim { return Err(anyhow!( "Vector {} has dimension {} but expected {}", - i, vec.len(), vector_dim + i, + vec.len(), + vector_dim )); } } @@ -2061,9 +2094,7 @@ impl GpuOps for WebGpuBackend { // Partial sort to get top k (more efficient than full sort for small k) let k = k.min(indexed.len()); if k < indexed.len() { - indexed.select_nth_unstable_by(k, |a, b| { - a.1.total_cmp(&b.1) - }); + indexed.select_nth_unstable_by(k, |a, b| a.1.total_cmp(&b.1)); indexed.truncate(k); } @@ -2351,7 +2382,10 @@ mod tests { #[test] fn test_metal_is_available() { - assert!(MetalBackend::is_available(), "Metal should be available on macOS"); + assert!( + MetalBackend::is_available(), + "Metal should be available on macOS" + ); } #[test] @@ -2373,16 +2407,19 @@ mod tests { let query = vec![1.0, 2.0, 3.0, 4.0]; let database = vec![ - vec![1.0, 2.0, 3.0, 4.0], // Same as query - vec![2.0, 3.0, 4.0, 5.0], // Different - vec![0.0, 0.0, 0.0, 0.0], // Origin + vec![1.0, 2.0, 3.0, 4.0], // Same as query + vec![2.0, 3.0, 4.0, 5.0], // Different + vec![0.0, 0.0, 0.0, 0.0], // Origin ]; let distances = backend.batch_euclidean_distance(&query, &database).unwrap(); assert_eq!(distances.len(), 3); assert!(distances[0] < 0.001, "Same vector should have distance ~0"); - assert!(distances[1] > 1.0, "Different vector should have larger distance"); + assert!( + distances[1] > 1.0, + "Different vector should have larger distance" + ); assert!(distances[2] > 5.0, "Origin should have large distance"); } @@ -2393,17 +2430,23 @@ mod tests { let query = vec![1.0, 0.0, 0.0, 0.0]; let database = vec![ - vec![1.0, 0.0, 0.0, 0.0], // Same direction - vec![0.0, 1.0, 0.0, 0.0], // Perpendicular - vec![-1.0, 0.0, 0.0, 0.0], // Opposite + vec![1.0, 0.0, 0.0, 0.0], // Same direction + vec![0.0, 1.0, 0.0, 0.0], // Perpendicular + vec![-1.0, 0.0, 0.0, 0.0], // Opposite ]; let similarities = backend.batch_cosine_similarity(&query, &database).unwrap(); assert_eq!(similarities.len(), 3); - assert!((similarities[0] - 1.0).abs() < 0.001, "Same direction should be 1.0"); + assert!( + (similarities[0] - 1.0).abs() < 0.001, + "Same direction should be 1.0" + ); assert!(similarities[1].abs() < 0.001, "Perpendicular should be 0.0"); - assert!((similarities[2] + 1.0).abs() < 0.001, "Opposite should be -1.0"); + assert!( + (similarities[2] + 1.0).abs() < 0.001, + "Opposite should be -1.0" + ); } #[test] @@ -2413,8 +2456,8 @@ mod tests { let query = vec![1.0, 2.0, 3.0, 4.0]; let database = vec![ - vec![1.0, 1.0, 1.0, 1.0], // Sum = 10 - vec![2.0, 2.0, 2.0, 2.0], // Sum = 20 + vec![1.0, 1.0, 1.0, 1.0], // Sum = 10 + vec![2.0, 2.0, 2.0, 2.0], // Sum = 20 ]; let products = backend.batch_dot_product(&query, &database).unwrap(); @@ -2430,8 +2473,8 @@ mod tests { let backend = MetalBackend::new(&config).unwrap(); let vectors = vec![ - vec![3.0, 4.0, 0.0, 0.0], // Magnitude 5 - vec![1.0, 0.0, 0.0, 0.0], // Already normalized + vec![3.0, 4.0, 0.0, 0.0], // Magnitude 5 + vec![1.0, 0.0, 0.0, 0.0], // Already normalized ]; let normalized = backend.batch_normalize(&vectors).unwrap(); @@ -2447,23 +2490,17 @@ mod tests { let config = GpuConfig::default(); let backend = MetalBackend::new(&config).unwrap(); - let a = vec![ - vec![1.0, 2.0], - vec![3.0, 4.0], - ]; - let b = vec![ - vec![5.0, 6.0], - vec![7.0, 8.0], - ]; + let a = vec![vec![1.0, 2.0], vec![3.0, 4.0]]; + let b = vec![vec![5.0, 6.0], vec![7.0, 8.0]]; let result = backend.matrix_multiply(&a, &b).unwrap(); assert_eq!(result.len(), 2); assert_eq!(result[0].len(), 2); - assert!((result[0][0] - 19.0).abs() < 0.001); // 1*5 + 2*7 - assert!((result[0][1] - 22.0).abs() < 0.001); // 1*6 + 2*8 - assert!((result[1][0] - 43.0).abs() < 0.001); // 3*5 + 4*7 - assert!((result[1][1] - 50.0).abs() < 0.001); // 3*6 + 4*8 + assert!((result[0][0] - 19.0).abs() < 0.001); // 1*5 + 2*7 + assert!((result[0][1] - 22.0).abs() < 0.001); // 1*6 + 2*8 + assert!((result[1][0] - 43.0).abs() < 0.001); // 3*5 + 4*7 + assert!((result[1][1] - 50.0).abs() < 0.001); // 3*6 + 4*8 } #[test] @@ -2475,7 +2512,7 @@ mod tests { let database = vec![ vec![0.0, 0.0, 0.0, 0.0], vec![1.0, 1.0, 0.0, 0.0], - vec![0.5, 0.5, 0.0, 0.0], // Exact match + vec![0.5, 0.5, 0.0, 0.0], // Exact match vec![10.0, 10.0, 0.0, 0.0], ]; @@ -2483,7 +2520,7 @@ mod tests { assert_eq!(indices.len(), 2); assert_eq!(distances.len(), 2); - assert_eq!(indices[0], 2); // Exact match should be first + assert_eq!(indices[0], 2); // Exact match should be first assert!(distances[0] < 0.001); } @@ -2498,7 +2535,11 @@ mod tests { let query: Vec = (0..dimension).map(|i| i as f32 * 0.01).collect(); let database: Vec> = (0..num_vectors) - .map(|i| (0..dimension).map(|j| ((i + j) % 100) as f32 * 0.01).collect()) + .map(|i| { + (0..dimension) + .map(|j| ((i + j) % 100) as f32 * 0.01) + .collect() + }) .collect(); let distances = backend.batch_euclidean_distance(&query, &database).unwrap(); @@ -2516,12 +2557,11 @@ mod tests { assert_eq!(executor.backend_type(), GpuBackend::Metal); let query = vec![1.0, 2.0, 3.0]; - let database = vec![ - vec![1.0, 2.0, 3.0], - vec![4.0, 5.0, 6.0], - ]; + let database = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]; - let distances = executor.batch_euclidean_distance(&query, &database).unwrap(); + let distances = executor + .batch_euclidean_distance(&query, &database) + .unwrap(); assert_eq!(distances.len(), 2); } @@ -2551,7 +2591,7 @@ mod tests { let query = vec![1.0, 2.0, 3.0]; let mismatched_db = vec![ - vec![1.0, 2.0], // Wrong dimension + vec![1.0, 2.0], // Wrong dimension ]; let result = backend.batch_euclidean_distance(&query, &mismatched_db); @@ -2594,21 +2634,24 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; let query = vec![1.0, 2.0, 3.0, 4.0]; let database = vec![ - vec![1.0, 2.0, 3.0, 4.0], // Same as query - vec![2.0, 3.0, 4.0, 5.0], // Different - vec![0.0, 0.0, 0.0, 0.0], // Origin + vec![1.0, 2.0, 3.0, 4.0], // Same as query + vec![2.0, 3.0, 4.0, 5.0], // Different + vec![0.0, 0.0, 0.0, 0.0], // Origin ]; let distances = backend.batch_euclidean_distance(&query, &database).unwrap(); assert_eq!(distances.len(), 3); assert!(distances[0] < 0.001, "Same vector should have distance ~0"); - assert!(distances[1] > 1.0, "Different vector should have larger distance"); + assert!( + distances[1] > 1.0, + "Different vector should have larger distance" + ); assert!(distances[2] > 5.0, "Origin should have large distance"); } @@ -2620,22 +2663,28 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; let query = vec![1.0, 0.0, 0.0, 0.0]; let database = vec![ - vec![1.0, 0.0, 0.0, 0.0], // Same direction - vec![0.0, 1.0, 0.0, 0.0], // Perpendicular - vec![-1.0, 0.0, 0.0, 0.0], // Opposite + vec![1.0, 0.0, 0.0, 0.0], // Same direction + vec![0.0, 1.0, 0.0, 0.0], // Perpendicular + vec![-1.0, 0.0, 0.0, 0.0], // Opposite ]; let similarities = backend.batch_cosine_similarity(&query, &database).unwrap(); assert_eq!(similarities.len(), 3); - assert!((similarities[0] - 1.0).abs() < 0.01, "Same direction should be 1.0"); + assert!( + (similarities[0] - 1.0).abs() < 0.01, + "Same direction should be 1.0" + ); assert!(similarities[1].abs() < 0.01, "Perpendicular should be 0.0"); - assert!((similarities[2] + 1.0).abs() < 0.01, "Opposite should be -1.0"); + assert!( + (similarities[2] + 1.0).abs() < 0.01, + "Opposite should be -1.0" + ); } #[test] @@ -2646,13 +2695,13 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; let query = vec![1.0, 2.0, 3.0, 4.0]; let database = vec![ - vec![1.0, 1.0, 1.0, 1.0], // Sum = 10 - vec![2.0, 2.0, 2.0, 2.0], // Sum = 20 + vec![1.0, 1.0, 1.0, 1.0], // Sum = 10 + vec![2.0, 2.0, 2.0, 2.0], // Sum = 20 ]; let products = backend.batch_dot_product(&query, &database).unwrap(); @@ -2670,12 +2719,12 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; let vectors = vec![ - vec![3.0, 4.0, 0.0, 0.0], // Magnitude 5 - vec![1.0, 0.0, 0.0, 0.0], // Already normalized + vec![3.0, 4.0, 0.0, 0.0], // Magnitude 5 + vec![1.0, 0.0, 0.0, 0.0], // Already normalized ]; let normalized = backend.batch_normalize(&vectors).unwrap(); @@ -2694,14 +2743,14 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; let query = vec![0.5, 0.5, 0.0, 0.0]; let database = vec![ vec![0.0, 0.0, 0.0, 0.0], vec![1.0, 1.0, 0.0, 0.0], - vec![0.5, 0.5, 0.0, 0.0], // Exact match + vec![0.5, 0.5, 0.0, 0.0], // Exact match vec![10.0, 10.0, 0.0, 0.0], ]; @@ -2709,7 +2758,7 @@ mod tests { assert_eq!(indices.len(), 2); assert_eq!(distances.len(), 2); - assert_eq!(indices[0], 2); // Exact match should be first + assert_eq!(indices[0], 2); // Exact match should be first assert!(distances[0] < 0.001); } @@ -2721,7 +2770,7 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; // Create a large batch to test GPU performance @@ -2730,7 +2779,11 @@ mod tests { let query: Vec = (0..dimension).map(|i| i as f32 * 0.01).collect(); let database: Vec> = (0..num_vectors) - .map(|i| (0..dimension).map(|j| ((i + j) % 100) as f32 * 0.01).collect()) + .map(|i| { + (0..dimension) + .map(|j| ((i + j) % 100) as f32 * 0.01) + .collect() + }) .collect(); let distances = backend.batch_euclidean_distance(&query, &database).unwrap(); @@ -2748,18 +2801,17 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; assert_eq!(executor.backend_type(), GpuBackend::WebGpu); let query = vec![1.0, 2.0, 3.0]; - let database = vec![ - vec![1.0, 2.0, 3.0], - vec![4.0, 5.0, 6.0], - ]; + let database = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]]; - let distances = executor.batch_euclidean_distance(&query, &database).unwrap(); + let distances = executor + .batch_euclidean_distance(&query, &database) + .unwrap(); assert_eq!(distances.len(), 2); } @@ -2771,7 +2823,7 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; // Empty database should return empty results @@ -2796,12 +2848,12 @@ mod tests { Err(_) => { println!("Skipping test: WebGPU not available"); return; - } + }, }; let query = vec![1.0, 2.0, 3.0]; let mismatched_db = vec![ - vec![1.0, 2.0], // Wrong dimension + vec![1.0, 2.0], // Wrong dimension ]; let result = backend.batch_euclidean_distance(&query, &mismatched_db); diff --git a/src/gpu/cuda_kernels.rs b/src/gpu/cuda_kernels.rs index 91e290a..f37fb5e 100644 --- a/src/gpu/cuda_kernels.rs +++ b/src/gpu/cuda_kernels.rs @@ -3,7 +3,6 @@ //! This module provides CUDA kernel implementations for distance calculations //! and other vector operations. - /// CUDA kernel source for Euclidean distance pub const EUCLIDEAN_DISTANCE_KERNEL: &str = r#" extern "C" __global__ void euclidean_distance_kernel( @@ -366,11 +365,17 @@ END: .map_err(|e| anyhow!("Failed to initialize CUDA device {}: {}", device_id, e))?; // Load PTX module - device.load_ptx( - cudarc::nvrtc::Ptx::from_src(Self::DISTANCE_PTX), - "distance_kernels", - &["euclidean_distance_kernel", "cosine_similarity_kernel", "dot_product_kernel"], - ).map_err(|e| anyhow!("Failed to load PTX: {}", e))?; + device + .load_ptx( + cudarc::nvrtc::Ptx::from_src(Self::DISTANCE_PTX), + "distance_kernels", + &[ + "euclidean_distance_kernel", + "cosine_similarity_kernel", + "dot_product_kernel", + ], + ) + .map_err(|e| anyhow!("Failed to load PTX: {}", e))?; let euclidean_kernel = device .get_func("distance_kernels", "euclidean_distance_kernel") @@ -404,13 +409,19 @@ END: use cudarc::driver::LaunchConfig; // Copy data to device - let query_dev = self.device.htod_sync_copy(query) + let query_dev = self + .device + .htod_sync_copy(query) .map_err(|e| anyhow!("Failed to copy query to device: {}", e))?; - let database_dev = self.device.htod_sync_copy(database) + let database_dev = self + .device + .htod_sync_copy(database) .map_err(|e| anyhow!("Failed to copy database to device: {}", e))?; // Allocate output buffer - let mut distances_dev = self.device.alloc_zeros::(num_vectors) + let mut distances_dev = self + .device + .alloc_zeros::(num_vectors) .map_err(|e| anyhow!("Failed to allocate output buffer: {}", e))?; // Configure launch @@ -424,14 +435,25 @@ END: // Launch kernel unsafe { - self.euclidean_kernel.clone().launch( - cfg, - (&query_dev, &database_dev, &mut distances_dev, num_vectors as u32, vector_dim as u32), - ).map_err(|e| anyhow!("Kernel launch failed: {}", e))?; + self.euclidean_kernel + .clone() + .launch( + cfg, + ( + &query_dev, + &database_dev, + &mut distances_dev, + num_vectors as u32, + vector_dim as u32, + ), + ) + .map_err(|e| anyhow!("Kernel launch failed: {}", e))?; } // Copy results back - let distances = self.device.dtoh_sync_copy(&distances_dev) + let distances = self + .device + .dtoh_sync_copy(&distances_dev) .map_err(|e| anyhow!("Failed to copy results from device: {}", e))?; Ok(distances) @@ -448,12 +470,18 @@ END: use cudarc::driver::LaunchAsync; use cudarc::driver::LaunchConfig; - let query_dev = self.device.htod_sync_copy(query) + let query_dev = self + .device + .htod_sync_copy(query) .map_err(|e| anyhow!("Failed to copy query: {}", e))?; - let database_dev = self.device.htod_sync_copy(database) + let database_dev = self + .device + .htod_sync_copy(database) .map_err(|e| anyhow!("Failed to copy database: {}", e))?; - let mut similarities_dev = self.device.alloc_zeros::(num_vectors) + let mut similarities_dev = self + .device + .alloc_zeros::(num_vectors) .map_err(|e| anyhow!("Failed to allocate output: {}", e))?; let threads_per_block = 256u32; @@ -465,13 +493,24 @@ END: }; unsafe { - self.cosine_kernel.clone().launch( - cfg, - (&query_dev, &database_dev, &mut similarities_dev, num_vectors as u32, vector_dim as u32), - ).map_err(|e| anyhow!("Kernel launch failed: {}", e))?; + self.cosine_kernel + .clone() + .launch( + cfg, + ( + &query_dev, + &database_dev, + &mut similarities_dev, + num_vectors as u32, + vector_dim as u32, + ), + ) + .map_err(|e| anyhow!("Kernel launch failed: {}", e))?; } - let similarities = self.device.dtoh_sync_copy(&similarities_dev) + let similarities = self + .device + .dtoh_sync_copy(&similarities_dev) .map_err(|e| anyhow!("Failed to copy results: {}", e))?; Ok(similarities) @@ -488,12 +527,18 @@ END: use cudarc::driver::LaunchAsync; use cudarc::driver::LaunchConfig; - let query_dev = self.device.htod_sync_copy(query) + let query_dev = self + .device + .htod_sync_copy(query) .map_err(|e| anyhow!("Failed to copy query: {}", e))?; - let database_dev = self.device.htod_sync_copy(database) + let database_dev = self + .device + .htod_sync_copy(database) .map_err(|e| anyhow!("Failed to copy database: {}", e))?; - let mut products_dev = self.device.alloc_zeros::(num_vectors) + let mut products_dev = self + .device + .alloc_zeros::(num_vectors) .map_err(|e| anyhow!("Failed to allocate output: {}", e))?; let threads_per_block = 256u32; @@ -505,13 +550,24 @@ END: }; unsafe { - self.dot_kernel.clone().launch( - cfg, - (&query_dev, &database_dev, &mut products_dev, num_vectors as u32, vector_dim as u32), - ).map_err(|e| anyhow!("Kernel launch failed: {}", e))?; + self.dot_kernel + .clone() + .launch( + cfg, + ( + &query_dev, + &database_dev, + &mut products_dev, + num_vectors as u32, + vector_dim as u32, + ), + ) + .map_err(|e| anyhow!("Kernel launch failed: {}", e))?; } - let products = self.device.dtoh_sync_copy(&products_dev) + let products = self + .device + .dtoh_sync_copy(&products_dev) .map_err(|e| anyhow!("Failed to copy results: {}", e))?; Ok(products) diff --git a/src/gpu/metal_executor.rs b/src/gpu/metal_executor.rs index 8c2288e..1575591 100644 --- a/src/gpu/metal_executor.rs +++ b/src/gpu/metal_executor.rs @@ -258,7 +258,9 @@ mod tests { let query = vec![1.0, 0.0]; let database = vec![0.0, 0.0, 1.0, 1.0]; - let distances = executor.euclidean_distance(&query, &database, 2, 2).unwrap(); + let distances = executor + .euclidean_distance(&query, &database, 2, 2) + .unwrap(); assert_eq!(distances.len(), 2); assert!((distances[0] - 1.0).abs() < 1e-6); // distance to [0, 0] is 1 diff --git a/src/gpu_indexing.rs b/src/gpu_indexing.rs index ee085bb..eecdaa8 100644 --- a/src/gpu_indexing.rs +++ b/src/gpu_indexing.rs @@ -31,7 +31,10 @@ //! ``` use std::collections::HashMap; -use std::sync::{RwLock, atomic::{AtomicUsize, Ordering}}; +use std::sync::{ + RwLock, + atomic::{AtomicUsize, Ordering}, +}; use std::time::{Duration, Instant}; use serde::{Deserialize, Serialize}; @@ -229,8 +232,11 @@ impl MemoryPool { fn allocate(&self, size: usize) -> Result { // Try to reuse existing buffer { - let mut buffers = self.buffers.write() - .map_err(|_| VecStoreError::LockError("failed to acquire write lock on GPU memory pool buffers".into()))?; + let mut buffers = self.buffers.write().map_err(|_| { + VecStoreError::LockError( + "failed to acquire write lock on GPU memory pool buffers".into(), + ) + })?; for buffer in &mut *buffers { if !buffer.in_use && buffer.size >= size { buffer.in_use = true; @@ -255,15 +261,20 @@ impl MemoryPool { in_use: true, }; - let mut buffers = self.buffers.write() - .map_err(|_| VecStoreError::LockError("failed to acquire write lock on GPU memory pool buffers for allocation".into()))?; + let mut buffers = self.buffers.write().map_err(|_| { + VecStoreError::LockError( + "failed to acquire write lock on GPU memory pool buffers for allocation".into(), + ) + })?; buffers.push(buffer); Ok(id) } fn free(&self, id: usize) { - let Ok(mut buffers) = self.buffers.write() else { return; }; + let Ok(mut buffers) = self.buffers.write() else { + return; + }; for buffer in &mut *buffers { if buffer.id == id { buffer.in_use = false; @@ -342,9 +353,15 @@ impl GpuHnswIndex { results } - fn search_layer(&self, query: &[f32], entry: usize, ef: usize, layer: usize) -> Vec<(usize, f32)> { - use std::collections::{BinaryHeap, HashSet}; + fn search_layer( + &self, + query: &[f32], + entry: usize, + ef: usize, + layer: usize, + ) -> Vec<(usize, f32)> { use std::cmp::Ordering; + use std::collections::{BinaryHeap, HashSet}; #[derive(Clone)] struct Candidate { @@ -378,8 +395,14 @@ impl GpuHnswIndex { let mut results: BinaryHeap = BinaryHeap::new(); let entry_dist = self.distance(query, entry); - candidates.push(Candidate { id: entry, distance: entry_dist }); - results.push(Candidate { id: entry, distance: entry_dist }); + candidates.push(Candidate { + id: entry, + distance: entry_dist, + }); + results.push(Candidate { + id: entry, + distance: entry_dist, + }); visited.insert(entry); while let Some(current) = candidates.pop() { @@ -399,8 +422,14 @@ impl GpuHnswIndex { let dist = self.distance(query, neighbor); if dist < worst_result || results.len() < ef { - candidates.push(Candidate { id: neighbor, distance: dist }); - results.push(Candidate { id: neighbor, distance: dist }); + candidates.push(Candidate { + id: neighbor, + distance: dist, + }); + results.push(Candidate { + id: neighbor, + distance: dist, + }); while results.len() > ef { results.pop(); @@ -410,7 +439,8 @@ impl GpuHnswIndex { } } - results.into_sorted_vec() + results + .into_sorted_vec() .into_iter() .map(|c| (c.id, c.distance)) .collect() @@ -418,18 +448,10 @@ impl GpuHnswIndex { fn distance(&self, query: &[f32], idx: usize) -> f32 { match self.params.metric { - GpuDistanceMetric::Cosine => { - 1.0 - cosine_similarity(query, &self.vectors[idx]) - } - GpuDistanceMetric::Euclidean => { - euclidean_distance(query, &self.vectors[idx]) - } - GpuDistanceMetric::DotProduct => { - -dot_product(query, &self.vectors[idx]) - } - GpuDistanceMetric::Manhattan => { - manhattan_distance(query, &self.vectors[idx]) - } + GpuDistanceMetric::Cosine => 1.0 - cosine_similarity(query, &self.vectors[idx]), + GpuDistanceMetric::Euclidean => euclidean_distance(query, &self.vectors[idx]), + GpuDistanceMetric::DotProduct => -dot_product(query, &self.vectors[idx]), + GpuDistanceMetric::Manhattan => manhattan_distance(query, &self.vectors[idx]), } } @@ -440,7 +462,9 @@ impl GpuHnswIndex { /// Get index statistics pub fn stats(&self) -> GpuIndexStats { - let total_edges: usize = self.layers.iter() + let total_edges: usize = self + .layers + .iter() .flat_map(|l| l.values()) .map(|neighbors| neighbors.len()) .sum(); @@ -562,7 +586,8 @@ impl GpuIndexBuilder { } // Entry point is the node with the highest layer - let entry_point = node_layers.iter() + let entry_point = node_layers + .iter() .enumerate() .max_by_key(|(_, l)| **l) .map(|(i, _)| i); @@ -576,9 +601,8 @@ impl GpuIndexBuilder { #[allow(clippy::needless_range_loop)] for layer in 0..=node_max_layer { // Find neighbors at this layer - let nodes_at_layer: Vec = (0..i) - .filter(|&j| node_layers[j] >= layer) - .collect(); + let nodes_at_layer: Vec = + (0..i).filter(|&j| node_layers[j] >= layer).collect(); if nodes_at_layer.is_empty() { continue; @@ -589,10 +613,16 @@ impl GpuIndexBuilder { .iter() .map(|&j| { let d = match params.metric { - GpuDistanceMetric::Cosine => 1.0 - cosine_similarity(&vectors[i], &vectors[j]), - GpuDistanceMetric::Euclidean => euclidean_distance(&vectors[i], &vectors[j]), + GpuDistanceMetric::Cosine => { + 1.0 - cosine_similarity(&vectors[i], &vectors[j]) + }, + GpuDistanceMetric::Euclidean => { + euclidean_distance(&vectors[i], &vectors[j]) + }, GpuDistanceMetric::DotProduct => -dot_product(&vectors[i], &vectors[j]), - GpuDistanceMetric::Manhattan => manhattan_distance(&vectors[i], &vectors[j]), + GpuDistanceMetric::Manhattan => { + manhattan_distance(&vectors[i], &vectors[j]) + }, }; (j, d) }) @@ -601,10 +631,8 @@ impl GpuIndexBuilder { distances.sort_by(|a, b| a.1.total_cmp(&b.1)); // Keep M neighbors - let neighbors: Vec = distances.iter() - .take(params.m) - .map(|(j, _)| *j) - .collect(); + let neighbors: Vec = + distances.iter().take(params.m).map(|(j, _)| *j).collect(); // Add bidirectional edges layers[layer].insert(i, neighbors.clone()); @@ -650,18 +678,23 @@ impl GpuIndexBuilder { } /// Compute batch distances on GPU - pub fn batch_distances(&self, queries: &[Vec], database: &[Vec], metric: GpuDistanceMetric) -> Vec> { + pub fn batch_distances( + &self, + queries: &[Vec], + database: &[Vec], + metric: GpuDistanceMetric, + ) -> Vec> { // In production, this would run on GPU - queries.iter() + queries + .iter() .map(|query| { - database.iter() - .map(|doc| { - match metric { - GpuDistanceMetric::Cosine => 1.0 - cosine_similarity(query, doc), - GpuDistanceMetric::Euclidean => euclidean_distance(query, doc), - GpuDistanceMetric::DotProduct => -dot_product(query, doc), - GpuDistanceMetric::Manhattan => manhattan_distance(query, doc), - } + database + .iter() + .map(|doc| match metric { + GpuDistanceMetric::Cosine => 1.0 - cosine_similarity(query, doc), + GpuDistanceMetric::Euclidean => euclidean_distance(query, doc), + GpuDistanceMetric::DotProduct => -dot_product(query, doc), + GpuDistanceMetric::Manhattan => manhattan_distance(query, doc), }) .collect() }) @@ -676,18 +709,16 @@ impl GpuIndexBuilder { /// Get available GPU devices pub fn available_devices() -> Vec { // In production, would actually enumerate devices - vec![ - GpuDeviceInfo { - index: 0, - name: "CPU (Fallback)".to_string(), - backend: GpuBackendType::CPU, - memory_bytes: 16 * 1024 * 1024 * 1024, // 16GB - available_memory_bytes: 8 * 1024 * 1024 * 1024, - compute_capability: None, - max_work_group_size: 1024, - max_shared_memory: 48 * 1024, - } - ] + vec![GpuDeviceInfo { + index: 0, + name: "CPU (Fallback)".to_string(), + backend: GpuBackendType::CPU, + memory_bytes: 16 * 1024 * 1024 * 1024, // 16GB + available_memory_bytes: 8 * 1024 * 1024 * 1024, + compute_capability: None, + max_work_group_size: 1024, + max_shared_memory: 48 * 1024, + }] } } @@ -753,9 +784,7 @@ mod tests { let builder = GpuIndexBuilder::new(config); let vectors: Vec> = (0..100) - .map(|i| { - (0..64).map(|j| ((i + j) as f32) / 100.0).collect() - }) + .map(|i| (0..64).map(|j| ((i + j) as f32) / 100.0).collect()) .collect(); let params = HnswParams { @@ -777,11 +806,7 @@ mod tests { let config = GpuConfig::default(); let builder = GpuIndexBuilder::new(config); - let vectors: Vec> = (0..50) - .map(|i| { - vec![i as f32 / 50.0; 32] - }) - .collect(); + let vectors: Vec> = (0..50).map(|i| vec![i as f32 / 50.0; 32]).collect(); let params = HnswParams::default(); let index = builder.build_hnsw(&vectors, params).unwrap(); @@ -799,10 +824,7 @@ mod tests { let config = GpuConfig::default(); let builder = GpuIndexBuilder::new(config); - let queries = vec![ - vec![1.0, 0.0, 0.0], - vec![0.0, 1.0, 0.0], - ]; + let queries = vec![vec![1.0, 0.0, 0.0], vec![0.0, 1.0, 0.0]]; let database = vec![ vec![1.0, 0.0, 0.0], @@ -831,9 +853,7 @@ mod tests { let config = GpuConfig::default(); let builder = GpuIndexBuilder::new(config); - let vectors: Vec> = (0..20) - .map(|i| vec![i as f32; 16]) - .collect(); + let vectors: Vec> = (0..20).map(|i| vec![i as f32; 16]).collect(); let index = builder.build_hnsw(&vectors, HnswParams::default()).unwrap(); let stats = index.build_stats(); diff --git a/src/graph_rag.rs b/src/graph_rag.rs index ca92a06..060c58d 100644 --- a/src/graph_rag.rs +++ b/src/graph_rag.rs @@ -63,7 +63,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet, VecDeque}; @@ -330,10 +330,7 @@ impl GraphRAG { .push(relation.clone()); // Add to incoming edges - self.incoming - .entry(to) - .or_default() - .push(relation); + self.incoming.entry(to).or_default().push(relation); Ok(()) } @@ -358,10 +355,7 @@ impl GraphRAG { .push(relation.clone()); // Add to incoming edges - self.incoming - .entry(to) - .or_default() - .push(relation); + self.incoming.entry(to).or_default().push(relation); Ok(()) } @@ -456,9 +450,10 @@ impl GraphRAG { } if let Some(ref filter) = query.relation_type_filter - && !filter.contains(&relation.relation_type) { - continue; - } + && !filter.contains(&relation.relation_type) + { + continue; + } if !visited.contains(&relation.to) { subgraph.push(relation.to.clone()); diff --git a/src/graph_vector.rs b/src/graph_vector.rs index 46c645f..bad6ddd 100644 --- a/src/graph_vector.rs +++ b/src/graph_vector.rs @@ -36,9 +36,9 @@ //! )?; //! ``` -use std::collections::{HashMap, HashSet, BinaryHeap}; -use std::cmp::Ordering; use serde::{Deserialize, Serialize}; +use std::cmp::Ordering; +use std::collections::{BinaryHeap, HashMap, HashSet}; use crate::error::VecStoreError; @@ -289,7 +289,8 @@ impl GraphVectorEngine { if self.nodes.contains_key(&id) { return Err(VecStoreError::InvalidInput(format!( - "Node {} already exists", id + "Node {} already exists", + id ))); } @@ -319,19 +320,24 @@ impl GraphVectorEngine { // Validate nodes exist if !self.nodes.contains_key(&edge.from) { return Err(VecStoreError::NotFound(format!( - "Source node {} not found", edge.from + "Source node {} not found", + edge.from ))); } if !self.nodes.contains_key(&edge.to) { return Err(VecStoreError::NotFound(format!( - "Target node {} not found", edge.to + "Target node {} not found", + edge.to ))); } let edge_idx = self.edges.len(); // Update adjacency lists - self.adjacency_out.get_mut(&edge.from).unwrap().push(edge_idx); + self.adjacency_out + .get_mut(&edge.from) + .unwrap() + .push(edge_idx); self.adjacency_in.get_mut(&edge.to).unwrap().push(edge_idx); // Update node references @@ -387,9 +393,10 @@ impl GraphVectorEngine { while let Some(current) = heap.pop() { // Skip if we've seen this node with a better score if let Some(existing) = visited.get(¤t.node_id) - && existing.score >= current.score { - continue; - } + && existing.score >= current.score + { + continue; + } // Get node info let node = match self.nodes.get(¤t.node_id) { @@ -401,8 +408,8 @@ impl GraphVectorEngine { let vector_score = self.cosine_similarity(query_vector, &node.vector); // Combined score - let combined_score = vector_score * config.vector_weight - + current.score * config.edge_weight; + let combined_score = + vector_score * config.vector_weight + current.score * config.edge_weight; if combined_score < config.min_score { continue; @@ -430,8 +437,7 @@ impl GraphVectorEngine { let edge = &self.edges[edge_idx]; // Check edge type filter - if !config.edge_types.is_empty() - && !config.edge_types.contains(&edge.edge_type) + if !config.edge_types.is_empty() && !config.edge_types.contains(&edge.edge_type) { continue; } @@ -469,7 +475,8 @@ impl GraphVectorEngine { } // Filter and sort results - let mut results: Vec = visited.into_values() + let mut results: Vec = visited + .into_values() .filter(|r| config.include_start || r.hops > 0) .filter(|r| r.score >= config.min_score) .collect(); @@ -489,7 +496,8 @@ impl GraphVectorEngine { ) -> Result, VecStoreError> { if !self.nodes.contains_key(start_id) { return Err(VecStoreError::NotFound(format!( - "Node {} not found", start_id + "Node {} not found", + start_id ))); } @@ -527,7 +535,11 @@ impl GraphVectorEngine { } /// Find the shortest path between two nodes - pub fn shortest_path(&self, from: &str, to: &str) -> Result>, VecStoreError> { + pub fn shortest_path( + &self, + from: &str, + to: &str, + ) -> Result>, VecStoreError> { if !self.nodes.contains_key(from) { return Err(VecStoreError::NotFound(format!("Node {} not found", from))); } @@ -587,7 +599,10 @@ impl GraphVectorEngine { direction: TraversalDirection, ) -> Result, VecStoreError> { if !self.nodes.contains_key(node_id) { - return Err(VecStoreError::NotFound(format!("Node {} not found", node_id))); + return Err(VecStoreError::NotFound(format!( + "Node {} not found", + node_id + ))); } let edges = self.get_edges_for_traversal(node_id, &direction); @@ -601,11 +616,7 @@ impl GraphVectorEngine { &edge.from }; - neighbors.push(( - neighbor_id.clone(), - edge.edge_type.clone(), - edge.weight, - )); + neighbors.push((neighbor_id.clone(), edge.edge_type.clone(), edge.weight)); } Ok(neighbors) @@ -668,7 +679,9 @@ impl GraphVectorEngine { } // Mark edges for removal (edges involving this node) - let edges_to_remove: Vec = self.edges.iter() + let edges_to_remove: Vec = self + .edges + .iter() .enumerate() .filter(|(_, e)| e.from == id || e.to == id) .map(|(i, _)| i) @@ -692,9 +705,10 @@ impl GraphVectorEngine { /// Update node vector pub fn update_vector(&mut self, id: &str, new_vector: Vec) -> Result<(), VecStoreError> { - let node = self.nodes.get_mut(id).ok_or_else(|| { - VecStoreError::NotFound(format!("Node {} not found", id)) - })?; + let node = self + .nodes + .get_mut(id) + .ok_or_else(|| VecStoreError::NotFound(format!("Node {} not found", id)))?; node.vector = new_vector; Ok(()) @@ -708,7 +722,9 @@ impl GraphVectorEngine { // === Helper Methods === fn find_similar_nodes(&self, query: &[f32], limit: usize) -> Vec<(String, f32)> { - let mut scores: Vec<_> = self.nodes.iter() + let mut scores: Vec<_> = self + .nodes + .iter() .map(|(id, node)| { let sim = self.cosine_similarity(query, &node.vector); (id.clone(), sim) @@ -728,12 +744,12 @@ impl GraphVectorEngine { if let Some(out_edges) = self.adjacency_out.get(node_id) { edges.extend(out_edges); } - } + }, TraversalDirection::Incoming => { if let Some(in_edges) = self.adjacency_in.get(node_id) { edges.extend(in_edges); } - } + }, TraversalDirection::Both => { if let Some(out_edges) = self.adjacency_out.get(node_id) { edges.extend(out_edges); @@ -741,7 +757,7 @@ impl GraphVectorEngine { if let Some(in_edges) = self.adjacency_in.get(node_id) { edges.extend(in_edges); } - } + }, } edges @@ -900,15 +916,29 @@ mod tests { let mut engine = GraphVectorEngine::new(); // Add nodes - engine.add_node(GraphNode::new("a", vec![1.0, 0.0, 0.0])).unwrap(); - engine.add_node(GraphNode::new("b", vec![0.9, 0.1, 0.0])).unwrap(); - engine.add_node(GraphNode::new("c", vec![0.0, 1.0, 0.0])).unwrap(); - engine.add_node(GraphNode::new("d", vec![0.0, 0.0, 1.0])).unwrap(); + engine + .add_node(GraphNode::new("a", vec![1.0, 0.0, 0.0])) + .unwrap(); + engine + .add_node(GraphNode::new("b", vec![0.9, 0.1, 0.0])) + .unwrap(); + engine + .add_node(GraphNode::new("c", vec![0.0, 1.0, 0.0])) + .unwrap(); + engine + .add_node(GraphNode::new("d", vec![0.0, 0.0, 1.0])) + .unwrap(); // Add edges: a -> b -> c, a -> d - engine.add_edge_simple("a", "b", EdgeType::References, 1.0).unwrap(); - engine.add_edge_simple("b", "c", EdgeType::References, 0.8).unwrap(); - engine.add_edge_simple("a", "d", EdgeType::RelatedTo, 0.5).unwrap(); + engine + .add_edge_simple("a", "b", EdgeType::References, 1.0) + .unwrap(); + engine + .add_edge_simple("b", "c", EdgeType::References, 0.8) + .unwrap(); + engine + .add_edge_simple("a", "d", EdgeType::RelatedTo, 0.5) + .unwrap(); engine } @@ -929,7 +959,9 @@ mod tests { fn test_get_neighbors() { let engine = create_test_graph(); - let neighbors = engine.get_neighbors("a", TraversalDirection::Outgoing).unwrap(); + let neighbors = engine + .get_neighbors("a", TraversalDirection::Outgoing) + .unwrap(); assert_eq!(neighbors.len(), 2); let neighbor_ids: Vec<_> = neighbors.iter().map(|(id, _, _)| id.as_str()).collect(); @@ -941,7 +973,9 @@ mod tests { fn test_find_reachable() { let engine = create_test_graph(); - let reachable = engine.find_reachable("a", 2, TraversalDirection::Outgoing).unwrap(); + let reachable = engine + .find_reachable("a", 2, TraversalDirection::Outgoing) + .unwrap(); // Should find a (0 hops), b and d (1 hop), c (2 hops) assert_eq!(reachable.len(), 4); @@ -990,14 +1024,26 @@ mod tests { let mut engine = GraphVectorEngine::new(); // Component 1: a - b - engine.add_node(GraphNode::new("a", vec![1.0, 0.0])).unwrap(); - engine.add_node(GraphNode::new("b", vec![0.0, 1.0])).unwrap(); - engine.add_edge_simple("a", "b", EdgeType::RelatedTo, 1.0).unwrap(); + engine + .add_node(GraphNode::new("a", vec![1.0, 0.0])) + .unwrap(); + engine + .add_node(GraphNode::new("b", vec![0.0, 1.0])) + .unwrap(); + engine + .add_edge_simple("a", "b", EdgeType::RelatedTo, 1.0) + .unwrap(); // Component 2: c - d (disconnected) - engine.add_node(GraphNode::new("c", vec![0.5, 0.5])).unwrap(); - engine.add_node(GraphNode::new("d", vec![0.5, 0.5])).unwrap(); - engine.add_edge_simple("c", "d", EdgeType::RelatedTo, 1.0).unwrap(); + engine + .add_node(GraphNode::new("c", vec![0.5, 0.5])) + .unwrap(); + engine + .add_node(GraphNode::new("d", vec![0.5, 0.5])) + .unwrap(); + engine + .add_edge_simple("c", "d", EdgeType::RelatedTo, 1.0) + .unwrap(); let stats = engine.stats(); assert_eq!(stats.component_count, 2); diff --git a/src/graphql_api.rs b/src/graphql_api.rs index 4067b55..c7ba948 100644 --- a/src/graphql_api.rs +++ b/src/graphql_api.rs @@ -32,11 +32,11 @@ //! } //! ``` -use std::collections::HashMap; use serde::{Deserialize, Serialize}; use serde_json::Value as JsonValue; +use std::collections::HashMap; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; // ============================================================================ // SCHEMA TYPES @@ -94,10 +94,17 @@ impl GqlCollection { pub fn to_sdl(&self) -> String { let mut sdl = format!("type {} {{\n", self.name); for field in &self.fields { - let desc = field.description.as_ref() + let desc = field + .description + .as_ref() .map(|d| format!(" \"{}\"\n", d)) .unwrap_or_default(); - sdl.push_str(&format!("{} {}: {}\n", desc, field.name, field.field_type.to_sdl())); + sdl.push_str(&format!( + "{} {}: {}\n", + desc, + field.name, + field.field_type.to_sdl() + )); } sdl.push_str(" _additional: _Additional\n"); sdl.push_str("}\n"); @@ -257,7 +264,9 @@ impl GraphQLExecutor { vector: Vec, properties: HashMap, ) -> Result<()> { - let col = self.collections.get(collection) + let col = self + .collections + .get(collection) .ok_or_else(|| VecStoreError::NotFound(format!("Collection: {}", collection)))?; if vector.len() != col.dimension { @@ -277,7 +286,9 @@ impl GraphQLExecutor { /// Execute a vector query pub fn execute(&self, query: VectorQuery) -> Result { - let _collection = self.collections.get(&query.collection) + let _collection = self + .collections + .get(&query.collection) .ok_or_else(|| VecStoreError::NotFound(format!("Collection: {}", query.collection)))?; let vectors = self.vectors.get(&query.collection).unwrap(); @@ -288,13 +299,14 @@ impl GraphQLExecutor { if let Some(near_vector) = &query.near_vector { for (id, (vec, props)) in vectors { let similarity = cosine_similarity(&near_vector.vector, vec); - let certainty = (similarity + 1.0) / 2.0; // Convert to 0-1 range + let certainty = (similarity + 1.0) / 2.0; // Convert to 0-1 range // Check certainty threshold if let Some(min_cert) = near_vector.certainty - && certainty < min_cert { - continue; - } + && certainty < min_cert + { + continue; + } // Check distance threshold if let Some(max_dist) = near_vector.distance { @@ -329,7 +341,8 @@ impl GraphQLExecutor { results = results.into_iter().skip(offset).take(limit).collect(); // Build result items - let items: Vec = results.into_iter() + let items: Vec = results + .into_iter() .map(|(id, score, props)| { let mut properties = HashMap::new(); for field in &query.fields { @@ -348,7 +361,10 @@ impl GraphQLExecutor { explanation: None, }; - QueryResultItem { properties, additional } + QueryResultItem { + properties, + additional, + } }) .collect(); @@ -368,57 +384,53 @@ impl GraphQLExecutor { } else { true } - } + }, WhereOperator::Or => { if let Some(operands) = &filter.operands { operands.iter().any(|f| self.evaluate_filter(f, props)) } else { false } - } + }, WhereOperator::Equal => { if let Some(path) = &filter.path && let Some(field) = path.first() - && let Some(value) = props.get(field) { - return self.compare_equal(value, filter); - } + && let Some(value) = props.get(field) + { + return self.compare_equal(value, filter); + } false - } + }, WhereOperator::NotEqual => { if let Some(path) = &filter.path && let Some(field) = path.first() - && let Some(value) = props.get(field) { - return !self.compare_equal(value, filter); - } + && let Some(value) = props.get(field) + { + return !self.compare_equal(value, filter); + } true - } - WhereOperator::GreaterThan => { - self.compare_numeric(filter, props, |a, b| a > b) - } - WhereOperator::GreaterThanEqual => { - self.compare_numeric(filter, props, |a, b| a >= b) - } - WhereOperator::LessThan => { - self.compare_numeric(filter, props, |a, b| a < b) - } - WhereOperator::LessThanEqual => { - self.compare_numeric(filter, props, |a, b| a <= b) - } + }, + WhereOperator::GreaterThan => self.compare_numeric(filter, props, |a, b| a > b), + WhereOperator::GreaterThanEqual => self.compare_numeric(filter, props, |a, b| a >= b), + WhereOperator::LessThan => self.compare_numeric(filter, props, |a, b| a < b), + WhereOperator::LessThanEqual => self.compare_numeric(filter, props, |a, b| a <= b), WhereOperator::Like => { if let (Some(path), Some(pattern)) = (&filter.path, &filter.value_string) && let Some(field) = path.first() - && let Some(JsonValue::String(s)) = props.get(field) { - return self.match_like(s, pattern); - } + && let Some(JsonValue::String(s)) = props.get(field) + { + return self.match_like(s, pattern); + } false - } + }, WhereOperator::IsNull => { if let Some(path) = &filter.path - && let Some(field) = path.first() { - return props.get(field).map(|v| v.is_null()).unwrap_or(true); - } + && let Some(field) = path.first() + { + return props.get(field).map(|v| v.is_null()).unwrap_or(true); + } false - } + }, _ => true, } } @@ -439,21 +451,27 @@ impl GraphQLExecutor { false } - fn compare_numeric(&self, filter: &WhereFilter, props: &HashMap, cmp: F) -> bool + fn compare_numeric( + &self, + filter: &WhereFilter, + props: &HashMap, + cmp: F, + ) -> bool where F: Fn(f64, f64) -> bool, { if let Some(path) = &filter.path && let Some(field) = path.first() - && let Some(value) = props.get(field) - && let Some(v) = value.as_f64() { - if let Some(n) = filter.value_number { - return cmp(v, n); - } - if let Some(i) = filter.value_int { - return cmp(v, i as f64); - } - } + && let Some(value) = props.get(field) + && let Some(v) = value.as_f64() + { + if let Some(n) = filter.value_number { + return cmp(v, n); + } + if let Some(i) = filter.value_int { + return cmp(v, i as f64); + } + } false } @@ -658,7 +676,8 @@ fn extract_collection(query: &str) -> Result { let rest = &query[start..]; if let Some(brace) = rest.find('{') { let after_brace = &rest[brace + 1..]; - let name: String = after_brace.trim() + let name: String = after_brace + .trim() .chars() .take_while(|c| c.is_alphanumeric() || *c == '_') .collect(); @@ -667,7 +686,9 @@ fn extract_collection(query: &str) -> Result { } } } - Err(VecStoreError::InvalidInput("Could not parse collection name".to_string())) + Err(VecStoreError::InvalidInput( + "Could not parse collection name".to_string(), + )) } fn extract_near_vector(query: &str) -> Option { @@ -675,19 +696,22 @@ fn extract_near_vector(query: &str) -> Option { // Very simplified - would need proper parsing if let Some(vec_start) = query[start..].find("vector:") && let Some(bracket_start) = query[start + vec_start..].find('[') - && let Some(bracket_end) = query[start + vec_start + bracket_start..].find(']') { - let vec_str = &query[start + vec_start + bracket_start + 1..start + vec_start + bracket_start + bracket_end]; - let vector: Vec = vec_str.split(',') - .filter_map(|s| s.trim().parse().ok()) - .collect(); - if !vector.is_empty() { - return Some(NearVectorInput { - vector, - certainty: extract_float(query, "certainty:"), - distance: extract_float(query, "distance:"), - }); - } - } + && let Some(bracket_end) = query[start + vec_start + bracket_start..].find(']') + { + let vec_str = &query[start + vec_start + bracket_start + 1 + ..start + vec_start + bracket_start + bracket_end]; + let vector: Vec = vec_str + .split(',') + .filter_map(|s| s.trim().parse().ok()) + .collect(); + if !vector.is_empty() { + return Some(NearVectorInput { + vector, + certainty: extract_float(query, "certainty:"), + distance: extract_float(query, "distance:"), + }); + } + } } None } @@ -695,7 +719,8 @@ fn extract_near_vector(query: &str) -> Option { fn extract_float(query: &str, key: &str) -> Option { if let Some(start) = query.find(key) { let rest = &query[start + key.len()..]; - let value: String = rest.trim() + let value: String = rest + .trim() .chars() .take_while(|c| c.is_numeric() || *c == '.' || *c == '-') .collect(); @@ -713,10 +738,7 @@ fn extract_where(_query: &str) -> Option { fn extract_limit(query: &str) -> Option { if let Some(start) = query.find("limit:") { let rest = &query[start + 6..]; - let value: String = rest.trim() - .chars() - .take_while(|c| c.is_numeric()) - .collect(); + let value: String = rest.trim().chars().take_while(|c| c.is_numeric()).collect(); value.parse().ok() } else { None @@ -778,14 +800,30 @@ mod tests { // Insert vectors let mut props1 = HashMap::new(); - props1.insert("title".to_string(), JsonValue::String("Tech News".to_string())); - props1.insert("category".to_string(), JsonValue::String("tech".to_string())); - executor.insert("Article", "doc1", vec![1.0, 0.0, 0.0, 0.0], props1).unwrap(); + props1.insert( + "title".to_string(), + JsonValue::String("Tech News".to_string()), + ); + props1.insert( + "category".to_string(), + JsonValue::String("tech".to_string()), + ); + executor + .insert("Article", "doc1", vec![1.0, 0.0, 0.0, 0.0], props1) + .unwrap(); let mut props2 = HashMap::new(); - props2.insert("title".to_string(), JsonValue::String("Sports News".to_string())); - props2.insert("category".to_string(), JsonValue::String("sports".to_string())); - executor.insert("Article", "doc2", vec![0.0, 1.0, 0.0, 0.0], props2).unwrap(); + props2.insert( + "title".to_string(), + JsonValue::String("Sports News".to_string()), + ); + props2.insert( + "category".to_string(), + JsonValue::String("sports".to_string()), + ); + executor + .insert("Article", "doc2", vec![0.0, 1.0, 0.0, 0.0], props2) + .unwrap(); // Execute query let query = VectorQuery { @@ -815,14 +853,12 @@ mod tests { let collection = GqlCollection { name: "Document".to_string(), description: None, - fields: vec![ - GqlField { - name: "content".to_string(), - field_type: GqlType::String, - description: None, - is_vector: false, - }, - ], + fields: vec![GqlField { + name: "content".to_string(), + field_type: GqlType::String, + description: None, + is_vector: false, + }], vector_field: None, dimension: 4, }; diff --git a/src/graphrag.rs b/src/graphrag.rs index 7773600..1d218aa 100644 --- a/src/graphrag.rs +++ b/src/graphrag.rs @@ -31,10 +31,10 @@ //! let results = graph_rag.query("vector database languages", 5)?; //! ``` -use std::collections::{HashMap, HashSet, VecDeque}; use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet, VecDeque}; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; /// Entity in the knowledge graph #[derive(Debug, Clone, Serialize, Deserialize)] @@ -187,11 +187,21 @@ pub struct GraphRAGConfig { pub use_entity_embeddings: bool, } -fn default_vector_weight() -> f32 { 0.7 } -fn default_graph_weight() -> f32 { 0.3 } -fn default_max_depth() -> usize { 2 } -fn default_max_expansion() -> usize { 10 } -fn default_true() -> bool { true } +fn default_vector_weight() -> f32 { + 0.7 +} +fn default_graph_weight() -> f32 { + 0.3 +} +fn default_max_depth() -> usize { + 2 +} +fn default_max_expansion() -> usize { + 10 +} +fn default_true() -> bool { + true +} impl Default for GraphRAGConfig { fn default() -> Self { @@ -250,10 +260,16 @@ impl GraphRAG { pub fn add_relationship(&mut self, rel: Relationship) -> Result<()> { // Verify entities exist if !self.entities.contains_key(&rel.from) { - return Err(VecStoreError::NotFound(format!("Entity not found: {}", rel.from))); + return Err(VecStoreError::NotFound(format!( + "Entity not found: {}", + rel.from + ))); } if !self.entities.contains_key(&rel.to) { - return Err(VecStoreError::NotFound(format!("Entity not found: {}", rel.to))); + return Err(VecStoreError::NotFound(format!( + "Entity not found: {}", + rel.to + ))); } // Add to outgoing edges @@ -263,10 +279,7 @@ impl GraphRAG { .push(rel.clone()); // Add to incoming edges - self.incoming - .entry(rel.to.clone()) - .or_default() - .push(rel); + self.incoming.entry(rel.to.clone()).or_default().push(rel); Ok(()) } @@ -287,7 +300,8 @@ impl GraphRAG { }); } - let linked_entities: HashSet = entity_ids.iter() + let linked_entities: HashSet = entity_ids + .iter() .filter(|&&id| self.entities.contains_key(id)) .map(|&s| s.to_string()) .collect(); @@ -300,23 +314,22 @@ impl GraphRAG { .insert(id.to_string()); } - self.documents.insert(id.to_string(), LinkedDocument { - id: id.to_string(), - content: content.to_string(), - embedding, - linked_entities, - metadata, - }); + self.documents.insert( + id.to_string(), + LinkedDocument { + id: id.to_string(), + content: content.to_string(), + embedding, + linked_entities, + metadata, + }, + ); Ok(()) } /// Query with graph expansion - pub fn query( - &self, - query_embedding: &[f32], - limit: usize, - ) -> Result> { + pub fn query(&self, query_embedding: &[f32], limit: usize) -> Result> { if query_embedding.len() != self.dimension { return Err(VecStoreError::DimensionMismatch { expected: self.dimension, @@ -335,7 +348,8 @@ impl GraphRAG { let expanded_entities = self.expand_entities(&relevant_entities); // Step 3: Find documents linked to expanded entities - let entity_docs: HashSet<&String> = expanded_entities.iter() + let entity_docs: HashSet<&String> = expanded_entities + .iter() .filter_map(|e| self.entity_documents.get(&e.id)) .flatten() .collect(); @@ -348,7 +362,9 @@ impl GraphRAG { // Calculate graph score based on entity overlap let graph_score = if entity_docs.contains(doc_id) { - let overlap = doc.linked_entities.iter() + let overlap = doc + .linked_entities + .iter() .filter(|e| expanded_entities.iter().any(|exp| exp.id == **e)) .count(); overlap as f32 / doc.linked_entities.len().max(1) as f32 @@ -356,16 +372,19 @@ impl GraphRAG { 0.0 }; - let combined_score = self.config.vector_weight * vector_score - + self.config.graph_weight * graph_score; + let combined_score = + self.config.vector_weight * vector_score + self.config.graph_weight * graph_score; // Get linked entities - let linked: Vec = doc.linked_entities.iter() + let linked: Vec = doc + .linked_entities + .iter() .filter_map(|id| self.entities.get(id).cloned()) .collect(); // Get related entities (from expansion, not directly linked) - let related: Vec = expanded_entities.iter() + let related: Vec = expanded_entities + .iter() .filter(|e| !doc.linked_entities.contains(&e.id)) .cloned() .collect(); @@ -394,16 +413,19 @@ impl GraphRAG { /// Find entities similar to query embedding fn find_relevant_entities(&self, query: &[f32], limit: usize) -> Vec { - let mut scored: Vec<(f32, &Entity)> = self.entities.values() + let mut scored: Vec<(f32, &Entity)> = self + .entities + .values() .filter_map(|e| { - e.embedding.as_ref().map(|emb| { - (Self::cosine_similarity(query, emb), e) - }) + e.embedding + .as_ref() + .map(|emb| (Self::cosine_similarity(query, emb), e)) }) .collect(); scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); - scored.into_iter() + scored + .into_iter() .take(limit) .map(|(_, e)| e.clone()) .collect() @@ -436,10 +458,11 @@ impl GraphRAG { if let Some(rels) = self.outgoing.get(&entity_id) { for rel in rels { if visited.insert(rel.to.clone()) - && let Some(entity) = self.entities.get(&rel.to) { - result.push(entity.clone()); - queue.push_back((rel.to.clone(), depth + 1)); - } + && let Some(entity) = self.entities.get(&rel.to) + { + result.push(entity.clone()); + queue.push_back((rel.to.clone(), depth + 1)); + } } } @@ -447,10 +470,11 @@ impl GraphRAG { if let Some(rels) = self.incoming.get(&entity_id) { for rel in rels { if visited.insert(rel.from.clone()) - && let Some(entity) = self.entities.get(&rel.from) { - result.push(entity.clone()); - queue.push_back((rel.from.clone(), depth + 1)); - } + && let Some(entity) = self.entities.get(&rel.from) + { + result.push(entity.clone()); + queue.push_back((rel.from.clone(), depth + 1)); + } } } } @@ -495,21 +519,24 @@ impl GraphRAG { /// Get all entities of a type pub fn get_entities_by_type(&self, entity_type: &str) -> Vec<&Entity> { - self.entities.values() + self.entities + .values() .filter(|e| e.entity_type == entity_type) .collect() } /// Get relationships from an entity pub fn get_outgoing(&self, entity_id: &str) -> Vec<&Relationship> { - self.outgoing.get(entity_id) + self.outgoing + .get(entity_id) .map(|rels| rels.iter().collect()) .unwrap_or_default() } /// Get relationships to an entity pub fn get_incoming(&self, entity_id: &str) -> Vec<&Relationship> { - self.incoming.get(entity_id) + self.incoming + .get(entity_id) .map(|rels| rels.iter().collect()) .unwrap_or_default() } @@ -541,7 +568,9 @@ mod tests { let mut graph = GraphRAG::new(64).unwrap(); graph.add_entity(Entity::new("rust", "Language")).unwrap(); - graph.add_entity(Entity::new("vecstore", "Database")).unwrap(); + graph + .add_entity(Entity::new("vecstore", "Database")) + .unwrap(); assert!(graph.get_entity("rust").is_some()); assert!(graph.get_entity("vecstore").is_some()); @@ -552,11 +581,13 @@ mod tests { let mut graph = GraphRAG::new(64).unwrap(); graph.add_entity(Entity::new("rust", "Language")).unwrap(); - graph.add_entity(Entity::new("vecstore", "Database")).unwrap(); + graph + .add_entity(Entity::new("vecstore", "Database")) + .unwrap(); - graph.add_relationship( - Relationship::new("vecstore", "rust", "written_in") - ).unwrap(); + graph + .add_relationship(Relationship::new("vecstore", "rust", "written_in")) + .unwrap(); let outgoing = graph.get_outgoing("vecstore"); assert_eq!(outgoing.len(), 1); @@ -568,16 +599,20 @@ mod tests { let mut graph = GraphRAG::new(64).unwrap(); graph.add_entity(Entity::new("rust", "Language")).unwrap(); - graph.add_entity(Entity::new("vecstore", "Database")).unwrap(); + graph + .add_entity(Entity::new("vecstore", "Database")) + .unwrap(); let embedding = vec![0.1f32; 64]; - graph.add_document( - "doc1", - "VecStore is written in Rust", - embedding.clone(), - &["rust", "vecstore"], - None, - ).unwrap(); + graph + .add_document( + "doc1", + "VecStore is written in Rust", + embedding.clone(), + &["rust", "vecstore"], + None, + ) + .unwrap(); let results = graph.query(&embedding, 5).unwrap(); assert_eq!(results.len(), 1); diff --git a/src/health.rs b/src/health.rs index e240204..406097d 100644 --- a/src/health.rs +++ b/src/health.rs @@ -363,13 +363,13 @@ impl HealthChecker { // Calculate memory utilization // Estimate total system memory (could be made configurable) let total_memory_estimate = 16 * 1024 * 1024 * 1024usize; // 16 GB default - let memory_utilization = (memory_bytes as f64 / total_memory_estimate as f64 * 100.0) - .clamp(0.0, 100.0); + let memory_utilization = + (memory_bytes as f64 / total_memory_estimate as f64 * 100.0).clamp(0.0, 100.0); // Calculate disk utilization (assuming 1TB max for simplicity) let total_disk_estimate = 1024 * 1024 * 1024 * 1024usize; // 1 TB - let disk_utilization = (disk_bytes as f64 / total_disk_estimate as f64 * 100.0) - .clamp(0.0, 100.0); + let disk_utilization = + (disk_bytes as f64 / total_disk_estimate as f64 * 100.0).clamp(0.0, 100.0); ResourceHealth { memory_bytes, @@ -482,17 +482,18 @@ impl HealthChecker { fn generate_performance_alerts(&self, perf: &PerformanceHealth, alerts: &mut Vec) { if let Some(latency) = perf.p95_query_latency_ms - && latency >= self.config.latency_warning_ms { - alerts.push(Alert { - severity: AlertSeverity::Warning, - category: AlertCategory::Performance, - message: format!("High query latency: {:.2}ms (p95)", latency), - value: Some(latency), - recommendation: Some( - "Check index parameters or consider upgrading hardware".to_string(), - ), - }); - } + && latency >= self.config.latency_warning_ms + { + alerts.push(Alert { + severity: AlertSeverity::Warning, + category: AlertCategory::Performance, + message: format!("High query latency: {:.2}ms (p95)", latency), + value: Some(latency), + recommendation: Some( + "Check index parameters or consider upgrading hardware".to_string(), + ), + }); + } if perf.performance_score < self.config.min_performance_score { alerts.push(Alert { @@ -633,12 +634,7 @@ pub fn print_health_report(report: &HealthReport) { AlertSeverity::Warning => "āš ļø", AlertSeverity::Critical => "āŒ", }; - println!( - " {} [{:?}] {}", - icon, - alert.category, - alert.message - ); + println!(" {} [{:?}] {}", icon, alert.category, alert.message); if let Some(rec) = &alert.recommendation { println!(" → {}", rec); } @@ -696,10 +692,12 @@ mod tests { // 60% deletion ratio should trigger Unhealthy status assert_eq!(report.status, HealthStatus::Unhealthy); - assert!(report - .alerts - .iter() - .any(|a| a.severity == AlertSeverity::Critical)); + assert!( + report + .alerts + .iter() + .any(|a| a.severity == AlertSeverity::Critical) + ); Ok(()) } diff --git a/src/hybrid_rerank.rs b/src/hybrid_rerank.rs index 6358222..152e5fc 100644 --- a/src/hybrid_rerank.rs +++ b/src/hybrid_rerank.rs @@ -464,7 +464,13 @@ impl InvertedIndex { idf } - fn search_bm25(&mut self, query_terms: &[String], k1: f32, b: f32, top_k: usize) -> Vec<(String, f32)> { + fn search_bm25( + &mut self, + query_terms: &[String], + k1: f32, + b: f32, + top_k: usize, + ) -> Vec<(String, f32)> { let mut scores: HashMap = HashMap::new(); for term in query_terms { @@ -546,7 +552,8 @@ impl DenseIndex { } fn search(&self, query: &[f32], top_k: usize) -> Vec<(String, f32)> { - let mut scores: Vec<_> = self.vectors + let mut scores: Vec<_> = self + .vectors .iter() .map(|(id, vec)| { let score = cosine_similarity(query, vec); @@ -594,14 +601,18 @@ impl HybridPipeline { pub fn index(&self, doc: HybridDocument) -> Result<()> { // Index in dense index { - let mut dense = self.dense_index.write() + let mut dense = self + .dense_index + .write() .map_err(|_| VecStoreError::LockError("dense_index lock poisoned".into()))?; dense.add_document(&doc.id, &doc.dense_vector); } // Index in sparse index if available if let Some(sparse) = &doc.sparse_vector { - let mut sparse_idx = self.sparse_index.write() + let mut sparse_idx = self + .sparse_index + .write() .map_err(|_| VecStoreError::LockError("sparse_index lock poisoned".into()))?; sparse_idx.add_document(&doc.id, sparse); } @@ -609,14 +620,18 @@ impl HybridPipeline { // Index in BM25 { let terms = tokenize(&doc.text); - let mut bm25 = self.bm25_index.write() + let mut bm25 = self + .bm25_index + .write() .map_err(|_| VecStoreError::LockError("bm25_index lock poisoned".into()))?; bm25.add_document(&doc.id, &terms); } // Store document { - let mut docs = self.documents.write() + let mut docs = self + .documents + .write() .map_err(|_| VecStoreError::LockError("documents lock poisoned".into()))?; docs.insert(doc.id.clone(), doc); } @@ -633,46 +648,54 @@ impl HybridPipeline { // Dense search if let Some(ref vector) = query.dense_vector - && weights.0 > 0.0 { - let dense = self.dense_index.read() - .map_err(|_| VecStoreError::LockError("dense_index lock poisoned".into()))?; - let results = dense.search(vector, self.config.max_candidates); - let signal_results: Vec = results - .into_iter() - .enumerate() - .map(|(rank, (id, score))| SignalResult { - id, - score, - rank, - signal: SignalType::Dense, - }) - .collect(); - all_signals.push(signal_results); - } + && weights.0 > 0.0 + { + let dense = self + .dense_index + .read() + .map_err(|_| VecStoreError::LockError("dense_index lock poisoned".into()))?; + let results = dense.search(vector, self.config.max_candidates); + let signal_results: Vec = results + .into_iter() + .enumerate() + .map(|(rank, (id, score))| SignalResult { + id, + score, + rank, + signal: SignalType::Dense, + }) + .collect(); + all_signals.push(signal_results); + } // Sparse search if let Some(ref sparse) = query.sparse_vector - && weights.1 > 0.0 { - let sparse_idx = self.sparse_index.read() - .map_err(|_| VecStoreError::LockError("sparse_index lock poisoned".into()))?; - let results = sparse_idx.search(sparse, self.config.max_candidates); - let signal_results: Vec = results - .into_iter() - .enumerate() - .map(|(rank, (id, score))| SignalResult { - id, - score, - rank, - signal: SignalType::Sparse, - }) - .collect(); - all_signals.push(signal_results); - } + && weights.1 > 0.0 + { + let sparse_idx = self + .sparse_index + .read() + .map_err(|_| VecStoreError::LockError("sparse_index lock poisoned".into()))?; + let results = sparse_idx.search(sparse, self.config.max_candidates); + let signal_results: Vec = results + .into_iter() + .enumerate() + .map(|(rank, (id, score))| SignalResult { + id, + score, + rank, + signal: SignalType::Sparse, + }) + .collect(); + all_signals.push(signal_results); + } // BM25 search if weights.2 > 0.0 && !query.text.is_empty() { let query_terms = tokenize(&query.text); - let mut bm25 = self.bm25_index.write() + let mut bm25 = self + .bm25_index + .write() .map_err(|_| VecStoreError::LockError("bm25_index lock poisoned".into()))?; let results = bm25.search_bm25(&query_terms, 1.2, 0.75, self.config.max_candidates); let signal_results: Vec = results @@ -693,7 +716,9 @@ impl HybridPipeline { // Apply reranking if configured if let Some(reranker) = &self.reranker { - let docs = self.documents.read() + let docs = self + .documents + .read() .map_err(|_| VecStoreError::LockError("documents lock poisoned".into()))?; fused = reranker.rerank(&query.text, fused, &docs); } @@ -715,11 +740,19 @@ impl HybridPipeline { o.bm25_weight.unwrap_or(self.config.bm25_weight), ) } else { - (self.config.dense_weight, self.config.sparse_weight, self.config.bm25_weight) + ( + self.config.dense_weight, + self.config.sparse_weight, + self.config.bm25_weight, + ) } } - fn fuse_results(&self, signals: &[Vec], weights: (f32, f32, f32)) -> Result> { + fn fuse_results( + &self, + signals: &[Vec], + weights: (f32, f32, f32), + ) -> Result> { match &self.config.fusion_strategy { FusionMethod::RRF { k } => self.fuse_rrf(signals, *k), FusionMethod::WeightedSum => self.fuse_weighted_sum(signals, weights), @@ -733,15 +766,16 @@ impl HybridPipeline { } fn fuse_rrf(&self, signals: &[Vec], k: f32) -> Result> { - let mut doc_scores: HashMap, HashMap)> = HashMap::new(); + let mut doc_scores: HashMap, HashMap)> = + HashMap::new(); for signal_results in signals { for result in signal_results { let rrf_score = 1.0 / (k + result.rank as f32 + 1.0); - let entry = doc_scores.entry(result.id.clone()).or_insert_with(|| { - (0.0, HashMap::new(), HashMap::new()) - }); + let entry = doc_scores + .entry(result.id.clone()) + .or_insert_with(|| (0.0, HashMap::new(), HashMap::new())); entry.0 += rrf_score; entry.1.insert(format!("{:?}", result.signal), result.score); @@ -752,23 +786,30 @@ impl HybridPipeline { self.build_results(doc_scores, "RRF") } - fn fuse_weighted_sum(&self, signals: &[Vec], weights: (f32, f32, f32)) -> Result> { + fn fuse_weighted_sum( + &self, + signals: &[Vec], + weights: (f32, f32, f32), + ) -> Result> { let weight_map: HashMap = [ (SignalType::Dense, weights.0), (SignalType::Sparse, weights.1), (SignalType::BM25, weights.2), - ].into_iter().collect(); + ] + .into_iter() + .collect(); - let mut doc_scores: HashMap, HashMap)> = HashMap::new(); + let mut doc_scores: HashMap, HashMap)> = + HashMap::new(); for signal_results in signals { for result in signal_results { let weight = weight_map.get(&result.signal).copied().unwrap_or(1.0); let weighted_score = result.score * weight; - let entry = doc_scores.entry(result.id.clone()).or_insert_with(|| { - (0.0, HashMap::new(), HashMap::new()) - }); + let entry = doc_scores + .entry(result.id.clone()) + .or_insert_with(|| (0.0, HashMap::new(), HashMap::new())); entry.0 += weighted_score; entry.1.insert(format!("{:?}", result.signal), result.score); @@ -780,13 +821,14 @@ impl HybridPipeline { } fn fuse_max_score(&self, signals: &[Vec]) -> Result> { - let mut doc_scores: HashMap, HashMap)> = HashMap::new(); + let mut doc_scores: HashMap, HashMap)> = + HashMap::new(); for signal_results in signals { for result in signal_results { - let entry = doc_scores.entry(result.id.clone()).or_insert_with(|| { - (0.0, HashMap::new(), HashMap::new()) - }); + let entry = doc_scores + .entry(result.id.clone()) + .or_insert_with(|| (0.0, HashMap::new(), HashMap::new())); if result.score > entry.0 { entry.0 = result.score; @@ -800,13 +842,16 @@ impl HybridPipeline { } fn fuse_geometric_mean(&self, signals: &[Vec]) -> Result> { - let mut doc_data: HashMap, HashMap, HashMap)> = HashMap::new(); + let mut doc_data: HashMap< + String, + (Vec, HashMap, HashMap), + > = HashMap::new(); for signal_results in signals { for result in signal_results { - let entry = doc_data.entry(result.id.clone()).or_insert_with(|| { - (Vec::new(), HashMap::new(), HashMap::new()) - }); + let entry = doc_data + .entry(result.id.clone()) + .or_insert_with(|| (Vec::new(), HashMap::new(), HashMap::new())); entry.0.push(result.score.max(0.001)); // Avoid zero entry.1.insert(format!("{:?}", result.signal), result.score); @@ -814,26 +859,30 @@ impl HybridPipeline { } } - let doc_scores: HashMap, HashMap)> = doc_data - .into_iter() - .map(|(id, (scores, signal_scores, signal_ranks))| { - let product: f32 = scores.iter().product(); - let geom_mean = product.powf(1.0 / scores.len() as f32); - (id, (geom_mean, signal_scores, signal_ranks)) - }) - .collect(); + let doc_scores: HashMap, HashMap)> = + doc_data + .into_iter() + .map(|(id, (scores, signal_scores, signal_ranks))| { + let product: f32 = scores.iter().product(); + let geom_mean = product.powf(1.0 / scores.len() as f32); + (id, (geom_mean, signal_scores, signal_ranks)) + }) + .collect(); self.build_results(doc_scores, "GeometricMean") } fn fuse_harmonic_mean(&self, signals: &[Vec]) -> Result> { - let mut doc_data: HashMap, HashMap, HashMap)> = HashMap::new(); + let mut doc_data: HashMap< + String, + (Vec, HashMap, HashMap), + > = HashMap::new(); for signal_results in signals { for result in signal_results { - let entry = doc_data.entry(result.id.clone()).or_insert_with(|| { - (Vec::new(), HashMap::new(), HashMap::new()) - }); + let entry = doc_data + .entry(result.id.clone()) + .or_insert_with(|| (Vec::new(), HashMap::new(), HashMap::new())); entry.0.push(result.score.max(0.001)); entry.1.insert(format!("{:?}", result.signal), result.score); @@ -841,26 +890,32 @@ impl HybridPipeline { } } - let doc_scores: HashMap, HashMap)> = doc_data - .into_iter() - .map(|(id, (scores, signal_scores, signal_ranks))| { - let sum_reciprocals: f32 = scores.iter().map(|s| 1.0 / s).sum(); - let harm_mean = scores.len() as f32 / sum_reciprocals; - (id, (harm_mean, signal_scores, signal_ranks)) - }) - .collect(); + let doc_scores: HashMap, HashMap)> = + doc_data + .into_iter() + .map(|(id, (scores, signal_scores, signal_ranks))| { + let sum_reciprocals: f32 = scores.iter().map(|s| 1.0 / s).sum(); + let harm_mean = scores.len() as f32 / sum_reciprocals; + (id, (harm_mean, signal_scores, signal_ranks)) + }) + .collect(); self.build_results(doc_scores, "HarmonicMean") } - fn fuse_voting(&self, signals: &[Vec], min_votes: usize) -> Result> { - let mut doc_votes: HashMap, HashMap)> = HashMap::new(); + fn fuse_voting( + &self, + signals: &[Vec], + min_votes: usize, + ) -> Result> { + let mut doc_votes: HashMap, HashMap)> = + HashMap::new(); for signal_results in signals { for result in signal_results { - let entry = doc_votes.entry(result.id.clone()).or_insert_with(|| { - (0, HashMap::new(), HashMap::new()) - }); + let entry = doc_votes + .entry(result.id.clone()) + .or_insert_with(|| (0, HashMap::new(), HashMap::new())); entry.0 += 1; entry.1.insert(format!("{:?}", result.signal), result.score); @@ -868,18 +923,23 @@ impl HybridPipeline { } } - let doc_scores: HashMap, HashMap)> = doc_votes - .into_iter() - .filter(|(_, (votes, _, _))| *votes >= min_votes) - .map(|(id, (votes, signal_scores, signal_ranks))| { - (id, (votes as f32, signal_scores, signal_ranks)) - }) - .collect(); + let doc_scores: HashMap, HashMap)> = + doc_votes + .into_iter() + .filter(|(_, (votes, _, _))| *votes >= min_votes) + .map(|(id, (votes, signal_scores, signal_ranks))| { + (id, (votes as f32, signal_scores, signal_ranks)) + }) + .collect(); self.build_results(doc_scores, "Voting") } - fn fuse_distribution(&self, signals: &[Vec], weights: (f32, f32, f32)) -> Result> { + fn fuse_distribution( + &self, + signals: &[Vec], + weights: (f32, f32, f32), + ) -> Result> { // Normalize scores within each signal, then combine let mut normalized_signals = Vec::new(); @@ -888,8 +948,14 @@ impl HybridPipeline { continue; } - let max_score = signal_results.iter().map(|r| r.score).fold(f32::NEG_INFINITY, f32::max); - let min_score = signal_results.iter().map(|r| r.score).fold(f32::INFINITY, f32::min); + let max_score = signal_results + .iter() + .map(|r| r.score) + .fold(f32::NEG_INFINITY, f32::max); + let min_score = signal_results + .iter() + .map(|r| r.score) + .fold(f32::INFINITY, f32::min); let range = (max_score - min_score).max(0.001); let normalized: Vec = signal_results @@ -913,7 +979,9 @@ impl HybridPipeline { doc_scores: HashMap, HashMap)>, method: &str, ) -> Result> { - let docs = self.documents.read() + let docs = self + .documents + .read() .map_err(|_| VecStoreError::LockError("documents lock poisoned".into()))?; let mut results: Vec = doc_scores @@ -952,11 +1020,17 @@ impl HybridPipeline { /// Get pipeline statistics pub fn stats(&self) -> Result { - let dense = self.dense_index.read() + let dense = self + .dense_index + .read() .map_err(|_| VecStoreError::LockError("dense_index lock poisoned".into()))?; - let sparse = self.sparse_index.read() + let sparse = self + .sparse_index + .read() .map_err(|_| VecStoreError::LockError("sparse_index lock poisoned".into()))?; - let bm25 = self.bm25_index.read() + let bm25 = self + .bm25_index + .read() .map_err(|_| VecStoreError::LockError("bm25_index lock poisoned".into()))?; Ok(PipelineStats { @@ -1018,11 +1092,11 @@ impl Reranker { } results.sort_by(|a, b| b.score.total_cmp(&a.score)); - } + }, _ => { // Other rerankers would call external APIs // For now, return unchanged - } + }, } results @@ -1070,10 +1144,7 @@ mod tests { id: format!("doc_{}", i), text: format!("This is document {} about machine learning", i), dense_vector: vec![0.1 * i as f32; 128], - sparse_vector: Some(SparseQueryVector::new( - vec![1, 2, 3], - vec![0.5, 0.3, 0.2], - )), + sparse_vector: Some(SparseQueryVector::new(vec![1, 2, 3], vec![0.5, 0.3, 0.2])), metadata: HashMap::new(), }; pipeline.index(doc).unwrap(); @@ -1107,8 +1178,7 @@ mod tests { pipeline.index(doc).unwrap(); } - let query = HybridSearchQuery::new("Document") - .with_dense(vec![2.0; 4]); + let query = HybridSearchQuery::new("Document").with_dense(vec![2.0; 4]); let results = pipeline.search(query).unwrap(); assert!(!results.is_empty()); @@ -1134,15 +1204,14 @@ mod tests { BooleanQuery::Phrase { terms, slop } => { assert_eq!(terms.len(), 2); assert_eq!(slop, 1); - } + }, _ => panic!("Expected Phrase query"), } } #[test] fn test_weight_overrides() { - let query = HybridSearchQuery::new("test") - .with_weights(0.5, 0.3, 0.2); + let query = HybridSearchQuery::new("test").with_weights(0.5, 0.3, 0.2); let overrides = query.weight_overrides.unwrap(); assert_eq!(overrides.dense_weight, Some(0.5)); diff --git a/src/import_export.rs b/src/import_export.rs index 484e0e9..87217c8 100644 --- a/src/import_export.rs +++ b/src/import_export.rs @@ -85,7 +85,7 @@ fn value_to_metadata(value: serde_json::Value) -> Metadata { serde_json::Value::Object(map) => { let fields = map.into_iter().collect(); Metadata { fields } - } + }, _ => Metadata { fields: HashMap::new(), }, diff --git a/src/incremental_index.rs b/src/incremental_index.rs index 81f7067..91e3567 100644 --- a/src/incremental_index.rs +++ b/src/incremental_index.rs @@ -194,7 +194,12 @@ impl IncrementalIndex { } /// Insert a vector - pub fn insert(&self, id: String, vector: Vec, metadata: Option) -> Result<()> { + pub fn insert( + &self, + id: String, + vector: Vec, + metadata: Option, + ) -> Result<()> { self.validate_vector(&vector)?; let op = Operation::Insert { @@ -218,18 +223,29 @@ impl IncrementalIndex { }; { - let mut vectors = self.delta.vectors.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let mut vectors = self + .delta + .vectors + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let is_new = !vectors.contains_key(&id); vectors.insert(id.clone(), entry); if is_new { self.delta.size.fetch_add(1, Ordering::Relaxed); - self.delta.insertion_order.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.push(id.clone()); + self.delta + .insertion_order + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .push(id.clone()); } } // Remove from tombstones if exists - self.tombstones.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.remove(&id); + self.tombstones + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .remove(&id); self.stats.inserts.fetch_add(1, Ordering::Relaxed); self.pending_ops.fetch_add(1, Ordering::Relaxed); @@ -241,7 +257,12 @@ impl IncrementalIndex { } /// Update a vector - pub fn update(&self, id: String, vector: Vec, metadata: Option) -> Result<()> { + pub fn update( + &self, + id: String, + vector: Vec, + metadata: Option, + ) -> Result<()> { self.validate_vector(&vector)?; // Check if exists @@ -269,13 +290,21 @@ impl IncrementalIndex { }; { - let mut vectors = self.delta.vectors.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let mut vectors = self + .delta + .vectors + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let is_new = !vectors.contains_key(&id); vectors.insert(id.clone(), entry); if is_new { self.delta.size.fetch_add(1, Ordering::Relaxed); - self.delta.insertion_order.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.push(id); + self.delta + .insertion_order + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .push(id); } } @@ -302,7 +331,11 @@ impl IncrementalIndex { // immutable base index instead, which can't be edited in place). Tombstoning a // delta-resident vector too would double-count it out of total_vectors(): once via // the delta_size decrement here, again via the tombstone subtraction there. - let mut vectors = self.delta.vectors.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let mut vectors = self + .delta + .vectors + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let removed_from_delta = vectors.remove(&id).is_some(); if removed_from_delta { self.delta.size.fetch_sub(1, Ordering::Relaxed); @@ -320,7 +353,10 @@ impl IncrementalIndex { deleted_at: current_timestamp(), version: self.wal_sequence.load(Ordering::Relaxed), }; - self.tombstones.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.insert(id.clone(), tombstone); + self.tombstones + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .insert(id.clone(), tombstone); } self.stats.deletes.fetch_add(1, Ordering::Relaxed); @@ -333,7 +369,10 @@ impl IncrementalIndex { } /// Batch insert - pub fn batch_insert(&self, items: Vec<(String, Vec, Option)>) -> Result { + pub fn batch_insert( + &self, + items: Vec<(String, Vec, Option)>, + ) -> Result { let mut succeeded = 0; let mut failed = 0; let mut errors = Vec::new(); @@ -344,26 +383,32 @@ impl IncrementalIndex { Err(e) => { failed += 1; errors.push((id, e.to_string())); - } + }, } } - Ok(BatchResult { succeeded, failed, errors }) + Ok(BatchResult { + succeeded, + failed, + errors, + }) } /// Check if vector exists pub fn exists(&self, id: &str) -> bool { // Check tombstones first if let Ok(tombstones) = self.tombstones.read() - && tombstones.contains_key(id) { - return false; - } + && tombstones.contains_key(id) + { + return false; + } // Check delta if let Ok(delta) = self.delta.vectors.read() - && delta.contains_key(id) { - return true; - } + && delta.contains_key(id) + { + return true; + } // Check base indexes if let Ok(bases) = self.base_indexes.read() { @@ -415,11 +460,18 @@ impl IncrementalIndex { self.validate_vector(query)?; self.stats.searches.fetch_add(1, Ordering::Relaxed); - let tombstones = self.tombstones.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let tombstones = self + .tombstones + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let mut results = Vec::new(); // Search delta index - let delta_vectors = self.delta.vectors.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let delta_vectors = self + .delta + .vectors + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; for (id, entry) in delta_vectors.iter() { if tombstones.contains_key(id) { continue; @@ -434,7 +486,12 @@ impl IncrementalIndex { } // Search base indexes - for base in self.base_indexes.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.iter() { + for base in self + .base_indexes + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .iter() + { for (id, entry) in &base.vectors { if tombstones.contains_key(id) || delta_vectors.contains_key(id) { continue; // Skip deleted or overridden @@ -460,7 +517,9 @@ impl IncrementalIndex { pub fn merge(&self) -> Result { // Check if merge already in progress if self.merge_in_progress.swap(true, Ordering::SeqCst) { - return Err(VecStoreError::InvalidInput("Merge already in progress".into())); + return Err(VecStoreError::InvalidInput( + "Merge already in progress".into(), + )); } let start = Instant::now(); @@ -468,29 +527,49 @@ impl IncrementalIndex { // Collect all vectors (delta + base - tombstones) let mut all_vectors = HashMap::new(); - let tombstones = self.tombstones.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let tombstones = self + .tombstones + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; // Add base vectors - for base in self.base_indexes.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.iter() { + for base in self + .base_indexes + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .iter() + { for (id, entry) in &base.vectors { if !tombstones.contains_key(id) { - all_vectors.insert(id.clone(), BaseEntry { - id: id.clone(), - vector: entry.vector.clone(), - metadata: entry.metadata.clone(), - }); + all_vectors.insert( + id.clone(), + BaseEntry { + id: id.clone(), + vector: entry.vector.clone(), + metadata: entry.metadata.clone(), + }, + ); } } } // Add/override with delta vectors - for (id, entry) in self.delta.vectors.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.iter() { + for (id, entry) in self + .delta + .vectors + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .iter() + { if !tombstones.contains_key(id) { - all_vectors.insert(id.clone(), BaseEntry { - id: id.clone(), - vector: entry.vector.clone(), - metadata: entry.metadata.clone(), - }); + all_vectors.insert( + id.clone(), + BaseEntry { + id: id.clone(), + vector: entry.vector.clone(), + metadata: entry.metadata.clone(), + }, + ); } } drop(tombstones); @@ -506,32 +585,59 @@ impl IncrementalIndex { // Replace indexes { - let mut bases = self.base_indexes.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let mut bases = self + .base_indexes + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; let _old_count = bases.len(); bases.clear(); bases.push(new_base); // Clear delta - self.delta.vectors.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.clear(); - self.delta.insertion_order.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.clear(); + self.delta + .vectors + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .clear(); + self.delta + .insertion_order + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .clear(); self.delta.size.store(0, Ordering::Relaxed); // Clear tombstones - let cleaned = self.tombstones.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.len(); - self.tombstones.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.clear(); - - self.stats.tombstones_cleaned.fetch_add(cleaned as u64, Ordering::Relaxed); + let cleaned = self + .tombstones + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .len(); + self.tombstones + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .clear(); + + self.stats + .tombstones_cleaned + .fetch_add(cleaned as u64, Ordering::Relaxed); } // Update state - *self.last_merge.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))? = Instant::now(); + *self + .last_merge + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? = Instant::now(); self.pending_ops.store(0, Ordering::Relaxed); self.merge_in_progress.store(false, Ordering::SeqCst); self.stats.merges.fetch_add(1, Ordering::Relaxed); Ok(MergeResult { vectors_merged: delta_size, - new_index_size: self.base_indexes.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?[0].vector_count, + new_index_size: self + .base_indexes + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?[0] + .vector_count, duration: start.elapsed(), tombstones_removed: 0, }) @@ -543,7 +649,11 @@ impl IncrementalIndex { return Ok(0); } - let tombstone_count = self.tombstones.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.len(); + let tombstone_count = self + .tombstones + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .len(); let total_vectors = self.total_vectors(); if total_vectors == 0 { @@ -581,11 +691,22 @@ impl IncrementalIndex { checksum: 0, // In production: compute actual checksum }; - self.wal_buffer.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.push_back(entry); - self.stats.wal_entries_written.fetch_add(1, Ordering::Relaxed); + self.wal_buffer + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .push_back(entry); + self.stats + .wal_entries_written + .fetch_add(1, Ordering::Relaxed); // Flush if buffer is large - if self.wal_buffer.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.len() > 1000 { + if self + .wal_buffer + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .len() + > 1000 + { self.flush_wal()?; } @@ -594,7 +715,10 @@ impl IncrementalIndex { fn flush_wal(&self) -> Result<()> { // In production: write to disk - let mut buffer = self.wal_buffer.write().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; + let mut buffer = self + .wal_buffer + .write() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))?; buffer.clear(); Ok(()) } @@ -603,8 +727,7 @@ impl IncrementalIndex { let delta_size = self.delta.size.load(Ordering::Relaxed); let pending = self.pending_ops.load(Ordering::Relaxed); - if delta_size >= self.config.delta_size_threshold - || pending >= self.config.merge_threshold + if delta_size >= self.config.delta_size_threshold || pending >= self.config.merge_threshold { // In production: spawn background task // For now, merge synchronously @@ -615,7 +738,11 @@ impl IncrementalIndex { } fn maybe_trigger_cleanup(&self) -> Result<()> { - let tombstone_count = self.tombstones.read().map_err(|_| VecStoreError::LockError("lock poisoned".into()))?.len(); + let tombstone_count = self + .tombstones + .read() + .map_err(|_| VecStoreError::LockError("lock poisoned".into()))? + .len(); let total = self.total_vectors(); if total > 0 { @@ -631,7 +758,9 @@ impl IncrementalIndex { /// Get total vector count pub fn total_vectors(&self) -> usize { let delta_size = self.delta.size.load(Ordering::Relaxed); - let base_size: usize = self.base_indexes.read() + let base_size: usize = self + .base_indexes + .read() .map(|b| b.iter().map(|b| b.vector_count).sum()) .unwrap_or(0); let tombstone_count = self.tombstones.read().map(|t| t.len()).unwrap_or(0); @@ -744,27 +873,38 @@ impl WalRecovery { for entry in entries { match entry.operation { - Operation::Insert { id, vector, metadata } => { + Operation::Insert { + id, + vector, + metadata, + } => { if index.insert(id, vector, metadata).is_ok() { replayed += 1; } else { errors += 1; } - } - Operation::Update { id, vector, metadata } => { + }, + Operation::Update { + id, + vector, + metadata, + } => { // For updates, treat as insert if not exists - if index.update(id.clone(), vector.clone(), metadata.clone()).is_err() { + if index + .update(id.clone(), vector.clone(), metadata.clone()) + .is_err() + { let _ = index.insert(id, vector, metadata); } replayed += 1; - } + }, Operation::Delete { id } => { if index.delete(id).is_ok() { replayed += 1; } else { errors += 1; } - } + }, Operation::BatchInsert { items } => { for (id, vector, metadata) in items { if index.insert(id, vector, metadata).is_ok() { @@ -773,7 +913,7 @@ impl WalRecovery { errors += 1; } } - } + }, } } @@ -801,9 +941,15 @@ mod tests { fn test_insert_and_search() { let index = IncrementalIndex::new(3, IncrementalConfig::default()); - index.insert("v1".to_string(), vec![1.0, 0.0, 0.0], None).unwrap(); - index.insert("v2".to_string(), vec![0.0, 1.0, 0.0], None).unwrap(); - index.insert("v3".to_string(), vec![0.0, 0.0, 1.0], None).unwrap(); + index + .insert("v1".to_string(), vec![1.0, 0.0, 0.0], None) + .unwrap(); + index + .insert("v2".to_string(), vec![0.0, 1.0, 0.0], None) + .unwrap(); + index + .insert("v3".to_string(), vec![0.0, 0.0, 1.0], None) + .unwrap(); let results = index.search(&[1.0, 0.1, 0.0], 2).unwrap(); @@ -815,12 +961,16 @@ mod tests { fn test_update() { let index = IncrementalIndex::new(3, IncrementalConfig::default()); - index.insert("v1".to_string(), vec![1.0, 0.0, 0.0], None).unwrap(); + index + .insert("v1".to_string(), vec![1.0, 0.0, 0.0], None) + .unwrap(); let entry1 = index.get("v1").unwrap(); assert_eq!(entry1.vector, vec![1.0, 0.0, 0.0]); - index.update("v1".to_string(), vec![0.0, 1.0, 0.0], None).unwrap(); + index + .update("v1".to_string(), vec![0.0, 1.0, 0.0], None) + .unwrap(); let entry2 = index.get("v1").unwrap(); assert_eq!(entry2.vector, vec![0.0, 1.0, 0.0]); @@ -830,7 +980,9 @@ mod tests { fn test_delete() { let index = IncrementalIndex::new(3, IncrementalConfig::default()); - index.insert("v1".to_string(), vec![1.0, 0.0, 0.0], None).unwrap(); + index + .insert("v1".to_string(), vec![1.0, 0.0, 0.0], None) + .unwrap(); assert!(index.exists("v1")); index.delete("v1".to_string()).unwrap(); @@ -848,11 +1000,9 @@ mod tests { // Insert vectors for i in 0..100 { - index.insert( - format!("v{}", i), - vec![i as f32, 0.0, 0.0], - None, - ).unwrap(); + index + .insert(format!("v{}", i), vec![i as f32, 0.0, 0.0], None) + .unwrap(); } let stats_before = index.get_stats(); @@ -883,7 +1033,9 @@ mod tests { // deletes -- a delta-resident vector is removed directly since delta is mutable (see // delete()'s comment), so deleting straight out of delta never produces a tombstone. for i in 0..10 { - index.insert(format!("v{}", i), vec![i as f32, 0.0, 0.0], None).unwrap(); + index + .insert(format!("v{}", i), vec![i as f32, 0.0, 0.0], None) + .unwrap(); } index.merge().unwrap(); diff --git a/src/inference.rs b/src/inference.rs index 010bc51..1b36a63 100644 --- a/src/inference.rs +++ b/src/inference.rs @@ -30,11 +30,11 @@ //! let results = engine.query_text("greeting", 10)?; //! ``` +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::{Arc, RwLock}; -use serde::{Deserialize, Serialize}; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; /// Embedding provider configuration #[derive(Debug, Clone, Serialize, Deserialize)] @@ -47,15 +47,9 @@ pub enum EmbeddingProvider { dimensions: Option, }, /// Cohere embeddings (embed-english-v3.0, embed-multilingual-v3.0) - Cohere { - model: String, - api_key: String, - }, + Cohere { model: String, api_key: String }, /// Voyage AI embeddings (voyage-3, voyage-3-lite, voyage-code-3) - VoyageAI { - model: String, - api_key: String, - }, + VoyageAI { model: String, api_key: String }, /// Google Vertex AI embeddings VertexAI { model: String, @@ -68,15 +62,9 @@ pub enum EmbeddingProvider { tokenizer_path: String, }, /// Ollama local models - Ollama { - model: String, - base_url: String, - }, + Ollama { model: String, base_url: String }, /// HuggingFace Inference API - HuggingFace { - model: String, - api_key: String, - }, + HuggingFace { model: String, api_key: String }, /// Custom HTTP endpoint Custom { endpoint: String, @@ -89,47 +77,39 @@ impl EmbeddingProvider { /// Get the default dimension for this provider/model pub fn default_dimension(&self) -> usize { match self { - EmbeddingProvider::OpenAI { model, dimensions, .. } => { - dimensions.unwrap_or_else(|| match model.as_str() { - "text-embedding-3-small" => 1536, - "text-embedding-3-large" => 3072, - "text-embedding-ada-002" => 1536, - _ => 1536, - }) - } - EmbeddingProvider::Cohere { model, .. } => { - match model.as_str() { - "embed-english-v3.0" => 1024, - "embed-multilingual-v3.0" => 1024, - "embed-english-light-v3.0" => 384, - _ => 1024, - } - } - EmbeddingProvider::VoyageAI { model, .. } => { - match model.as_str() { - "voyage-3" => 1024, - "voyage-3-lite" => 512, - "voyage-code-3" => 1024, - _ => 1024, - } - } - EmbeddingProvider::VertexAI { model, .. } => { - match model.as_str() { - "textembedding-gecko@003" => 768, - "text-embedding-004" => 768, - "text-multilingual-embedding-002" => 768, - _ => 768, - } - } + EmbeddingProvider::OpenAI { + model, dimensions, .. + } => dimensions.unwrap_or_else(|| match model.as_str() { + "text-embedding-3-small" => 1536, + "text-embedding-3-large" => 3072, + "text-embedding-ada-002" => 1536, + _ => 1536, + }), + EmbeddingProvider::Cohere { model, .. } => match model.as_str() { + "embed-english-v3.0" => 1024, + "embed-multilingual-v3.0" => 1024, + "embed-english-light-v3.0" => 384, + _ => 1024, + }, + EmbeddingProvider::VoyageAI { model, .. } => match model.as_str() { + "voyage-3" => 1024, + "voyage-3-lite" => 512, + "voyage-code-3" => 1024, + _ => 1024, + }, + EmbeddingProvider::VertexAI { model, .. } => match model.as_str() { + "textembedding-gecko@003" => 768, + "text-embedding-004" => 768, + "text-multilingual-embedding-002" => 768, + _ => 768, + }, EmbeddingProvider::LocalONNX { .. } => 384, // Default for MiniLM - EmbeddingProvider::Ollama { model, .. } => { - match model.as_str() { - "nomic-embed-text" => 768, - "mxbai-embed-large" => 1024, - "all-minilm" => 384, - _ => 768, - } - } + EmbeddingProvider::Ollama { model, .. } => match model.as_str() { + "nomic-embed-text" => 768, + "mxbai-embed-large" => 1024, + "all-minilm" => 384, + _ => 768, + }, EmbeddingProvider::HuggingFace { model, .. } => { if model.contains("minilm") || model.contains("MiniLM") { 384 @@ -138,7 +118,7 @@ impl EmbeddingProvider { } else { 768 } - } + }, EmbeddingProvider::Custom { .. } => 768, // Default } } @@ -180,10 +160,18 @@ pub struct InferenceConfig { pub timeout_ms: u64, } -fn default_batch_size() -> usize { 100 } -fn default_true() -> bool { true } -fn default_cache_size() -> usize { 10000 } -fn default_timeout() -> u64 { 30000 } +fn default_batch_size() -> usize { + 100 +} +fn default_true() -> bool { + true +} +fn default_cache_size() -> usize { + 10000 +} +fn default_timeout() -> u64 { + 30000 +} impl InferenceConfig { /// Create a new inference configuration @@ -280,14 +268,15 @@ impl EmbeddingCache { /// Get an embedding from cache pub fn get(&mut self, key: &str, model: &str) -> Option> { if let Some(entry) = self.entries.get(key) - && entry.model == model { - // Move to end (most recently used) - if let Some(pos) = self.order.iter().position(|k| k == key) { - self.order.remove(pos); - self.order.push(key.to_string()); - } - return Some(entry.embedding.clone()); + && entry.model == model + { + // Move to end (most recently used) + if let Some(pos) = self.order.iter().position(|k| k == key) { + self.order.remove(pos); + self.order.push(key.to_string()); } + return Some(entry.embedding.clone()); + } None } @@ -301,11 +290,14 @@ impl EmbeddingCache { } } - self.entries.insert(key.clone(), CacheEntry { - embedding, - model, - created_at: std::time::Instant::now(), - }); + self.entries.insert( + key.clone(), + CacheEntry { + embedding, + model, + created_at: std::time::Instant::now(), + }, + ); self.order.push(key); } @@ -382,16 +374,19 @@ impl InferenceEngine { /// Get current statistics pub fn stats(&self) -> InferenceStats { - let Ok(guard) = self.stats.read() else { return InferenceStats::default(); }; + let Ok(guard) = self.stats.read() else { + return InferenceStats::default(); + }; guard.clone() } /// Embed a single text pub fn embed(&self, text: &str) -> Result> { let embeddings = self.embed_batch(&[text.to_string()])?; - embeddings.into_iter().next().ok_or_else(|| { - VecStoreError::InvalidInput("No embedding returned".to_string()) - }) + embeddings + .into_iter() + .next() + .ok_or_else(|| VecStoreError::InvalidInput("No embedding returned".to_string())) } /// Embed a batch of texts @@ -413,20 +408,31 @@ impl InferenceEngine { // Check cache first if self.config.cache_enabled { - let mut cache = self.cache.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on embedding cache".into()))?; + let mut cache = self.cache.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on embedding cache".into()) + })?; for (i, text) in texts.iter().enumerate() { let cache_key = Self::cache_key(text); if let Some(embedding) = cache.get(&cache_key, &model_name) { results[i] = Some(embedding); - self.stats.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on inference stats".into()))? + self.stats + .write() + .map_err(|_| { + VecStoreError::LockError( + "Failed to acquire write lock on inference stats".into(), + ) + })? .cache_hits += 1; } else { uncached_indices.push(i); uncached_texts.push(text.clone()); - self.stats.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on inference stats".into()))? + self.stats + .write() + .map_err(|_| { + VecStoreError::LockError( + "Failed to acquire write lock on inference stats".into(), + ) + })? .cache_misses += 1; } } @@ -441,8 +447,11 @@ impl InferenceEngine { // Update cache and results if self.config.cache_enabled { - let mut cache = self.cache.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on embedding cache".into()))?; + let mut cache = self.cache.write().map_err(|_| { + VecStoreError::LockError( + "Failed to acquire write lock on embedding cache".into(), + ) + })?; for (i, embedding) in new_embeddings.into_iter().enumerate() { let orig_idx = uncached_indices[i]; let cache_key = Self::cache_key(&uncached_texts[i]); @@ -458,13 +467,15 @@ impl InferenceEngine { // Update stats { - let mut stats = self.stats.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on inference stats".into()))?; + let mut stats = self.stats.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on inference stats".into()) + })?; stats.total_requests += 1; stats.total_embeddings += texts.len() as u64; } - results.into_iter() + results + .into_iter() .map(|r| r.ok_or_else(|| VecStoreError::InvalidInput("Missing embedding".to_string()))) .collect() } @@ -482,30 +493,30 @@ impl InferenceEngine { /// Generate embeddings using the configured provider fn generate_embeddings(&self, texts: &[String]) -> Result>> { match &self.config.provider { - EmbeddingProvider::OpenAI { model, api_key, dimensions } => { - self.embed_openai(texts, model, api_key, *dimensions) - } + EmbeddingProvider::OpenAI { + model, + api_key, + dimensions, + } => self.embed_openai(texts, model, api_key, *dimensions), EmbeddingProvider::Cohere { model, api_key } => { self.embed_cohere(texts, model, api_key) - } + }, EmbeddingProvider::VoyageAI { model, api_key } => { self.embed_voyageai(texts, model, api_key) - } + }, EmbeddingProvider::Ollama { model, base_url } => { self.embed_ollama(texts, model, base_url) - } - EmbeddingProvider::LocalONNX { .. } => { - self.embed_local_placeholder(texts) - } - EmbeddingProvider::VertexAI { .. } => { - self.embed_local_placeholder(texts) - } + }, + EmbeddingProvider::LocalONNX { .. } => self.embed_local_placeholder(texts), + EmbeddingProvider::VertexAI { .. } => self.embed_local_placeholder(texts), EmbeddingProvider::HuggingFace { model, api_key } => { self.embed_huggingface(texts, model, api_key) - } - EmbeddingProvider::Custom { endpoint, api_key, headers } => { - self.embed_custom(texts, endpoint, api_key.as_deref(), headers) - } + }, + EmbeddingProvider::Custom { + endpoint, + api_key, + headers, + } => self.embed_custom(texts, endpoint, api_key.as_deref(), headers), } } @@ -535,7 +546,9 @@ impl InferenceEngine { .header("Content-Type", "application/json") .json(&body) .send() - .map_err(|e| VecStoreError::Internal(format!("OpenAI embeddings request failed: {}", e)))?; + .map_err(|e| { + VecStoreError::Internal(format!("OpenAI embeddings request failed: {}", e)) + })?; if !response.status().is_success() { let status = response.status(); @@ -546,18 +559,20 @@ impl InferenceEngine { ))); } - let json: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse OpenAI response: {}", e)))?; + let json: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse OpenAI response: {}", e)) + })?; let embeddings: Vec> = json["data"] .as_array() .ok_or_else(|| VecStoreError::Internal("No data in OpenAI response".to_string()))? .iter() .filter_map(|item| { - item["embedding"] - .as_array() - .map(|arr| arr.iter().filter_map(|v| v.as_f64().map(|f| f as f32)).collect()) + item["embedding"].as_array().map(|arr| { + arr.iter() + .filter_map(|v| v.as_f64().map(|f| f as f32)) + .collect() + }) }) .collect(); @@ -574,9 +589,10 @@ impl InferenceEngine { dimensions: Option, ) -> Result>> { let dim = dimensions.unwrap_or(self.dimension); - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, dim) - }).collect()) + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, dim)) + .collect()) } /// Cohere embedding implementation @@ -586,9 +602,10 @@ impl InferenceEngine { _model: &str, _api_key: &str, ) -> Result>> { - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, self.dimension) - }).collect()) + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, self.dimension)) + .collect()) } /// Voyage AI embedding implementation @@ -598,19 +615,15 @@ impl InferenceEngine { _model: &str, _api_key: &str, ) -> Result>> { - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, self.dimension) - }).collect()) + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, self.dimension)) + .collect()) } /// Ollama embedding implementation #[cfg(feature = "ollama")] - fn embed_ollama( - &self, - texts: &[String], - model: &str, - base_url: &str, - ) -> Result>> { + fn embed_ollama(&self, texts: &[String], model: &str, base_url: &str) -> Result>> { let client = reqwest::blocking::Client::new(); let url = format!("{}/api/embeddings", base_url.trim_end_matches('/')); @@ -628,7 +641,9 @@ impl InferenceEngine { .header("Content-Type", "application/json") .json(&body) .send() - .map_err(|e| VecStoreError::Internal(format!("Ollama embeddings request failed: {}", e)))?; + .map_err(|e| { + VecStoreError::Internal(format!("Ollama embeddings request failed: {}", e)) + })?; if !response.status().is_success() { let status = response.status(); @@ -639,13 +654,15 @@ impl InferenceEngine { ))); } - let json: serde_json::Value = response - .json() - .map_err(|e| VecStoreError::Internal(format!("Failed to parse Ollama response: {}", e)))?; + let json: serde_json::Value = response.json().map_err(|e| { + VecStoreError::Internal(format!("Failed to parse Ollama response: {}", e)) + })?; let embedding: Vec = json["embedding"] .as_array() - .ok_or_else(|| VecStoreError::Internal("No embedding in Ollama response".to_string()))? + .ok_or_else(|| { + VecStoreError::Internal("No embedding in Ollama response".to_string()) + })? .iter() .filter_map(|v| v.as_f64().map(|f| f as f32)) .collect(); @@ -664,9 +681,10 @@ impl InferenceEngine { _model: &str, _base_url: &str, ) -> Result>> { - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, self.dimension) - }).collect()) + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, self.dimension)) + .collect()) } /// HuggingFace embedding implementation @@ -676,9 +694,10 @@ impl InferenceEngine { _model: &str, _api_key: &str, ) -> Result>> { - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, self.dimension) - }).collect()) + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, self.dimension)) + .collect()) } /// Custom endpoint embedding implementation @@ -689,9 +708,10 @@ impl InferenceEngine { _api_key: Option<&str>, _headers: &HashMap, ) -> Result>> { - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, self.dimension) - }).collect()) + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, self.dimension)) + .collect()) } /// Fallback embedding when ONNX/VertexAI features are not enabled. @@ -703,10 +723,13 @@ impl InferenceEngine { /// These embeddings preserve some semantic similarity (same words produce similar hashes) /// but are not suitable for production semantic search. fn embed_local_placeholder(&self, texts: &[String]) -> Result>> { - tracing::debug!("Using placeholder embeddings - enable 'embeddings' feature for real ONNX inference"); - Ok(texts.iter().map(|text| { - Self::deterministic_embedding(text, self.dimension) - }).collect()) + tracing::debug!( + "Using placeholder embeddings - enable 'embeddings' feature for real ONNX inference" + ); + Ok(texts + .iter() + .map(|text| Self::deterministic_embedding(text, self.dimension)) + .collect()) } /// Create a deterministic embedding from text (for testing/fallback) @@ -741,13 +764,20 @@ impl InferenceEngine { /// Clear the embedding cache pub fn clear_cache(&self) { - let Ok(mut cache) = self.cache.write() else { return; }; + let Ok(mut cache) = self.cache.write() else { + return; + }; cache.clear(); } /// Get cache statistics pub fn cache_stats(&self) -> CacheStats { - let Ok(cache) = self.cache.read() else { return CacheStats { size: 0, max_size: 0 }; }; + let Ok(cache) = self.cache.read() else { + return CacheStats { + size: 0, + max_size: 0, + }; + }; cache.stats() } } diff --git a/src/ivf_pq.rs b/src/ivf_pq.rs index 1b2dc3c..72d8727 100644 --- a/src/ivf_pq.rs +++ b/src/ivf_pq.rs @@ -51,7 +51,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use rayon::prelude::*; use serde::{Deserialize, Serialize}; diff --git a/src/k8s_operator.rs b/src/k8s_operator.rs index b49e82d..a7a7064 100644 --- a/src/k8s_operator.rs +++ b/src/k8s_operator.rs @@ -303,7 +303,12 @@ impl VecStoreOperator { } /// Handle cluster creation - pub fn handle_create(&self, name: &str, namespace: &str, spec: VecStoreClusterSpec) -> Result<()> { + pub fn handle_create( + &self, + name: &str, + namespace: &str, + spec: VecStoreClusterSpec, + ) -> Result<()> { let cluster = ManagedCluster { name: name.to_string(), namespace: namespace.to_string(), @@ -325,8 +330,11 @@ impl VecStoreOperator { }; let key = format!("{}/{}", namespace, name); - self.clusters.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on clusters".into()))? + self.clusters + .write() + .map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on clusters".into()) + })? .insert(key.clone(), cluster); // Queue reconciliation @@ -336,10 +344,16 @@ impl VecStoreOperator { } /// Handle cluster update - pub fn handle_update(&self, name: &str, namespace: &str, spec: VecStoreClusterSpec) -> Result<()> { + pub fn handle_update( + &self, + name: &str, + namespace: &str, + spec: VecStoreClusterSpec, + ) -> Result<()> { let key = format!("{}/{}", namespace, name); - let mut clusters = self.clusters.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on clusters".into()))?; + let mut clusters = self.clusters.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on clusters".into()) + })?; if let Some(cluster) = clusters.get_mut(&key) { cluster.spec = spec; @@ -347,7 +361,10 @@ impl VecStoreOperator { self.queue_reconcile(name, namespace, ReconcileTrigger::Update); Ok(()) } else { - Err(VecStoreError::NotFound(format!("Cluster {} not found", key))) + Err(VecStoreError::NotFound(format!( + "Cluster {} not found", + key + ))) } } @@ -357,8 +374,12 @@ impl VecStoreOperator { self.queue_reconcile(name, namespace, ReconcileTrigger::Delete); // Mark for deletion (actual removal after cleanup) - if let Some(cluster) = self.clusters.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on clusters".into()))? + if let Some(cluster) = self + .clusters + .write() + .map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on clusters".into()) + })? .get_mut(&key) { cluster.status.phase = ClusterPhase::Terminating; @@ -375,7 +396,9 @@ impl VecStoreOperator { trigger, queued_at: Instant::now(), }; - let Ok(mut queue) = self.reconcile_queue.write() else { return; }; + let Ok(mut queue) = self.reconcile_queue.write() else { + return; + }; queue.push(request); } @@ -389,22 +412,25 @@ impl VecStoreOperator { // Update metrics use std::sync::atomic::Ordering; self.metrics.reconcile_total.fetch_add(1, Ordering::Relaxed); - self.metrics.reconcile_duration_sum.fetch_add( - start.elapsed().as_millis() as u64, - Ordering::Relaxed, - ); + self.metrics + .reconcile_duration_sum + .fetch_add(start.elapsed().as_millis() as u64, Ordering::Relaxed); if result.is_err() { - self.metrics.reconcile_errors.fetch_add(1, Ordering::Relaxed); + self.metrics + .reconcile_errors + .fetch_add(1, Ordering::Relaxed); } result } fn do_reconcile(&self, key: &str) -> Result { - let mut clusters = self.clusters.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire write lock on clusters".into()))?; - let cluster = clusters.get_mut(key) + let mut clusters = self.clusters.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire write lock on clusters".into()) + })?; + let cluster = clusters + .get_mut(key) .ok_or_else(|| VecStoreError::NotFound(format!("Cluster {} not found", key)))?; // Check if deletion is requested @@ -454,9 +480,12 @@ impl VecStoreOperator { // Check HA configuration if let Some(ref ha) = cluster.spec.ha_config - && ha.auto_failover && cluster.status.leader.is_none() && cluster.status.ready_replicas > 0 { - actions.push(ReconcileAction::ElectLeader); - } + && ha.auto_failover + && cluster.status.leader.is_none() + && cluster.status.ready_replicas > 0 + { + actions.push(ReconcileAction::ElectLeader); + } // Check backup schedule (Rust 1.92 if-let chain) if let Some(ref backup) = cluster.spec.backup @@ -478,48 +507,72 @@ impl VecStoreOperator { cluster.status.ready_replicas += count; cluster.status.phase = ClusterPhase::Running; - self.add_condition(cluster, ConditionType::ScalingUp, true, - "ScalingUp", &format!("Scaling up by {} replicas", count)); - } + self.add_condition( + cluster, + ConditionType::ScalingUp, + true, + "ScalingUp", + &format!("Scaling up by {} replicas", count), + ); + }, ReconcileAction::ScaleDown { count } => { cluster.status.phase = ClusterPhase::Scaling; cluster.status.total_replicas -= count; cluster.status.ready_replicas -= count; cluster.status.phase = ClusterPhase::Running; - self.add_condition(cluster, ConditionType::ScalingDown, true, - "ScalingDown", &format!("Scaling down by {} replicas", count)); - } + self.add_condition( + cluster, + ConditionType::ScalingDown, + true, + "ScalingDown", + &format!("Scaling down by {} replicas", count), + ); + }, ReconcileAction::UpdateConfig => { cluster.status.phase = ClusterPhase::Updating; // In real operator: rolling update of pods cluster.status.phase = ClusterPhase::Running; - self.add_condition(cluster, ConditionType::Progressing, true, - "ConfigUpdated", "Configuration updated successfully"); - } + self.add_condition( + cluster, + ConditionType::Progressing, + true, + "ConfigUpdated", + "Configuration updated successfully", + ); + }, ReconcileAction::ElectLeader => { // Simple leader election (in real operator: use Raft or lease) cluster.status.leader = Some(format!("{}-0", cluster.name)); - self.add_condition(cluster, ConditionType::Ready, true, - "LeaderElected", "Leader election completed"); - } + self.add_condition( + cluster, + ConditionType::Ready, + true, + "LeaderElected", + "Leader election completed", + ); + }, ReconcileAction::TriggerBackup => { // In real operator: create backup job cluster.status.last_backup = Some(chrono::Utc::now().to_rfc3339()); - self.add_condition(cluster, ConditionType::BackupReady, true, - "BackupCompleted", "Backup completed successfully"); - } + self.add_condition( + cluster, + ConditionType::BackupReady, + true, + "BackupCompleted", + "Backup completed successfully", + ); + }, ReconcileAction::Repair { reason } => { cluster.status.phase = ClusterPhase::Updating; // In real operator: repair degraded pods cluster.status.phase = ClusterPhase::Running; - self.add_condition(cluster, ConditionType::Degraded, false, - "Repaired", reason); - } + self.add_condition(cluster, ConditionType::Degraded, false, "Repaired", reason); + }, } Ok(()) @@ -543,19 +596,29 @@ impl VecStoreOperator { fn backup_due(&self, cluster: &ManagedCluster) -> bool { if let Some(ref last) = cluster.status.last_backup - && let Ok(last_time) = chrono::DateTime::parse_from_rfc3339(last) { - let now = chrono::Utc::now(); - let elapsed = now.signed_duration_since(last_time); - // Simple check: backup if more than 24 hours - return elapsed.num_hours() >= 24; - } + && let Ok(last_time) = chrono::DateTime::parse_from_rfc3339(last) + { + let now = chrono::Utc::now(); + let elapsed = now.signed_duration_since(last_time); + // Simple check: backup if more than 24 hours + return elapsed.num_hours() >= 24; + } true // No backup yet } - fn add_condition(&self, cluster: &mut ManagedCluster, ctype: ConditionType, - status: bool, reason: &str, message: &str) { + fn add_condition( + &self, + cluster: &mut ManagedCluster, + ctype: ConditionType, + status: bool, + reason: &str, + message: &str, + ) { // Remove existing condition of same type - cluster.status.conditions.retain(|c| c.condition_type != ctype); + cluster + .status + .conditions + .retain(|c| c.condition_type != ctype); cluster.status.conditions.push(ClusterCondition { condition_type: ctype, @@ -569,14 +632,20 @@ impl VecStoreOperator { /// Get cluster status pub fn get_status(&self, name: &str, namespace: &str) -> Option { let key = format!("{}/{}", namespace, name); - let Ok(clusters) = self.clusters.read() else { return None; }; + let Ok(clusters) = self.clusters.read() else { + return None; + }; clusters.get(&key).map(|c| c.status.clone()) } /// List all managed clusters pub fn list_clusters(&self) -> Vec<(String, String, VecStoreClusterStatus)> { - let Ok(clusters) = self.clusters.read() else { return Vec::new(); }; - clusters.values().map(|c| (c.namespace.clone(), c.name.clone(), c.status.clone())) + let Ok(clusters) = self.clusters.read() else { + return Vec::new(); + }; + clusters + .values() + .map(|c| (c.namespace.clone(), c.name.clone(), c.status.clone())) .collect() } @@ -735,12 +804,14 @@ spec: kind: VecStoreCluster shortNames: - vsc -"#.to_string() +"# + .to_string() } /// Generate StatefulSet manifest for a cluster pub fn generate_statefulset(spec: &VecStoreClusterSpec, name: &str, namespace: &str) -> String { - format!(r#"apiVersion: apps/v1 + format!( + r#"apiVersion: apps/v1 kind: StatefulSet metadata: name: {name} @@ -827,7 +898,8 @@ spec: /// Generate headless service manifest pub fn generate_headless_service(name: &str, namespace: &str) -> String { - format!(r#"apiVersion: v1 + format!( + r#"apiVersion: v1 kind: Service metadata: name: {name}-headless @@ -853,7 +925,8 @@ spec: /// Generate client service manifest pub fn generate_client_service(name: &str, namespace: &str) -> String { - format!(r#"apiVersion: v1 + format!( + r#"apiVersion: v1 kind: Service metadata: name: {name} @@ -882,7 +955,8 @@ spec: /// Generate HPA manifest pub fn generate_hpa(spec: &VecStoreClusterSpec, name: &str, namespace: &str) -> Option { spec.autoscaling.as_ref().map(|autoscaling| { - format!(r#"apiVersion: autoscaling/v2 + format!( + r#"apiVersion: autoscaling/v2 kind: HorizontalPodAutoscaler metadata: name: {name} @@ -918,7 +992,8 @@ spec: /// Generate PodDisruptionBudget manifest pub fn generate_pdb(spec: &VecStoreClusterSpec, name: &str, namespace: &str) -> Option { spec.ha_config.as_ref().map(|ha| { - format!(r#"apiVersion: policy/v1 + format!( + r#"apiVersion: policy/v1 kind: PodDisruptionBudget metadata: name: {name} @@ -939,7 +1014,8 @@ spec: /// Generate ServiceMonitor for Prometheus pub fn generate_service_monitor(name: &str, namespace: &str) -> String { - format!(r#"apiVersion: monitoring.coreos.com/v1 + format!( + r#"apiVersion: monitoring.coreos.com/v1 kind: ServiceMonitor metadata: name: {name} @@ -1067,12 +1143,14 @@ securityContext: nodeSelector: {} tolerations: [] affinity: {} -"#.to_string() +"# + .to_string() } /// Generate Chart.yaml pub fn generate_chart() -> String { - format!(r#"apiVersion: v2 + format!( + r#"apiVersion: v2 name: vecstore description: A Helm chart for VecStore vector database type: application @@ -1090,7 +1168,9 @@ sources: maintainers: - name: VecStore Team email: team@vecstore.io -"#, env!("CARGO_PKG_VERSION")) +"#, + env!("CARGO_PKG_VERSION") + ) } } @@ -1141,7 +1221,9 @@ mod tests { let operator = VecStoreOperator::new(OperatorConfig::default()); let spec = test_spec(); - operator.handle_create("test-cluster", "default", spec).unwrap(); + operator + .handle_create("test-cluster", "default", spec) + .unwrap(); let status = operator.get_status("test-cluster", "default").unwrap(); assert_eq!(status.phase, ClusterPhase::Pending); @@ -1152,7 +1234,9 @@ mod tests { let operator = VecStoreOperator::new(OperatorConfig::default()); let spec = test_spec(); - operator.handle_create("test-cluster", "default", spec).unwrap(); + operator + .handle_create("test-cluster", "default", spec) + .unwrap(); let result = operator.reconcile("test-cluster", "default").unwrap(); assert!(result.requeue); diff --git a/src/kafka_connector.rs b/src/kafka_connector.rs index 38b6e47..81c3651 100644 --- a/src/kafka_connector.rs +++ b/src/kafka_connector.rs @@ -35,7 +35,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; diff --git a/src/learned_index.rs b/src/learned_index.rs index ed8540b..8e9f09a 100644 --- a/src/learned_index.rs +++ b/src/learned_index.rs @@ -28,9 +28,9 @@ //! println!("Recommended beam width: {}", params.beam_width); //! ``` +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::time::Duration; -use serde::{Deserialize, Serialize}; use crate::error::VecStoreError; @@ -210,7 +210,12 @@ impl CostModel { &observation.params_used, &DataDistribution { count: 0, - dimensions: observation.profile.vector.as_ref().map(|v| v.len()).unwrap_or(128), + dimensions: observation + .profile + .vector + .as_ref() + .map(|v| v.len()) + .unwrap_or(128), centroid: Vec::new(), variance: Vec::new(), intrinsic_dim: 0.0, @@ -261,7 +266,9 @@ impl RecallPredictor { profile.filter_selectivity as f64, ]; - let raw: f64 = self.coefficients.iter() + let raw: f64 = self + .coefficients + .iter() .zip(features.iter()) .map(|(c, f)| c * f) .sum(); @@ -271,14 +278,22 @@ impl RecallPredictor { } /// Update predictor with observation - pub fn update(&mut self, params: &SearchParams, profile: &QueryProfile, actual_recall: f64, learning_rate: f64) { + pub fn update( + &mut self, + params: &SearchParams, + profile: &QueryProfile, + actual_recall: f64, + learning_rate: f64, + ) { let predicted = self.predict(params, profile); let error = actual_recall - predicted; - let features = [1.0, + let features = [ + 1.0, params.ef_search as f64, params.beam_width as f64, - profile.filter_selectivity as f64]; + profile.filter_selectivity as f64, + ]; // Gradient descent update for (i, f) in features.iter().enumerate() { @@ -353,27 +368,43 @@ impl WorkloadAnalyzer { for obs in &self.history { let key = format!( "k={}_filter={}", - if obs.profile.k <= 10 { "small" } else if obs.profile.k <= 100 { "medium" } else { "large" }, + if obs.profile.k <= 10 { + "small" + } else if obs.profile.k <= 100 { + "medium" + } else { + "large" + }, obs.profile.has_filter ); patterns.entry(key).or_default().push(obs); } - let total_time = self.history.last() - .and_then(|last| self.history.first().map(|first| { - last.timestamp.duration_since(first.timestamp).unwrap_or_default().as_secs_f64() - })) + let total_time = self + .history + .last() + .and_then(|last| { + self.history.first().map(|first| { + last.timestamp + .duration_since(first.timestamp) + .unwrap_or_default() + .as_secs_f64() + }) + }) .unwrap_or(1.0); - self.patterns = patterns.into_iter() + self.patterns = patterns + .into_iter() .map(|(name, obs)| { let count = obs.len(); let frequency = count as f64 / total_time.max(1.0); let avg_k = obs.iter().map(|o| o.profile.k as f64).sum::() / count as f64; - let filter_pct = obs.iter().filter(|o| o.profile.has_filter).count() as f64 / count as f64; + let filter_pct = + obs.iter().filter(|o| o.profile.has_filter).count() as f64 / count as f64; // Find best params (lowest latency meeting recall target) - let best = obs.iter() + let best = obs + .iter() .filter(|o| o.recall.unwrap_or(1.0) >= 0.9) .min_by(|a, b| a.latency_ms.partial_cmp(&b.latency_ms).unwrap()); @@ -399,10 +430,25 @@ impl WorkloadAnalyzer { let total = self.history.len(); let avg_latency = self.history.iter().map(|o| o.latency_ms).sum::() / total as f64; let avg_k = self.history.iter().map(|o| o.profile.k as f64).sum::() / total as f64; - let filter_ratio = self.history.iter().filter(|o| o.profile.has_filter).count() as f64 / total as f64; - - let p50_latency = percentile(&self.history.iter().map(|o| o.latency_ms).collect::>(), 0.5); - let p99_latency = percentile(&self.history.iter().map(|o| o.latency_ms).collect::>(), 0.99); + let filter_ratio = + self.history.iter().filter(|o| o.profile.has_filter).count() as f64 / total as f64; + + let p50_latency = percentile( + &self + .history + .iter() + .map(|o| o.latency_ms) + .collect::>(), + 0.5, + ); + let p99_latency = percentile( + &self + .history + .iter() + .map(|o| o.latency_ms) + .collect::>(), + 0.99, + ); WorkloadSummary { total_queries: total, @@ -510,7 +556,8 @@ impl LearnedIndex { // Update models if self.config.continuous_learning { - self.cost_model.update(&observation, self.config.learning_rate); + self.cost_model + .update(&observation, self.config.learning_rate); if let Some(actual_recall) = recall { self.recall_predictor.update( @@ -675,7 +722,8 @@ impl LearnedIndex { let mut min_dist = f32::MAX; for (j, v2) in vectors.iter().enumerate() { if i != j { - let dist: f32 = v1.iter() + let dist: f32 = v1 + .iter() .zip(v2.iter()) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -721,9 +769,13 @@ impl LearnedIndex { } // Adjust beam width similarly - let avg_recall = self.workload.history.iter() + let avg_recall = self + .workload + .history + .iter() .filter_map(|o| o.recall) - .sum::() / self.workload.history.len().max(1) as f64; + .sum::() + / self.workload.history.len().max(1) as f64; if avg_recall < self.config.target_recall { self.current_params.beam_width = (self.current_params.beam_width as f64 * 1.2) as usize; @@ -760,7 +812,8 @@ impl LearnedIndex { if profile.has_filter && profile.filter_selectivity < 0.5 { // Low selectivity = need more candidates let mut params = base.clone(); - params.ef_search = (base.ef_search as f32 / profile.filter_selectivity.max(0.1)) as usize; + params.ef_search = + (base.ef_search as f32 / profile.filter_selectivity.max(0.1)) as usize; params.ef_search = params.ef_search.min(500); candidates.push(params); } @@ -814,11 +867,18 @@ pub struct BenchmarkResult { impl ParamBenchmark { /// Create a new benchmark pub fn new() -> Self { - Self { results: Vec::new() } + Self { + results: Vec::new(), + } } /// Run benchmark with given parameters - pub fn run(&mut self, params: SearchParams, query_fn: F, iterations: usize) -> BenchmarkResult + pub fn run( + &mut self, + params: SearchParams, + query_fn: F, + iterations: usize, + ) -> BenchmarkResult where F: Fn(&SearchParams) -> (Duration, Option), { @@ -859,10 +919,17 @@ impl ParamBenchmark { } /// Find best parameters meeting constraints - pub fn best_params(&self, min_recall: Option, max_latency: Option) -> Option<&BenchmarkResult> { - self.results.iter() + pub fn best_params( + &self, + min_recall: Option, + max_latency: Option, + ) -> Option<&BenchmarkResult> { + self.results + .iter() .filter(|r| { - let recall_ok = min_recall.map(|mr| r.recall.unwrap_or(1.0) >= mr).unwrap_or(true); + let recall_ok = min_recall + .map(|mr| r.recall.unwrap_or(1.0) >= mr) + .unwrap_or(true); let latency_ok = max_latency.map(|ml| r.avg_latency_ms <= ml).unwrap_or(true); recall_ok && latency_ok }) @@ -906,7 +973,9 @@ mod tests { let mut index = LearnedIndex::new(config); for i in 0..100 { - index.add(format!("v{}", i), vec![i as f32, 0.0, 0.0]).unwrap(); + index + .add(format!("v{}", i), vec![i as f32, 0.0, 0.0]) + .unwrap(); } index.update_distribution(); @@ -930,13 +999,7 @@ mod tests { category: None, }; - index.observe_query( - profile, - 5.0, - Some(0.95), - 10, - SearchParams::default(), - ); + index.observe_query(profile, 5.0, Some(0.95), 10, SearchParams::default()); let summary = index.workload_summary(); assert_eq!(summary.total_queries, 1); diff --git a/src/lib.rs b/src/lib.rs index ee18f4f..8ac449d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -295,65 +295,65 @@ pub mod object_storage; pub mod pq; // 2026 Competitive Features - Phase 2 -pub mod inference; // Integrated Inference API (embed at query time) -pub mod sql; // SQL Interface for vector operations -pub mod rabitq; // RaBitQ 1-bit quantization (30x compression) -pub mod generative; // Native Generative Search (RAG in single call) -pub mod llm_cache; // Semantic LLM Cache (query deduplication) -pub mod multilang; // Multi-Language Stemming (17+ languages) -pub mod typo_tolerance; // Typo-Tolerant Search (fuzzy matching) -pub mod graphrag; // GraphRAG Integration (knowledge graphs) -pub mod streaming; // Streaming Index Updates (real-time) -pub mod edge_runtime; // Edge Runtime (Workers/Vercel) -pub mod ml_ranking; // ONNX/ML Models in Ranking +pub mod edge_runtime; // Edge Runtime (Workers/Vercel) +pub mod generative; // Native Generative Search (RAG in single call) +pub mod graphrag; // GraphRAG Integration (knowledge graphs) +pub mod inference; // Integrated Inference API (embed at query time) +pub mod llm_cache; // Semantic LLM Cache (query deduplication) +pub mod ml_ranking; +pub mod multilang; // Multi-Language Stemming (17+ languages) +pub mod rabitq; // RaBitQ 1-bit quantization (30x compression) +pub mod sql; // SQL Interface for vector operations +pub mod streaming; // Streaming Index Updates (real-time) +pub mod typo_tolerance; // Typo-Tolerant Search (fuzzy matching) // ONNX/ML Models in Ranking // 2026 Competitive Features - Phase 3 (Market Domination) -pub mod sparse_native; // Native SPLADE/BM25 sparse vectors -pub mod serverless; // Pinecone-style auto-scaling -pub mod timetravel; // Milvus-style point-in-time queries -pub mod payload_index; // Qdrant-style B-tree/Hash/Geo indexes -pub mod partition_keys; // Multi-tenant partition routing -pub mod graphql_api; // Weaviate-style GraphQL interface +pub mod dynamic_schema; // Runtime schema flexibility +pub mod embedded_onnx; +pub mod graphql_api; // Weaviate-style GraphQL interface +pub mod partition_keys; // Multi-tenant partition routing +pub mod payload_index; // Qdrant-style B-tree/Hash/Geo indexes pub mod recommendations; // Item-to-item recommendations -pub mod zero_copy; // LanceDB-style versioning -pub mod dynamic_schema; // Runtime schema flexibility -pub mod embedded_onnx; // Chroma-style local embeddings +pub mod serverless; // Pinecone-style auto-scaling +pub mod sparse_native; // Native SPLADE/BM25 sparse vectors +pub mod timetravel; // Milvus-style point-in-time queries +pub mod zero_copy; // LanceDB-style versioning // Chroma-style local embeddings // 2026 Competitive Features - Phase 4 (Market Leadership) -pub mod agentic; // Agentic Query Framework (multi-step retrieval) -pub mod object_store; // Object Storage Backend (S3/GCS, 100x cost reduction) -pub mod hybrid_rerank; // Hybrid Reranking Pipeline (dense+sparse+BM25) -pub mod managed_embeddings; // Managed Embedding Service (pay-as-you-go) -pub mod advanced_quant; // Advanced Quantization (1.5-bit, 2-bit, asymmetric) -pub mod multimodal_search; // Multimodal Search (text↔image↔video) +pub mod advanced_quant; // Advanced Quantization (1.5-bit, 2-bit, asymmetric) +pub mod agentic; // Agentic Query Framework (multi-step retrieval) +pub mod change_streams; // Change Streams (real-time subscriptions) +pub mod discovery; // Discovery API (vector-space exploration) pub mod enterprise_protect; // Enterprise Protection (CMEK, backup/restore) -pub mod discovery; // Discovery API (vector-space exploration) -pub mod change_streams; // Change Streams (real-time subscriptions) -pub mod gpu_indexing; // GPU Index Building (10x faster ingestion) +pub mod gpu_indexing; +pub mod hybrid_rerank; // Hybrid Reranking Pipeline (dense+sparse+BM25) +pub mod managed_embeddings; // Managed Embedding Service (pay-as-you-go) +pub mod multimodal_search; // Multimodal Search (text↔image↔video) +pub mod object_store; // Object Storage Backend (S3/GCS, 100x cost reduction) // GPU Index Building (10x faster ingestion) // 2026 Competitive Features - Phase 5 (Production Readiness) -pub mod otel; // OpenTelemetry OTLP Export (distributed tracing) -pub mod query_explain; // Visual Query Explain (EXPLAIN ANALYZE) -pub mod adaptive_cache; // Multi-Tier Adaptive Cache (L1/L2/L3) -pub mod backpressure; // Streaming Backpressure (flow control) -pub mod rest_api; // REST API with OpenAPI (spec generation) -pub mod raft_consensus; // Production Raft Consensus (leader election) -pub mod sdk_gen; // Multi-Language SDK Generation (TS/Python/Go) +pub mod adaptive_cache; // Multi-Tier Adaptive Cache (L1/L2/L3) +pub mod backpressure; // Streaming Backpressure (flow control) +pub mod otel; // OpenTelemetry OTLP Export (distributed tracing) +pub mod query_explain; // Visual Query Explain (EXPLAIN ANALYZE) +pub mod raft_consensus; // Production Raft Consensus (leader election) +pub mod rest_api; // REST API with OpenAPI (spec generation) +pub mod sdk_gen; // Multi-Language SDK Generation (TS/Python/Go) // 2026 Competitive Features - Phase 6 (Cloud-Native Operations) -pub mod k8s_operator; // Kubernetes Operator (CRD, Helm charts) -pub mod query_router; // Query Router / Load Balancer (circuit breaker) -pub mod data_lineage; // Data Lineage Tracking (provenance, impact analysis) -pub mod ab_testing; // A/B Testing Framework (experiments, variants) +pub mod ab_testing; // A/B Testing Framework (experiments, variants) +pub mod data_lineage; // Data Lineage Tracking (provenance, impact analysis) +pub mod federation; pub mod incremental_index; // Incremental Index Updates (delta indexes, WAL) -pub mod federation; // Federation (cross-store queries, multi-cluster) +pub mod k8s_operator; // Kubernetes Operator (CRD, Helm charts) +pub mod query_router; // Query Router / Load Balancer (circuit breaker) // Federation (cross-store queries, multi-cluster) // 2026 Competitive Features - Phase 7 (Advanced Analytics & ML) -pub mod neural_ranker; // Neural Ranker (cross-encoder, late interaction, SPLADE) -pub mod auto_tuning; // Auto-Tuning Engine (Bayesian optimization, hyperparameter search) pub mod anomaly_detection; // Anomaly Detection (isolation forest, LOF, drift monitoring) -pub mod query_analytics; // Query Analytics Dashboard (performance insights, trends) -pub mod cost_optimizer; // Cost Optimizer (resource right-sizing, auto-scaling) +pub mod auto_tuning; // Auto-Tuning Engine (Bayesian optimization, hyperparameter search) +pub mod cost_optimizer; +pub mod neural_ranker; // Neural Ranker (cross-encoder, late interaction, SPLADE) +pub mod query_analytics; // Query Analytics Dashboard (performance insights, trends) // Cost Optimizer (resource right-sizing, auto-scaling) #[cfg(feature = "async")] pub mod kafka_connector; @@ -368,11 +368,11 @@ pub use namespace::{Namespace, NamespaceId, NamespaceQuotas, NamespaceStatus, Re pub use namespace_manager::{AggregateStats, NamespaceManager, NamespaceStats}; pub use schema::{FieldSchema, FieldType, Schema, ValidationError}; pub use store::{ - make_record, parse_filter, BatchError, BatchOperation, BatchResult, CompactionConfig, - CompactionResult, Config, Distance, ExplainedNeighbor, FilterExpr, FilterOp, FilterParseError, - HNSWSearchParams, HybridQuery, Metadata, Neighbor, PQConfig, PQVectorStore, PrefetchQuery, - ProductQuantizer, QuantizationConfig, QuantizationStats, Query, QueryEstimate, QueryExplanation, - QueryPlan, QueryStage, QueryStep, Record, VecStore, VecStoreBuilder, + BatchError, BatchOperation, BatchResult, CompactionConfig, CompactionResult, Config, Distance, + ExplainedNeighbor, FilterExpr, FilterOp, FilterParseError, HNSWSearchParams, HybridQuery, + Metadata, Neighbor, PQConfig, PQVectorStore, PrefetchQuery, ProductQuantizer, + QuantizationConfig, QuantizationStats, Query, QueryEstimate, QueryExplanation, QueryPlan, + QueryStage, QueryStep, Record, VecStore, VecStoreBuilder, make_record, parse_filter, }; pub use text_splitter::{ RecursiveCharacterTextSplitter, TextChunk, TextSplitter, TokenTextSplitter, @@ -380,8 +380,8 @@ pub use text_splitter::{ // Export fuzzy search types pub use fuzzy::{ - damerau_levenshtein_distance, levenshtein_distance, similarity_score, suggest_corrections, - BKTree, FuzzyMatcher, + BKTree, FuzzyMatcher, damerau_levenshtein_distance, levenshtein_distance, similarity_score, + suggest_corrections, }; // Export auto-tuning types @@ -389,8 +389,8 @@ pub use autotuning::{AutoTuner, HnswParams, PerformanceConstraints, TuningGoal}; // Export compression types pub use compression::{ - decode_rle, decode_varint, encode_rle, encode_varint, CompressedNeighborList, - CompressionConfig, CompressionLevel, CompressionMethod, CompressionStats, + CompressedNeighborList, CompressionConfig, CompressionLevel, CompressionMethod, + CompressionStats, decode_rle, decode_varint, encode_rle, encode_varint, }; // Export real-time indexing types @@ -453,8 +453,9 @@ pub use benchmark::{ // Export health check types pub use health::{ - print_health_report, Alert, AlertCategory, AlertSeverity, DatabaseHealth, HealthCheckConfig, - HealthChecker, HealthReport, HealthStatus, IndexHealth, PerformanceHealth, ResourceHealth, + Alert, AlertCategory, AlertSeverity, DatabaseHealth, HealthCheckConfig, HealthChecker, + HealthReport, HealthStatus, IndexHealth, PerformanceHealth, ResourceHealth, + print_health_report, }; // Export metadata indexing types @@ -485,7 +486,7 @@ pub use anomaly::{ }; // Export dimensionality reduction types -pub use dim_reduction::{ReductionStats, PCA}; +pub use dim_reduction::{PCA, ReductionStats}; // Export recommender system types pub use recommender::{ @@ -560,7 +561,7 @@ pub use geospatial::{ }; // Export advanced filter types -pub use advanced_filter::{parse_advanced_filter, AdvancedFilter, FilterBuilder}; +pub use advanced_filter::{AdvancedFilter, FilterBuilder, parse_advanced_filter}; // Export profiler types pub use profiler::{ProfileStage, ProfileSummary, ProfilerConfig, QueryProfile, QueryProfiler}; @@ -568,9 +569,8 @@ pub use profiler::{ProfileStage, ProfileSummary, ProfilerConfig, QueryProfile, Q #[cfg(feature = "async")] // Export Kafka connector types pub use kafka_connector::{ - ConsumerStats, KafkaConfig, KafkaConsumer, KafkaProducer, - Operation as KafkaOperation, PipelineStats as KafkaPipelineStats, - ProducerStats, StreamingPipeline, VectorMessage, + ConsumerStats, KafkaConfig, KafkaConsumer, KafkaProducer, Operation as KafkaOperation, + PipelineStats as KafkaPipelineStats, ProducerStats, StreamingPipeline, VectorMessage, }; #[cfg(feature = "async")] @@ -579,9 +579,9 @@ pub use python_async::{AsyncPyVecStore, AsyncSearchResult}; // Export vectors module types pub use vectors::{ - bm25_score, bm25_score_simple, hybrid_search_score, normalize_scores, normalize_scores_zscore, BM25Config, BM25Stats, FusionStrategy, HybridQuery as HybridQueryV2, HybridSearchConfig, - KMeans, Vector, VectorOps, + KMeans, Vector, VectorOps, bm25_score, bm25_score_simple, hybrid_search_score, + normalize_scores, normalize_scores_zscore, }; // Export SIMD-accelerated distance functions @@ -613,114 +613,122 @@ pub use embeddings::{CandleEmbedder, CandleModel}; pub use wasm::{WasmSearchResult, WasmVecStore}; // Export 2026 Competitive Features - Phase 2 -pub use inference::{EmbeddingProvider, InferenceConfig, InferenceEngine, InferenceStats}; -pub use sql::{DataType, Row, TableSchema, VectorSQL}; -pub use rabitq::{BinaryCode, RaBitQ, RaBitQConfig, RaBitQIndex}; +pub use edge_runtime::{EdgeConfig, EdgeSearchResult, EdgeVectorStore}; pub use generative::{ Citation, GenerativeConfig, GenerativeResponse, GenerativeSearch, LLMProvider, RetrievedDocument, TokenUsage, }; -pub use llm_cache::{CacheConfig, CacheResult, SemanticCache}; -pub use multilang::{AutoLangAnalyzer, Language, MultiLangAnalyzer}; -pub use typo_tolerance::{MatchType, TypoConfig, TypoSearchResult, TypoTolerantSearch}; pub use graphrag::{Entity as GraphEntity, GraphRAG as GraphRAGV2, Relationship}; -pub use streaming::{ConflictStrategy, StreamConfig, StreamingIndex, StreamingStats}; -pub use edge_runtime::{EdgeConfig, EdgeSearchResult, EdgeVectorStore}; +pub use inference::{EmbeddingProvider, InferenceConfig, InferenceEngine, InferenceStats}; +pub use llm_cache::{CacheConfig, CacheResult, SemanticCache}; pub use ml_ranking::{ MLRanker, RankedResult, RankingCandidate, RankingConfig, RankingFeatures, RankingModel, RankingPipeline, RankingStage, }; +pub use multilang::{AutoLangAnalyzer, Language, MultiLangAnalyzer}; +pub use rabitq::{BinaryCode, RaBitQ, RaBitQConfig, RaBitQIndex}; +pub use sql::{DataType, Row, TableSchema, VectorSQL}; +pub use streaming::{ConflictStrategy, StreamConfig, StreamingIndex, StreamingStats}; +pub use typo_tolerance::{MatchType, TypoConfig, TypoSearchResult, TypoTolerantSearch}; // Export 2026 Competitive Features - Phase 3 (Market Domination) -pub use sparse_native::{ - HybridSparseIndex, InvertedSparseIndex, SparseVector as NativeSparseVector, - SpladeEncoder as NativeSpladeEncoder, -}; -pub use serverless::{ - Autoscaler, ColdStartOptimizer, LoadBalancer, ReplicaInfo, ServerlessCluster, ServerlessConfig, - ServerlessStats, -}; -pub use timetravel::{ - SnapshotManager, TimeTravelIndex, TimeTravelResult, Timestamp, VersionEntry, - VersionHistory as TimeTravelVersionHistory, +pub use dynamic_schema::{ + DynamicCollection, FieldDef, FieldType as DynamicFieldType, Schema as DynamicSchema, + SchemaBuilder, SchemaChange, SchemaEvolution, }; -pub use payload_index::{ - Filter, GeoIndex as PayloadGeoIndex, GeoPoint as PayloadGeoPoint, HashIndex as PayloadHashIndex, - BTreeIndex as PayloadBTreeIndex, FullTextIndex, PayloadIndexManager, IndexType as PayloadIndexType, +pub use embedded_onnx::{EmbeddedStore, EmbeddingFunction, EmbeddingModel, ModelConfig, ModelType}; +pub use graphql_api::{ + GqlCollection, GqlField, GqlType, GraphQLExecutor, NearTextInput, NearVectorInput, VectorQuery, + WhereFilter, WhereOperator, }; pub use partition_keys::{ HashAlgorithm, Partition, PartitionConfig as PartitionKeyConfig, PartitionRouter, - PartitionedIndex, PartitionState, + PartitionState, PartitionedIndex, }; -pub use graphql_api::{ - GqlCollection, GqlField, GqlType, GraphQLExecutor, NearTextInput, NearVectorInput, VectorQuery, - WhereFilter, WhereOperator, +pub use payload_index::{ + BTreeIndex as PayloadBTreeIndex, Filter, FullTextIndex, GeoIndex as PayloadGeoIndex, + GeoPoint as PayloadGeoPoint, HashIndex as PayloadHashIndex, IndexType as PayloadIndexType, + PayloadIndexManager, }; pub use recommendations::{ RecommendRequest, RecommendResult, RecommendStrategy, RecommendationEngine, UserProfile as RecommendUserProfile, }; -pub use zero_copy::{ - Branch, BranchName, Commit, CommitHash, Diff, MergeResult, VersionedStore as CowVersionedStore, +pub use serverless::{ + Autoscaler, ColdStartOptimizer, LoadBalancer, ReplicaInfo, ServerlessCluster, ServerlessConfig, + ServerlessStats, }; -pub use dynamic_schema::{ - DynamicCollection, FieldDef, FieldType as DynamicFieldType, Schema as DynamicSchema, - SchemaBuilder, SchemaChange, SchemaEvolution, +pub use sparse_native::{ + HybridSparseIndex, InvertedSparseIndex, SparseVector as NativeSparseVector, + SpladeEncoder as NativeSpladeEncoder, +}; +pub use timetravel::{ + SnapshotManager, TimeTravelIndex, TimeTravelResult, Timestamp, VersionEntry, + VersionHistory as TimeTravelVersionHistory, }; -pub use embedded_onnx::{ - EmbeddedStore, EmbeddingFunction, EmbeddingModel, ModelConfig, ModelType, +pub use zero_copy::{ + Branch, BranchName, Commit, CommitHash, Diff, MergeResult, VersionedStore as CowVersionedStore, }; // Export 2026 Competitive Features - Phase 4 (Market Leadership) +pub use advanced_quant::{ + AdaptiveQuantizer, AsymmetricQuantizer, BinaryQuantizer as AdvancedBinaryQuantizer, BitWidth, + FourBitQuantizer, QuantConfig, QuantParams, QuantStrategy, QuantizedIndex, QuantizedIndexStats, + QuantizedVector, TernaryQuantizer, TwoBitQuantizer, +}; pub use agentic::{ - AgentConfig, AgentExecutor, AgentExecutionResult, AgentState, AgentStats, + AgentConfig, AgentExecutionResult, AgentExecutor, AgentState, AgentStats, DecompositionStrategy, FusedResult, FusionStrategy as AgenticFusionStrategy, QueryPlan as AgenticQueryPlan, QueryPlanner, SubQuery, ToolCall, ToolDefinition, ToolResult, }; -pub use object_store::{ - CacheStats, CacheTier, CostEstimate, ObjectMetadata, ObjectStoreBackend, ObjectStoreConfig, - ObjectStoreStats, PreWarmingConfig, PreWarmStrategy, StorageProvider, StoredVector, - TieredStorageManager, TieredStorageStats, +pub use change_streams::{ + ChangeEvent, ChangeFilter, ChangeMetadata, ChangeSource, ChangeStream, ChangeStreamManager, + ChangeStreamStats, DocumentSnapshot, Operation, ResumeToken, Subscription, +}; +pub use discovery::{ + DiscoveryConfig, DiscoveryIndex, DiscoveryMode, DiscoveryQuery, DiscoveryResult, + DiscoveryStats, ExplorationResult, Region, +}; +pub use enterprise_protect::{ + BackupConfig, BackupManager, BackupMetadata, BackupStats, CMEKConfig, ComplianceChecker, + ComplianceFramework, ComplianceStatus, ComplianceSummary, EncryptedData, EncryptionManager, + EnterpriseConfig, KMSProvider, KeyRotationResult, RetentionAction, RetentionPolicy, +}; +pub use gpu_indexing::{ + BuildStats, GpuBackendType, GpuConfig as GpuIndexConfig, GpuDeviceInfo as GpuIndexDeviceInfo, + GpuDistanceMetric, GpuHnswIndex, GpuIndexBuilder, GpuIndexStats, HnswParams as GpuHnswParams, + MemoryPoolStats, }; pub use hybrid_rerank::{ BooleanQuery, FusionMethod, HybridConfig, HybridDocument, HybridPipeline, HybridResult, - HybridSearchQuery, PipelineStats, RerankerConfig, RerankerType, ScoreExplanation, - SignalType, SparseQueryVector, + HybridSearchQuery, PipelineStats, RerankerConfig, RerankerType, ScoreExplanation, SignalType, + SparseQueryVector, }; pub use managed_embeddings::{ - EmbeddingCollection as ManagedEmbeddingCollection, EmbeddingProvider as ManagedEmbeddingProvider, EmbeddingRequest, - EmbeddingResponse, EmbeddingService, EmbeddingServiceConfig, - ModelConfig as EmbeddingModelConfig, QueryResult as EmbeddingQueryResult, TaskType, UsageStats, -}; -pub use advanced_quant::{ - AdaptiveQuantizer, AsymmetricQuantizer, BinaryQuantizer as AdvancedBinaryQuantizer, - BitWidth, FourBitQuantizer, QuantConfig, QuantizedIndex, QuantizedIndexStats, - QuantizedVector, QuantParams, QuantStrategy, TernaryQuantizer, TwoBitQuantizer, + EmbeddingCollection as ManagedEmbeddingCollection, + EmbeddingProvider as ManagedEmbeddingProvider, EmbeddingRequest, EmbeddingResponse, + EmbeddingService, EmbeddingServiceConfig, ModelConfig as EmbeddingModelConfig, + QueryResult as EmbeddingQueryResult, TaskType, UsageStats, }; pub use multimodal_search::{ Modality as SearchModality, MultimodalConfig, MultimodalDocument, MultimodalIndex, MultimodalQuery, MultimodalResult, MultimodalStats, }; -pub use enterprise_protect::{ - BackupConfig, BackupManager, BackupMetadata, BackupStats, CMEKConfig, ComplianceChecker, - ComplianceFramework, ComplianceStatus, ComplianceSummary, EncryptedData, EncryptionManager, - EnterpriseConfig, KMSProvider, KeyRotationResult, RetentionAction, RetentionPolicy, -}; -pub use discovery::{ - DiscoveryConfig, DiscoveryIndex, DiscoveryMode, DiscoveryQuery, DiscoveryResult, - DiscoveryStats, ExplorationResult, Region, -}; -pub use change_streams::{ - ChangeEvent, ChangeFilter, ChangeMetadata, ChangeSource, ChangeStream, ChangeStreamManager, - ChangeStreamStats, DocumentSnapshot, Operation, ResumeToken, Subscription, -}; -pub use gpu_indexing::{ - BuildStats, GpuBackendType, GpuConfig as GpuIndexConfig, GpuDeviceInfo as GpuIndexDeviceInfo, - GpuDistanceMetric, GpuHnswIndex, GpuIndexBuilder, GpuIndexStats, - HnswParams as GpuHnswParams, MemoryPoolStats, +pub use object_store::{ + CacheStats, CacheTier, CostEstimate, ObjectMetadata, ObjectStoreBackend, ObjectStoreConfig, + ObjectStoreStats, PreWarmStrategy, PreWarmingConfig, StorageProvider, StoredVector, + TieredStorageManager, TieredStorageStats, }; // Export 2026 Competitive Features - Phase 5 (Production Readiness) +pub use adaptive_cache::{ + AdaptiveCache, BloomFilter, CacheConfig as AdaptiveCacheConfig, CacheEntry, CacheKeyGenerator, + CacheStats as AdaptiveCacheStats, CachedResult, TierStats, +}; +pub use backpressure::{ + BackpressureController, BackpressureState, BatchProcessor, Checkpoint, DeadLetterRecord, + LagMetrics, StreamConfig as BackpressureStreamConfig, StreamRecord, StreamStats, +}; pub use otel::{ BatchConfig as OtelBatchConfig, ExporterType, MetricPoint as OtelMetricPoint, MetricType as OtelMetricType, OtelConfig, OtlpExportData, OtlpSpan, SamplingStrategy, @@ -731,91 +739,90 @@ pub use query_explain::{ ExecutionStats, ExplainFormat, ExplainOptions, ExplainResult, IndexStats as ExplainIndexStats, OptimizationSuggestion, PlanNode, PlanNodeType, QueryExplainer, SuggestionType, }; -pub use adaptive_cache::{ - AdaptiveCache, BloomFilter, CacheConfig as AdaptiveCacheConfig, CacheEntry, - CacheKeyGenerator, CacheStats as AdaptiveCacheStats, CachedResult, TierStats, -}; -pub use backpressure::{ - BackpressureController, BackpressureState, BatchProcessor, Checkpoint, - DeadLetterRecord, LagMetrics, StreamConfig as BackpressureStreamConfig, StreamRecord, - StreamStats, -}; -pub use rest_api::{ - ApiInfo, ApiRateLimiter, HttpMethod, OpenApiGenerator, Request as RestRequest, - Response as RestResponse, RestConfig, Route as RestRoute, Router as RestRouter, -}; pub use raft_consensus::{ AppendEntriesRequest, AppendEntriesResponse, ClusterStatus, EntryType, InstallSnapshotRequest, LeaderState, LogEntry, PersistentState, RaftConfig, RaftError, RaftNode, RaftState, SnapshotMetadata, VolatileState, VoteRequest, VoteResponse, }; +pub use rest_api::{ + ApiInfo, ApiRateLimiter, HttpMethod, OpenApiGenerator, Request as RestRequest, + Response as RestResponse, RestConfig, Route as RestRoute, Router as RestRouter, +}; pub use sdk_gen::{ FieldDef as SdkFieldDef, FieldType as SdkFieldType, Language as SdkLanguage, MethodDef, ParamDef, ParamLocation, SdkGenerator, TypeDef, }; // Export 2026 Competitive Features - Phase 6 (Cloud-Native Operations) -pub use k8s_operator::{ - AutoscalingConfig, BackupConfig as K8sBackupConfig, ClusterCondition, ClusterPhase, - ConditionType, DistanceMetric as K8sDistanceMetric, HAConfig, HelmChartGenerator, - IndexConfig as K8sIndexConfig, ManifestGenerator, MonitoringConfig as K8sMonitoringConfig, - OperatorConfig, OperatorMetricsSnapshot, ReconcileAction, ReconcileResult, - ResourceRequirements, ResourceSpec, VecStoreClusterSpec, VecStoreClusterStatus, - VecStoreOperator, -}; -pub use query_router::{ - CircuitState as RouterCircuitState, FilterComplexity, HealthStatus as RouterHealthStatus, - LoadBalanceStrategy, NodeCapabilities, NodeStats, QueryCharacteristics, QueryPriority, - QueryRouter, ReplicaNode, RequestContext, RetryConfig, RetryHandler, RouterConfig, - RouterStats, RoutingDecision, RoutingReason, -}; -pub use data_lineage::{ - AccessRecord, AccessType, DerivationOperation, ImpactReport, LineageConfig, - LineageRecord, LineageStats, LineageTracker, QualityMetrics as LineageQualityMetrics, - Transformation, TransformationType, VectorId, VectorSource, -}; pub use ab_testing::{ ABTestManager, ABTestStats, ExperimentConfig, ExperimentResults, ExperimentStatus, ExperimentSummary, ExperimentType, IndexParams as ABIndexParams, Metric as ABMetric, - MetricSummary, MetricType as ABMetricType, RerankingConfig as ABRerankingConfig, - RequestResult as ABRequestResult, SearchParams as ABSearchParams, TrafficAllocation, + MetricSummary, MetricType as ABMetricType, RequestResult as ABRequestResult, + RerankingConfig as ABRerankingConfig, SearchParams as ABSearchParams, TrafficAllocation, Variant, VariantAssignment, VariantComparison, VariantConfig, VariantResult, }; +pub use data_lineage::{ + AccessRecord, AccessType, DerivationOperation, ImpactReport, LineageConfig, LineageRecord, + LineageStats, LineageTracker, QualityMetrics as LineageQualityMetrics, Transformation, + TransformationType, VectorId, VectorSource, +}; +pub use federation::{ + FederatedQuery, FederatedQueryResult, FederatedResult, Federation, FederationBuilder, + FederationConfig, FederationMember, FederationStatus, MemberCapabilities, MemberHealth, + MemberStatus, MemberType, MergeStrategy, +}; pub use incremental_index::{ BatchResult as IncrementalBatchResult, IncrementalConfig, IncrementalIndex, IndexStats as IncrementalIndexStats, MergeResult as IncrementalMergeResult, Operation as IncrementalOperation, RecoveryResult, SearchResult as IncrementalSearchResult, SearchSource, VectorEntry, WalRecovery, }; -pub use federation::{ - FederatedQuery, FederatedQueryResult, FederatedResult, Federation, FederationBuilder, - FederationConfig, FederationMember, FederationStatus, MemberCapabilities, MemberHealth, - MemberStatus, MemberType, MergeStrategy, +pub use k8s_operator::{ + AutoscalingConfig, BackupConfig as K8sBackupConfig, ClusterCondition, ClusterPhase, + ConditionType, DistanceMetric as K8sDistanceMetric, HAConfig, HelmChartGenerator, + IndexConfig as K8sIndexConfig, ManifestGenerator, MonitoringConfig as K8sMonitoringConfig, + OperatorConfig, OperatorMetricsSnapshot, ReconcileAction, ReconcileResult, + ResourceRequirements, ResourceSpec, VecStoreClusterSpec, VecStoreClusterStatus, + VecStoreOperator, +}; +pub use query_router::{ + CircuitState as RouterCircuitState, FilterComplexity, HealthStatus as RouterHealthStatus, + LoadBalanceStrategy, NodeCapabilities, NodeStats, QueryCharacteristics, QueryPriority, + QueryRouter, ReplicaNode, RequestContext, RetryConfig, RetryHandler, RouterConfig, RouterStats, + RoutingDecision, RoutingReason, }; // Phase 7 exports -pub use neural_ranker::{ - ActivationFn, CombinationMethod, CrossEncoder, CrossEncoderConfig, EnsembleConfig, - EnsembleRanker, LateInteractionConfig, LateInteractionModel, LearnedSparseConfig, - LearnedSparseModel, ModelType as NeuralModelType, NeuralRanker, NeuralRankerBuilder, - NeuralRankerConfig, RankerStatsSummary, Ranker as NeuralRankerTrait, RerankCandidate, - RerankResponse, RerankResult, SimilarityFn, SparseRepresentation, +pub use anomaly_detection::{ + AnomalyConfig, AnomalyDetector as VectorAnomalyDetector, + AnomalyDetectorBuilder as VectorAnomalyDetectorBuilder, AnomalyRecord, AnomalyType, + BatchDetectionResult, DetectionMethod, DetectionResult, DetectionStats, DetectionSummary, + DriftConfig, DriftDetector, DriftEvent, DriftMethod, DriftSeverity, }; pub use auto_tuning::{ AcquisitionFunction, AutoTuneConfig, AutoTuner as HyperparameterTuner, AutoTunerBuilder as HyperparameterTunerBuilder, Constraint as TuneConstraint, ConstraintType as TuneConstraintType, IndexTuningPreset, OptimizeDirection, - Parameter as TuneParameter, ParameterSpace, ParameterType as TuneParameterType, - ParameterValue, SearchStrategy, Trial, TrialState, TuneMetric, TuneResult, - TunerStatsSummary, WorkloadAnalyzer as TuneWorkloadAnalyzer, WorkloadConfig as TuneWorkloadConfig, - WorkloadRecommendation as TuneWorkloadRecommendation, - hnsw_parameter_space, ivf_parameter_space, + Parameter as TuneParameter, ParameterSpace, ParameterType as TuneParameterType, ParameterValue, + SearchStrategy, Trial, TrialState, TuneMetric, TuneResult, TunerStatsSummary, + WorkloadAnalyzer as TuneWorkloadAnalyzer, WorkloadConfig as TuneWorkloadConfig, + WorkloadRecommendation as TuneWorkloadRecommendation, hnsw_parameter_space, + ivf_parameter_space, }; -pub use anomaly_detection::{ - AnomalyConfig, AnomalyDetector as VectorAnomalyDetector, - AnomalyDetectorBuilder as VectorAnomalyDetectorBuilder, AnomalyRecord, AnomalyType, - BatchDetectionResult, DetectionMethod, DetectionResult, DetectionStats, DetectionSummary, - DriftConfig, DriftDetector, DriftEvent, DriftMethod, DriftSeverity, +pub use cost_optimizer::{ + AutoScalerConfig, BudgetStatus, CostAnalysis, CostAwareAutoScaler, + CostBreakdown as CostOptimizerBreakdown, CostForecast, CostOptimizer, CostOptimizerBuilder, + CostOptimizerConfig, CostPeriod, CostRecommendation, CostRecord, CostTrend, EffortLevel, + OptimizationAction, OptimizationCategory, OptimizerStatsSummary, RecommendationStatus, + ResourcePricing, ResourceUsage as CostResourceUsage, RiskLevel, ScaleAction, ScalingDecision, + TrendDirection as CostTrendDirection, +}; +pub use neural_ranker::{ + ActivationFn, CombinationMethod, CrossEncoder, CrossEncoderConfig, EnsembleConfig, + EnsembleRanker, LateInteractionConfig, LateInteractionModel, LearnedSparseConfig, + LearnedSparseModel, ModelType as NeuralModelType, NeuralRanker, NeuralRankerBuilder, + NeuralRankerConfig, Ranker as NeuralRankerTrait, RankerStatsSummary, RerankCandidate, + RerankResponse, RerankResult, SimilarityFn, SparseRepresentation, }; pub use query_analytics::{ AggregateStats as AnalyticsAggregateStats, Alert as AnalyticsAlert, @@ -827,18 +834,10 @@ pub use query_analytics::{ RecommendationCategory, ReportConfig, ReportFormat, TimeRange, TimeSeriesPoint, TrendData, TrendDirection as AnalyticsTrendDirection, }; -pub use cost_optimizer::{ - AutoScalerConfig, BudgetStatus, CostAnalysis, CostAwareAutoScaler, - CostBreakdown as CostOptimizerBreakdown, CostForecast, CostOptimizer, CostOptimizerBuilder, - CostOptimizerConfig, CostPeriod, CostRecommendation, CostRecord, CostTrend, EffortLevel, - OptimizationAction, OptimizationCategory, OptimizerStatsSummary, RecommendationStatus, - ResourcePricing, ResourceUsage as CostResourceUsage, RiskLevel, ScaleAction, ScalingDecision, - TrendDirection as CostTrendDirection, -}; /// Initialize tracing subscriber for logging pub fn init_tracing() { - use tracing_subscriber::{fmt, EnvFilter}; + use tracing_subscriber::{EnvFilter, fmt}; let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); diff --git a/src/lineage.rs b/src/lineage.rs index 74cb372..85cb551 100644 --- a/src/lineage.rs +++ b/src/lineage.rs @@ -35,10 +35,10 @@ //! println!("Transformations: {:?}", lineage.transformations); //! ``` -use std::collections::HashMap; -use std::sync::{Arc, RwLock}; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; use crate::error::VecStoreError; @@ -207,26 +207,18 @@ pub enum TransformationType { to_dims: usize, }, /// Vector was quantized - Quantized { - bits: u8, - }, + Quantized { bits: u8 }, /// Product quantization was applied ProductQuantized { num_subvectors: usize, bits_per_subvector: u8, }, /// Vector was averaged with others - Averaged { - source_count: usize, - }, + Averaged { source_count: usize }, /// Vector was concatenated with others - Concatenated { - source_ids: Vec, - }, + Concatenated { source_ids: Vec }, /// Vector was updated/modified - Updated { - reason: String, - }, + Updated { reason: String }, /// Custom transformation Custom { name: String, @@ -433,7 +425,11 @@ impl LineageTracker { } /// Register a new vector with its origin - pub fn register(&self, vector_id: impl Into, origin: VectorOrigin) -> Result<(), VecStoreError> { + pub fn register( + &self, + vector_id: impl Into, + origin: VectorOrigin, + ) -> Result<(), VecStoreError> { let id = vector_id.into(); let lineage = VectorLineage { vector_id: id.clone(), @@ -447,9 +443,10 @@ impl LineageTracker { access_log: Vec::new(), }; - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; store.insert(id, lineage); Ok(()) @@ -464,9 +461,10 @@ impl LineageTracker { ) -> Result<(), VecStoreError> { let id = vector_id.into(); - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; // Update parent's child_ids for parent_id in &parent_ids { @@ -509,9 +507,10 @@ impl LineageTracker { before_hash: Option, after_hash: Option, ) -> Result<(), VecStoreError> { - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; let lineage = store.get_mut(vector_id).ok_or_else(|| { VecStoreError::NotFound(format!("Vector {} not found in lineage", vector_id)) @@ -550,9 +549,10 @@ impl LineageTracker { return Ok(()); } - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; if let Some(lineage) = store.get_mut(vector_id) { lineage.access_log.push(AccessRecord { @@ -568,9 +568,10 @@ impl LineageTracker { /// Get lineage for a vector pub fn get_lineage(&self, vector_id: &str) -> Result { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; store.get(vector_id).cloned().ok_or_else(|| { VecStoreError::NotFound(format!("Vector {} not found in lineage", vector_id)) @@ -579,9 +580,10 @@ impl LineageTracker { /// Get full lineage tree (including ancestors) pub fn get_lineage_tree(&self, vector_id: &str) -> Result, VecStoreError> { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; let mut result = Vec::new(); let mut to_visit = vec![vector_id.to_string()]; @@ -612,9 +614,10 @@ impl LineageTracker { vector_id: &str, tags: Vec, ) -> Result<(), VecStoreError> { - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; let lineage = store.get_mut(vector_id).ok_or_else(|| { VecStoreError::NotFound(format!("Vector {} not found in lineage", vector_id)) @@ -634,9 +637,10 @@ impl LineageTracker { /// Check retention and get vectors that need action pub fn check_retention(&self) -> Result, VecStoreError> { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; let now = Utc::now(); let mut records = Vec::new(); @@ -645,7 +649,10 @@ impl LineageTracker { for (policy_name, policy) in &self.retention_policies { // Check if policy applies to this vector let applies = policy.applies_to_tags.is_empty() - || lineage.compliance_tags.iter().any(|t| policy.applies_to_tags.contains(t)); + || lineage + .compliance_tags + .iter() + .any(|t| policy.applies_to_tags.contains(t)); if !applies { continue; @@ -681,9 +688,10 @@ impl LineageTracker { period_start: DateTime, period_end: DateTime, ) -> Result { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; let mut by_source_type: HashMap = HashMap::new(); let mut by_model: HashMap = HashMap::new(); @@ -734,7 +742,8 @@ impl LineageTracker { let mut issues = Vec::new(); // Check for expired vectors - let expired: Vec<_> = retention_status.iter() + let expired: Vec<_> = retention_status + .iter() .filter(|r| r.is_expired) .map(|r| r.vector_id.clone()) .collect(); @@ -743,14 +752,18 @@ impl LineageTracker { issues.push(ComplianceIssue { severity: IssueSeverity::Warning, category: "Retention".to_string(), - description: format!("{} vectors have exceeded their retention period", expired.len()), + description: format!( + "{} vectors have exceeded their retention period", + expired.len() + ), affected_vectors: expired, recommendation: "Review and apply retention policies".to_string(), }); } // Check for vectors without model info - let missing_model: Vec<_> = store.iter() + let missing_model: Vec<_> = store + .iter() .filter(|(_, l)| l.origin.embedding_model.is_none()) .map(|(id, _)| id.clone()) .collect(); @@ -759,7 +772,10 @@ impl LineageTracker { issues.push(ComplianceIssue { severity: IssueSeverity::Info, category: "Provenance".to_string(), - description: format!("{} vectors are missing model attribution", missing_model.len()), + description: format!( + "{} vectors are missing model attribution", + missing_model.len() + ), affected_vectors: missing_model, recommendation: "Update vectors with embedding model information".to_string(), }); @@ -786,11 +802,13 @@ impl LineageTracker { /// Find vectors by source pub fn find_by_source(&self, source_id: &str) -> Result, VecStoreError> { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; - let results: Vec = store.iter() + let results: Vec = store + .iter() .filter(|(_, lineage)| lineage.origin.source_id == source_id) .map(|(id, _)| id.clone()) .collect(); @@ -800,13 +818,18 @@ impl LineageTracker { /// Find vectors by model pub fn find_by_model(&self, model_name: &str) -> Result, VecStoreError> { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; - let results: Vec = store.iter() + let results: Vec = store + .iter() .filter(|(_, lineage)| { - lineage.origin.embedding_model.as_ref() + lineage + .origin + .embedding_model + .as_ref() .map(|m| m.name == model_name) .unwrap_or(false) }) @@ -818,11 +841,13 @@ impl LineageTracker { /// Get vectors with specific compliance tag pub fn find_by_tag(&self, tag: &str) -> Result, VecStoreError> { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; - let results: Vec = store.iter() + let results: Vec = store + .iter() .filter(|(_, lineage)| lineage.compliance_tags.contains(&tag.to_string())) .map(|(id, _)| id.clone()) .collect(); @@ -833,22 +858,21 @@ impl LineageTracker { /// Export lineage data for a vector pub fn export_lineage(&self, vector_id: &str) -> Result { let lineage = self.get_lineage(vector_id)?; - serde_json::to_string_pretty(&lineage).map_err(|e| { - VecStoreError::Serialization(e.to_string()) - }) + serde_json::to_string_pretty(&lineage) + .map_err(|e| VecStoreError::Serialization(e.to_string())) } /// Import lineage data pub fn import_lineage(&self, json: &str) -> Result { - let lineage: VectorLineage = serde_json::from_str(json).map_err(|e| { - VecStoreError::Serialization(e.to_string()) - })?; + let lineage: VectorLineage = + serde_json::from_str(json).map_err(|e| VecStoreError::Serialization(e.to_string()))?; let id = lineage.vector_id.clone(); - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; store.insert(id.clone(), lineage); Ok(id) @@ -856,9 +880,10 @@ impl LineageTracker { /// Delete lineage for a vector pub fn delete(&self, vector_id: &str) -> Result<(), VecStoreError> { - let mut store = self.lineage.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire write lock".into()) - })?; + let mut store = self + .lineage + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire write lock".into()))?; store.remove(vector_id); Ok(()) @@ -866,9 +891,10 @@ impl LineageTracker { /// Get total count of tracked vectors pub fn count(&self) -> Result { - let store = self.lineage.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire read lock".into()) - })?; + let store = self + .lineage + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire read lock".into()))?; Ok(store.len()) } @@ -903,8 +929,7 @@ mod tests { fn test_register_and_get_lineage() { let tracker = LineageTracker::new(); - let origin = VectorOrigin::new("doc1") - .from_model("test-model"); + let origin = VectorOrigin::new("doc1").from_model("test-model"); tracker.register("vec1", origin).unwrap(); @@ -921,8 +946,12 @@ mod tests { let origin = VectorOrigin::new("doc1"); tracker.register("vec1", origin).unwrap(); - tracker.add_transformation("vec1", TransformationType::Normalized).unwrap(); - tracker.add_transformation("vec1", TransformationType::Quantized { bits: 8 }).unwrap(); + tracker + .add_transformation("vec1", TransformationType::Normalized) + .unwrap(); + tracker + .add_transformation("vec1", TransformationType::Quantized { bits: 8 }) + .unwrap(); let lineage = tracker.get_lineage("vec1").unwrap(); assert_eq!(lineage.transformations.len(), 2); @@ -934,16 +963,22 @@ mod tests { let tracker = LineageTracker::new(); // Register parent vectors - tracker.register("parent1", VectorOrigin::new("doc1")).unwrap(); - tracker.register("parent2", VectorOrigin::new("doc2")).unwrap(); + tracker + .register("parent1", VectorOrigin::new("doc1")) + .unwrap(); + tracker + .register("parent2", VectorOrigin::new("doc2")) + .unwrap(); // Register derived vector let origin = VectorOrigin::new("derived_source"); - tracker.register_derived( - "child1", - origin, - vec!["parent1".to_string(), "parent2".to_string()], - ).unwrap(); + tracker + .register_derived( + "child1", + origin, + vec!["parent1".to_string(), "parent2".to_string()], + ) + .unwrap(); // Check child let child_lineage = tracker.get_lineage("child1").unwrap(); @@ -961,12 +996,14 @@ mod tests { let origin = VectorOrigin::new("doc1"); tracker.register("vec1", origin).unwrap(); - tracker.log_access( - "vec1", - AccessType::SearchResult, - Some("user123".to_string()), - Some("query456".to_string()), - ).unwrap(); + tracker + .log_access( + "vec1", + AccessType::SearchResult, + Some("user123".to_string()), + Some("query456".to_string()), + ) + .unwrap(); let lineage = tracker.get_lineage("vec1").unwrap(); assert_eq!(lineage.access_log.len(), 1); @@ -980,10 +1017,9 @@ mod tests { let origin = VectorOrigin::new("doc1"); tracker.register("vec1", origin).unwrap(); - tracker.add_compliance_tags("vec1", vec![ - "GDPR".to_string(), - "PII".to_string(), - ]).unwrap(); + tracker + .add_compliance_tags("vec1", vec!["GDPR".to_string(), "PII".to_string()]) + .unwrap(); let lineage = tracker.get_lineage("vec1").unwrap(); assert!(lineage.compliance_tags.contains(&"GDPR".to_string())); @@ -1006,9 +1042,15 @@ mod tests { fn test_find_by_model() { let tracker = LineageTracker::new(); - tracker.register("vec1", VectorOrigin::new("doc1").from_model("model-a")).unwrap(); - tracker.register("vec2", VectorOrigin::new("doc2").from_model("model-a")).unwrap(); - tracker.register("vec3", VectorOrigin::new("doc3").from_model("model-b")).unwrap(); + tracker + .register("vec1", VectorOrigin::new("doc1").from_model("model-a")) + .unwrap(); + tracker + .register("vec2", VectorOrigin::new("doc2").from_model("model-a")) + .unwrap(); + tracker + .register("vec3", VectorOrigin::new("doc3").from_model("model-b")) + .unwrap(); let results = tracker.find_by_model("model-a").unwrap(); assert_eq!(results.len(), 2); @@ -1019,9 +1061,23 @@ mod tests { let tracker = LineageTracker::new(); // Build a tree: grandparent -> parent -> child - tracker.register("gp", VectorOrigin::new("gp_source")).unwrap(); - tracker.register_derived("parent", VectorOrigin::new("parent_source"), vec!["gp".to_string()]).unwrap(); - tracker.register_derived("child", VectorOrigin::new("child_source"), vec!["parent".to_string()]).unwrap(); + tracker + .register("gp", VectorOrigin::new("gp_source")) + .unwrap(); + tracker + .register_derived( + "parent", + VectorOrigin::new("parent_source"), + vec!["gp".to_string()], + ) + .unwrap(); + tracker + .register_derived( + "child", + VectorOrigin::new("child_source"), + vec!["parent".to_string()], + ) + .unwrap(); let tree = tracker.get_lineage_tree("child").unwrap(); assert_eq!(tree.len(), 3); @@ -1035,7 +1091,9 @@ mod tests { .from_model("test-model") .with_type(SourceType::Document); tracker.register("vec1", origin).unwrap(); - tracker.add_transformation("vec1", TransformationType::Normalized).unwrap(); + tracker + .add_transformation("vec1", TransformationType::Normalized) + .unwrap(); // Export let json = tracker.export_lineage("vec1").unwrap(); @@ -1054,12 +1112,25 @@ mod tests { fn test_compliance_report() { let tracker = LineageTracker::new(); - tracker.register("vec1", VectorOrigin::new("doc1").from_model("model-a")).unwrap(); - tracker.register("vec2", VectorOrigin::new("doc2").from_model("model-b")).unwrap(); + tracker + .register("vec1", VectorOrigin::new("doc1").from_model("model-a")) + .unwrap(); + tracker + .register("vec2", VectorOrigin::new("doc2").from_model("model-b")) + .unwrap(); tracker.register("vec3", VectorOrigin::new("doc3")).unwrap(); // No model - tracker.log_access("vec1", AccessType::Read, Some("user1".to_string()), None).unwrap(); - tracker.log_access("vec1", AccessType::SearchResult, Some("user2".to_string()), None).unwrap(); + tracker + .log_access("vec1", AccessType::Read, Some("user1".to_string()), None) + .unwrap(); + tracker + .log_access( + "vec1", + AccessType::SearchResult, + Some("user2".to_string()), + None, + ) + .unwrap(); let start = Utc::now() - chrono::Duration::hours(1); let end = Utc::now() + chrono::Duration::hours(1); diff --git a/src/llm_cache.rs b/src/llm_cache.rs index a52dd09..6c091cd 100644 --- a/src/llm_cache.rs +++ b/src/llm_cache.rs @@ -37,9 +37,9 @@ //! assert_eq!(response1, response2); // Same response from cache //! ``` +use serde::{Deserialize, Serialize}; use std::sync::{Arc, RwLock}; use std::time::{Duration, Instant}; -use serde::{Deserialize, Serialize}; use crate::error::{Result, VecStoreError}; @@ -72,11 +72,17 @@ pub const DEFAULT_SIMILARITY_THRESHOLD: f32 = 0.95; pub const DEFAULT_MAX_SIZE: usize = 10000; #[inline] -const fn default_similarity_threshold() -> f32 { DEFAULT_SIMILARITY_THRESHOLD } +const fn default_similarity_threshold() -> f32 { + DEFAULT_SIMILARITY_THRESHOLD +} #[inline] -const fn default_max_size() -> usize { DEFAULT_MAX_SIZE } +const fn default_max_size() -> usize { + DEFAULT_MAX_SIZE +} #[inline] -const fn default_true() -> bool { true } +const fn default_true() -> bool { + true +} impl Default for CacheConfig { fn default() -> Self { @@ -255,7 +261,9 @@ impl SemanticCache { /// Find a similar cached query fn find_similar(&self, _query: &str, embedding: &[f32]) -> Result> { - let entries = self.entries.read() + let entries = self + .entries + .read() .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?; let ttl = Duration::from_secs(self.config.ttl_seconds); @@ -271,9 +279,10 @@ impl SemanticCache { let similarity = Self::cosine_similarity(embedding, &entry.embedding); if similarity >= self.config.similarity_threshold - && (best_match.is_none() || similarity > best_match.unwrap().1) { - best_match = Some((i, similarity)); - } + && (best_match.is_none() || similarity > best_match.unwrap().1) + { + best_match = Some((i, similarity)); + } } if let Some((idx, similarity)) = best_match { @@ -296,13 +305,16 @@ impl SemanticCache { response: &str, metadata: Option, ) -> Result<()> { - let mut entries = self.entries.write() + let mut entries = self + .entries + .write() .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?; // Evict if at capacity while entries.len() >= self.config.max_size { // Remove least recently accessed - if let Some(idx) = entries.iter() + if let Some(idx) = entries + .iter() .enumerate() .min_by_key(|(_, e)| e.last_accessed) .map(|(i, _)| i) @@ -330,7 +342,9 @@ impl SemanticCache { return Ok(()); } - let mut analytics = self.analytics.write() + let mut analytics = self + .analytics + .write() .map_err(|_| VecStoreError::LockError("analytics lock poisoned".into()))?; analytics.total_queries += 1; analytics.cache_hits += 1; @@ -345,8 +359,11 @@ impl SemanticCache { analytics.avg_hit_similarity = (analytics.avg_hit_similarity * (n - 1.0) + hit.similarity) / n; - analytics.cache_size = self.entries.read() - .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?.len(); + analytics.cache_size = self + .entries + .read() + .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))? + .len(); Ok(()) } @@ -356,31 +373,43 @@ impl SemanticCache { return Ok(()); } - let mut analytics = self.analytics.write() + let mut analytics = self + .analytics + .write() .map_err(|_| VecStoreError::LockError("analytics lock poisoned".into()))?; analytics.total_queries += 1; analytics.cache_misses += 1; - analytics.cache_size = self.entries.read() - .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?.len(); + analytics.cache_size = self + .entries + .read() + .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))? + .len(); Ok(()) } /// Get cache analytics pub fn analytics(&self) -> Result { - Ok(self.analytics.read() - .map_err(|_| VecStoreError::LockError("analytics lock poisoned".into()))?.clone()) + Ok(self + .analytics + .read() + .map_err(|_| VecStoreError::LockError("analytics lock poisoned".into()))? + .clone()) } /// Clear the cache pub fn clear(&self) -> Result<()> { - self.entries.write() - .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?.clear(); + self.entries + .write() + .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))? + .clear(); Ok(()) } /// Invalidate expired entries pub fn cleanup(&self) -> Result { - let mut entries = self.entries.write() + let mut entries = self + .entries + .write() .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?; let ttl = Duration::from_secs(self.config.ttl_seconds); @@ -404,14 +433,20 @@ impl SemanticCache { /// Get cache size pub fn len(&self) -> Result { - Ok(self.entries.read() - .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?.len()) + Ok(self + .entries + .read() + .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))? + .len()) } /// Check if empty pub fn is_empty(&self) -> Result { - Ok(self.entries.read() - .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))?.is_empty()) + Ok(self + .entries + .read() + .map_err(|_| VecStoreError::LockError("entries lock poisoned".into()))? + .is_empty()) } } @@ -514,20 +549,20 @@ mod tests { // First query - miss let embedding1 = vec![0.1f32; 64]; - let result1 = cache.get_or_generate( - "What is AI?", - &embedding1, - || Ok("AI is artificial intelligence.".to_string()), - ).unwrap(); + let result1 = cache + .get_or_generate("What is AI?", &embedding1, || { + Ok("AI is artificial intelligence.".to_string()) + }) + .unwrap(); assert!(!result1.cache_hit); // Same query - hit - let result2 = cache.get_or_generate( - "What is AI?", - &embedding1, - || Ok("Should not be called".to_string()), - ).unwrap(); + let result2 = cache + .get_or_generate("What is AI?", &embedding1, || { + Ok("Should not be called".to_string()) + }) + .unwrap(); assert!(result2.cache_hit); assert_eq!(result1.response, result2.response); @@ -541,10 +576,14 @@ mod tests { let embedding = vec![0.1f32; 64]; // Miss - cache.get_or_generate("Query 1", &embedding, || Ok("Response 1".to_string())).unwrap(); + cache + .get_or_generate("Query 1", &embedding, || Ok("Response 1".to_string())) + .unwrap(); // Hit - cache.get_or_generate("Query 1", &embedding, || Ok("Response 2".to_string())).unwrap(); + cache + .get_or_generate("Query 1", &embedding, || Ok("Response 2".to_string())) + .unwrap(); let analytics = cache.analytics().unwrap(); assert_eq!(analytics.total_queries, 2); diff --git a/src/lsh.rs b/src/lsh.rs index dee69a9..e9f83c2 100644 --- a/src/lsh.rs +++ b/src/lsh.rs @@ -49,7 +49,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; use serde::{Deserialize, Serialize}; @@ -90,7 +90,9 @@ struct HashFunction { impl HashFunction { /// Create a new random hash function fn new(dimension: usize, rng: &mut StdRng) -> Self { - let projection: Vec = (0..dimension).map(|_| rng.random_range(-1.0..1.0)).collect(); + let projection: Vec = (0..dimension) + .map(|_| rng.random_range(-1.0..1.0)) + .collect(); Self { projection } } @@ -105,11 +107,7 @@ impl HashFunction { .map(|(a, b)| a * b) .sum(); - if dot_product >= 0.0 { - 1 - } else { - 0 - } + if dot_product >= 0.0 { 1 } else { 0 } } } @@ -163,9 +161,7 @@ impl HashTable { fn get_candidates(&self, query: &[f32]) -> Vec<(String, Vec)> { let hash_code = self.compute_hash(query); - self.buckets - .get(&hash_code).cloned() - .unwrap_or_default() + self.buckets.get(&hash_code).cloned().unwrap_or_default() } /// Remove a vector by ID diff --git a/src/managed_embeddings.rs b/src/managed_embeddings.rs index 0696ab1..35dda6f 100644 --- a/src/managed_embeddings.rs +++ b/src/managed_embeddings.rs @@ -401,16 +401,20 @@ impl EmbeddingCache { self.evict_oldest(); } - self.entries.insert(key, CacheEntry { - embedding, - model: model.to_string(), - cached_at: Instant::now(), - hits: 0, - }); + self.entries.insert( + key, + CacheEntry { + embedding, + model: model.to_string(), + cached_at: Instant::now(), + hits: 0, + }, + ); } fn evict_oldest(&mut self) { - let oldest = self.entries + let oldest = self + .entries .iter() .min_by_key(|(_, e)| e.cached_at) .map(|(k, _)| k.clone()); @@ -508,29 +512,43 @@ impl UsageTracker { } } - fn record(&mut self, provider: &str, model: &str, tokens: usize, embeddings: usize, cost: f64, latency_ms: u64) { + fn record( + &mut self, + provider: &str, + model: &str, + tokens: usize, + embeddings: usize, + cost: f64, + latency_ms: u64, + ) { self.total_tokens += tokens as u64; self.total_requests += 1; self.total_embeddings += embeddings as u64; self.total_latency_ms += latency_ms; // By provider - let provider_usage = self.by_provider.entry(provider.to_string()).or_insert(ProviderUsage { - tokens: 0, - requests: 0, - errors: 0, - cost: 0.0, - }); + let provider_usage = + self.by_provider + .entry(provider.to_string()) + .or_insert(ProviderUsage { + tokens: 0, + requests: 0, + errors: 0, + cost: 0.0, + }); provider_usage.tokens += tokens as u64; provider_usage.requests += 1; provider_usage.cost += cost; // By model - let model_usage = self.by_model.entry(model.to_string()).or_insert(ModelUsage { - tokens: 0, - embeddings: 0, - cost: 0.0, - }); + let model_usage = self + .by_model + .entry(model.to_string()) + .or_insert(ModelUsage { + tokens: 0, + embeddings: 0, + cost: 0.0, + }); model_usage.tokens += tokens as u64; model_usage.embeddings += embeddings as u64; model_usage.cost += cost; @@ -545,12 +563,15 @@ impl UsageTracker { } fn record_error(&mut self, provider: &str) { - let provider_usage = self.by_provider.entry(provider.to_string()).or_insert(ProviderUsage { - tokens: 0, - requests: 0, - errors: 0, - cost: 0.0, - }); + let provider_usage = + self.by_provider + .entry(provider.to_string()) + .or_insert(ProviderUsage { + tokens: 0, + requests: 0, + errors: 0, + cost: 0.0, + }); provider_usage.errors += 1; } @@ -608,7 +629,10 @@ impl EmbeddingService { /// Embed texts pub fn embed(&self, request: EmbeddingRequest) -> Result { let start = Instant::now(); - let model = request.model_override.as_ref().unwrap_or(&self.config.primary_model); + let model = request + .model_override + .as_ref() + .unwrap_or(&self.config.primary_model); let mut embeddings = Vec::with_capacity(request.texts.len()); let mut cache_hits = 0; @@ -616,7 +640,9 @@ impl EmbeddingService { // Check cache for each text if self.config.enable_cache && !request.skip_cache { - let mut cache = self.cache.write() + let mut cache = self + .cache + .write() .map_err(|_| VecStoreError::LockError("cache lock poisoned".into()))?; for (i, text) in request.texts.iter().enumerate() { @@ -625,18 +651,25 @@ impl EmbeddingService { if let Some(cached) = cache.get(&cache_key) { embeddings.push((i, cached)); cache_hits += 1; - self.usage.write() + self.usage + .write() .map_err(|_| VecStoreError::LockError("usage lock poisoned".into()))? .record_cache_hit(); } else { texts_to_embed.push((i, text.clone())); - self.usage.write() + self.usage + .write() .map_err(|_| VecStoreError::LockError("usage lock poisoned".into()))? .record_cache_miss(); } } } else { - texts_to_embed = request.texts.iter().enumerate().map(|(i, t)| (i, t.clone())).collect(); + texts_to_embed = request + .texts + .iter() + .enumerate() + .map(|(i, t)| (i, t.clone())) + .collect(); } // Embed remaining texts @@ -644,7 +677,9 @@ impl EmbeddingService { if !texts_to_embed.is_empty() { // Rate limiting { - let mut limiter = self.rate_limiter.write() + let mut limiter = self + .rate_limiter + .write() .map_err(|_| VecStoreError::LockError("rate_limiter lock poisoned".into()))?; if !limiter.acquire(1.0) { // Would need to wait in real async implementation @@ -653,12 +688,15 @@ impl EmbeddingService { // Call embedding provider let texts: Vec<&str> = texts_to_embed.iter().map(|(_, t)| t.as_str()).collect(); - let (new_embeddings, tokens) = self.call_provider(&texts, model, request.task_type.as_ref())?; + let (new_embeddings, tokens) = + self.call_provider(&texts, model, request.task_type.as_ref())?; total_tokens = tokens; // Store in cache if self.config.enable_cache && !request.skip_cache { - let mut cache = self.cache.write() + let mut cache = self + .cache + .write() .map_err(|_| VecStoreError::LockError("cache lock poisoned".into()))?; for ((i, text), emb) in texts_to_embed.iter().zip(&new_embeddings) { let cache_key = self.cache_key(text, &model.model, request.task_type.as_ref()); @@ -678,13 +716,10 @@ impl EmbeddingService { // Apply dimension override if specified let final_embeddings = if let Some(dim) = request.dimension_override { - final_embeddings.into_iter().map(|e| { - if e.len() > dim { - e[..dim].to_vec() - } else { - e - } - }).collect() + final_embeddings + .into_iter() + .map(|e| if e.len() > dim { e[..dim].to_vec() } else { e }) + .collect() } else { final_embeddings }; @@ -695,7 +730,8 @@ impl EmbeddingService { // Record usage if self.config.track_usage && total_tokens > 0 { let provider_name = format!("{:?}", model.provider); - self.usage.write() + self.usage + .write() .map_err(|_| VecStoreError::LockError("usage lock poisoned".into()))? .record( &provider_name, @@ -721,39 +757,45 @@ impl EmbeddingService { pub fn embed_one(&self, text: &str) -> Result> { let request = EmbeddingRequest::new(vec![text.to_string()]); let response = self.embed(request)?; - response.embeddings.into_iter().next().ok_or_else(|| { - VecStoreError::EmbeddingError("No embedding returned".to_string()) - }) + response + .embeddings + .into_iter() + .next() + .ok_or_else(|| VecStoreError::EmbeddingError("No embedding returned".to_string())) } /// Embed for query (uses query task type) pub fn embed_query(&self, text: &str) -> Result> { - let request = EmbeddingRequest::new(vec![text.to_string()]) - .with_task_type(TaskType::Query); + let request = EmbeddingRequest::new(vec![text.to_string()]).with_task_type(TaskType::Query); let response = self.embed(request)?; - response.embeddings.into_iter().next().ok_or_else(|| { - VecStoreError::EmbeddingError("No embedding returned".to_string()) - }) + response + .embeddings + .into_iter() + .next() + .ok_or_else(|| VecStoreError::EmbeddingError("No embedding returned".to_string())) } /// Embed for documents (uses document task type) pub fn embed_documents(&self, texts: Vec) -> Result>> { - let request = EmbeddingRequest::new(texts) - .with_task_type(TaskType::Document); + let request = EmbeddingRequest::new(texts).with_task_type(TaskType::Document); let response = self.embed(request)?; Ok(response.embeddings) } /// Get usage statistics pub fn usage(&self) -> Result { - let usage = self.usage.read() + let usage = self + .usage + .read() .map_err(|_| VecStoreError::LockError("usage lock poisoned".into()))?; Ok(usage.stats()) } /// Reset usage tracking pub fn reset_usage(&self) -> Result<()> { - let mut usage = self.usage.write() + let mut usage = self + .usage + .write() .map_err(|_| VecStoreError::LockError("usage lock poisoned".into()))?; *usage = UsageTracker::new(); Ok(()) @@ -771,7 +813,9 @@ impl EmbeddingService { /// Clear embedding cache pub fn clear_cache(&self) -> Result<()> { - let mut cache = self.cache.write() + let mut cache = self + .cache + .write() .map_err(|_| VecStoreError::LockError("cache lock poisoned".into()))?; *cache = EmbeddingCache::new(self.config.max_cache_entries, self.config.cache_ttl_seconds); Ok(()) @@ -779,7 +823,9 @@ impl EmbeddingService { /// Get cache statistics pub fn cache_stats(&self) -> Result<(usize, u64)> { - let cache = self.cache.read() + let cache = self + .cache + .read() .map_err(|_| VecStoreError::LockError("cache lock poisoned".into()))?; let stats = cache.stats(); Ok((stats.entries, stats.total_hits)) @@ -800,7 +846,12 @@ impl EmbeddingService { format!("{:x}", hasher.finish()) } - fn call_provider(&self, texts: &[&str], model: &ModelConfig, task_type: Option<&TaskType>) -> Result<(Vec>, usize)> { + fn call_provider( + &self, + texts: &[&str], + model: &ModelConfig, + task_type: Option<&TaskType>, + ) -> Result<(Vec>, usize)> { // Simulated embedding - in production would call actual APIs match &model.provider { EmbeddingProvider::OpenAI => self.embed_openai(texts, model), @@ -819,19 +870,34 @@ impl EmbeddingService { Ok((embeddings, total_tokens)) } - fn embed_cohere(&self, texts: &[&str], model: &ModelConfig, _task_type: Option<&TaskType>) -> Result<(Vec>, usize)> { + fn embed_cohere( + &self, + texts: &[&str], + model: &ModelConfig, + _task_type: Option<&TaskType>, + ) -> Result<(Vec>, usize)> { let total_tokens: usize = texts.iter().map(|t| t.split_whitespace().count()).sum(); let embeddings = self.generate_embeddings(texts, model.dimension); Ok((embeddings, total_tokens)) } - fn embed_voyage(&self, texts: &[&str], model: &ModelConfig, _task_type: Option<&TaskType>) -> Result<(Vec>, usize)> { + fn embed_voyage( + &self, + texts: &[&str], + model: &ModelConfig, + _task_type: Option<&TaskType>, + ) -> Result<(Vec>, usize)> { let total_tokens: usize = texts.iter().map(|t| t.split_whitespace().count()).sum(); let embeddings = self.generate_embeddings(texts, model.dimension); Ok((embeddings, total_tokens)) } - fn embed_jina(&self, texts: &[&str], model: &ModelConfig, _task_type: Option<&TaskType>) -> Result<(Vec>, usize)> { + fn embed_jina( + &self, + texts: &[&str], + model: &ModelConfig, + _task_type: Option<&TaskType>, + ) -> Result<(Vec>, usize)> { let total_tokens: usize = texts.iter().map(|t| t.split_whitespace().count()).sum(); let embeddings = self.generate_embeddings(texts, model.dimension); Ok((embeddings, total_tokens)) @@ -847,8 +913,14 @@ impl EmbeddingService { /// /// Generates deterministic hash-based embeddings for testing and development. /// Production deployments should use supported providers (OpenAI, Cohere, Voyage, Jina, LocalONNX). - fn embed_simulated(&self, texts: &[&str], model: &ModelConfig) -> Result<(Vec>, usize)> { - tracing::warn!("Using simulated embeddings for unsupported provider - not suitable for production"); + fn embed_simulated( + &self, + texts: &[&str], + model: &ModelConfig, + ) -> Result<(Vec>, usize)> { + tracing::warn!( + "Using simulated embeddings for unsupported provider - not suitable for production" + ); let total_tokens: usize = texts.iter().map(|t| t.split_whitespace().count()).sum(); let embeddings = self.generate_embeddings(texts, model.dimension); Ok((embeddings, total_tokens)) @@ -856,26 +928,31 @@ impl EmbeddingService { fn generate_embeddings(&self, texts: &[&str], dimension: usize) -> Vec> { // Generate deterministic embeddings based on text hash for testing/fallback - texts.iter().map(|text| { - let mut embedding = vec![0.0f32; dimension]; - let bytes = text.as_bytes(); - - for (i, chunk) in bytes.chunks(4).enumerate() { - let idx = i % dimension; - let val: u32 = chunk.iter().fold(0u32, |acc, &b| acc.wrapping_add(b as u32)); - embedding[idx] = ((val as f32) / 255.0 - 0.5) * 2.0; - } + texts + .iter() + .map(|text| { + let mut embedding = vec![0.0f32; dimension]; + let bytes = text.as_bytes(); + + for (i, chunk) in bytes.chunks(4).enumerate() { + let idx = i % dimension; + let val: u32 = chunk + .iter() + .fold(0u32, |acc, &b| acc.wrapping_add(b as u32)); + embedding[idx] = ((val as f32) / 255.0 - 0.5) * 2.0; + } - // Normalize - let norm: f32 = embedding.iter().map(|x| x * x).sum::().sqrt(); - if norm > 0.0 { - for x in &mut embedding { - *x /= norm; + // Normalize + let norm: f32 = embedding.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for x in &mut embedding { + *x /= norm; + } } - } - embedding - }).collect() + embedding + }) + .collect() } } @@ -897,14 +974,23 @@ impl EmbeddingCollection { } /// Add document with automatic embedding - pub fn add(&self, id: &str, text: &str, metadata: HashMap) -> Result<()> { + pub fn add( + &self, + id: &str, + text: &str, + metadata: HashMap, + ) -> Result<()> { let embedding = self.service.embed_one(text)?; - let mut vectors = self.vectors.write() + let mut vectors = self + .vectors + .write() .map_err(|_| VecStoreError::LockError("vectors lock poisoned".into()))?; vectors.insert(id.to_string(), (embedding, metadata)); - let mut texts = self.texts.write() + let mut texts = self + .texts + .write() .map_err(|_| VecStoreError::LockError("texts lock poisoned".into()))?; texts.insert(id.to_string(), text.to_string()); @@ -912,13 +998,20 @@ impl EmbeddingCollection { } /// Add multiple documents - pub fn add_batch(&self, documents: Vec<(String, String, HashMap)>) -> Result { + pub fn add_batch( + &self, + documents: Vec<(String, String, HashMap)>, + ) -> Result { let texts: Vec = documents.iter().map(|(_, t, _)| t.clone()).collect(); let embeddings = self.service.embed_documents(texts.clone())?; - let mut vectors = self.vectors.write() + let mut vectors = self + .vectors + .write() .map_err(|_| VecStoreError::LockError("vectors lock poisoned".into()))?; - let mut text_store = self.texts.write() + let mut text_store = self + .texts + .write() .map_err(|_| VecStoreError::LockError("texts lock poisoned".into()))?; for ((id, text, metadata), embedding) in documents.into_iter().zip(embeddings) { @@ -933,9 +1026,13 @@ impl EmbeddingCollection { pub fn query(&self, text: &str, top_k: usize) -> Result> { let query_embedding = self.service.embed_query(text)?; - let vectors = self.vectors.read() + let vectors = self + .vectors + .read() .map_err(|_| VecStoreError::LockError("vectors lock poisoned".into()))?; - let texts = self.texts.read() + let texts = self + .texts + .read() .map_err(|_| VecStoreError::LockError("texts lock poisoned".into()))?; let mut results: Vec = vectors @@ -959,16 +1056,22 @@ impl EmbeddingCollection { /// Get vector by ID pub fn get(&self, id: &str) -> Result>> { - let vectors = self.vectors.read() + let vectors = self + .vectors + .read() .map_err(|_| VecStoreError::LockError("vectors lock poisoned".into()))?; Ok(vectors.get(id).map(|(v, _)| v.clone())) } /// Delete by ID pub fn delete(&self, id: &str) -> Result { - let mut vectors = self.vectors.write() + let mut vectors = self + .vectors + .write() .map_err(|_| VecStoreError::LockError("vectors lock poisoned".into()))?; - let mut texts = self.texts.write() + let mut texts = self + .texts + .write() .map_err(|_| VecStoreError::LockError("texts lock poisoned".into()))?; texts.remove(id); Ok(vectors.remove(id).is_some()) @@ -976,7 +1079,9 @@ impl EmbeddingCollection { /// Count documents pub fn count(&self) -> Result { - let vectors = self.vectors.read() + let vectors = self + .vectors + .read() .map_err(|_| VecStoreError::LockError("vectors lock poisoned".into()))?; Ok(vectors.len()) } @@ -1020,10 +1125,12 @@ mod tests { let config = EmbeddingServiceConfig::default(); let service = EmbeddingService::new(config); - let response = service.embed(EmbeddingRequest::new(vec![ - "Hello world".to_string(), - "Machine learning".to_string(), - ])).unwrap(); + let response = service + .embed(EmbeddingRequest::new(vec![ + "Hello world".to_string(), + "Machine learning".to_string(), + ])) + .unwrap(); assert_eq!(response.embeddings.len(), 2); assert_eq!(response.embeddings[0].len(), 1536); // OpenAI small dimension @@ -1035,10 +1142,14 @@ mod tests { let service = EmbeddingService::new(config); // First request - let response1 = service.embed(EmbeddingRequest::new(vec!["Hello".to_string()])).unwrap(); + let response1 = service + .embed(EmbeddingRequest::new(vec!["Hello".to_string()])) + .unwrap(); // Second request (should be cached) - let response2 = service.embed(EmbeddingRequest::new(vec!["Hello".to_string()])).unwrap(); + let response2 = service + .embed(EmbeddingRequest::new(vec!["Hello".to_string()])) + .unwrap(); assert_eq!(response2.cache_hits, 1); assert_eq!(response1.embeddings[0], response2.embeddings[0]); @@ -1049,7 +1160,9 @@ mod tests { let config = EmbeddingServiceConfig::default(); let service = EmbeddingService::new(config); - service.embed(EmbeddingRequest::new(vec!["Test".to_string()])).unwrap(); + service + .embed(EmbeddingRequest::new(vec!["Test".to_string()])) + .unwrap(); let usage = service.usage().unwrap(); assert!(usage.total_requests > 0); @@ -1061,8 +1174,12 @@ mod tests { let service = Arc::new(EmbeddingService::new(config)); let collection = EmbeddingCollection::new(service); - collection.add("doc1", "Hello world", HashMap::new()).unwrap(); - collection.add("doc2", "Machine learning tutorial", HashMap::new()).unwrap(); + collection + .add("doc1", "Hello world", HashMap::new()) + .unwrap(); + collection + .add("doc2", "Machine learning tutorial", HashMap::new()) + .unwrap(); let results = collection.query("Hello", 10).unwrap(); assert!(!results.is_empty()); @@ -1089,15 +1206,19 @@ mod tests { }; let service = EmbeddingService::new(config); - let doc_embedding = service.embed( - EmbeddingRequest::new(vec!["Document text".to_string()]) - .with_task_type(TaskType::Document) - ).unwrap(); - - let query_embedding = service.embed( - EmbeddingRequest::new(vec!["Query text".to_string()]) - .with_task_type(TaskType::Query) - ).unwrap(); + let doc_embedding = service + .embed( + EmbeddingRequest::new(vec!["Document text".to_string()]) + .with_task_type(TaskType::Document), + ) + .unwrap(); + + let query_embedding = service + .embed( + EmbeddingRequest::new(vec!["Query text".to_string()]) + .with_task_type(TaskType::Query), + ) + .unwrap(); assert_eq!(doc_embedding.embeddings[0].len(), 1024); assert_eq!(query_embedding.embeddings[0].len(), 1024); diff --git a/src/matryoshka.rs b/src/matryoshka.rs index e6d054d..359ba97 100644 --- a/src/matryoshka.rs +++ b/src/matryoshka.rs @@ -44,7 +44,7 @@ //! let results = store.search(&query, 10)?; //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -174,10 +174,11 @@ impl MatryoshkaEmbedding { // Try to get from full embedding if let Some(ref full) = self.full - && dim <= full.len() { - // Would need to compute truncated - return None for now - return None; - } + && dim <= full.len() + { + // Would need to compute truncated - return None for now + return None; + } // Return closest smaller dimension self.truncated @@ -314,9 +315,10 @@ impl MatryoshkaStore { // Re-score candidates at higher dimension for (id, score) in &mut candidates { if let Some(emb) = self.embeddings.get(id) - && let Some(v) = emb.at_dimension(dim) { - *score = self.cosine_similarity(&refined_query, v); - } + && let Some(v) = emb.at_dimension(dim) + { + *score = self.cosine_similarity(&refined_query, v); + } } // Re-sort after refinement @@ -755,7 +757,8 @@ impl AdaptiveSearch { fn estimate_optimal_dimension(&self, query: &[f32]) -> usize { // Simple heuristic: use query variance to estimate complexity let mean: f32 = query.iter().sum::() / query.len() as f32; - let variance: f32 = query.iter().map(|x| (x - mean).powi(2)).sum::() / query.len() as f32; + let variance: f32 = + query.iter().map(|x| (x - mean).powi(2)).sum::() / query.len() as f32; // Higher variance queries need more dimensions if variance > 0.5 { diff --git a/src/metadata_index.rs b/src/metadata_index.rs index 928898c..96c771e 100644 --- a/src/metadata_index.rs +++ b/src/metadata_index.rs @@ -92,7 +92,7 @@ impl IndexedValue { } else { IndexedValue::Null } - } + }, serde_json::Value::Bool(b) => IndexedValue::Bool(*b), _ => IndexedValue::Null, } @@ -132,10 +132,7 @@ impl BTreeIndex { /// Insert value pub fn insert(&mut self, value: IndexedValue, id: String) { - self.tree - .entry(value) - .or_default() - .insert(id); + self.tree.entry(value).or_default().insert(id); } /// Remove value @@ -157,14 +154,14 @@ impl BTreeIndex { if let Some(ids) = self.tree.get(value) { result.extend(ids.iter().cloned()); } - } + }, "!=" => { for (k, ids) in &self.tree { if k != value { result.extend(ids.iter().cloned()); } } - } + }, ">" => { for (_k, ids) in self.tree.range(( std::ops::Bound::Excluded(value.clone()), @@ -172,7 +169,7 @@ impl BTreeIndex { )) { result.extend(ids.iter().cloned()); } - } + }, ">=" => { for (_k, ids) in self.tree.range(( std::ops::Bound::Included(value.clone()), @@ -180,7 +177,7 @@ impl BTreeIndex { )) { result.extend(ids.iter().cloned()); } - } + }, "<" => { for (_k, ids) in self.tree.range(( std::ops::Bound::Unbounded, @@ -188,7 +185,7 @@ impl BTreeIndex { )) { result.extend(ids.iter().cloned()); } - } + }, "<=" => { for (_k, ids) in self.tree.range(( std::ops::Bound::Unbounded, @@ -196,8 +193,8 @@ impl BTreeIndex { )) { result.extend(ids.iter().cloned()); } - } - _ => {} + }, + _ => {}, } result @@ -233,10 +230,7 @@ impl HashIndex { /// Insert value pub fn insert(&mut self, value: IndexedValue, id: String) { - self.map - .entry(value) - .or_default() - .insert(id); + self.map.entry(value).or_default().insert(id); } /// Remove value @@ -386,15 +380,15 @@ impl MetadataIndex { match self { MetadataIndex::BTree(idx) => { idx.insert(IndexedValue::from_json(value), id); - } + }, MetadataIndex::Hash(idx) => { idx.insert(IndexedValue::from_json(value), id); - } + }, MetadataIndex::Inverted(idx) => { if let serde_json::Value::String(text) = value { idx.insert(text, id); } - } + }, } Ok(()) } @@ -404,15 +398,15 @@ impl MetadataIndex { match self { MetadataIndex::BTree(idx) => { idx.remove(&IndexedValue::from_json(value), id); - } + }, MetadataIndex::Hash(idx) => { idx.remove(&IndexedValue::from_json(value), id); - } + }, MetadataIndex::Inverted(idx) => { if let serde_json::Value::String(text) = value { idx.remove(text, id); } - } + }, } Ok(()) } @@ -512,7 +506,7 @@ impl MetadataIndexManager { } else { HashSet::new() } - } + }, }; Some(result) @@ -527,7 +521,7 @@ impl MetadataIndexManager { let indexed_values: Vec = values.iter().map(IndexedValue::from_json).collect(); Some(idx.query_in(&indexed_values)) - } + }, _ => None, } } @@ -546,17 +540,17 @@ impl MetadataIndexManager { let unique = idx.tree.len(); let total = idx.all_ids().len(); (IndexType::BTree, unique, total) - } + }, MetadataIndex::Hash(idx) => { let unique = idx.map.len(); let total = idx.all_ids().len(); (IndexType::Hash, unique, total) - } + }, MetadataIndex::Inverted(idx) => { let unique = idx.index.len(); let total: usize = idx.index.values().map(|s| s.len()).sum(); (IndexType::Inverted, unique, total) - } + }, }; Some(IndexStats { diff --git a/src/metrics.rs b/src/metrics.rs index 6fe7568..26ce272 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -42,8 +42,8 @@ //! let metrics_text = prom.encode()?; //! ``` -use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; /// Configuration for metrics collection @@ -556,7 +556,9 @@ struct EmbeddedMetricsInner { impl EmbeddedMetrics { /// Create a new embedded metrics registry pub fn new() -> Self { - Self::with_buckets(vec![1.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 5000.0]) + Self::with_buckets(vec![ + 1.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 5000.0, + ]) } /// Create with custom latency buckets (in milliseconds) @@ -608,18 +610,24 @@ impl EmbeddedMetrics { /// Set the estimated memory usage pub fn set_memory_bytes(&self, bytes: u64) { - self.inner.estimated_memory_bytes.store(bytes, Ordering::Relaxed); + self.inner + .estimated_memory_bytes + .store(bytes, Ordering::Relaxed); } /// Set HNSW index stats pub fn set_hnsw_stats(&self, layers: u64, connections_per_layer: u64) { self.inner.hnsw_layers.store(layers, Ordering::Relaxed); - self.inner.hnsw_connections_per_layer.store(connections_per_layer, Ordering::Relaxed); + self.inner + .hnsw_connections_per_layer + .store(connections_per_layer, Ordering::Relaxed); } /// Set quantization stats pub fn set_quantization_stats(&self, quantized_count: u64, bits: u64) { - self.inner.quantized_vectors.store(quantized_count, Ordering::Relaxed); + self.inner + .quantized_vectors + .store(quantized_count, Ordering::Relaxed); self.inner.quantization_bits.store(bits, Ordering::Relaxed); } @@ -649,9 +657,8 @@ impl EmbeddedMetrics { if l.is_empty() { String::new() } else { - let pairs: Vec = l.iter() - .map(|(k, v)| format!("{}=\"{}\"", k, v)) - .collect(); + let pairs: Vec = + l.iter().map(|(k, v)| format!("{}=\"{}\"", k, v)).collect(); format!("{{{}}}", pairs.join(",")) } } else { @@ -727,7 +734,8 @@ impl EmbeddedMetrics { )); output.push_str(&format!( "vecstore_query_latency_ms_sum{} {:.2}\n", - label_str, snapshot.avg_query_latency_micros * snapshot.total_queries as f64 / 1000.0 + label_str, + snapshot.avg_query_latency_micros * snapshot.total_queries as f64 / 1000.0 )); output.push_str(&format!( "vecstore_query_latency_ms_count{} {}\n\n", @@ -780,7 +788,10 @@ impl EmbeddedMetrics { )); let layers = self.inner.hnsw_layers.load(Ordering::Relaxed); - let conns = self.inner.hnsw_connections_per_layer.load(Ordering::Relaxed); + let conns = self + .inner + .hnsw_connections_per_layer + .load(Ordering::Relaxed); output.push_str(&format!( "# HELP vecstore_hnsw_layers Number of HNSW layers\n\ # TYPE vecstore_hnsw_layers gauge\n\ diff --git a/src/migration.rs b/src/migration.rs index 3431fc5..5128abd 100644 --- a/src/migration.rs +++ b/src/migration.rs @@ -39,7 +39,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; @@ -253,7 +253,7 @@ impl Migrator { Err(e) => { stats.add_failure(format!("Line {}: Failed to read line: {}", line_num + 1, e)); continue; - } + }, }; if line.trim().is_empty() { @@ -265,15 +265,16 @@ impl Migrator { Ok(record) => { // Validate dimension if let Some(max_dim) = self.max_dimension - && record.vector.len() > max_dim { - stats.add_failure(format!( - "Line {}: Vector dimension {} exceeds maximum {}", - line_num + 1, - record.vector.len(), - max_dim - )); - continue; - } + && record.vector.len() > max_dim + { + stats.add_failure(format!( + "Line {}: Vector dimension {} exceeds maximum {}", + line_num + 1, + record.vector.len(), + max_dim + )); + continue; + } // Enforce dimension if specified let record = if let Some(target_dim) = self.enforce_dimension { @@ -285,14 +286,14 @@ impl Migrator { stats.total_vectors += record.vector.len(); records.push(record); stats.add_success(); - } + }, Err(e) => { if self.strict_mode { return Err(anyhow!("Line {}: {}", line_num + 1, e)); } else { stats.add_failure(format!("Line {}: {}", line_num + 1, e)); } - } + }, } } @@ -315,7 +316,7 @@ impl Migrator { metadata: record.metadata, timestamp: None, }) - } + }, SourceDatabase::Weaviate => { let record: WeaviateRecord = serde_json::from_str(line)?; @@ -325,7 +326,7 @@ impl Migrator { metadata: record.properties, timestamp: None, }) - } + }, SourceDatabase::Qdrant => { let record: QdrantRecord = serde_json::from_str(line)?; @@ -341,7 +342,7 @@ impl Migrator { metadata: record.payload, timestamp: None, }) - } + }, SourceDatabase::ChromaDB => { let record: ChromaRecord = serde_json::from_str(line)?; @@ -358,7 +359,7 @@ impl Migrator { metadata, timestamp: None, }) - } + }, SourceDatabase::Milvus => { let record: MilvusRecord = serde_json::from_str(line)?; @@ -374,13 +375,13 @@ impl Migrator { metadata: record.fields, timestamp: None, }) - } + }, SourceDatabase::Generic => { // Try to parse as generic migration record serde_json::from_str(line) .map_err(|e| anyhow!("Failed to parse generic record: {}", e)) - } + }, } } @@ -483,7 +484,7 @@ impl Migrator { Err(e) => { stats.add_failure(format!("Line {}: Failed to read: {}", line_num + 2, e)); continue; - } + }, }; let parts: Vec<&str> = line.split(',').collect(); @@ -506,7 +507,7 @@ impl Migrator { e )); continue; - } + }, } } diff --git a/src/ml_ranking.rs b/src/ml_ranking.rs index 664685e..24195ae 100644 --- a/src/ml_ranking.rs +++ b/src/ml_ranking.rs @@ -36,16 +36,16 @@ use serde::{Deserialize, Serialize}; use crate::error::Result; -#[cfg(feature = "embeddings")] -use std::path::Path; -#[cfg(feature = "embeddings")] -use std::sync::{Arc, Mutex, RwLock}; #[cfg(feature = "embeddings")] use crate::error::VecStoreError; #[cfg(feature = "embeddings")] use ndarray::Array2; #[cfg(feature = "embeddings")] -use ort::session::{builder::GraphOptimizationLevel, Session}; +use ort::session::{Session, builder::GraphOptimizationLevel}; +#[cfg(feature = "embeddings")] +use std::path::Path; +#[cfg(feature = "embeddings")] +use std::sync::{Arc, Mutex, RwLock}; /// Default batch size for ONNX inference const DEFAULT_BATCH_SIZE: usize = 32; @@ -69,10 +69,16 @@ impl OnnxSessionCache { } /// Get or create a session for a model path - fn get_or_create_session(&self, model_path: &str, num_threads: usize) -> Result>> { + fn get_or_create_session( + &self, + model_path: &str, + num_threads: usize, + ) -> Result>> { // Check if session already exists { - let sessions = self.sessions.read() + let sessions = self + .sessions + .read() .map_err(|e| VecStoreError::LockError(e.to_string()))?; if let Some(session) = sessions.get(model_path) { return Ok(Arc::clone(session)); @@ -82,13 +88,20 @@ impl OnnxSessionCache { // Create new session using ort 2.0 API let path = Path::new(model_path); if !path.exists() { - return Err(VecStoreError::NotFound(format!("Model file not found: {}", model_path))); + return Err(VecStoreError::NotFound(format!( + "Model file not found: {}", + model_path + ))); } let session = Session::builder() - .map_err(|e| VecStoreError::Internal(format!("Failed to create session builder: {}", e)))? + .map_err(|e| { + VecStoreError::Internal(format!("Failed to create session builder: {}", e)) + })? .with_optimization_level(GraphOptimizationLevel::Level3) - .map_err(|e| VecStoreError::Internal(format!("Failed to set optimization level: {}", e)))? + .map_err(|e| { + VecStoreError::Internal(format!("Failed to set optimization level: {}", e)) + })? .with_intra_threads(num_threads) .map_err(|e| VecStoreError::Internal(format!("Failed to set thread count: {}", e)))? .commit_from_file(path) @@ -98,7 +111,9 @@ impl OnnxSessionCache { // Cache the session { - let mut sessions = self.sessions.write() + let mut sessions = self + .sessions + .write() .map_err(|e| VecStoreError::LockError(e.to_string()))?; // Evict oldest if cache is full @@ -124,13 +139,16 @@ impl OnnxSessionCache { .map_err(|e| VecStoreError::Internal(format!("Failed to create dummy input: {}", e)))?; // Create Tensor object for ort 2.0 - let input_tensor = ort::value::Tensor::from_array(input_array) - .map_err(|e| VecStoreError::Internal(format!("Failed to create input tensor: {}", e)))?; + let input_tensor = ort::value::Tensor::from_array(input_array).map_err(|e| { + VecStoreError::Internal(format!("Failed to create input tensor: {}", e)) + })?; // Run warmup inference using ort 2.0 inputs! macro - let mut session_guard = session.lock() + let mut session_guard = session + .lock() .map_err(|e| VecStoreError::LockError(e.to_string()))?; - let _ = session_guard.run(ort::inputs![input_tensor]) + let _ = session_guard + .run(ort::inputs![input_tensor]) .map_err(|e| VecStoreError::Internal(format!("Warmup inference failed: {}", e)))?; Ok(()) @@ -138,7 +156,9 @@ impl OnnxSessionCache { /// Clear all cached sessions fn clear(&self) -> Result<()> { - let mut sessions = self.sessions.write() + let mut sessions = self + .sessions + .write() .map_err(|e| VecStoreError::LockError(e.to_string()))?; sessions.clear(); Ok(()) @@ -156,7 +176,9 @@ fn get_session_cache() -> Result<&'static OnnxSessionCache> { let cache = OnnxSessionCache::new()?; // Ignore the result if another thread initialized first let _ = CACHE.set(cache); - CACHE.get().ok_or_else(|| VecStoreError::Internal("Failed to initialize session cache".to_string())) + CACHE + .get() + .ok_or_else(|| VecStoreError::Internal("Failed to initialize session cache".to_string())) } /// Ranking configuration @@ -188,11 +210,21 @@ pub struct RankingConfig { pub warmup_models: bool, } -fn default_top_k() -> usize { 100 } -fn default_final_k() -> usize { 10 } -fn default_true() -> bool { true } -fn default_batch_size() -> usize { DEFAULT_BATCH_SIZE } -fn default_num_threads() -> usize { 4 } +fn default_top_k() -> usize { + 100 +} +fn default_final_k() -> usize { + 10 +} +fn default_true() -> bool { + true +} +fn default_batch_size() -> usize { + DEFAULT_BATCH_SIZE +} +fn default_num_threads() -> usize { + 4 +} impl Default for RankingConfig { fn default() -> Self { @@ -263,8 +295,7 @@ pub enum ScoreNormalization { } /// Model type for ONNX inference -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] -#[derive(Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] pub enum OnnxModelType { /// Generic ranking model (outputs a single relevance score) #[default] @@ -279,7 +310,6 @@ pub enum OnnxModelType { NeuralNetwork, } - /// ONNX model configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct OnnxModelConfig { @@ -308,8 +338,12 @@ pub struct OnnxModelConfig { pub output_index: usize, } -fn default_input_name() -> String { "input".to_string() } -fn default_output_name() -> String { "output".to_string() } +fn default_input_name() -> String { + "input".to_string() +} +fn default_output_name() -> String { + "output".to_string() +} impl OnnxModelConfig { /// Create a new ONNX model config @@ -409,14 +443,9 @@ pub enum RankingModel { num_features: usize, }, /// Linear model - Linear { - weights: Vec, - bias: f32, - }, + Linear { weights: Vec, bias: f32 }, /// Custom scoring function - Custom { - name: String, - }, + Custom { name: String }, } impl RankingModel { @@ -631,9 +660,10 @@ impl MLRanker { let candidates_with_features: Vec<(RankingCandidate, RankingFeatures)> = candidates .into_iter() .map(|c| { - let features = c.features.clone().unwrap_or_else(|| { - self.compute_features(query, &c) - }); + let features = c + .features + .clone() + .unwrap_or_else(|| self.compute_features(query, &c)); (c, features) }) .collect(); @@ -660,19 +690,21 @@ impl MLRanker { ) -> Result> { if self.models.is_empty() { // No models - just use original scores - return Ok(candidates.iter().map(|(c, features)| { - RankedResult { + return Ok(candidates + .iter() + .map(|(c, features)| RankedResult { id: c.id.clone(), score: c.original_score, original_score: c.original_score, feature_contributions: Some(self.get_feature_contributions(features)), metadata: c.metadata.clone(), - } - }).collect()); + }) + .collect()); } // Collect all feature vectors - let feature_vectors: Vec> = candidates.iter() + let feature_vectors: Vec> = candidates + .iter() .map(|(_, features)| features.to_vector()) .collect(); @@ -685,16 +717,17 @@ impl MLRanker { } // Combine scores for each candidate - let results: Vec = candidates.iter() + let results: Vec = candidates + .iter() .enumerate() .map(|(idx, (c, features))| { - let model_scores: Vec = all_model_scores.iter() - .map(|scores| scores[idx]) - .collect(); + let model_scores: Vec = + all_model_scores.iter().map(|scores| scores[idx]).collect(); let final_score = if self.config.ensemble_weights.len() == model_scores.len() { // Weighted ensemble - model_scores.iter() + model_scores + .iter() .zip(&self.config.ensemble_weights) .map(|(s, w)| s * w) .sum() @@ -723,30 +756,40 @@ impl MLRanker { feature_vectors: &[Vec], ) -> Result> { match model { - RankingModel::Linear { weights, bias } => { - Ok(feature_vectors.iter().map(|fv| { - let score: f32 = fv.iter() + RankingModel::Linear { weights, bias } => Ok(feature_vectors + .iter() + .map(|fv| { + let score: f32 = fv + .iter() .zip(weights.iter()) .map(|(f, w)| f * w) - .sum::() + bias; + .sum::() + + bias; self.sigmoid(score) - }).collect()) - } - RankingModel::ONNX(config) => { - self.score_batch_onnx(&config.path, feature_vectors, config.apply_sigmoid, config.apply_softmax, config.output_index) - } + }) + .collect()), + RankingModel::ONNX(config) => self.score_batch_onnx( + &config.path, + feature_vectors, + config.apply_sigmoid, + config.apply_softmax, + config.output_index, + ), RankingModel::XGBoost { path, .. } => { // XGBoost exported to ONNX - apply sigmoid for binary classification self.score_batch_onnx(path, feature_vectors, true, false, 0) - } + }, RankingModel::LightGBM { path, .. } => { // LightGBM exported to ONNX - apply sigmoid for binary classification self.score_batch_onnx(path, feature_vectors, true, false, 0) - } + }, RankingModel::Custom { .. } => { // Custom models just return original scores (first feature) - Ok(feature_vectors.iter().map(|fv| fv.first().copied().unwrap_or(0.0)).collect()) - } + Ok(feature_vectors + .iter() + .map(|fv| fv.first().copied().unwrap_or(0.0)) + .collect()) + }, } } @@ -772,9 +815,16 @@ impl MLRanker { // Process in batches for chunk in feature_vectors.chunks(batch_size) { - let mut session_guard = session.lock() + let mut session_guard = session + .lock() .map_err(|e| VecStoreError::LockError(e.to_string()))?; - let batch_scores = self.run_onnx_batch(&mut session_guard, chunk, apply_sigmoid, apply_softmax, output_index)?; + let batch_scores = self.run_onnx_batch( + &mut session_guard, + chunk, + apply_sigmoid, + apply_softmax, + output_index, + )?; all_scores.extend(batch_scores); } @@ -799,26 +849,27 @@ impl MLRanker { let feature_dim = batch[0].len(); // Flatten features into a single vector - let flat_features: Vec = batch.iter() - .flat_map(|v| v.iter().copied()) - .collect(); + let flat_features: Vec = batch.iter().flat_map(|v| v.iter().copied()).collect(); // Create input array let input_array = Array2::from_shape_vec((batch_size, feature_dim), flat_features) .map_err(|e| VecStoreError::Internal(format!("Failed to create input array: {}", e)))?; // Create Tensor object for ort 2.0 - let input_tensor = ort::value::Tensor::from_array(input_array) - .map_err(|e| VecStoreError::Internal(format!("Failed to create input tensor: {}", e)))?; + let input_tensor = ort::value::Tensor::from_array(input_array).map_err(|e| { + VecStoreError::Internal(format!("Failed to create input tensor: {}", e)) + })?; // Run inference using ort 2.0 inputs! macro - let outputs = session.run(ort::inputs![input_tensor]) + let outputs = session + .run(ort::inputs![input_tensor]) .map_err(|e| VecStoreError::Internal(format!("ONNX inference failed: {}", e)))?; // Extract output tensor using ort 2.0 API let output_tensor = &outputs[0]; - let output_view = output_tensor.try_extract_array::() + let output_view = output_tensor + .try_extract_array::() .map_err(|e| VecStoreError::Internal(format!("Failed to extract output: {}", e)))?; let output_array = output_view.to_owned(); @@ -850,12 +901,14 @@ impl MLRanker { // Apply softmax across classes, then extract target class let num_classes = output_shape[1]; let actual_index = output_index.min(num_classes - 1); - (0..batch_size).map(|i| { - let row: Vec = (0..num_classes).map(|j| output_array[[i, j]]).collect(); - let max_val = row.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); - let exp_sum: f32 = row.iter().map(|&x| (x - max_val).exp()).sum(); - (row[actual_index] - max_val).exp() / exp_sum - }).collect() + (0..batch_size) + .map(|i| { + let row: Vec = (0..num_classes).map(|j| output_array[[i, j]]).collect(); + let max_val = row.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); + let exp_sum: f32 = row.iter().map(|&x| (x - max_val).exp()).sum(); + (row[actual_index] - max_val).exp() / exp_sum + }) + .collect() } else if apply_sigmoid { scores.iter().map(|&x| self.sigmoid(x)).collect() } else { @@ -876,7 +929,8 @@ impl MLRanker { _output_index: usize, ) -> Result> { // Without embeddings feature, just return the first feature (vector_score) - Ok(feature_vectors.iter() + Ok(feature_vectors + .iter() .map(|fv| fv.first().copied().unwrap_or(0.0)) .collect()) } @@ -890,7 +944,8 @@ impl MLRanker { if !query.is_empty() && !candidate.vector.is_empty() { // Dot product - let dot: f32 = query.iter() + let dot: f32 = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| q * d) .sum(); @@ -913,45 +968,54 @@ impl MLRanker { qd_features.push(cosine); // Max element-wise product - let max_sim: f32 = query.iter() + let max_sim: f32 = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| q * d) .fold(f32::NEG_INFINITY, f32::max); qd_features.push(max_sim); // Min element-wise product - let min_sim: f32 = query.iter() + let min_sim: f32 = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| q * d) .fold(f32::INFINITY, f32::min); qd_features.push(min_sim); // Mean element-wise product - let mean_sim: f32 = query.iter() + let mean_sim: f32 = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| q * d) - .sum::() / query.len().max(1) as f32; + .sum::() + / query.len().max(1) as f32; qd_features.push(mean_sim); // Std of element-wise products - let products: Vec = query.iter() + let products: Vec = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| q * d) .collect(); - let variance: f32 = products.iter() + let variance: f32 = products + .iter() .map(|&p| (p - mean_sim).powi(2)) - .sum::() / products.len().max(1) as f32; + .sum::() + / products.len().max(1) as f32; qd_features.push(variance.sqrt()); // L1 distance (Manhattan) - let l1_dist: f32 = query.iter() + let l1_dist: f32 = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| (q - d).abs()) .sum(); qd_features.push(l1_dist); // L2 distance (Euclidean) - let l2_dist: f32 = query.iter() + let l2_dist: f32 = query + .iter() .zip(candidate.vector.iter()) .map(|(q, d)| (q - d).powi(2)) .sum::() @@ -959,8 +1023,7 @@ impl MLRanker { qd_features.push(l2_dist); } - RankingFeatures::new(vector_score) - .with_qd_features(qd_features) + RankingFeatures::new(vector_score).with_qd_features(qd_features) } /// Get feature contributions for explainability @@ -973,13 +1036,21 @@ impl MLRanker { } let qd_feature_names = [ - "dot_product", "query_norm", "doc_norm", "cosine_similarity", - "max_product", "min_product", "mean_product", "std_product", - "l1_distance", "l2_distance" + "dot_product", + "query_norm", + "doc_norm", + "cosine_similarity", + "max_product", + "min_product", + "mean_product", + "std_product", + "l1_distance", + "l2_distance", ]; for (i, &f) in features.qd_features.iter().enumerate() { - let name = qd_feature_names.get(i) + let name = qd_feature_names + .get(i) .map(|s| s.to_string()) .unwrap_or_else(|| format!("qd_feature_{}", i)); contributions.insert(name, f); @@ -1003,36 +1074,47 @@ impl MLRanker { } match self.config.normalization { - ScoreNormalization::None => {} + ScoreNormalization::None => {}, ScoreNormalization::MinMax => { - let min = results.iter().map(|r| r.score).fold(f32::INFINITY, f32::min); - let max = results.iter().map(|r| r.score).fold(f32::NEG_INFINITY, f32::max); + let min = results + .iter() + .map(|r| r.score) + .fold(f32::INFINITY, f32::min); + let max = results + .iter() + .map(|r| r.score) + .fold(f32::NEG_INFINITY, f32::max); let range = max - min; if range > 0.0 { for r in results { r.score = (r.score - min) / range; } } - } + }, ScoreNormalization::ZScore => { let mean: f32 = results.iter().map(|r| r.score).sum::() / results.len() as f32; - let variance: f32 = results.iter() + let variance: f32 = results + .iter() .map(|r| (r.score - mean).powi(2)) - .sum::() / results.len() as f32; + .sum::() + / results.len() as f32; let std = variance.sqrt(); if std > 0.0 { for r in results { r.score = (r.score - mean) / std; } } - } + }, ScoreNormalization::Softmax => { - let max = results.iter().map(|r| r.score).fold(f32::NEG_INFINITY, f32::max); + let max = results + .iter() + .map(|r| r.score) + .fold(f32::NEG_INFINITY, f32::max); let exp_sum: f32 = results.iter().map(|r| (r.score - max).exp()).sum(); for r in results { r.score = (r.score - max).exp() / exp_sum; } - } + }, } } @@ -1074,12 +1156,7 @@ impl RankingPipeline { } /// Add a ranking stage - pub fn add_stage( - &mut self, - name: impl Into, - ranker: MLRanker, - top_k: usize, - ) { + pub fn add_stage(&mut self, name: impl Into, ranker: MLRanker, top_k: usize) { self.stages.push(RankingStage { name: name.into(), ranker, @@ -1099,7 +1176,8 @@ impl RankingPipeline { let results = stage.ranker.rerank(query, candidates)?; // Convert results back to candidates for next stage - candidates = results.into_iter() + candidates = results + .into_iter() .take(stage.top_k) .map(|r| RankingCandidate { id: r.id, @@ -1115,7 +1193,8 @@ impl RankingPipeline { if let Some(last_stage) = self.stages.last_mut() { last_stage.ranker.rerank(query, candidates) } else { - Ok(candidates.into_iter() + Ok(candidates + .into_iter() .map(|c| RankedResult { id: c.id, score: c.original_score, @@ -1281,9 +1360,7 @@ mod tests { #[test] fn test_no_models_uses_original_scores() { - let config = RankingConfig::new() - .with_final_k(10) - .with_warmup(false); + let config = RankingConfig::new().with_final_k(10).with_warmup(false); let mut ranker = MLRanker::new(config).unwrap(); let candidates: Vec = vec![ @@ -1314,9 +1391,7 @@ mod tests { #[test] fn test_batch_size_config() { - let config = RankingConfig::new() - .with_batch_size(64) - .with_num_threads(8); + let config = RankingConfig::new().with_batch_size(64).with_num_threads(8); assert_eq!(config.batch_size, 64); assert_eq!(config.num_threads, 8); @@ -1403,12 +1478,8 @@ mod tests { #[test] fn test_pipeline_execution() { - let config1 = RankingConfig::new() - .with_final_k(5) - .with_warmup(false); - let config2 = RankingConfig::new() - .with_final_k(3) - .with_warmup(false); + let config1 = RankingConfig::new().with_final_k(5).with_warmup(false); + let config2 = RankingConfig::new().with_final_k(3).with_warmup(false); let ranker1 = MLRanker::new(config1).unwrap(); let ranker2 = MLRanker::new(config2).unwrap(); diff --git a/src/monitoring.rs b/src/monitoring.rs index 3fab4b2..ae79a51 100644 --- a/src/monitoring.rs +++ b/src/monitoring.rs @@ -101,7 +101,7 @@ impl AlertCondition { AlertCondition::NotEqual => (value - threshold).abs() >= 1e-6, AlertCondition::PercentageChange => { (value / threshold - 1.0).abs() > 0.1 // 10% change - } + }, } } } @@ -312,9 +312,10 @@ impl Monitor { // Check cooldown period if let Some(last_time) = self.last_alert_time.get(&rule.name) && let Ok(elapsed) = SystemTime::now().duration_since(*last_time) - && elapsed < self.config.alert_cooldown { - continue; - } + && elapsed < self.config.alert_cooldown + { + continue; + } // Evaluate condition if rule.condition.evaluate(value, rule.threshold) { diff --git a/src/multi_vector.rs b/src/multi_vector.rs index 5a57b01..9ceca93 100644 --- a/src/multi_vector.rs +++ b/src/multi_vector.rs @@ -57,7 +57,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -213,10 +213,7 @@ impl MultiVectorIndex { let token_vec = &self.token_vectors[token_id]; let sim = cosine_similarity(query_vec, token_vec); - doc_scores - .entry(doc_id.clone()) - .or_default() - .push(sim); + doc_scores.entry(doc_id.clone()).or_default().push(sim); } } @@ -227,7 +224,7 @@ impl MultiVectorIndex { let score = match self.aggregation { AggregationMethod::MaxSim => { sims.iter().copied().fold(f32::NEG_INFINITY, f32::max) - } + }, AggregationMethod::AvgSim => sims.iter().sum::() / sims.len() as f32, AggregationMethod::SumSim => sims.iter().sum(), AggregationMethod::FirstToken => sims.first().copied().unwrap_or(0.0), @@ -607,7 +604,7 @@ impl OptimizedMultiVectorIndex { let score = match self.aggregation { AggregationMethod::MaxSim => { sims.iter().copied().fold(f32::NEG_INFINITY, f32::max) - } + }, AggregationMethod::AvgSim => sims.iter().sum::() / sims.len() as f32, AggregationMethod::SumSim => sims.iter().sum(), AggregationMethod::FirstToken => sims.first().copied().unwrap_or(0.0), @@ -729,7 +726,7 @@ impl LateInteractionScorer { ScoreMode::SumMaxSim => maxsim_per_query.iter().sum(), ScoreMode::AvgMaxSim => { maxsim_per_query.iter().sum::() / maxsim_per_query.len() as f32 - } + }, ScoreMode::MaxMaxSim => maxsim_per_query .iter() .copied() diff --git a/src/multilang.rs b/src/multilang.rs index 36aeb4b..7f6b97e 100644 --- a/src/multilang.rs +++ b/src/multilang.rs @@ -21,9 +21,8 @@ //! // ["gato", "corr", "rapid"] - stemmed, no stopwords //! ``` -use std::collections::{HashMap, HashSet}; use serde::{Deserialize, Serialize}; - +use std::collections::{HashMap, HashSet}; /// Supported languages #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] @@ -161,15 +160,33 @@ impl Language { /// Get all supported languages pub fn all() -> Vec { vec![ - Language::English, Language::Spanish, Language::French, - Language::German, Language::Italian, Language::Portuguese, - Language::Dutch, Language::Swedish, Language::Norwegian, - Language::Danish, Language::Finnish, Language::Russian, - Language::Arabic, Language::Turkish, Language::Chinese, - Language::Japanese, Language::Korean, Language::Greek, - Language::Hebrew, Language::Hindi, Language::Thai, - Language::Vietnamese, Language::Polish, Language::Czech, - Language::Hungarian, Language::Romanian, Language::Indonesian, + Language::English, + Language::Spanish, + Language::French, + Language::German, + Language::Italian, + Language::Portuguese, + Language::Dutch, + Language::Swedish, + Language::Norwegian, + Language::Danish, + Language::Finnish, + Language::Russian, + Language::Arabic, + Language::Turkish, + Language::Chinese, + Language::Japanese, + Language::Korean, + Language::Greek, + Language::Hebrew, + Language::Hindi, + Language::Thai, + Language::Vietnamese, + Language::Polish, + Language::Czech, + Language::Hungarian, + Language::Romanian, + Language::Indonesian, ] } } @@ -177,69 +194,79 @@ impl Language { /// Stopwords for each language fn get_stopwords(lang: Language) -> HashSet<&'static str> { match lang { - Language::English => { - ["a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for", - "of", "with", "by", "from", "as", "is", "was", "are", "were", "been", - "be", "have", "has", "had", "do", "does", "did", "will", "would", - "could", "should", "may", "might", "must", "can", "this", "that", - "these", "those", "i", "you", "he", "she", "it", "we", "they", - "what", "which", "who", "when", "where", "why", "how", "all", - "each", "every", "both", "few", "more", "most", "other", "some", - "such", "no", "not", "only", "same", "so", "than", "too", "very"] - .iter().cloned().collect() - } - Language::Spanish => { - ["de", "la", "que", "el", "en", "y", "a", "los", "del", "se", - "las", "por", "un", "para", "con", "no", "una", "su", "al", - "es", "lo", "como", "mĆ”s", "pero", "sus", "le", "ya", "o", - "fue", "este", "ha", "sĆ­", "porque", "esta", "son", "entre", - "estĆ”", "cuando", "muy", "sin", "sobre", "ser", "tiene", "tambiĆ©n"] - .iter().cloned().collect() - } - Language::French => { - ["le", "la", "les", "de", "du", "des", "un", "une", "et", "en", - "Ć ", "au", "aux", "ce", "que", "qui", "dans", "pour", "sur", - "par", "avec", "ne", "pas", "plus", "ou", "son", "sa", "ses", - "il", "elle", "nous", "vous", "ils", "elles", "je", "tu", - "est", "sont", "ĆŖtre", "avoir", "fait", "faire", "comme"] - .iter().cloned().collect() - } - Language::German => { - ["der", "die", "das", "und", "in", "zu", "den", "von", "mit", - "ist", "nicht", "ein", "eine", "als", "auch", "es", "an", - "auf", "für", "sich", "des", "dem", "werden", "bei", "haben", - "wird", "sind", "oder", "nach", "am", "um", "aus", "nur", - "wie", "über", "so", "wenn", "aber", "noch", "durch", "kann"] - .iter().cloned().collect() - } - Language::Italian => { - ["il", "lo", "la", "i", "gli", "le", "un", "uno", "una", "di", - "a", "da", "in", "con", "su", "per", "tra", "fra", "che", - "e", "ma", "o", "se", "come", "non", "più", "anche", "solo", - "questo", "quello", "suo", "loro", "essere", "avere", "fare"] - .iter().cloned().collect() - } - Language::Russian => { - ["Šø", "в", "не", "на", "я", "что", "он", "с", "ŃŃ‚Š¾", "а", - "как", "она", "по", "но", "они", "Šŗ", "у", "ты", "ŠøŠ·", "мы", - "за", "от", "о", "же", "все", "так", "его", "её", "ŠøŃ…"] - .iter().cloned().collect() - } - Language::Arabic => { - ["في", "من", "على", "؄لى", "عن", "Ł…Ų¹", "هذا", "هذه", "Ų§Ł„ŲŖŁŠ", - "Ų§Ł„Ų°ŁŠ", "ŁƒŲ§Ł†", "لم", "لا", "Ł…Ų§", "Ł‡Łˆ", "Ł‡ŁŠ", "أن", "أو", - "و", "Ų«Ł…", "ŲØŲ¹ŲÆ", "قبل", "ŁƒŁ„", "ŲØŁŠŁ†", "حتى", "Ų„Ų°Ų§"] - .iter().cloned().collect() - } + Language::English => [ + "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", + "by", "from", "as", "is", "was", "are", "were", "been", "be", "have", "has", "had", + "do", "does", "did", "will", "would", "could", "should", "may", "might", "must", "can", + "this", "that", "these", "those", "i", "you", "he", "she", "it", "we", "they", "what", + "which", "who", "when", "where", "why", "how", "all", "each", "every", "both", "few", + "more", "most", "other", "some", "such", "no", "not", "only", "same", "so", "than", + "too", "very", + ] + .iter() + .cloned() + .collect(), + Language::Spanish => [ + "de", "la", "que", "el", "en", "y", "a", "los", "del", "se", "las", "por", "un", + "para", "con", "no", "una", "su", "al", "es", "lo", "como", "mĆ”s", "pero", "sus", "le", + "ya", "o", "fue", "este", "ha", "sĆ­", "porque", "esta", "son", "entre", "estĆ”", + "cuando", "muy", "sin", "sobre", "ser", "tiene", "tambiĆ©n", + ] + .iter() + .cloned() + .collect(), + Language::French => [ + "le", "la", "les", "de", "du", "des", "un", "une", "et", "en", "Ć ", "au", "aux", "ce", + "que", "qui", "dans", "pour", "sur", "par", "avec", "ne", "pas", "plus", "ou", "son", + "sa", "ses", "il", "elle", "nous", "vous", "ils", "elles", "je", "tu", "est", "sont", + "ĆŖtre", "avoir", "fait", "faire", "comme", + ] + .iter() + .cloned() + .collect(), + Language::German => [ + "der", "die", "das", "und", "in", "zu", "den", "von", "mit", "ist", "nicht", "ein", + "eine", "als", "auch", "es", "an", "auf", "für", "sich", "des", "dem", "werden", "bei", + "haben", "wird", "sind", "oder", "nach", "am", "um", "aus", "nur", "wie", "über", "so", + "wenn", "aber", "noch", "durch", "kann", + ] + .iter() + .cloned() + .collect(), + Language::Italian => [ + "il", "lo", "la", "i", "gli", "le", "un", "uno", "una", "di", "a", "da", "in", "con", + "su", "per", "tra", "fra", "che", "e", "ma", "o", "se", "come", "non", "più", "anche", + "solo", "questo", "quello", "suo", "loro", "essere", "avere", "fare", + ] + .iter() + .cloned() + .collect(), + Language::Russian => [ + "Šø", "в", "не", "на", "я", "что", "он", "с", "ŃŃ‚Š¾", "а", "как", "она", "по", "но", + "они", "Šŗ", "у", "ты", "ŠøŠ·", "мы", "за", "от", "о", "же", "все", "так", "его", "её", + "ŠøŃ…", + ] + .iter() + .cloned() + .collect(), + Language::Arabic => [ + "في", "من", "على", "؄لى", "عن", "Ł…Ų¹", "هذا", "هذه", "Ų§Ł„ŲŖŁŠ", "Ų§Ł„Ų°ŁŠ", "ŁƒŲ§Ł†", "لم", "لا", + "Ł…Ų§", "Ł‡Łˆ", "Ł‡ŁŠ", "أن", "أو", "و", "Ų«Ł…", "ŲØŲ¹ŲÆ", "قبل", "ŁƒŁ„", "ŲØŁŠŁ†", "حتى", "Ų„Ų°Ų§", + ] + .iter() + .cloned() + .collect(), Language::Chinese | Language::Japanese | Language::Korean => { // CJK languages typically don't use stopword removal HashSet::new() - } + }, _ => { // Default minimal stopwords ["a", "an", "the", "and", "or", "but", "in", "on", "at", "to"] - .iter().cloned().collect() - } + .iter() + .cloned() + .collect() + }, } } @@ -255,58 +282,120 @@ impl StemmingRules { match lang { Language::English => Self { suffixes: vec![ - ("ational", "ate"), ("tional", "tion"), ("enci", "ence"), - ("anci", "ance"), ("izer", "ize"), ("isation", "ize"), - ("ization", "ize"), ("ation", "ate"), ("ator", "ate"), - ("alism", "al"), ("iveness", "ive"), ("fulness", "ful"), - ("ousness", "ous"), ("aliti", "al"), ("iviti", "ive"), - ("biliti", "ble"), ("alli", "al"), ("entli", "ent"), - ("eli", "e"), ("ousli", "ous"), ("ling", ""), - ("ement", ""), ("ment", ""), ("ness", ""), ("ing", ""), - ("ies", "y"), ("es", ""), ("ed", ""), ("s", ""), + ("ational", "ate"), + ("tional", "tion"), + ("enci", "ence"), + ("anci", "ance"), + ("izer", "ize"), + ("isation", "ize"), + ("ization", "ize"), + ("ation", "ate"), + ("ator", "ate"), + ("alism", "al"), + ("iveness", "ive"), + ("fulness", "ful"), + ("ousness", "ous"), + ("aliti", "al"), + ("iviti", "ive"), + ("biliti", "ble"), + ("alli", "al"), + ("entli", "ent"), + ("eli", "e"), + ("ousli", "ous"), + ("ling", ""), + ("ement", ""), + ("ment", ""), + ("ness", ""), + ("ing", ""), + ("ies", "y"), + ("es", ""), + ("ed", ""), + ("s", ""), ], min_stem_length: 3, }, Language::Spanish => Self { suffixes: vec![ - ("amiento", ""), ("imientos", ""), ("imiento", ""), - ("aciones", ""), ("adores", ""), ("amente", ""), - ("idades", ""), ("mente", ""), ("ables", ""), - ("ibles", ""), ("ación", ""), ("ador", ""), - ("ante", ""), ("ando", ""), ("endo", ""), - ("ido", ""), ("ido", ""), ("ar", ""), ("er", ""), - ("ir", ""), ("as", ""), ("es", ""), ("os", ""), + ("amiento", ""), + ("imientos", ""), + ("imiento", ""), + ("aciones", ""), + ("adores", ""), + ("amente", ""), + ("idades", ""), + ("mente", ""), + ("ables", ""), + ("ibles", ""), + ("ación", ""), + ("ador", ""), + ("ante", ""), + ("ando", ""), + ("endo", ""), + ("ido", ""), + ("ido", ""), + ("ar", ""), + ("er", ""), + ("ir", ""), + ("as", ""), + ("es", ""), + ("os", ""), ], min_stem_length: 3, }, Language::French => Self { suffixes: vec![ - ("issements", ""), ("issement", ""), ("atrices", ""), - ("ateurs", ""), ("ations", ""), ("logies", ""), - ("usion", "u"), ("ution", "u"), ("ences", ""), - ("ances", ""), ("ments", ""), ("ement", ""), - ("euses", ""), ("ables", ""), ("istes", ""), - ("ation", ""), ("ique", ""), ("isme", ""), - ("able", ""), ("iste", ""), ("ment", ""), - ("ence", ""), ("ance", ""), ("euse", ""), - ("eux", ""), ("ant", ""), ("ent", ""), + ("issements", ""), + ("issement", ""), + ("atrices", ""), + ("ateurs", ""), + ("ations", ""), + ("logies", ""), + ("usion", "u"), + ("ution", "u"), + ("ences", ""), + ("ances", ""), + ("ments", ""), + ("ement", ""), + ("euses", ""), + ("ables", ""), + ("istes", ""), + ("ation", ""), + ("ique", ""), + ("isme", ""), + ("able", ""), + ("iste", ""), + ("ment", ""), + ("ence", ""), + ("ance", ""), + ("euse", ""), + ("eux", ""), + ("ant", ""), + ("ent", ""), ("e", ""), ], min_stem_length: 3, }, Language::German => Self { suffixes: vec![ - ("ungen", ""), ("heit", ""), ("keit", ""), - ("lich", ""), ("isch", ""), ("ung", ""), - ("ig", ""), ("ik", ""), ("en", ""), ("er", ""), - ("em", ""), ("es", ""), ("e", ""), ("s", ""), + ("ungen", ""), + ("heit", ""), + ("keit", ""), + ("lich", ""), + ("isch", ""), + ("ung", ""), + ("ig", ""), + ("ik", ""), + ("en", ""), + ("er", ""), + ("em", ""), + ("es", ""), + ("e", ""), + ("s", ""), ], min_stem_length: 3, }, _ => Self { - suffixes: vec![ - ("ing", ""), ("ed", ""), ("es", ""), ("s", ""), - ], + suffixes: vec![("ing", ""), ("ed", ""), ("es", ""), ("s", "")], min_stem_length: 3, }, } @@ -406,13 +495,13 @@ impl MultiLangAnalyzer { &text[start..start + c.len_utf8()] }) .collect() - } + }, _ => { // Word-based tokenization for most languages text.split(|c: char| !c.is_alphanumeric() && c != '\'') .filter(|s| !s.is_empty()) .collect() - } + }, } } @@ -521,13 +610,22 @@ impl AutoLangAnalyzer { // Check common words for European languages let common_words: HashMap<&str, Language> = [ - ("the", Language::English), ("and", Language::English), - ("der", Language::German), ("und", Language::German), - ("le", Language::French), ("et", Language::French), - ("el", Language::Spanish), ("y", Language::Spanish), - ("il", Language::Italian), ("che", Language::Italian), - ("het", Language::Dutch), ("en", Language::Dutch), - ].iter().cloned().collect(); + ("the", Language::English), + ("and", Language::English), + ("der", Language::German), + ("und", Language::German), + ("le", Language::French), + ("et", Language::French), + ("el", Language::Spanish), + ("y", Language::Spanish), + ("il", Language::Italian), + ("che", Language::Italian), + ("het", Language::Dutch), + ("en", Language::Dutch), + ] + .iter() + .cloned() + .collect(); for word in text_lower.split_whitespace() { if let Some(&lang) = common_words.get(word) { diff --git a/src/multimodal.rs b/src/multimodal.rs index 934c76a..a872ba8 100644 --- a/src/multimodal.rs +++ b/src/multimodal.rs @@ -56,7 +56,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use rayon::prelude::*; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -501,7 +501,7 @@ impl MultiModalIndex { } Ok(fused) - } + }, MultiModalFusion::WeightedAverage(weights) => { let embeddings: Vec<&Vec> = query_embeddings.values().collect(); @@ -515,7 +515,7 @@ impl MultiModalIndex { } Ok(fused) - } + }, MultiModalFusion::Concatenate => { // Concatenate all embeddings @@ -524,7 +524,7 @@ impl MultiModalIndex { fused.extend_from_slice(embedding); } Ok(fused) - } + }, MultiModalFusion::Max => { let embeddings: Vec<&Vec> = query_embeddings.values().collect(); @@ -538,12 +538,12 @@ impl MultiModalIndex { } Ok(fused) - } + }, MultiModalFusion::LateFusion { .. } => { // Should not reach here (handled separately) Err(anyhow!("Late fusion not applicable in early fusion search")) - } + }, } } diff --git a/src/multimodal_search.rs b/src/multimodal_search.rs index fd86fd1..1adf40d 100644 --- a/src/multimodal_search.rs +++ b/src/multimodal_search.rs @@ -296,7 +296,9 @@ impl Embedder for MockEmbedder { for (i, chunk) in bytes.chunks(4).enumerate() { let idx = i % self.dimension; - let val: u32 = chunk.iter().fold(0u32, |acc, &b| acc.wrapping_add(b as u32)); + let val: u32 = chunk + .iter() + .fold(0u32, |acc, &b| acc.wrapping_add(b as u32)); embedding[idx] = ((val as f32) / 255.0 - 0.5) * 2.0; } @@ -362,15 +364,17 @@ impl MultimodalIndex { // Add to documents { - let mut docs = self.documents.write() - .map_err(|_| VecStoreError::LockError("failed to acquire documents write lock".into()))?; + let mut docs = self.documents.write().map_err(|_| { + VecStoreError::LockError("failed to acquire documents write lock".into()) + })?; docs.insert(id.clone(), doc); } // Add to modality index { - let mut indexes = self.modality_indexes.write() - .map_err(|_| VecStoreError::LockError("failed to acquire modality_indexes write lock".into()))?; + let mut indexes = self.modality_indexes.write().map_err(|_| { + VecStoreError::LockError("failed to acquire modality_indexes write lock".into()) + })?; indexes.entry(modality).or_insert_with(Vec::new).push(id); } @@ -378,7 +382,12 @@ impl MultimodalIndex { } /// Insert text document - pub fn insert_text(&self, id: &str, text: &str, metadata: HashMap) -> Result<()> { + pub fn insert_text( + &self, + id: &str, + text: &str, + metadata: HashMap, + ) -> Result<()> { let embedding = self.embedder.embed_text(text)?; let doc = MultimodalDocument { @@ -396,7 +405,13 @@ impl MultimodalIndex { } /// Insert image document - pub fn insert_image(&self, id: &str, image_ref: &str, caption: Option<&str>, metadata: HashMap) -> Result<()> { + pub fn insert_image( + &self, + id: &str, + image_ref: &str, + caption: Option<&str>, + metadata: HashMap, + ) -> Result<()> { let embedding = self.embedder.embed_image(image_ref)?; let mut aux = HashMap::new(); @@ -432,19 +447,30 @@ impl MultimodalIndex { return Err(VecStoreError::InvalidInput("No query provided".to_string())); }; - let docs = self.documents.read() - .map_err(|_| VecStoreError::LockError("failed to acquire documents read lock".into()))?; - let indexes = self.modality_indexes.read() - .map_err(|_| VecStoreError::LockError("failed to acquire modality_indexes read lock".into()))?; + let docs = self.documents.read().map_err(|_| { + VecStoreError::LockError("failed to acquire documents read lock".into()) + })?; + let indexes = self.modality_indexes.read().map_err(|_| { + VecStoreError::LockError("failed to acquire modality_indexes read lock".into()) + })?; // Determine which documents to search let candidate_ids: Vec<&String> = if let Some(target) = &query.target_modality { // Search only target modality - indexes.get(target).map(|ids| ids.iter().collect()).unwrap_or_default() + indexes + .get(target) + .map(|ids| ids.iter().collect()) + .unwrap_or_default() } else if let Some(filter) = &query.modality_filter { // Search filtered modalities - filter.iter() - .flat_map(|m| indexes.get(m).map(|ids| ids.iter()).unwrap_or_else(|| [].iter())) + filter + .iter() + .flat_map(|m| { + indexes + .get(m) + .map(|ids| ids.iter()) + .unwrap_or_else(|| [].iter()) + }) .collect() } else { // Search all @@ -474,16 +500,17 @@ impl MultimodalIndex { let final_score = if let Some(ref cs) = cross_scores { match &self.config.fusion_strategy { FusionStrategy::WeightedSum => { - let aux_avg: f32 = cs.values().sum::() / cs.len().max(1) as f32; + let aux_avg: f32 = + cs.values().sum::() / cs.len().max(1) as f32; score * 0.7 + aux_avg * 0.3 - } + }, FusionStrategy::MaxScore => { cs.values().fold(score, |max, &s| max.max(s)) - } + }, FusionStrategy::Average => { let sum: f32 = cs.values().sum::() + score; sum / (cs.len() + 1) as f32 - } + }, FusionStrategy::LateFusion { weights } => { let mut weighted_sum = score; for (name, &s) in cs { @@ -491,16 +518,17 @@ impl MultimodalIndex { weighted_sum += s * w; } weighted_sum - } + }, } } else { score }; // Apply metadata filters - let passes_filters = query.filters.iter().all(|(key, value)| { - doc.metadata.get(key) == Some(value) - }); + let passes_filters = query + .filters + .iter() + .all(|(key, value)| doc.metadata.get(key) == Some(value)); if passes_filters { Some(MultimodalResult { @@ -553,16 +581,22 @@ impl MultimodalIndex { /// Get document by ID pub fn get(&self, id: &str) -> Option { - let Ok(docs) = self.documents.read() else { return None; }; + let Ok(docs) = self.documents.read() else { + return None; + }; docs.get(id).cloned() } /// Delete document pub fn delete(&self, id: &str) -> bool { - let Ok(mut docs) = self.documents.write() else { return false; }; + let Ok(mut docs) = self.documents.write() else { + return false; + }; if let Some(doc) = docs.remove(id) { - let Ok(mut indexes) = self.modality_indexes.write() else { return false; }; + let Ok(mut indexes) = self.modality_indexes.write() else { + return false; + }; if let Some(ids) = indexes.get_mut(&doc.modality) { ids.retain(|x| x != id); } @@ -643,10 +677,19 @@ mod tests { let index = MultimodalIndex::new(config); // Insert text - index.insert_text("txt1", "A beautiful sunset over the ocean", HashMap::new()).unwrap(); + index + .insert_text("txt1", "A beautiful sunset over the ocean", HashMap::new()) + .unwrap(); // Insert image - index.insert_image("img1", "/path/to/sunset.jpg", Some("sunset photo"), HashMap::new()).unwrap(); + index + .insert_image( + "img1", + "/path/to/sunset.jpg", + Some("sunset photo"), + HashMap::new(), + ) + .unwrap(); // Search let results = index.search(MultimodalQuery::text("sunset")).unwrap(); @@ -660,12 +703,14 @@ mod tests { // Insert images for i in 0..5 { - index.insert_image( - &format!("img_{}", i), - &format!("/images/{}.jpg", i), - Some(&format!("Photo number {}", i)), - HashMap::new(), - ).unwrap(); + index + .insert_image( + &format!("img_{}", i), + &format!("/images/{}.jpg", i), + Some(&format!("Photo number {}", i)), + HashMap::new(), + ) + .unwrap(); } // Text to image search @@ -682,14 +727,17 @@ mod tests { let config = MultimodalConfig::default(); let index = MultimodalIndex::new(config); - index.insert_text("txt1", "Hello world", HashMap::new()).unwrap(); - index.insert_image("img1", "/img.jpg", None, HashMap::new()).unwrap(); + index + .insert_text("txt1", "Hello world", HashMap::new()) + .unwrap(); + index + .insert_image("img1", "/img.jpg", None, HashMap::new()) + .unwrap(); // Filter to only text - let results = index.search( - MultimodalQuery::text("hello") - .with_modality_filter(vec![Modality::Text]) - ).unwrap(); + let results = index + .search(MultimodalQuery::text("hello").with_modality_filter(vec![Modality::Text])) + .unwrap(); assert!(results.iter().all(|r| r.modality == Modality::Text)); } @@ -699,9 +747,15 @@ mod tests { let config = MultimodalConfig::default(); let index = MultimodalIndex::new(config); - index.insert_text("txt1", "Text document", HashMap::new()).unwrap(); - index.insert_text("txt2", "Another text", HashMap::new()).unwrap(); - index.insert_image("img1", "/img.jpg", None, HashMap::new()).unwrap(); + index + .insert_text("txt1", "Text document", HashMap::new()) + .unwrap(); + index + .insert_text("txt2", "Another text", HashMap::new()) + .unwrap(); + index + .insert_image("img1", "/img.jpg", None, HashMap::new()) + .unwrap(); let stats = index.stats(); assert_eq!(stats.total_documents, 3); diff --git a/src/namespace.rs b/src/namespace.rs index 526b3ba..1fa749c 100644 --- a/src/namespace.rs +++ b/src/namespace.rs @@ -3,7 +3,7 @@ //! Provides namespace isolation, per-namespace quotas, and resource management //! for SaaS deployments. -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::time::{SystemTime, UNIX_EPOCH}; @@ -177,33 +177,36 @@ impl ResourceUsage { pub fn check_quotas(&self, quotas: &NamespaceQuotas) -> Result<()> { // Check vector count if let Some(max_vectors) = quotas.max_vectors - && self.vector_count >= max_vectors { - return Err(anyhow!( - "Vector quota exceeded: {} / {} vectors", - self.vector_count, - max_vectors - )); - } + && self.vector_count >= max_vectors + { + return Err(anyhow!( + "Vector quota exceeded: {} / {} vectors", + self.vector_count, + max_vectors + )); + } // Check storage if let Some(max_storage) = quotas.max_storage_bytes - && self.storage_bytes >= max_storage { - return Err(anyhow!( - "Storage quota exceeded: {} / {} bytes", - self.storage_bytes, - max_storage - )); - } + && self.storage_bytes >= max_storage + { + return Err(anyhow!( + "Storage quota exceeded: {} / {} bytes", + self.storage_bytes, + max_storage + )); + } // Check concurrent queries if let Some(max_concurrent) = quotas.max_concurrent_queries - && self.active_queries >= max_concurrent { - return Err(anyhow!( - "Concurrent query limit exceeded: {} / {}", - self.active_queries, - max_concurrent - )); - } + && self.active_queries >= max_concurrent + { + return Err(anyhow!( + "Concurrent query limit exceeded: {} / {}", + self.active_queries, + max_concurrent + )); + } Ok(()) } @@ -223,9 +226,10 @@ impl ResourceUsage { // Check rate limit if let Some(max_rps) = quotas.max_requests_per_second - && self.requests_current_window as f64 >= max_rps { - return Err(anyhow!("Rate limit exceeded: {} requests/second", max_rps)); - } + && self.requests_current_window as f64 >= max_rps + { + return Err(anyhow!("Rate limit exceeded: {} requests/second", max_rps)); + } self.requests_current_window += 1; self.total_requests += 1; @@ -277,24 +281,26 @@ impl Namespace { // Check if adding 'count' vectors would exceed quota if let Some(max_vectors) = self.quotas.max_vectors - && self.usage.vector_count + count > max_vectors { - return Err(anyhow!( - "Upsert would exceed vector quota: {} + {} > {}", - self.usage.vector_count, - count, - max_vectors - )); - } + && self.usage.vector_count + count > max_vectors + { + return Err(anyhow!( + "Upsert would exceed vector quota: {} + {} > {}", + self.usage.vector_count, + count, + max_vectors + )); + } // Check batch size if let Some(max_batch) = self.quotas.max_batch_size - && count > max_batch { - return Err(anyhow!( - "Batch size exceeds limit: {} > {}", - count, - max_batch - )); - } + && count > max_batch + { + return Err(anyhow!( + "Batch size exceeds limit: {} > {}", + count, + max_batch + )); + } Ok(()) } @@ -305,15 +311,15 @@ impl Namespace { match self.status { NamespaceStatus::PendingDeletion => { return Err(anyhow!("Namespace is pending deletion")); - } + }, NamespaceStatus::Suspended => { return Err(anyhow!( "Namespace is suspended. Contact administrator to reactivate." )); - } + }, NamespaceStatus::Active | NamespaceStatus::ReadOnly => { // OK to query - } + }, } // Check concurrent query limit @@ -321,13 +327,14 @@ impl Namespace { // Check result limit if let Some(max_results) = self.quotas.max_results_per_query - && k > max_results { - return Err(anyhow!( - "Query limit exceeds maximum: {} > {}", - k, - max_results - )); - } + && k > max_results + { + return Err(anyhow!( + "Query limit exceeds maximum: {} > {}", + k, + max_results + )); + } Ok(()) } diff --git a/src/namespace_manager.rs b/src/namespace_manager.rs index 52bf7dd..60b3425 100644 --- a/src/namespace_manager.rs +++ b/src/namespace_manager.rs @@ -5,7 +5,7 @@ use crate::namespace::{Namespace, NamespaceId, NamespaceQuotas, NamespaceStatus}; use crate::store::{Metadata, Neighbor, Query, VecStore}; -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::{Arc, PoisonError, RwLock}; diff --git a/src/neural_ranker.rs b/src/neural_ranker.rs index e145993..98cd08a 100644 --- a/src/neural_ranker.rs +++ b/src/neural_ranker.rs @@ -408,10 +408,14 @@ impl NeuralRanker { max_length: self.config.max_seq_length, }; - *self.model.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire model write lock".into()))? = Some(model); - *self.tokenizer.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire tokenizer write lock".into()))? = Some(tokenizer); + *self + .model + .write() + .map_err(|_| VecStoreError::LockError("Failed to acquire model write lock".into()))? = + Some(model); + *self.tokenizer.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire tokenizer write lock".into()) + })? = Some(tokenizer); Ok(()) } @@ -499,10 +503,13 @@ impl NeuralRanker { pub fn rerank(&self, query: &str, candidates: Vec) -> Result { let start = Instant::now(); self.stats.total_reranks.fetch_add(1, Ordering::Relaxed); - self.stats.total_candidates.fetch_add(candidates.len() as u64, Ordering::Relaxed); + self.stats + .total_candidates + .fetch_add(candidates.len() as u64, Ordering::Relaxed); // Limit candidates - let top_candidates: Vec<_> = candidates.into_iter() + let top_candidates: Vec<_> = candidates + .into_iter() .take(self.config.rerank_top_k) .collect(); @@ -517,8 +524,9 @@ impl NeuralRanker { // Check cache let neural_score = { - let cache = self.cache.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire cache read lock".into()))?; + let cache = self.cache.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire cache read lock".into()) + })?; if let Some(entry) = cache.entries.get(&cache_key) { cache_hit = true; self.stats.cache_hits.fetch_add(1, Ordering::Relaxed); @@ -531,14 +539,18 @@ impl NeuralRanker { let score = self.compute_score(query, &candidate.text)?; // Cache result - let mut cache = self.cache.write() - .map_err(|_| VecStoreError::LockError("Failed to acquire cache write lock".into()))?; + let mut cache = self.cache.write().map_err(|_| { + VecStoreError::LockError("Failed to acquire cache write lock".into()) + })?; if cache.entries.len() < cache.max_size { - cache.entries.insert(cache_key, CacheEntry { - embedding: vec![score], - created_at: Instant::now(), - access_count: 1, - }); + cache.entries.insert( + cache_key, + CacheEntry { + embedding: vec![score], + created_at: Instant::now(), + access_count: 1, + }, + ); } score @@ -624,7 +636,9 @@ impl NeuralRanker { // Use token ID and position to seed the embedding // This creates consistent, position-aware embeddings - let mut seed = (token_id as u64).wrapping_mul(31).wrapping_add(position as u64); + let mut seed = (token_id as u64) + .wrapping_mul(31) + .wrapping_add(position as u64); for (i, emb_val) in embedding.iter_mut().enumerate().take(dim) { // LCG for deterministic pseudo-random values @@ -738,14 +752,22 @@ impl NeuralRanker { /// Tokenize text fn tokenize(&self, text: &str) -> Result> { - let tokenizer = self.tokenizer.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire tokenizer read lock".into()))?; - let tokenizer = tokenizer.as_ref() + let tokenizer = self.tokenizer.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire tokenizer read lock".into()) + })?; + let tokenizer = tokenizer + .as_ref() .ok_or(VecStoreError::IndexNotInitialized)?; let words: Vec<&str> = text.split_whitespace().collect(); - let tokens: Vec = words.iter() - .map(|w| *tokenizer.vocab.get(*w).unwrap_or(&tokenizer.special_tokens.unk_token)) + let tokens: Vec = words + .iter() + .map(|w| { + *tokenizer + .vocab + .get(*w) + .unwrap_or(&tokenizer.special_tokens.unk_token) + }) .take(tokenizer.max_length) .collect(); @@ -778,7 +800,9 @@ impl NeuralRanker { /// Update running average latency fn update_latency(&self, latency_ms: f64) { - let Ok(mut avg) = self.stats.avg_latency_ms.write() else { return; }; + let Ok(mut avg) = self.stats.avg_latency_ms.write() else { + return; + }; let count = self.stats.total_reranks.load(Ordering::Relaxed) as f64; *avg = (*avg * (count - 1.0) + latency_ms) / count; } @@ -793,7 +817,11 @@ impl NeuralRanker { cache_hit_rate: { let hits = self.stats.cache_hits.load(Ordering::Relaxed) as f64; let misses = self.stats.cache_misses.load(Ordering::Relaxed) as f64; - if hits + misses > 0.0 { hits / (hits + misses) } else { 0.0 } + if hits + misses > 0.0 { + hits / (hits + misses) + } else { + 0.0 + } }, avg_latency_ms: self.stats.avg_latency_ms.read().map(|g| *g).unwrap_or(0.0), gpu_inferences: self.stats.gpu_inferences.load(Ordering::Relaxed), @@ -803,7 +831,9 @@ impl NeuralRanker { /// Clear embedding cache pub fn clear_cache(&self) { - let Ok(mut cache) = self.cache.write() else { return; }; + let Ok(mut cache) = self.cache.write() else { + return; + }; cache.entries.clear(); cache.hits = 0; cache.misses = 0; @@ -877,9 +907,7 @@ impl CrossEncoder { /// Batch score multiple pairs pub fn batch_score(&self, pairs: &[(String, String)]) -> Result> { - pairs.iter() - .map(|(q, d)| self.score(q, d)) - .collect() + pairs.iter().map(|(q, d)| self.score(q, d)).collect() } } @@ -896,24 +924,28 @@ impl LateInteractionModel { /// Encode a query pub fn encode_query(&self, query: &str) -> Result>> { - let tokens: Vec<&str> = query.split_whitespace() + let tokens: Vec<&str> = query + .split_whitespace() .take(self.config.query_max_tokens) .collect(); // Generate embeddings per token - Ok(tokens.iter() + Ok(tokens + .iter() .map(|_| vec![0.0; self.config.embedding_dim]) .collect()) } /// Encode a document pub fn encode_document(&self, doc: &str) -> Result>> { - let tokens: Vec<&str> = doc.split_whitespace() + let tokens: Vec<&str> = doc + .split_whitespace() .take(self.config.doc_max_tokens) .collect(); // Generate embeddings per token - let embeddings: Vec> = tokens.iter() + let embeddings: Vec> = tokens + .iter() .map(|_| vec![0.0; self.config.embedding_dim]) .collect(); @@ -944,7 +976,8 @@ impl LateInteractionModel { fn max_sim(&self, query: &[Vec], doc: &[Vec]) -> f32 { let mut total = 0.0; for q in query { - let max = doc.iter() + let max = doc + .iter() .map(|d| cosine_similarity(q, d)) .fold(f32::MIN, f32::max); total += max; @@ -957,21 +990,19 @@ impl LateInteractionModel { } fn avg_max(&self, query: &[Vec], doc: &[Vec]) -> f32 { - if query.is_empty() { return 0.0; } + if query.is_empty() { + return 0.0; + } self.max_sim(query, doc) / query.len() as f32 } fn soft_max(&self, query: &[Vec], doc: &[Vec]) -> f32 { let mut total = 0.0; for q in query { - let sims: Vec = doc.iter() - .map(|d| cosine_similarity(q, d)) - .collect(); + let sims: Vec = doc.iter().map(|d| cosine_similarity(q, d)).collect(); let max = sims.iter().fold(f32::MIN, |a, &b| a.max(b)); let exp_sum: f32 = sims.iter().map(|&s| (s - max).exp()).sum(); - let softmax_weighted: f32 = sims.iter() - .map(|&s| s * (s - max).exp() / exp_sum) - .sum(); + let softmax_weighted: f32 = sims.iter().map(|&s| s * (s - max).exp() / exp_sum).sum(); total += softmax_weighted; } total @@ -979,17 +1010,21 @@ impl LateInteractionModel { /// Index document embeddings pub fn index_document(&self, doc_id: &str, embeddings: Vec>) { - let Ok(mut doc_emb) = self.doc_embeddings.write() else { return; }; + let Ok(mut doc_emb) = self.doc_embeddings.write() else { + return; + }; doc_emb.insert(doc_id.to_string(), embeddings); } /// Retrieve and score against indexed documents pub fn retrieve(&self, query: &str, top_k: usize) -> Result> { let query_embeddings = self.encode_query(query)?; - let docs = self.doc_embeddings.read() - .map_err(|_| VecStoreError::LockError("Failed to acquire doc_embeddings read lock".into()))?; + let docs = self.doc_embeddings.read().map_err(|_| { + VecStoreError::LockError("Failed to acquire doc_embeddings read lock".into()) + })?; - let mut scores: Vec<(String, f32)> = docs.iter() + let mut scores: Vec<(String, f32)> = docs + .iter() .map(|(id, emb)| (id.clone(), self.score(&query_embeddings, emb))) .collect(); @@ -1013,7 +1048,9 @@ impl LearnedSparseModel { /// Encode text to sparse representation pub fn encode(&self, text: &str) -> Result { let tokens: Vec<&str> = text.split_whitespace().collect(); - let vocab = self.vocab.read() + let vocab = self + .vocab + .read() .map_err(|_| VecStoreError::LockError("Failed to acquire vocab read lock".into()))?; let mut term_ids = Vec::new(); @@ -1132,7 +1169,9 @@ impl EnsembleRanker { /// Score using ensemble pub fn score(&self, query: &str, document: &str) -> Result { - let scores: Vec = self.rankers.iter() + let scores: Vec = self + .rankers + .iter() .filter_map(|(r, _)| r.score(query, document).ok()) .collect(); @@ -1143,34 +1182,30 @@ impl EnsembleRanker { match &self.config.combination { CombinationMethod::WeightedSum => { let total_weight: f32 = self.rankers.iter().map(|(_, w)| w).sum(); - let weighted_sum: f32 = scores.iter() + let weighted_sum: f32 = scores + .iter() .zip(self.rankers.iter()) .map(|(s, (_, w))| s * w) .sum(); Ok(weighted_sum / total_weight) - } - CombinationMethod::Max => { - Ok(scores.iter().fold(f32::MIN, |a, &b| a.max(b))) - } - CombinationMethod::Min => { - Ok(scores.iter().fold(f32::MAX, |a, &b| a.min(b))) - } - CombinationMethod::Average => { - Ok(scores.iter().sum::() / scores.len() as f32) - } + }, + CombinationMethod::Max => Ok(scores.iter().fold(f32::MIN, |a, &b| a.max(b))), + CombinationMethod::Min => Ok(scores.iter().fold(f32::MAX, |a, &b| a.min(b))), + CombinationMethod::Average => Ok(scores.iter().sum::() / scores.len() as f32), CombinationMethod::RRF(k) => { // Reciprocal rank fusion - let rrf: f32 = scores.iter() + let rrf: f32 = scores + .iter() .enumerate() .map(|(rank, _)| 1.0 / (k + rank as f32 + 1.0)) .sum(); Ok(rrf) - } + }, CombinationMethod::CombMNZ => { let non_zero = scores.iter().filter(|&&s| s > 0.0).count() as f32; let sum: f32 = scores.iter().sum(); Ok(sum * non_zero) - } + }, } } } diff --git a/src/object_storage.rs b/src/object_storage.rs index dd5e5a5..95f8adc 100644 --- a/src/object_storage.rs +++ b/src/object_storage.rs @@ -50,7 +50,7 @@ //! let results = store.search(&query, 10).await?; //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::{HashMap, VecDeque}; use std::path::PathBuf; @@ -399,14 +399,11 @@ impl VectorChunk { /// Get vector by ID pub fn get(&self, id: &str) -> Option> { - self.vector_ids - .iter() - .position(|vid| vid == id) - .map(|idx| { - let start = idx * self.dimension; - let end = start + self.dimension; - self.vectors[start..end].to_vec() - }) + self.vector_ids.iter().position(|vid| vid == id).map(|idx| { + let start = idx * self.dimension; + let end = start + self.dimension; + self.vectors[start..end].to_vec() + }) } /// Serialize chunk @@ -421,8 +418,8 @@ impl VectorChunk { /// Compress chunk data pub fn compress(&self, level: u32) -> Result> { - use flate2::write::GzEncoder; use flate2::Compression; + use flate2::write::GzEncoder; use std::io::Write; let serialized = self.serialize()?; @@ -761,9 +758,7 @@ impl TieredStore { let to_demote: Vec = { let mut lru = self.hot_lru.write().await; let demote_count = (hot_bytes - hot_capacity_bytes) / (self.dimension * 4) + 1; - (0..demote_count) - .filter_map(|_| lru.pop_front()) - .collect() + (0..demote_count).filter_map(|_| lru.pop_front()).collect() }; // Demote vectors to warm tier @@ -823,9 +818,7 @@ impl TieredStore { let chunks_to_flush: Vec<(String, VectorChunk)> = { let warm = self.warm_chunks.read().await; - warm.iter() - .map(|(k, v)| (k.clone(), v.clone())) - .collect() + warm.iter().map(|(k, v)| (k.clone(), v.clone())).collect() }; for (chunk_id, chunk) in chunks_to_flush { diff --git a/src/object_store.rs b/src/object_store.rs index d14b446..45b734f 100644 --- a/src/object_store.rs +++ b/src/object_store.rs @@ -53,15 +53,9 @@ pub enum StorageProvider { region: String, }, /// Google Cloud Storage - GCS { - bucket: String, - prefix: String, - }, + GCS { bucket: String, prefix: String }, /// Azure Blob Storage - AzureBlob { - container: String, - prefix: String, - }, + AzureBlob { container: String, prefix: String }, /// MinIO (S3-compatible) MinIO { endpoint: String, @@ -69,9 +63,7 @@ pub enum StorageProvider { prefix: String, }, /// Local filesystem (for testing) - Local { - path: PathBuf, - }, + Local { path: PathBuf }, } /// Cache tier configuration @@ -95,7 +87,10 @@ pub enum CacheStorageType { /// SSD-backed cache SSD { path: PathBuf }, /// Hybrid memory + SSD - Hybrid { ssd_path: PathBuf, memory_ratio: f32 }, + Hybrid { + ssd_path: PathBuf, + memory_ratio: f32, + }, } /// Object store configuration @@ -128,14 +123,12 @@ impl ObjectStoreConfig { prefix: prefix.to_string(), region: "us-east-1".to_string(), }, - cache_tiers: vec![ - CacheTier { - name: "memory".to_string(), - max_size_bytes: 1024 * 1024 * 1024, // 1GB - ttl_seconds: 300, - storage_type: CacheStorageType::Memory, - }, - ], + cache_tiers: vec![CacheTier { + name: "memory".to_string(), + max_size_bytes: 1024 * 1024 * 1024, // 1GB + ttl_seconds: 300, + storage_type: CacheStorageType::Memory, + }], read_through_cache: true, write_through_cache: true, pre_warming: None, @@ -152,14 +145,12 @@ impl ObjectStoreConfig { bucket: bucket.to_string(), prefix: prefix.to_string(), }, - cache_tiers: vec![ - CacheTier { - name: "memory".to_string(), - max_size_bytes: 1024 * 1024 * 1024, - ttl_seconds: 300, - storage_type: CacheStorageType::Memory, - }, - ], + cache_tiers: vec![CacheTier { + name: "memory".to_string(), + max_size_bytes: 1024 * 1024 * 1024, + ttl_seconds: 300, + storage_type: CacheStorageType::Memory, + }], read_through_cache: true, write_through_cache: true, pre_warming: None, @@ -362,15 +353,17 @@ impl MemoryCache { fn evict_one(&mut self) { // LRU eviction - let oldest_key = self.entries + let oldest_key = self + .entries .iter() .min_by_key(|(_, e)| e.last_access) .map(|(k, _)| k.clone()); if let Some(key) = oldest_key - && let Some(entry) = self.entries.remove(&key) { - self.current_size = self.current_size.saturating_sub(entry.data.len() as u64); - } + && let Some(entry) = self.entries.remove(&key) + { + self.current_size = self.current_size.saturating_sub(entry.data.len() as u64); + } } fn invalidate(&mut self, key: &str) { @@ -381,7 +374,8 @@ impl MemoryCache { fn stats(&self) -> CacheStats { let now = Instant::now(); - let valid_entries: Vec<_> = self.entries + let valid_entries: Vec<_> = self + .entries .values() .filter(|e| now.duration_since(e.cached_at) < self.ttl) .collect(); @@ -484,13 +478,22 @@ impl AccessTracker { } fn most_recent(&self, n: usize) -> Vec<&str> { - self.recent.iter().rev().take(n).map(|s| s.as_str()).collect() + self.recent + .iter() + .rev() + .take(n) + .map(|s| s.as_str()) + .collect() } fn most_frequent(&self, n: usize) -> Vec<(&str, u64)> { let mut items: Vec<_> = self.frequency.iter().collect(); items.sort_by_key(|(_, count)| std::cmp::Reverse(*count)); - items.into_iter().take(n).map(|(k, v)| (k.as_str(), *v)).collect() + items + .into_iter() + .take(n) + .map(|(k, v)| (k.as_str(), *v)) + .collect() } } @@ -549,7 +552,12 @@ impl ObjectStoreBackend { } /// Put a vector - pub fn put(&self, id: &str, vector: &[f32], metadata: &HashMap) -> Result<()> { + pub fn put( + &self, + id: &str, + vector: &[f32], + metadata: &HashMap, + ) -> Result<()> { let start = Instant::now(); // Serialize @@ -644,10 +652,14 @@ impl ObjectStoreBackend { // Update cache if self.config.read_through_cache { let index_guard = self.metadata_index.read().map_err(|_| { - VecStoreError::LockError("Failed to acquire metadata index read lock".to_string()) + VecStoreError::LockError( + "Failed to acquire metadata index read lock".to_string(), + ) })?; - let meta = index_guard.get(id).cloned().unwrap_or_else(|| { - ObjectMetadata { + let meta = index_guard + .get(id) + .cloned() + .unwrap_or_else(|| ObjectMetadata { key: id.to_string(), size_bytes: data.len() as u64, created_at: unix_timestamp(), @@ -655,12 +667,13 @@ impl ObjectStoreBackend { access_count: 1, compression: self.config.compression.clone(), custom: HashMap::new(), - } - }); + }); drop(index_guard); let mut cache = self.memory_cache.write().map_err(|_| { - VecStoreError::LockError("Failed to acquire memory cache write lock".to_string()) + VecStoreError::LockError( + "Failed to acquire memory cache write lock".to_string(), + ) })?; cache.put(id.to_string(), data.clone(), meta); } @@ -709,7 +722,10 @@ impl ObjectStoreBackend { let mut success_count = 0; for vector in vectors { - if self.put(&vector.id, &vector.vector, &vector.metadata).is_ok() { + if self + .put(&vector.id, &vector.vector, &vector.metadata) + .is_ok() + { success_count += 1; } } @@ -735,26 +751,42 @@ impl ObjectStoreBackend { let ids_to_warm: Vec = match strategy { PreWarmStrategy::MostRecent => { let tracker = self.access_tracker.read().map_err(|_| { - VecStoreError::LockError("Failed to acquire access tracker read lock".to_string()) + VecStoreError::LockError( + "Failed to acquire access tracker read lock".to_string(), + ) })?; - tracker.most_recent(count).into_iter().map(|s| s.to_string()).collect() - } + tracker + .most_recent(count) + .into_iter() + .map(|s| s.to_string()) + .collect() + }, PreWarmStrategy::MostFrequent => { let tracker = self.access_tracker.read().map_err(|_| { - VecStoreError::LockError("Failed to acquire access tracker read lock".to_string()) + VecStoreError::LockError( + "Failed to acquire access tracker read lock".to_string(), + ) })?; - tracker.most_frequent(count).into_iter().map(|(s, _)| s.to_string()).collect() - } + tracker + .most_frequent(count) + .into_iter() + .map(|(s, _)| s.to_string()) + .collect() + }, PreWarmStrategy::Predictive => { // Use frequency as a simple predictor let tracker = self.access_tracker.read().map_err(|_| { - VecStoreError::LockError("Failed to acquire access tracker read lock".to_string()) + VecStoreError::LockError( + "Failed to acquire access tracker read lock".to_string(), + ) })?; - tracker.most_frequent(count).into_iter().map(|(s, _)| s.to_string()).collect() - } - PreWarmStrategy::Explicit { ids } => { - ids.iter().take(count).cloned().collect() - } + tracker + .most_frequent(count) + .into_iter() + .map(|(s, _)| s.to_string()) + .collect() + }, + PreWarmStrategy::Explicit { ids } => ids.iter().take(count).cloned().collect(), }; let mut warmed = 0; @@ -884,8 +916,8 @@ impl ObjectStoreBackend { // Internal methods fn serialize(&self, vector: &StoredVector) -> Result> { - let json = serde_json::to_vec(vector) - .map_err(|e| VecStoreError::Serialization(e.to_string()))?; + let json = + serde_json::to_vec(vector).map_err(|e| VecStoreError::Serialization(e.to_string()))?; // Apply compression match &self.config.compression { @@ -893,7 +925,7 @@ impl ObjectStoreBackend { CompressionType::Zstd { level: _ } => { // Simplified - in production would use actual zstd Ok(json) // Placeholder - } + }, CompressionType::Gzip { level: _ } => Ok(json), CompressionType::Lz4 => Ok(json), CompressionType::Snappy => Ok(json), @@ -914,55 +946,67 @@ impl ObjectStoreBackend { match &self.config.provider { StorageProvider::Local { path: _ } => { let mut storage = self.local_storage.write().map_err(|_| { - VecStoreError::LockError("Failed to acquire local storage write lock".to_string()) + VecStoreError::LockError( + "Failed to acquire local storage write lock".to_string(), + ) })?; storage.insert(key.to_string(), data.to_vec()); Ok(()) - } + }, StorageProvider::S3 { .. } => { // In production: use aws-sdk-s3 let mut storage = self.local_storage.write().map_err(|_| { - VecStoreError::LockError("Failed to acquire local storage write lock".to_string()) + VecStoreError::LockError( + "Failed to acquire local storage write lock".to_string(), + ) })?; storage.insert(key.to_string(), data.to_vec()); Ok(()) - } + }, StorageProvider::GCS { .. } => { let mut storage = self.local_storage.write().map_err(|_| { - VecStoreError::LockError("Failed to acquire local storage write lock".to_string()) + VecStoreError::LockError( + "Failed to acquire local storage write lock".to_string(), + ) })?; storage.insert(key.to_string(), data.to_vec()); Ok(()) - } + }, StorageProvider::AzureBlob { .. } => { let mut storage = self.local_storage.write().map_err(|_| { - VecStoreError::LockError("Failed to acquire local storage write lock".to_string()) + VecStoreError::LockError( + "Failed to acquire local storage write lock".to_string(), + ) })?; storage.insert(key.to_string(), data.to_vec()); Ok(()) - } + }, StorageProvider::MinIO { .. } => { let mut storage = self.local_storage.write().map_err(|_| { - VecStoreError::LockError("Failed to acquire local storage write lock".to_string()) + VecStoreError::LockError( + "Failed to acquire local storage write lock".to_string(), + ) })?; storage.insert(key.to_string(), data.to_vec()); Ok(()) - } + }, } } fn read_from_backend(&self, key: &str) -> Result>> { match &self.config.provider { - StorageProvider::Local { path: _ } | - StorageProvider::S3 { .. } | - StorageProvider::GCS { .. } | - StorageProvider::AzureBlob { .. } | - StorageProvider::MinIO { .. } => { + StorageProvider::Local { path: _ } + | StorageProvider::S3 { .. } + | StorageProvider::GCS { .. } + | StorageProvider::AzureBlob { .. } + | StorageProvider::MinIO { .. } => { let storage = self.local_storage.read().map_err(|_| { - VecStoreError::LockError("Failed to acquire local storage read lock".to_string()) + VecStoreError::LockError( + "Failed to acquire local storage read lock".to_string(), + ) })?; Ok(storage.get(key).cloned()) - } + }, } } @@ -1027,8 +1071,8 @@ pub struct TieringConfig { impl Default for TieringConfig { fn default() -> Self { Self { - hot_to_warm_seconds: 3600, // 1 hour - warm_to_cold_seconds: 86400, // 24 hours + hot_to_warm_seconds: 3600, // 1 hour + warm_to_cold_seconds: 86400, // 24 hours hot_access_threshold: 10, } } @@ -1059,7 +1103,12 @@ impl TieredStorageManager { } /// Put to appropriate tier - pub fn put(&self, id: &str, vector: &[f32], metadata: &HashMap) -> Result<()> { + pub fn put( + &self, + id: &str, + vector: &[f32], + metadata: &HashMap, + ) -> Result<()> { // New data goes to hot tier self.hot_tier.put(id, vector, metadata)?; @@ -1080,19 +1129,23 @@ impl TieredStorageManager { // Try warm if let Some(warm) = &self.warm_tier - && let Some(vector) = warm.get(id)? { - // Promote to hot on access - self.hot_tier.put(&vector.id, &vector.vector, &vector.metadata)?; - return Ok(Some(vector)); - } + && let Some(vector) = warm.get(id)? + { + // Promote to hot on access + self.hot_tier + .put(&vector.id, &vector.vector, &vector.metadata)?; + return Ok(Some(vector)); + } // Try cold if let Some(cold) = &self.cold_tier - && let Some(vector) = cold.get(id)? { - // Promote to hot on access - self.hot_tier.put(&vector.id, &vector.vector, &vector.metadata)?; - return Ok(Some(vector)); - } + && let Some(vector) = cold.get(id)? + { + // Promote to hot on access + self.hot_tier + .put(&vector.id, &vector.vector, &vector.metadata)?; + return Ok(Some(vector)); + } Ok(None) } @@ -1151,8 +1204,8 @@ mod tests { #[test] fn test_cache_behavior() { - let mut config = ObjectStoreConfig::local("/tmp/vecstore_cache_test") - .with_cache_tier(CacheTier { + let mut config = + ObjectStoreConfig::local("/tmp/vecstore_cache_test").with_cache_tier(CacheTier { name: "memory".to_string(), max_size_bytes: 1024 * 1024, ttl_seconds: 300, @@ -1212,7 +1265,9 @@ mod tests { // Add some data for i in 0..100 { - backend.put(&format!("vec_{}", i), &vec![0.1; 128], &HashMap::new()).unwrap(); + backend + .put(&format!("vec_{}", i), &vec![0.1; 128], &HashMap::new()) + .unwrap(); } let cost = backend.estimate_cost(); diff --git a/src/otel.rs b/src/otel.rs index 4ea4c64..b1e2279 100644 --- a/src/otel.rs +++ b/src/otel.rs @@ -27,7 +27,10 @@ //! ``` use std::collections::HashMap; -use std::sync::{Arc, RwLock, atomic::{AtomicU64, AtomicUsize, Ordering}}; +use std::sync::{ + Arc, RwLock, + atomic::{AtomicU64, AtomicUsize, Ordering}, +}; use std::time::{Duration, Instant, UNIX_EPOCH}; use serde::{Deserialize, Serialize}; @@ -119,9 +122,15 @@ impl OtelConfig { Self { service_name: "vecstore".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), - trace_exporter: ExporterType::OtlpGrpc { endpoint: endpoint.to_string() }, - metrics_exporter: ExporterType::OtlpGrpc { endpoint: endpoint.to_string() }, - logs_exporter: ExporterType::OtlpGrpc { endpoint: endpoint.to_string() }, + trace_exporter: ExporterType::OtlpGrpc { + endpoint: endpoint.to_string(), + }, + metrics_exporter: ExporterType::OtlpGrpc { + endpoint: endpoint.to_string(), + }, + logs_exporter: ExporterType::OtlpGrpc { + endpoint: endpoint.to_string(), + }, sampling: SamplingStrategy::Ratio(0.1), resource_attributes: HashMap::new(), batch_config: BatchConfig::default(), @@ -134,7 +143,9 @@ impl OtelConfig { Self { service_name: "vecstore".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), - trace_exporter: ExporterType::Jaeger { endpoint: endpoint.to_string() }, + trace_exporter: ExporterType::Jaeger { + endpoint: endpoint.to_string(), + }, metrics_exporter: ExporterType::Prometheus { port: 9090 }, logs_exporter: ExporterType::Console, sampling: SamplingStrategy::Ratio(0.1), @@ -149,7 +160,9 @@ impl OtelConfig { Self { service_name: "vecstore".to_string(), service_version: env!("CARGO_PKG_VERSION").to_string(), - trace_exporter: ExporterType::Zipkin { endpoint: endpoint.to_string() }, + trace_exporter: ExporterType::Zipkin { + endpoint: endpoint.to_string(), + }, metrics_exporter: ExporterType::Prometheus { port: 9090 }, logs_exporter: ExporterType::Console, sampling: SamplingStrategy::Ratio(0.1), @@ -182,7 +195,8 @@ impl OtelConfig { /// Add resource attribute pub fn with_attribute(mut self, key: &str, value: &str) -> Self { - self.resource_attributes.insert(key.to_string(), value.to_string()); + self.resource_attributes + .insert(key.to_string(), value.to_string()); self } @@ -304,7 +318,9 @@ impl Span { /// Set error status pub fn set_error(&mut self, message: &str) { - self.status = SpanStatus::Error { message: message.to_string() }; + self.status = SpanStatus::Error { + message: message.to_string(), + }; } /// Set OK status @@ -314,7 +330,9 @@ impl Span { /// Get duration pub fn duration(&self) -> Duration { - self.end_time.unwrap_or_else(Instant::now).duration_since(self.start_time) + self.end_time + .unwrap_or_else(Instant::now) + .duration_since(self.start_time) } } @@ -375,7 +393,9 @@ pub struct VecStoreMetrics { impl VecStoreMetrics { /// Create new metrics pub fn new() -> Self { - let boundaries = vec![1.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 5000.0]; + let boundaries = vec![ + 1.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 5000.0, + ]; let buckets: Vec = (0..=boundaries.len()).map(|_| AtomicU64::new(0)).collect(); Self { @@ -396,10 +416,13 @@ impl VecStoreMetrics { /// Record query pub fn record_query(&self, latency_ms: f64) { self.query_count.fetch_add(1, Ordering::Relaxed); - self.query_latency_sum.fetch_add((latency_ms * 1000.0) as u64, Ordering::Relaxed); + self.query_latency_sum + .fetch_add((latency_ms * 1000.0) as u64, Ordering::Relaxed); // Update histogram - let Ok(buckets) = self.histogram_buckets.read() else { return; }; + let Ok(buckets) = self.histogram_buckets.read() else { + return; + }; for (i, &boundary) in self.histogram_boundaries.iter().enumerate() { if latency_ms <= boundary { buckets[i].fetch_add(1, Ordering::Relaxed); @@ -419,7 +442,10 @@ impl VecStoreMetrics { /// Record delete pub fn record_delete(&self, count: usize) { self.delete_count.fetch_add(count as u64, Ordering::Relaxed); - self.vector_count.fetch_sub(count.min(self.vector_count.load(Ordering::Relaxed)), Ordering::Relaxed); + self.vector_count.fetch_sub( + count.min(self.vector_count.load(Ordering::Relaxed)), + Ordering::Relaxed, + ); } /// Record cache hit @@ -502,10 +528,14 @@ impl VecStoreMetrics { )); // Histogram - output.push_str("# HELP vecstore_query_latency_ms Query latency in milliseconds\n\ - # TYPE vecstore_query_latency_ms histogram\n"); + output.push_str( + "# HELP vecstore_query_latency_ms Query latency in milliseconds\n\ + # TYPE vecstore_query_latency_ms histogram\n", + ); - let Ok(buckets) = self.histogram_buckets.read() else { return output; }; + let Ok(buckets) = self.histogram_buckets.read() else { + return output; + }; let mut cumulative = 0u64; for (i, boundary) in self.histogram_boundaries.iter().enumerate() { cumulative += buckets[i].load(Ordering::Relaxed); @@ -515,11 +545,18 @@ impl VecStoreMetrics { )); } cumulative += buckets[self.histogram_boundaries.len()].load(Ordering::Relaxed); - output.push_str(&format!("vecstore_query_latency_ms_bucket{{le=\"+Inf\"}} {}\n", cumulative)); - output.push_str(&format!("vecstore_query_latency_ms_sum {}\n", - self.query_latency_sum.load(Ordering::Relaxed) as f64 / 1000.0)); - output.push_str(&format!("vecstore_query_latency_ms_count {}\n", - self.query_count.load(Ordering::Relaxed))); + output.push_str(&format!( + "vecstore_query_latency_ms_bucket{{le=\"+Inf\"}} {}\n", + cumulative + )); + output.push_str(&format!( + "vecstore_query_latency_ms_sum {}\n", + self.query_latency_sum.load(Ordering::Relaxed) as f64 / 1000.0 + )); + output.push_str(&format!( + "vecstore_query_latency_ms_count {}\n", + self.query_count.load(Ordering::Relaxed) + )); output } @@ -600,12 +637,16 @@ impl TelemetryProvider { let span_id = span.span_id.clone(); { - let Ok(mut current) = self.current_trace.write() else { return span_id; }; + let Ok(mut current) = self.current_trace.write() else { + return span_id; + }; *current = Some(trace_id.clone()); } { - let Ok(mut spans) = self.spans.write() else { return span_id; }; + let Ok(mut spans) = self.spans.write() else { + return span_id; + }; spans.insert(span_id.clone(), span); } @@ -614,7 +655,9 @@ impl TelemetryProvider { /// Start child span pub fn start_span(&self, name: &str, parent_id: Option<&str>) -> String { - let trace_id = self.current_trace.read() + let trace_id = self + .current_trace + .read() .ok() .and_then(|guard| guard.clone()) .unwrap_or_else(generate_trace_id); @@ -623,7 +666,9 @@ impl TelemetryProvider { let span_id = span.span_id.clone(); { - let Ok(mut spans) = self.spans.write() else { return span_id; }; + let Ok(mut spans) = self.spans.write() else { + return span_id; + }; spans.insert(span_id.clone(), span); } @@ -633,20 +678,26 @@ impl TelemetryProvider { /// End span pub fn end_span(&self, span_id: &str) { let span = { - let Ok(mut spans) = self.spans.write() else { return; }; + let Ok(mut spans) = self.spans.write() else { + return; + }; spans.remove(span_id) }; if let Some(mut span) = span { span.end(); - let Ok(mut completed) = self.completed_spans.write() else { return; }; + let Ok(mut completed) = self.completed_spans.write() else { + return; + }; completed.push(span); } } /// Set span attribute pub fn set_span_attribute(&self, span_id: &str, key: &str, value: SpanAttribute) { - let Ok(mut spans) = self.spans.write() else { return; }; + let Ok(mut spans) = self.spans.write() else { + return; + }; if let Some(span) = spans.get_mut(span_id) { span.set_attribute(key, value); } @@ -654,7 +705,9 @@ impl TelemetryProvider { /// Set span error pub fn set_span_error(&self, span_id: &str, message: &str) { - let Ok(mut spans) = self.spans.write() else { return; }; + let Ok(mut spans) = self.spans.write() else { + return; + }; if let Some(span) = spans.get_mut(span_id) { span.set_error(message); } @@ -666,15 +719,25 @@ impl TelemetryProvider { // Create a span for the query let span_id = self.start_span(operation, None); - self.set_span_attribute(&span_id, "db.system", SpanAttribute::String("vecstore".to_string())); - self.set_span_attribute(&span_id, "db.operation", SpanAttribute::String(operation.to_string())); + self.set_span_attribute( + &span_id, + "db.system", + SpanAttribute::String("vecstore".to_string()), + ); + self.set_span_attribute( + &span_id, + "db.operation", + SpanAttribute::String(operation.to_string()), + ); self.set_span_attribute(&span_id, "db.latency_ms", SpanAttribute::Float(latency_ms)); self.end_span(&span_id); } /// Flush completed spans pub fn flush(&self) -> Vec { - let Ok(mut completed) = self.completed_spans.write() else { return Vec::new(); }; + let Ok(mut completed) = self.completed_spans.write() else { + return Vec::new(); + }; std::mem::take(&mut *completed) } @@ -686,22 +749,27 @@ impl TelemetryProvider { OtlpExportData { service_name: self.config.service_name.clone(), service_version: self.config.service_version.clone(), - spans: spans.into_iter().map(|s| OtlpSpan { - trace_id: s.trace_id, - span_id: s.span_id, - parent_span_id: s.parent_id, - name: s.name, - start_time_unix_nano: 0, // Would be real timestamp - end_time_unix_nano: 0, - attributes: s.attributes.into_iter() - .map(|(k, v)| (k, format!("{:?}", v))) - .collect(), - status: match s.status { - SpanStatus::Ok => "OK".to_string(), - SpanStatus::Error { message } => format!("ERROR: {}", message), - SpanStatus::Unset => "UNSET".to_string(), - }, - }).collect(), + spans: spans + .into_iter() + .map(|s| OtlpSpan { + trace_id: s.trace_id, + span_id: s.span_id, + parent_span_id: s.parent_id, + name: s.name, + start_time_unix_nano: 0, // Would be real timestamp + end_time_unix_nano: 0, + attributes: s + .attributes + .into_iter() + .map(|(k, v)| (k, format!("{:?}", v))) + .collect(), + status: match s.status { + SpanStatus::Ok => "OK".to_string(), + SpanStatus::Error { message } => format!("ERROR: {}", message), + SpanStatus::Unset => "UNSET".to_string(), + }, + }) + .collect(), metrics, } } @@ -818,7 +886,10 @@ impl TraceContext { /// Convert to W3C traceparent header pub fn to_traceparent(&self) -> String { - format!("00-{}-{}-{:02x}", self.trace_id, self.span_id, self.trace_flags) + format!( + "00-{}-{}-{:02x}", + self.trace_id, self.span_id, self.trace_flags + ) } } diff --git a/src/partition_keys.rs b/src/partition_keys.rs index 15d06bb..8c3f6b8 100644 --- a/src/partition_keys.rs +++ b/src/partition_keys.rs @@ -32,11 +32,11 @@ //! let all_results = index.search_all_partitions(&query, 10)?; //! ``` +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::{Arc, RwLock}; -use serde::{Deserialize, Serialize}; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; // ============================================================================ // CONFIGURATION @@ -241,7 +241,11 @@ impl Partition { pub fn get(&self, id: &str) -> Option<(Vec, Option)> { let vectors = self.vectors.read().ok()?; vectors.get(id).map(|v| { - let meta = self.vector_metadata.read().ok().and_then(|m| m.get(id).cloned()); + let meta = self + .vector_metadata + .read() + .ok() + .and_then(|m| m.get(id).cloned()); (v.clone(), meta) }) } @@ -252,7 +256,8 @@ impl Partition { return Vec::new(); }; - let mut results: Vec<_> = vectors.iter() + let mut results: Vec<_> = vectors + .iter() .map(|(id, vec)| { let score = cosine_similarity(query, vec); (id.clone(), score) @@ -264,7 +269,8 @@ impl Partition { let partition_id = self.meta.read().map(|m| m.id).unwrap_or(0); - results.into_iter() + results + .into_iter() .map(|(id, score)| PartitionSearchResult { id, score, @@ -275,17 +281,20 @@ impl Partition { /// Get partition metadata pub fn meta(&self) -> PartitionMeta { - self.meta.read().map(|m| m.clone()).unwrap_or_else(|_| PartitionMeta { - id: 0, - key_value: String::new(), - state: PartitionState::Active, - created_at: 0, - last_accessed: 0, - vector_count: 0, - deleted_count: 0, - ttl_expiry: 0, - metadata: None, - }) + self.meta + .read() + .map(|m| m.clone()) + .unwrap_or_else(|_| PartitionMeta { + id: 0, + key_value: String::new(), + state: PartitionState::Active, + created_at: 0, + last_accessed: 0, + vector_count: 0, + deleted_count: 0, + ttl_expiry: 0, + metadata: None, + }) } /// Check if compaction is needed @@ -296,7 +305,8 @@ impl Partition { if meta.vector_count == 0 { return false; } - let delete_ratio = meta.deleted_count as f32 / (meta.vector_count + meta.deleted_count) as f32; + let delete_ratio = + meta.deleted_count as f32 / (meta.vector_count + meta.deleted_count) as f32; delete_ratio > threshold } @@ -381,18 +391,18 @@ impl PartitionedIndex { h ^= h >> 47; } h % self.config.max_partitions as u64 - } + }, HashAlgorithm::XXHash => { use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); key_value.hash(&mut hasher); hasher.finish() % self.config.max_partitions as u64 - } + }, HashAlgorithm::Modulo => { // For testing - just use first char key_value.bytes().next().unwrap_or(0) as u64 % self.config.max_partitions as u64 - } + }, } } @@ -441,23 +451,20 @@ impl PartitionedIndex { /// Extract partition key from metadata fn extract_partition_key(&self, metadata: &serde_json::Value) -> Option { - metadata.get(&self.config.partition_key) + metadata + .get(&self.config.partition_key) .and_then(|v| v.as_str()) .map(String::from) } /// Insert or update a vector - pub fn upsert( - &self, - id: &str, - vector: Vec, - metadata: serde_json::Value, - ) -> Result<()> { - let key_value = self.extract_partition_key(&metadata) - .ok_or_else(|| VecStoreError::InvalidInput(format!( + pub fn upsert(&self, id: &str, vector: Vec, metadata: serde_json::Value) -> Result<()> { + let key_value = self.extract_partition_key(&metadata).ok_or_else(|| { + VecStoreError::InvalidInput(format!( "Missing partition key: {}", self.config.partition_key - )))?; + )) + })?; let partition = self.get_or_create_partition(&key_value)?; partition.upsert(id, vector, Some(metadata))?; @@ -510,7 +517,7 @@ impl PartitionedIndex { ) -> Result> { if !self.config.allow_cross_partition { return Err(VecStoreError::InvalidInput( - "Cross-partition queries not allowed".to_string() + "Cross-partition queries not allowed".to_string(), )); } @@ -559,7 +566,10 @@ impl PartitionedIndex { /// List all partitions pub fn list_partitions(&self) -> Vec { - self.partitions.read().map(|p| p.values().map(|p| p.meta()).collect()).unwrap_or_default() + self.partitions + .read() + .map(|p| p.values().map(|p| p.meta()).collect()) + .unwrap_or_default() } /// Drop a partition @@ -668,12 +678,18 @@ impl PartitionRouter { /// Get hot partitions pub fn hot_partitions(&self) -> Vec { - self.hot_partitions.read().map(|h| h.clone()).unwrap_or_default() + self.hot_partitions + .read() + .map(|h| h.clone()) + .unwrap_or_default() } /// Get cold partitions (candidates for archival) pub fn cold_partitions(&self) -> Vec { - self.cold_partitions.read().map(|c| c.clone()).unwrap_or_default() + self.cold_partitions + .read() + .map(|c| c.clone()) + .unwrap_or_default() } /// Reset access counts (call periodically) @@ -721,8 +737,12 @@ mod tests { fn test_partition() { let partition = Partition::new(1, "tenant_a", 4); - partition.upsert("doc1", vec![1.0, 0.0, 0.0, 0.0], None).unwrap(); - partition.upsert("doc2", vec![0.0, 1.0, 0.0, 0.0], None).unwrap(); + partition + .upsert("doc1", vec![1.0, 0.0, 0.0, 0.0], None) + .unwrap(); + partition + .upsert("doc2", vec![0.0, 1.0, 0.0, 0.0], None) + .unwrap(); let results = partition.search(&[1.0, 0.0, 0.0, 0.0], 10); assert_eq!(results.len(), 2); @@ -738,13 +758,25 @@ mod tests { let index = PartitionedIndex::new(4, config).unwrap(); // Insert to different partitions - index.upsert("doc1", vec![1.0, 0.0, 0.0, 0.0], serde_json::json!({ - "tenant_id": "acme" - })).unwrap(); - - index.upsert("doc2", vec![0.0, 1.0, 0.0, 0.0], serde_json::json!({ - "tenant_id": "globex" - })).unwrap(); + index + .upsert( + "doc1", + vec![1.0, 0.0, 0.0, 0.0], + serde_json::json!({ + "tenant_id": "acme" + }), + ) + .unwrap(); + + index + .upsert( + "doc2", + vec![0.0, 1.0, 0.0, 0.0], + serde_json::json!({ + "tenant_id": "globex" + }), + ) + .unwrap(); // Search within partition let results = index.search(&[1.0, 0.0, 0.0, 0.0], 10, "acme").unwrap(); @@ -759,14 +791,19 @@ mod tests { #[test] fn test_partition_listing() { - let config = PartitionConfig::new() - .with_partition_key("tenant_id"); + let config = PartitionConfig::new().with_partition_key("tenant_id"); let index = PartitionedIndex::new(4, config).unwrap(); - index.upsert("doc1", vec![1.0; 4], serde_json::json!({"tenant_id": "a"})).unwrap(); - index.upsert("doc2", vec![1.0; 4], serde_json::json!({"tenant_id": "b"})).unwrap(); - index.upsert("doc3", vec![1.0; 4], serde_json::json!({"tenant_id": "c"})).unwrap(); + index + .upsert("doc1", vec![1.0; 4], serde_json::json!({"tenant_id": "a"})) + .unwrap(); + index + .upsert("doc2", vec![1.0; 4], serde_json::json!({"tenant_id": "b"})) + .unwrap(); + index + .upsert("doc3", vec![1.0; 4], serde_json::json!({"tenant_id": "c"})) + .unwrap(); let partitions = index.list_partitions(); assert_eq!(partitions.len(), 3); diff --git a/src/partitioning.rs b/src/partitioning.rs index c25780e..d7a98ee 100644 --- a/src/partitioning.rs +++ b/src/partitioning.rs @@ -152,11 +152,11 @@ impl PartitionedStore { self.partitions.insert(partition_id.clone(), store); self.partition_info.insert(partition_id, info); - } + }, Err(_) => { // Skip corrupted partitions continue; - } + }, } } } @@ -210,13 +210,14 @@ impl PartitionedStore { // Check partition size limit if let Some(max_size) = self.config.max_vectors_per_partition && let Some(info) = self.partition_info.get(partition_id) - && info.vector_count >= max_size { - return Err(anyhow::anyhow!( - "Partition '{}' has reached maximum size of {} vectors", - partition_id, - max_size - )); - } + && info.vector_count >= max_size + { + return Err(anyhow::anyhow!( + "Partition '{}' has reached maximum size of {} vectors", + partition_id, + max_size + )); + } let partition = self.get_or_create_partition(partition_id)?; partition.upsert(id, vector, metadata)?; diff --git a/src/payload_index.rs b/src/payload_index.rs index 7b408ea..504280e 100644 --- a/src/payload_index.rs +++ b/src/payload_index.rs @@ -38,12 +38,12 @@ //! let matching_ids = manager.query(&filter)?; //! ``` -use std::collections::{HashMap, HashSet, BTreeMap}; -use std::sync::RwLock; -use serde::{Deserialize, Serialize}; use ordered_float::OrderedFloat; +use serde::{Deserialize, Serialize}; +use std::collections::{BTreeMap, HashMap, HashSet}; +use std::sync::RwLock; -use crate::error::{VecStoreError, Result}; +use crate::error::{Result, VecStoreError}; // ============================================================================ // INDEX TYPES @@ -102,12 +102,13 @@ impl HashIndex { pub fn insert(&mut self, doc_id: &str, value: &str) { // Remove old value if exists if let Some(old_value) = self.reverse.remove(doc_id) - && let Some(docs) = self.index.get_mut(&old_value) { - docs.remove(doc_id); - if docs.is_empty() { - self.index.remove(&old_value); - } + && let Some(docs) = self.index.get_mut(&old_value) + { + docs.remove(doc_id); + if docs.is_empty() { + self.index.remove(&old_value); } + } // Insert new value self.index @@ -120,12 +121,13 @@ impl HashIndex { /// Remove a document pub fn remove(&mut self, doc_id: &str) { if let Some(value) = self.reverse.remove(doc_id) - && let Some(docs) = self.index.get_mut(&value) { - docs.remove(doc_id); - if docs.is_empty() { - self.index.remove(&value); - } + && let Some(docs) = self.index.get_mut(&value) + { + docs.remove(doc_id); + if docs.is_empty() { + self.index.remove(&value); } + } } /// Get documents with exact value @@ -186,12 +188,13 @@ impl BTreeIndex { // Remove old value if exists if let Some(old_value) = self.reverse.remove(doc_id) - && let Some(docs) = self.index.get_mut(&old_value) { - docs.remove(doc_id); - if docs.is_empty() { - self.index.remove(&old_value); - } + && let Some(docs) = self.index.get_mut(&old_value) + { + docs.remove(doc_id); + if docs.is_empty() { + self.index.remove(&old_value); } + } // Insert new value self.index @@ -204,12 +207,13 @@ impl BTreeIndex { /// Remove a document pub fn remove(&mut self, doc_id: &str) { if let Some(value) = self.reverse.remove(doc_id) - && let Some(docs) = self.index.get_mut(&value) { - docs.remove(doc_id); - if docs.is_empty() { - self.index.remove(&value); - } + && let Some(docs) = self.index.get_mut(&value) + { + docs.remove(doc_id); + if docs.is_empty() { + self.index.remove(&value); } + } } /// Get documents in range [min, max] @@ -291,8 +295,7 @@ impl GeoPoint { let dlat = (other.lat - self.lat).to_radians(); let dlon = (other.lon - self.lon).to_radians(); - let a = (dlat / 2.0).sin().powi(2) - + lat1.cos() * lat2.cos() * (dlon / 2.0).sin().powi(2); + let a = (dlat / 2.0).sin().powi(2) + lat1.cos() * lat2.cos() * (dlon / 2.0).sin().powi(2); let c = 2.0 * a.sqrt().asin(); r * c @@ -411,12 +414,12 @@ impl GeoIndex { for doc_id in docs { if let Some(point) = self.points.get(doc_id) && point.lat >= min_lat - && point.lat <= max_lat - && point.lon >= min_lon - && point.lon <= max_lon - { - results.insert(doc_id.clone()); - } + && point.lat <= max_lat + && point.lon >= min_lon + && point.lon <= max_lon + { + results.insert(doc_id.clone()); + } } } } @@ -585,22 +588,22 @@ impl PayloadIndexManager { self.hash_indexes .write()? .insert(field.to_string(), HashIndex::new(field)); - } + }, IndexType::BTree => { self.btree_indexes .write()? .insert(field.to_string(), BTreeIndex::new(field)); - } + }, IndexType::Geo => { self.geo_indexes .write()? .insert(field.to_string(), GeoIndex::new(field, 0.1)); - } + }, IndexType::FullText => { self.fulltext_indexes .write()? .insert(field.to_string(), FullTextIndex::new(field)); - } + }, } Ok(()) } @@ -629,9 +632,10 @@ impl PayloadIndexManager { for (field, value) in map { // Hash index if let Some(idx) = self.hash_indexes.write()?.get_mut(field) - && let Some(s) = value.as_str() { - idx.insert(doc_id, s); - } + && let Some(s) = value.as_str() + { + idx.insert(doc_id, s); + } // B-tree index if let Some(idx) = self.btree_indexes.write()?.get_mut(field) { @@ -645,18 +649,20 @@ impl PayloadIndexManager { // Geo index if let Some(idx) = self.geo_indexes.write()?.get_mut(field) && let serde_json::Value::Object(geo) = value - && let (Some(lat), Some(lon)) = ( - geo.get("lat").and_then(|v| v.as_f64()), - geo.get("lon").and_then(|v| v.as_f64()), - ) { - idx.insert(doc_id, GeoPoint::new(lat, lon)); - } + && let (Some(lat), Some(lon)) = ( + geo.get("lat").and_then(|v| v.as_f64()), + geo.get("lon").and_then(|v| v.as_f64()), + ) + { + idx.insert(doc_id, GeoPoint::new(lat, lon)); + } // Full-text index if let Some(idx) = self.fulltext_indexes.write()?.get_mut(field) - && let Some(s) = value.as_str() { - idx.insert(doc_id, s); - } + && let Some(s) = value.as_str() + { + idx.insert(doc_id, s); + } } } Ok(()) @@ -703,7 +709,7 @@ impl PayloadIndexManager { field ))) } - } + }, Filter::In { field, values } => { let guard = self.hash_indexes.read()?; if let Some(idx) = guard.get(field) { @@ -714,7 +720,7 @@ impl PayloadIndexManager { field ))) } - } + }, Filter::Range { field, min, max } => { let guard = self.btree_indexes.read()?; if let Some(idx) = guard.get(field) { @@ -725,7 +731,7 @@ impl PayloadIndexManager { field ))) } - } + }, Filter::Gt { field, value } => { let guard = self.btree_indexes.read()?; if let Some(idx) = guard.get(field) { @@ -736,7 +742,7 @@ impl PayloadIndexManager { field ))) } - } + }, Filter::Lt { field, value } => { let guard = self.btree_indexes.read()?; if let Some(idx) = guard.get(field) { @@ -747,8 +753,12 @@ impl PayloadIndexManager { field ))) } - } - Filter::GeoRadius { field, center, radius_km } => { + }, + Filter::GeoRadius { + field, + center, + radius_km, + } => { let guard = self.geo_indexes.read()?; if let Some(idx) = guard.get(field) { let results = idx.within_radius(center, *radius_km); @@ -759,7 +769,7 @@ impl PayloadIndexManager { field ))) } - } + }, Filter::FullText { field, query } => { let guard = self.fulltext_indexes.read()?; if let Some(idx) = guard.get(field) { @@ -770,7 +780,7 @@ impl PayloadIndexManager { field ))) } - } + }, Filter::And(filters) => { let mut result: Option> = None; for f in filters { @@ -781,21 +791,21 @@ impl PayloadIndexManager { }); } Ok(result.unwrap_or_default()) - } + }, Filter::Or(filters) => { let mut result = HashSet::new(); for f in filters { result.extend(self.evaluate_filter(f)?); } Ok(result) - } + }, Filter::Not(_inner) => { // Would need all document IDs to compute NOT // For now, return empty (would need enhancement) Err(VecStoreError::InvalidInput( "NOT filter requires full document list".to_string(), )) - } + }, } } @@ -852,7 +862,11 @@ pub enum Filter { /// Less than Lt { field: String, value: f64 }, /// Geo radius query - GeoRadius { field: String, center: GeoPoint, radius_km: f64 }, + GeoRadius { + field: String, + center: GeoPoint, + radius_km: f64, + }, /// Full-text search FullText { field: String, query: String }, /// AND of filters @@ -996,10 +1010,15 @@ mod tests { manager.create_index("category", IndexType::Hash).unwrap(); manager.create_index("price", IndexType::BTree).unwrap(); - manager.index_payload("doc1", &serde_json::json!({ - "category": "electronics", - "price": 299.99 - })).unwrap(); + manager + .index_payload( + "doc1", + &serde_json::json!({ + "category": "electronics", + "price": 299.99 + }), + ) + .unwrap(); let filter = Filter::and(vec![ Filter::eq("category", "electronics"), diff --git a/src/pq.rs b/src/pq.rs index cda8cf4..bf7acb7 100644 --- a/src/pq.rs +++ b/src/pq.rs @@ -49,9 +49,9 @@ //! let distance = pq.asymmetric_distance(&query, &codes)?; //! ``` -use anyhow::{anyhow, Result}; -use rand::seq::SliceRandom; +use anyhow::{Result, anyhow}; use rand::Rng; +use rand::seq::SliceRandom; use serde::{Deserialize, Serialize}; // ============================================================================ @@ -296,7 +296,11 @@ impl ProductQuantizer { .iter() .map(|v| { let mut recon = vec![0.0; d]; - for (m, codebook) in temp_codebooks.iter().enumerate().take(self.config.num_subspaces) { + for (m, codebook) in temp_codebooks + .iter() + .enumerate() + .take(self.config.num_subspaces) + { let start_dim = m * self.subspace_dim; let end_dim = start_dim + self.subspace_dim; let subvec = &v[start_dim..end_dim]; @@ -499,7 +503,8 @@ impl ProductQuantizer { /// Find nearest centroid index #[inline] fn find_nearest_centroid(&self, vector: &[f32], centroids: &[Vec]) -> usize { - self.find_nearest_centroid_with_distance(vector, centroids).0 + self.find_nearest_centroid_with_distance(vector, centroids) + .0 } /// Find nearest centroid with distance @@ -639,11 +644,11 @@ impl ProductQuantizer { match self.config.metric { PQMetric::L2 => { distance += self.l2_distance_squared(query_subvec, centroid); - } + }, PQMetric::InnerProduct => { let ip: f32 = query_subvec.iter().zip(centroid).map(|(a, b)| a * b).sum(); distance -= ip; // Negate for "lower is better" - } + }, } } @@ -680,9 +685,11 @@ impl ProductQuantizer { .iter() .map(|centroid| match self.config.metric { PQMetric::L2 => self.l2_distance_squared(query_subvec, centroid), - PQMetric::InnerProduct => { - -query_subvec.iter().zip(centroid).map(|(a, b)| a * b).sum::() - } + PQMetric::InnerProduct => -query_subvec + .iter() + .zip(centroid) + .map(|(a, b)| a * b) + .sum::(), }) .collect(); diff --git a/src/privacy.rs b/src/privacy.rs index c6eadf4..41477fe 100644 --- a/src/privacy.rs +++ b/src/privacy.rs @@ -31,10 +31,10 @@ //! let private_query = engine.privatize_query(&query_vector)?; //! ``` -use std::collections::HashMap; -use std::sync::{Arc, RwLock}; use rand_distr::{Distribution, Normal}; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; use crate::error::VecStoreError; @@ -63,8 +63,7 @@ impl Distribution for Laplace { } /// Differential privacy mechanism -#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)] -#[derive(Default)] +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Default)] pub enum DPMechanism { /// Laplace mechanism (pure DP) Laplace, @@ -77,7 +76,6 @@ pub enum DPMechanism { RandomizedResponse, } - /// Privacy configuration #[derive(Debug, Clone, Serialize, Deserialize)] pub struct PrivacyConfig { @@ -185,8 +183,14 @@ impl PrivacyBudget { } /// Spend budget for a specific entity - pub fn spend_entity(&mut self, entity_id: &str, epsilon: f64, delta: f64) -> Result<(), VecStoreError> { - let entry = self.entity_budgets + pub fn spend_entity( + &mut self, + entity_id: &str, + epsilon: f64, + delta: f64, + ) -> Result<(), VecStoreError> { + let entry = self + .entity_budgets .entry(entity_id.to_string()) .or_insert_with(|| EntityBudget { epsilon_spent: 0.0, @@ -252,9 +256,10 @@ impl PrivacyEngine { pub fn privatize(&self, vector: &[f32]) -> Result, VecStoreError> { // Track budget if enabled if self.config.track_budget { - let mut budget = self.budget.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let mut budget = self + .budget + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; budget.spend(self.config.epsilon, self.config.delta)?; } @@ -272,9 +277,9 @@ impl PrivacyEngine { DPMechanism::DiscreteLaplace => self.add_discrete_laplace_noise(&clipped)?, DPMechanism::RandomizedResponse => { return Err(VecStoreError::InvalidInput( - "RandomizedResponse not applicable to continuous vectors".into() + "RandomizedResponse not applicable to continuous vectors".into(), )); - } + }, }; Ok(noisy) @@ -286,9 +291,10 @@ impl PrivacyEngine { let query_epsilon = self.config.epsilon * 0.1; // Use 10% of vector budget if self.config.track_budget { - let mut budget = self.budget.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let mut budget = self + .budget + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; budget.spend(query_epsilon, self.config.delta * 0.1)?; } @@ -311,9 +317,10 @@ impl PrivacyEngine { let batch_epsilon = self.config.epsilon * n.sqrt(); // Sublinear in batch size if self.config.track_budget { - let mut budget = self.budget.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let mut budget = self + .budget + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; budget.spend(batch_epsilon, self.config.delta * n)?; } @@ -374,7 +381,8 @@ impl PrivacyEngine { // Round to reduce precision (increases k-anonymity) let precision = 1.0 / k_anonymity as f32; - let anonymized: Vec = noisy.iter() + let anonymized: Vec = noisy + .iter() .map(|&v| (v / precision).round() * precision) .collect(); @@ -388,11 +396,11 @@ impl PrivacyEngine { // Gaussian mechanism: sigma = sensitivity * sqrt(2 * ln(1.25/delta)) / epsilon let c = (2.0 * (1.25 / self.config.delta).ln()).sqrt(); self.config.sensitivity * c / self.config.epsilon - } + }, DPMechanism::Laplace => { // Laplace mechanism: scale = sensitivity / epsilon self.config.sensitivity / self.config.epsilon - } + }, _ => self.config.sensitivity / self.config.epsilon, } } @@ -410,7 +418,8 @@ impl PrivacyEngine { }; // Signal magnitude - let signal_mag: f64 = sample_vector.iter() + let signal_mag: f64 = sample_vector + .iter() .map(|&v| (v as f64) * (v as f64)) .sum::() .sqrt(); @@ -437,27 +446,30 @@ impl PrivacyEngine { /// Get remaining privacy budget pub fn remaining_budget(&self) -> Result<(f64, f64), VecStoreError> { - let budget = self.budget.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let budget = self + .budget + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; Ok(budget.remaining()) } /// Check if privacy budget is exhausted pub fn is_budget_exhausted(&self) -> Result { - let budget = self.budget.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let budget = self + .budget + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; Ok(budget.is_exhausted()) } /// Reset privacy budget pub fn reset_budget(&self) -> Result<(), VecStoreError> { - let mut budget = self.budget.write().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let mut budget = self + .budget + .write() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; *budget = PrivacyBudget::new(self.config.max_budget, self.config.delta * 100.0); Ok(()) @@ -465,9 +477,10 @@ impl PrivacyEngine { /// Get privacy budget snapshot pub fn get_budget_snapshot(&self) -> Result { - let budget = self.budget.read().map_err(|_| { - VecStoreError::Internal("Failed to acquire budget lock".into()) - })?; + let budget = self + .budget + .read() + .map_err(|_| VecStoreError::Internal("Failed to acquire budget lock".into()))?; Ok(budget.clone()) } @@ -479,13 +492,17 @@ impl PrivacyEngine { self.add_gaussian_noise_with_scale(vector, scale) } - fn add_gaussian_noise_with_scale(&self, vector: &[f32], scale: f64) -> Result, VecStoreError> { - let normal = Normal::new(0.0, scale).map_err(|e| { - VecStoreError::InvalidInput(format!("Invalid noise scale: {}", e)) - })?; + fn add_gaussian_noise_with_scale( + &self, + vector: &[f32], + scale: f64, + ) -> Result, VecStoreError> { + let normal = Normal::new(0.0, scale) + .map_err(|e| VecStoreError::InvalidInput(format!("Invalid noise scale: {}", e)))?; let mut rng = rand::rng(); - let noisy: Vec = vector.iter() + let noisy: Vec = vector + .iter() .map(|&v| v + normal.sample(&mut rng) as f32) .collect(); @@ -494,12 +511,12 @@ impl PrivacyEngine { fn add_laplace_noise(&self, vector: &[f32]) -> Result, VecStoreError> { let scale = self.config.sensitivity / self.config.epsilon; - let laplace = Laplace::new(0.0, scale).map_err(|e| { - VecStoreError::InvalidInput(format!("Invalid noise scale: {}", e)) - })?; + let laplace = Laplace::new(0.0, scale) + .map_err(|e| VecStoreError::InvalidInput(format!("Invalid noise scale: {}", e)))?; let mut rng = rand::rng(); - let noisy: Vec = vector.iter() + let noisy: Vec = vector + .iter() .map(|&v| v + laplace.sample(&mut rng) as f32) .collect(); @@ -509,12 +526,12 @@ impl PrivacyEngine { fn add_discrete_laplace_noise(&self, vector: &[f32]) -> Result, VecStoreError> { // Discretized Laplace for integer-like values let scale = self.config.sensitivity / self.config.epsilon; - let laplace = Laplace::new(0.0, scale).map_err(|e| { - VecStoreError::InvalidInput(format!("Invalid noise scale: {}", e)) - })?; + let laplace = Laplace::new(0.0, scale) + .map_err(|e| VecStoreError::InvalidInput(format!("Invalid noise scale: {}", e)))?; let mut rng = rand::rng(); - let noisy: Vec = vector.iter() + let noisy: Vec = vector + .iter() .map(|&v| { let noise = laplace.sample(&mut rng).round() as f32; v + noise @@ -527,7 +544,8 @@ impl PrivacyEngine { /// Clip vector to have at most the given L2 norm fn clip_vector(vector: &[f32], max_norm: f64) -> Vec { - let current_norm: f64 = vector.iter() + let current_norm: f64 = vector + .iter() .map(|&v| (v as f64) * (v as f64)) .sum::() .sqrt(); @@ -672,7 +690,8 @@ mod tests { assert_eq!(private.len(), original.len()); // Should be different (with high probability) - let diff: f32 = original.iter() + let diff: f32 = original + .iter() .zip(private.iter()) .map(|(a, b)| (a - b).abs()) .sum(); @@ -685,7 +704,8 @@ mod tests { let clipped = clip_vector(&large_vec, 1.0); // Should have norm <= 1 - let norm: f64 = clipped.iter() + let norm: f64 = clipped + .iter() .map(|&v| (v as f64) * (v as f64)) .sum::() .sqrt(); @@ -828,7 +848,8 @@ mod tests { // Should be perturbed assert_eq!(private_query.len(), 3); - let diff: f32 = query.iter() + let diff: f32 = query + .iter() .zip(private_query.iter()) .map(|(a, b)| (a - b).abs()) .sum(); @@ -840,11 +861,7 @@ mod tests { let config = PrivacyConfig::default(); let engine = PrivacyEngine::new(config); - let vectors = vec![ - vec![1.0, 0.0], - vec![0.0, 1.0], - vec![0.5, 0.5], - ]; + let vectors = vec![vec![1.0, 0.0], vec![0.0, 1.0], vec![0.5, 0.5]]; let private = engine.privatize_batch(&vectors).unwrap(); diff --git a/src/profiler.rs b/src/profiler.rs index 70f2d4a..564a74c 100644 --- a/src/profiler.rs +++ b/src/profiler.rs @@ -48,7 +48,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::time::Instant; diff --git a/src/protocol.rs b/src/protocol.rs index 152dbdd..1a2580a 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -42,7 +42,7 @@ //! # } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::collections::HashMap; @@ -214,7 +214,7 @@ impl ProtocolAdapter { Protocol::Milvus => self.parse_milvus(json), Protocol::Universal => { serde_json::from_str(json).map_err(|e| anyhow!("Invalid JSON: {}", e)) - } + }, } } @@ -482,7 +482,7 @@ impl ProtocolAdapter { Ok(UniversalResponse::Upsert { upserted_count: count, }) - } + }, UniversalRequest::Query { vector, @@ -521,7 +521,7 @@ impl ProtocolAdapter { .collect(); Ok(UniversalResponse::Query { matches }) - } + }, UniversalRequest::Delete { ids } => { let mut count = 0; @@ -533,7 +533,7 @@ impl ProtocolAdapter { Ok(UniversalResponse::Delete { deleted_count: count, }) - } + }, UniversalRequest::Fetch { ids } => { let mut vectors = HashMap::new(); @@ -552,7 +552,7 @@ impl ProtocolAdapter { } Ok(UniversalResponse::Fetch { vectors }) - } + }, } } @@ -564,7 +564,7 @@ impl ProtocolAdapter { // Universal and other protocols fall back to the native JSON format Protocol::Universal | Protocol::Weaviate | Protocol::ChromaDB | Protocol::Milvus => { serde_json::to_string(&response).map_err(|e| anyhow!("Serialization error: {}", e)) - } + }, } } @@ -575,18 +575,18 @@ impl ProtocolAdapter { serde_json::json!({ "upsertedCount": upserted_count }) - } + }, UniversalResponse::Query { matches } => { serde_json::json!({ "matches": matches, "namespace": "" }) - } + }, UniversalResponse::Delete { deleted_count } => { serde_json::json!({ "deletedCount": deleted_count }) - } + }, _ => serde_json::json!(response), }; @@ -605,7 +605,7 @@ impl ProtocolAdapter { "status": "ok", "time": 0.0 }) - } + }, UniversalResponse::Query { matches } => { let results: Vec = matches .into_iter() @@ -623,7 +623,7 @@ impl ProtocolAdapter { "status": "ok", "time": 0.0 }) - } + }, _ => serde_json::json!(response), }; @@ -688,7 +688,7 @@ impl ProtocolAdapter { conditions.push(FilterExpr::And(sub_filters)); } } - } + }, "$or" => { if let Some(arr) = value.as_array() { let sub_filters: Vec = arr @@ -707,7 +707,7 @@ impl ProtocolAdapter { conditions.push(FilterExpr::Or(sub_filters)); } } - } + }, field => { // Check if value is an operator object if let Some(obj) = value.as_object() { @@ -740,7 +740,7 @@ impl ProtocolAdapter { value: value.clone(), }); } - } + }, } } diff --git a/src/python.rs b/src/python.rs index 9d3a15c..e3e1c0e 100644 --- a/src/python.rs +++ b/src/python.rs @@ -6,12 +6,12 @@ // This module provides Python-friendly wrappers around the Rust API. use crate::collection::{Collection, VecDatabase}; -use crate::store::{parse_filter, FilterExpr, HybridQuery, Metadata, Query, VecStore}; +use crate::store::{FilterExpr, HybridQuery, Metadata, Query, VecStore, parse_filter}; use crate::text_splitter::{RecursiveCharacterTextSplitter, TextSplitter}; +use pyo3::IntoPyObjectExt; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; use pyo3::types::{PyDict, PyList}; -use pyo3::IntoPyObjectExt; use std::collections::HashMap; use std::path::PathBuf; @@ -108,21 +108,21 @@ fn metadata_to_pydict(py: Python, metadata: &Metadata) -> PyResult> { match value { serde_json::Value::String(s) => { dict.set_item(key, s.as_str())?; - } + }, serde_json::Value::Number(n) => { if let Some(i) = n.as_i64() { dict.set_item(key, i)?; } else if let Some(f) = n.as_f64() { dict.set_item(key, f)?; } - } + }, serde_json::Value::Bool(b) => { dict.set_item(key, b)?; - } + }, serde_json::Value::Null => { dict.set_item(key, py.None())?; - } - _ => {} // Skip complex types + }, + _ => {}, // Skip complex types } } Ok(dict.unbind()) @@ -803,7 +803,10 @@ impl PyLangChainVectorStore { for (i, (text, embedding)) in texts.into_iter().zip(embeddings).enumerate() { let id = if let Some(ref id_list) = ids { - id_list.get(i).cloned().unwrap_or_else(|| format!("doc_{}", i)) + id_list + .get(i) + .cloned() + .unwrap_or_else(|| format!("doc_{}", i)) } else { format!("doc_{}", i) }; @@ -1103,7 +1106,7 @@ fn json_to_pyobject(py: Python, value: &serde_json::Value) -> PyResult } else { Ok(py.None()) } - } + }, serde_json::Value::String(s) => s.into_py_any(py), serde_json::Value::Array(arr) => { let list = PyList::new( @@ -1113,14 +1116,14 @@ fn json_to_pyobject(py: Python, value: &serde_json::Value) -> PyResult .collect::>>()?, )?; Ok(list.unbind().into_any()) - } + }, serde_json::Value::Object(obj) => { let dict = PyDict::new(py); for (k, v) in obj { dict.set_item(k, json_to_pyobject(py, v)?)?; } Ok(dict.unbind().into_any()) - } + }, } } diff --git a/src/python_async.rs b/src/python_async.rs index 2a47e63..9c3c013 100644 --- a/src/python_async.rs +++ b/src/python_async.rs @@ -31,7 +31,7 @@ //! } //! ``` -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Arc; diff --git a/src/quantization.rs b/src/quantization.rs index 31ac25f..781d7f9 100644 --- a/src/quantization.rs +++ b/src/quantization.rs @@ -17,7 +17,7 @@ //! | SQ4 | 8x | 95-97% | 3-4x | //! | BQ | 32x | 85-95% | 4-8x | -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use serde::{Deserialize, Serialize}; use std::collections::HashMap; @@ -689,7 +689,8 @@ impl ScalarQuantizer2 { let decoded_a = self.decode(a).unwrap(); let decoded_b = self.decode(b).unwrap(); - decoded_a.iter() + decoded_a + .iter() .zip(&decoded_b) .map(|(x, y)| (x - y).powi(2)) .sum::() @@ -740,9 +741,7 @@ impl TernaryQuantizer { let threshold = abs_values[threshold_idx]; // Compute scale as mean of non-zero absolute values - let non_zero: Vec = abs_values.into_iter() - .filter(|&v| v > threshold) - .collect(); + let non_zero: Vec = abs_values.into_iter().filter(|&v| v > threshold).collect(); let scale = if non_zero.is_empty() { 1.0 } else { @@ -825,8 +824,16 @@ impl TernaryQuantizer { let tb = (encoded_b[byte_idx] >> bit_offset) & 0b11; // Convert to signed: 00->0, 01->+1, 10->-1 - let va = match ta { 0b01 => 1, 0b10 => -1, _ => 0 }; - let vb = match tb { 0b01 => 1, 0b10 => -1, _ => 0 }; + let va = match ta { + 0b01 => 1, + 0b10 => -1, + _ => 0, + }; + let vb = match tb { + 0b01 => 1, + 0b10 => -1, + _ => 0, + }; sum += (va * vb) as f32 * self.scales[i] * self.scales[i]; } @@ -836,8 +843,7 @@ impl TernaryQuantizer { /// Memory footprint pub fn memory_usage(&self, num_vectors: usize) -> usize { - num_vectors * self.dimension.div_ceil(4) - + self.dimension * 8 // threshold + scale per dimension + num_vectors * self.dimension.div_ceil(4) + self.dimension * 8 // threshold + scale per dimension } } @@ -894,7 +900,8 @@ impl AsymmetricQuantizer { for (id, score) in &mut candidates { if let Some(full_vec) = full_vecs.get(id) { // Compute exact distance - *score = query.iter() + *score = query + .iter() .zip(full_vec) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -920,7 +927,8 @@ impl Quantizer for ScalarQuantizer8 { #[inline] fn asymmetric_distance(&self, query: &[f32], encoded: &[u8]) -> f32 { let decoded = self.decode(encoded).unwrap(); - query.iter() + query + .iter() .zip(&decoded) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -942,7 +950,8 @@ impl Quantizer for ScalarQuantizer4 { #[inline] fn asymmetric_distance(&self, query: &[f32], encoded: &[u8]) -> f32 { let decoded = self.decode(encoded).unwrap(); - query.iter() + query + .iter() .zip(&decoded) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -964,7 +973,8 @@ impl Quantizer for ScalarQuantizer2 { #[inline] fn asymmetric_distance(&self, query: &[f32], encoded: &[u8]) -> f32 { let decoded = self.decode(encoded).unwrap(); - query.iter() + query + .iter() .zip(&decoded) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -1022,7 +1032,8 @@ impl Quantizer for TernaryQuantizer { #[inline] fn asymmetric_distance(&self, query: &[f32], encoded: &[u8]) -> f32 { let decoded = self.decode(encoded).unwrap(); - query.iter() + query + .iter() .zip(&decoded) .map(|(a, b)| (a - b).powi(2)) .sum::() @@ -1055,7 +1066,11 @@ pub struct QuantizationComparison { } /// Run comparison of all quantization methods -pub fn compare_quantizers(vectors: &[Vec], _test_queries: &[Vec], _k: usize) -> Vec { +pub fn compare_quantizers( + vectors: &[Vec], + _test_queries: &[Vec], + _k: usize, +) -> Vec { let mut results = Vec::new(); // SQ8 @@ -1124,10 +1139,11 @@ fn compute_reconstruction_mse(sq8: &ScalarQuantizer8, vectors: &[Vec]) -> f let mut total_mse = 0.0; for v in vectors { if let Ok(encoded) = sq8.encode(v) - && let Ok(decoded) = sq8.decode(&encoded) { - let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); - total_mse += mse / v.len() as f32; - } + && let Ok(decoded) = sq8.decode(&encoded) + { + let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); + total_mse += mse / v.len() as f32; + } } total_mse / vectors.len() as f32 } @@ -1136,10 +1152,11 @@ fn compute_reconstruction_mse_sq4(sq4: &ScalarQuantizer4, vectors: &[Vec]) let mut total_mse = 0.0; for v in vectors { if let Ok(encoded) = sq4.encode(v) - && let Ok(decoded) = sq4.decode(&encoded) { - let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); - total_mse += mse / v.len() as f32; - } + && let Ok(decoded) = sq4.decode(&encoded) + { + let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); + total_mse += mse / v.len() as f32; + } } total_mse / vectors.len() as f32 } @@ -1148,10 +1165,11 @@ fn compute_reconstruction_mse_sq2(sq2: &ScalarQuantizer2, vectors: &[Vec]) let mut total_mse = 0.0; for v in vectors { if let Ok(encoded) = sq2.encode(v) - && let Ok(decoded) = sq2.decode(&encoded) { - let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); - total_mse += mse / v.len() as f32; - } + && let Ok(decoded) = sq2.decode(&encoded) + { + let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); + total_mse += mse / v.len() as f32; + } } total_mse / vectors.len() as f32 } @@ -1160,10 +1178,11 @@ fn compute_reconstruction_mse_ternary(ternary: &TernaryQuantizer, vectors: &[Vec let mut total_mse = 0.0; for v in vectors { if let Ok(encoded) = ternary.encode(v) - && let Ok(decoded) = ternary.decode(&encoded) { - let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); - total_mse += mse / v.len() as f32; - } + && let Ok(decoded) = ternary.decode(&encoded) + { + let mse: f32 = v.iter().zip(&decoded).map(|(a, b)| (a - b).powi(2)).sum(); + total_mse += mse / v.len() as f32; + } } total_mse / vectors.len() as f32 } @@ -1215,7 +1234,9 @@ mod ultra_low_bit_tests { for (i, &val) in decoded.iter().enumerate() { assert!( val == 0.0 || val.abs() - quantizer.scales[i].abs() < 0.001, - "Unexpected value {} at dim {}", val, i + "Unexpected value {} at dim {}", + val, + i ); } } @@ -1263,12 +1284,9 @@ mod ultra_low_bit_tests { .collect(); // Search - let results = asymmetric.search( - &vectors[0], - &encoded_db, - 5, - Some(&full_vectors), - ).unwrap(); + let results = asymmetric + .search(&vectors[0], &encoded_db, 5, Some(&full_vectors)) + .unwrap(); assert_eq!(results.len(), 5); // First result should be the query itself @@ -1286,8 +1304,10 @@ mod ultra_low_bit_tests { // Verify compression ratios increase for result in &comparison { - println!("{}: {:.1}x compression, {:.3} MSE", - result.method, result.compression_ratio, result.mse); + println!( + "{}: {:.1}x compression, {:.3} MSE", + result.method, result.compression_ratio, result.mse + ); } } } diff --git a/src/query_analytics.rs b/src/query_analytics.rs index d057ada..c7a4c16 100644 --- a/src/query_analytics.rs +++ b/src/query_analytics.rs @@ -1,9 +1,9 @@ // Query Analytics Dashboard - Comprehensive analytics for vector search operations // Query patterns, performance insights, and usage analytics -use std::collections::{HashMap, VecDeque, BTreeMap}; -use std::sync::atomic::{AtomicU64, Ordering}; +use std::collections::{BTreeMap, HashMap, VecDeque}; use std::sync::RwLock; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; use serde::{Deserialize, Serialize}; @@ -37,9 +37,9 @@ impl Default for AnalyticsConfig { max_events: 100000, realtime_streaming: false, aggregation_intervals: vec![ - Duration::from_secs(60), // 1 minute - Duration::from_secs(3600), // 1 hour - Duration::from_secs(86400), // 1 day + Duration::from_secs(60), // 1 minute + Duration::from_secs(3600), // 1 hour + Duration::from_secs(86400), // 1 day ], log_queries: true, privacy_mode: false, @@ -317,33 +317,40 @@ impl QueryAnalytics { match event.event_type { EventType::Search | EventType::BatchSearch => { self.stats.search_events.fetch_add(1, Ordering::Relaxed); - } + }, EventType::Insert | EventType::BatchInsert => { self.stats.insert_events.fetch_add(1, Ordering::Relaxed); - } + }, EventType::Update => { self.stats.update_events.fetch_add(1, Ordering::Relaxed); - } + }, EventType::Delete => { self.stats.delete_events.fetch_add(1, Ordering::Relaxed); - } - _ => {} + }, + _ => {}, } // Update running average latency { - let Ok(mut avg) = self.stats.avg_latency_ms.write() else { return; }; + let Ok(mut avg) = self.stats.avg_latency_ms.write() else { + return; + }; let count = self.stats.total_events.load(Ordering::Relaxed) as f64; *avg = (*avg * (count - 1.0) + event.latency_ms) / count; } // Store event { - let Ok(mut events) = self.events.write() else { return; }; + let Ok(mut events) = self.events.write() else { + return; + }; if events.len() >= self.config.max_events { events.pop_front(); } - events.push_back(QueryEvent { id: event_id, ..event.clone() }); + events.push_back(QueryEvent { + id: event_id, + ..event.clone() + }); } // Update aggregates @@ -364,18 +371,21 @@ impl QueryAnalytics { return false; } - use std::hash::{Hash, Hasher}; use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); Instant::now().hash(&mut hasher); (hasher.finish() as f64 / u64::MAX as f64) < self.config.sample_rate } fn update_aggregates(&self, event: &QueryEvent) { - let Ok(mut aggregates) = self.aggregates.write() else { return; }; + let Ok(mut aggregates) = self.aggregates.write() else { + return; + }; // Overall aggregate - let overall = aggregates.entry("overall".to_string()) + let overall = aggregates + .entry("overall".to_string()) .or_insert_with(|| AggregateStats { min_latency_ms: f64::INFINITY, ..Default::default() @@ -394,7 +404,8 @@ impl QueryAnalytics { overall.cache_hit_rate = overall.cache_hits as f64 / overall.count as f64; // Per-collection aggregate - let collection = aggregates.entry(format!("collection:{}", event.collection)) + let collection = aggregates + .entry(format!("collection:{}", event.collection)) .or_insert_with(|| AggregateStats { min_latency_ms: f64::INFINITY, ..Default::default() @@ -414,38 +425,41 @@ impl QueryAnalytics { } fn update_patterns(&self, event: &QueryEvent) { - let Ok(mut patterns) = self.patterns.write() else { return; }; + let Ok(mut patterns) = self.patterns.write() else { + return; + }; // Track dimension distribution - *patterns.dimension_distribution + *patterns + .dimension_distribution .entry(event.dimensions) .or_insert(0) += 1; // Track k distribution - *patterns.k_distribution - .entry(event.top_k) - .or_insert(0) += 1; + *patterns.k_distribution.entry(event.top_k).or_insert(0) += 1; // Track query patterns (collection + filter count) let pattern_key = format!("{}:filters={}", event.collection, event.filter_count); - *patterns.frequent_patterns - .entry(pattern_key) - .or_insert(0) += 1; + *patterns.frequent_patterns.entry(pattern_key).or_insert(0) += 1; // Track temporal patterns (hourly) let hour = event.timestamp / 3600; - patterns.temporal_patterns + patterns + .temporal_patterns .entry(hour) .or_insert_with(Vec::new) .push(event.collection.clone()); } fn update_timeseries(&self, event: &QueryEvent) { - let Ok(mut ts) = self.timeseries.write() else { return; }; + let Ok(mut ts) = self.timeseries.write() else { + return; + }; let bucket = event.timestamp / 60; // Minute buckets // Latency series - let latency_series = ts.series + let latency_series = ts + .series .entry("latency".to_string()) .or_insert_with(Vec::new); @@ -471,7 +485,8 @@ impl QueryAnalytics { } // Throughput series - let throughput_series = ts.series + let throughput_series = ts + .series .entry("throughput".to_string()) .or_insert_with(Vec::new); @@ -513,7 +528,8 @@ impl QueryAnalytics { let overall = aggregates.get("overall").cloned().unwrap_or_default(); let uptime = self.stats.started_at.elapsed(); - let qps = self.stats.total_events.load(Ordering::Relaxed) as f64 / uptime.as_secs_f64().max(1.0); + let qps = + self.stats.total_events.load(Ordering::Relaxed) as f64 / uptime.as_secs_f64().max(1.0); let summary = DashboardSummary { total_queries: overall.count, @@ -522,13 +538,15 @@ impl QueryAnalytics { p99_latency_ms: overall.p99_latency_ms, cache_hit_rate: overall.cache_hit_rate, error_rate: overall.error_rate, - active_collections: aggregates.keys() + active_collections: aggregates + .keys() .filter(|k| k.starts_with("collection:")) .count(), uptime, }; - let recent_queries: Vec = events.iter() + let recent_queries: Vec = events + .iter() .rev() .take(20) .map(|e| QuerySummary { @@ -542,7 +560,8 @@ impl QueryAnalytics { let performance_trends = self.get_performance_trends(); - let mut top_collections: Vec = aggregates.iter() + let mut top_collections: Vec = aggregates + .iter() .filter(|(k, _)| k.starts_with("collection:")) .map(|(k, stats)| CollectionStats { name: k.replace("collection:", ""), @@ -597,7 +616,9 @@ impl QueryAnalytics { } fn get_performance_trends(&self) -> Vec { - let Ok(ts) = self.timeseries.read() else { return vec![]; }; + let Ok(ts) = self.timeseries.read() else { + return vec![]; + }; let mut trends = Vec::new(); for (metric_name, points) in &ts.series { @@ -681,18 +702,24 @@ impl QueryAnalytics { /// Get aggregate statistics pub fn get_aggregates(&self) -> HashMap { - let Ok(aggregates) = self.aggregates.read() else { return HashMap::new(); }; + let Ok(aggregates) = self.aggregates.read() else { + return HashMap::new(); + }; aggregates.clone() } /// Get time series data pub fn get_timeseries(&self, metric: &str, last_n_minutes: usize) -> Vec { - let Ok(ts) = self.timeseries.read() else { return vec![]; }; + let Ok(ts) = self.timeseries.read() else { + return vec![]; + }; let cutoff = current_timestamp() / 60 - last_n_minutes as u64; - ts.series.get(metric) + ts.series + .get(metric) .map(|points| { - points.iter() + points + .iter() .filter(|p| p.timestamp >= cutoff) .cloned() .collect() @@ -702,20 +729,26 @@ impl QueryAnalytics { /// Generate a report pub fn generate_report(&self, config: &ReportConfig) -> Result { - let events = self.events.read() + let events = self + .events + .read() .map_err(|_| VecStoreError::LockError("events lock poisoned".into()))?; - let filtered: Vec<_> = events.iter() - .filter(|e| e.timestamp >= config.time_range.start && e.timestamp <= config.time_range.end) + let filtered: Vec<_> = events + .iter() + .filter(|e| { + e.timestamp >= config.time_range.start && e.timestamp <= config.time_range.end + }) .collect(); match config.format { ReportFormat::Json => { let summary: HashMap = HashMap::new(); Ok(serde_json::to_string_pretty(&summary).unwrap_or_default()) - } + }, ReportFormat::Csv => { - let mut csv = String::from("timestamp,collection,latency_ms,result_count,event_type\n"); + let mut csv = + String::from("timestamp,collection,latency_ms,result_count,event_type\n"); for e in filtered { csv.push_str(&format!( "{},{},{},{},{:?}\n", @@ -723,7 +756,7 @@ impl QueryAnalytics { )); } Ok(csv) - } + }, ReportFormat::Html => { let html = format!( r#" @@ -740,9 +773,11 @@ impl QueryAnalytics { filtered.len() ); Ok(html) - } + }, ReportFormat::Prometheus => { - let aggregates = self.aggregates.read() + let aggregates = self + .aggregates + .read() .map_err(|_| VecStoreError::LockError("aggregates lock poisoned".into()))?; let mut output = String::new(); @@ -757,14 +792,12 @@ impl QueryAnalytics { # HELP vecstore_cache_hit_rate Cache hit rate\n\ # TYPE vecstore_cache_hit_rate gauge\n\ vecstore_cache_hit_rate {:.4}\n", - overall.count, - overall.avg_latency_ms, - overall.cache_hit_rate + overall.count, overall.avg_latency_ms, overall.cache_hit_rate )); } Ok(output) - } + }, } } @@ -778,18 +811,24 @@ impl QueryAnalytics { }; }; - let mut frequent: Vec<(String, u64)> = patterns.frequent_patterns.iter() + let mut frequent: Vec<(String, u64)> = patterns + .frequent_patterns + .iter() .map(|(k, v)| (k.clone(), *v)) .collect(); frequent.sort_by(|a, b| b.1.cmp(&a.1)); frequent.truncate(10); - let mut dimensions: Vec<(usize, u64)> = patterns.dimension_distribution.iter() + let mut dimensions: Vec<(usize, u64)> = patterns + .dimension_distribution + .iter() .map(|(k, v)| (*k, *v)) .collect(); dimensions.sort_by(|a, b| b.1.cmp(&a.1)); - let mut k_values: Vec<(usize, u64)> = patterns.k_distribution.iter() + let mut k_values: Vec<(usize, u64)> = patterns + .k_distribution + .iter() .map(|(k, v)| (*k, *v)) .collect(); k_values.sort_by(|a, b| b.1.cmp(&a.1)); @@ -827,22 +866,32 @@ impl QueryAnalytics { error_rate: overall.error_rate, queries_per_second: { let uptime = self.stats.started_at.elapsed().as_secs_f64(); - if uptime > 0.0 { overall.count as f64 / uptime } else { 0.0 } + if uptime > 0.0 { + overall.count as f64 / uptime + } else { + 0.0 + } }, } } /// Reset analytics pub fn reset(&self) { - let Ok(mut events) = self.events.write() else { return; }; + let Ok(mut events) = self.events.write() else { + return; + }; events.clear(); drop(events); - let Ok(mut aggregates) = self.aggregates.write() else { return; }; + let Ok(mut aggregates) = self.aggregates.write() else { + return; + }; aggregates.clear(); drop(aggregates); - let Ok(mut patterns) = self.patterns.write() else { return; }; + let Ok(mut patterns) = self.patterns.write() else { + return; + }; *patterns = PatternAnalyzer { frequent_patterns: HashMap::new(), temporal_patterns: BTreeMap::new(), @@ -857,7 +906,9 @@ impl QueryAnalytics { self.stats.update_events.store(0, Ordering::Relaxed); self.stats.delete_events.store(0, Ordering::Relaxed); - let Ok(mut avg) = self.stats.avg_latency_ms.write() else { return; }; + let Ok(mut avg) = self.stats.avg_latency_ms.write() else { + return; + }; *avg = 0.0; } } @@ -892,9 +943,7 @@ pub struct LatencyHistogram { impl LatencyHistogram { pub fn new(bucket_boundaries: &[f64]) -> Self { - let buckets = bucket_boundaries.iter() - .map(|&b| (b, 0)) - .collect(); + let buckets = bucket_boundaries.iter().map(|&b| (b, 0)).collect(); Self { buckets, @@ -970,7 +1019,8 @@ impl QueryRecommender { recommendations.push(Recommendation { category: RecommendationCategory::Performance, title: "High Average Latency".to_string(), - description: "Consider optimizing index parameters or adding more replicas".to_string(), + description: "Consider optimizing index parameters or adding more replicas" + .to_string(), impact: Impact::High, effort: Impact::Medium, }); @@ -978,33 +1028,35 @@ impl QueryRecommender { // Check dimension distribution for index optimization if let Some((most_common_dim, _)) = patterns.dimension_distribution.first() - && *most_common_dim > 512 { - recommendations.push(Recommendation { - category: RecommendationCategory::IndexOptimization, - title: "High Dimensional Vectors".to_string(), - description: format!( - "Most queries use {}-dimensional vectors. Consider PQ or dimension reduction.", - most_common_dim - ), - impact: Impact::Medium, - effort: Impact::High, - }); - } + && *most_common_dim > 512 + { + recommendations.push(Recommendation { + category: RecommendationCategory::IndexOptimization, + title: "High Dimensional Vectors".to_string(), + description: format!( + "Most queries use {}-dimensional vectors. Consider PQ or dimension reduction.", + most_common_dim + ), + impact: Impact::Medium, + effort: Impact::High, + }); + } // Check k distribution if let Some((most_common_k, _)) = patterns.k_distribution.first() - && *most_common_k > 100 { - recommendations.push(Recommendation { - category: RecommendationCategory::QueryOptimization, - title: "Large Top-K Values".to_string(), - description: format!( - "Many queries request top-{}. Consider pagination instead.", - most_common_k - ), - impact: Impact::Medium, - effort: Impact::Low, - }); - } + && *most_common_k > 100 + { + recommendations.push(Recommendation { + category: RecommendationCategory::QueryOptimization, + title: "Large Top-K Values".to_string(), + description: format!( + "Many queries request top-{}. Consider pagination instead.", + most_common_k + ), + impact: Impact::Medium, + effort: Impact::Low, + }); + } recommendations } @@ -1099,9 +1151,7 @@ mod tests { #[test] fn test_analytics_creation() { - let analytics = QueryAnalyticsBuilder::new() - .sample_rate(1.0) - .build(); + let analytics = QueryAnalyticsBuilder::new().sample_rate(1.0).build(); let dashboard = analytics.get_dashboard(); assert_eq!(dashboard.summary.total_queries, 0); @@ -1109,9 +1159,7 @@ mod tests { #[test] fn test_record_event() { - let analytics = QueryAnalyticsBuilder::new() - .sample_rate(1.0) - .build(); + let analytics = QueryAnalyticsBuilder::new().sample_rate(1.0).build(); let event = QueryEvent { id: 0, @@ -1137,9 +1185,7 @@ mod tests { #[test] fn test_timeseries() { - let analytics = QueryAnalyticsBuilder::new() - .sample_rate(1.0) - .build(); + let analytics = QueryAnalyticsBuilder::new().sample_rate(1.0).build(); for i in 0..10 { let event = QueryEvent { @@ -1166,9 +1212,7 @@ mod tests { #[test] fn test_report_generation() { - let analytics = QueryAnalyticsBuilder::new() - .sample_rate(1.0) - .build(); + let analytics = QueryAnalyticsBuilder::new().sample_rate(1.0).build(); let config = ReportConfig { time_range: TimeRange { diff --git a/src/query_explain.rs b/src/query_explain.rs index 4c18d8d..9785cb6 100644 --- a/src/query_explain.rs +++ b/src/query_explain.rs @@ -81,47 +81,27 @@ pub enum PlanNodeType { /// Root query node Query, /// Vector similarity search - VectorScan { - index_type: String, - metric: String, - }, + VectorScan { index_type: String, metric: String }, /// Metadata filter - Filter { - expression: String, - }, + Filter { expression: String }, /// Hybrid search fusion - Fusion { - strategy: String, - }, + Fusion { strategy: String }, /// Reranking stage - Rerank { - method: String, - }, + Rerank { method: String }, /// Result limit/offset - Limit { - limit: usize, - offset: usize, - }, + Limit { limit: usize, offset: usize }, /// Aggregation - Aggregate { - function: String, - }, + Aggregate { function: String }, /// Cache lookup CacheLookup, /// Parallel execution - Parallel { - workers: usize, - }, + Parallel { workers: usize }, /// Index scan - IndexScan { - index_name: String, - }, + IndexScan { index_name: String }, /// Sequential scan (fallback) SeqScan, /// Sort operation - Sort { - keys: Vec, - }, + Sort { keys: Vec }, } /// Execution plan node @@ -181,9 +161,7 @@ impl PlanNode { /// Total time including children pub fn total_time_ms(&self) -> f64 { let self_time = self.actual_time_ms.unwrap_or(0.0); - let children_time: f64 = self.children.iter() - .map(|c| c.total_time_ms()) - .sum(); + let children_time: f64 = self.children.iter().map(|c| c.total_time_ms()).sum(); self_time + children_time } @@ -196,26 +174,19 @@ impl PlanNode { // Node type and basic info let node_name = match &self.node_type { PlanNodeType::Query => "Query".to_string(), - PlanNodeType::VectorScan { index_type, metric } => - format!("Vector Scan ({}, {})", index_type, metric), - PlanNodeType::Filter { expression } => - format!("Filter ({})", expression), - PlanNodeType::Fusion { strategy } => - format!("Fusion ({})", strategy), - PlanNodeType::Rerank { method } => - format!("Rerank ({})", method), - PlanNodeType::Limit { limit, offset } => - format!("Limit {} Offset {}", limit, offset), - PlanNodeType::Aggregate { function } => - format!("Aggregate ({})", function), + PlanNodeType::VectorScan { index_type, metric } => { + format!("Vector Scan ({}, {})", index_type, metric) + }, + PlanNodeType::Filter { expression } => format!("Filter ({})", expression), + PlanNodeType::Fusion { strategy } => format!("Fusion ({})", strategy), + PlanNodeType::Rerank { method } => format!("Rerank ({})", method), + PlanNodeType::Limit { limit, offset } => format!("Limit {} Offset {}", limit, offset), + PlanNodeType::Aggregate { function } => format!("Aggregate ({})", function), PlanNodeType::CacheLookup => "Cache Lookup".to_string(), - PlanNodeType::Parallel { workers } => - format!("Parallel ({} workers)", workers), - PlanNodeType::IndexScan { index_name } => - format!("Index Scan ({})", index_name), + PlanNodeType::Parallel { workers } => format!("Parallel ({} workers)", workers), + PlanNodeType::IndexScan { index_name } => format!("Index Scan ({})", index_name), PlanNodeType::SeqScan => "Sequential Scan".to_string(), - PlanNodeType::Sort { keys } => - format!("Sort ({})", keys.join(", ")), + PlanNodeType::Sort { keys } => format!("Sort ({})", keys.join(", ")), }; output.push_str(&format!("{}{}{}", prefix, arrow, node_name)); @@ -273,7 +244,8 @@ impl PlanNode { pub fn format_html(&self) -> String { let mut html = String::new(); - html.push_str(r#" + html.push_str( + r#"