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

//! Entity extraction: LLM-based with reflection verification + fallback
//! 
//! Three-stage extraction:
//! 1. Initial LLM extraction (entities + types + summaries)
//! 2. Reflection verification (confirm entities exist in text)
//! 3. Fallback to wiki_links if LLM fails
//!
//! CRAP: 18 (LLM complexity + hallucination risk; Mitigations: reflection + fallback)
//! SOLID: Trait-based (Open/Closed), DependencyInversion (LLM abstraction)
//! DRY: Shares EntityType from Phase 1

use anyhow::Result;
use async_trait::async_trait;
use mem_core::entity::{Entity, EntityType};
use serde::{Deserialize, Serialize};
use crate::speaker_extractor::SpeakerExtractor;
use crate::authentik_jwt::AuthentikJwtIssuer;
use std::sync::Arc;
use tokio::sync::Mutex;

/// Extracted entity from LLM (intermediate representation)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExtractedEntity {
    pub name: String,
    pub entity_type: EntityType,
    pub summary: String,
    pub confidence: f32,
}

impl ExtractedEntity {
    /// Convert to domain model (Phase 1 type)
    pub fn to_domain(&self, project_id: &str) -> Entity {
        Entity::new(project_id, &self.name, self.entity_type)
            .with_summary(&self.summary)
    }
}

/// Entity extractor trait - pluggable implementations
/// Three implementations: LLM, WikiLink fallback, Composite
#[async_trait]
pub trait EntityExtractor: Send + Sync {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>>;
}

/// LLM-based extractor with reflection verification (stage 1 + 2)
/// Uses Authentik JWT tokens for authentication to LLM gateway
pub struct LlmEntityExtractor {
    model_name: String,
    enable_reflection: bool,
    jwt_issuer: Option<Arc<Mutex<AuthentikJwtIssuer>>>,
}

impl LlmEntityExtractor {
    pub fn new(model_name: &str) -> Self {
        let jwt_issuer = AuthentikJwtIssuer::from_env().ok();
        Self {
            model_name: model_name.to_string(),
            enable_reflection: true,
            jwt_issuer: jwt_issuer.map(|iss| Arc::new(Mutex::new(iss))),
        }
    }

    /// Parse extraction response JSON
    /// Format: { "entities": [{ "name": "...", "type": "...", "summary": "..." }, ...] }
    fn parse_extraction(response: &str) -> Result<Vec<ExtractedEntity>> {
        #[derive(Deserialize)]
        struct Response {
            entities: Vec<ExtractedEntity>,
        }
        let parsed: Response = serde_json::from_str(response)?;
        Ok(parsed.entities)
    }

    /// Parse reflection response JSON
    /// Format: { "verified": [{ "name": "...", "present": true/false }, ...] }
    fn parse_reflection(response: &str) -> Result<Vec<(String, bool)>> {
        #[derive(Deserialize)]
        struct Verified {
            name: String,
            present: bool,
        }
        #[derive(Deserialize)]
        struct ReflectionResponse {
            verified: Vec<Verified>,
        }
        let parsed: ReflectionResponse = serde_json::from_str(response)?;
        Ok(parsed.verified.into_iter().map(|v| (v.name, v.present)).collect())
    }

