kreuzberg 4.0.8 → 4.1.0

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

@@ -0,0 +1,238 @@
+//! Validator registry implementation.
+
+use crate::Result;
+use crate::plugins::Validator;
+use indexmap::IndexMap;
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
+/// Registry for validator plugins.
+///
+/// Manages validators with priority-based execution order.
+pub struct ValidatorRegistry {
+    validators: BTreeMap<i32, IndexMap<String, Arc<dyn Validator>>>,
+}
+
+impl ValidatorRegistry {
+    /// Create a new empty validator registry.
+    pub fn new() -> Self {
+        Self {
+            validators: BTreeMap::new(),
+        }
+    }
+
+    /// Register a validator.
+    ///
+    /// # Arguments
+    ///
+    /// * `validator` - The validator to register
+    pub fn register(&mut self, validator: Arc<dyn Validator>) -> Result<()> {
+        let name = validator.name().to_string();
+        let priority = validator.priority();
+
+        super::validate_plugin_name(&name)?;
+
+        validator.initialize()?;
+
+        self.validators.entry(priority).or_default().insert(name, validator);
+
+        Ok(())
+    }
+
+    /// Get all validators in priority order.
+    ///
+    /// # Returns
+    ///
+    /// Vector of validators in priority order (highest first).
+    pub fn get_all(&self) -> Vec<Arc<dyn Validator>> {
+        let mut result = Vec::new();
+
+        for (_priority, validators) in self.validators.iter().rev() {
+            for validator in validators.values() {
+                result.push(Arc::clone(validator));
+            }
+        }
+
+        result
+    }
+
+    /// List all registered validator names.
+    pub fn list(&self) -> Vec<String> {
+        let mut names = std::collections::HashSet::new();
+        for validators in self.validators.values() {
+            names.extend(validators.keys().cloned());
+        }
+        names.into_iter().collect()
+    }
+
+    /// Remove a validator from the registry.
+    pub fn remove(&mut self, name: &str) -> Result<()> {
+        let mut validator_to_shutdown: Option<Arc<dyn Validator>> = None;
+
+        for validators in self.validators.values_mut() {
+            if let Some(validator) = validators.shift_remove(name)
+                && validator_to_shutdown.is_none()
+            {
+                validator_to_shutdown = Some(validator);
+            }
+        }
+
+        if let Some(validator) = validator_to_shutdown {
+            validator.shutdown()?;
+        }
+
+        self.validators.retain(|_, validators| !validators.is_empty());
+
+        Ok(())
+    }
+
+    /// Shutdown all validators 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 ValidatorRegistry {
+    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 MockValidator {
+        name: String,
+        priority: i32,
+    }
+
+    impl Plugin for MockValidator {
+        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 Validator for MockValidator {
+        async fn validate(&self, _: &ExtractionResult, _: &ExtractionConfig) -> Result<()> {
+            Ok(())
+        }
+
+        fn priority(&self) -> i32 {
+            self.priority
+        }
+    }
+
+    #[test]
+    fn test_validator_registry() {
+        let mut registry = ValidatorRegistry::new();
+
+        let high_priority = Arc::new(MockValidator {
+            name: "high-priority".to_string(),
+            priority: 100,
+        });
+
+        let low_priority = Arc::new(MockValidator {
+            name: "low-priority".to_string(),
+            priority: 10,
+        });
+
+        registry.register(high_priority).unwrap();
+        registry.register(low_priority).unwrap();
+
+        let validators = registry.get_all();
+        assert_eq!(validators.len(), 2);
+        assert_eq!(validators[0].name(), "high-priority");
+        assert_eq!(validators[1].name(), "low-priority");
+    }
+
+    #[test]
+    fn test_validator_registry_remove() {
+        let mut registry = ValidatorRegistry::new();
+
+        let validator = Arc::new(MockValidator {
+            name: "test-validator".to_string(),
+            priority: 50,
+        });
+
+        registry.register(validator).unwrap();
+        assert_eq!(registry.get_all().len(), 1);
+
+        registry.remove("test-validator").unwrap();
+        assert_eq!(registry.get_all().len(), 0);
+    }
+
+    #[test]
+    fn test_validator_registry_default() {
+        let registry = ValidatorRegistry::default();
+        assert_eq!(registry.get_all().len(), 0);
+    }
+
+    #[test]
+    fn test_validator_registry_invalid_name_empty() {
+        let mut registry = ValidatorRegistry::new();
+
+        let validator = Arc::new(MockValidator {
+            name: "".to_string(),
+            priority: 50,
+        });
+
+        let result = registry.register(validator);
+        assert!(matches!(result, Err(KreuzbergError::Validation { .. })));
+    }
+
+    #[test]
+    fn test_validator_registry_invalid_name_whitespace() {
+        let mut registry = ValidatorRegistry::new();
+
+        let validator = Arc::new(MockValidator {
+            name: "my validator".to_string(),
+            priority: 50,
+        });
+
+        let result = registry.register(validator);
+        assert!(matches!(result, Err(KreuzbergError::Validation { .. })));
+    }
+
+    #[test]
+    fn test_validator_registry_shutdown_all() {
+        let mut registry = ValidatorRegistry::new();
+
+        let validator1 = Arc::new(MockValidator {
+            name: "validator1".to_string(),
+            priority: 100,
+        });
+
+        let validator2 = Arc::new(MockValidator {
+            name: "validator2".to_string(),
+            priority: 50,
+        });
+
+        registry.register(validator1).unwrap();
+        registry.register(validator2).unwrap();
+
+        assert_eq!(registry.get_all().len(), 2);
+
+        registry.shutdown_all().unwrap();
+        assert_eq!(registry.get_all().len(), 0);
+    }
+}

data/vendor/kreuzberg/src/plugins/validator/mod.rs

@@ -0,0 +1,424 @@
+//! Validator plugin system.
+//!
+//! This module provides the trait and registry for implementing custom validators.
+
+mod registry;
+mod r#trait;
+
+// Re-export trait for backward compatibility
+pub use r#trait::Validator;
+
+// Re-export registry functions for backward compatibility
+pub use registry::{clear_validators, list_validators, register_validator, unregister_validator};
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::KreuzbergError;
+    use crate::Result;
+    use crate::core::config::ExtractionConfig;
+    use crate::plugins::Plugin;
+    use crate::types::ExtractionResult;
+    use async_trait::async_trait;
+    use std::collections::HashMap;
+
+    struct MockValidator {
+        should_fail: bool,
+    }
+
+    impl Plugin for MockValidator {
+        fn name(&self) -> &str {
+            "mock-validator"
+        }
+
+        fn version(&self) -> String {
+            "1.0.0".to_string()
+        }
+
+        fn initialize(&self) -> Result<()> {
+            Ok(())
+        }
+
+        fn shutdown(&self) -> Result<()> {
+            Ok(())
+        }
+    }
+
+    #[async_trait]
+    impl Validator for MockValidator {
+        async fn validate(&self, _result: &ExtractionResult, _config: &ExtractionConfig) -> Result<()> {
+            if self.should_fail {
+                Err(KreuzbergError::validation("Validation failed".to_string()))
+            } else {
+                Ok(())
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_validator_success() {
+        let validator = MockValidator { should_fail: false };
+
+        let result = ExtractionResult {
+            content: "test content".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,
+        };
+
+        let config = ExtractionConfig::default();
+        assert!(validator.validate(&result, &config).await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_validator_failure() {
+        let validator = MockValidator { should_fail: true };
+
+        let result = ExtractionResult {
+            content: "test content".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,
+        };
+
+        let config = ExtractionConfig::default();
+        let validation_result = validator.validate(&result, &config).await;
+
+        assert!(matches!(validation_result, Err(KreuzbergError::Validation { .. })));
+    }
+
+    #[test]
+    fn test_validator_should_validate_default() {
+        let validator = MockValidator { should_fail: false };
+
+        let result = 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,
+        };
+
+        let config = ExtractionConfig::default();
+
+        assert!(validator.should_validate(&result, &config));
+    }
+
+    #[test]
+    fn test_validator_priority_default() {
+        let validator = MockValidator { should_fail: false };
+        assert_eq!(validator.priority(), 50);
+    }
+
+    #[tokio::test]
+    async fn test_validator_plugin_interface() {
+        let validator = MockValidator { should_fail: false };
+
+        assert_eq!(validator.name(), "mock-validator");
+        assert_eq!(validator.version(), "1.0.0");
+        assert!(validator.initialize().is_ok());
+        assert!(validator.shutdown().is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_validator_empty_content() {
+        let validator = MockValidator { should_fail: false };
+
+        let result = ExtractionResult {
+            content: String::new(),
+            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,
+        };
+
+        let config = ExtractionConfig::default();
+        assert!(validator.validate(&result, &config).await.is_ok());
+    }
+
+    #[test]
+    fn test_validator_should_validate_conditional() {
+        struct PdfOnlyValidator;
+
+        impl Plugin for PdfOnlyValidator {
+            fn name(&self) -> &str {
+                "pdf-only"
+            }
+            fn version(&self) -> String {
+                "1.0.0".to_string()
+            }
+            fn initialize(&self) -> Result<()> {
+                Ok(())
+            }
+            fn shutdown(&self) -> Result<()> {
+                Ok(())
+            }
+        }
+
+        #[async_trait]
+        impl Validator for PdfOnlyValidator {
+            async fn validate(&self, _result: &ExtractionResult, _config: &ExtractionConfig) -> Result<()> {
+                Ok(())
+            }
+
+            fn should_validate(&self, result: &ExtractionResult, _config: &ExtractionConfig) -> bool {
+                result.mime_type == "application/pdf"
+            }
+        }
+
+        let validator = PdfOnlyValidator;
+        let config = ExtractionConfig::default();
+
+        let pdf_result = ExtractionResult {
+            content: "test".to_string(),
+            mime_type: "application/pdf".to_string(),
+            metadata: crate::types::Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            djot_content: None,
+            pages: None,
+            elements: None,
+        };
+
+        let txt_result = 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,
+        };
+
+        assert!(validator.should_validate(&pdf_result, &config));
+        assert!(!validator.should_validate(&txt_result, &config));
+    }
+
+    #[test]
+    fn test_validator_priority_ranges() {
+        struct HighPriorityValidator;
+        struct LowPriorityValidator;
+
+        impl Plugin for HighPriorityValidator {
+            fn name(&self) -> &str {
+                "high-priority"
+            }
+            fn version(&self) -> String {
+                "1.0.0".to_string()
+            }
+            fn initialize(&self) -> Result<()> {
+                Ok(())
+            }
+            fn shutdown(&self) -> Result<()> {
+                Ok(())
+            }
+        }
+
+        impl Plugin for LowPriorityValidator {
+            fn name(&self) -> &str {
+                "low-priority"
+            }
+            fn version(&self) -> String {
+                "1.0.0".to_string()
+            }
+            fn initialize(&self) -> Result<()> {
+                Ok(())
+            }
+            fn shutdown(&self) -> Result<()> {
+                Ok(())
+            }
+        }
+
+        #[async_trait]
+        impl Validator for HighPriorityValidator {
+            async fn validate(&self, _result: &ExtractionResult, _config: &ExtractionConfig) -> Result<()> {
+                Ok(())
+            }
+
+            fn priority(&self) -> i32 {
+                100
+            }
+        }
+
+        #[async_trait]
+        impl Validator for LowPriorityValidator {
+            async fn validate(&self, _result: &ExtractionResult, _config: &ExtractionConfig) -> Result<()> {
+                Ok(())
+            }
+
+            fn priority(&self) -> i32 {
+                10
+            }
+        }
+
+        let high = HighPriorityValidator;
+        let low = LowPriorityValidator;
+
+        assert_eq!(high.priority(), 100);
+        assert_eq!(low.priority(), 10);
+        assert!(high.priority() > low.priority());
+    }
+
+    #[tokio::test]
+    async fn test_validator_error_message() {
+        let validator = MockValidator { should_fail: true };
+
+        let result = 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,
+        };
+
+        let config = ExtractionConfig::default();
+        let err = validator.validate(&result, &config).await.unwrap_err();
+
+        match err {
+            KreuzbergError::Validation { message: msg, .. } => {
+                assert_eq!(msg, "Validation failed");
+            }
+            _ => panic!("Expected Validation error"),
+        }
+    }
+
+    #[tokio::test]
+    async fn test_validator_with_metadata() {
+        let validator = MockValidator { should_fail: false };
+
+        let mut additional = HashMap::new();
+        additional.insert("quality_score".to_string(), serde_json::json!(0.95));
+
+        let result = ExtractionResult {
+            content: "test".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: crate::types::Metadata {
+                additional,
+                ..Default::default()
+            },
+            pages: None,
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        let config = ExtractionConfig::default();
+        assert!(validator.validate(&result, &config).await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_validator_with_tables() {
+        use crate::types::Table;
+
+        let validator = MockValidator { should_fail: false };
+
+        let table = Table {
+            cells: vec![vec!["A".to_string(), "B".to_string()]],
+            markdown: "| A | B |".to_string(),
+            page_number: 0,
+        };
+
+        let result = ExtractionResult {
+            content: "test".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: crate::types::Metadata::default(),
+            tables: vec![table],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            djot_content: None,
+            pages: None,
+            elements: None,
+        };
+
+        let config = ExtractionConfig::default();
+        assert!(validator.validate(&result, &config).await.is_ok());
+    }
+
+    #[tokio::test]
+    async fn test_validator_different_mime_types() {
+        let validator = MockValidator { should_fail: false };
+        let config = ExtractionConfig::default();
+
+        let mime_types = vec![
+            "text/plain",
+            "application/pdf",
+            "application/json",
+            "text/html",
+            "image/png",
+        ];
+
+        for mime_type in mime_types {
+            let result = ExtractionResult {
+                content: "test".to_string(),
+                mime_type: mime_type.to_string(),
+                metadata: crate::types::Metadata::default(),
+                tables: vec![],
+                detected_languages: None,
+                chunks: None,
+                images: None,
+                djot_content: None,
+                pages: None,
+                elements: None,
+            };
+
+            assert!(validator.validate(&result, &config).await.is_ok());
+        }
+    }
+
+    #[tokio::test]
+    async fn test_validator_long_content() {
+        let validator = MockValidator { should_fail: false };
+
+        let result = ExtractionResult {
+            content: "test content ".repeat(10000),
+            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,
+        };
+
+        let config = ExtractionConfig::default();
+        assert!(validator.validate(&result, &config).await.is_ok());
+    }
+}