feat(phase3-4): Complete hybrid retrieval + LLM optimization pipeline

Phase 3: Hybrid Retrieval - HybridRetriever: TF-IDF prefilter + semantic rerank + RRF fusion - WikiScopedFilter: BFS wiki-graph traversal - RetrievalRoute: Direct | WikiScoped | ReferenceOnly - 10 unit tests Phase 4: LLM Call Optimization - ChunkOptimizer: unified pipeline (threshold + budget + dedup) - ScoreThresholdFilter: configurable min_score (default 0.6) - BudgetSelector: greedy selection within byte budget - ShingleDeduplicator: Jaccard similarity dedup - 8 unit tests QueryRouter (Phase 3+4 Integration) - Bridges WikiLinkGraph + HybridRetriever + ChunkOptimizer - RouterConfig: max_hops, thresholds, budget, RRF weights - WikiGraphBuilder: construct graph from markdown docs - 11 unit tests Integration Tests (it_phase3_phase4.rs) - 19 end-to-end tests covering full pipeline - Wiki-link parsing, graph traversal, route selection - TF-IDF prefilter, RRF fusion, chunk optimization - Edge cases (empty, no matches, config customization) Total: 107 tests passing (was 32)
2026-08-31 22:42:34 -07:00
parent b71831557d
commit cd76424baa
4 changed files with 1039 additions and 59 deletions
@@ -25,6 +25,7 @@ pub mod query_filter;
 pub mod advanced_ranking;
 pub mod result_compressor;
 pub mod federation;
+pub mod query_router;

 pub use endpoints::{IngestQueue, IngestRequest, JobStatus};
 pub use ingest_worker::IngestWorker;
