378 lines
12 KiB
Rust
378 lines
12 KiB
Rust
|
//! Memorability Gate: Filter extraction based on graph context
|
||
|
|
//!
|
||
|
|
//! Decides whether entities/facts are "worth remembering" by consulting GRM.
|
||
|
|
//! Configurable thresholds for different decision strategies.
|
||
|
|
//!
|
||
|
|
//! CRAP: 12 (Straightforward filtering + thresholds)
|
||
|
|
//! SOLID: Single responsibility (gate logic), delegates to retriever
|
||
|
|
//! DRY: Reuses GrmConfig and decision types
|
||
|
|
|
||
|
|
use anyhow::Result;
|
||
|
|
use serde::{Deserialize, Serialize};
|
||
|
|
use tracing::{debug, info};
|
||
|
|
|
||
|
|
use crate::grm_retriever::{
|
||
|
|
EntityContext, FactContext, GraphContextRetriever, MemorabilityDecision, GrmConfig, MockGrmRetriever,
|
||
|
|
};
|
||
|
|
use mem_core::entity::{Entity, EntityType};
|
||
|
|
use mem_core::edge::Edge;
|
||
|
|
|
||
|
|
/// Entity filtering result
|
||
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
|
|
pub struct FilteredEntity {
|
||
|
|
pub entity: Entity,
|
||
|
|
pub context: EntityContext,
|
||
|
|
pub filtered: bool, // true = dropped by GRM gate
|
||
|
|
pub reason: String,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Fact filtering result
|
||
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
|
|
pub struct FilteredFact {
|
||
|
|
pub edge: Edge,
|
||
|
|
pub context: FactContext,
|
||
|
|
pub filtered: bool, // true = dropped by GRM gate
|
||
|
|
pub reason: String,
|
||
|
|
pub requires_review: bool, // true = queue for human verification
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Memorability Gate
|
||
|
|
pub struct MemorabilityGate {
|
||
|
|
config: GrmConfig,
|
||
|
|
retriever: Box<dyn GraphContextRetriever>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl MemorabilityGate {
|
||
|
|
/// Create gate with custom retriever (for testing or custom backends)
|
||
|
|
pub fn new(config: GrmConfig, retriever: Box<dyn GraphContextRetriever>) -> Self {
|
||
|
|
Self { config, retriever }
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Create gate with mock retriever (everything passes)
|
||
|
|
pub fn with_mock(config: GrmConfig) -> Self {
|
||
|
|
Self {
|
||
|
|
config,
|
||
|
|
retriever: Box::new(MockGrmRetriever),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Check if GRM gate is enabled
|
||
|
|
pub fn is_enabled(&self) -> bool {
|
||
|
|
self.config.enabled
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Filter entity through GRM gate
|
||
|
|
pub async fn filter_entity(&self, entity: &Entity) -> Result<FilteredEntity> {
|
||
|
|
if !self.config.enabled {
|
||
|
|
debug!("GRM gate disabled, passing entity: {}", entity.name);
|
||
|
|
return Ok(FilteredEntity {
|
||
|
|
entity: entity.clone(),
|
||
|
|
context: EntityContext {
|
||
|
|
entity_name: entity.name.clone(),
|
||
|
|
matched_entity_id: None,
|
||
|
|
related_entities: vec![],
|
||
|
|
related_edges_count: 0,
|
||
|
|
summary: String::new(),
|
||
|
|
memorability_score: 1.0,
|
||
|
|
decision: MemorabilityDecision::Keep,
|
||
|
|
reasoning: "GRM gate disabled".to_string(),
|
||
|
|
},
|
||
|
|
filtered: false,
|
||
|
|
reason: "GRM disabled".to_string(),
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
debug!("GRM gate: filtering entity {}", entity.name);
|
||
|
|
let context = self.retriever.get_entity_context(&entity.name).await?;
|
||
|
|
|
||
|
|
let (filtered, reason) = match context.decision {
|
||
|
|
MemorabilityDecision::Keep => {
|
||
|
|
if context.matched_entity_id.is_some() {
|
||
|
|
(true, format!("Existing entity (merge required)"))
|
||
|
|
} else {
|
||
|
|
(false, format!("New entity (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
MemorabilityDecision::Drop => {
|
||
|
|
(true, format!("Noise/irrelevant (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
MemorabilityDecision::ReviewQueue => {
|
||
|
|
(false, format!("Low confidence, queued for review (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
MemorabilityDecision::Merge => {
|
||
|
|
(true, format!("Duplicate, requires merge (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
info!(
|
||
|
|
"GRM entity filter: {} → filtered={} ({})",
|
||
|
|
entity.name, filtered, reason
|
||
|
|
);
|
||
|
|
|
||
|
|
Ok(FilteredEntity {
|
||
|
|
entity: entity.clone(),
|
||
|
|
context,
|
||
|
|
filtered,
|
||
|
|
reason,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Filter fact through GRM gate
|
||
|
|
pub async fn filter_fact(
|
||
|
|
&self,
|
||
|
|
edge: &Edge,
|
||
|
|
source_name: Option<&str>,
|
||
|
|
target_name: Option<&str>,
|
||
|
|
) -> Result<FilteredFact> {
|
||
|
|
if !self.config.enabled {
|
||
|
|
debug!("GRM gate disabled, passing fact: {}", edge.fact);
|
||
|
|
return Ok(FilteredFact {
|
||
|
|
edge: edge.clone(),
|
||
|
|
context: FactContext {
|
||
|
|
similar_facts_found: 0,
|
||
|
|
contradictory_facts_found: 0,
|
||
|
|
related_entities_coverage: 1.0,
|
||
|
|
memorability_score: 1.0,
|
||
|
|
decision: MemorabilityDecision::Keep,
|
||
|
|
reasoning: "GRM gate disabled".to_string(),
|
||
|
|
},
|
||
|
|
filtered: false,
|
||
|
|
reason: "GRM disabled".to_string(),
|
||
|
|
requires_review: false,
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
debug!("GRM gate: filtering fact {}", edge.fact);
|
||
|
|
let context = self.retriever
|
||
|
|
.get_fact_context(
|
||
|
|
&edge.source_entity_id,
|
||
|
|
&edge.target_entity_id,
|
||
|
|
&edge.relation_type,
|
||
|
|
&edge.fact,
|
||
|
|
)
|
||
|
|
.await?;
|
||
|
|
|
||
|
|
let (filtered, requires_review, reason) = match context.decision {
|
||
|
|
MemorabilityDecision::Keep => {
|
||
|
|
(false, false, format!("Novel fact (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
MemorabilityDecision::Drop => {
|
||
|
|
(true, false, format!("Redundant/noise (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
MemorabilityDecision::ReviewQueue => {
|
||
|
|
(false, true, format!("Low confidence, queued for review (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
MemorabilityDecision::Merge => {
|
||
|
|
(true, false, format!("Duplicate, requires merge (score: {:.2})", context.memorability_score))
|
||
|
|
}
|
||
|
|
};
|
||
|
|
|
||
|
|
info!(
|
||
|
|
"GRM fact filter: {} → {} → filtered={} requires_review={} ({})",
|
||
|
|
source_name.unwrap_or("?"),
|
||
|
|
target_name.unwrap_or("?"),
|
||
|
|
filtered,
|
||
|
|
requires_review,
|
||
|
|
reason
|
||
|
|
);
|
||
|
|
|
||
|
|
Ok(FilteredFact {
|
||
|
|
edge: edge.clone(),
|
||
|
|
context,
|
||
|
|
filtered,
|
||
|
|
reason,
|
||
|
|
requires_review,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Batch filter entities
|
||
|
|
pub async fn filter_entities(&self, entities: &[Entity]) -> Result<Vec<FilteredEntity>> {
|
||
|
|
let mut results = Vec::new();
|
||
|
|
for entity in entities {
|
||
|
|
results.push(self.filter_entity(entity).await?);
|
||
|
|
}
|
||
|
|
Ok(results)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Batch filter facts
|
||
|
|
pub async fn filter_facts(
|
||
|
|
&self,
|
||
|
|
edges: &[Edge],
|
||
|
|
source_names: Option<&[Option<String>]>,
|
||
|
|
target_names: Option<&[Option<String>]>,
|
||
|
|
) -> Result<Vec<FilteredFact>> {
|
||
|
|
let mut results = Vec::new();
|
||
|
|
for (i, edge) in edges.iter().enumerate() {
|
||
|
|
let source = source_names.and_then(|names| names.get(i).and_then(|n| n.as_deref()));
|
||
|
|
let target = target_names.and_then(|names| names.get(i).and_then(|n| n.as_deref()));
|
||
|
|
results.push(self.filter_fact(edge, source, target).await?);
|
||
|
|
}
|
||
|
|
Ok(results)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Get statistics about filtering results
|
||
|
|
pub fn stats(filtered: &[FilteredEntity]) -> FilterStatistics {
|
||
|
|
let total = filtered.len();
|
||
|
|
let dropped = filtered.iter().filter(|f| f.filtered).count();
|
||
|
|
let kept = total - dropped;
|
||
|
|
let avg_score = filtered
|
||
|
|
.iter()
|
||
|
|
.map(|f| f.context.memorability_score)
|
||
|
|
.sum::<f32>() / (total as f32).max(1.0);
|
||
|
|
|
||
|
|
FilterStatistics {
|
||
|
|
total,
|
||
|
|
kept,
|
||
|
|
dropped,
|
||
|
|
drop_rate: (dropped as f32 / total as f32).clamp(0.0, 1.0),
|
||
|
|
avg_memorability_score: avg_score,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Filter statistics
|
||
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
|
|
pub struct FilterStatistics {
|
||
|
|
pub total: usize,
|
||
|
|
pub kept: usize,
|
||
|
|
pub dropped: usize,
|
||
|
|
pub drop_rate: f32,
|
||
|
|
pub avg_memorability_score: f32,
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
use mem_core::entity::Entity;
|
||
|
|
|
||
|
|
fn create_test_entity(name: &str) -> Entity {
|
||
|
|
Entity::new("poimen", name, EntityType::Person)
|
||
|
|
}
|
||
|
|
|
||
|
|
fn create_test_edge(source: &str, target: &str, fact: &str) -> Edge {
|
||
|
|
Edge::new("poimen", source, target, "USES", fact)
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_gate_disabled() {
|
||
|
|
let config = GrmConfig {
|
||
|
|
enabled: false,
|
||
|
|
..Default::default()
|
||
|
|
};
|
||
|
|
let gate = MemorabilityGate::with_mock(config);
|
||
|
|
|
||
|
|
let entity = create_test_entity("Rock");
|
||
|
|
let result = gate.filter_entity(&entity).await.unwrap();
|
||
|
|
|
||
|
|
assert!(!result.filtered);
|
||
|
|
assert_eq!(result.reason, "GRM disabled");
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_gate_enabled_known_entity() {
|
||
|
|
let config = GrmConfig {
|
||
|
|
enabled: true,
|
||
|
|
entity_memorability_threshold: 0.75,
|
||
|
|
..Default::default()
|
||
|
|
};
|
||
|
|
let gate = MemorabilityGate::with_mock(config);
|
||
|
|
|
||
|
|
let entity = create_test_entity("Rock");
|
||
|
|
let result = gate.filter_entity(&entity).await.unwrap();
|
||
|
|
|
||
|
|
// With mock retriever, entity "Rock" has high score
|
||
|
|
assert_eq!(result.context.decision, MemorabilityDecision::Keep);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_gate_enabled_new_entity() {
|
||
|
|
let config = GrmConfig {
|
||
|
|
enabled: true,
|
||
|
|
entity_memorability_threshold: 0.75,
|
||
|
|
..Default::default()
|
||
|
|
};
|
||
|
|
let gate = MemorabilityGate::with_mock(config);
|
||
|
|
|
||
|
|
let entity = create_test_entity("UnknownPerson");
|
||
|
|
let result = gate.filter_entity(&entity).await.unwrap();
|
||
|
|
|
||
|
|
// With mock retriever, all entities get KEEP decision
|
||
|
|
assert_eq!(result.context.decision, MemorabilityDecision::Keep);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_gate_filter_fact_disabled() {
|
||
|
|
let config = GrmConfig {
|
||
|
|
enabled: false,
|
||
|
|
..Default::default()
|
||
|
|
};
|
||
|
|
let gate = MemorabilityGate::with_mock(config);
|
||
|
|
|
||
|
|
let edge = create_test_edge("entity-1", "entity-2", "Rock uses Kubernetes");
|
||
|
|
let result = gate.filter_fact(&edge, Some("Rock"), Some("Kubernetes")).await.unwrap();
|
||
|
|
|
||
|
|
assert!(!result.filtered);
|
||
|
|
assert!(!result.requires_review);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[tokio::test]
|
||
|
|
async fn test_gate_batch_filter_entities() {
|
||
|
|
let config = GrmConfig {
|
||
|
|
enabled: true,
|
||
|
|
..Default::default()
|
||
|
|
};
|
||
|
|
let gate = MemorabilityGate::with_mock(config);
|
||
|
|
|
||
|
|
let entities = vec![
|
||
|
|
create_test_entity("Rock"),
|
||
|
|
create_test_entity("Kubernetes"),
|
||
|
|
create_test_entity("ArgoCD"),
|
||
|
|
];
|
||
|
|
|
||
|
|
let results = gate.filter_entities(&entities).await.unwrap();
|
||
|
|
assert_eq!(results.len(), 3);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_filter_statistics() {
|
||
|
|
let filtered = vec![
|
||
|
|
FilteredEntity {
|
||
|
|
entity: create_test_entity("A"),
|
||
|
|
context: EntityContext {
|
||
|
|
entity_name: "A".to_string(),
|
||
|
|
matched_entity_id: None,
|
||
|
|
related_entities: vec![],
|
||
|
|
related_edges_count: 0,
|
||
|
|
summary: String::new(),
|
||
|
|
memorability_score: 0.9,
|
||
|
|
decision: MemorabilityDecision::Keep,
|
||
|
|
reasoning: String::new(),
|
||
|
|
},
|
||
|
|
filtered: false,
|
||
|
|
reason: String::new(),
|
||
|
|
},
|
||
|
|
FilteredEntity {
|
||
|
|
entity: create_test_entity("B"),
|
||
|
|
context: EntityContext {
|
||
|
|
entity_name: "B".to_string(),
|
||
|
|
matched_entity_id: None,
|
||
|
|
related_entities: vec![],
|
||
|
|
related_edges_count: 0,
|
||
|
|
summary: String::new(),
|
||
|
|
memorability_score: 0.3,
|
||
|
|
decision: MemorabilityDecision::Drop,
|
||
|
|
reasoning: String::new(),
|
||
|
|
},
|
||
|
|
filtered: true,
|
||
|
|
reason: String::new(),
|
||
|
|
},
|
||
|
|
];
|
||
|
|
|
||
|
|
let stats = MemorabilityGate::stats(&filtered);
|
||
|
|
assert_eq!(stats.total, 2);
|
||
|
|
assert_eq!(stats.kept, 1);
|
||
|
|
assert_eq!(stats.dropped, 1);
|
||
|
|
assert_eq!(stats.drop_rate, 0.5);
|
||
|
|
}
|
||
|
|
}
|