kreuzberg 4.0.8 → 4.1.0

data/ext/kreuzberg_rb/native/src/metadata.rs

@@ -0,0 +1,34 @@
+//! Metadata handling and document format detection
+//!
+//! Provides utilities for MIME type detection, format validation, and extension mapping.
+
+use crate::error_handling::runtime_error;
+use magnus::Error;
+
+/// Detect MIME type from bytes
+pub fn detect_mime_type_from_bytes(bytes: String) -> Result<String, Error> {
+    let bytes_vec = bytes.into_bytes();
+    kreuzberg::core::mime::detect_mime_type_from_bytes(&bytes_vec)
+        .map_err(|e| runtime_error(format!("Failed to detect MIME type: {}", e)))
+}
+
+/// Detect MIME type from file path
+pub fn detect_mime_type_from_path_native(path: String) -> Result<String, Error> {
+    kreuzberg::core::mime::detect_mime_type(&path, true)
+        .map_err(|e| runtime_error(format!("Failed to detect MIME type from path: {}", e)))
+}
+
+/// Validate MIME type
+pub fn validate_mime_type_native(mime_type: String) -> Result<String, Error> {
+    if kreuzberg::core::mime::validate_mime_type(&mime_type).is_ok() {
+        Ok(mime_type)
+    } else {
+        Err(runtime_error(format!("Unsupported MIME type: {}", mime_type)))
+    }
+}
+
+/// Get file extensions for a given MIME type
+pub fn get_extensions_for_mime_native(mime_type: String) -> Result<Vec<String>, Error> {
+    kreuzberg::core::mime::get_extensions_for_mime(&mime_type)
+        .map_err(|e| runtime_error(format!("Failed to get extensions: {}", e)))
+}

data/ext/kreuzberg_rb/native/src/plugins/mod.rs

@@ -0,0 +1,92 @@
+//! Plugin management for Kreuzberg
+//!
+//! Handles registration and management of custom plugins including post-processors,
+//! validators, and OCR backends.
+
+pub mod post_processor;
+pub mod validator;
+pub mod ocr_backend;
+
+pub use post_processor::register_post_processor;
+pub use validator::register_validator;
+pub use ocr_backend::{register_ocr_backend, unregister_ocr_backend, list_ocr_backends, clear_ocr_backends};
+
+// Plugin registry functions
+pub use kreuzberg::get_post_processor_registry;
+
+use magnus::Error;
+use kreuzberg::plugins::{
+    unregister_validator as kz_unregister_validator,
+    clear_validators as kz_clear_validators,
+    list_validators as kz_list_validators,
+    list_post_processors as kz_list_post_processors,
+    list_extractors as kz_list_extractors,
+    unregister_extractor as kz_unregister_extractor,
+    clear_extractors as kz_clear_extractors,
+};
+
+/// Unregister a post-processor plugin by name
+pub fn unregister_post_processor(name: String) -> Result<(), Error> {
+    let registry = get_post_processor_registry();
+    registry
+        .write()
+        .map_err(|e| crate::error_handling::runtime_error(format!("Failed to acquire registry lock: {}", e)))?
+        .remove(&name)
+        .map_err(crate::error_handling::kreuzberg_error)?;
+
+    Ok(())
+}
+
+/// Unregister a validator plugin by name
+pub fn unregister_validator(name: String) -> Result<(), Error> {
+    kz_unregister_validator(&name)
+        .map_err(crate::error_handling::kreuzberg_error)
+}
+
+/// Clear all post-processors
+pub fn clear_post_processors() -> Result<(), Error> {
+    let registry = get_post_processor_registry();
+    registry
+        .write()
+        .map_err(|e| crate::error_handling::runtime_error(format!("Failed to acquire registry lock: {}", e)))?
+        .shutdown_all()
+        .map_err(crate::error_handling::kreuzberg_error)?;
+
+    Ok(())
+}
+
+/// Clear all validators
+pub fn clear_validators() -> Result<(), Error> {
+    kz_clear_validators()
+        .map_err(crate::error_handling::kreuzberg_error)
+}
+
+/// List registered post-processors
+pub fn list_post_processors() -> Result<Vec<String>, Error> {
+    kz_list_post_processors()
+        .map_err(crate::error_handling::kreuzberg_error)
+}
+
+/// List registered validators
+pub fn list_validators() -> Result<Vec<String>, Error> {
+    kz_list_validators()
+        .map_err(crate::error_handling::kreuzberg_error)
+}
+
+/// List registered document extractors
+pub fn list_document_extractors() -> Result<Vec<String>, Error> {
+    kz_list_extractors()
+        .map_err(crate::error_handling::kreuzberg_error)
+}
+
+/// Unregister a document extractor
+pub fn unregister_document_extractor(name: String) -> Result<(), Error> {
+    kz_unregister_extractor(&name)
+        .map_err(crate::error_handling::kreuzberg_error)
+}
+
+/// Clear all document extractors
+pub fn clear_document_extractors() -> Result<(), Error> {
+    kz_clear_extractors()
+        .map_err(crate::error_handling::kreuzberg_error)
+}

