sochdb_query/

embedding_provider.rs

// SPDX-License-Identifier: AGPL-3.0-or-later
// SochDB - LLM-Optimized Embedded Database
// Copyright (C) 2026 Sushanth Reddy Vanagala (https://github.com/sushanthpy)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Automatic Embedding Generation (Task 2)
//!
//! This module provides colocated embedding resolution for text-to-vector conversion.
//! It enables first-class text search by automatically generating embeddings.
//!
//! ## Design
//!
//! ```text
//! search_text(collection, text, k)
//!     │
//!     ▼
//! ┌─────────────────┐
//! │ EmbeddingProvider │
//! │  ├─ LRU Cache    │
//! │  └─ ONNX Runtime │
//! └─────────────────┘
//!     │
//!     ▼
//! search_by_embedding(collection, embedding, k)
//! ```
//!
//! ## Providers
//!
//! - `LocalProvider`: Uses FastEmbed/ONNX for offline embedding
//! - `CachedProvider`: LRU cache wrapper for any provider
//! - `MockProvider`: For testing
//!
//! ## Complexity
//!
//! - Embedding generation: O(n) where n = text length (transformer inference)
//! - Cache lookup: O(1) expected (hash-based LRU)
//! - Batch embedding: O(k) compute with ~O(1) ONNX session overhead

use moka::sync::Cache;
use std::sync::Arc;

// ============================================================================
// Embedding Provider Trait
// ============================================================================

/// Error type for embedding operations
#[derive(Debug, Clone)]
pub enum EmbeddingError {
    /// Model not loaded or unavailable
    ModelNotAvailable(String),
    /// Text too long for model
    TextTooLong { max_length: usize, actual: usize },
    /// Dimension mismatch
    DimensionMismatch { expected: usize, actual: usize },
    /// Provider error
    ProviderError(String),
    /// Cache error
    CacheError(String),
}

impl std::fmt::Display for EmbeddingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::ModelNotAvailable(model) => write!(f, "Embedding model not available: {}", model),
            Self::TextTooLong { max_length, actual } => {
                write!(f, "Text too long: {} > {} max", actual, max_length)
            }
            Self::DimensionMismatch { expected, actual } => {
                write!(
                    f,
                    "Dimension mismatch: expected {}, got {}",
                    expected, actual
                )
            }
            Self::ProviderError(msg) => write!(f, "Provider error: {}", msg),
            Self::CacheError(msg) => write!(f, "Cache error: {}", msg),
        }
    }
}

impl std::error::Error for EmbeddingError {}

/// Result type for embedding operations
pub type EmbeddingResult<T> = Result<T, EmbeddingError>;

/// Embedding provider trait
pub trait EmbeddingProvider: Send + Sync {
    /// Get the model name
    fn model_name(&self) -> &str;

    /// Get the embedding dimension
    fn dimension(&self) -> usize;

    /// Maximum text length (in characters or tokens)
    fn max_length(&self) -> usize;

    /// Whether this provider produces SEMANTICALLY meaningful embeddings.
    ///
    /// Real model-backed providers return `true` (the default). The hash-based
    /// `MockEmbeddingProvider` returns `false` — callers use this to avoid fusing
    /// a meaningless cosine lane into hybrid results (which would corrupt recall
    /// silently). Defaulted so real providers need not implement it.
    fn is_semantic(&self) -> bool {
        true
    }

    /// Generate embedding for a single text (document side).
    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>>;

    /// Embed a QUERY for asymmetric retrieval. Defaults to `embed` (symmetric).
    /// Providers for instruction-tuned models (e.g. BGE) override this to apply
    /// the query-side instruction so query and document embeddings align — a
    /// large recall lever for those models.
    fn embed_query(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        self.embed(text)
    }

    /// Generate embeddings for multiple texts (batch)
    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
        // Default implementation: sequential embedding
        texts.iter().map(|t| self.embed(t)).collect()
    }

    /// Normalize an embedding vector (L2 normalization)
    fn normalize(&self, embedding: &mut [f32]) {
        let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 1e-10 {
            for x in embedding.iter_mut() {
                *x /= norm;
            }
        }
    }
}

