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

//! Graph Retrieval Memory (GRM) Context Retriever
//!
//! Query existing graph to validate & enrich entity/fact extraction.
//! Confirms "memorability" before committing to storage.
//!
//! CRAP: 18 (Database queries + scoring logic)
//! SOLID: Single responsibility (retrieve context), delegates scoring
//! DRY: Reuses entity/edge types from mem_core

use anyhow::Result;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tracing::{debug, info};
use mem_core::entity::Entity;
use mem_core::edge::Edge;

/// Memorability decision for entity or fact
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub enum MemorabilityDecision {
    /// Entity/fact already exists, merge with it
    Merge,
    /// New entity/fact, worth storing
    Keep,
    /// Noise or irrelevant, skip
    Drop,
    /// Low confidence, queue for human review
    ReviewQueue,
}

/// Context about an entity from the graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EntityContext {
    pub entity_name: String,
    pub matched_entity_id: Option<String>,  // If found in graph
    pub related_entities: Vec<(String, String)>,  // (id, name)
    pub related_edges_count: usize,
    pub summary: String,                    // "Rock: DevOps expert with K8s/ArgoCD expertise"
    pub memorability_score: f32,            // 0-1
    pub decision: MemorabilityDecision,
    pub reasoning: String,
}

/// Context about a fact from the graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FactContext {
    pub similar_facts_found: usize,
    pub contradictory_facts_found: usize,
    pub related_entities_coverage: f32,    // Fraction of entities that exist
    pub memorability_score: f32,            // 0-1
    pub decision: MemorabilityDecision,
    pub reasoning: String,
}

/// Graph Retrieval Memory configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GrmConfig {
    pub enabled: bool,                     // Enable/disable GRM gate
    pub entity_similarity_threshold: f32,  // Default: 0.7
    pub max_entity_context_size: usize,    // Default: 10
    pub max_related_edges: usize,          // Default: 20
    pub entity_memorability_threshold: f32, // Default: 0.75 (>= continue, < review)
    pub fact_memorability_threshold: f32,  // Default: 0.75
    pub fact_drop_threshold: f32,          // Default: 0.50 (< drop)
}

impl Default for GrmConfig {
    fn default() -> Self {
        Self {
            enabled: false,  // Disabled by default (Phase 2.5 TBD)
            entity_similarity_threshold: 0.7,
            max_entity_context_size: 10,
            max_related_edges: 20,
            entity_memorability_threshold: 0.75,
            fact_memorability_threshold: 0.75,
            fact_drop_threshold: 0.50,
        }
    }
}

/// Graph Context Retriever trait
#[async_trait]
pub trait GraphContextRetriever: Send + Sync {
    /// Get context for an entity from the graph
    async fn get_entity_context(
        &self,
        entity_name: &str,
    ) -> Result<EntityContext>;

    /// Get context for a fact from the graph
    async fn get_fact_context(
        &self,
        source_entity_id: &str,
        target_entity_id: &str,
        relation_type: &str,
        fact_text: &str,
    ) -> Result<FactContext>;
}

/// Mock GRM Retriever for testing (always returns KEEP)
#[derive(Debug, Clone)]
pub struct MockGrmRetriever;

#[async_trait]
impl GraphContextRetriever for MockGrmRetriever {
    async fn get_entity_context(&self, entity_name: &str) -> Result<EntityContext> {
        debug!("MockGrmRetriever: get_entity_context({})", entity_name);
        
        Ok(EntityContext {
            entity_name: entity_name.to_string(),
            matched_entity_id: None,
            related_entities: vec![],
            related_edges_count: 0,
            summary: format!("Mock entity: {}", entity_name),
            memorability_score: 0.95,
            decision: MemorabilityDecision::Keep,
            reasoning: "Mock: no graph available".to_string(),
        })
    }

    async fn get_fact_context(
        &self,
        _source: &str,
        _target: &str,
        _relation: &str,
        fact_text: &str,
    ) -> Result<FactContext> {
        debug!("MockGrmRetriever: get_fact_context({})", fact_text);
        
        Ok(FactContext {
            similar_facts_found: 0,
            contradictory_facts_found: 0,
            related_entities_coverage: 1.0,
            memorability_score: 0.95,
            decision: MemorabilityDecision::Keep,
            reasoning: "Mock: no graph available".to_string(),
        })
    }
}

