kreuzberg 4.0.0.rc1

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);
+    }
+}

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

@@ -0,0 +1,294 @@
+//! RAKE (Rapid Automatic Keyword Extraction) backend implementation.
+
+use super::config::{KeywordConfig, RakeParams};
+use super::types::{Keyword, KeywordAlgorithm};
+use crate::Result;
+use crate::stopwords::{STOPWORDS, get_stopwords};
+use rake::*;
+use std::collections::HashSet;
+
+/// Extract keywords using RAKE algorithm.
+///
+/// RAKE is a co-occurrence based keyword extraction method that:
+/// - Identifies candidate keywords using stop words as delimiters
+/// - Calculates word scores based on frequency and degree
+/// - Combines scores for multi-word phrases
+///
+/// # 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 keyword extraction fails.
+pub fn extract_keywords_rake(text: &str, config: &KeywordConfig) -> Result<Vec<Keyword>> {
+    let params = config.rake_params.as_ref().cloned().unwrap_or_default();
+    let lang = config.language.as_deref().unwrap_or("en");
+    let normalized_lang = normalize_language_code(lang);
+
+    let mut keywords = run_rake_with_stopwords(text, config, &params, stopwords_for_language(lang));
+
+    if keywords.is_empty() && normalized_lang != "en" {
+        keywords = run_rake_with_stopwords(text, config, &params, stopwords_for_language("en"));
+    }
+
+    Ok(keywords)
+}
+
+fn stopwords_for_language(lang: &str) -> StopWords {
+    if let Some(words) = get_stopwords(lang) {
+        let set: HashSet<String> = words.iter().cloned().collect();
+        return StopWords::from(set);
+    }
+
+    let english: HashSet<String> = STOPWORDS
+        .get("en")
+        .map(|set| set.iter().cloned().collect())
+        .unwrap_or_default();
+    StopWords::from(english)
+}
+
+fn run_rake_with_stopwords(
+    text: &str,
+    config: &KeywordConfig,
+    params: &RakeParams,
+    stopwords: StopWords,
+) -> Vec<Keyword> {
+    let rake = Rake::new(stopwords);
+    let results = rake.run(text);
+
+    let filtered_results: Vec<_> = results
+        .into_iter()
+        .filter_map(|keyword_score| {
+            let keyword = keyword_score.keyword.clone();
+
+            if keyword.len() < params.min_word_length {
+                return None;
+            }
+
+            let word_count = keyword.split_whitespace().count();
+            if word_count > params.max_words_per_phrase {
+                return None;
+            }
+
+            if word_count < config.ngram_range.0 || word_count > config.ngram_range.1 {
+                return None;
+            }
+
+            Some((keyword, keyword_score.score))
+        })
+        .collect();
+
+    if filtered_results.is_empty() {
+        return Vec::new();
+    }
+
+    let min_score = filtered_results.iter().map(|(_, s)| *s).fold(f64::INFINITY, f64::min);
+    let max_score = filtered_results
+        .iter()
+        .map(|(_, s)| *s)
+        .fold(f64::NEG_INFINITY, f64::max);
+
+    let mut keywords: Vec<_> = filtered_results
+        .into_iter()
+        .map(|(keyword, raw_score)| {
+            let normalized_score = if max_score > min_score {
+                ((raw_score - min_score) / (max_score - min_score)).clamp(0.0, 1.0)
+            } else {
+                1.0
+            };
+
+            Keyword::new(keyword, normalized_score as f32, KeywordAlgorithm::Rake)
+        })
+        .collect();
+
+    if config.min_score > 0.0 {
+        keywords.retain(|k| k.score >= config.min_score);
+    }
+
+    keywords.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap_or(std::cmp::Ordering::Equal));
+
+    keywords.truncate(config.max_keywords);
+
+    keywords
+}
+
+fn normalize_language_code(lang: &str) -> String {
+    let lower = lang.to_lowercase();
+    if let Some(pos) = lower.find(['-', '_']) {
+        lower[..pos].to_string()
+    } else if lower.len() >= 2 {
+        lower[..2].to_string()
+    } else {
+        lower
+    }
+}
+
+#[cfg(all(test, feature = "keywords"))]
+mod tests {
+    use super::*;
+    use crate::keywords::config::RakeParams;
+
+    #[test]
+    fn test_rake_extraction_basic() {
+        let text = "Rust is a systems programming language. \
+                    Rust provides memory safety and performance. \
+                    Memory safety is achieved without garbage collection.";
+
+        let config = KeywordConfig::rake();
+
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords");
+        assert!(
+            keywords.len() <= config.max_keywords,
+            "Should respect max_keywords limit"
+        );
+
+        for i in 1..keywords.len() {
+            assert!(
+                keywords[i - 1].score >= keywords[i].score,
+                "Keywords should be sorted by score"
+            );
+        }
+
+        for keyword in &keywords {
+            assert_eq!(keyword.algorithm, KeywordAlgorithm::Rake);
+        }
+    }
+
+    #[test]
+    fn test_rake_extraction_with_min_score() {
+        let text = "Rust programming language provides memory safety without garbage collection.";
+
+        let config = KeywordConfig::rake().with_min_score(0.3);
+
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            assert!(
+                keyword.score >= config.min_score,
+                "Keyword score {} should be >= min_score {}",
+                keyword.score,
+                config.min_score
+            );
+        }
+    }
+
+    #[test]
+    fn test_rake_extraction_with_ngram_range() {
+        let text = "Machine learning models require large datasets for training.";
+
+        let config = KeywordConfig::rake().with_ngram_range(1, 1);
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            assert_eq!(
+                keyword.text.split_whitespace().count(),
+                1,
+                "Should only extract unigrams"
+            );
+        }
+    }
+
+    #[test]
+    fn test_rake_extraction_empty_text() {
+        let config = KeywordConfig::rake();
+        let keywords = extract_keywords_rake("", &config).unwrap();
+        assert!(keywords.is_empty(), "Empty text should yield no keywords");
+    }
+
+    #[test]
+    fn test_rake_extraction_with_custom_params() {
+        let text = "Natural language processing enables computers to understand human language.";
+
+        let params = RakeParams {
+            min_word_length: 3,
+            max_words_per_phrase: 2,
+        };
+
+        let config = KeywordConfig::rake().with_rake_params(params);
+
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        for keyword in &keywords {
+            for word in keyword.text.split_whitespace() {
+                assert!(word.len() >= 3, "Word '{}' should have min length 3", word);
+            }
+        }
+
+        for keyword in &keywords {
+            assert!(
+                keyword.text.split_whitespace().count() <= 2,
+                "Keyword '{}' should have max 2 words",
+                keyword.text
+            );
+        }
+    }
+
+    #[cfg_attr(coverage, ignore = "coverage instrumentation interferes with stopword heuristics")]
+    #[test]
+    fn test_rake_multilingual() {
+        let spanish_text = "El idioma español es una lengua romance.";
+        let config = KeywordConfig::rake().with_language("es");
+        let keywords = extract_keywords_rake(spanish_text, &config).unwrap();
+        assert!(!keywords.is_empty(), "Should extract Spanish keywords");
+
+        assert!(
+            keywords
+                .iter()
+                .any(|k| k.text.contains("idioma") || k.text.contains("español") || k.text.contains("lengua"))
+        );
+
+        let english_text = "Natural language processing is a subfield of artificial intelligence.";
+        let config = KeywordConfig::rake().with_language("fr");
+        let keywords = extract_keywords_rake(english_text, &config).unwrap();
+        dbg!(&keywords);
+        assert!(
+            !keywords.is_empty(),
+            "Should fall back to English stopwords and extract keywords"
+        );
+    }
+
+    #[test]
+    fn test_rake_score_normalization() {
+        let text = "Rust is a systems programming language that provides memory safety and \
+                    thread safety without garbage collection. Rust uses a ownership system \
+                    with rules that the compiler checks at compile time.";
+
+        let config = KeywordConfig::rake();
+        let keywords = extract_keywords_rake(text, &config).unwrap();
+
+        assert!(!keywords.is_empty(), "Should extract keywords");
+
+        for keyword in &keywords {
+            assert!(
+                keyword.score >= 0.0 && keyword.score <= 1.0,
+                "Keyword '{}' score {} should be in range [0.0, 1.0]",
+                keyword.text,
+                keyword.score
+            );
+        }
+
+        if !keywords.is_empty() {
+            let max_score = keywords.iter().map(|k| k.score).fold(0.0f32, f32::max);
+            assert!(
+                (max_score - 1.0).abs() < 0.001,
+                "Max score should be 1.0, got {}",
+                max_score
+            );
+        }
+
+        for i in 1..keywords.len() {
+            assert!(
+                keywords[i - 1].score >= keywords[i].score,
+                "Keywords should be sorted by score"
+            );
+        }
+    }
+}

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

