sochdb_query/
embedding_provider.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2// SochDB - LLM-Optimized Embedded Database
3// Copyright (C) 2026 Sushanth Reddy Vanagala (https://github.com/sushanthpy)
4//
5// This program is free software: you can redistribute it and/or modify
6// it under the terms of the GNU Affero General Public License as published by
7// the Free Software Foundation, either version 3 of the License, or
8// (at your option) any later version.
9//
10// This program is distributed in the hope that it will be useful,
11// but WITHOUT ANY WARRANTY; without even the implied warranty of
12// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
13// GNU Affero General Public License for more details.
14//
15// You should have received a copy of the GNU Affero General Public License
16// along with this program. If not, see <https://www.gnu.org/licenses/>.
17
18//! Automatic Embedding Generation (Task 2)
19//!
20//! This module provides colocated embedding resolution for text-to-vector conversion.
21//! It enables first-class text search by automatically generating embeddings.
22//!
23//! ## Design
24//!
25//! ```text
26//! search_text(collection, text, k)
27//!     │
28//!     ▼
29//! ┌─────────────────┐
30//! │ EmbeddingProvider │
31//! │  ├─ LRU Cache    │
32//! │  └─ ONNX Runtime │
33//! └─────────────────┘
34//!     │
35//!     ▼
36//! search_by_embedding(collection, embedding, k)
37//! ```
38//!
39//! ## Providers
40//!
41//! - `LocalProvider`: Uses FastEmbed/ONNX for offline embedding
42//! - `CachedProvider`: LRU cache wrapper for any provider
43//! - `MockProvider`: For testing
44//!
45//! ## Complexity
46//!
47//! - Embedding generation: O(n) where n = text length (transformer inference)
48//! - Cache lookup: O(1) expected (hash-based LRU)
49//! - Batch embedding: O(k) compute with ~O(1) ONNX session overhead
50
51use moka::sync::Cache;
52use std::sync::Arc;
53
54// ============================================================================
55// Embedding Provider Trait
56// ============================================================================
57
58/// Error type for embedding operations
59#[derive(Debug, Clone)]
60pub enum EmbeddingError {
61    /// Model not loaded or unavailable
62    ModelNotAvailable(String),
63    /// Text too long for model
64    TextTooLong { max_length: usize, actual: usize },
65    /// Dimension mismatch
66    DimensionMismatch { expected: usize, actual: usize },
67    /// Provider error
68    ProviderError(String),
69    /// Cache error
70    CacheError(String),
71}
72
73impl std::fmt::Display for EmbeddingError {
74    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
75        match self {
76            Self::ModelNotAvailable(model) => write!(f, "Embedding model not available: {}", model),
77            Self::TextTooLong { max_length, actual } => {
78                write!(f, "Text too long: {} > {} max", actual, max_length)
79            }
80            Self::DimensionMismatch { expected, actual } => {
81                write!(
82                    f,
83                    "Dimension mismatch: expected {}, got {}",
84                    expected, actual
85                )
86            }
87            Self::ProviderError(msg) => write!(f, "Provider error: {}", msg),
88            Self::CacheError(msg) => write!(f, "Cache error: {}", msg),
89        }
90    }
91}
92
93impl std::error::Error for EmbeddingError {}
94
95/// Result type for embedding operations
96pub type EmbeddingResult<T> = Result<T, EmbeddingError>;
97
98/// Embedding provider trait
99pub trait EmbeddingProvider: Send + Sync {
100    /// Get the model name
101    fn model_name(&self) -> &str;
102
103    /// Get the embedding dimension
104    fn dimension(&self) -> usize;
105
106    /// Maximum text length (in characters or tokens)
107    fn max_length(&self) -> usize;
108
109    /// Generate embedding for a single text
110    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>>;
111
112    /// Generate embeddings for multiple texts (batch)
113    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
114        // Default implementation: sequential embedding
115        texts.iter().map(|t| self.embed(t)).collect()
116    }
117
118    /// Normalize an embedding vector (L2 normalization)
119    fn normalize(&self, embedding: &mut [f32]) {
120        let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
121        if norm > 1e-10 {
122            for x in embedding.iter_mut() {
123                *x /= norm;
124            }
125        }
126    }
127}
128
129// ============================================================================
130// Embedding Configuration
131// ============================================================================
132
133/// Configuration for embedding providers
134#[derive(Debug, Clone)]
135pub struct EmbeddingConfig {
136    /// Model identifier (e.g., "all-MiniLM-L6-v2")
137    pub model: String,
138
139    /// Model path (for local ONNX models)
140    pub model_path: Option<String>,
141
142    /// Embedding dimension
143    pub dimension: usize,
144
145    /// Maximum text length
146    pub max_length: usize,
147
148    /// Whether to normalize embeddings
149    pub normalize: bool,
150
151    /// Batch size for embedding generation
152    pub batch_size: usize,
153
154    /// Cache size (number of embeddings to cache)
155    pub cache_size: usize,
156
157    /// Cache TTL in seconds (0 = no expiry)
158    pub cache_ttl_secs: u64,
159}
160
161impl Default for EmbeddingConfig {
162    fn default() -> Self {
163        Self {
164            model: "all-MiniLM-L6-v2".to_string(),
165            model_path: None,
166            dimension: 384, // MiniLM dimension
167            max_length: 512,
168            normalize: true,
169            batch_size: 32,
170            cache_size: 10_000,
171            cache_ttl_secs: 3600, // 1 hour
172        }
173    }
174}
175
176impl EmbeddingConfig {
177    /// Create config for sentence-transformers models
178    pub fn sentence_transformer(model: &str) -> Self {
179        let dimension = match model {
180            "all-MiniLM-L6-v2" => 384,
181            "all-MiniLM-L12-v2" => 384,
182            "all-mpnet-base-v2" => 768,
183            "paraphrase-MiniLM-L6-v2" => 384,
184            "multi-qa-MiniLM-L6-cos-v1" => 384,
185            _ => 384, // Default
186        };
187
188        Self {
189            model: model.to_string(),
190            dimension,
191            ..Default::default()
192        }
193    }
194
195    /// Create config for OpenAI-compatible models
196    pub fn openai(model: &str) -> Self {
197        let dimension = match model {
198            "text-embedding-ada-002" => 1536,
199            "text-embedding-3-small" => 1536,
200            "text-embedding-3-large" => 3072,
201            _ => 1536,
202        };
203
204        Self {
205            model: model.to_string(),
206            dimension,
207            max_length: 8192,
208            ..Default::default()
209        }
210    }
211}
212
213// ============================================================================
214// Mock Embedding Provider (for testing)
215// ============================================================================
216
217/// Mock embedding provider for testing
218pub struct MockEmbeddingProvider {
219    config: EmbeddingConfig,
220    /// Deterministic embeddings based on text hash
221    use_hash: bool,
222}
223
224impl MockEmbeddingProvider {
225    /// Create a new mock provider
226    pub fn new(dimension: usize) -> Self {
227        Self {
228            config: EmbeddingConfig {
229                model: "mock".to_string(),
230                dimension,
231                ..Default::default()
232            },
233            use_hash: true,
234        }
235    }
236
237    /// Create with custom config
238    pub fn with_config(config: EmbeddingConfig) -> Self {
239        Self {
240            config,
241            use_hash: true,
242        }
243    }
244
245    /// Generate a deterministic embedding from text
246    fn hash_embed(&self, text: &str) -> Vec<f32> {
247        use std::collections::hash_map::DefaultHasher;
248        use std::hash::{Hash, Hasher};
249
250        let mut embedding = Vec::with_capacity(self.config.dimension);
251
252        // Generate pseudo-random values based on text hash
253        for i in 0..self.config.dimension {
254            let mut hasher = DefaultHasher::new();
255            text.hash(&mut hasher);
256            i.hash(&mut hasher);
257            let hash = hasher.finish();
258
259            // Convert to f32 in range [-1, 1]
260            let value = ((hash as f64) / (u64::MAX as f64) * 2.0 - 1.0) as f32;
261            embedding.push(value);
262        }
263
264        embedding
265    }
266}
267
268impl EmbeddingProvider for MockEmbeddingProvider {
269    fn model_name(&self) -> &str {
270        &self.config.model
271    }
272
273    fn dimension(&self) -> usize {
274        self.config.dimension
275    }
276
277    fn max_length(&self) -> usize {
278        self.config.max_length
279    }
280
281    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
282        if text.len() > self.config.max_length {
283            return Err(EmbeddingError::TextTooLong {
284                max_length: self.config.max_length,
285                actual: text.len(),
286            });
287        }
288
289        let mut embedding = if self.use_hash {
290            self.hash_embed(text)
291        } else {
292            vec![0.0; self.config.dimension]
293        };
294
295        if self.config.normalize {
296            self.normalize(&mut embedding);
297        }
298
299        Ok(embedding)
300    }
301}
302
303// ============================================================================
304// Cached Embedding Provider
305// ============================================================================
306
307/// LRU-cached embedding provider wrapper
308pub struct CachedEmbeddingProvider<P: EmbeddingProvider> {
309    /// Inner provider
310    inner: P,
311
312    /// LRU cache: text hash -> embedding
313    cache: Cache<u64, Vec<f32>>,
314
315    /// Cache statistics
316    stats: Arc<CacheStats>,
317}
318
319/// Cache statistics
320#[derive(Debug, Default)]
321pub struct CacheStats {
322    /// Number of cache hits
323    pub hits: std::sync::atomic::AtomicUsize,
324    /// Number of cache misses
325    pub misses: std::sync::atomic::AtomicUsize,
326    /// Number of embeddings cached
327    pub size: std::sync::atomic::AtomicUsize,
328}
329
330impl CacheStats {
331    /// Get hit rate
332    pub fn hit_rate(&self) -> f64 {
333        let hits = self.hits.load(std::sync::atomic::Ordering::Relaxed);
334        let misses = self.misses.load(std::sync::atomic::Ordering::Relaxed);
335        let total = hits + misses;
336        if total == 0 {
337            0.0
338        } else {
339            hits as f64 / total as f64
340        }
341    }
342}
343
344impl<P: EmbeddingProvider> CachedEmbeddingProvider<P> {
345    /// Create a new cached provider
346    pub fn new(inner: P, cache_size: usize) -> Self {
347        Self {
348            inner,
349            cache: Cache::new(cache_size as u64),
350            stats: Arc::new(CacheStats::default()),
351        }
352    }
353
354    /// Create with TTL
355    pub fn with_ttl(inner: P, cache_size: usize, ttl_secs: u64) -> Self {
356        let cache = Cache::builder()
357            .max_capacity(cache_size as u64)
358            .time_to_live(std::time::Duration::from_secs(ttl_secs))
359            .build();
360
361        Self {
362            inner,
363            cache,
364            stats: Arc::new(CacheStats::default()),
365        }
366    }
367
368    /// Get cache statistics
369    pub fn stats(&self) -> &Arc<CacheStats> {
370        &self.stats
371    }
372
373    /// Compute hash for cache key
374    fn text_hash(text: &str) -> u64 {
375        use std::collections::hash_map::DefaultHasher;
376        use std::hash::{Hash, Hasher};
377
378        let mut hasher = DefaultHasher::new();
379        text.hash(&mut hasher);
380        hasher.finish()
381    }
382}
383
384impl<P: EmbeddingProvider> EmbeddingProvider for CachedEmbeddingProvider<P> {
385    fn model_name(&self) -> &str {
386        self.inner.model_name()
387    }
388
389    fn dimension(&self) -> usize {
390        self.inner.dimension()
391    }
392
393    fn max_length(&self) -> usize {
394        self.inner.max_length()
395    }
396
397    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
398        let hash = Self::text_hash(text);
399
400        // Check cache
401        if let Some(cached) = self.cache.get(&hash) {
402            self.stats
403                .hits
404                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
405            return Ok(cached);
406        }
407
408        self.stats
409            .misses
410            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
411
412        // Generate embedding
413        let embedding = self.inner.embed(text)?;
414
415        // Cache result
416        self.cache.insert(hash, embedding.clone());
417        self.stats
418            .size
419            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
420
421        Ok(embedding)
422    }
423
424    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
425        let mut results = Vec::with_capacity(texts.len());
426        let mut uncached: Vec<(usize, &str)> = Vec::new();
427
428        // Check cache for each text
429        for (i, text) in texts.iter().enumerate() {
430            let hash = Self::text_hash(text);
431            if let Some(cached) = self.cache.get(&hash) {
432                self.stats
433                    .hits
434                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
435                results.push((i, cached));
436            } else {
437                self.stats
438                    .misses
439                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
440                uncached.push((i, *text));
441            }
442        }
443
444        // Generate embeddings for uncached texts
445        if !uncached.is_empty() {
446            let uncached_texts: Vec<&str> = uncached.iter().map(|(_, t)| *t).collect();
447            let embeddings = self.inner.embed_batch(&uncached_texts)?;
448
449            for ((i, text), embedding) in uncached.iter().zip(embeddings.into_iter()) {
450                let hash = Self::text_hash(text);
451                self.cache.insert(hash, embedding.clone());
452                self.stats
453                    .size
454                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
455                results.push((*i, embedding));
456            }
457        }
458
459        // Sort by original index
460        results.sort_by_key(|(i, _)| *i);
461        Ok(results.into_iter().map(|(_, e)| e).collect())
462    }
463}
464
465// ============================================================================
466// Local ONNX Provider (Stub)
467// ============================================================================
468
469/// Local ONNX-based embedding provider
470///
471/// This is a stub implementation. In production, this would use:
472/// - ort (ONNX Runtime) for model inference
473/// - fastembed-rs for pre-packaged models
474/// - tokenizers for text preprocessing
475#[derive(Debug)]
476pub struct LocalOnnxProvider {
477    config: EmbeddingConfig,
478    /// Model weights (placeholder)
479    #[allow(dead_code)]
480    model_loaded: bool,
481}
482
483impl LocalOnnxProvider {
484    /// Create a new local ONNX provider
485    pub fn new(config: EmbeddingConfig) -> EmbeddingResult<Self> {
486        // In production: Load ONNX model from path
487        Ok(Self {
488            config,
489            model_loaded: false,
490        })
491    }
492
493    /// Load a pre-trained model by name
494    pub fn load_pretrained(model_name: &str) -> EmbeddingResult<Self> {
495        let config = EmbeddingConfig::sentence_transformer(model_name);
496        Self::new(config)
497    }
498}
499
500impl EmbeddingProvider for LocalOnnxProvider {
501    fn model_name(&self) -> &str {
502        &self.config.model
503    }
504
505    fn dimension(&self) -> usize {
506        self.config.dimension
507    }
508
509    fn max_length(&self) -> usize {
510        self.config.max_length
511    }
512
513    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
514        // Stub: Return mock embedding
515        // In production: Run ONNX inference
516        let mock = MockEmbeddingProvider::with_config(self.config.clone());
517        mock.embed(text)
518    }
519}
520
521// ============================================================================
522// FastEmbed (ONNX) semantic provider  —  feature = "fastembed"
523// ============================================================================
524
525/// Real semantic embedding provider backed by fastembed-rs (ONNX runtime).
526///
527/// Only compiled with the `fastembed` feature (the `ort` native dep is heavy).
528/// The model is downloaded + cached on first construction. `TextEmbedding::embed`
529/// takes `&mut self`, so the model is held behind a `Mutex` to satisfy the
530/// `&self` trait contract (embedding is the bottleneck anyway, so the lock is
531/// not the limiting factor).
532#[cfg(feature = "fastembed")]
533pub struct FastEmbedProvider {
534    model: std::sync::Mutex<fastembed::TextEmbedding>,
535    model_name: String,
536    dimension: usize,
537    max_length: usize,
538}
539
540#[cfg(feature = "fastembed")]
541impl std::fmt::Debug for FastEmbedProvider {
542    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
543        f.debug_struct("FastEmbedProvider")
544            .field("model_name", &self.model_name)
545            .field("dimension", &self.dimension)
546            .finish()
547    }
548}
549
550#[cfg(feature = "fastembed")]
551impl FastEmbedProvider {
552    /// Construct from a short model alias, e.g. `bge-small-en`, `all-minilm`,
553    /// `bge-base-en`, `bge-large-en`. Downloads + caches the ONNX model on first
554    /// use (set `FASTEMBED_CACHE_DIR` to control the cache location).
555    pub fn new(model_alias: &str) -> EmbeddingResult<Self> {
556        use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
557        let (model, dimension) = match model_alias.to_ascii_lowercase().as_str() {
558            "bge-small-en" | "bge-small-en-v1.5" | "bge-small" => {
559                (EmbeddingModel::BGESmallENV15, 384)
560            }
561            "all-minilm" | "all-minilm-l6-v2" | "minilm" => (EmbeddingModel::AllMiniLML6V2, 384),
562            "bge-base-en" | "bge-base-en-v1.5" | "bge-base" => (EmbeddingModel::BGEBaseENV15, 768),
563            "bge-large-en" | "bge-large-en-v1.5" | "bge-large" => {
564                (EmbeddingModel::BGELargeENV15, 1024)
565            }
566            other => {
567                return Err(EmbeddingError::ModelNotAvailable(format!(
568                    "fastembed:{other}"
569                )));
570            }
571        };
572        let model =
573            TextEmbedding::try_new(InitOptions::new(model).with_show_download_progress(false))
574                .map_err(|e| {
575                    EmbeddingError::ProviderError(format!("fastembed init failed: {e}"))
576                })?;
577        Ok(Self {
578            model: std::sync::Mutex::new(model),
579            model_name: format!("fastembed:{model_alias}"),
580            dimension,
581            max_length: 512,
582        })
583    }
584}
585
586#[cfg(feature = "fastembed")]
587impl EmbeddingProvider for FastEmbedProvider {
588    fn model_name(&self) -> &str {
589        &self.model_name
590    }
591    fn dimension(&self) -> usize {
592        self.dimension
593    }
594    fn max_length(&self) -> usize {
595        self.max_length
596    }
597    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
598        let mut m = self
599            .model
600            .lock()
601            .map_err(|_| EmbeddingError::ProviderError("embedder mutex poisoned".into()))?;
602        let mut out = m
603            .embed(vec![text], None)
604            .map_err(|e| EmbeddingError::ProviderError(format!("fastembed embed failed: {e}")))?;
605        out.pop()
606            .ok_or_else(|| EmbeddingError::ProviderError("fastembed returned no embedding".into()))
607    }
608    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
609        let mut m = self
610            .model
611            .lock()
612            .map_err(|_| EmbeddingError::ProviderError("embedder mutex poisoned".into()))?;
613        m.embed(texts.to_vec(), None)
614            .map_err(|e| EmbeddingError::ProviderError(format!("fastembed batch failed: {e}")))
615    }
616}
617
618/// Build an embedding provider from the `SOCHDB_EMBEDDER` environment variable.
619///
620/// Accepted values:
621/// * `fastembed:<model>` — real semantic embeddings (requires the `fastembed`
622///   feature; e.g. `fastembed:bge-small-en`).
623/// * `mock` / `hash` / unset — deterministic [`MockEmbeddingProvider`] (384-d).
624///
625/// Falls back to the mock provider (with a warning) when `fastembed` is requested
626/// but the binary was built without the feature, or the model fails to load — so
627/// the server always boots rather than hard-failing on the embedder.
628pub fn embedder_from_env() -> std::sync::Arc<dyn EmbeddingProvider> {
629    let spec = std::env::var("SOCHDB_EMBEDDER").unwrap_or_default();
630    embedder_from_spec(&spec)
631}
632
633/// Like [`embedder_from_env`] but from an explicit spec string.
634pub fn embedder_from_spec(spec: &str) -> std::sync::Arc<dyn EmbeddingProvider> {
635    let spec = spec.trim();
636    if let Some(model) = spec.strip_prefix("fastembed:") {
637        #[cfg(feature = "fastembed")]
638        {
639            match FastEmbedProvider::new(model) {
640                Ok(p) => {
641                    tracing::info!("memory embedder: fastembed:{model} (dim={})", p.dimension());
642                    return std::sync::Arc::new(p);
643                }
644                Err(e) => {
645                    tracing::warn!("fastembed:{model} unavailable ({e}); using mock embedder");
646                }
647            }
648        }
649        #[cfg(not(feature = "fastembed"))]
650        {
651            tracing::warn!(
652                "SOCHDB_EMBEDDER=fastembed:{model} but this binary was built without the \
653                 `fastembed` feature; using mock embedder"
654            );
655        }
656    } else {
657        tracing::info!("memory embedder: mock (384-d) [SOCHDB_EMBEDDER={spec:?}]");
658    }
659    std::sync::Arc::new(MockEmbeddingProvider::new(384))
660}
661
662// ============================================================================
663// Embedding-Enabled Vector Index
664// ============================================================================
665
666/// Vector index with automatic text embedding
667pub struct EmbeddingVectorIndex<V, P>
668where
669    V: crate::context_query::VectorIndex,
670    P: EmbeddingProvider,
671{
672    /// Underlying vector index
673    index: Arc<V>,
674
675    /// Embedding provider
676    provider: Arc<P>,
677}
678
679impl<V, P> EmbeddingVectorIndex<V, P>
680where
681    V: crate::context_query::VectorIndex,
682    P: EmbeddingProvider,
683{
684    /// Create a new embedding-enabled vector index
685    pub fn new(index: Arc<V>, provider: Arc<P>) -> Self {
686        Self { index, provider }
687    }
688
689    /// Search by text (automatically generates embedding)
690    pub fn search_text(
691        &self,
692        collection: &str,
693        text: &str,
694        k: usize,
695        min_score: Option<f32>,
696    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
697        // Generate embedding
698        let embedding = self.provider.embed(text).map_err(|e| e.to_string())?;
699
700        // Search by embedding
701        self.index
702            .search_by_embedding(collection, &embedding, k, min_score)
703    }
704
705    /// Search by embedding (pass-through)
706    pub fn search_embedding(
707        &self,
708        collection: &str,
709        embedding: &[f32],
710        k: usize,
711        min_score: Option<f32>,
712    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
713        // Validate dimension
714        if embedding.len() != self.provider.dimension() {
715            return Err(format!(
716                "Embedding dimension mismatch: expected {}, got {}",
717                self.provider.dimension(),
718                embedding.len()
719            ));
720        }
721
722        self.index
723            .search_by_embedding(collection, embedding, k, min_score)
724    }
725
726    /// Get the embedding provider
727    pub fn provider(&self) -> &Arc<P> {
728        &self.provider
729    }
730
731    /// Get the underlying index
732    pub fn index(&self) -> &Arc<V> {
733        &self.index
734    }
735}
736
737impl<V, P> crate::context_query::VectorIndex for EmbeddingVectorIndex<V, P>
738where
739    V: crate::context_query::VectorIndex,
740    P: EmbeddingProvider,
741{
742    fn search_by_embedding(
743        &self,
744        collection: &str,
745        embedding: &[f32],
746        k: usize,
747        min_score: Option<f32>,
748    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
749        self.search_embedding(collection, embedding, k, min_score)
750    }
751
752    fn search_by_text(
753        &self,
754        collection: &str,
755        text: &str,
756        k: usize,
757        min_score: Option<f32>,
758    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
759        self.search_text(collection, text, k, min_score)
760    }
761
762    fn stats(&self, collection: &str) -> Option<crate::context_query::VectorIndexStats> {
763        self.index.stats(collection)
764    }
765}
766
767// ============================================================================
768// Convenience Functions
769// ============================================================================
770
771/// Create a cached mock embedding provider for testing
772pub fn create_mock_provider(
773    dimension: usize,
774    cache_size: usize,
775) -> CachedEmbeddingProvider<MockEmbeddingProvider> {
776    let mock = MockEmbeddingProvider::new(dimension);
777    CachedEmbeddingProvider::new(mock, cache_size)
778}
779
780/// Create an embedding-enabled vector index with mock provider
781pub fn create_embedding_index<V: crate::context_query::VectorIndex>(
782    index: Arc<V>,
783    dimension: usize,
784) -> EmbeddingVectorIndex<V, CachedEmbeddingProvider<MockEmbeddingProvider>> {
785    let provider = Arc::new(create_mock_provider(dimension, 10_000));
786    EmbeddingVectorIndex::new(index, provider)
787}
788
789// ============================================================================
790// Tests
791// ============================================================================
792
793#[cfg(test)]
794mod tests {
795    use super::*;
796
797    #[test]
798    fn test_mock_embedding_deterministic() {
799        let provider = MockEmbeddingProvider::new(384);
800
801        let emb1 = provider.embed("hello world").unwrap();
802        let emb2 = provider.embed("hello world").unwrap();
803
804        assert_eq!(emb1, emb2);
805        assert_eq!(emb1.len(), 384);
806    }
807
808    #[test]
809    fn test_mock_embedding_different_texts() {
810        let provider = MockEmbeddingProvider::new(384);
811
812        let emb1 = provider.embed("hello").unwrap();
813        let emb2 = provider.embed("world").unwrap();
814
815        assert_ne!(emb1, emb2);
816    }
817
818    #[test]
819    fn test_cached_provider() {
820        let mock = MockEmbeddingProvider::new(128);
821        let cached = CachedEmbeddingProvider::new(mock, 100);
822
823        // First call - miss
824        let _ = cached.embed("test text").unwrap();
825        assert_eq!(
826            cached
827                .stats()
828                .hits
829                .load(std::sync::atomic::Ordering::Relaxed),
830            0
831        );
832        assert_eq!(
833            cached
834                .stats()
835                .misses
836                .load(std::sync::atomic::Ordering::Relaxed),
837            1
838        );
839
840        // Second call - hit
841        let _ = cached.embed("test text").unwrap();
842        assert_eq!(
843            cached
844                .stats()
845                .hits
846                .load(std::sync::atomic::Ordering::Relaxed),
847            1
848        );
849        assert_eq!(
850            cached
851                .stats()
852                .misses
853                .load(std::sync::atomic::Ordering::Relaxed),
854            1
855        );
856
857        assert!(cached.stats().hit_rate() > 0.4);
858    }
859
860    #[test]
861    fn test_batch_embedding() {
862        let mock = MockEmbeddingProvider::new(128);
863        let cached = CachedEmbeddingProvider::new(mock, 100);
864
865        let texts = vec!["hello", "world", "test"];
866        let embeddings = cached.embed_batch(&texts).unwrap();
867
868        assert_eq!(embeddings.len(), 3);
869        for emb in &embeddings {
870            assert_eq!(emb.len(), 128);
871        }
872    }
873
874    #[test]
875    fn test_normalization() {
876        let provider = MockEmbeddingProvider::new(3);
877        let emb = provider.embed("test").unwrap();
878
879        // Check L2 norm is approximately 1
880        let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
881        assert!((norm - 1.0).abs() < 1e-5);
882    }
883
884    #[test]
885    fn test_text_too_long() {
886        let config = EmbeddingConfig {
887            max_length: 10,
888            ..Default::default()
889        };
890        let provider = MockEmbeddingProvider::with_config(config);
891
892        let result = provider.embed("this is a very long text that exceeds the limit");
893        assert!(matches!(result, Err(EmbeddingError::TextTooLong { .. })));
894    }
895
896    #[cfg(feature = "fastembed")]
897    #[test]
898    fn fastembed_provider_real_semantic_embeddings() {
899        // Downloads + caches bge-small-en on first run. Validates: correct dim,
900        // non-trivial (non-zero) vectors, and real semantics (synonyms rank more
901        // similar than unrelated text) — i.e. NOT the mock fallback.
902        let p = FastEmbedProvider::new("bge-small-en").expect("load bge-small-en");
903        assert_eq!(p.dimension(), 384);
904        let cat = p.embed("a cat sat on the mat").unwrap();
905        let feline = p.embed("a feline rested on the rug").unwrap();
906        let finance = p.embed("quarterly financial earnings report").unwrap();
907        assert_eq!(cat.len(), 384);
908        assert!(
909            cat.iter().any(|&x| x.abs() > 1e-6),
910            "embedding must be non-zero"
911        );
912        let cos = |x: &[f32], y: &[f32]| {
913            let d: f32 = x.iter().zip(y).map(|(a, b)| a * b).sum();
914            let nx: f32 = x.iter().map(|a| a * a).sum::<f32>().sqrt();
915            let ny: f32 = y.iter().map(|a| a * a).sum::<f32>().sqrt();
916            d / (nx * ny + 1e-9)
917        };
918        assert!(
919            cos(&cat, &feline) > cos(&cat, &finance),
920            "synonyms ({}) should outrank unrelated text ({})",
921            cos(&cat, &feline),
922            cos(&cat, &finance)
923        );
924
925        // factory path must yield a real provider (not mock) under the feature
926        let from_spec = embedder_from_spec("fastembed:bge-small-en");
927        assert!(from_spec.model_name().starts_with("fastembed:"));
928    }
929}
sochdb_query/embedding_provider.rs

sochdb_query/
embedding_provider.rs