// ============================================================================
// Embedding Configuration
// ============================================================================

/// Configuration for embedding providers
#[derive(Debug, Clone)]
pub struct EmbeddingConfig {
    /// Model identifier (e.g., "all-MiniLM-L6-v2")
    pub model: String,

    /// Model path (for local ONNX models)
    pub model_path: Option<String>,

    /// Embedding dimension
    pub dimension: usize,

    /// Maximum text length
    pub max_length: usize,

    /// Whether to normalize embeddings
    pub normalize: bool,

    /// Batch size for embedding generation
    pub batch_size: usize,

    /// Cache size (number of embeddings to cache)
    pub cache_size: usize,

    /// Cache TTL in seconds (0 = no expiry)
    pub cache_ttl_secs: u64,
}

impl Default for EmbeddingConfig {
    fn default() -> Self {
        Self {
            model: "all-MiniLM-L6-v2".to_string(),
            model_path: None,
            dimension: 384, // MiniLM dimension
            max_length: 512,
            normalize: true,
            batch_size: 32,
            cache_size: 10_000,
            cache_ttl_secs: 3600, // 1 hour
        }
    }
}

impl EmbeddingConfig {
    /// Create config for sentence-transformers models
    pub fn sentence_transformer(model: &str) -> Self {
        let dimension = match model {
            "all-MiniLM-L6-v2" => 384,
            "all-MiniLM-L12-v2" => 384,
            "all-mpnet-base-v2" => 768,
            "paraphrase-MiniLM-L6-v2" => 384,
            "multi-qa-MiniLM-L6-cos-v1" => 384,
            _ => 384, // Default
        };

        Self {
            model: model.to_string(),
            dimension,
            ..Default::default()
        }
    }

    /// Create config for OpenAI-compatible models
    pub fn openai(model: &str) -> Self {
        let dimension = match model {
            "text-embedding-ada-002" => 1536,
            "text-embedding-3-small" => 1536,
            "text-embedding-3-large" => 3072,
            _ => 1536,
        };

        Self {
            model: model.to_string(),
            dimension,
            max_length: 8192,
            ..Default::default()
        }
    }
}

// ============================================================================
// Mock Embedding Provider (for testing)
// ============================================================================

/// Mock embedding provider for testing
pub struct MockEmbeddingProvider {
    config: EmbeddingConfig,
    /// Deterministic embeddings based on text hash
    use_hash: bool,
}

impl MockEmbeddingProvider {
    /// Create a new mock provider
    pub fn new(dimension: usize) -> Self {
        Self {
            config: EmbeddingConfig {
                model: "mock".to_string(),
                dimension,
                ..Default::default()
            },
            use_hash: true,
        }
    }

    /// Create with custom config
    pub fn with_config(config: EmbeddingConfig) -> Self {
        Self {
            config,
            use_hash: true,
        }
    }

    /// Generate a deterministic embedding from text
    fn hash_embed(&self, text: &str) -> Vec<f32> {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut embedding = Vec::with_capacity(self.config.dimension);

        // Generate pseudo-random values based on text hash
        for i in 0..self.config.dimension {
            let mut hasher = DefaultHasher::new();
            text.hash(&mut hasher);
            i.hash(&mut hasher);
            let hash = hasher.finish();

            // Convert to f32 in range [-1, 1]
            let value = ((hash as f64) / (u64::MAX as f64) * 2.0 - 1.0) as f32;
            embedding.push(value);
        }

        embedding
    }
}

impl EmbeddingProvider for MockEmbeddingProvider {
    fn model_name(&self) -> &str {
        &self.config.model
    }

    fn dimension(&self) -> usize {
        self.config.dimension
    }

    fn max_length(&self) -> usize {
        self.config.max_length
    }

    /// Hash-based embeddings have no semantic meaning — cosine over them is
    /// noise, so hybrid retrieval must NOT fuse the vector lane when this
    /// provider is in use.
    fn is_semantic(&self) -> bool {
        false
    }

    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        if text.len() > self.config.max_length {
            return Err(EmbeddingError::TextTooLong {
                max_length: self.config.max_length,
                actual: text.len(),
            });
        }

