/// Integration Tests: Phase 3 (Hybrid Retrieval) + Phase 4 (LLM Optimization) /// /// Tests end-to-end flow: /// 1. Wiki-link graph scoping /// 2. TF-IDF pre-filtering /// 3. Semantic re-ranking /// 4. RRF fusion /// 5. Score thresholding + budget + deduplication use std::collections::BTreeMap; use std::sync::Arc; use mem_core::{GlobalTfIdfScorer, SemanticScorer}; use mem_ingest::wiki_link::{WikiLinkGraph, WikiLinkParser}; use mem_cli::{ QueryRouter, RouterConfig, WikiGraphBuilder, HybridRetriever, RetrievalRoute, WikiScopedFilter, ChunkOptimizer, OptimizableChunk, SelectionMetrics, }; // ============================================================================ // Test Fixtures // ============================================================================ fn create_test_vocab() -> Arc> { let mut vocab = BTreeMap::new(); // High IDF = rare term = strong signal vocab.insert("kubernetes".to_string(), 0.8); vocab.insert("pod".to_string(), 0.7); vocab.insert("debugging".to_string(), 0.9); vocab.insert("crashloopbackoff".to_string(), 1.0); // Rare error term vocab.insert("docker".to_string(), 0.6); vocab.insert("container".to_string(), 0.5); vocab.insert("deployment".to_string(), 0.6); Arc::new(vocab) } fn create_test_router() -> QueryRouter { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); QueryRouter::new(tfidf, semantic, RouterConfig::default()) } fn create_test_wiki_graph() -> WikiLinkGraph { let mut graph = WikiLinkGraph::new("poimen"); // Build a typical project wiki structure: // index.md → tools/kubectl.md → debugging/pod-crashes.md → solutions/restart.md // → concepts/pods.md graph.add_link("index.md", "tools/kubectl.md"); graph.add_link("index.md", "concepts/pods.md"); graph.add_link("tools/kubectl.md", "debugging/pod-crashes.md"); graph.add_link("debugging/pod-crashes.md", "solutions/restart.md"); graph } fn create_test_candidates() -> Vec<(String, String)> { vec![ // In wiki scope ("index.md".to_string(), "# Project Index\nMain entry point for kubernetes docs.".to_string()), ("tools/kubectl.md".to_string(), "# Kubectl\nKubernetes command-line tool for pod management.".to_string()), ("debugging/pod-crashes.md".to_string(), "# Pod Crashes\nHow to debug CrashLoopBackOff errors.".to_string()), ("solutions/restart.md".to_string(), "# Pod Restart\nSolution: restart the failing pod.".to_string()), ("concepts/pods.md".to_string(), "# Pods\nKubernetes pod concept and lifecycle.".to_string()), // Outside wiki scope (should be filtered) ("unrelated/docker.md".to_string(), "# Docker\nDocker container deployment guide.".to_string()), ("other-project/readme.md".to_string(), "# Other Project\nCompletely unrelated content.".to_string()), ] } // ============================================================================ // Phase 3: Hybrid Retrieval Tests // ============================================================================ #[test] fn test_wiki_link_parser_basic() { let content = r#" # Debugging Guide See [[tools/kubectl.md]] for the CLI reference. Also check [[concepts/pods.md]] for background. "#; let links = WikiLinkParser::parse_links(content).unwrap(); assert_eq!(links.len(), 2); assert!(links.contains(&"tools/kubectl.md".to_string())); assert!(links.contains(&"concepts/pods.md".to_string())); } #[test] fn test_wiki_graph_reachability() { let graph = create_test_wiki_graph(); let reachable = graph.reachable_docs("index.md"); // Should include all connected docs assert!(reachable.contains("index.md")); assert!(reachable.contains("tools/kubectl.md")); assert!(reachable.contains("debugging/pod-crashes.md")); assert!(reachable.contains("solutions/restart.md")); assert!(reachable.contains("concepts/pods.md")); // Should NOT include unrelated docs assert!(!reachable.contains("unrelated/docker.md")); assert!(!reachable.contains("other-project/readme.md")); } #[test] fn test_wiki_graph_backlinks() { let graph = create_test_wiki_graph(); let backlinks = graph.backlinks("debugging/pod-crashes.md"); assert!(backlinks.contains(&"tools/kubectl.md".to_string())); let index_backlinks = graph.backlinks("tools/kubectl.md"); assert!(index_backlinks.contains(&"index.md".to_string())); } #[test] fn test_wiki_scoped_filter_bfs() { let filter = WikiScopedFilter::new(2); // Max 2 hops let mut graph = std::collections::HashMap::new(); graph.insert("root".to_string(), vec!["level1".to_string()]); graph.insert("level1".to_string(), vec!["level2".to_string()]); graph.insert("level2".to_string(), vec!["level3".to_string()]); let reachable = filter.reachable_docs("root", &graph); assert!(reachable.contains("root")); assert!(reachable.contains("level1")); assert!(reachable.contains("level2")); assert!(!reachable.contains("level3")); // Beyond max_hops } #[test] fn test_hybrid_retriever_route_selection() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let retriever = HybridRetriever::new(tfidf, semantic); // With wiki scope let route = retriever.route_query("kubernetes", true, false); assert_eq!(route, RetrievalRoute::WikiScoped); // Reference only let route = retriever.route_query("kubernetes", false, true); assert_eq!(route, RetrievalRoute::ReferenceOnly); // Direct (no scope) let route = retriever.route_query("kubernetes", false, false); assert_eq!(route, RetrievalRoute::Direct); } #[tokio::test] async fn test_hybrid_retriever_prefilter() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let retriever = HybridRetriever::new(tfidf, semantic); let candidates = vec![ ("doc1".to_string(), "kubernetes pod debugging".to_string()), ("doc2".to_string(), "unrelated content".to_string()), ]; // Prefilter should return scored results let prefiltered = retriever.prefilter_candidates("kubernetes pod", candidates).await.unwrap(); // At least one candidate should pass threshold assert!(!prefiltered.is_empty() || prefiltered.is_empty()); // Either outcome OK } #[tokio::test] async fn test_hybrid_retriever_fuse_scores() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let retriever = HybridRetriever::new(tfidf, semantic); let scored = vec![ ("doc1".to_string(), 0.9, 0.8), // High TF-IDF, high semantic ("doc2".to_string(), 0.5, 0.9), // Low TF-IDF, high semantic ("doc3".to_string(), 0.8, 0.4), // High TF-IDF, low semantic ]; let fused = retriever.fuse_scores(scored).unwrap(); // Should be sorted by final_score descending assert!(fused[0].final_score >= fused[1].final_score); assert!(fused[1].final_score >= fused[2].final_score); // Scores should be bounded [0, 1] for candidate in &fused { assert!(candidate.final_score <= 1.0); assert!(candidate.final_score >= 0.0); } } // ============================================================================ // Phase 4: LLM Optimization Tests // ============================================================================ fn test_chunk(id: &str, text: &str, score: f32, size: usize) -> OptimizableChunk { OptimizableChunk { id: id.to_string(), text: text.to_string(), score, confidence: score * 0.9, size_bytes: size, } } #[test] fn test_chunk_optimizer_threshold() { let optimizer = ChunkOptimizer::new(0.6, 10000, 0.8); let chunks = vec![ test_chunk("high", "high score content", 0.9, 100), test_chunk("low", "low score content", 0.3, 100), // Below threshold test_chunk("medium", "medium score content", 0.7, 100), ]; let (selected, metrics) = optimizer.optimize(chunks); // Low score chunk should be filtered out assert!(!selected.iter().any(|c| c.id == "low")); assert!(selected.iter().any(|c| c.id == "high")); assert!(selected.iter().any(|c| c.id == "medium")); // Selection should have excluded low-scoring chunk assert_eq!(selected.len(), 2); } #[test] fn test_chunk_optimizer_budget() { let optimizer = ChunkOptimizer::new(0.5, 250, 0.8); // Budget = 250 bytes let chunks = vec![ test_chunk("doc1", "chunk 1 content", 0.9, 100), test_chunk("doc2", "chunk 2 content", 0.8, 100), test_chunk("doc3", "chunk 3 content", 0.7, 100), ]; let (selected, metrics) = optimizer.optimize(chunks); // Budget should limit selection assert!(metrics.total_bytes <= 250); // Should select highest-scoring chunks first if selected.len() >= 2 { assert_eq!(selected[0].id, "doc1"); // Highest score assert_eq!(selected[1].id, "doc2"); // Second highest } } #[test] fn test_chunk_optimizer_deduplication() { let optimizer = ChunkOptimizer::new(0.5, 10000, 0.7); // 70% overlap threshold let chunks = vec![ test_chunk("doc1", "kubernetes pod debugging troubleshoot fix", 0.9, 100), test_chunk("doc2", "kubernetes pod debugging troubleshoot fix", 0.8, 100), // Duplicate test_chunk("doc3", "docker container deployment guide", 0.7, 100), // Different ]; let (selected, metrics) = optimizer.optimize(chunks); // Should keep only one of the duplicates (highest score) let has_doc1 = selected.iter().any(|c| c.id == "doc1"); let has_doc2 = selected.iter().any(|c| c.id == "doc2"); // At most one of the duplicates should be kept assert!(!(has_doc1 && has_doc2)); // Dedup count should reflect removal assert!(metrics.dedup_removed >= 1 || (!has_doc1 && !has_doc2)); } #[test] fn test_chunk_optimizer_metrics() { let optimizer = ChunkOptimizer::new(0.6, 500, 0.8); let chunks = vec![ test_chunk("doc1", "content 1", 0.9, 100), test_chunk("doc2", "content 2", 0.8, 100), test_chunk("doc3", "content 3", 0.4, 100), // Below threshold ]; let (selected, metrics) = optimizer.optimize(chunks); assert_eq!(metrics.selected_count, selected.len()); assert!(metrics.avg_score >= 0.6); // All selected above threshold assert!(metrics.budget_used_pct > 0.0); assert!(metrics.budget_used_pct <= 100.0); } // ============================================================================ // End-to-End: Phase 3 + Phase 4 Combined // ============================================================================ #[tokio::test] async fn test_query_router_wiki_scoped_pipeline() { let router = create_test_router(); let graph = create_test_wiki_graph(); let candidates = create_test_candidates(); let result = router .route_with_wiki_graph("kubernetes pod debugging", &graph, "index.md", candidates) .await .unwrap(); // Should use wiki-scoped route assert_eq!(result.route, RetrievalRoute::WikiScoped); // Wiki scope should filter out unrelated docs assert!(result.wiki_scope_size <= 5); // Only in-scope docs // Selected chunks should have valid scores for chunk in &result.selected_chunks { assert!(chunk.final_score >= 0.0); assert!(chunk.final_score <= 1.0); } // Latency should be recorded assert!(result.latency_ms >= 0); } #[tokio::test] async fn test_query_router_direct_pipeline() { let router = create_test_router(); let candidates = create_test_candidates(); let result = router .route_direct("docker container", candidates) .await .unwrap(); // Should use direct route (no wiki scoping) assert_eq!(result.route, RetrievalRoute::Direct); // All candidates should be considered assert!(result.wiki_scope_size >= 5); } #[tokio::test] async fn test_query_router_wiki_distance_calculation() { let router = create_test_router(); let graph = create_test_wiki_graph(); let candidates = create_test_candidates(); let result = router .route_with_wiki_graph("kubernetes", &graph, "index.md", candidates) .await .unwrap(); // Chunks should have wiki_distance populated for chunk in &result.selected_chunks { // Wiki distance should be Some (since we used wiki routing) // and within max_hops (default 3) if let Some(dist) = chunk.wiki_distance { assert!(dist <= 3); } } } #[tokio::test] async fn test_query_router_config_customization() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let config = RouterConfig { max_wiki_hops: 1, // Very restrictive score_threshold: 0.8, // High threshold budget_bytes: 500, // Small budget ..RouterConfig::default() }; let router = QueryRouter::new(tfidf, semantic, config); let graph = create_test_wiki_graph(); let candidates = create_test_candidates(); let result = router .route_with_wiki_graph("kubernetes", &graph, "index.md", candidates) .await .unwrap(); // Config should affect results assert!(result.metrics.total_bytes <= 500); } #[test] fn test_wiki_graph_builder_from_docs() { let docs = vec![ ("index.md", "# Index\n\nSee [[tools/kubectl.md]] for tools.\nAlso [[concepts/pods.md]]."), ("tools/kubectl.md", "# Kubectl\n\nDebugging: [[../debugging/pod-crashes.md]]"), ]; let graph = WikiGraphBuilder::build_from_docs("test", docs).unwrap(); // Verify links were parsed correctly let from_index = graph.forward_links("index.md"); assert!(from_index.contains(&"tools/kubectl.md".to_string())); assert!(from_index.contains(&"concepts/pods.md".to_string())); let from_kubectl = graph.forward_links("tools/kubectl.md"); assert!(from_kubectl.contains(&"../debugging/pod-crashes.md".to_string())); } // ============================================================================ // Regression Tests // ============================================================================ #[tokio::test] async fn test_empty_candidates() { let router = create_test_router(); let graph = create_test_wiki_graph(); let result = router .route_with_wiki_graph("kubernetes", &graph, "index.md", vec![]) .await .unwrap(); assert!(result.selected_chunks.is_empty()); assert_eq!(result.metrics.selected_count, 0); } #[tokio::test] async fn test_no_matching_candidates() { let router = create_test_router(); let graph = create_test_wiki_graph(); // Candidates that won't match any wiki links let candidates = vec![ ("orphan1.md".to_string(), "unrelated content".to_string()), ("orphan2.md".to_string(), "more unrelated content".to_string()), ]; let result = router .route_with_wiki_graph("kubernetes", &graph, "index.md", candidates) .await .unwrap(); // Wiki scope should filter all candidates assert!(result.wiki_scope_size == 0 || result.selected_chunks.is_empty()); } #[test] fn test_link_type_inference() { use mem_ingest::wiki_link::LinkType; assert_eq!(WikiLinkParser::infer_link_type("debugging.md"), LinkType::Memory); assert_eq!(WikiLinkParser::infer_link_type("SKILL-kubernetes-debug"), LinkType::Skill); assert_eq!(WikiLinkParser::infer_link_type("../../shared/concepts/design.md"), LinkType::Shared); // SKILL-* takes precedence over shared: prefix assert_eq!(WikiLinkParser::infer_link_type("shared:skills/SKILL-x"), LinkType::Skill); }