/// Postgres-backed GRM Retriever (to be implemented in Phase 2.5)
#[derive(Debug, Clone)]
pub struct PostgresGrmRetriever {
    config: GrmConfig,
    // pool: PgPool,  // TODO (Phase 2.5): Add database connection
}

impl PostgresGrmRetriever {
    pub fn new(config: GrmConfig) -> Self {
        Self { config }
    }

    /// Score entity memorability (0-1)
    /// Higher = more memorable (more related facts, exact match, etc.)
    fn score_entity_memorability(
        &self,
        matched: bool,
        related_edges_count: usize,
    ) -> f32 {
        if matched {
            // Existing entity: very memorable
            // Bonus: more related edges = more established
            let edge_bonus = (related_edges_count as f32 / 10.0).min(0.2);
            0.8 + edge_bonus  // 0.8-1.0
        } else {
            // New entity: less memorable unless connecting to existing graph
            if related_edges_count > 0 {
                0.6 + (related_edges_count as f32 / 20.0).min(0.2)  // 0.6-0.8
            } else {
                0.5  // Isolated entity
            }
        }
    }

    /// Score fact memorability (0-1)
    /// Higher = more memorable (novel fact, no contradictions, etc.)
    fn score_fact_memorability(
        &self,
        similar_facts: usize,
        contradictions: usize,
        entity_coverage: f32,
        extraction_confidence: Option<f32>,
    ) -> f32 {
        let mut score = 0.5;

        // Novel fact: +0.3 (no similar facts)
        score += if similar_facts == 0 { 0.3 } else { -0.1 * (similar_facts as f32).min(3.0) };

        // No contradictions: +0.2
        score += if contradictions == 0 { 0.2 } else { -0.15 * (contradictions as f32) };

        // Entity coverage: +0.2 (both entities exist in graph)
        score += entity_coverage * 0.2;

        // Extraction confidence: +0.1 (if provided)
        if let Some(conf) = extraction_confidence {
            score += conf * 0.1;
        }

        score.clamp(0.0, 1.0)
    }
}

#[async_trait]
impl GraphContextRetriever for PostgresGrmRetriever {
    async fn get_entity_context(&self, entity_name: &str) -> Result<EntityContext> {
        debug!("PostgresGrmRetriever: get_entity_context({})", entity_name);

        // TODO (Phase 2.5): Implement actual database query
        // SELECT id, name, summary FROM memory_entity
        // WHERE name_embedding <-> query_embedding < (1 - threshold)
        // LIMIT max_entity_context_size

        // For now, return mock
        let matched = entity_name.to_lowercase().contains("rock");
        let related_edges_count = if matched { 23 } else { 0 };
        let memorability_score = self.score_entity_memorability(matched, related_edges_count);

        let decision = if memorability_score >= self.config.entity_memorability_threshold {
            if matched {
                MemorabilityDecision::Merge
            } else {
                MemorabilityDecision::Keep
            }
        } else {
            MemorabilityDecision::ReviewQueue
        };

        Ok(EntityContext {
            entity_name: entity_name.to_string(),
            matched_entity_id: if matched {
                Some("entity-rock-001".to_string())
            } else {
                None
            },
            related_entities: if matched {
                vec![
                    ("entity-k8s-001".to_string(), "Kubernetes".to_string()),
                    ("entity-argo-001".to_string(), "ArgoCD".to_string()),
                ]
            } else {
                vec![]
            },
            related_edges_count,
            summary: if matched {
                "Rock: DevOps engineer, expertise in Kubernetes, ArgoCD, GitOps".to_string()
            } else {
                format!("New entity: {}", entity_name)
            },
            memorability_score,
            decision,
            reasoning: format!(
                "matched={}, related_edges={}, score={}",
                matched, related_edges_count, memorability_score
            ),
        })
    }