        let mut embedding = if self.use_hash {
            self.hash_embed(text)
        } else {
            vec![0.0; self.config.dimension]
        };

        if self.config.normalize {
            self.normalize(&mut embedding);
        }

        Ok(embedding)
    }
}

// ============================================================================
// Cached Embedding Provider
// ============================================================================

/// LRU-cached embedding provider wrapper
pub struct CachedEmbeddingProvider<P: EmbeddingProvider> {
    /// Inner provider
    inner: P,

    /// LRU cache: text hash -> embedding
    cache: Cache<u64, Vec<f32>>,

    /// Cache statistics
    stats: Arc<CacheStats>,
}

/// Cache statistics
#[derive(Debug, Default)]
pub struct CacheStats {
    /// Number of cache hits
    pub hits: std::sync::atomic::AtomicUsize,
    /// Number of cache misses
    pub misses: std::sync::atomic::AtomicUsize,
    /// Number of embeddings cached
    pub size: std::sync::atomic::AtomicUsize,
}

impl CacheStats {
    /// Get hit rate
    pub fn hit_rate(&self) -> f64 {
        let hits = self.hits.load(std::sync::atomic::Ordering::Relaxed);
        let misses = self.misses.load(std::sync::atomic::Ordering::Relaxed);
        let total = hits + misses;
        if total == 0 {
            0.0
        } else {
            hits as f64 / total as f64
        }
    }
}

impl<P: EmbeddingProvider> CachedEmbeddingProvider<P> {
    /// Create a new cached provider
    pub fn new(inner: P, cache_size: usize) -> Self {
        Self {
            inner,
            cache: Cache::new(cache_size as u64),
            stats: Arc::new(CacheStats::default()),
        }
    }

    /// Create with TTL
    pub fn with_ttl(inner: P, cache_size: usize, ttl_secs: u64) -> Self {
        let cache = Cache::builder()
            .max_capacity(cache_size as u64)
            .time_to_live(std::time::Duration::from_secs(ttl_secs))
            .build();

        Self {
            inner,
            cache,
            stats: Arc::new(CacheStats::default()),
        }
    }

    /// Get cache statistics
    pub fn stats(&self) -> &Arc<CacheStats> {
        &self.stats
    }

    /// Compute hash for cache key
    fn text_hash(text: &str) -> u64 {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut hasher = DefaultHasher::new();
        text.hash(&mut hasher);
        hasher.finish()
    }
}

impl<P: EmbeddingProvider> EmbeddingProvider for CachedEmbeddingProvider<P> {
    fn model_name(&self) -> &str {
        self.inner.model_name()
    }

    fn dimension(&self) -> usize {
        self.inner.dimension()
    }

    fn max_length(&self) -> usize {
        self.inner.max_length()
    }

    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        let hash = Self::text_hash(text);

        // Check cache
        if let Some(cached) = self.cache.get(&hash) {
            self.stats
                .hits
                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
            return Ok(cached);
        }

        self.stats
            .misses
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        // Generate embedding
        let embedding = self.inner.embed(text)?;

        // Cache result
        self.cache.insert(hash, embedding.clone());
        self.stats
            .size
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        Ok(embedding)
    }

    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
        let mut results = Vec::with_capacity(texts.len());
        let mut uncached: Vec<(usize, &str)> = Vec::new();

        // Check cache for each text
        for (i, text) in texts.iter().enumerate() {
            let hash = Self::text_hash(text);
            if let Some(cached) = self.cache.get(&hash) {
                self.stats
                    .hits
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                results.push((i, cached));
            } else {
                self.stats
                    .misses
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                uncached.push((i, *text));
            }
        }

        // Generate embeddings for uncached texts
        if !uncached.is_empty() {
            let uncached_texts: Vec<&str> = uncached.iter().map(|(_, t)| *t).collect();
            let embeddings = self.inner.embed_batch(&uncached_texts)?;

            for ((i, text), embedding) in uncached.iter().zip(embeddings.into_iter()) {
                let hash = Self::text_hash(text);
                self.cache.insert(hash, embedding.clone());
                self.stats
                    .size
                    .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                results.push((*i, embedding));
            }
        }

        // Sort by original index
        results.sort_by_key(|(i, _)| *i);
        Ok(results.into_iter().map(|(_, e)| e).collect())
    }
}

