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

//! Fact extraction: Identify relationships between entities
//!
//! Three implementations:
//! 1. SimpleFactExtractor: Pattern-based (verbs + wiki links)
//! 2. LlmFactExtractor: LLM-based extraction with entity context
//! 3. Fallback chain: LLM → Simple pattern matching
//!
//! Aligned with Zep paper §2.2.2: Facts as edges between entity pairs,
//! with temporal extraction and dedup against existing edges.

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 entity context (Zep §2.2.2: facts between known entities)
    async fn extract_with_context(
        &self,
        text: &str,
        _entity_contexts: &[crate::grm_retriever::EntityContext],
    ) -> Result<Vec<ExtractedFact>> {
        self.extract(text).await
    }
}

/// Simple fact extractor based on verb patterns
/// Pattern: [[Entity1]] verb [[Entity2]]
pub struct SimpleFactExtractor;

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

        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();

        let verbs = ["uses", "manages", "runs", "deployed_to", "works_with",
                      "depends_on", "contains", "extends", "implements", "connects_to"];

        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 (Zep §2.2.2 alignment)
/// Extracts relationships between entity pairs using LLM
pub struct LlmFactExtractor {
    model_name: String,
    jwt_issuer: Option<std::sync::Arc<tokio::sync::Mutex<crate::authentik_jwt::AuthentikJwtIssuer>>>,
}

impl LlmFactExtractor {
    pub fn new(model_name: &str) -> Self {
        let jwt_issuer = crate::authentik_jwt::AuthentikJwtIssuer::from_env().ok();
        Self {
            model_name: model_name.to_string(),
            jwt_issuer: jwt_issuer.map(|iss| std::sync::Arc::new(tokio::sync::Mutex::new(iss))),
        }
    }

    /// Clean LLM response: strip thinking tags, markdown fences, extract JSON
    fn clean_llm_response(text: &str) -> String {
        let mut result = text.to_string();
        while let Some(start) = result.find("<think>") {
            if let Some(end) = result.find("</think>") {
                result = format!("{}{}", &result[..start], &result[end + 8..]);
            } else { break; }
        }
        result = result.replace("```json", "").replace("```", "");
        let trimmed = result.trim();
        if let Some(start) = trimmed.find('{') {
            if let Some(end) = trimmed.rfind('}') {
                return trimmed[start..=end].to_string();
            }
        }
        if let Some(start) = trimmed.find('[') {
            if let Some(end) = trimmed.rfind(']') {
                return format!("{{\"facts\": {}}}", &trimmed[start..=end]);
            }
        }
        trimmed.to_string()
    }

    async fn call_llm(&self, prompt: &str) -> Result<String> {
        let endpoint = std::env::var("LLM_ENDPOINT")
            .unwrap_or_else(|_| "http://localhost:11434/v1/chat/completions".to_string());

        // Get auth header: Authentik JWT if configured, else API key
        let auth_header = if let Some(jwt_issuer) = &self.jwt_issuer {
            let issuer = jwt_issuer.lock().await;
            match issuer.get_access_token().await {
                Ok(token) => format!("Bearer {}", token),
                Err(e) => {
                    tracing::warn!(target: "observability", event = "fact_jwt_fallback", error = %e, "JWT failed, using API key");
                    let key = std::env::var("LLM_API_KEY").unwrap_or_else(|_| "default-key".to_string());
                    format!("Bearer {}", key)
                }
            }
        } else {
            let key = std::env::var("LLM_API_KEY")
                .or_else(|_| std::env::var("MEM_API_KEY"))
                .unwrap_or_else(|_| "default-key".to_string());
            format!("Bearer {}", key)
        };

        let start = std::time::Instant::now();
        let client = reqwest::Client::new();
        let payload = serde_json::json!({
            "model": self.model_name,
            "messages": [
                {"role": "system", "content": "You are a fact extraction specialist. Extract relationships between entities from text. Output ONLY valid JSON."},
                {"role": "user", "content": prompt}
            ],
            "max_tokens": 12000,
            "temperature": 0.1
        });

        let response = client
            .post(&endpoint)
            .header("Authorization", &auth_header)
            .header("Content-Type", "application/json")
            .json(&payload)
            .timeout(std::time::Duration::from_secs(120))
            .send()
            .await?;

        let status = response.status();
        if !status.is_success() {
            let body = response.text().await.unwrap_or_default();
            tracing::warn!(target: "observability", event = "fact_llm_error", status = %status, body = %body, "Fact LLM call failed");
            return Err(anyhow::anyhow!("LLM API error: {}", status));
        }

        let elapsed = start.elapsed();
        let data: serde_json::Value = response.json().await?;
        
        // Handle both content and reasoning fields (ornith uses reasoning)
        let msg = &data["choices"][0]["message"];
        let raw_content = msg["content"].as_str().unwrap_or("").to_string();
        let raw_reasoning = msg["reasoning"].as_str().unwrap_or("").to_string();
        let raw = if !raw_content.trim().is_empty() { &raw_content } else { &raw_reasoning };
        let cleaned = Self::clean_llm_response(raw);

        let tokens = &data["usage"];
        tracing::info!(
            target: "observability",
            event = "llm_fact_call",
            model = %self.model_name,
            endpoint = %endpoint,
            raw_len = raw.len(),
            cleaned_len = cleaned.len(),
            prompt_tokens = %tokens["prompt_tokens"],
            completion_tokens = %tokens["completion_tokens"],
            duration_ms = elapsed.as_millis() as u64,
            has_reasoning = !raw_reasoning.is_empty(),
            "LLM fact extraction call complete"
        );
        Ok(cleaned)
    }
}