data/ext/kreuzberg_rb/native/src/plugins/ocr_backend.rs

@@ -0,0 +1,159 @@
+//! OCR backend plugin registration and management
+
+use crate::error_handling::{kreuzberg_error, runtime_error};
+use crate::gc_guarded_value::GcGuardedValue;
+use magnus::{Error, Ruby, TryConvert, Value};
+use magnus::value::ReprValue;
+use kreuzberg::plugins::{
+    register_ocr_backend as kz_register_ocr_backend,
+    unregister_ocr_backend as kz_unregister_ocr_backend,
+    list_ocr_backends as kz_list_ocr_backends,
+    clear_ocr_backends as kz_clear_ocr_backends,
+    OcrBackend, OcrBackendType, Plugin,
+};
+use kreuzberg::types::{ExtractionResult, Metadata};
+use kreuzberg::{OcrConfig, KreuzbergError};
+use async_trait::async_trait;
+use std::path::Path;
+use std::sync::Arc;
+
+/// Ruby OCR backend wrapper that implements the OcrBackend trait
+struct RubyOcrBackend {
+    name: String,
+    backend: GcGuardedValue,
+}
+
+// SAFETY: Ruby's GC is handled by GcGuardedValue, and we ensure all Ruby
+// calls happen through proper Magnus/Ruby FFI boundaries
+unsafe impl Send for RubyOcrBackend {}
+unsafe impl Sync for RubyOcrBackend {}
+
+impl Plugin for RubyOcrBackend {
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn version(&self) -> String {
+        "1.0.0".to_string()
+    }
+
+    fn initialize(&self) -> kreuzberg::Result<()> {
+        Ok(())
+    }
+
+    fn shutdown(&self) -> kreuzberg::Result<()> {
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl OcrBackend for RubyOcrBackend {
+    async fn process_image(&self, image_bytes: &[u8], config: &OcrConfig) -> kreuzberg::Result<ExtractionResult> {
+        let backend_name = self.name.clone();
+        let backend = self.backend.value();
+        let image_data = image_bytes.to_vec();
+        let ocr_config = config.clone();
+
+        tokio::task::block_in_place(|| {
+            let ruby = Ruby::get().expect("Ruby not initialized");
+
+            // Convert image bytes to Ruby string (binary)
+            let ruby_bytes = ruby.str_from_slice(&image_data);
+
+            // Convert config to Ruby hash
+            let config_hash = ruby.hash_new();
+            config_hash.aset("backend", ocr_config.backend.as_str())
+                .map_err(|e| KreuzbergError::Plugin {
+                    message: format!("Failed to set backend in config: {}", e),
+                    plugin_name: backend_name.clone(),
+                })?;
+            config_hash.aset("language", ocr_config.language.as_str())
+                .map_err(|e| KreuzbergError::Plugin {
+                    message: format!("Failed to set language in config: {}", e),
+                    plugin_name: backend_name.clone(),
+                })?;
+
+            // Call Ruby backend's process_image method
+            let result: magnus::Value = backend
+                .funcall("process_image", (ruby_bytes, config_hash))
+                .map_err(|e| KreuzbergError::Plugin {
+                    message: format!("Ruby OCR backend failed: {}", e),
+                    plugin_name: backend_name.clone(),
+                })?;
+
+            // Convert result to String
+            let content = String::try_convert(result)
+                .map_err(|e| KreuzbergError::Plugin {
+                    message: format!("OCR backend must return a String: {}", e),
+                    plugin_name: backend_name.clone(),
+                })?;
+
+            Ok(ExtractionResult {
+                content,
+                mime_type: "text/plain".to_string(),
+                metadata: Metadata::default(),
+                tables: vec![],
+                detected_languages: None,
+                chunks: None,
+                images: None,
+                djot_content: None,
+                pages: None,
+                elements: None,
+            })
+        })
+    }
+
+    async fn process_file(&self, path: &Path, config: &OcrConfig) -> kreuzberg::Result<ExtractionResult> {
+        let bytes = std::fs::read(path)?;
+        self.process_image(&bytes, config).await
+    }
+
+    fn supports_language(&self, _lang: &str) -> bool {
+        // Ruby backends are assumed to support all languages by default
+        // A more sophisticated implementation could call back to Ruby
+        true
+    }
+
+    fn backend_type(&self) -> OcrBackendType {
+        OcrBackendType::Custom
+    }
+}
+
+/// Register an OCR backend plugin
+pub fn register_ocr_backend(name: String, backend: Value) -> Result<(), Error> {
+    let _ruby = Ruby::get().expect("Ruby not initialized");
+
+    // Validate that the backend has the required methods
+    if !backend.respond_to("name", true)? {
+        return Err(runtime_error("OCR backend must implement #name method"));
+    }
+    if !backend.respond_to("process_image", true)? {
+        return Err(runtime_error("OCR backend must implement #process_image(image_bytes, config) method"));
+    }
+
+    let backend_impl = Arc::new(RubyOcrBackend {
+        name: name.clone(),
+        backend: GcGuardedValue::new(backend),
+    });
+
+    kz_register_ocr_backend(backend_impl)
+        .map_err(kreuzberg_error)
+}
+
+/// Unregister an OCR backend
+pub fn unregister_ocr_backend(_name: String) -> Result<(), Error> {
+    kz_unregister_ocr_backend(_name.as_str())
+        .map_err(kreuzberg_error)
+}
+
+/// List registered OCR backends
+pub fn list_ocr_backends() -> Result<Vec<String>, Error> {
+    kz_list_ocr_backends()
+        .map_err(kreuzberg_error)
+}
+
+/// Clear all OCR backends
+pub fn clear_ocr_backends() -> Result<(), Error> {
+    kz_clear_ocr_backends()
+        .map_err(kreuzberg_error)
+}

data/ext/kreuzberg_rb/native/src/plugins/post_processor.rs

@@ -0,0 +1,126 @@
+//! Post-processor plugin registration and management
+
+use crate::{error_handling::{kreuzberg_error, runtime_error}, gc_guarded_value::GcGuardedValue, helpers::get_kw};
+use magnus::{Error, Ruby, Value, scan_args::scan_args, TryConvert};
+use magnus::value::ReprValue;
+use std::sync::Arc;
+
+/// Register a post-processor plugin
+pub fn register_post_processor(args: &[Value]) -> Result<(), Error> {
+    let _ruby = Ruby::get().expect("Ruby not initialized");
+    let args = scan_args::<(String, Value), (Option<i32>,), (), (), (), ()>(args)?;
+    let (name, processor) = args.required;
+    let (priority,) = args.optional;
+    let priority = priority.unwrap_or(50);
+
+    if !processor.respond_to("call", true)? {
+        return Err(runtime_error("Post-processor must be a Proc or respond to 'call'"));
+    }
+
+    use async_trait::async_trait;
+    use kreuzberg::plugins::{Plugin, PostProcessor, ProcessingStage};
+
+    struct RubyPostProcessor {
+        name: String,
+        processor: GcGuardedValue,
+    }
+
+    unsafe impl Send for RubyPostProcessor {}
+    unsafe impl Sync for RubyPostProcessor {}
+
+    impl Plugin for RubyPostProcessor {
+        fn name(&self) -> &str {
+            &self.name
+        }
+
+        fn version(&self) -> String {
+            "1.0.0".to_string()
+        }
+
+        fn initialize(&self) -> kreuzberg::Result<()> {
+            Ok(())
+        }
+
+        fn shutdown(&self) -> kreuzberg::Result<()> {
+            Ok(())
+        }
+    }
+
+    #[async_trait]
+    impl PostProcessor for RubyPostProcessor {
+        async fn process(
+            &self,
+            result: &mut kreuzberg::ExtractionResult,
+            _config: &kreuzberg::ExtractionConfig,
+        ) -> kreuzberg::Result<()> {
+            let processor_name = self.name.clone();
+            let processor = self.processor.value();
+            let result_clone = result.clone();
+
+            let updated_result = tokio::task::block_in_place(|| {
+                let ruby = Ruby::get().expect("Ruby not initialized");
+                let result_hash = crate::result::extraction_result_to_ruby(&ruby, result_clone.clone()).map_err(|e| {
+                    kreuzberg::KreuzbergError::Plugin {
+                        message: format!("Failed to convert result to Ruby: {}", e),
+                        plugin_name: processor_name.clone(),
+                    }
+                })?;
+
+                let modified = processor
+                    .funcall::<_, _, magnus::Value>("call", (result_hash,))
+                    .map_err(|e| kreuzberg::KreuzbergError::Plugin {
+                        message: format!("Ruby post-processor failed: {}", e),
+                        plugin_name: processor_name.clone(),
+                    })?;
+
+                let modified_hash =
+                    magnus::RHash::try_convert(modified).map_err(|e| kreuzberg::KreuzbergError::Plugin {
+                        message: format!("Post-processor must return a Hash: {}", e),
+                        plugin_name: processor_name.clone(),
+                    })?;
+
+                let mut updated_result = result_clone;
+
+                if let Some(content_val) = get_kw(&ruby, modified_hash, "content") {
+                    let new_content =
+                        String::try_convert(content_val).map_err(|e| kreuzberg::KreuzbergError::Plugin {
+                            message: format!("Failed to convert content: {}", e),
+                            plugin_name: processor_name.clone(),
+                        })?;
+                    updated_result.content = new_content;
+                }
+
+                if let Some(mime_val) = get_kw(&ruby, modified_hash, "mime_type") {
+                    let new_mime = String::try_convert(mime_val).map_err(|e| kreuzberg::KreuzbergError::Plugin {
+                        message: format!("Failed to convert mime_type: {}", e),
+                        plugin_name: processor_name.clone(),
+                    })?;
+                    updated_result.mime_type = new_mime;
+                }
+
+                Ok::<kreuzberg::ExtractionResult, kreuzberg::KreuzbergError>(updated_result)
+            })?;
+
+            *result = updated_result;
+            Ok(())
+        }
+
+        fn processing_stage(&self) -> ProcessingStage {
+            ProcessingStage::Late
+        }
+    }
+
+    let processor_impl = Arc::new(RubyPostProcessor {
+        name: name.clone(),
+        processor: GcGuardedValue::new(processor),
+    });
+
+    let registry = kreuzberg::get_post_processor_registry();
+    registry
+        .write()
+        .map_err(|e| runtime_error(format!("Failed to acquire registry lock: {}", e)))?
+        .register(processor_impl, priority)
+        .map_err(kreuzberg_error)?;
+
+    Ok(())
+}

data/ext/kreuzberg_rb/native/src/plugins/validator.rs

@@ -0,0 +1,99 @@
+//! Validator plugin registration and management
+
+use crate::{error_handling::{kreuzberg_error, runtime_error}, gc_guarded_value::GcGuardedValue};
+use magnus::{Error, Value, scan_args::scan_args, Ruby};
+use magnus::value::ReprValue;
+use std::sync::Arc;
+
+/// Register a validator plugin
+pub fn register_validator(args: &[Value]) -> Result<(), Error> {
+    let _ruby = Ruby::get().expect("Ruby not initialized");
+    let args = scan_args::<(String, Value), (Option<i32>,), (), (), (), ()>(args)?;
+    let (name, validator) = args.required;
+    let (priority,) = args.optional;
+    let priority = priority.unwrap_or(50);
+
+    if !validator.respond_to("call", true)? {
+        return Err(runtime_error("Validator must be a Proc or respond to 'call'"));
+    }
+
+    use async_trait::async_trait;
+    use kreuzberg::plugins::{Plugin, Validator};
+
+    struct RubyValidator {
+        name: String,
+        validator: GcGuardedValue,
+        priority: i32,
+    }
+
+    unsafe impl Send for RubyValidator {}
+    unsafe impl Sync for RubyValidator {}
+
+    impl Plugin for RubyValidator {
+        fn name(&self) -> &str {
+            &self.name
+        }
+
+        fn version(&self) -> String {
+            "1.0.0".to_string()
+        }
+
+        fn initialize(&self) -> kreuzberg::Result<()> {
+            Ok(())
+        }
+
+        fn shutdown(&self) -> kreuzberg::Result<()> {
+            Ok(())
+        }
+    }
+
+    #[async_trait]
+    impl Validator for RubyValidator {
+        async fn validate(
+            &self,
+            result: &kreuzberg::ExtractionResult,
+            _config: &kreuzberg::ExtractionConfig,
+        ) -> kreuzberg::Result<()> {
+            let validator_name = self.name.clone();
+            let validator = self.validator.value();
+            let result_clone = result.clone();
+
+            tokio::task::block_in_place(|| {
+                let ruby = Ruby::get().expect("Ruby not initialized");
+                let result_hash =
+                    crate::result::extraction_result_to_ruby(&ruby, result_clone).map_err(|e| kreuzberg::KreuzbergError::Plugin {
+                        message: format!("Failed to convert result to Ruby: {}", e),
+                        plugin_name: validator_name.clone(),
+                    })?;
+
+                validator
+                    .funcall::<_, _, magnus::Value>("call", (result_hash,))
+                    .map_err(|e| kreuzberg::KreuzbergError::Validation {
+                        message: format!("Validation failed: {}", e),
+                        source: None,
+                    })?;
+
+                Ok(())
+            })
+        }
+
+        fn priority(&self) -> i32 {
+            self.priority
+        }
+    }
+
+    let validator_impl = Arc::new(RubyValidator {
+        name: name.clone(),
+        validator: GcGuardedValue::new(validator),
+        priority,
+    });
+
+    let registry = kreuzberg::get_validator_registry();
+    registry
+        .write()
+        .map_err(|e| runtime_error(format!("Failed to acquire registry lock: {}", e)))?
+        .register(validator_impl)
+        .map_err(kreuzberg_error)?;
+
+    Ok(())
+}