diff --git a/Cargo.lock b/Cargo.lock index 3a2c896a1..361ce58a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1631,7 +1631,7 @@ checksum = "c07782be35f9e1140080c6b96f0d44b739e2278479f64e02fdab4e32dfd8b081" dependencies = [ "bitflags 1.3.2", "core-foundation 0.9.4", - "core-graphics-types", + "core-graphics-types 0.1.3", "foreign-types 0.5.0", "libc", ] @@ -1647,6 +1647,17 @@ dependencies = [ "libc", ] +[[package]] +name = "core-graphics-types" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d44a101f213f6c4cdc1853d4b78aef6db6bdfa3468798cc1d9912f4735013eb" +dependencies = [ + "bitflags 2.13.0", + "core-foundation 0.10.1", + "libc", +] + [[package]] name = "core-text" version = "20.1.0" @@ -4914,6 +4925,51 @@ dependencies = [ "regex-automata", ] +[[package]] +name = "lattice-embed" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81c40a349ed412ca8a4b89ad4f0c19cca28e9de969a570bada665e6df195d745" +dependencies = [ + "async-trait", + "blake3", + "chrono", + "lattice-inference", + "lru 0.12.5", + "parking_lot 0.12.5", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", +] + +[[package]] +name = "lattice-inference" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cf759944705182b487cc3a1d99687fa518f031c8239857d351297c9cecabd42" +dependencies = [ + "axum 0.8.9", + "clap", + "futures", + "half", + "image 0.25.10", + "indexmap 2.12.1", + "memmap2", + "metal 0.33.0", + "objc", + "rayon", + "rustc-hash 2.1.2", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.18", + "tokio", + "tracing", + "ureq 2.12.1", +] + [[package]] name = "lazy_static" version = "1.5.0" @@ -5346,7 +5402,22 @@ checksum = "7ecfd3296f8c56b7c1f6fbac3c71cefa9d78ce009850c45000015f206dc7fa21" dependencies = [ "bitflags 2.13.0", "block", - "core-graphics-types", + "core-graphics-types 0.1.3", + "foreign-types 0.5.0", + "log", + "objc", + "paste", +] + +[[package]] +name = "metal" +version = "0.33.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7047791b5bc903b8cd963014b355f71dc9864a9a0b727057676c1dcae5cbc15" +dependencies = [ + "bitflags 2.13.0", + "block", + "core-graphics-types 0.2.0", "foreign-types 0.5.0", "log", "objc", @@ -9093,6 +9164,7 @@ dependencies = [ "dashmap 6.2.1", "hf-hub", "hnsw_rs", + "lattice-embed", "memmap2", "mockall", "ndarray 0.16.1", @@ -9112,6 +9184,7 @@ dependencies = [ "tempfile", "thiserror 2.0.18", "tokenizers", + "tokio", "tracing", "tracing-subscriber", "uuid", @@ -10909,9 +10982,10 @@ dependencies = [ "futures-core", "half", "hf-hub", + "lattice-inference", "md5", "memmap2", - "metal", + "metal 0.29.0", "ndarray 0.16.1", "objc", "objc2", @@ -13411,7 +13485,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7adf545a99a086d362efc739e7cf4317c18cbeda22706000fd434d70ea3d95" dependencies = [ "half", - "metal", + "metal 0.29.0", "objc", "serde", "thiserror 1.0.69", @@ -14069,7 +14143,7 @@ dependencies = [ "block", "bytemuck", "cfg_aliases 0.1.1", - "core-graphics-types", + "core-graphics-types 0.1.3", "glow", "glutin_wgl_sys", "gpu-alloc", @@ -14080,7 +14154,7 @@ dependencies = [ "libc", "libloading 0.8.9", "log", - "metal", + "metal 0.29.0", "naga", "ndk-sys", "objc", diff --git a/crates/ruvector-core/Cargo.toml b/crates/ruvector-core/Cargo.toml index 1c3ed0a63..a21806066 100644 --- a/crates/ruvector-core/Cargo.toml +++ b/crates/ruvector-core/Cargo.toml @@ -55,6 +55,14 @@ tokenizers = { version = "0.20", default-features = false, features = ["onig"], # HuggingFace Hub for model downloads hf-hub = { version = "0.4", optional = true } +# Native (pure-Rust) local embeddings via lattice-embed (not available in WASM). +# NOTE: lattice-embed 0.5.1 requires Rust >= 1.93 (edition 2024). Cargo cannot +# express a per-feature `rust-version`, so enabling the `lattice-embeddings` +# feature raises the effective MSRV above this crate's workspace-inherited +# 1.77 for anyone who turns it on. The default build is unaffected. +lattice-embed = { version = "0.5.1", optional = true } +tokio = { workspace = true, optional = true } + [dev-dependencies] criterion = { workspace = true } proptest = { workspace = true } @@ -110,6 +118,7 @@ uuid-support = [] # Deprecated: uuid is now always included real-embeddings = [] # Feature flag for embedding provider API (use ApiEmbedding for production) api-embeddings = ["reqwest"] # API-based embeddings (not available in WASM) onnx-embeddings = ["ort", "tokenizers", "hf-hub"] # ONNX-based local embeddings (not available in WASM) +lattice-embeddings = ["dep:lattice-embed", "dep:tokio"] # Native pure-Rust local embeddings via lattice-embed (not available in WASM) [lib] crate-type = ["rlib"] diff --git a/crates/ruvector-core/src/embeddings.rs b/crates/ruvector-core/src/embeddings.rs index c7e9c396f..9849b22b1 100644 --- a/crates/ruvector-core/src/embeddings.rs +++ b/crates/ruvector-core/src/embeddings.rs @@ -6,6 +6,7 @@ //! //! - **HashEmbedding**: Fast hash-based placeholder (default, not semantic) //! - **OnnxEmbedding**: Real semantic embeddings using ONNX Runtime (feature: `onnx-embeddings`) ✅ RECOMMENDED +//! - **LatticeEmbedding**: Real semantic embeddings using lattice-embed, pure-Rust native inference (feature: `lattice-embeddings`) //! - **CandleEmbedding**: Real embeddings using candle-transformers (feature: `real-embeddings`) //! - **ApiEmbedding**: External API calls (OpenAI, Anthropic, Cohere, etc.) //! @@ -33,7 +34,11 @@ //! ``` use crate::error::Result; -#[cfg(any(feature = "real-embeddings", feature = "api-embeddings"))] +#[cfg(any( + feature = "real-embeddings", + feature = "api-embeddings", + feature = "lattice-embeddings" +))] use crate::error::RuvectorError; use std::sync::Arc; @@ -718,6 +723,363 @@ pub mod onnx { #[cfg(feature = "onnx-embeddings")] pub use onnx::OnnxEmbedding; +// ============================================================================ +// Lattice Embeddings (pure-Rust, native, no C++ FFI / no ONNX Runtime) +// ============================================================================ + +/// Native embedding provider backed by [`lattice-embed`](https://crates.io/crates/lattice-embed), +/// a pure-Rust transformer inference engine (SIMD matmul, safetensors weight +/// loading, no ONNX Runtime, no C++ FFI). +/// +/// Requires feature flag: `lattice-embeddings` +/// +/// ## Supported models +/// - `bge-small-en-v1.5` / `BAAI/bge-small-en-v1.5` (384 dims, default, recommended for `.rvf` packs) +/// - `bge-base-en-v1.5` / `BAAI/bge-base-en-v1.5` (768 dims) +/// - `bge-large-en-v1.5` / `BAAI/bge-large-en-v1.5` (1024 dims) +/// - `multilingual-e5-small` / `intfloat/multilingual-e5-small` (384 dims) +/// - `multilingual-e5-base` / `intfloat/multilingual-e5-base` (768 dims) +/// - `all-minilm-l6-v2` / `sentence-transformers/all-MiniLM-L6-v2` (384 dims) +/// - `paraphrase-multilingual-minilm-l12-v2` / `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2` (384 dims) +/// - `qwen3-embedding-0.6b` / `Qwen/Qwen3-Embedding-0.6B` (1024 dims) +/// - `qwen3-embedding-4b` / `Qwen/Qwen3-Embedding-4B` (2560 dims) +/// +/// Model-id parsing is delegated to `lattice_embed::EmbeddingModel`'s own +/// `FromStr` impl (case-insensitive, accepts display names, short names, and +/// HuggingFace ids) rather than re-implementing the mapping here, so this +/// provider stays in sync with lattice-embed's canonical model table. +/// +/// ## CPU / native, no GPU +/// This provider uses lattice-embed's default `native` feature (CPU-only, +/// SIMD-accelerated). It does **not** enable lattice-embed's `metal-gpu` +/// feature. +/// +/// ## Minimum Supported Rust Version +/// Enabling the `lattice-embeddings` feature raises the effective MSRV for +/// this crate to Rust 1.93 (edition 2024), since `lattice-embed` requires it. +/// Cargo has no mechanism to express a per-feature `rust-version`, so this is +/// not reflected in `rust-version.workspace = true` above — it only applies +/// when this feature is enabled. The crate's default build (feature +/// disabled) keeps the workspace MSRV of 1.77. +/// +/// ## Model download +/// BERT-family models (BGE, E5, MiniLM) download automatically from +/// HuggingFace into `~/.lattice/models` on first use. Qwen3-Embedding models +/// must be placed at `~/.lattice/models/qwen3-embedding-{0.6b,4b}/` manually +/// (or pointed to via `LATTICE_QWEN_MODEL_DIR`) before first use. +/// +/// ## Asymmetric retrieval (query vs. passage prefixing) +/// BGE, E5, and Qwen3-Embedding are asymmetric retrievers: the query side is +/// prefixed with a retrieval instruction, the document side is not. +/// [`EmbeddingProvider::embed`] always takes the **passage/document** side (no +/// query instruction) via `lattice_embed::EmbeddingService::embed_passage`. +/// Use the inherent [`LatticeEmbedding::embed_query`] method for query text — +/// it applies the model's query instruction via +/// `EmbeddingService::embed_query`: BGE v1.5 prefixes queries with +/// `"Represent this sentence for searching relevant passages: "`, E5 with +/// `"query: "`, and Qwen3-Embedding with its search instruction. For all three +/// families `embed_query` and `embed` therefore produce different vectors, +/// which is what makes asymmetric retrieval correct. MiniLM is genuinely +/// symmetric (contrastive training on raw text, no prefix), so its two methods +/// are equivalent. +/// +/// # Example +/// ```rust,no_run +/// use ruvector_core::embeddings::{EmbeddingProvider, LatticeEmbedding}; +/// +/// let provider = LatticeEmbedding::from_pretrained("bge-small-en-v1.5")?; +/// +/// // Document side: no query instruction. +/// let doc_embedding = provider.embed("The cat sat on the mat.")?; +/// assert_eq!(doc_embedding.len(), 384); +/// +/// // Query side: applies the model's query instruction, if any. +/// let query_embedding = provider.embed_query("Where did the cat sit?")?; +/// assert_eq!(query_embedding.len(), 384); +/// # Ok::<(), Box>(()) +/// ``` +#[cfg(feature = "lattice-embeddings")] +pub mod lattice_native { + use super::*; + use lattice_embed::{ + EmbeddingModel as LatticeEmbeddingModel, EmbeddingService, NativeEmbeddingService, + }; + use std::sync::mpsc; + use std::sync::Mutex; + use std::thread; + + /// Which side of asymmetric retrieval a queued embedding request is for. + enum EmbedKind { + Query, + Passage, + } + + /// A single embedding request sent to the worker thread, with a + /// per-request reply channel for the result. + struct EmbedRequest { + kind: EmbedKind, + text: String, + reply_tx: mpsc::Sender, String>>, + } + + /// See the [module-level docs](self) for the full provider description. + /// + /// # Threading model + /// `lattice-embed`'s [`EmbeddingService`] is `async`-only (no sync/blocking + /// API), but [`EmbeddingProvider::embed`] is a sync method that ruvector-core + /// callers may invoke from anywhere, including from inside an existing Tokio + /// runtime (e.g. an async server handler). Bridging via a stored + /// `Runtime::block_on` would panic in that case (`block_on` cannot be + /// called from within an already-running runtime). Instead, the runtime and + /// the embedding service live on a dedicated worker thread with no ambient + /// async context of its own; `embed` / `embed_query` send a request over a + /// channel and block on `Receiver::recv`, which is safe to call from any + /// context, sync or async. + pub struct LatticeEmbedding { + model: LatticeEmbeddingModel, + model_id: &'static str, + dimensions: usize, + request_tx: Mutex>, + // Keeps the worker thread's handle alive for the lifetime of this + // provider. Not joined on drop (that would block); dropping + // `request_tx` closes the channel, which ends the worker's `recv` + // loop and lets the thread exit on its own. + _worker: thread::JoinHandle<()>, + } + + impl LatticeEmbedding { + /// Load a pre-trained embedding model by id. + /// + /// Accepts display names (`"bge-small-en-v1.5"`), short names + /// (`"bge-small"`, `"small"`), and HuggingFace ids + /// (`"BAAI/bge-small-en-v1.5"`) — see [`lattice_embed::EmbeddingModel`]'s + /// `FromStr` impl for the full accepted set. Returns an error for any + /// unrecognized id, and for any id that resolves to a model + /// [`lattice_embed`]'s native service cannot run locally (e.g. the + /// remote-only OpenAI variants). + /// + /// # Example + /// ```rust,no_run + /// use ruvector_core::embeddings::LatticeEmbedding; + /// + /// let provider = LatticeEmbedding::from_pretrained("bge-small-en-v1.5")?; + /// # Ok::<(), Box>(()) + /// ``` + pub fn from_pretrained(model_id: &str) -> Result { + let model: LatticeEmbeddingModel = model_id.parse().map_err(|e: String| { + RuvectorError::ModelLoadError(format!( + "unknown lattice-embed model id '{model_id}': {e}" + )) + })?; + Self::with_model(model) + } + + /// Load a pre-trained embedding model from an already-resolved + /// [`lattice_embed::EmbeddingModel`] variant. + pub fn with_model(model: LatticeEmbeddingModel) -> Result { + if !model.is_local() { + return Err(RuvectorError::ModelLoadError(format!( + "'{model}' cannot be loaded natively: lattice-embed's \ + NativeEmbeddingService only supports models it can run \ + on-device. Remote/API-only models (e.g. the OpenAI \ + text-embedding-* family) are not supported by LatticeEmbedding." + ))); + } + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|e| { + RuvectorError::ModelLoadError(format!( + "failed to build tokio runtime for LatticeEmbedding: {e}" + )) + })?; + let service = NativeEmbeddingService::with_model(model); + + let (request_tx, request_rx) = mpsc::channel::(); + let worker = thread::Builder::new() + .name("lattice-embed-worker".to_string()) + .spawn(move || { + // No ambient Tokio runtime exists on this thread, so + // `block_on` here can never panic on nested-runtime + // grounds regardless of the caller's own context. + for request in request_rx { + let outcome = runtime.block_on(async { + match request.kind { + EmbedKind::Query => { + service.embed_query(&[request.text], model).await + } + EmbedKind::Passage => { + service.embed_passage(&[request.text], model).await + } + } + }); + let mapped = + outcome + .map_err(|e| e.to_string()) + .and_then(|mut embeddings| { + embeddings.pop().ok_or_else(|| { + "lattice-embed returned no embedding".to_string() + }) + }); + // Ignore send errors: they only occur if the caller + // already dropped its reply receiver. + let _ = request.reply_tx.send(mapped); + } + }) + .map_err(|e| { + RuvectorError::ModelLoadError(format!( + "failed to spawn LatticeEmbedding worker thread: {e}" + )) + })?; + + Ok(Self { + model, + model_id: model.model_id(), + dimensions: model.dimensions(), + request_tx: Mutex::new(request_tx), + _worker: worker, + }) + } + + /// Get the dimensionality of embeddings produced by the loaded model. + pub fn dimensions(&self) -> usize { + self.dimensions + } + + /// Embed **query** text, applying the model's query-side prompt + /// instruction if it uses one (BGE v1.5's `"Represent this sentence + /// for searching relevant passages: "` prefix, E5's `"query: "` + /// prefix, Qwen3's search-query instruction). For those asymmetric + /// models this produces a different vector than + /// [`EmbeddingProvider::embed`]; only MiniLM is symmetric, so its two + /// methods are equivalent. + /// + /// This is what makes asymmetric retrieval correct: index documents + /// via [`EmbeddingProvider::embed`] (passage side, no prefix) and + /// embed the search query via this method (query side, prefixed). + /// + /// Safe to call from any context, including from inside a Tokio + /// runtime — see the [threading model](LatticeEmbedding#threading-model). + pub fn embed_query(&self, text: &str) -> Result> { + self.send_request(EmbedKind::Query, text) + } + + /// Send an embedding request to the worker thread and block on the + /// reply. Never calls `block_on` on the caller's thread, so this is + /// safe to invoke from inside an existing async runtime. + fn send_request(&self, kind: EmbedKind, text: &str) -> Result> { + let (reply_tx, reply_rx) = mpsc::channel(); + let request = EmbedRequest { + kind, + text: text.to_string(), + reply_tx, + }; + + self.request_tx + .lock() + .map_err(|_| { + RuvectorError::ModelInferenceError( + "lattice-embed embedding worker request channel poisoned".to_string(), + ) + })? + .send(request) + .map_err(|_| { + RuvectorError::ModelInferenceError( + "lattice-embed embedding worker unavailable".to_string(), + ) + })?; + + reply_rx + .recv() + .map_err(|_| { + RuvectorError::ModelInferenceError( + "lattice-embed embedding worker unavailable".to_string(), + ) + })? + .map_err(|e| { + RuvectorError::ModelInferenceError(format!( + "lattice-embed embedding failed: {e}" + )) + }) + } + } + + impl EmbeddingProvider for LatticeEmbedding { + /// Embed **passage/document** text (no query instruction applied). + /// + /// Use [`LatticeEmbedding::embed_query`] for the query side of + /// asymmetric retrieval. Safe to call from any context, including + /// from inside a Tokio runtime — see the + /// [threading model](LatticeEmbedding#threading-model). + fn embed(&self, text: &str) -> Result> { + self.send_request(EmbedKind::Passage, text) + } + + fn dimensions(&self) -> usize { + self.dimensions + } + + fn name(&self) -> &str { + self.model_id + } + } + + #[cfg(test)] + mod tests { + use super::*; + + #[test] + fn from_pretrained_rejects_remote_only_models() { + // "text-embedding-3-small" and "openai" both parse successfully + // to `EmbeddingModel::TextEmbedding3Small` (see lattice-embed's + // `FromStr` impl) but that variant is remote/API-only — + // `NativeEmbeddingService` cannot run it. Both aliases must be + // rejected at construction time, not on first `embed()` call. + assert!( + LatticeEmbedding::from_pretrained("text-embedding-3-small").is_err(), + "remote-only model 'text-embedding-3-small' must be rejected at construction" + ); + assert!( + LatticeEmbedding::from_pretrained("openai").is_err(), + "remote-only model alias 'openai' must be rejected at construction" + ); + } + + #[test] + fn from_pretrained_accepts_native_local_model() { + assert!( + LatticeEmbedding::from_pretrained("bge-small-en-v1.5").is_ok(), + "native local model 'bge-small-en-v1.5' must construct successfully" + ); + } + + /// Regression test for the nested-runtime panic: `embed` / `embed_query` + /// used to call `Runtime::block_on` on a `Runtime` stored on the + /// provider, which panics when invoked from inside an already-running + /// Tokio runtime. The worker-thread bridge has no ambient runtime on + /// the calling side, so both calls must succeed here instead. + #[tokio::test] + async fn embed_from_inside_async_runtime_does_not_panic() { + let provider = LatticeEmbedding::from_pretrained("bge-small-en-v1.5") + .expect("bge-small-en-v1.5 is a native local model"); + + let doc = provider + .embed("a nested-runtime regression test") + .expect("embed must not panic or error from inside a Tokio runtime"); + assert_eq!(doc.len(), provider.dimensions()); + + let query = provider + .embed_query("a nested-runtime regression test") + .expect("embed_query must not panic or error from inside a Tokio runtime"); + assert_eq!(query.len(), provider.dimensions()); + } + } +} + +#[cfg(feature = "lattice-embeddings")] +pub use lattice_native::LatticeEmbedding; + /// Type-erased embedding provider for dynamic dispatch pub type BoxedEmbeddingProvider = Arc; @@ -839,4 +1201,63 @@ mod tests { } } } + + #[cfg(feature = "lattice-embeddings")] + mod lattice_tests { + use super::*; + use crate::embeddings::LatticeEmbedding; + + /// Pure model-id mapping test — no network, no model load. + /// `LatticeEmbedding::from_pretrained` delegates to + /// `lattice_embed::EmbeddingModel::from_str`; this test locks in that + /// bge-small resolves from both its display name and its HuggingFace + /// id, and that an unrecognized id errors instead of silently + /// defaulting. + #[test] + fn test_lattice_from_pretrained_model_id_mapping() { + let by_display_name = LatticeEmbedding::from_pretrained("bge-small-en-v1.5").unwrap(); + assert_eq!(by_display_name.dimensions(), 384); + assert_eq!(EmbeddingProvider::dimensions(&by_display_name), 384); + + let by_hf_id = LatticeEmbedding::from_pretrained("BAAI/bge-small-en-v1.5").unwrap(); + assert_eq!(by_hf_id.dimensions(), 384); + + let unknown = LatticeEmbedding::from_pretrained("not-a-real-model"); + assert!( + unknown.is_err(), + "unknown model id should error, not default" + ); + } + + #[test] + fn test_lattice_from_pretrained_minilm_mapping() { + let by_short = LatticeEmbedding::from_pretrained("all-minilm-l6-v2").unwrap(); + assert_eq!(by_short.dimensions(), 384); + + let by_hf_id = + LatticeEmbedding::from_pretrained("sentence-transformers/all-MiniLM-L6-v2") + .unwrap(); + assert_eq!(by_hf_id.dimensions(), 384); + } + + /// Real end-to-end embedding test. Requires the bge-small-en-v1.5 + /// model to be downloaded from HuggingFace on first use (~130MB) — + /// network access, not run in CI. Run manually with: + /// cargo test -p ruvector-core --features lattice-embeddings -- --ignored lattice_tests + #[test] + #[ignore] + fn test_lattice_embedding_real() { + let provider = LatticeEmbedding::from_pretrained("bge-small-en-v1.5").unwrap(); + + let embedding = provider.embed("hello world").unwrap(); + assert_eq!(embedding.len(), 384); + assert!(embedding.iter().all(|v| v.is_finite())); + + let norm: f32 = embedding.iter().map(|x| x * x).sum::().sqrt(); + assert!( + (norm - 1.0).abs() < 1e-3, + "embedding should be L2-normalized, got norm={norm}" + ); + } + } } diff --git a/crates/ruvector-core/src/lib.rs b/crates/ruvector-core/src/lib.rs index f46c72945..5663b9f96 100644 --- a/crates/ruvector-core/src/lib.rs +++ b/crates/ruvector-core/src/lib.rs @@ -102,6 +102,9 @@ pub use embeddings::CandleEmbedding; #[cfg(feature = "onnx-embeddings")] pub use embeddings::OnnxEmbedding; +#[cfg(feature = "lattice-embeddings")] +pub use embeddings::LatticeEmbedding; + // Compile-time warning about AgenticDB limitations #[cfg(feature = "storage")] #[allow(deprecated, clippy::let_unit_value)]