// ============================================================================
// Local ONNX Provider (Stub)
// ============================================================================

/// Local ONNX-based embedding provider
///
/// This is a stub implementation. In production, this would use:
/// - ort (ONNX Runtime) for model inference
/// - fastembed-rs for pre-packaged models
/// - tokenizers for text preprocessing
#[derive(Debug)]
pub struct LocalOnnxProvider {
    config: EmbeddingConfig,
    /// Model weights (placeholder)
    #[allow(dead_code)]
    model_loaded: bool,
}

impl LocalOnnxProvider {
    /// Create a new local ONNX provider
    pub fn new(config: EmbeddingConfig) -> EmbeddingResult<Self> {
        // In production: Load ONNX model from path
        Ok(Self {
            config,
            model_loaded: false,
        })
    }

    /// Load a pre-trained model by name
    pub fn load_pretrained(model_name: &str) -> EmbeddingResult<Self> {
        let config = EmbeddingConfig::sentence_transformer(model_name);
        Self::new(config)
    }
}

impl EmbeddingProvider for LocalOnnxProvider {
    fn model_name(&self) -> &str {
        &self.config.model
    }

    fn dimension(&self) -> usize {
        self.config.dimension
    }

    fn max_length(&self) -> usize {
        self.config.max_length
    }

    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        // Stub: Return mock embedding
        // In production: Run ONNX inference
        let mock = MockEmbeddingProvider::with_config(self.config.clone());
        mock.embed(text)
    }
}

// ============================================================================
// FastEmbed (ONNX) semantic provider  —  feature = "fastembed"
// ============================================================================

/// Real semantic embedding provider backed by fastembed-rs (ONNX runtime).
///
/// Only compiled with the `fastembed` feature (the `ort` native dep is heavy).
/// The model is downloaded + cached on first construction. `TextEmbedding::embed`
/// takes `&mut self`, so the model is held behind a `Mutex` to satisfy the
/// `&self` trait contract (embedding is the bottleneck anyway, so the lock is
/// not the limiting factor).
#[cfg(feature = "fastembed")]
pub struct FastEmbedProvider {
    model: std::sync::Mutex<fastembed::TextEmbedding>,
    model_name: String,
    dimension: usize,
    max_length: usize,
}

#[cfg(feature = "fastembed")]
impl std::fmt::Debug for FastEmbedProvider {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FastEmbedProvider")
            .field("model_name", &self.model_name)
            .field("dimension", &self.dimension)
            .finish()
    }
}

#[cfg(feature = "fastembed")]
impl FastEmbedProvider {
    /// Construct from a short model alias, e.g. `bge-small-en`, `all-minilm`,
    /// `bge-base-en`, `bge-large-en`. Downloads + caches the ONNX model on first
    /// use (set `FASTEMBED_CACHE_DIR` to control the cache location).
    pub fn new(model_alias: &str) -> EmbeddingResult<Self> {
        use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
        let (model, dimension) = match model_alias.to_ascii_lowercase().as_str() {
            "bge-small-en" | "bge-small-en-v1.5" | "bge-small" => {
                (EmbeddingModel::BGESmallENV15, 384)
            }
            "all-minilm" | "all-minilm-l6-v2" | "minilm" => (EmbeddingModel::AllMiniLML6V2, 384),
            "bge-base-en" | "bge-base-en-v1.5" | "bge-base" => (EmbeddingModel::BGEBaseENV15, 768),
            "bge-large-en" | "bge-large-en-v1.5" | "bge-large" => {
                (EmbeddingModel::BGELargeENV15, 1024)
            }
            other => {
                return Err(EmbeddingError::ModelNotAvailable(format!(
                    "fastembed:{other}"
                )));
            }
        };
        let model =
            TextEmbedding::try_new(InitOptions::new(model).with_show_download_progress(false))
                .map_err(|e| {
                    EmbeddingError::ProviderError(format!("fastembed init failed: {e}"))
                })?;
        Ok(Self {
            model: std::sync::Mutex::new(model),
            model_name: format!("fastembed:{model_alias}"),
            dimension,
            max_length: 512,
        })
    }
}