@@ -0,0 +1,68 @@
+//! Shared types for keyword extraction.
+
+use serde::{Deserialize, Serialize};
+
+/// Keyword algorithm selection.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "lowercase")]
+pub enum KeywordAlgorithm {
+    /// YAKE (Yet Another Keyword Extractor) - statistical approach
+    #[cfg(feature = "keywords-yake")]
+    Yake,
+
+    /// RAKE (Rapid Automatic Keyword Extraction) - co-occurrence based
+    #[cfg(feature = "keywords-rake")]
+    Rake,
+}
+
+impl Default for KeywordAlgorithm {
+    fn default() -> Self {
+        #[cfg(feature = "keywords-yake")]
+        return Self::Yake;
+
+        #[cfg(all(feature = "keywords-rake", not(feature = "keywords-yake")))]
+        return Self::Rake;
+
+        #[cfg(not(any(feature = "keywords-yake", feature = "keywords-rake")))]
+        compile_error!("At least one keyword extraction feature must be enabled");
+    }
+}
+
+/// Extracted keyword with metadata.
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct Keyword {
+    /// The keyword text.
+    pub text: String,
+
+    /// Relevance score (higher is better, algorithm-specific range).
+    pub score: f32,
+
+    /// Algorithm that extracted this keyword.
+    pub algorithm: KeywordAlgorithm,
+
+    /// Optional positions where keyword appears in text (character offsets).
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub positions: Option<Vec<usize>>,
+}
+
+impl Keyword {
+    /// Create a new keyword.
+    pub fn new(text: String, score: f32, algorithm: KeywordAlgorithm) -> Self {
+        Self {
+            text,
+            score,
+            algorithm,
+            positions: None,
+        }
+    }
+
+    /// Create a new keyword with positions.
+    pub fn with_positions(text: String, score: f32, algorithm: KeywordAlgorithm, positions: Vec<usize>) -> Self {
+        Self {
+            text,
+            score,
+            algorithm,
+            positions: Some(positions),
+        }
+    }
+}