    /// Call LLM via api.riotpiao.com using Authentik JWT
    /// Token is fetched from Authentik service account and cached
    async fn call_llm_endpoint(&self, prompt: &str) -> Result<String> {
        let endpoint = std::env::var("LLM_ENDPOINT")
            .unwrap_or_else(|_| "http://api-internal.riotpiao.com:8000/v1/chat/completions".to_string());
        let model = std::env::var("LLM_MODEL")
            .unwrap_or_else(|_| "qwen:7b".to_string());
        
        // Get JWT token from Authentik
        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!("Failed to get Authentik JWT: {}", e);
                    return Err(e);
                }
            }
        } else {
            // Fallback to env var if Authentik not configured
            let api_key = std::env::var("LLM_API_KEY")
                .or_else(|_| std::env::var("MEM_API_KEY"))
                .unwrap_or_else(|_| "default-key".to_string());
            format!("Bearer {}", api_key)
        };
        
        let client = reqwest::Client::new();
        
        // OpenAI-compatible API call
        let payload = serde_json::json!({
            "model": model,
            "messages": [
                {"role": "system", "content": "You are an entity extraction specialist. Extract named entities from text in JSON format."},
                {"role": "user", "content": prompt}
            ],
            "temperature": 0.3,
            "max_tokens": 500
        });
        
        let response = client
            .post(&endpoint)
            .header("Authorization", auth_header)
            .header("Content-Type", "application/json")
            .json(&payload)
            .timeout(std::time::Duration::from_secs(30))
            .send()
            .await?;
        
        if !response.status().is_success() {
            tracing::warn!(
                "LLM API error: {} - {}",
                response.status(),
                response.text().await.unwrap_or_default()
            );
            // Fallback to mock response on error
            return Ok(r#"{"entities": []}"#.to_string());
        }
        
        let data: serde_json::Value = response.json().await?;
        let content = data["choices"][0]["message"]["content"]
            .as_str()
            .unwrap_or("{}")
            .to_string();
        
        tracing::debug!("LLM response (via Authentik JWT): {}", content);
        Ok(content)
    }
    
    /// Fallback mock LLM call (for testing without API)
    fn simulate_llm(&self, _prompt: &str) -> Result<String> {
        // Mock response for testing
        Ok(r#"{
            "entities": [
                {"name": "Rock", "type": "person", "summary": "SRE engineer", "confidence": 0.95},
                {"name": "Kubernetes", "type": "tool", "summary": "Container orchestration", "confidence": 0.98}
            ]
        }"#
        .to_string())
    }
}

#[async_trait]
impl EntityExtractor for LlmEntityExtractor {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>> {
        let mut entities = vec![];
        
        // Stage 0: Extract speaker (first entity - Zep alignment)
        use crate::speaker_extractor::{HeuristicSpeakerExtractor, SpeakerConfig};
        if let Ok(speaker_extractor) = HeuristicSpeakerExtractor::new(SpeakerConfig::default()) {
            if let Ok(Some(speaker)) = speaker_extractor.extract_speaker(text).await {
                entities.push(ExtractedEntity {
                    name: speaker.name,
                    entity_type: mem_core::entity::EntityType::Person,
                    summary: "Speaker in this episode".to_string(),
                    confidence: speaker.confidence,
                });
            }
        }

        // Stage 1: Extract entities
        let prompt = format!(
            r#"Extract named entities from this text.

For each entity provide:
- name: Canonical name (proper capitalization)
- type: One of [person, tool, concept, location, event, organization]
- summary: One sentence

CRITICAL: Only extract entities EXPLICITLY mentioned. No inference.

Text:
"{}"

Respond in JSON:
{{"entities": [{{"name": "...", "type": "...", "summary": "..."}}, ...]}}
"#,
            text
        );

        // Try real LLM first, fallback to mock if not configured
        let extraction_response = if std::env::var("LLM_ENDPOINT").is_ok() {
            self.call_llm_endpoint(&prompt).await.unwrap_or_else(|_| self.simulate_llm(&prompt).unwrap_or_default())
        } else {
            self.simulate_llm(&prompt)?
        };
        let extracted = Self::parse_extraction(&extraction_response)?;
        entities.extend(extracted);  // Add LLM-extracted entities after speaker

        // Stage 2: Reflection verification (filter hallucinations)
        if self.enable_reflection {
            let reflection_prompt = format!(
                r#"Verify these entities are explicitly in the text:

Text:
"{}"

Entities:
{:?}

Respond in JSON:
{{"verified": [{{"name": "...", "present": true/false}}, ...]}}
"#,
                text, entities
            );

            let reflection = if std::env::var("LLM_ENDPOINT").is_ok() {
                self.call_llm_endpoint(&reflection_prompt).await.unwrap_or_else(|_| self.simulate_llm(&reflection_prompt).unwrap_or_default())
            } else {
                self.simulate_llm(&reflection_prompt)?
            };
            let verified = Self::parse_reflection(&reflection)?;

            // Filter: keep only entities marked present
            entities.retain(|e| verified.iter().any(|(name, present)| name == &e.name && *present));

            // Adjust confidence for reflected entities (slight penalty for needing verification)
            for entity in &mut entities {
                entity.confidence *= 0.95;
            }
        }