#[cfg(feature = "fastembed")]
impl EmbeddingProvider for FastEmbedProvider {
    fn model_name(&self) -> &str {
        &self.model_name
    }
    fn dimension(&self) -> usize {
        self.dimension
    }
    fn max_length(&self) -> usize {
        self.max_length
    }
    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        let mut m = self
            .model
            .lock()
            .map_err(|_| EmbeddingError::ProviderError("embedder mutex poisoned".into()))?;
        let mut out = m
            .embed(vec![text], None)
            .map_err(|e| EmbeddingError::ProviderError(format!("fastembed embed failed: {e}")))?;
        out.pop()
            .ok_or_else(|| EmbeddingError::ProviderError("fastembed returned no embedding".into()))
    }
    fn embed_query(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        // BGE v1.5 models are asymmetric: queries are embedded WITH the retrieval
        // instruction prefix, documents without it. Embedding queries as plain
        // documents (no prefix) misaligns them with the indexed docs and depresses
        // recall — more so for the more instruction-tuned bge-base. MiniLM and
        // other symmetric models need no prefix and fall through to `embed`.
        if self.model_name.contains("bge") {
            self.embed(&format!(
                "Represent this sentence for searching relevant passages: {text}"
            ))
        } else {
            self.embed(text)
        }
    }
    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
        let mut m = self
            .model
            .lock()
            .map_err(|_| EmbeddingError::ProviderError("embedder mutex poisoned".into()))?;
        m.embed(texts.to_vec(), None)
            .map_err(|e| EmbeddingError::ProviderError(format!("fastembed batch failed: {e}")))
    }
}

/// Build an embedding provider from the `SOCHDB_EMBEDDER` environment variable.
///
/// Accepted values:
/// * `fastembed:<model>` — real semantic embeddings (requires the `fastembed`
///   feature; e.g. `fastembed:bge-small-en`).
/// * `mock` / `hash` / unset — deterministic [`MockEmbeddingProvider`] (384-d).
///
/// When `SOCHDB_EMBEDDER=fastembed:...` is set but cannot be honored (binary
/// built without the feature, or the model fails to load), this PANICS at
/// startup rather than silently falling back to the non-semantic mock embedder
/// — a silent fallback corrupts retrieval quality without any error. An unset
/// or `mock`/`hash` spec uses the mock provider normally.
pub fn embedder_from_env() -> std::sync::Arc<dyn EmbeddingProvider> {
    let spec = std::env::var("SOCHDB_EMBEDDER").unwrap_or_default();
    embedder_from_spec(&spec)
}

/// Like [`embedder_from_env`] but from an explicit spec string.
pub fn embedder_from_spec(spec: &str) -> std::sync::Arc<dyn EmbeddingProvider> {
    let spec = spec.trim();
    if let Some(model) = spec.strip_prefix("fastembed:") {
        #[cfg(feature = "fastembed")]
        {
            match FastEmbedProvider::new(model) {
                Ok(p) => {
                    tracing::info!("memory embedder: fastembed:{model} (dim={})", p.dimension());
                    return std::sync::Arc::new(p);
                }
                Err(e) => panic!(
                    "SOCHDB_EMBEDDER=fastembed:{model} was requested but the model failed to \
                     load: {e}. Refusing to start with a non-semantic mock embedder, which would \
                     silently corrupt retrieval. Fix the model/cache (FASTEMBED_CACHE_DIR) or \
                     unset SOCHDB_EMBEDDER to use the mock explicitly."
                ),
            }
        }
        #[cfg(not(feature = "fastembed"))]
        {
            panic!(
                "SOCHDB_EMBEDDER=fastembed:{model} was requested but this binary was built \
                 WITHOUT the `fastembed` feature. Rebuild with `--features fastembed`, or unset \
                 SOCHDB_EMBEDDER to use the mock embedder explicitly. Refusing to silently fall \
                 back to a non-semantic mock."
            );
        }
    } else {
        tracing::info!("memory embedder: mock (384-d) [SOCHDB_EMBEDDER={spec:?}]");
    }
    std::sync::Arc::new(MockEmbeddingProvider::new(384))
}

