kreuzberg 4.0.0.pre.rc.6

data/vendor/kreuzberg/src/embeddings.rs

@@ -0,0 +1,432 @@
+//! Embedding generation support for RAG (Retrieval-Augmented Generation) systems.
+//!
+//! This module provides text embedding generation using ONNX models via fastembed-rs.
+//! Embeddings can be generated for text chunks to enable semantic search and RAG pipelines.
+//!
+//! # Features
+//!
+//! - Multiple pre-configured models optimized for different use cases
+//! - Preset configurations for common RAG scenarios
+//! - Full customization of model location and parameters
+//! - Batch processing for efficient embedding generation
+//! - Optional GPU acceleration via ONNX Runtime execution providers
+//!
+//! # Example
+//!
+//! ```rust,ignore
+//! use kreuzberg::{extract_file, ExtractionConfig, ChunkingConfig, EmbeddingConfig};
+//!
+//! let config = ExtractionConfig {
+//!     chunking: Some(ChunkingConfig {
+//!         preset: Some("balanced".to_string()),
+//!         embedding: Some(EmbeddingConfig::default()),
+//!         ..Default::default()
+//!     }),
+//!     ..Default::default()
+//! };
+//!
+//! let result = extract_file("document.pdf", None, &config).await?;
+//! for chunk in result.chunks.unwrap() {
+//!     if let Some(embedding) = chunk.embedding {
+//!         println!("Chunk has {} dimension embedding", embedding.len());
+//!     }
+//! }
+//! ```
+
+#[cfg(feature = "embeddings")]
+use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
+
+#[cfg(feature = "embeddings")]
+use std::sync::{Arc, Mutex, RwLock};
+
+#[cfg(feature = "embeddings")]
+use std::collections::HashMap;
+
+#[cfg(feature = "embeddings")]
+use lazy_static::lazy_static;
+
+/// Wrapper for TextEmbedding that prevents cleanup during process shutdown.
+///
+/// # Problem
+///
+/// When the process terminates, global static objects are dropped. The `TextEmbedding`
+/// objects from fastembed contain ONNX Runtime sessions (via `ort v2.0.0-rc.10`), and
+/// during their `Drop` implementation, ONNX Runtime's C++ destructor tries to acquire
+/// mutexes for cleanup.
+///
+/// At process shutdown time, the C++ runtime may have already begun tearing down
+/// threading infrastructure, causing mutex operations to fail with:
+/// "mutex lock failed: Invalid argument"
+///
+/// This manifests as:
+/// ```text
+/// libc++abi: terminating due to uncaught exception of type std::__1::system_error:
+/// mutex lock failed: Invalid argument
+/// ```
+///
+/// This is a known issue in `ort` (see pykeio/ort#441), fixed in later versions via commit
+/// 317be20 ("fix: let `Environment` drop"), but we're using v2.0.0-rc.10 through fastembed
+/// v5.3.1 which predates the fix.
+///
+/// # Solution
+///
+/// We use `Box::leak` to intentionally leak `TextEmbedding` objects during process
+/// shutdown, preventing their `Drop` implementation from running. This is acceptable because:
+///
+/// 1. The OS will reclaim all process memory anyway
+/// 2. Avoiding the crash is more important than cleanup
+/// 3. This only affects process termination, not runtime behavior
+/// 4. Models are long-lived and would survive until process exit anyway
+/// 5. The memory leak is bounded (one model per unique config)
+///
+/// # Remaining Issue
+///
+/// Even with this fix, you may still see the mutex error during final process cleanup.
+/// This is because `ort` v2.0.0-rc.10 also holds the ONNX Runtime `Environment` as a
+/// static variable, which gets dropped during C++ static destruction after Rust cleanup.
+/// This error occurs *after* all Rust code has finished and can be safely ignored - all
+/// tests pass before the error occurs.
+///
+/// The error will be resolved when fastembed upgrades to ort >= 2.0.0 (post-rc.10) which
+/// contains the proper fix.
+///
+/// # Safety
+///
+/// The leak is contained to process shutdown and does not affect runtime behavior.
+/// All normal usage patterns (creating embeddings, caching models) work identically.
+/// We use static references to the leaked models, which is safe because:
+/// - The pointers are never null (we leak valid Box<TextEmbedding>)
+/// - The models live until process exit
+/// - We never manually deallocate the leaked memory
+/// - Mutex provides interior mutability for the embed() method
+///
+/// Thread-safe wrapper for leaked TextEmbedding that allows interior mutability.
+///
+/// This wrapper holds a raw pointer to a leaked `TextEmbedding` and provides
+/// safe access through the Mutex lock in MODEL_CACHE.
+#[cfg(feature = "embeddings")]
+pub(crate) struct LeakedModel {
+    ptr: *mut TextEmbedding,
+}
+
+#[cfg(feature = "embeddings")]
+impl LeakedModel {
+    fn new(model: TextEmbedding) -> Self {
+        Self {
+            ptr: Box::into_raw(Box::new(model)),
+        }
+    }
+
+    /// Get a mutable reference to the model.
+    ///
+    /// # Safety
+    ///
+    /// This is safe to call only when:
+    /// 1. The caller has exclusive access (guaranteed by Mutex in MODEL_CACHE)
+    /// 2. The pointer is valid (guaranteed by Box::into_raw and never deallocating)
+    #[allow(unsafe_code, clippy::mut_from_ref)]
+    unsafe fn get_mut(&self) -> &mut TextEmbedding {
+        unsafe { &mut *self.ptr }
+    }
+}
+
+#[cfg(feature = "embeddings")]
+#[allow(unsafe_code)]
+unsafe impl Send for LeakedModel {}
+#[cfg(feature = "embeddings")]
+#[allow(unsafe_code)]
+unsafe impl Sync for LeakedModel {}
+
+#[cfg(feature = "embeddings")]
+type CachedEmbedding = Arc<Mutex<LeakedModel>>;
+
+#[cfg(feature = "embeddings")]
+lazy_static! {
+    static ref MODEL_CACHE: RwLock<HashMap<String, CachedEmbedding>> = RwLock::new(HashMap::new());
+}
+
+/// Get or initialize a text embedding model from cache.
+///
+/// This function ensures models are initialized only once and reused across
+/// the application, avoiding redundant downloads and initialization overhead.
+#[cfg(feature = "embeddings")]
+#[allow(private_interfaces)]
+pub fn get_or_init_model(
+    model: EmbeddingModel,
+    cache_dir: Option<std::path::PathBuf>,
+) -> crate::Result<CachedEmbedding> {
+    let cache_directory = cache_dir.unwrap_or_else(|| {
+        let mut path = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from("."));
+        path.push(".kreuzberg");
+        path.push("embeddings");
+        path
+    });
+
+    let model_key = format!("{:?}_{}", model, cache_directory.display());
+
+    {
+        match MODEL_CACHE.read() {
+            Ok(cache) => {
+                if let Some(cached_model) = cache.get(&model_key) {
+                    return Ok(Arc::clone(cached_model));
+                }
+            }
+            Err(poison_error) => {
+                let cache = poison_error.get_ref();
+                if let Some(cached_model) = cache.get(&model_key) {
+                    return Ok(Arc::clone(cached_model));
+                }
+            }
+        }
+    }
+
+    {
+        let mut cache = match MODEL_CACHE.write() {
+            Ok(guard) => guard,
+            Err(poison_error) => poison_error.into_inner(),
+        };
+
+        if let Some(cached_model) = cache.get(&model_key) {
+            return Ok(Arc::clone(cached_model));
+        }
+
+        let mut init_options = InitOptions::new(model);
+        init_options = init_options.with_cache_dir(cache_directory);
+
+        let embedding_model = TextEmbedding::try_new(init_options).map_err(|e| crate::KreuzbergError::Plugin {
+            message: format!("Failed to initialize embedding model: {}", e),
+            plugin_name: "embeddings".to_string(),
+        })?;
+
+        let leaked_model = LeakedModel::new(embedding_model);
+        let arc_model = Arc::new(Mutex::new(leaked_model));
+        cache.insert(model_key, Arc::clone(&arc_model));
+
+        Ok(arc_model)
+    }
+}
+
+/// Preset configurations for common RAG use cases.
+///
+/// Each preset combines chunk size, overlap, and embedding model
+/// to provide an optimized configuration for specific scenarios.
+#[derive(Debug, Clone)]
+pub struct EmbeddingPreset {
+    pub name: &'static str,
+    pub chunk_size: usize,
+    pub overlap: usize,
+    #[cfg(feature = "embeddings")]
+    pub model: EmbeddingModel,
+    #[cfg(not(feature = "embeddings"))]
+    pub model_name: &'static str,
+    pub dimensions: usize,
+    pub description: &'static str,
+}
+
+/// All available embedding presets.
+pub const EMBEDDING_PRESETS: &[EmbeddingPreset] = &[
+    EmbeddingPreset {
+        name: "fast",
+        chunk_size: 512,
+        overlap: 50,
+        #[cfg(feature = "embeddings")]
+        model: EmbeddingModel::AllMiniLML6V2Q,
+        #[cfg(not(feature = "embeddings"))]
+        model_name: "AllMiniLML6V2Q",
+        dimensions: 384,
+        description: "Fast embedding with quantized model (384 dims, ~22M params). Best for: Quick prototyping, development, resource-constrained environments.",
+    },
+    EmbeddingPreset {
+        name: "balanced",
+        chunk_size: 1024,
+        overlap: 100,
+        #[cfg(feature = "embeddings")]
+        model: EmbeddingModel::BGEBaseENV15,
+        #[cfg(not(feature = "embeddings"))]
+        model_name: "BGEBaseENV15",
+        dimensions: 768,
+        description: "Balanced quality and speed (768 dims, ~109M params). Best for: General-purpose RAG, production deployments, English documents.",
+    },
+    EmbeddingPreset {
+        name: "quality",
+        chunk_size: 2000,
+        overlap: 200,
+        #[cfg(feature = "embeddings")]
+        model: EmbeddingModel::BGELargeENV15,
+        #[cfg(not(feature = "embeddings"))]
+        model_name: "BGELargeENV15",
+        dimensions: 1024,
+        description: "High quality with larger context (1024 dims, ~335M params). Best for: Complex documents, maximum accuracy, sufficient compute resources.",
+    },
+    EmbeddingPreset {
+        name: "multilingual",
+        chunk_size: 1024,
+        overlap: 100,
+        #[cfg(feature = "embeddings")]
+        model: EmbeddingModel::MultilingualE5Base,
+        #[cfg(not(feature = "embeddings"))]
+        model_name: "MultilingualE5Base",
+        dimensions: 768,
+        description: "Multilingual support (768 dims, 100+ languages). Best for: International documents, mixed-language content, global applications.",
+    },
+];
+
+/// Get a preset by name.
+pub fn get_preset(name: &str) -> Option<&'static EmbeddingPreset> {
+    EMBEDDING_PRESETS.iter().find(|p| p.name == name)
+}
+
+/// List all available preset names.
+pub fn list_presets() -> Vec<&'static str> {
+    EMBEDDING_PRESETS.iter().map(|p| p.name).collect()
+}
+
+/// Generate embeddings for text chunks using the specified configuration.
+///
+/// This function modifies chunks in-place, populating their `embedding` field
+/// with generated embedding vectors. It uses batch processing for efficiency.
+///
+/// # Arguments
+///
+/// * `chunks` - Mutable reference to vector of chunks to generate embeddings for
+/// * `config` - Embedding configuration specifying model and parameters
+///
+/// # Returns
+///
+/// Returns `Ok(())` if embeddings were generated successfully, or an error if
+/// model initialization or embedding generation fails.
+///
+/// # Example
+///
+/// ```rust,ignore
+/// let mut chunks = vec![
+///     Chunk { content: "Hello world".to_string(), embedding: None, metadata: ... },
+///     Chunk { content: "Second chunk".to_string(), embedding: None, metadata: ... },
+/// ];
+/// let config = EmbeddingConfig::default();
+/// generate_embeddings_for_chunks(&mut chunks, &config)?;
+/// // Now chunks have embeddings populated
+/// ```
+#[cfg(feature = "embeddings")]
+pub fn generate_embeddings_for_chunks(
+    chunks: &mut [crate::types::Chunk],
+    config: &crate::core::config::EmbeddingConfig,
+) -> crate::Result<()> {
+    if chunks.is_empty() {
+        return Ok(());
+    }
+
+    let fastembed_model = match &config.model {
+        crate::core::config::EmbeddingModelType::Preset { name } => {
+            let preset = get_preset(name).ok_or_else(|| crate::KreuzbergError::Plugin {
+                message: format!("Unknown embedding preset: {}", name),
+                plugin_name: "embeddings".to_string(),
+            })?;
+            preset.model.clone()
+        }
+        #[cfg(feature = "embeddings")]
+        crate::core::config::EmbeddingModelType::FastEmbed { model, .. } => match model.as_str() {
+            "AllMiniLML6V2Q" => fastembed::EmbeddingModel::AllMiniLML6V2Q,
+            "BGEBaseENV15" => fastembed::EmbeddingModel::BGEBaseENV15,
+            "BGELargeENV15" => fastembed::EmbeddingModel::BGELargeENV15,
+            "MultilingualE5Base" => fastembed::EmbeddingModel::MultilingualE5Base,
+            _ => {
+                return Err(crate::KreuzbergError::Plugin {
+                    message: format!("Unknown fastembed model: {}", model),
+                    plugin_name: "embeddings".to_string(),
+                });
+            }
+        },
+        crate::core::config::EmbeddingModelType::Custom { .. } => {
+            return Err(crate::KreuzbergError::Plugin {
+                message: "Custom ONNX models are not yet supported for embedding generation".to_string(),
+                plugin_name: "embeddings".to_string(),
+            });
+        }
+    };
+
+    let model = get_or_init_model(fastembed_model, config.cache_dir.clone())?;
+
+    let texts: Vec<String> = chunks.iter().map(|chunk| chunk.content.clone()).collect();
+
+    let embeddings_result = {
+        let locked_model = model.lock().map_err(|e| crate::KreuzbergError::Plugin {
+            message: format!("Failed to acquire model lock: {}", e),
+            plugin_name: "embeddings".to_string(),
+        })?;
+
+        #[allow(unsafe_code)]
+        let model_mut = unsafe { locked_model.get_mut() };
+
+        model_mut
+            .embed(texts, Some(config.batch_size))
+            .map_err(|e| crate::KreuzbergError::Plugin {
+                message: format!("Failed to generate embeddings: {}", e),
+                plugin_name: "embeddings".to_string(),
+            })?
+    };
+
+    for (chunk, mut embedding) in chunks.iter_mut().zip(embeddings_result.into_iter()) {
+        if config.normalize {
+            let magnitude: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
+            if magnitude > 0.0 {
+                embedding.iter_mut().for_each(|x| *x /= magnitude);
+            }
+        }
+
+        chunk.embedding = Some(embedding);
+    }
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_get_preset() {
+        assert!(get_preset("balanced").is_some());
+        assert!(get_preset("fast").is_some());
+        assert!(get_preset("quality").is_some());
+        assert!(get_preset("multilingual").is_some());
+        assert!(get_preset("nonexistent").is_none());
+    }
+
+    #[test]
+    fn test_list_presets() {
+        let presets = list_presets();
+        assert_eq!(presets.len(), 4);
+        assert!(presets.contains(&"fast"));
+        assert!(presets.contains(&"balanced"));
+        assert!(presets.contains(&"quality"));
+        assert!(presets.contains(&"multilingual"));
+    }
+
+    #[test]
+    fn test_preset_dimensions() {
+        let balanced = get_preset("balanced").unwrap();
+        assert_eq!(balanced.dimensions, 768);
+
+        let fast = get_preset("fast").unwrap();
+        assert_eq!(fast.dimensions, 384);
+
+        let quality = get_preset("quality").unwrap();
+        assert_eq!(quality.dimensions, 1024);
+    }
+
+    #[test]
+    fn test_preset_chunk_sizes() {
+        let fast = get_preset("fast").unwrap();
+        assert_eq!(fast.chunk_size, 512);
+        assert_eq!(fast.overlap, 50);
+
+        let quality = get_preset("quality").unwrap();
+        assert_eq!(quality.chunk_size, 2000);
+        assert_eq!(quality.overlap, 200);
+    }
+
+    #[cfg(feature = "embeddings")]
+    #[test]
+    fn test_lock_poisoning_recovery_semantics() {}
+}