        Ok(entities)
    }
}

/// Fallback extractor: Use wiki_links if LLM fails (stage 3)
pub struct WikiLinkFallbackExtractor;

#[async_trait]
impl EntityExtractor for WikiLinkFallbackExtractor {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>> {
        // Extract [[wiki_link]] patterns from text
        let mut entities = vec![];
        let re = regex::Regex::new(r"\[\[([^\]]+)\]\]")?;
        
        for cap in re.captures_iter(text) {
            if let Some(name) = cap.get(1) {
                let name_str = name.as_str();
                entities.push(ExtractedEntity {
                    name: name_str.to_string(),
                    entity_type: EntityType::Unknown,
                    summary: format!("Mentioned in episode"),
                    confidence: 0.7, // Lower confidence for fallback
                });
            }
        }

        Ok(entities)
    }
}

/// Composite extractor: LLM first, fallback to wiki_links (all stages)
pub struct CompositeEntityExtractor {
    primary: Box<dyn EntityExtractor>,
    fallback: Box<dyn EntityExtractor>,
}

impl CompositeEntityExtractor {
    pub fn new(primary: Box<dyn EntityExtractor>, fallback: Box<dyn EntityExtractor>) -> Self {
        Self { primary, fallback }
    }

    /// Default: LLM with wiki_links fallback
    pub fn default_llm() -> Self {
        Self::new(
            Box::new(LlmEntityExtractor::new("reasoning")),
            Box::new(WikiLinkFallbackExtractor),
        )
    }
}

#[async_trait]
impl EntityExtractor for CompositeEntityExtractor {
    async fn extract(&self, text: &str) -> Result<Vec<ExtractedEntity>> {
        match self.primary.extract(text).await {
            Ok(entities) if !entities.is_empty() => {
                tracing::debug!("LLM extraction succeeded: {} entities", entities.len());
                Ok(entities)
            }
            Ok(_) => {
                tracing::warn!("LLM extraction returned empty, using fallback");
                self.fallback.extract(text).await
            }
            Err(e) => {
                tracing::warn!("LLM extraction failed: {}, using fallback", e);
                self.fallback.extract(text).await
            }
        }
    }
}

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

    #[tokio::test]
    async fn test_wiki_link_extraction() {
        let extractor = WikiLinkFallbackExtractor;
        let text = "Rock uses [[Kubernetes]] and [[ArgoCD]] for GitOps";

        let entities = extractor.extract(text).await.unwrap();
        assert_eq!(entities.len(), 2);
        assert!(entities.iter().any(|e| e.name == "Kubernetes"));
        assert!(entities.iter().any(|e| e.name == "ArgoCD"));
    }

    #[tokio::test]
    async fn test_extracted_entity_to_domain() {
        let extracted = ExtractedEntity {
            name: "Test Entity".to_string(),
            entity_type: EntityType::Tool,
            summary: "A test entity".to_string(),
            confidence: 0.95,
        };

        let domain = extracted.to_domain("proj1");
        assert_eq!(domain.name, "Test Entity");
        assert_eq!(domain.entity_type, EntityType::Tool);
    }

    #[tokio::test]
    async fn test_composite_fallback() {
        let primary = Box::new(WikiLinkFallbackExtractor);
        let fallback = Box::new(WikiLinkFallbackExtractor);

        let composite = CompositeEntityExtractor::new(primary, fallback);
        let text = "[[Entity1]] and [[Entity2]]";

        let entities = composite.extract(text).await.unwrap();
        assert!(entities.len() > 0);
    }
}