// ============================================================================
// Embedding-Enabled Vector Index
// ============================================================================

/// Vector index with automatic text embedding
pub struct EmbeddingVectorIndex<V, P>
where
    V: crate::context_query::VectorIndex,
    P: EmbeddingProvider,
{
    /// Underlying vector index
    index: Arc<V>,

    /// Embedding provider
    provider: Arc<P>,
}

impl<V, P> EmbeddingVectorIndex<V, P>
where
    V: crate::context_query::VectorIndex,
    P: EmbeddingProvider,
{
    /// Create a new embedding-enabled vector index
    pub fn new(index: Arc<V>, provider: Arc<P>) -> Self {
        Self { index, provider }
    }

    /// Search by text (automatically generates embedding)
    pub fn search_text(
        &self,
        collection: &str,
        text: &str,
        k: usize,
        min_score: Option<f32>,
    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
        // Generate embedding
        let embedding = self.provider.embed(text).map_err(|e| e.to_string())?;

        // Search by embedding
        self.index
            .search_by_embedding(collection, &embedding, k, min_score)
    }

    /// Search by embedding (pass-through)
    pub fn search_embedding(
        &self,
        collection: &str,
        embedding: &[f32],
        k: usize,
        min_score: Option<f32>,
    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
        // Validate dimension
        if embedding.len() != self.provider.dimension() {
            return Err(format!(
                "Embedding dimension mismatch: expected {}, got {}",
                self.provider.dimension(),
                embedding.len()
            ));
        }

        self.index
            .search_by_embedding(collection, embedding, k, min_score)
    }

    /// Get the embedding provider
    pub fn provider(&self) -> &Arc<P> {
        &self.provider
    }

    /// Get the underlying index
    pub fn index(&self) -> &Arc<V> {
        &self.index
    }
}

impl<V, P> crate::context_query::VectorIndex for EmbeddingVectorIndex<V, P>
where
    V: crate::context_query::VectorIndex,
    P: EmbeddingProvider,
{
    fn search_by_embedding(
        &self,
        collection: &str,
        embedding: &[f32],
        k: usize,
        min_score: Option<f32>,
    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
        self.search_embedding(collection, embedding, k, min_score)
    }

    fn search_by_text(
        &self,
        collection: &str,
        text: &str,
        k: usize,
        min_score: Option<f32>,
    ) -> Result<Vec<crate::context_query::VectorSearchResult>, String> {
        self.search_text(collection, text, k, min_score)
    }

    fn stats(&self, collection: &str) -> Option<crate::context_query::VectorIndexStats> {
        self.index.stats(collection)
    }
}

// ============================================================================
// Convenience Functions
// ============================================================================

/// Create a cached mock embedding provider for testing
pub fn create_mock_provider(
    dimension: usize,
    cache_size: usize,
) -> CachedEmbeddingProvider<MockEmbeddingProvider> {
    let mock = MockEmbeddingProvider::new(dimension);
    CachedEmbeddingProvider::new(mock, cache_size)
}