    async fn get_fact_context(
        &self,
        _source: &str,
        _target: &str,
        _relation: &str,
        fact_text: &str,
    ) -> Result<FactContext> {
        debug!("PostgresGrmRetriever: get_fact_context({})", fact_text);

        // TODO (Phase 2.5): Implement actual database query
        // SELECT COUNT(*) FROM memory_edge
        // WHERE source_id = ? AND target_id = ?
        //   AND fact_embedding <-> query_embedding < (1 - similarity_threshold)
        // AND (t_invalid IS NULL OR t_invalid > NOW())

        let is_duplicate = fact_text.to_lowercase().contains("kubernetes");
        let similar_facts = if is_duplicate { 3 } else { 0 };
        let entity_coverage = 0.9;
        let memorability_score =
            self.score_fact_memorability(similar_facts, 0, entity_coverage, Some(0.9));

        let decision = if memorability_score < self.config.fact_drop_threshold {
            MemorabilityDecision::Drop
        } else if memorability_score >= self.config.fact_memorability_threshold {
            if is_duplicate {
                MemorabilityDecision::Merge
            } else {
                MemorabilityDecision::Keep
            }
        } else {
            MemorabilityDecision::ReviewQueue
        };

        Ok(FactContext {
            similar_facts_found: similar_facts,
            contradictory_facts_found: 0,
            related_entities_coverage: entity_coverage,
            memorability_score,
            decision,
            reasoning: format!(
                "similar={}, contradictions=0, entity_coverage={}, score={}",
                similar_facts, entity_coverage, memorability_score
            ),
        })
    }
}

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

    #[test]
    fn test_grm_config_defaults() {
        let config = GrmConfig::default();
        assert!(!config.enabled);
        assert_eq!(config.entity_similarity_threshold, 0.7);
        assert_eq!(config.max_entity_context_size, 10);
    }

    #[tokio::test]
    async fn test_mock_grm_retriever() {
        let retriever = MockGrmRetriever;
        let context = retriever.get_entity_context("Rock").await.unwrap();
        assert_eq!(context.entity_name, "Rock");
        assert_eq!(context.decision, MemorabilityDecision::Keep);
    }

    #[tokio::test]
    async fn test_postgres_grm_retriever_known_entity() {
        let config = GrmConfig::default();
        let retriever = PostgresGrmRetriever::new(config);
        
        let context = retriever.get_entity_context("Rock").await.unwrap();
        assert_eq!(context.entity_name, "Rock");
        assert!(context.matched_entity_id.is_some());
        assert_eq!(context.related_edges_count, 23);
        assert!(context.memorability_score > 0.8);
    }

    #[tokio::test]
    async fn test_postgres_grm_retriever_new_entity() {
        let config = GrmConfig::default();
        let retriever = PostgresGrmRetriever::new(config);
        
        let context = retriever.get_entity_context("UnknownPerson").await.unwrap();
        assert_eq!(context.entity_name, "UnknownPerson");
        assert!(context.matched_entity_id.is_none());
        assert_eq!(context.related_edges_count, 0);
    }

    #[tokio::test]
    async fn test_fact_context_duplicate() {
        let config = GrmConfig::default();
        let retriever = PostgresGrmRetriever::new(config);
        
        let context = retriever
            .get_fact_context("entity-1", "entity-2", "USES", "Rock uses Kubernetes")
            .await
            .unwrap();
        
        assert!(context.similar_facts_found > 0);
        assert_eq!(context.contradictory_facts_found, 0);
    }

    #[test]
    fn test_entity_memorability_scoring() {
        let config = GrmConfig::default();
        let retriever = PostgresGrmRetriever::new(config);
        
        // Existing entity with many related edges
        let score_high = retriever.score_entity_memorability(true, 20);
        assert!(score_high > 0.9);
        
        // New entity with no related edges
        let score_low = retriever.score_entity_memorability(false, 0);
        assert_eq!(score_low, 0.5);
        
        // New entity with some related edges
        let score_mid = retriever.score_entity_memorability(false, 5);
        assert!(score_mid > 0.5 && score_mid <= 0.8);
    }

    #[test]
    fn test_fact_memorability_scoring() {
        let config = GrmConfig::default();
        let retriever = PostgresGrmRetriever::new(config);
        
        // Novel fact with high entity coverage
        let score_high = retriever.score_fact_memorability(0, 0, 1.0, Some(0.95));
        assert!(score_high > 0.8);
        
        // Duplicate fact
        let score_low = retriever.score_fact_memorability(3, 1, 0.5, Some(0.6));
        assert!(score_low < 0.7);
    }
}