//! 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, // 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; /// 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; } /// 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 { 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 { 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 { 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 { 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 { 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); } }