/// Create an embedding-enabled vector index with mock provider
pub fn create_embedding_index<V: crate::context_query::VectorIndex>(
    index: Arc<V>,
    dimension: usize,
) -> EmbeddingVectorIndex<V, CachedEmbeddingProvider<MockEmbeddingProvider>> {
    let provider = Arc::new(create_mock_provider(dimension, 10_000));
    EmbeddingVectorIndex::new(index, provider)
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_mock_embedding_deterministic() {
        let provider = MockEmbeddingProvider::new(384);

        let emb1 = provider.embed("hello world").unwrap();
        let emb2 = provider.embed("hello world").unwrap();

        assert_eq!(emb1, emb2);
        assert_eq!(emb1.len(), 384);
    }

    #[test]
    fn test_mock_embedding_different_texts() {
        let provider = MockEmbeddingProvider::new(384);

        let emb1 = provider.embed("hello").unwrap();
        let emb2 = provider.embed("world").unwrap();

        assert_ne!(emb1, emb2);
    }

    #[test]
    fn test_cached_provider() {
        let mock = MockEmbeddingProvider::new(128);
        let cached = CachedEmbeddingProvider::new(mock, 100);

        // First call - miss
        let _ = cached.embed("test text").unwrap();
        assert_eq!(
            cached
                .stats()
                .hits
                .load(std::sync::atomic::Ordering::Relaxed),
            0
        );
        assert_eq!(
            cached
                .stats()
                .misses
                .load(std::sync::atomic::Ordering::Relaxed),
            1
        );

        // Second call - hit
        let _ = cached.embed("test text").unwrap();
        assert_eq!(
            cached
                .stats()
                .hits
                .load(std::sync::atomic::Ordering::Relaxed),
            1
        );
        assert_eq!(
            cached
                .stats()
                .misses
                .load(std::sync::atomic::Ordering::Relaxed),
            1
        );

        assert!(cached.stats().hit_rate() > 0.4);
    }

    #[test]
    fn test_batch_embedding() {
        let mock = MockEmbeddingProvider::new(128);
        let cached = CachedEmbeddingProvider::new(mock, 100);

        let texts = vec!["hello", "world", "test"];
        let embeddings = cached.embed_batch(&texts).unwrap();

        assert_eq!(embeddings.len(), 3);
        for emb in &embeddings {
            assert_eq!(emb.len(), 128);
        }
    }

    #[test]
    fn test_normalization() {
        let provider = MockEmbeddingProvider::new(3);
        let emb = provider.embed("test").unwrap();

        // Check L2 norm is approximately 1
        let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_text_too_long() {
        let config = EmbeddingConfig {
            max_length: 10,
            ..Default::default()
        };
        let provider = MockEmbeddingProvider::with_config(config);

        let result = provider.embed("this is a very long text that exceeds the limit");
        assert!(matches!(result, Err(EmbeddingError::TextTooLong { .. })));
    }

    #[cfg(feature = "fastembed")]
    #[test]
    fn fastembed_provider_real_semantic_embeddings() {
        // Downloads + caches bge-small-en on first run. Validates: correct dim,
        // non-trivial (non-zero) vectors, and real semantics (synonyms rank more
        // similar than unrelated text) — i.e. NOT the mock fallback.
        let p = FastEmbedProvider::new("bge-small-en").expect("load bge-small-en");
        assert_eq!(p.dimension(), 384);
        let cat = p.embed("a cat sat on the mat").unwrap();
        let feline = p.embed("a feline rested on the rug").unwrap();
        let finance = p.embed("quarterly financial earnings report").unwrap();
        assert_eq!(cat.len(), 384);
        assert!(
            cat.iter().any(|&x| x.abs() > 1e-6),
            "embedding must be non-zero"
        );
        let cos = |x: &[f32], y: &[f32]| {
            let d: f32 = x.iter().zip(y).map(|(a, b)| a * b).sum();
            let nx: f32 = x.iter().map(|a| a * a).sum::<f32>().sqrt();
            let ny: f32 = y.iter().map(|a| a * a).sum::<f32>().sqrt();
            d / (nx * ny + 1e-9)
        };
        assert!(
            cos(&cat, &feline) > cos(&cat, &finance),
            "synonyms ({}) should outrank unrelated text ({})",
            cos(&cat, &feline),
            cos(&cat, &finance)
        );

        // factory path must yield a real provider (not mock) under the feature
        let from_spec = embedder_from_spec("fastembed:bge-small-en");
        assert!(from_spec.model_name().starts_with("fastembed:"));
    }
}