@@ -35,3 +36,4 @@ pub use chunk_metadata::{MetadataExtractor, MetadataBooster, ChunkMetadata, Chun
 pub use cache_alignment::{LruChunkCache, KvCacheAligner, CacheLocalityAnalyzer, RetrievalProfiler, CacheMetrics};
 pub use query_orchestrator::{QueryOrchestrator, QueryResult, OptimizedChunk, QueryContext, MemoryProjection};
 pub use query_filter::{QueryFilter, FilterableDocument, FilterEngine, FilterStatistics};
+pub use query_router::{QueryRouter, RouterConfig, RoutedResult, SelectedChunk, WikiGraphBuilder};
@@ -0,0 +1,485 @@
+/// Query Router: Unified Phase 3+4 pipeline
+///
+/// Bridges wiki-link graph (Phase 1) with hybrid retrieval (Phase 3)
+/// and LLM optimization (Phase 4) into a single query flow.
+///
+/// Pipeline:
+/// 1. Wiki-scope filtering (via WikiLinkGraph)
+/// 2. TF-IDF pre-filtering
+/// 3. Semantic re-ranking
+/// 4. RRF fusion
+/// 5. Score thresholding + budget selection + deduplication
+
+use anyhow::Result;
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use mem_ingest::wiki_link::{WikiLinkGraph, WikiLinkParser};
+use mem_core::{DocumentScorer, GlobalTfIdfScorer, SemanticScorer};
+
+use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate};
+use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics};
+
+/// Query routing configuration
+#[derive(Debug, Clone)]
+pub struct RouterConfig {
+    pub max_wiki_hops: u32,
+    pub tfidf_threshold: f32,
+    pub prefilter_limit: usize,
+    pub score_threshold: f32,
+    pub budget_bytes: usize,
+    pub dedup_threshold: f32,
+    pub rrf_tfidf_weight: f32,
+    pub rrf_semantic_weight: f32,
+}
+
+impl Default for RouterConfig {
+    fn default() -> Self {
+        Self {
+            max_wiki_hops: 3,
+            tfidf_threshold: 0.3,
+            prefilter_limit: 50,
+            score_threshold: 0.6,
+            budget_bytes: 8192,
+            dedup_threshold: 0.8,
+            rrf_tfidf_weight: 0.4,
+            rrf_semantic_weight: 0.6,
+        }
+    }
+}
+
+/// Query routing result with full metrics
+#[derive(Debug, Clone)]
+pub struct RoutedResult {
+    pub selected_chunks: Vec<SelectedChunk>,
+    pub route: RetrievalRoute,
+    pub wiki_scope_size: usize,
+    pub prefilter_size: usize,
+    pub metrics: SelectionMetrics,
+    pub latency_ms: u64,
+}
+
+/// Selected chunk with all scores
+#[derive(Debug, Clone)]
+pub struct SelectedChunk {
+    pub id: String,
+    pub text: String,
+    pub tfidf_score: f32,
+    pub semantic_score: f32,
+    pub final_score: f32,
+    pub wiki_distance: Option<u32>,
+}
+
+/// Query Router: end-to-end Phase 3+4 pipeline
+pub struct QueryRouter {
+    wiki_filter: WikiScopedFilter,
+    retriever: HybridRetriever,
+    optimizer: ChunkOptimizer,
+    config: RouterConfig,
+}
+
+impl QueryRouter {
+    pub fn new(
+        tfidf_scorer: Arc<GlobalTfIdfScorer>,
+        semantic_scorer: Arc<SemanticScorer>,
+        config: RouterConfig,
+    ) -> Self {
+        let wiki_filter = WikiScopedFilter::new(config.max_wiki_hops);
+        let retriever = HybridRetriever::new(tfidf_scorer, semantic_scorer);
+        let optimizer = ChunkOptimizer::new(
+            config.score_threshold,
+            config.budget_bytes,
+            config.dedup_threshold,
+        );
+
+        Self {
+            wiki_filter,
+            retriever,
+            optimizer,
+            config,
+        }
+    }
+
+    /// Execute full query pipeline with wiki-link graph scoping
+    pub async fn route_with_wiki_graph(
+        &self,
+        query: &str,
+        wiki_graph: &WikiLinkGraph,
+        root_doc: &str,
+        all_candidates: Vec<(String, String)>, // (doc_id, text)
+    ) -> Result<RoutedResult> {
+        let start = std::time::Instant::now();
+
+        // Phase 1: Wiki-scope reduction
+        let wiki_reachable = wiki_graph.reachable_docs(root_doc);
+        let wiki_scope_size = wiki_reachable.len();
+
+        // Convert wiki-graph to HashMap for WikiScopedFilter
+        let graph_map = self.wiki_graph_to_hashmap(wiki_graph, root_doc);
+
+        // Filter candidates by wiki scope
+        let scoped_candidates: Vec<_> = all_candidates
+            .into_iter()
+            .filter(|(doc_id, _)| wiki_reachable.contains(doc_id))
+            .collect();
+
+        // Phase 3: Hybrid retrieval
+        let route = self.retriever.route_query(query, !wiki_reachable.is_empty(), false);
+        let ranked = self.retriever.retrieve(query, scoped_candidates, route.clone()).await?;
+        let prefilter_size = ranked.len();
+
+        // Convert to optimizable chunks
+        let optimizable: Vec<OptimizableChunk> = ranked
+            .into_iter()
+            .map(|r| {
+                let size = r.text.len();
+                OptimizableChunk {
+                    id: r.doc_id,
+                    text: r.text,
+                    score: r.final_score,
+                    confidence: r.semantic_score,
+                    size_bytes: size,
+                }
+            })
+            .collect();
+
+        // Phase 4: LLM optimization (threshold + budget + dedup)
+        let (selected_opt, metrics) = self.optimizer.optimize(optimizable);
+
+        // Build final result with wiki distances
+        let selected_chunks: Vec<SelectedChunk> = selected_opt
+            .into_iter()
+            .map(|chunk| {
+                let wiki_distance = self.calculate_wiki_distance(&chunk.id, root_doc, &graph_map);
+                SelectedChunk {
+                    id: chunk.id,
+                    text: chunk.text,
+                    tfidf_score: chunk.score * self.config.rrf_tfidf_weight,
+                    semantic_score: chunk.score * self.config.rrf_semantic_weight,
+                    final_score: chunk.score,
+                    wiki_distance,
+                }
+            })
+            .collect();
+
+        let latency_ms = start.elapsed().as_millis() as u64;
+
+        Ok(RoutedResult {
+            selected_chunks,
+            route,
+            wiki_scope_size,
+            prefilter_size,
+            metrics,
+            latency_ms,
+        })
+    }
+
+    /// Execute query without wiki-graph (direct retrieval)
+    pub async fn route_direct(
+        &self,
+        query: &str,
+        all_candidates: Vec<(String, String)>,
+    ) -> Result<RoutedResult> {
+        let start = std::time::Instant::now();
+
+        // Direct retrieval (no wiki scoping)
+        let route = RetrievalRoute::Direct;
+        let ranked = self.retriever.retrieve(query, all_candidates.clone(), route.clone()).await?;
+        let prefilter_size = ranked.len();
+
+        // Convert to optimizable chunks
+        let optimizable: Vec<OptimizableChunk> = ranked
+            .into_iter()
+            .map(|r| {
+                let size = r.text.len();
+                OptimizableChunk {
+                    id: r.doc_id,
+                    text: r.text,
+                    score: r.final_score,
+                    confidence: r.semantic_score,
+                    size_bytes: size,
+                }
+            })
+            .collect();
+
+        // Phase 4: LLM optimization
+        let (selected_opt, metrics) = self.optimizer.optimize(optimizable);
+
+        let selected_chunks: Vec<SelectedChunk> = selected_opt
+            .into_iter()
+            .map(|chunk| SelectedChunk {
+                id: chunk.id,
+                text: chunk.text,
+                tfidf_score: chunk.score * self.config.rrf_tfidf_weight,
+                semantic_score: chunk.score * self.config.rrf_semantic_weight,
+                final_score: chunk.score,
+                wiki_distance: None,
+            })
+            .collect();
+
+        let latency_ms = start.elapsed().as_millis() as u64;
+
+        Ok(RoutedResult {
+            selected_chunks,
+            route,
+            wiki_scope_size: all_candidates.len(),
+            prefilter_size,
+            metrics,
+            latency_ms,
+        })
+    }
+
+    /// Convert WikiLinkGraph to HashMap for distance calculation
+    fn wiki_graph_to_hashmap(
+        &self,
+        wiki_graph: &WikiLinkGraph,
+        root_doc: &str,
+    ) -> HashMap<String, Vec<String>> {
+        let reachable = wiki_graph.reachable_docs(root_doc);
+        let mut graph_map = HashMap::new();
+
+        for doc in &reachable {
+            let forward = wiki_graph.forward_links(doc);
+            graph_map.insert(doc.clone(), forward);
+        }
+
+        graph_map
+    }
+
+    /// Calculate wiki distance using BFS
+    fn calculate_wiki_distance(
+        &self,
+        doc_id: &str,
+        root_doc: &str,
+        graph: &HashMap<String, Vec<String>>,
+    ) -> Option<u32> {
+        if doc_id == root_doc {
+            return Some(0);
+        }
+
+        let mut visited = std::collections::HashSet::new();
+        let mut queue = std::collections::VecDeque::new();
+
+        queue.push_back((root_doc.to_string(), 0u32));
+        visited.insert(root_doc.to_string());
+
+        while let Some((current, distance)) = queue.pop_front() {
+            if current == doc_id {
+                return Some(distance);
+            }
+
+            if distance >= self.config.max_wiki_hops {
+                continue;
+            }
+
+            if let Some(neighbors) = graph.get(&current) {
+                for neighbor in neighbors {
+                    if !visited.contains(neighbor) {
+                        visited.insert(neighbor.clone());
+                        queue.push_back((neighbor.clone(), distance + 1));
+                    }
+                }
+            }
+        }
+
+        None // Not reachable
+    }
+
+    pub fn config(&self) -> &RouterConfig {
+        &self.config
+    }
+}
+
+/// Build wiki-link graph from markdown content
+pub struct WikiGraphBuilder;
+
+impl WikiGraphBuilder {
+    /// Build graph from list of (doc_id, content) pairs
+    pub fn build_from_docs(
+        project: &str,
+        docs: Vec<(&str, &str)>,
+    ) -> Result<WikiLinkGraph> {
+        let mut graph = WikiLinkGraph::new(project);
+
+        for (doc_id, content) in docs {
+            let links = WikiLinkParser::parse_links(content)?;
+            for target in links {
+                graph.add_link(doc_id, &target);
+            }
+        }
+
+        Ok(graph)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::collections::BTreeMap;
+
+    fn create_test_router() -> QueryRouter {
+        let vocab = Arc::new(BTreeMap::new());
+        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("test");
+        graph.add_link("index.md", "tools/kubectl.md");
+        graph.add_link("tools/kubectl.md", "debugging/pod-crashes.md");
+        graph.add_link("debugging/pod-crashes.md", "solutions/restart-pod.md");
+        graph
+    }
+
+    #[test]
+    fn test_router_config_default() {
+        let config = RouterConfig::default();
+        assert_eq!(config.max_wiki_hops, 3);
+        assert_eq!(config.score_threshold, 0.6);
+        assert_eq!(config.budget_bytes, 8192);
+    }
+
+    #[test]
+    fn test_wiki_graph_to_hashmap() {
+        let router = create_test_router();
+        let graph = create_test_wiki_graph();
+
+        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");
+
+        assert!(hashmap.contains_key("index.md"));
+        assert!(hashmap.contains_key("tools/kubectl.md"));
+        assert!(hashmap.contains_key("debugging/pod-crashes.md"));
+    }
+
+    #[test]
+    fn test_calculate_wiki_distance_root() {
+        let router = create_test_router();
+        let graph = create_test_wiki_graph();
+        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");
+
+        let distance = router.calculate_wiki_distance("index.md", "index.md", &hashmap);
+        assert_eq!(distance, Some(0));
+    }
+
+    #[test]
+    fn test_calculate_wiki_distance_direct_child() {
+        let router = create_test_router();
+        let graph = create_test_wiki_graph();
+        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");
+
+        let distance = router.calculate_wiki_distance("tools/kubectl.md", "index.md", &hashmap);
+        assert_eq!(distance, Some(1));
+    }
+
+    #[test]
+    fn test_calculate_wiki_distance_grandchild() {
+        let router = create_test_router();
+        let graph = create_test_wiki_graph();
+        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");
+
+        let distance = router.calculate_wiki_distance("debugging/pod-crashes.md", "index.md", &hashmap);
+        assert_eq!(distance, Some(2));
+    }
+
+    #[test]
+    fn test_calculate_wiki_distance_unreachable() {
+        let router = create_test_router();
+        let graph = create_test_wiki_graph();
+        let hashmap = router.wiki_graph_to_hashmap(&graph, "index.md");
+
+        let distance = router.calculate_wiki_distance("unknown.md", "index.md", &hashmap);
+        assert_eq!(distance, None);
+    }
+
+    #[tokio::test]
+    async fn test_route_direct() {
+        let router = create_test_router();
+        let candidates = vec![
+            ("doc1".to_string(), "kubernetes pod debugging".to_string()),
+            ("doc2".to_string(), "docker container deployment".to_string()),
+        ];
+
+        let result = router.route_direct("kubernetes", candidates).await.unwrap();
+
+        assert_eq!(result.route, RetrievalRoute::Direct);
+        assert!(result.latency_ms >= 0);
+    }
+
+    #[tokio::test]
+    async fn test_route_with_wiki_graph() {
+        let router = create_test_router();
+        let graph = create_test_wiki_graph();
+
+        let candidates = vec![
+            ("index.md".to_string(), "main index".to_string()),
+            ("tools/kubectl.md".to_string(), "kubectl tool".to_string()),
+            ("debugging/pod-crashes.md".to_string(), "debugging content".to_string()),
+            ("unrelated.md".to_string(), "not in graph".to_string()),
+        ];
+
+        let result = router
+            .route_with_wiki_graph("kubectl", &graph, "index.md", candidates)
+            .await
+            .unwrap();
+
+        // Should filter out "unrelated.md" (not reachable from index.md)
+        assert!(result.wiki_scope_size <= 4);
+        assert_eq!(result.route, RetrievalRoute::WikiScoped);
+    }
+
+    #[test]
+    fn test_wiki_graph_builder() {
+        let docs = vec![
+            ("index.md", "# Index\nSee [[tools/kubectl.md]] for tools."),
+            ("tools/kubectl.md", "# Kubectl\nSee [[debugging.md]] for debugging."),
+        ];
+
+        let graph = WikiGraphBuilder::build_from_docs("test", docs).unwrap();
+
+        let reachable = graph.reachable_docs("index.md");
+        assert!(reachable.contains("index.md"));
+        assert!(reachable.contains("tools/kubectl.md"));
+        assert!(reachable.contains("debugging.md"));
+    }
+
+    #[test]
+    fn test_selected_chunk_structure() {
+        let chunk = SelectedChunk {
+            id: "doc1".to_string(),
+            text: "content".to_string(),
+            tfidf_score: 0.4,
+            semantic_score: 0.6,
+            final_score: 0.9,
+            wiki_distance: Some(1),
+        };
+
+        assert_eq!(chunk.id, "doc1");
+        assert!(chunk.final_score <= 1.0);
+        assert_eq!(chunk.wiki_distance, Some(1));
+    }
+
+    #[test]
+    fn test_routed_result_structure() {
+        let result = RoutedResult {
+            selected_chunks: vec![],
+            route: RetrievalRoute::WikiScoped,
+            wiki_scope_size: 10,
+            prefilter_size: 5,
+            metrics: SelectionMetrics {
+                selected_count: 3,
+                rejected_count: 2,
+                total_bytes: 1000,
+                budget_used_pct: 12.5,
+                avg_score: 0.8,
+                dedup_removed: 0,
+            },
+            latency_ms: 50,
+        };
+
+        assert_eq!(result.wiki_scope_size, 10);
+        assert_eq!(result.prefilter_size, 5);
+        assert_eq!(result.metrics.selected_count, 3);
+    }
+}