/// Full Pipeline: Complete Phase 1-6 Integration /// /// Unified orchestration of all phases: /// - Phase 1: Wiki-link graph (mem-ingest) /// - Phase 2: Scoring pipeline (mem-core) /// - Phase 3: Hybrid retrieval (QueryRouter) /// - Phase 4: LLM optimization (ChunkOptimizer) /// - Phase 5: Metadata enhancement (MetadataBooster) /// - Phase 6: Cache alignment (KvCacheAligner) /// /// This module provides: /// - `FullPipeline`: complete query orchestration /// - `PipelineConfig`: unified configuration /// - `PipelineResult`: comprehensive result with all metrics use anyhow::Result; use std::collections::HashMap; use std::sync::Arc; use mem_core::{GlobalTfIdfScorer, SemanticScorer}; use mem_ingest::wiki_link::WikiLinkGraph; use crate::query_router::{QueryRouter, RouterConfig, RoutedResult, SelectedChunk}; use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkMetadata, ChunkCategory, QueryIntent}; use crate::cache_alignment::{KvCacheAligner, CachedChunk, CacheLocalityAnalyzer, RetrievalProfiler, CacheMetrics}; /// Unified pipeline configuration #[derive(Debug, Clone)] pub struct PipelineConfig { // Phase 1: Wiki-link pub project: String, pub wiki_root_doc: String, pub max_wiki_hops: u32, // Phase 3: Hybrid retrieval pub tfidf_threshold: f32, pub prefilter_limit: usize, pub rrf_tfidf_weight: f32, pub rrf_semantic_weight: f32, // Phase 4: LLM optimization pub score_threshold: f32, pub budget_bytes: usize, pub dedup_threshold: f32, // Phase 5: Metadata pub enable_metadata_boost: bool, pub category_boost_factor: f32, // Phase 6: Cache pub cache_capacity: usize, pub context_window: usize, pub chunk_avg_tokens: usize, pub preload_top_k: usize, } impl Default for PipelineConfig { fn default() -> Self { Self { // Phase 1 project: "default".to_string(), wiki_root_doc: "index.md".to_string(), max_wiki_hops: 3, // Phase 3 tfidf_threshold: 0.3, prefilter_limit: 50, rrf_tfidf_weight: 0.4, rrf_semantic_weight: 0.6, // Phase 4 score_threshold: 0.6, budget_bytes: 8192, dedup_threshold: 0.8, // Phase 5 enable_metadata_boost: true, category_boost_factor: 1.5, // Phase 6 cache_capacity: 1000, context_window: 4096, chunk_avg_tokens: 100, preload_top_k: 5, } } } /// Fully enriched chunk with all phase metadata #[derive(Debug, Clone)] pub struct EnrichedChunk { // Core pub id: String, pub text: String, // Phase 3: Retrieval scores pub tfidf_score: f32, pub semantic_score: f32, pub rrf_score: f32, // Phase 4: Optimization pub pre_boost_score: f32, pub final_score: f32, // Phase 5: Metadata pub category: ChunkCategory, pub heading: Option, pub key_terms: Vec, pub metadata_boost: f32, pub query_intent_match: bool, // Phase 6: Cache pub wiki_distance: Option, pub cache_slot: u32, pub cache_priority: f32, } /// Pipeline execution metrics #[derive(Debug, Clone)] pub struct PipelineMetrics { // Phase counts pub wiki_scope_docs: usize, pub prefilter_candidates: usize, pub post_optimization_count: usize, // Phase 4 metrics pub rejected_by_threshold: usize, pub rejected_by_budget: usize, pub dedup_removed: usize, pub budget_used_bytes: usize, pub budget_used_pct: f32, // Phase 5 metrics pub metadata_boosts_applied: usize, pub avg_boost: f32, // Phase 6 metrics pub cache_hits: u64, pub cache_misses: u64, pub cache_hit_ratio: f32, pub preloaded_chunks: usize, // Timing pub phase_timings: Vec<(String, u64)>, pub total_latency_ms: u64, } impl PipelineMetrics { pub fn new() -> Self { Self { wiki_scope_docs: 0, prefilter_candidates: 0, post_optimization_count: 0, rejected_by_threshold: 0, rejected_by_budget: 0, dedup_removed: 0, budget_used_bytes: 0, budget_used_pct: 0.0, metadata_boosts_applied: 0, avg_boost: 0.0, cache_hits: 0, cache_misses: 0, cache_hit_ratio: 0.0, preloaded_chunks: 0, phase_timings: Vec::new(), total_latency_ms: 0, } } } /// Complete pipeline result #[derive(Debug, Clone)] pub struct PipelineResult { pub query: String, pub query_intent: QueryIntent, pub chunks: Vec, pub metrics: PipelineMetrics, } /// Full Pipeline: orchestrates all phases pub struct FullPipeline { router: QueryRouter, booster: MetadataBooster, aligner: KvCacheAligner, profiler: RetrievalProfiler, config: PipelineConfig, } impl FullPipeline { pub fn new( tfidf_scorer: Arc, semantic_scorer: Arc, config: PipelineConfig, ) -> Self { let router_config = RouterConfig { max_wiki_hops: config.max_wiki_hops, tfidf_threshold: config.tfidf_threshold, prefilter_limit: config.prefilter_limit, score_threshold: config.score_threshold, budget_bytes: config.budget_bytes, dedup_threshold: config.dedup_threshold, rrf_tfidf_weight: config.rrf_tfidf_weight, rrf_semantic_weight: config.rrf_semantic_weight, }; let router = QueryRouter::new(tfidf_scorer, semantic_scorer, router_config); let booster = MetadataBooster::new(); let aligner = KvCacheAligner::new( config.context_window, config.chunk_avg_tokens, config.cache_capacity, ); let profiler = RetrievalProfiler::new(); Self { router, booster, aligner, profiler, config, } } /// Execute full pipeline with wiki-graph pub async fn execute_with_wiki( &self, query: &str, wiki_graph: &WikiLinkGraph, candidates: Vec<(String, String)>, ) -> Result { let start = std::time::Instant::now(); let mut metrics = PipelineMetrics::new(); // Phase 5: Infer query intent let t0 = std::time::Instant::now(); let query_intent = MetadataExtractor::infer_query_intent(query); metrics.phase_timings.push(("intent_inference".to_string(), t0.elapsed().as_millis() as u64)); // Phase 1-4: Wiki-scoped hybrid retrieval + optimization let t1 = std::time::Instant::now(); let routed = self.router .route_with_wiki_graph(query, wiki_graph, &self.config.wiki_root_doc, candidates) .await?; metrics.phase_timings.push(("routing_retrieval".to_string(), t1.elapsed().as_millis() as u64)); metrics.wiki_scope_docs = routed.wiki_scope_size; metrics.prefilter_candidates = routed.prefilter_size; metrics.post_optimization_count = routed.selected_chunks.len(); metrics.dedup_removed = routed.metrics.dedup_removed; metrics.budget_used_bytes = routed.metrics.total_bytes; metrics.budget_used_pct = routed.metrics.budget_used_pct; metrics.rejected_by_threshold = routed.metrics.rejected_count; // Phase 5: Apply metadata boost let t2 = std::time::Instant::now(); let mut enriched_chunks = Vec::new(); let mut total_boost = 0.0; let mut boosts_applied = 0; for chunk in routed.selected_chunks { let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text); let mut boost = 0.0; let mut intent_match = false; if self.config.enable_metadata_boost { boost = self.booster.calculate_boost(query_intent, &metadata); if boost > 0.0 { boosts_applied += 1; total_boost += boost; intent_match = true; } } let boosted_score = self.booster.apply_boost(chunk.final_score, boost); enriched_chunks.push(EnrichedChunk { id: chunk.id.clone(), text: chunk.text.clone(), tfidf_score: chunk.tfidf_score, semantic_score: chunk.semantic_score, rrf_score: chunk.final_score, pre_boost_score: chunk.final_score, final_score: boosted_score, category: metadata.category, heading: metadata.heading, key_terms: metadata.key_terms, metadata_boost: boost, query_intent_match: intent_match, wiki_distance: chunk.wiki_distance, cache_slot: 0, cache_priority: 0.0, }); } metrics.metadata_boosts_applied = boosts_applied; metrics.avg_boost = if boosts_applied > 0 { total_boost / boosts_applied as f32 } else { 0.0 }; metrics.phase_timings.push(("metadata_boost".to_string(), t2.elapsed().as_millis() as u64)); // Re-sort by boosted score enriched_chunks.sort_by(|a, b| { b.final_score.partial_cmp(&a.final_score).unwrap_or(std::cmp::Ordering::Equal) }); // Phase 6: Cache alignment let t3 = std::time::Instant::now(); let cached: Vec = enriched_chunks .iter() .enumerate() .map(|(i, chunk)| CachedChunk { chunk_id: chunk.id.clone(), text: chunk.text.clone(), score: chunk.final_score, cache_distance: chunk.wiki_distance.unwrap_or(u32::MAX), access_count: 1, last_accessed_slot: i as u32, }) .collect(); // Assign cache slots let slots = self.aligner.assign_slots(&cached); for chunk in &mut enriched_chunks { if let Some((_, slot)) = slots.iter().find(|(id, _)| id == &chunk.id) { chunk.cache_slot = *slot; } // Cache priority: higher score + closer wiki distance = higher priority let dist_factor = 1.0 / (1.0 + chunk.wiki_distance.unwrap_or(10) as f32); chunk.cache_priority = chunk.final_score * dist_factor; } // Preload hot chunks let preload_chunks: Vec<_> = enriched_chunks .iter() .take(self.config.preload_top_k) .map(|c| (c.id.as_str(), c.text.as_str())) .collect(); self.aligner.preload_hot_chunks(preload_chunks)?; metrics.preloaded_chunks = self.config.preload_top_k.min(enriched_chunks.len()); let cache_metrics = self.aligner.get_metrics(); metrics.cache_hits = cache_metrics.hits; metrics.cache_misses = cache_metrics.misses; metrics.cache_hit_ratio = cache_metrics.hit_ratio(); metrics.phase_timings.push(("cache_alignment".to_string(), t3.elapsed().as_millis() as u64)); metrics.total_latency_ms = start.elapsed().as_millis() as u64; Ok(PipelineResult { query: query.to_string(), query_intent, chunks: enriched_chunks, metrics, }) } /// Execute pipeline without wiki-graph (direct mode) pub async fn execute_direct( &self, query: &str, candidates: Vec<(String, String)>, ) -> Result { let start = std::time::Instant::now(); let mut metrics = PipelineMetrics::new(); // Phase 5: Infer query intent let t0 = std::time::Instant::now(); let query_intent = MetadataExtractor::infer_query_intent(query); metrics.phase_timings.push(("intent_inference".to_string(), t0.elapsed().as_millis() as u64)); // Phase 3-4: Direct retrieval + optimization let t1 = std::time::Instant::now(); let routed = self.router.route_direct(query, candidates).await?; metrics.phase_timings.push(("routing_retrieval".to_string(), t1.elapsed().as_millis() as u64)); metrics.wiki_scope_docs = routed.wiki_scope_size; metrics.prefilter_candidates = routed.prefilter_size; metrics.post_optimization_count = routed.selected_chunks.len(); metrics.dedup_removed = routed.metrics.dedup_removed; metrics.budget_used_bytes = routed.metrics.total_bytes; metrics.budget_used_pct = routed.metrics.budget_used_pct; // Phase 5: Apply metadata boost let t2 = std::time::Instant::now(); let mut enriched_chunks = Vec::new(); let mut total_boost = 0.0; let mut boosts_applied = 0; for chunk in routed.selected_chunks { let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text); let mut boost = 0.0; let mut intent_match = false; if self.config.enable_metadata_boost { boost = self.booster.calculate_boost(query_intent, &metadata); if boost > 0.0 { boosts_applied += 1; total_boost += boost; intent_match = true; } } let boosted_score = self.booster.apply_boost(chunk.final_score, boost); enriched_chunks.push(EnrichedChunk { id: chunk.id.clone(), text: chunk.text.clone(), tfidf_score: chunk.tfidf_score, semantic_score: chunk.semantic_score, rrf_score: chunk.final_score, pre_boost_score: chunk.final_score, final_score: boosted_score, category: metadata.category, heading: metadata.heading, key_terms: metadata.key_terms, metadata_boost: boost, query_intent_match: intent_match, wiki_distance: None, cache_slot: 0, cache_priority: 0.0, }); } metrics.metadata_boosts_applied = boosts_applied; metrics.avg_boost = if boosts_applied > 0 { total_boost / boosts_applied as f32 } else { 0.0 }; metrics.phase_timings.push(("metadata_boost".to_string(), t2.elapsed().as_millis() as u64)); // Re-sort by boosted score enriched_chunks.sort_by(|a, b| { b.final_score.partial_cmp(&a.final_score).unwrap_or(std::cmp::Ordering::Equal) }); // Phase 6: Cache alignment (simplified without wiki distances) let t3 = std::time::Instant::now(); let cached: Vec = enriched_chunks .iter() .enumerate() .map(|(i, chunk)| CachedChunk { chunk_id: chunk.id.clone(), text: chunk.text.clone(), score: chunk.final_score, cache_distance: i as u32, // Use position as distance proxy access_count: 1, last_accessed_slot: i as u32, }) .collect(); let slots = self.aligner.assign_slots(&cached); for chunk in &mut enriched_chunks { if let Some((_, slot)) = slots.iter().find(|(id, _)| id == &chunk.id) { chunk.cache_slot = *slot; } chunk.cache_priority = chunk.final_score; } let preload_chunks: Vec<_> = enriched_chunks .iter() .take(self.config.preload_top_k) .map(|c| (c.id.as_str(), c.text.as_str())) .collect(); self.aligner.preload_hot_chunks(preload_chunks)?; metrics.preloaded_chunks = self.config.preload_top_k.min(enriched_chunks.len()); let cache_metrics = self.aligner.get_metrics(); metrics.cache_hits = cache_metrics.hits; metrics.cache_misses = cache_metrics.misses; metrics.cache_hit_ratio = cache_metrics.hit_ratio(); metrics.phase_timings.push(("cache_alignment".to_string(), t3.elapsed().as_millis() as u64)); metrics.total_latency_ms = start.elapsed().as_millis() as u64; Ok(PipelineResult { query: query.to_string(), query_intent, chunks: enriched_chunks, metrics, }) } /// Get config pub fn config(&self) -> &PipelineConfig { &self.config } /// Get profiler summary pub fn profiler_summary(&self) -> Vec<(String, u64)> { self.profiler.summary() } } /// Builder for FullPipeline with sensible defaults pub struct PipelineBuilder { tfidf_scorer: Option>, semantic_scorer: Option>, config: PipelineConfig, } impl PipelineBuilder { pub fn new() -> Self { Self { tfidf_scorer: None, semantic_scorer: None, config: PipelineConfig::default(), } } pub fn with_scorers( mut self, tfidf: Arc, semantic: Arc, ) -> Self { self.tfidf_scorer = Some(tfidf); self.semantic_scorer = Some(semantic); self } pub fn with_project(mut self, project: &str) -> Self { self.config.project = project.to_string(); self } pub fn with_wiki_root(mut self, root_doc: &str) -> Self { self.config.wiki_root_doc = root_doc.to_string(); self } pub fn with_budget(mut self, bytes: usize) -> Self { self.config.budget_bytes = bytes; self } pub fn with_score_threshold(mut self, threshold: f32) -> Self { self.config.score_threshold = threshold; self } pub fn with_metadata_boost(mut self, enabled: bool) -> Self { self.config.enable_metadata_boost = enabled; self } pub fn with_cache_capacity(mut self, capacity: usize) -> Self { self.config.cache_capacity = capacity; self } pub fn with_config(mut self, config: PipelineConfig) -> Self { self.config = config; self } pub fn build(self) -> Result { let tfidf = self.tfidf_scorer .ok_or_else(|| anyhow::anyhow!("TF-IDF scorer required"))?; let semantic = self.semantic_scorer .ok_or_else(|| anyhow::anyhow!("Semantic scorer required"))?; Ok(FullPipeline::new(tfidf, semantic, self.config)) } } #[cfg(test)] mod tests { use super::*; use std::collections::BTreeMap; fn create_test_vocab() -> Arc> { let mut vocab = BTreeMap::new(); vocab.insert("kubernetes".to_string(), 0.8); vocab.insert("pod".to_string(), 0.7); vocab.insert("error".to_string(), 0.9); vocab.insert("fix".to_string(), 0.85); vocab.insert("solution".to_string(), 0.8); Arc::new(vocab) } fn create_test_pipeline() -> FullPipeline { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); FullPipeline::new(tfidf, semantic, PipelineConfig::default()) } fn create_test_wiki_graph() -> WikiLinkGraph { let mut graph = WikiLinkGraph::new("test"); graph.add_link("index.md", "tools/kubectl.md"); graph.add_link("tools/kubectl.md", "debugging/pod-errors.md"); graph.add_link("debugging/pod-errors.md", "solutions/restart.md"); graph } fn create_test_candidates() -> Vec<(String, String)> { vec![ ("index.md".to_string(), "# Index\nKubernetes documentation.".to_string()), ("tools/kubectl.md".to_string(), "# Kubectl\nTool for kubernetes pod management.".to_string()), ("debugging/pod-errors.md".to_string(), "# Pod Errors\nError: CrashLoopBackOff. Fix by checking logs.".to_string()), ("solutions/restart.md".to_string(), "# Restart Solution\nSolution: restart the failing pod.".to_string()), ("unrelated.md".to_string(), "# Unrelated\nDocker container guide.".to_string()), ] } #[test] fn test_pipeline_config_default() { let config = PipelineConfig::default(); assert_eq!(config.max_wiki_hops, 3); assert_eq!(config.score_threshold, 0.6); assert_eq!(config.budget_bytes, 8192); assert!(config.enable_metadata_boost); } #[test] fn test_pipeline_metrics_new() { let metrics = PipelineMetrics::new(); assert_eq!(metrics.wiki_scope_docs, 0); assert_eq!(metrics.total_latency_ms, 0); assert!(metrics.phase_timings.is_empty()); } #[test] fn test_enriched_chunk_structure() { let chunk = EnrichedChunk { id: "doc1".to_string(), text: "content".to_string(), tfidf_score: 0.4, semantic_score: 0.6, rrf_score: 0.5, pre_boost_score: 0.5, final_score: 0.6, category: ChunkCategory::Solution, heading: Some("Fix Pods".to_string()), key_terms: vec!["kubernetes".to_string()], metadata_boost: 0.1, query_intent_match: true, wiki_distance: Some(2), cache_slot: 0, cache_priority: 0.8, }; assert_eq!(chunk.id, "doc1"); assert!(chunk.query_intent_match); assert_eq!(chunk.wiki_distance, Some(2)); } #[test] fn test_pipeline_builder() { let vocab = create_test_vocab(); let tfidf = Arc::new(GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(SemanticScorer::new()); let pipeline = PipelineBuilder::new() .with_scorers(tfidf, semantic) .with_project("test-project") .with_budget(4096) .with_score_threshold(0.7) .with_metadata_boost(true) .build() .unwrap(); assert_eq!(pipeline.config().project, "test-project"); assert_eq!(pipeline.config().budget_bytes, 4096); assert_eq!(pipeline.config().score_threshold, 0.7); } #[test] fn test_pipeline_builder_missing_scorers() { let result = PipelineBuilder::new().build(); assert!(result.is_err()); } #[tokio::test] async fn test_execute_with_wiki() { let pipeline = create_test_pipeline(); let graph = create_test_wiki_graph(); let candidates = create_test_candidates(); let result = pipeline .execute_with_wiki("fix kubernetes pod error", &graph, candidates) .await .unwrap(); assert_eq!(result.query, "fix kubernetes pod error"); assert_eq!(result.query_intent, QueryIntent::FixError); assert!(result.metrics.total_latency_ms >= 0); assert!(!result.metrics.phase_timings.is_empty()); } #[tokio::test] async fn test_execute_direct() { let pipeline = create_test_pipeline(); let candidates = create_test_candidates(); let result = pipeline .execute_direct("kubernetes deployment", candidates) .await .unwrap(); assert_eq!(result.query, "kubernetes deployment"); assert!(result.metrics.wiki_scope_docs > 0); } #[tokio::test] async fn test_metadata_boost_applied() { let pipeline = create_test_pipeline(); let candidates = vec![ ("error-doc.md".to_string(), "# Error\nPod error CrashLoopBackOff fix solution.".to_string()), ("concept-doc.md".to_string(), "# Concept\nKubernetes pod design pattern.".to_string()), ]; let result = pipeline .execute_direct("fix pod error", candidates) .await .unwrap(); // FixError query should boost error/solution chunks assert_eq!(result.query_intent, QueryIntent::FixError); // Check that metadata boost was applied for chunk in &result.chunks { if chunk.category == ChunkCategory::Error || chunk.category == ChunkCategory::Solution { // These should have intent match if chunk.text.contains("error") || chunk.text.contains("solution") { // Boost might be applied depending on category detection } } } } #[tokio::test] async fn test_cache_preload() { let pipeline = create_test_pipeline(); let candidates = create_test_candidates(); let result = pipeline .execute_direct("kubernetes", candidates) .await .unwrap(); // Should have preloaded some chunks assert!(result.metrics.preloaded_chunks <= pipeline.config().preload_top_k); } #[tokio::test] async fn test_wiki_distance_calculation() { let pipeline = create_test_pipeline(); let graph = create_test_wiki_graph(); let candidates = create_test_candidates(); let result = pipeline .execute_with_wiki("kubernetes", &graph, candidates) .await .unwrap(); // Chunks should have wiki_distance populated for chunk in &result.chunks { // Wiki distances should be within max_hops or None if unreachable if let Some(dist) = chunk.wiki_distance { assert!(dist <= pipeline.config().max_wiki_hops); } } } #[tokio::test] async fn test_phase_timings() { let pipeline = create_test_pipeline(); let candidates = create_test_candidates(); let result = pipeline .execute_direct("test query", candidates) .await .unwrap(); // Should have timing for all phases let phase_names: Vec<_> = result.metrics.phase_timings.iter().map(|(n, _)| n.as_str()).collect(); assert!(phase_names.contains(&"intent_inference")); assert!(phase_names.contains(&"routing_retrieval")); assert!(phase_names.contains(&"metadata_boost")); assert!(phase_names.contains(&"cache_alignment")); } #[tokio::test] async fn test_empty_candidates() { let pipeline = create_test_pipeline(); let result = pipeline .execute_direct("query", vec![]) .await .unwrap(); assert!(result.chunks.is_empty()); assert_eq!(result.metrics.post_optimization_count, 0); } #[test] fn test_cache_priority_calculation() { // Higher score + closer wiki distance = higher priority let chunk_close = EnrichedChunk { id: "close".to_string(), text: "".to_string(), tfidf_score: 0.0, semantic_score: 0.0, rrf_score: 0.0, pre_boost_score: 0.0, final_score: 0.8, category: ChunkCategory::Unknown, heading: None, key_terms: vec![], metadata_boost: 0.0, query_intent_match: false, wiki_distance: Some(1), cache_slot: 0, cache_priority: 0.8 * (1.0 / 2.0), // score * 1/(1+dist) }; let chunk_far = EnrichedChunk { wiki_distance: Some(5), cache_priority: 0.8 * (1.0 / 6.0), ..chunk_close.clone() }; assert!(chunk_close.cache_priority > chunk_far.cache_priority); } }