kreuzberg 4.0.8 → 4.1.0

data/vendor/kreuzberg/src/plugins/registry/ocr.rs

@@ -0,0 +1,293 @@
+//! OCR backend registry.
+
+use crate::plugins::OcrBackend;
+use crate::{KreuzbergError, Result};
+use std::collections::HashMap;
+use std::sync::Arc;
+
+/// Registry for OCR backend plugins.
+///
+/// Manages OCR backends with backend type and language-based selection.
+///
+/// # Thread Safety
+///
+/// The registry is thread-safe and can be accessed concurrently from multiple threads.
+///
+/// # Example
+///
+/// ```rust,no_run
+/// use kreuzberg::plugins::registry::OcrBackendRegistry;
+/// use std::sync::Arc;
+///
+/// let registry = OcrBackendRegistry::new();
+/// // Register OCR backends
+/// // registry.register(Arc::new(TesseractBackend::new()));
+/// ```
+pub struct OcrBackendRegistry {
+    pub(super) backends: HashMap<String, Arc<dyn OcrBackend>>,
+}
+
+impl OcrBackendRegistry {
+    /// Create a new OCR backend registry with default backends.
+    ///
+    /// Registers the Tesseract backend by default if the "ocr" feature is enabled.
+    pub fn new() -> Self {
+        #[cfg(feature = "ocr")]
+        let mut registry = Self {
+            backends: HashMap::new(),
+        };
+
+        #[cfg(not(feature = "ocr"))]
+        let registry = Self {
+            backends: HashMap::new(),
+        };
+
+        #[cfg(feature = "ocr")]
+        {
+            use crate::ocr::tesseract_backend::TesseractBackend;
+            if let Ok(backend) = TesseractBackend::new() {
+                let _ = registry.register(Arc::new(backend));
+            }
+        }
+
+        registry
+    }
+
+    /// Create a new empty OCR backend registry without default backends.
+    ///
+    /// This is useful for testing or when you want full control over backend registration.
+    pub fn new_empty() -> Self {
+        Self {
+            backends: HashMap::new(),
+        }
+    }
+
+    /// Register an OCR backend.
+    ///
+    /// # Arguments
+    ///
+    /// * `backend` - The OCR backend to register
+    ///
+    /// # Returns
+    ///
+    /// - `Ok(())` if registration succeeded
+    /// - `Err(...)` if initialization failed
+    ///
+    /// # Example
+    ///
+    /// ```rust,no_run
+    /// # use kreuzberg::plugins::registry::OcrBackendRegistry;
+    /// # use std::sync::Arc;
+    /// let mut registry = OcrBackendRegistry::new();
+    /// // let backend = Arc::new(MyOcrBackend::new());
+    /// // registry.register(backend)?;
+    /// # Ok::<(), kreuzberg::KreuzbergError>(())
+    /// ```
+    pub fn register(&mut self, backend: Arc<dyn OcrBackend>) -> Result<()> {
+        let name = backend.name().to_string();
+
+        super::validate_plugin_name(&name)?;
+
+        backend.initialize()?;
+
+        self.backends.insert(name, backend);
+        Ok(())
+    }
+
+    /// Get an OCR backend by name.
+    ///
+    /// # Arguments
+    ///
+    /// * `name` - Backend name
+    ///
+    /// # Returns
+    ///
+    /// The backend if found, or an error if not registered.
+    pub fn get(&self, name: &str) -> Result<Arc<dyn OcrBackend>> {
+        self.backends.get(name).cloned().ok_or_else(|| KreuzbergError::Plugin {
+            message: format!("OCR backend '{}' not registered", name),
+            plugin_name: name.to_string(),
+        })
+    }
+
+    /// Get an OCR backend that supports a specific language.
+    ///
+    /// Returns the first backend that supports the language.
+    ///
+    /// # Arguments
+    ///
+    /// * `language` - Language code (e.g., "eng", "deu")
+    ///
+    /// # Returns
+    ///
+    /// The first backend supporting the language, or an error if none found.
+    pub fn get_for_language(&self, language: &str) -> Result<Arc<dyn OcrBackend>> {
+        self.backends
+            .values()
+            .find(|backend| backend.supports_language(language))
+            .cloned()
+            .ok_or_else(|| KreuzbergError::Plugin {
+                message: format!("No OCR backend supports language '{}'", language),
+                plugin_name: language.to_string(),
+            })
+    }
+
+    /// List all registered backend names.
+    pub fn list(&self) -> Vec<String> {
+        self.backends.keys().cloned().collect()
+    }
+
+    /// Remove a backend from the registry.
+    ///
+    /// Calls `shutdown()` on the backend before removing.
+    pub fn remove(&mut self, name: &str) -> Result<()> {
+        if let Some(backend) = self.backends.remove(name) {
+            backend.shutdown()?;
+        }
+        Ok(())
+    }
+
+    /// Shutdown all backends and clear the registry.
+    pub fn shutdown_all(&mut self) -> Result<()> {
+        let names: Vec<_> = self.backends.keys().cloned().collect();
+        for name in names {
+            self.remove(&name)?;
+        }
+        Ok(())
+    }
+}
+
+impl Default for OcrBackendRegistry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::core::config::OcrConfig;
+    use crate::plugins::{OcrBackend, Plugin};
+    use crate::types::ExtractionResult;
+    use async_trait::async_trait;
+
+    struct MockOcrBackend {
+        name: String,
+        languages: Vec<String>,
+    }
+
+    impl Plugin for MockOcrBackend {
+        fn name(&self) -> &str {
+            &self.name
+        }
+        fn version(&self) -> String {
+            "1.0.0".to_string()
+        }
+        fn initialize(&self) -> Result<()> {
+            Ok(())
+        }
+        fn shutdown(&self) -> Result<()> {
+            Ok(())
+        }
+    }
+
+    #[async_trait]
+    impl OcrBackend for MockOcrBackend {
+        async fn process_image(&self, _: &[u8], _: &OcrConfig) -> Result<ExtractionResult> {
+            Ok(ExtractionResult {
+                content: "test".to_string(),
+                mime_type: "text/plain".to_string(),
+                metadata: crate::types::Metadata::default(),
+                tables: vec![],
+                detected_languages: None,
+                chunks: None,
+                images: None,
+                djot_content: None,
+                pages: None,
+                elements: None,
+            })
+        }
+
+        fn supports_language(&self, lang: &str) -> bool {
+            self.languages.iter().any(|l| l == lang)
+        }
+
+        fn backend_type(&self) -> crate::plugins::ocr::OcrBackendType {
+            crate::plugins::ocr::OcrBackendType::Custom
+        }
+    }
+
+    #[test]
+    fn test_ocr_backend_registry() {
+        let mut registry = OcrBackendRegistry::new_empty();
+
+        let backend = Arc::new(MockOcrBackend {
+            name: "test-ocr".to_string(),
+            languages: vec!["eng".to_string(), "deu".to_string()],
+        });
+
+        registry.register(backend).unwrap();
+
+        let retrieved = registry.get("test-ocr").unwrap();
+        assert_eq!(retrieved.name(), "test-ocr");
+
+        let eng_backend = registry.get_for_language("eng").unwrap();
+        assert_eq!(eng_backend.name(), "test-ocr");
+
+        let names = registry.list();
+        assert_eq!(names.len(), 1);
+        assert!(names.contains(&"test-ocr".to_string()));
+    }
+
+    #[test]
+    fn test_ocr_backend_registry_new_empty() {
+        let registry = OcrBackendRegistry::new_empty();
+        assert_eq!(registry.list().len(), 0);
+    }
+
+    #[test]
+    fn test_ocr_backend_get_missing() {
+        let registry = OcrBackendRegistry::new_empty();
+        let result = registry.get("nonexistent");
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_ocr_backend_get_for_language_missing() {
+        let registry = OcrBackendRegistry::new_empty();
+        let result = registry.get_for_language("fra");
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_ocr_backend_remove() {
+        let mut registry = OcrBackendRegistry::new_empty();
+        let backend = Arc::new(MockOcrBackend {
+            name: "test-backend".to_string(),
+            languages: vec!["eng".to_string()],
+        });
+        registry.register(backend).unwrap();
+
+        registry.remove("test-backend").unwrap();
+        assert_eq!(registry.list().len(), 0);
+    }
+
+    #[test]
+    fn test_ocr_backend_shutdown_all() {
+        let mut registry = OcrBackendRegistry::new_empty();
+        let backend1 = Arc::new(MockOcrBackend {
+            name: "backend1".to_string(),
+            languages: vec!["eng".to_string()],
+        });
+        let backend2 = Arc::new(MockOcrBackend {
+            name: "backend2".to_string(),
+            languages: vec!["deu".to_string()],
+        });
+
+        registry.register(backend1).unwrap();
+        registry.register(backend2).unwrap();
+
+        registry.shutdown_all().unwrap();
+        assert_eq!(registry.list().len(), 0);
+    }
+}

data/vendor/kreuzberg/src/plugins/registry/processor.rs

@@ -0,0 +1,304 @@
+//! Post-processor registry implementation.
+
+use crate::Result;
+use crate::plugins::{PostProcessor, ProcessingStage};
+use std::collections::{BTreeMap, HashMap};
+use std::sync::Arc;
+
+/// Registry for post-processor plugins.
+///
+/// Manages post-processors organized by processing stage.
+pub struct PostProcessorRegistry {
+    processors: HashMap<ProcessingStage, BTreeMap<i32, Vec<Arc<dyn PostProcessor>>>>,
+    name_index: HashMap<String, (ProcessingStage, i32)>,
+}
+
+impl PostProcessorRegistry {
+    /// Create a new empty post-processor registry.
+    pub fn new() -> Self {
+        Self {
+            processors: HashMap::new(),
+            name_index: HashMap::new(),
+        }
+    }
+
+    /// Register a post-processor.
+    ///
+    /// # Arguments
+    ///
+    /// * `processor` - The post-processor to register
+    /// * `priority` - Execution priority (higher = runs first within stage)
+    pub fn register(&mut self, processor: Arc<dyn PostProcessor>, priority: i32) -> Result<()> {
+        let name = processor.name().to_string();
+        let stage = processor.processing_stage();
+
+        super::validate_plugin_name(&name)?;
+
+        processor.initialize()?;
+
+        if self.name_index.contains_key(&name) {
+            self.remove(&name)?;
+        }
+
+        self.processors
+            .entry(stage)
+            .or_default()
+            .entry(priority)
+            .or_default()
+            .push(Arc::clone(&processor));
+
+        self.name_index.insert(name, (stage, priority));
+
+        Ok(())
+    }
+
+    /// Get all processors for a specific stage, in priority order.
+    ///
+    /// # Arguments
+    ///
+    /// * `stage` - The processing stage
+    ///
+    /// # Returns
+    ///
+    /// Vector of processors in priority order (highest first).
+    pub fn get_for_stage(&self, stage: ProcessingStage) -> Vec<Arc<dyn PostProcessor>> {
+        let mut result = Vec::new();
+
+        if let Some(priority_map) = self.processors.get(&stage) {
+            for (_priority, processors) in priority_map.iter().rev() {
+                for processor in processors {
+                    result.push(Arc::clone(processor));
+                }
+            }
+        }
+
+        result
+    }
+
+    /// List all registered processor names.
+    pub fn list(&self) -> Vec<String> {
+        self.name_index.keys().cloned().collect()
+    }
+
+    /// Remove a processor from the registry.
+    pub fn remove(&mut self, name: &str) -> Result<()> {
+        let (stage, priority) = match self.name_index.remove(name) {
+            Some(location) => location,
+            None => return Ok(()),
+        };
+
+        let processor_to_shutdown = if let Some(priority_map) = self.processors.get_mut(&stage) {
+            let processor = priority_map.get_mut(&priority).and_then(|processors| {
+                processors
+                    .iter()
+                    .position(|p| p.name() == name)
+                    .map(|pos| processors.remove(pos))
+            });
+
+            if let Some(processors) = priority_map.get(&priority)
+                && processors.is_empty()
+            {
+                priority_map.remove(&priority);
+            }
+
+            if priority_map.is_empty() {
+                self.processors.remove(&stage);
+            }
+            processor
+        } else {
+            None
+        };
+
+        if let Some(processor) = processor_to_shutdown {
+            processor.shutdown()?;
+        }
+
+        Ok(())
+    }
+
+    /// Shutdown all processors and clear the registry.
+    pub fn shutdown_all(&mut self) -> Result<()> {
+        let names = self.list();
+        for name in names {
+            self.remove(&name)?;
+        }
+        Ok(())
+    }
+}
+
+impl Default for PostProcessorRegistry {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::KreuzbergError;
+    use crate::core::config::ExtractionConfig;
+    use crate::plugins::Plugin;
+    use crate::types::ExtractionResult;
+    use async_trait::async_trait;
+
+    struct MockPostProcessor {
+        name: String,
+        stage: ProcessingStage,
+    }
+
+    impl Plugin for MockPostProcessor {
+        fn name(&self) -> &str {
+            &self.name
+        }
+        fn version(&self) -> String {
+            "1.0.0".to_string()
+        }
+        fn initialize(&self) -> Result<()> {
+            Ok(())
+        }
+        fn shutdown(&self) -> Result<()> {
+            Ok(())
+        }
+    }
+
+    #[async_trait]
+    impl PostProcessor for MockPostProcessor {
+        async fn process(&self, _result: &mut ExtractionResult, _: &ExtractionConfig) -> Result<()> {
+            Ok(())
+        }
+
+        fn processing_stage(&self) -> ProcessingStage {
+            self.stage
+        }
+    }
+
+    #[test]
+    fn test_post_processor_registry() {
+        let mut registry = PostProcessorRegistry::new();
+
+        let early = Arc::new(MockPostProcessor {
+            name: "early-processor".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        let middle = Arc::new(MockPostProcessor {
+            name: "middle-processor".to_string(),
+            stage: ProcessingStage::Middle,
+        });
+
+        registry.register(early, 100).unwrap();
+        registry.register(middle, 50).unwrap();
+
+        let early_processors = registry.get_for_stage(ProcessingStage::Early);
+        assert_eq!(early_processors.len(), 1);
+        assert_eq!(early_processors[0].name(), "early-processor");
+
+        let middle_processors = registry.get_for_stage(ProcessingStage::Middle);
+        assert_eq!(middle_processors.len(), 1);
+
+        let names = registry.list();
+        assert_eq!(names.len(), 2);
+    }
+
+    #[test]
+    fn test_post_processor_registry_remove() {
+        let mut registry = PostProcessorRegistry::new();
+
+        let processor = Arc::new(MockPostProcessor {
+            name: "test-processor".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        registry.register(processor, 50).unwrap();
+        assert_eq!(registry.get_for_stage(ProcessingStage::Early).len(), 1);
+
+        registry.remove("test-processor").unwrap();
+        assert_eq!(registry.get_for_stage(ProcessingStage::Early).len(), 0);
+    }
+
+    #[test]
+    fn test_post_processor_registry_default() {
+        let registry = PostProcessorRegistry::default();
+        assert_eq!(registry.list().len(), 0);
+    }
+
+    #[test]
+    fn test_post_processor_registry_invalid_name_empty() {
+        let mut registry = PostProcessorRegistry::new();
+
+        let processor = Arc::new(MockPostProcessor {
+            name: "".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        let result = registry.register(processor, 50);
+        assert!(matches!(result, Err(KreuzbergError::Validation { .. })));
+    }
+
+    #[test]
+    fn test_post_processor_registry_invalid_name_whitespace() {
+        let mut registry = PostProcessorRegistry::new();
+
+        let processor = Arc::new(MockPostProcessor {
+            name: "my processor".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        let result = registry.register(processor, 50);
+        assert!(matches!(result, Err(KreuzbergError::Validation { .. })));
+    }
+
+    #[test]
+    fn test_post_processor_registry_shutdown_all() {
+        let mut registry = PostProcessorRegistry::new();
+
+        let early = Arc::new(MockPostProcessor {
+            name: "early".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        let late = Arc::new(MockPostProcessor {
+            name: "late".to_string(),
+            stage: ProcessingStage::Late,
+        });
+
+        registry.register(early, 100).unwrap();
+        registry.register(late, 50).unwrap();
+
+        assert_eq!(registry.list().len(), 2);
+
+        registry.shutdown_all().unwrap();
+        assert_eq!(registry.list().len(), 0);
+    }
+
+    #[test]
+    fn test_post_processor_registry_priority_order() {
+        let mut registry = PostProcessorRegistry::new();
+
+        let low = Arc::new(MockPostProcessor {
+            name: "low-priority".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        let high = Arc::new(MockPostProcessor {
+            name: "high-priority".to_string(),
+            stage: ProcessingStage::Early,
+        });
+
+        registry.register(low, 10).unwrap();
+        registry.register(high, 100).unwrap();
+
+        let processors = registry.get_for_stage(ProcessingStage::Early);
+        assert_eq!(processors.len(), 2);
+        assert_eq!(processors[0].name(), "high-priority");
+        assert_eq!(processors[1].name(), "low-priority");
+    }
+
+    #[test]
+    fn test_post_processor_registry_empty_stage() {
+        let registry = PostProcessorRegistry::new();
+
+        let processors = registry.get_for_stage(ProcessingStage::Late);
+        assert_eq!(processors.len(), 0);
+    }
+}