xberg 1.0.7

//! Summarisation post-processor.
//!
//! Runs at [`ProcessingStage::Middle`]. Triggers when
//! [`ExtractionConfig::summarization`] is `Some`. Dispatches on
//! [`SummaryStrategy`]:
//!
//! - [`SummaryStrategy::Extractive`] — pure-Rust TextRank via
//!   [`crate::text::summarization::textrank::summarize`]. Deterministic and
//!   available in every build.
//! - [`SummaryStrategy::Abstractive`] — LLM-driven via
//!   [`crate::text::summarization::llm::summarize_with_llm`]. Requires both
//!   the `summarization-llm` cargo feature and an `llm` config slot on
//!   [`crate::core::config::SummarizationConfig`]. When the feature is absent
//!   or `llm` is `None`, the processor returns a validation error so callers
//!   notice the misconfiguration.

use std::sync::Arc;

use crate::Result;
use crate::core::config::ExtractionConfig;
use crate::plugins::{Plugin, PostProcessor, ProcessingStage, register_post_processor};
use crate::types::ExtractedDocument;
use crate::types::summary::{DocumentSummary, SummaryStrategy};
use async_trait::async_trait;

/// Post-processor that produces a [`DocumentSummary`] for the extraction result.
#[cfg_attr(alef, alef(skip))]
#[derive(Debug, Clone, Copy, Default)]
pub struct SummarizationProcessor;

impl Plugin for SummarizationProcessor {
    fn name(&self) -> &str {
        "summarization"
    }

    fn version(&self) -> String {
        env!("CARGO_PKG_VERSION").to_string()
    }

    fn initialize(&self) -> Result<()> {
        Ok(())
    }

    fn shutdown(&self) -> Result<()> {
        Ok(())
    }
}

#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
impl PostProcessor for SummarizationProcessor {
    async fn process(&self, result: &mut ExtractedDocument, config: &ExtractionConfig) -> Result<()> {
        let summarization_config = match &config.summarization {
            Some(cfg) => cfg,
            None => return Ok(()),
        };

        if result.content.trim().is_empty() {
            return Ok(());
        }

        match summarization_config.strategy {
            SummaryStrategy::Extractive => run_extractive(result, summarization_config.max_tokens),
            SummaryStrategy::Abstractive => run_abstractive(result, summarization_config).await,
        }
    }

    fn processing_stage(&self) -> ProcessingStage {
        ProcessingStage::Middle
    }

    fn should_process(&self, _result: &ExtractedDocument, config: &ExtractionConfig) -> bool {
        config.summarization.is_some()
    }

    fn estimated_duration_ms(&self, result: &ExtractedDocument) -> u64 {
        let word_count = result.content.split_whitespace().count() as u64;
        word_count / 200 + 50
    }
}

fn run_extractive(result: &mut ExtractedDocument, max_tokens: Option<u32>) -> Result<()> {
    let language = result
        .detected_languages
        .as_ref()
        .and_then(|langs| langs.first())
        .map(String::as_str);

    let summary_text =
        crate::text::summarization::textrank::summarize(&result.content, language, max_tokens).unwrap_or_default();

    if summary_text.is_empty() {
        return Ok(());
    }

    let token_count = crate::text::summarization::textrank::token_count(&summary_text);
    result.summary = Some(DocumentSummary {
        text: summary_text,
        strategy: SummaryStrategy::Extractive,
        token_count: Some(token_count),
    });
    Ok(())
}

#[cfg(feature = "summarization-llm")]
async fn run_abstractive(
    result: &mut ExtractedDocument,
    summarization_config: &crate::core::config::SummarizationConfig,
) -> Result<()> {
    use crate::types::ProcessingWarning;
    use std::borrow::Cow;

    let llm_config = summarization_config.llm.as_ref().ok_or_else(|| {
        crate::XbergError::validation(
            "Abstractive summarisation requires `SummarizationConfig.llm` to be set".to_string(),
        )
    })?;

    let outcome = crate::text::summarization::llm::summarize_with_llm(
        &result.content,
        llm_config,
        summarization_config.max_tokens,
    )
    .await;

    let (text, usage) = match outcome {
        Ok(pair) => pair,
        Err(e) => {
            result.processing_warnings.push(ProcessingWarning {
                source: Cow::Borrowed("summarization_abstractive"),
                message: Cow::Owned(format!("LLM summarisation failed: {e}")),
            });
            return Ok(());
        }
    };

    if text.is_empty() {
        return Ok(());
    }

    let token_count = crate::text::summarization::textrank::token_count(&text);
    result.summary = Some(DocumentSummary {
        text,
        strategy: SummaryStrategy::Abstractive,
        token_count: Some(token_count),
    });

    if let Some(u) = usage {
        result.llm_usage.get_or_insert_with(Vec::new).push(u);
    }

    Ok(())
}

#[cfg(not(feature = "summarization-llm"))]
async fn run_abstractive(
    _result: &mut ExtractedDocument,
    _summarization_config: &crate::core::config::SummarizationConfig,
) -> Result<()> {
    Err(crate::XbergError::validation(
        "Abstractive summarisation requires the `summarization-llm` cargo feature".to_string(),
    ))
}

/// Register the default summarisation post-processor with the global registry.
#[cfg_attr(alef, alef(skip))]
pub fn register() -> Result<()> {
    register_post_processor(Arc::new(SummarizationProcessor))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::config::SummarizationConfig;
    use std::borrow::Cow;

    const PARAGRAPH: &str = "Machine learning is a branch of artificial intelligence. \
        It focuses on building systems that learn from data. \
        Deep learning is a subset of machine learning. \
        Neural networks are inspired by the human brain.";

    #[tokio::test]
    async fn extractive_processor_populates_summary() {
        let processor = SummarizationProcessor;
        let config = ExtractionConfig {
            summarization: Some(SummarizationConfig {
                strategy: SummaryStrategy::Extractive,
                max_tokens: Some(40),
                llm: None,
            }),
            ..Default::default()
        };
        let mut result = ExtractedDocument {
            content: PARAGRAPH.to_string(),
            mime_type: Cow::Borrowed("text/plain"),
            detected_languages: Some(vec!["en".to_string()]),
            ..Default::default()
        };

        processor.process(&mut result, &config).await.unwrap();
        let summary = result.summary.as_ref().expect("summary populated");
        assert_eq!(summary.strategy, SummaryStrategy::Extractive);
        assert!(!summary.text.is_empty());
    }

    #[tokio::test]
    async fn no_config_means_no_summary() {
        let processor = SummarizationProcessor;
        let config = ExtractionConfig::default();
        let mut result = ExtractedDocument {
            content: PARAGRAPH.to_string(),
            mime_type: Cow::Borrowed("text/plain"),
            ..Default::default()
        };
        processor.process(&mut result, &config).await.unwrap();
        assert!(result.summary.is_none());
    }

    #[tokio::test]
    async fn empty_content_skips() {
        let processor = SummarizationProcessor;
        let config = ExtractionConfig {
            summarization: Some(SummarizationConfig::default()),
            ..Default::default()
        };
        let mut result = ExtractedDocument {
            content: String::new(),
            mime_type: Cow::Borrowed("text/plain"),
            ..Default::default()
        };
        processor.process(&mut result, &config).await.unwrap();
        assert!(result.summary.is_none());
    }

    #[test]
    fn stage_is_middle() {
        assert_eq!(SummarizationProcessor.processing_stage(), ProcessingStage::Middle);
    }

    #[test]
    fn plugin_metadata() {
        let p = SummarizationProcessor;
        assert_eq!(p.name(), "summarization");
        assert!(!p.version().is_empty());
    }
}