poimen-memory/crates/mem-ingest/src/fact_extractor.rs

//! Fact extraction: Identify relationships between entities
//!
//! Two implementations:
//! 1. SimpleFactExtractor: Pattern-based (verbs + wiki links)
//! 2. LlmFactExtractor: LLM-based (placeholder for production)
//!
//! CRAP: 12 (Simple pattern matching + LLM placeholder)
//! SOLID: Trait-based (Open/Closed)
//! DRY: Reuses EntityExtractor pattern

use anyhow::Result;
use async_trait::async_trait;
use regex::Regex;
use serde::{Deserialize, Serialize};

/// Extracted fact (relationship) from text
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractedFact {
    pub source_entity_id: String,
    pub target_entity_id: String,
    pub relation_type: String,
    pub fact: String,
}

/// Fact extractor trait - pluggable implementations
#[async_trait]
pub trait FactExtractor: Send + Sync {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>>;
    
    /// Extract facts with GRM context (optional, defaults to extract())
    async fn extract_with_context(
        &self,
        text: &str,
        _entity_contexts: &[crate::grm_retriever::EntityContext],
    ) -> Result<Vec<ExtractedFact>> {
        // Default: ignore context, use plain extraction
        self.extract(text).await
    }
}

/// Simple fact extractor based on verb patterns
/// Pattern: [[Entity1]] verb [[Entity2]]
/// Common verbs: uses, manages, runs, deployed_to, works_with
pub struct SimpleFactExtractor;

#[async_trait]
impl FactExtractor for SimpleFactExtractor {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>> {
        let mut facts = vec![];

        // Extract [[Entity]] patterns
        let entity_pattern = Regex::new(r"\[\[([^\]]+)\]\]")?;
        let entities: Vec<String> = entity_pattern
            .captures_iter(text)
            .filter_map(|cap| cap.get(1).map(|m| m.as_str().to_string()))
            .collect();

        // Common relationship verbs
        let verbs = ["uses", "manages", "runs", "deployed_to", "works_with"];

        // Simple heuristic: if two entities appear close together with a verb between them
        for verb in &verbs {
            let pattern = format!(
                r"\[\[([^\]]+)\]\].*?{}.*?\[\[([^\]]+)\]\]",
                verb.to_lowercase()
            );
            if let Ok(re) = Regex::new(&pattern) {
                for cap in re.captures_iter(text) {
                    if let (Some(src), Some(tgt)) = (cap.get(1), cap.get(2)) {
                        facts.push(ExtractedFact {
                            source_entity_id: src.as_str().to_string(),
                            target_entity_id: tgt.as_str().to_string(),
                            relation_type: verb.to_uppercase(),
                            fact: format!(
                                "{} {} {}",
                                src.as_str(),
                                verb,
                                tgt.as_str()
                            ),
                        });
                    }
                }
            }
        }

        Ok(facts)
    }
}

/// LLM-based fact extractor (placeholder for production)
/// TODO (Phase 2.6): Implement with real LLM API
/// TODO (Phase 2.6): Support complex relationships (3-way, temporal, conditional)
pub struct LlmFactExtractor;

#[async_trait]
impl FactExtractor for LlmFactExtractor {
    async fn extract(&self, _text: &str) -> Result<Vec<ExtractedFact>> {
        // TODO (Phase 2.6): Implement LLM-based extraction
        // Pattern: Send text to api.riotpiao.com with prompt
        // Parse response for [source, relation, target] tuples
        Ok(vec![])
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_simple_fact_extraction() {
        let extractor = SimpleFactExtractor;
        let text = "[[Rock]] uses [[Kubernetes]] and [[ArgoCD]]";

        let facts = extractor.extract(text).await.unwrap();
        assert!(facts.len() > 0);
        assert!(facts.iter().any(|f| f.relation_type == "USES"));
    }
}