kreuzberg 4.0.8 → 4.1.0

data/vendor/kreuzberg/src/core/mod.rs

@@ -41,8 +41,11 @@ pub mod mime;
 pub mod pipeline;
 pub mod server_config;
 
+#[cfg(feature = "pdf")]
+pub use config::HierarchyConfig;
 pub use config::{
-    ChunkingConfig, ExtractionConfig, ImageExtractionConfig, LanguageDetectionConfig, OcrConfig, TokenReductionConfig,
+    ChunkingConfig, EmbeddingConfig, EmbeddingModelType, ExtractionConfig, ImageExtractionConfig,
+    LanguageDetectionConfig, OcrConfig, OutputFormat, PageConfig, PostProcessorConfig, TokenReductionConfig,
 };
 pub use config_validation::{
     validate_binarization_method, validate_chunking_params, validate_confidence, validate_dpi, validate_language_code,

data/vendor/kreuzberg/src/core/pipeline/cache.rs

@@ -0,0 +1,60 @@
+//! Processor caching to reduce lock contention.
+//!
+//! This module manages the caching of post-processors by processing stage,
+//! eliminating repeated registry lock acquisitions.
+
+use crate::Result;
+use crate::plugins::{PostProcessor, ProcessingStage};
+use once_cell::sync::Lazy;
+use std::sync::Arc;
+use std::sync::RwLock as StdRwLock;
+
+/// Cached post-processors for each stage to reduce lock contention.
+///
+/// This cache is populated once during the first pipeline run and reused
+/// for all subsequent extractions, eliminating 3 of 4 registry lock acquisitions
+/// per extraction.
+pub(super) struct ProcessorCache {
+    pub(super) early: Arc<Vec<Arc<dyn PostProcessor>>>,
+    pub(super) middle: Arc<Vec<Arc<dyn PostProcessor>>>,
+    pub(super) late: Arc<Vec<Arc<dyn PostProcessor>>>,
+}
+
+impl ProcessorCache {
+    /// Create a new processor cache by fetching from the registry.
+    pub(super) fn new() -> Result<Self> {
+        let processor_registry = crate::plugins::registry::get_post_processor_registry();
+        let registry = processor_registry
+            .read()
+            .map_err(|e| crate::KreuzbergError::Other(format!("Post-processor registry lock poisoned: {}", e)))?;
+
+        Ok(Self {
+            early: Arc::new(registry.get_for_stage(ProcessingStage::Early)),
+            middle: Arc::new(registry.get_for_stage(ProcessingStage::Middle)),
+            late: Arc::new(registry.get_for_stage(ProcessingStage::Late)),
+        })
+    }
+
+    /// Get processors for a specific stage from cache.
+    #[allow(dead_code)]
+    pub(super) fn get_for_stage(&self, stage: ProcessingStage) -> Arc<Vec<Arc<dyn PostProcessor>>> {
+        match stage {
+            ProcessingStage::Early => Arc::clone(&self.early),
+            ProcessingStage::Middle => Arc::clone(&self.middle),
+            ProcessingStage::Late => Arc::clone(&self.late),
+        }
+    }
+}
+
+/// Lazy processor cache - initialized on first use, then cached.
+pub(super) static PROCESSOR_CACHE: Lazy<StdRwLock<Option<ProcessorCache>>> = Lazy::new(|| StdRwLock::new(None));
+
+/// Clear the processor cache (primarily for testing when registry changes).
+#[allow(dead_code)]
+pub fn clear_processor_cache() -> Result<()> {
+    let mut cache = PROCESSOR_CACHE
+        .write()
+        .map_err(|e| crate::KreuzbergError::Other(format!("Processor cache lock poisoned: {}", e)))?;
+    *cache = None;
+    Ok(())
+}

data/vendor/kreuzberg/src/core/pipeline/execution.rs

@@ -0,0 +1,89 @@
+//! Core processor execution logic.
+//!
+//! This module handles the execution of post-processors and validators
+//! in the correct order.
+
+use crate::core::config::ExtractionConfig;
+use crate::plugins::ProcessingStage;
+use crate::types::ExtractionResult;
+use crate::{KreuzbergError, Result};
+
+/// Execute all registered post-processors by stage.
+pub(super) async fn execute_processors(
+    result: &mut ExtractionResult,
+    config: &ExtractionConfig,
+    pp_config: &Option<&crate::core::config::PostProcessorConfig>,
+    early_processors: std::sync::Arc<Vec<std::sync::Arc<dyn crate::plugins::PostProcessor>>>,
+    middle_processors: std::sync::Arc<Vec<std::sync::Arc<dyn crate::plugins::PostProcessor>>>,
+    late_processors: std::sync::Arc<Vec<std::sync::Arc<dyn crate::plugins::PostProcessor>>>,
+) -> Result<()> {
+    for (_stage, processors_arc) in [
+        (ProcessingStage::Early, early_processors),
+        (ProcessingStage::Middle, middle_processors),
+        (ProcessingStage::Late, late_processors),
+    ] {
+        for processor in processors_arc.iter() {
+            let processor_name = processor.name();
+
+            let should_run = should_processor_run(pp_config, processor_name);
+
+            if should_run && processor.should_process(result, config) {
+                match processor.process(result, config).await {
+                    Ok(_) => {}
+                    Err(err @ KreuzbergError::Io(_))
+                    | Err(err @ KreuzbergError::LockPoisoned(_))
+                    | Err(err @ KreuzbergError::Plugin { .. }) => {
+                        return Err(err);
+                    }
+                    Err(err) => {
+                        result.metadata.additional.insert(
+                            format!("processing_error_{processor_name}"),
+                            serde_json::Value::String(err.to_string()),
+                        );
+                    }
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
+/// Determine if a processor should run based on configuration.
+fn should_processor_run(pp_config: &Option<&crate::core::config::PostProcessorConfig>, processor_name: &str) -> bool {
+    if let Some(config) = pp_config {
+        if let Some(ref enabled_set) = config.enabled_set {
+            enabled_set.contains(processor_name)
+        } else if let Some(ref disabled_set) = config.disabled_set {
+            !disabled_set.contains(processor_name)
+        } else if let Some(ref enabled) = config.enabled_processors {
+            enabled.iter().any(|name| name == processor_name)
+        } else if let Some(ref disabled) = config.disabled_processors {
+            !disabled.iter().any(|name| name == processor_name)
+        } else {
+            true
+        }
+    } else {
+        true
+    }
+}
+
+/// Execute all registered validators.
+pub(super) async fn execute_validators(result: &ExtractionResult, config: &ExtractionConfig) -> Result<()> {
+    let validator_registry = crate::plugins::registry::get_validator_registry();
+    let validators = {
+        let registry = validator_registry
+            .read()
+            .map_err(|e| crate::KreuzbergError::Other(format!("Validator registry lock poisoned: {}", e)))?;
+        registry.get_all()
+    };
+
+    if !validators.is_empty() {
+        for validator in validators {
+            if validator.should_validate(result, config) {
+                validator.validate(result, config).await?;
+            }
+        }
+    }
+
+    Ok(())
+}

data/vendor/kreuzberg/src/core/pipeline/features.rs

@@ -0,0 +1,108 @@
+//! Feature processing logic.
+//!
+//! This module handles feature-specific processing like chunking,
+//! embedding generation, and language detection.
+
+use crate::Result;
+use crate::core::config::ExtractionConfig;
+use crate::types::ExtractionResult;
+
+/// Execute chunking if configured.
+pub(super) fn execute_chunking(result: &mut ExtractionResult, config: &ExtractionConfig) -> Result<()> {
+    #[cfg(feature = "chunking")]
+    if let Some(ref chunking_config) = config.chunking {
+        let chunk_config = crate::chunking::ChunkingConfig {
+            max_characters: chunking_config.max_chars,
+            overlap: chunking_config.max_overlap,
+            trim: true,
+            chunker_type: crate::chunking::ChunkerType::Text,
+        };
+
+        let page_boundaries = result.metadata.pages.as_ref().and_then(|ps| ps.boundaries.as_deref());
+
+        match crate::chunking::chunk_text(&result.content, &chunk_config, page_boundaries) {
+            Ok(chunking_result) => {
+                result.chunks = Some(chunking_result.chunks);
+
+                if let Some(ref chunks) = result.chunks {
+                    result.metadata.additional.insert(
+                        "chunk_count".to_string(),
+                        serde_json::Value::Number(serde_json::Number::from(chunks.len())),
+                    );
+                }
+
+                #[cfg(feature = "embeddings")]
+                if let Some(ref embedding_config) = chunking_config.embedding
+                    && let Some(ref mut chunks) = result.chunks
+                {
+                    match crate::embeddings::generate_embeddings_for_chunks(chunks, embedding_config) {
+                        Ok(()) => {
+                            result
+                                .metadata
+                                .additional
+                                .insert("embeddings_generated".to_string(), serde_json::Value::Bool(true));
+                        }
+                        Err(e) => {
+                            result
+                                .metadata
+                                .additional
+                                .insert("embedding_error".to_string(), serde_json::Value::String(e.to_string()));
+                        }
+                    }
+                }
+
+                #[cfg(not(feature = "embeddings"))]
+                if chunking_config.embedding.is_some() {
+                    result.metadata.additional.insert(
+                        "embedding_error".to_string(),
+                        serde_json::Value::String("Embeddings feature not enabled".to_string()),
+                    );
+                }
+            }
+            Err(e) => {
+                result
+                    .metadata
+                    .additional
+                    .insert("chunking_error".to_string(), serde_json::Value::String(e.to_string()));
+            }
+        }
+    }
+
+    #[cfg(not(feature = "chunking"))]
+    if config.chunking.is_some() {
+        result.metadata.additional.insert(
+            "chunking_error".to_string(),
+            serde_json::Value::String("Chunking feature not enabled".to_string()),
+        );
+    }
+
+    Ok(())
+}
+
+/// Execute language detection if configured.
+pub(super) fn execute_language_detection(result: &mut ExtractionResult, config: &ExtractionConfig) -> Result<()> {
+    #[cfg(feature = "language-detection")]
+    if let Some(ref lang_config) = config.language_detection {
+        match crate::language_detection::detect_languages(&result.content, lang_config) {
+            Ok(detected) => {
+                result.detected_languages = detected;
+            }
+            Err(e) => {
+                result.metadata.additional.insert(
+                    "language_detection_error".to_string(),
+                    serde_json::Value::String(e.to_string()),
+                );
+            }
+        }
+    }
+
+    #[cfg(not(feature = "language-detection"))]
+    if config.language_detection.is_some() {
+        result.metadata.additional.insert(
+            "language_detection_error".to_string(),
+            serde_json::Value::String("Language detection feature not enabled".to_string()),
+        );
+    }
+
+    Ok(())
+}

data/vendor/kreuzberg/src/core/pipeline/format.rs

@@ -0,0 +1,392 @@
+//! Output format conversion for extraction results.
+//!
+//! This module handles conversion of extraction results to various output formats
+//! (Plain, Djot, Markdown, HTML) with proper error handling and metadata recording.
+
+use crate::core::config::OutputFormat;
+use crate::types::ExtractionResult;
+
+/// Apply output format conversion to the extraction result.
+///
+/// This function converts the result's content field based on the configured output format:
+/// - `Plain`: No conversion (default)
+/// - `Djot`: Use djot_content if available, otherwise keep plain text
+/// - `Markdown`: Convert to Markdown format (uses djot as it's similar)
+/// - `Html`: Convert to HTML format
+///
+/// Skips conversion if content was already formatted during extraction (e.g., HTML extractor
+/// already produced djot or markdown output).
+///
+/// # Arguments
+///
+/// * `result` - The extraction result to modify
+/// * `output_format` - The desired output format
+pub fn apply_output_format(result: &mut ExtractionResult, output_format: OutputFormat) {
+    // Check if content was already formatted during extraction
+    let already_formatted = match result.mime_type.as_str() {
+        "text/markdown" if output_format == OutputFormat::Markdown => true,
+        "text/djot" if output_format == OutputFormat::Djot => true,
+        _ => false,
+    };
+
+    if already_formatted {
+        return; // Skip re-conversion
+    }
+
+    match output_format {
+        OutputFormat::Plain => {
+            // Default - no conversion needed
+        }
+        OutputFormat::Djot => {
+            // Convert the extraction result to djot markup
+            match crate::extractors::djot_format::extraction_result_to_djot(result) {
+                Ok(djot_markup) => {
+                    result.content = djot_markup;
+                }
+                Err(e) => {
+                    // Keep original content on error, record error in metadata
+                    result.metadata.additional.insert(
+                        "output_format_error".to_string(),
+                        serde_json::Value::String(format!("Failed to convert to djot: {}", e)),
+                    );
+                }
+            }
+        }
+        OutputFormat::Markdown => {
+            // Djot is syntactically similar to Markdown, so we use djot output as a
+            // reasonable approximation. Full Markdown conversion would require a
+            // dedicated converter that handles the syntactic differences (e.g.,
+            // emphasis markers are swapped: djot uses _ for emphasis and * for strong,
+            // while CommonMark uses * for emphasis and ** for strong).
+            if result.djot_content.is_some() {
+                match crate::extractors::djot_format::extraction_result_to_djot(result) {
+                    Ok(djot_markup) => {
+                        result.content = djot_markup;
+                    }
+                    Err(e) => {
+                        // Keep original content on error, record error in metadata
+                        result.metadata.additional.insert(
+                            "output_format_error".to_string(),
+                            serde_json::Value::String(format!("Failed to convert to markdown: {}", e)),
+                        );
+                    }
+                }
+            }
+            // For non-djot documents, content remains as-is
+        }
+        OutputFormat::Html => {
+            // Convert to HTML format
+            if result.djot_content.is_some() {
+                // First generate djot markup, then convert to HTML
+                match crate::extractors::djot_format::extraction_result_to_djot(result) {
+                    Ok(djot_markup) => {
+                        match crate::extractors::djot_format::djot_to_html(&djot_markup) {
+                            Ok(html) => {
+                                result.content = html;
+                            }
+                            Err(e) => {
+                                // Keep original content on error, record error in metadata
+                                result.metadata.additional.insert(
+                                    "output_format_error".to_string(),
+                                    serde_json::Value::String(format!("Failed to convert djot to HTML: {}", e)),
+                                );
+                            }
+                        }
+                    }
+                    Err(e) => {
+                        // Keep original content on error, record error in metadata
+                        result.metadata.additional.insert(
+                            "output_format_error".to_string(),
+                            serde_json::Value::String(format!("Failed to generate djot for HTML conversion: {}", e)),
+                        );
+                    }
+                }
+            } else {
+                // For non-djot documents, wrap plain text in basic HTML
+                let escaped_content = html_escape(&result.content);
+                result.content = format!("<pre>{}</pre>", escaped_content);
+            }
+        }
+    }
+}
+
+/// Escape HTML special characters in a string.
+fn html_escape(s: &str) -> String {
+    s.replace('&', "&amp;")
+        .replace('<', "&lt;")
+        .replace('>', "&gt;")
+        .replace('"', "&quot;")
+        .replace('\'', "&#39;")
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::types::Metadata;
+
+    #[test]
+    fn test_apply_output_format_plain() {
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Plain);
+
+        // Plain format should not modify content
+        assert_eq!(result.content, "Hello World");
+    }
+
+    #[test]
+    fn test_apply_output_format_djot_with_djot_content() {
+        use crate::types::{BlockType, DjotContent, FormattedBlock, InlineElement, InlineType};
+
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/djot".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            elements: None,
+            djot_content: Some(DjotContent {
+                plain_text: "Hello World".to_string(),
+                blocks: vec![FormattedBlock {
+                    block_type: BlockType::Heading,
+                    level: Some(1),
+                    inline_content: vec![InlineElement {
+                        element_type: InlineType::Text,
+                        content: "Hello World".to_string(),
+                        attributes: None,
+                        metadata: None,
+                    }],
+                    attributes: None,
+                    language: None,
+                    code: None,
+                    children: vec![],
+                }],
+                metadata: Metadata::default(),
+                tables: vec![],
+                images: vec![],
+                links: vec![],
+                footnotes: vec![],
+                attributes: std::collections::HashMap::new(),
+            }),
+        };
+
+        apply_output_format(&mut result, OutputFormat::Djot);
+
+        // The content should still be present (the function preserves content)
+        assert!(!result.content.is_empty());
+    }
+
+    #[test]
+    fn test_apply_output_format_djot_without_djot_content() {
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Djot);
+
+        // Without djot_content, content is converted to djot paragraphs
+        // extraction_result_to_djot creates paragraphs from plain text
+        assert!(result.content.contains("Hello World"));
+    }
+
+    #[test]
+    fn test_apply_output_format_html() {
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Html);
+
+        // For non-djot documents, HTML wraps content in <pre> tags
+        assert!(result.content.contains("<pre>"));
+        assert!(result.content.contains("Hello World"));
+        assert!(result.content.contains("</pre>"));
+    }
+
+    #[test]
+    fn test_apply_output_format_html_escapes_special_chars() {
+        let mut result = ExtractionResult {
+            content: "<script>alert('XSS')</script>".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Html);
+
+        // HTML special characters should be escaped
+        assert!(result.content.contains("&lt;"));
+        assert!(result.content.contains("&gt;"));
+        assert!(!result.content.contains("<script>"));
+    }
+
+    #[test]
+    fn test_apply_output_format_markdown() {
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Markdown);
+
+        // For non-djot documents without djot_content, markdown keeps content as-is
+        assert_eq!(result.content, "Hello World");
+    }
+
+    #[test]
+    fn test_apply_output_format_preserves_metadata() {
+        let mut additional = std::collections::HashMap::new();
+        additional.insert("custom_key".to_string(), serde_json::json!("custom_value"));
+        let metadata = Metadata {
+            title: Some("Test Title".to_string()),
+            additional,
+            ..Default::default()
+        };
+
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata,
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Djot);
+
+        // Metadata should be preserved
+        assert_eq!(result.metadata.title, Some("Test Title".to_string()));
+        assert_eq!(
+            result.metadata.additional.get("custom_key"),
+            Some(&serde_json::json!("custom_value"))
+        );
+    }
+
+    #[test]
+    fn test_apply_output_format_preserves_tables() {
+        use crate::types::Table;
+
+        let table = Table {
+            cells: vec![vec!["A".to_string(), "B".to_string()]],
+            markdown: "| A | B |".to_string(),
+            page_number: 1,
+        };
+
+        let mut result = ExtractionResult {
+            content: "Hello World".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![table],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            djot_content: None,
+            elements: None,
+        };
+
+        apply_output_format(&mut result, OutputFormat::Html);
+
+        // Tables should be preserved
+        assert_eq!(result.tables.len(), 1);
+        assert_eq!(result.tables[0].cells[0][0], "A");
+    }
+
+    #[test]
+    fn test_apply_output_format_preserves_djot_content() {
+        use crate::types::{BlockType, DjotContent, FormattedBlock, InlineElement, InlineType};
+
+        let djot_content = DjotContent {
+            plain_text: "test".to_string(),
+            blocks: vec![FormattedBlock {
+                block_type: BlockType::Paragraph,
+                level: None,
+                inline_content: vec![InlineElement {
+                    element_type: InlineType::Text,
+                    content: "test".to_string(),
+                    attributes: None,
+                    metadata: None,
+                }],
+                attributes: None,
+                language: None,
+                code: None,
+                children: vec![],
+            }],
+            metadata: Metadata::default(),
+            tables: vec![],
+            images: vec![],
+            links: vec![],
+            footnotes: vec![],
+            attributes: std::collections::HashMap::new(),
+        };
+
+        let mut result = ExtractionResult {
+            content: "test".to_string(),
+            mime_type: "text/djot".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+            pages: None,
+            elements: None,
+            djot_content: Some(djot_content),
+        };
+
+        apply_output_format(&mut result, OutputFormat::Djot);
+
+        // djot_content should still be present after format application
+        assert!(result.djot_content.is_some());
+        assert_eq!(result.djot_content.as_ref().unwrap().blocks.len(), 1);
+    }
+}