#[async_trait]
impl FactExtractor for LlmFactExtractor {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedFact>> {
        self.extract_with_context(text, &[]).await
    }

    async fn extract_with_context(
        &self,
        text: &str,
        entity_contexts: &[crate::grm_retriever::EntityContext],
    ) -> Result<Vec<ExtractedFact>> {
        // Build entity list for prompt
        let entity_names: Vec<&str> = entity_contexts
            .iter()
            .map(|e| e.entity_name.as_str())
            .collect();

        if entity_names.is_empty() {
            tracing::debug!("No entities provided, skipping fact extraction");
            return Ok(vec![]);
        }

        let prompt = format!(
            r#"Extract relationships (facts) between these entities from the text.

Entities: {:?}

Text:
"{}"

For each relationship provide:
- source: Entity name (must be from the list above)
- target: Entity name (must be from the list above)  
- relation: Verb/predicate describing the relationship (e.g., "uses", "manages", "is_part_of", "deployed_on")
- fact: One-sentence natural language description

CRITICAL: Only extract relationships EXPLICITLY stated or strongly implied. Source and target must both be from the entity list.

Respond in JSON:
{{"facts": [{{"source": "...", "target": "...", "relation": "...", "fact": "..."}}, ...]}}
"#,
            entity_names, text
        );

        let llm_ok = std::env::var("LLM_ENDPOINT").is_ok();
        let response = if llm_ok {
            match self.call_llm(&prompt).await {
                Ok(r) => r,
                Err(e) => {
                    tracing::warn!("Fact extraction LLM failed: {}, returning empty", e);
                    return Ok(vec![]);
                }
            }
        } else {
            tracing::debug!("LLM_ENDPOINT not set, skipping LLM fact extraction");
            return Ok(vec![]);
        };

        // Parse response
        #[derive(Deserialize)]
        struct FactResponse {
            facts: Vec<RawFact>,
        }
        #[derive(Deserialize)]
        struct RawFact {
            source: String,
            target: String,
            relation: String,
            fact: String,
        }

        // Try parsing, if trailing chars error try trimming to valid JSON
        let parsed = match serde_json::from_str::<FactResponse>(&response) {
            Ok(r) => Ok(r),
            Err(e) if e.to_string().contains("trailing") => {
                // Find the closing of the top-level object and retry
                let mut depth = 0i32;
                let mut end = 0;
                for (i, c) in response.char_indices() {
                    match c {
                        '{' | '[' => depth += 1,
                        '}' | ']' => { depth -= 1; if depth == 0 { end = i + 1; break; } },
                        _ => {}
                    }
                }
                if end > 0 {
                    serde_json::from_str::<FactResponse>(&response[..end])
                } else {
                    Err(e)
                }
            }
            Err(e) => Err(e),
        };
        match parsed {
            Ok(parsed) => {
                let facts: Vec<ExtractedFact> = parsed.facts
                    .into_iter()
                    .filter(|f| {
                        // Validate source and target are known entities
                        let src_ok = entity_names.iter().any(|e| e.eq_ignore_ascii_case(&f.source));
                        let tgt_ok = entity_names.iter().any(|e| e.eq_ignore_ascii_case(&f.target));
                        if !src_ok || !tgt_ok {
                            tracing::debug!(
                                "Dropping fact with unknown entity: {} -> {}",
                                f.source, f.target
                            );
                        }
                        src_ok && tgt_ok && f.source != f.target
                    })
                    .map(|f| ExtractedFact {
                        source_entity_id: f.source,
                        target_entity_id: f.target,
                        relation_type: f.relation.to_uppercase(),
                        fact: f.fact,
                    })
                    .collect();

                tracing::info!(
                    "LLM fact extraction: {} facts from {} entities",
                    facts.len(), entity_names.len()
                );
                Ok(facts)
            }
            Err(e) => {
                tracing::warn!("Fact extraction JSON parse failed: {}", e);
                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.is_empty());
        assert!(facts.iter().any(|f| f.relation_type == "USES"));
    }

    #[tokio::test]
    async fn test_simple_no_wiki_links() {
        let extractor = SimpleFactExtractor;
        let text = "Kubernetes uses etcd for storage";
        let facts = extractor.extract(text).await.unwrap();
        assert!(facts.is_empty()); // No [[wiki links]]
    }

    #[test]
    fn test_clean_llm_response() {
        let input = r#"<think>reasoning here</think>{"facts": [{"source": "A", "target": "B", "relation": "uses", "fact": "A uses B"}]}"#;
        let cleaned = LlmFactExtractor::clean_llm_response(input);
        assert!(cleaned.starts_with("{"));
        assert!(cleaned.contains("facts"));
    }

    #[test]
    fn test_strip_thinking_no_tags() {
        let input = r#"{"facts": []}"#;
        let cleaned = LlmFactExtractor::clean_llm_response(input);
        assert_eq!(cleaned, input);
    }

    #[tokio::test]
    async fn test_llm_fact_no_entities_returns_empty() {
        let extractor = LlmFactExtractor::new("test");
        let facts = extractor.extract_with_context("some text", &[]).await.unwrap();
        assert!(facts.is_empty());
    }
}