kreuzberg 4.0.0.pre.rc.6

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

@@ -0,0 +1,237 @@
+//! Keyword extraction module.
+//!
+//! Provides unified keyword extraction interface supporting multiple algorithms:
+//! - YAKE (Yet Another Keyword Extractor) - statistical approach
+//! - RAKE (Rapid Automatic Keyword Extraction) - co-occurrence based
+//!
+//! # Feature Flags
+//!
+//! - `keywords-yake`: Enable YAKE algorithm
+//! - `keywords-rake`: Enable RAKE algorithm
+//! - `keywords`: Enable both algorithms (default in `full` feature)
+//!
+//! # Examples
+//!
+//! ```rust,no_run
+//! # use kreuzberg::keywords::{extract_keywords, KeywordConfig};
+//! let text = "Rust is a systems programming language focused on safety and performance.";
+//!
+//! // Use default algorithm (YAKE if available)
+//! let config = KeywordConfig::default();
+//! let keywords = extract_keywords(text, &config).unwrap();
+//!
+//! for keyword in keywords {
+//!     println!("{}: {:.3}", keyword.text, keyword.score);
+//! }
+//! ```
+//!
+//! ```rust,no_run
+//! # #[cfg(feature = "keywords-rake")]
+//! # {
+//! # use kreuzberg::keywords::{extract_keywords, KeywordConfig};
+//! // Use RAKE algorithm explicitly
+//! let text = "Machine learning models require large datasets.";
+//! let config = KeywordConfig::rake()
+//!     .with_max_keywords(5)
+//!     .with_min_score(0.3);
+//!
+//! let keywords = extract_keywords(text, &config).unwrap();
+//! # }
+//! ```
+
+use crate::Result;
+use crate::plugins::registry::get_post_processor_registry;
+use once_cell::sync::Lazy;
+use std::sync::Arc;
+
+pub mod config;
+pub mod processor;
+pub mod types;
+
+#[cfg(feature = "keywords-yake")]
+mod yake;
+
+#[cfg(feature = "keywords-rake")]
+mod rake;
+
+pub use config::KeywordConfig;
+pub use processor::KeywordExtractor;
+
+#[cfg(feature = "keywords-rake")]
+pub use config::RakeParams;
+
+#[cfg(feature = "keywords-yake")]
+pub use config::YakeParams;
+pub use types::{Keyword, KeywordAlgorithm};
+
+/// Extract keywords from text using the specified algorithm.
+///
+/// This is the unified entry point for keyword extraction. The algorithm
+/// used is determined by `config.algorithm`.
+///
+/// # Arguments
+///
+/// * `text` - The text to extract keywords from
+/// * `config` - Keyword extraction configuration
+///
+/// # Returns
+///
+/// A vector of keywords sorted by relevance (highest score first).
+///
+/// # Errors
+///
+/// Returns an error if:
+/// - The specified algorithm feature is not enabled
+/// - Keyword extraction fails
+///
+/// # Examples
+///
+/// ```rust,no_run
+/// # use kreuzberg::keywords::{extract_keywords, KeywordConfig};
+/// let text = "Document intelligence with Rust provides memory safety.";
+/// let config = KeywordConfig::default()
+///     .with_max_keywords(10)
+///     .with_language("en");
+///
+/// let keywords = extract_keywords(text, &config)?;
+///
+/// for keyword in keywords {
+///     println!("{}: {:.3}", keyword.text, keyword.score);
+/// }
+/// # Ok::<(), kreuzberg::KreuzbergError>(())
+/// ```
+pub fn extract_keywords(text: &str, config: &KeywordConfig) -> Result<Vec<Keyword>> {
+    match config.algorithm {
+        #[cfg(feature = "keywords-yake")]
+        KeywordAlgorithm::Yake => yake::extract_keywords_yake(text, config),
+
+        #[cfg(feature = "keywords-rake")]
+        KeywordAlgorithm::Rake => rake::extract_keywords_rake(text, config),
+
+        #[cfg(not(any(feature = "keywords-yake", feature = "keywords-rake")))]
+        _ => Err(crate::KreuzbergError::Other(
+            "No keyword extraction algorithm feature enabled".to_string(),
+        )),
+    }
+}
+
+/// Lazy-initialized flag that ensures keyword processor is registered exactly once.
+///
+/// This static is accessed on first use to automatically register the
+/// keyword extraction processor with the plugin registry.
+static PROCESSOR_INITIALIZED: Lazy<Result<()>> = Lazy::new(register_keyword_processor);
+
+/// Ensure the keyword processor is registered.
+///
+/// This function is called automatically when needed.
+/// It's safe to call multiple times - registration only happens once.
+pub fn ensure_initialized() -> Result<()> {
+    PROCESSOR_INITIALIZED
+        .as_ref()
+        .map(|_| ())
+        .map_err(|e| crate::KreuzbergError::Plugin {
+            message: format!("Failed to register keyword processor: {}", e),
+            plugin_name: "keyword-extraction".to_string(),
+        })
+}
+
+/// Register the keyword extraction processor with the global registry.
+///
+/// This function should be called once at application startup to register
+/// the keyword extraction post-processor.
+///
+/// **Note:** This is called automatically on first use.
+/// Explicit calling is optional.
+///
+/// # Example
+///
+/// ```rust
+/// # #[cfg(any(feature = "keywords-yake", feature = "keywords-rake"))]
+/// use kreuzberg::keywords::register_keyword_processor;
+///
+/// # #[cfg(any(feature = "keywords-yake", feature = "keywords-rake"))]
+/// # fn main() -> kreuzberg::Result<()> {
+/// register_keyword_processor()?;
+/// # Ok(())
+/// # }
+/// # #[cfg(not(any(feature = "keywords-yake", feature = "keywords-rake")))]
+/// # fn main() {}
+/// ```
+pub fn register_keyword_processor() -> Result<()> {
+    let registry = get_post_processor_registry();
+    let mut registry = registry
+        .write()
+        .map_err(|e| crate::KreuzbergError::Other(format!("Post-processor registry lock poisoned: {}", e)))?;
+
+    registry.register(Arc::new(KeywordExtractor), 50)?;
+
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_extract_keywords_default_algorithm() {
+        let text = "Rust programming language provides memory safety and performance.";
+        let config = KeywordConfig::default();
+
+        let keywords = extract_keywords(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords");
+        assert!(keywords.len() <= config.max_keywords);
+    }
+
+    #[cfg(feature = "keywords-yake")]
+    #[test]
+    fn test_extract_keywords_yake() {
+        let text = "Natural language processing using Rust is efficient and safe.";
+        let config = KeywordConfig::yake();
+
+        let keywords = extract_keywords(text, &config).unwrap();
+
+        assert!(!keywords.is_empty());
+        assert_eq!(keywords[0].algorithm, KeywordAlgorithm::Yake);
+    }
+
+    #[cfg(feature = "keywords-rake")]
+    #[test]
+    fn test_extract_keywords_rake() {
+        let text = "Natural language processing using Rust is efficient and safe.";
+        let config = KeywordConfig::rake();
+
+        let keywords = extract_keywords(text, &config).unwrap();
+
+        assert!(!keywords.is_empty());
+        assert_eq!(keywords[0].algorithm, KeywordAlgorithm::Rake);
+    }
+
+    #[cfg(all(feature = "keywords-yake", feature = "keywords-rake"))]
+    #[test]
+    fn test_compare_algorithms() {
+        let text = "Machine learning and artificial intelligence are transforming technology. \
+                    Deep learning models require substantial computational resources.";
+
+        let yake_config = KeywordConfig::yake().with_max_keywords(5);
+        let yake_keywords = extract_keywords(text, &yake_config).unwrap();
+
+        let rake_config = KeywordConfig::rake().with_max_keywords(5);
+        let rake_keywords = extract_keywords(text, &rake_config).unwrap();
+
+        assert!(!yake_keywords.is_empty());
+        assert!(!rake_keywords.is_empty());
+
+        assert!(yake_keywords.iter().all(|k| k.algorithm == KeywordAlgorithm::Yake));
+        assert!(rake_keywords.iter().all(|k| k.algorithm == KeywordAlgorithm::Rake));
+
+        println!(
+            "YAKE keywords: {:?}",
+            yake_keywords.iter().map(|k| &k.text).collect::<Vec<_>>()
+        );
+        println!(
+            "RAKE keywords: {:?}",
+            rake_keywords.iter().map(|k| &k.text).collect::<Vec<_>>()
+        );
+    }
+}

data/vendor/kreuzberg/src/keywords/processor.rs

@@ -0,0 +1,267 @@
+//! Keyword extraction post-processor.
+//!
+//! This module provides a PostProcessor plugin that extracts keywords from
+//! extraction results and stores them in metadata.
+
+use crate::plugins::{Plugin, PostProcessor, ProcessingStage};
+use crate::{ExtractionConfig, ExtractionResult, KreuzbergError, Result};
+use async_trait::async_trait;
+
+/// Post-processor that extracts keywords from document content.
+///
+/// This processor:
+/// - Runs in the Middle processing stage
+/// - Only processes when `config.keywords` is configured
+/// - Stores extracted keywords in `metadata.additional["keywords"]`
+/// - Uses the configured algorithm (YAKE or RAKE)
+///
+/// # Example
+///
+/// ```rust,no_run
+/// use kreuzberg::plugins::{Plugin, PostProcessor};
+/// use kreuzberg::keywords::processor::KeywordExtractor;
+///
+/// let processor = KeywordExtractor;
+/// assert_eq!(processor.name(), "keyword-extraction");
+/// ```
+#[derive(Debug, Clone, Copy)]
+pub struct KeywordExtractor;
+
+impl Plugin for KeywordExtractor {
+    fn name(&self) -> &str {
+        "keyword-extraction"
+    }
+
+    fn version(&self) -> String {
+        env!("CARGO_PKG_VERSION").to_string()
+    }
+
+    fn initialize(&self) -> Result<()> {
+        Ok(())
+    }
+
+    fn shutdown(&self) -> Result<()> {
+        Ok(())
+    }
+}
+
+#[async_trait]
+impl PostProcessor for KeywordExtractor {
+    async fn process(&self, result: &mut ExtractionResult, config: &ExtractionConfig) -> Result<()> {
+        let keyword_config = match &config.keywords {
+            Some(cfg) => cfg,
+            None => return Ok(()),
+        };
+
+        let word_count = result.content.split_whitespace().count();
+        if word_count < 10 {
+            return Ok(());
+        }
+
+        let keywords = super::extract_keywords(&result.content, keyword_config)
+            .map_err(|e| KreuzbergError::Other(format!("Keyword extraction failed: {}", e)))?;
+
+        result
+            .metadata
+            .additional
+            .insert("keywords".to_string(), serde_json::to_value(&keywords)?);
+
+        Ok(())
+    }
+
+    fn processing_stage(&self) -> ProcessingStage {
+        ProcessingStage::Middle
+    }
+
+    fn should_process(&self, _result: &ExtractionResult, config: &ExtractionConfig) -> bool {
+        config.keywords.is_some()
+    }
+
+    fn estimated_duration_ms(&self, result: &ExtractionResult) -> u64 {
+        let word_count = result.content.split_whitespace().count();
+        (word_count as u64) / 100 + 10
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::keywords::KeywordConfig;
+    use crate::types::Metadata;
+
+    const TEST_TEXT: &str = r#"
+Machine learning is a branch of artificial intelligence that focuses on
+building systems that can learn from data. Deep learning is a subset of
+machine learning that uses neural networks with multiple layers.
+    "#;
+
+    #[tokio::test]
+    #[cfg(feature = "keywords-yake")]
+    async fn test_keyword_processor_with_yake() {
+        let processor = KeywordExtractor;
+        let config = ExtractionConfig {
+            keywords: Some(KeywordConfig::yake()),
+            ..Default::default()
+        };
+
+        let mut result = ExtractionResult {
+            content: TEST_TEXT.to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        processor.process(&mut result, &config).await.unwrap();
+
+        assert!(result.metadata.additional.contains_key("keywords"));
+        let keywords = result.metadata.additional.get("keywords").unwrap();
+        assert!(keywords.is_array());
+        let kw_array = keywords.as_array().unwrap();
+        assert!(!kw_array.is_empty());
+    }
+
+    #[tokio::test]
+    #[cfg(feature = "keywords-rake")]
+    async fn test_keyword_processor_with_rake() {
+        let processor = KeywordExtractor;
+        let config = ExtractionConfig {
+            keywords: Some(KeywordConfig::rake()),
+            ..Default::default()
+        };
+
+        let mut result = ExtractionResult {
+            content: TEST_TEXT.to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        processor.process(&mut result, &config).await.unwrap();
+
+        assert!(result.metadata.additional.contains_key("keywords"));
+        let keywords = result.metadata.additional.get("keywords").unwrap();
+        assert!(keywords.is_array());
+        let kw_array = keywords.as_array().unwrap();
+        assert!(!kw_array.is_empty());
+    }
+
+    #[tokio::test]
+    async fn test_keyword_processor_no_config() {
+        let processor = KeywordExtractor;
+        let config = ExtractionConfig::default();
+
+        let mut result = ExtractionResult {
+            content: TEST_TEXT.to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        processor.process(&mut result, &config).await.unwrap();
+
+        assert!(!result.metadata.additional.contains_key("keywords"));
+    }
+
+    #[tokio::test]
+    #[cfg(feature = "keywords-yake")]
+    async fn test_keyword_processor_short_content() {
+        let processor = KeywordExtractor;
+        let config = ExtractionConfig {
+            keywords: Some(KeywordConfig::yake()),
+            ..Default::default()
+        };
+
+        let mut result = ExtractionResult {
+            content: "Short text".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        processor.process(&mut result, &config).await.unwrap();
+
+        assert!(!result.metadata.additional.contains_key("keywords"));
+    }
+
+    #[test]
+    fn test_keyword_processor_plugin_interface() {
+        let processor = KeywordExtractor;
+        assert_eq!(processor.name(), "keyword-extraction");
+        assert!(!processor.version().is_empty());
+        assert!(processor.initialize().is_ok());
+        assert!(processor.shutdown().is_ok());
+    }
+
+    #[test]
+    fn test_keyword_processor_stage() {
+        let processor = KeywordExtractor;
+        assert_eq!(processor.processing_stage(), ProcessingStage::Middle);
+    }
+
+    #[test]
+    #[cfg(feature = "keywords-yake")]
+    fn test_keyword_processor_should_process() {
+        let processor = KeywordExtractor;
+
+        let result = ExtractionResult {
+            content: TEST_TEXT.to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        let config_with_keywords = ExtractionConfig {
+            keywords: Some(KeywordConfig::yake()),
+            ..Default::default()
+        };
+        assert!(processor.should_process(&result, &config_with_keywords));
+
+        let config_without_keywords = ExtractionConfig::default();
+        assert!(!processor.should_process(&result, &config_without_keywords));
+    }
+
+    #[test]
+    fn test_keyword_processor_estimated_duration() {
+        let processor = KeywordExtractor;
+
+        let short_result = ExtractionResult {
+            content: "Short text with just a few words".to_string(),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        let long_result = ExtractionResult {
+            content: "word ".repeat(1000),
+            mime_type: "text/plain".to_string(),
+            metadata: Metadata::default(),
+            tables: vec![],
+            detected_languages: None,
+            chunks: None,
+            images: None,
+        };
+
+        let short_duration = processor.estimated_duration_ms(&short_result);
+        let long_duration = processor.estimated_duration_ms(&long_result);
+
+        assert!(long_duration > short_duration);
+    }
+}