diff --git a/src/rag/graph.rs b/src/rag/graph.rs index 166a3e7..a63025c 100644 --- a/src/rag/graph.rs +++ b/src/rag/graph.rs @@ -732,4 +732,143 @@ mod tests { assert!(docs.contains(&DocumentId(0))); assert!(docs.contains(&DocumentId(1))); } + + #[test] + fn compact_preserves_edges_between_survivors() { + let mut kg = KnowledgeGraph::default(); + kg.merge(doc(0), extraction(vec![entity("A", "CONCEPT")], vec![])); + kg.merge( + doc(1), + extraction( + vec![entity("B", "CONCEPT"), entity("C", "CONCEPT")], + vec![rel("B", "C", "linked", 0.8)], + ), + ); + kg.remove_documents(&[doc(0)]); + let b_raw = kg.entity_index["b"]; + let c_raw = kg.entity_index["c"]; + let b_idx = NodeIndex::new(b_raw as usize); + let c_idx = NodeIndex::new(c_raw as usize); + assert_eq!( + kg.graph.edges_connecting(b_idx, c_idx).count(), + 1, + "B→C edge should survive compaction" + ); + } + + #[test] + fn expand_two_hops_reaches_transitive_neighbor() { + let mut kg = KnowledgeGraph::default(); + kg.merge( + doc(0), + extraction( + vec![ + entity("A", "CONCEPT"), + entity("B", "CONCEPT"), + entity("C", "CONCEPT"), + ], + vec![rel("A", "B", "uses", 1.0), rel("B", "C", "uses", 0.5)], + ), + ); + let a_raw = kg.entity_index["a"]; + let c_raw = kg.entity_index["c"]; + + let one_hop = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1); + assert!( + !one_hop.contains_key(&c_raw), + "C should not be reachable at 1 hop" + ); + + let two_hop = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 2); + assert!( + two_hop.contains_key(&c_raw), + "C should be reachable at 2 hops" + ); + let c_score = two_hop[&c_raw]; + assert!( + (c_score - 0.5).abs() < 1e-6, + "C score should be 1.0 * 1.0 * 0.5 = 0.5, got {c_score}" + ); + } + + #[test] + fn merge_clamps_edge_weight_above_one() { + let mut kg = KnowledgeGraph::default(); + kg.merge( + doc(0), + extraction( + vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")], + vec![rel("A", "B", "uses", 1.5)], + ), + ); + let a_raw = kg.entity_index["a"]; + let b_raw = kg.entity_index["b"]; + let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1); + let b_score = result[&b_raw]; + assert!( + (b_score - 1.0).abs() < 1e-6, + "weight 1.5 clamped to 1.0: b_score should be 1.0, got {b_score}" + ); + } + + #[test] + fn merge_clamps_edge_weight_below_zero() { + let mut kg = KnowledgeGraph::default(); + kg.merge( + doc(0), + extraction( + vec![entity("A", "CONCEPT"), entity("B", "CONCEPT")], + vec![rel("A", "B", "uses", -0.5)], + ), + ); + let a_raw = kg.entity_index["a"]; + let b_raw = kg.entity_index["b"]; + let result = kg.expand_neighbors_scored(&[(a_raw, 1.0)], 1); + let b_score = result.get(&b_raw).copied().unwrap_or(0.0); + assert!( + b_score.abs() < 1e-6, + "weight -0.5 clamped to 0.0: b_score should be 0.0, got {b_score}" + ); + } + + #[test] + fn merge_fills_missing_description_from_later_chunk() { + let mut kg = KnowledgeGraph::default(); + kg.merge( + doc(0), + extraction(vec![entity("Python", "TECHNOLOGY")], vec![]), + ); + kg.merge( + doc(1), + ExtractionResult { + entities: vec![ExtractedEntity { + name: "python".to_string(), + entity_type: "TECHNOLOGY".to_string(), + description: Some("A general-purpose language".to_string()), + }], + relationships: vec![], + }, + ); + let raw = kg.entity_index["python"]; + let desc = &kg.graph[NodeIndex::new(raw as usize)].description; + assert_eq!( + desc.as_deref(), + Some("A general-purpose language"), + "description should be backfilled from later chunk" + ); + } + + #[test] + fn remove_all_documents_empties_graph() { + let mut kg = KnowledgeGraph::default(); + kg.merge(doc(0), extraction(vec![entity("A", "CONCEPT")], vec![])); + kg.merge(doc(1), extraction(vec![entity("B", "CONCEPT")], vec![])); + kg.remove_documents(&[doc(0), doc(1)]); + assert_eq!(kg.graph.node_count(), 0, "all nodes should be removed"); + assert_eq!(kg.entity_index.len(), 0, "entity index should be empty"); + assert!( + kg.document_entities.is_empty(), + "document_entities should be empty" + ); + } } diff --git a/src/rag/mod.rs b/src/rag/mod.rs index d2275ff..b06928d 100644 --- a/src/rag/mod.rs +++ b/src/rag/mod.rs @@ -1781,4 +1781,91 @@ mod tests { assert_eq!(file_idx, 0); assert_eq!(doc_idx, 0); } + + #[test] + fn rag_data_del_removes_graph_entities() { + use super::graph::{ExtractedEntity, ExtractionResult}; + let mut data = RagData::new( + "m".into(), + 100, + 10, + None, + 5, + None, + GraphRagConfig::default(), + ); + let file = RagFile { + hash: "abc".into(), + path: "test.txt".into(), + documents: vec![RagDocument::new("Python is great")], + }; + data.files.insert(0, file); + let doc_id = DocumentId::new(0, 0); + data.knowledge_graph.merge( + doc_id, + ExtractionResult { + entities: vec![ExtractedEntity { + name: "Python".to_string(), + entity_type: "TECHNOLOGY".to_string(), + description: None, + }], + relationships: vec![], + }, + ); + assert!( + data.knowledge_graph.entity_index.contains_key("python"), + "entity should exist before del" + ); + data.del(vec![0]); + assert!( + !data.knowledge_graph.entity_index.contains_key("python"), + "entity should be removed after del" + ); + } + + #[test] + fn reciprocal_rank_fusion_empty_lists() { + let result = super::reciprocal_rank_fusion(vec![], vec![], 5); + assert!(result.is_empty(), "empty input should produce empty output"); + } + + #[test] + fn reciprocal_rank_fusion_deduplicates_across_signals() { + let doc_a = DocumentId::new(0, 0); + let doc_b = DocumentId::new(0, 1); + let result = super::reciprocal_rank_fusion( + vec![vec![doc_a, doc_b], vec![doc_a, doc_b]], + vec![1.0, 1.0], + 5, + ); + let unique: std::collections::HashSet<_> = result.iter().collect(); + assert_eq!( + unique.len(), + result.len(), + "each document should appear at most once" + ); + assert_eq!(result.len(), 2); + } + + #[test] + fn reciprocal_rank_fusion_respects_top_k() { + let docs: Vec = (0..10).map(|i| DocumentId::new(0, i)).collect(); + let result = super::reciprocal_rank_fusion(vec![docs], vec![1.0], 3); + assert_eq!(result.len(), 3, "result should be capped at top_k=3"); + } + + #[test] + fn reciprocal_rank_fusion_weights_affect_ranking() { + let doc_a = DocumentId::new(0, 0); + let doc_b = DocumentId::new(0, 1); + let result = super::reciprocal_rank_fusion( + vec![vec![doc_a, doc_b], vec![doc_b, doc_a]], + vec![10.0, 1.0], + 2, + ); + assert_eq!( + result[0], doc_a, + "higher-weight signal's top doc should rank first" + ); + } }