/// Query Orchestrator: Unified interface combining all phases 1-6 /// /// Orchestrates: /// - Phase 1: Wiki-link graph traversal /// - Phase 2: Scoring pipeline /// - Phase 3: Hybrid retrieval (TF-IDF + semantic + RRF) /// - Phase 4: LLM optimization (threshold, budget, dedup) /// - Phase 5: Metadata enhancement (category + intent boost) /// - Phase 6: Cache alignment (locality + pre-load) use anyhow::Result; use std::collections::HashMap; use std::sync::Arc; use mem_core::DocumentScorer; use crate::hybrid_retrieval::{HybridRetriever, RetrievalRoute, WikiScopedFilter, RankedCandidate}; use crate::chunk_optimizer::{ChunkOptimizer, OptimizableChunk, SelectionMetrics}; use crate::chunk_metadata::{MetadataExtractor, MetadataBooster, QueryIntent}; use crate::cache_alignment::{KvCacheAligner, CachedChunk, CacheLocalityAnalyzer, RetrievalProfiler}; /// Complete query result with all metadata #[derive(Debug, Clone)] pub struct QueryResult { pub query: String, pub selected_chunks: Vec, pub selection_metrics: SelectionMetrics, pub cache_metrics: crate::cache_alignment::CacheMetrics, pub profiling: Vec<(String, u64)>, // stage -> duration_ms pub total_latency_ms: u64, pub query_intent: QueryIntent, } /// Chunk with all enrichments #[derive(Debug, Clone)] pub struct OptimizedChunk { pub id: String, pub text: String, pub tfidf_score: f32, pub semantic_score: f32, pub metadata_boost: f32, pub final_score: f32, pub category: crate::chunk_metadata::ChunkCategory, pub cache_distance: u32, pub cache_slot: u32, } /// Query execution context pub struct QueryContext { pub project: String, pub wiki_root_doc: String, pub max_wiki_hops: u32, pub budget_bytes: usize, pub score_threshold: f32, pub dedup_threshold: f32, pub cache_capacity: usize, } impl Default for QueryContext { fn default() -> Self { Self { project: "default".to_string(), wiki_root_doc: "index.md".to_string(), max_wiki_hops: 3, budget_bytes: 8192, score_threshold: 0.6, dedup_threshold: 0.8, cache_capacity: 1000, } } } /// Orchestrator: combines all phases pub struct QueryOrchestrator { retriever: Arc, optimizer: Arc, booster: Arc, aligner: Arc, profiler: Arc, } impl QueryOrchestrator { pub fn new( tfidf_scorer: Arc, semantic_scorer: Arc, context: &QueryContext, ) -> Self { let retriever = Arc::new(HybridRetriever::new( tfidf_scorer.clone(), semantic_scorer.clone(), )); let optimizer = Arc::new(ChunkOptimizer::new( context.score_threshold, context.budget_bytes, context.dedup_threshold, )); let booster = Arc::new(MetadataBooster::new()); let aligner = Arc::new(KvCacheAligner::new(4096, 100, context.cache_capacity)); let profiler = Arc::new(RetrievalProfiler::new()); Self { retriever, optimizer, booster, aligner, profiler, } } /// End-to-end query execution pub async fn execute( &self, query: &str, all_candidates: Vec<(String, String)>, // (doc_id, text) context: &QueryContext, ) -> Result { let start = std::time::Instant::now(); // Step 1: Infer query intent (Phase 5) let query_intent = MetadataExtractor::infer_query_intent(query); self.profiler.record("infer_intent", 1); // Step 2: Route retrieval (Phase 3) let has_wiki_scope = !context.wiki_root_doc.is_empty(); let route = self.retriever.route_query(query, has_wiki_scope, false); self.profiler.record("route_selection", 2); // Step 3: Hybrid retrieval (Phase 3) let start_retrieval = std::time::Instant::now(); let ranked = self .retriever .retrieve(query, all_candidates, route.clone()) .await?; let retrieval_time = start_retrieval.elapsed().as_millis() as u64; self.profiler.record("hybrid_retrieval", retrieval_time); // Step 4: Convert to optimizable chunks let mut optimizable: Vec = ranked .into_iter() .map(|r| { let text_len = r.text.len(); OptimizableChunk { id: r.doc_id, text: r.text, score: r.final_score, confidence: r.semantic_score, // Confidence from semantic size_bytes: text_len, } }) .collect(); // Step 5: Metadata enhancement (Phase 5) let start_metadata = std::time::Instant::now(); for chunk in &mut optimizable { let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text); let boost = self.booster.calculate_boost(query_intent, &metadata); chunk.score = self.booster.apply_boost(chunk.score, boost); } let metadata_time = start_metadata.elapsed().as_millis() as u64; self.profiler.record("metadata_boost", metadata_time); // Step 6: LLM optimization (Phase 4) let start_optimize = std::time::Instant::now(); let (selected_opt, selection_metrics) = self.optimizer.optimize(optimizable.clone()); let optimize_time = start_optimize.elapsed().as_millis() as u64; self.profiler.record("llm_optimize", optimize_time); // Step 7: Cache alignment (Phase 6) let start_cache = std::time::Instant::now(); let cached: Vec = selected_opt .iter() .enumerate() .map(|(i, chunk)| CachedChunk { chunk_id: chunk.id.clone(), text: chunk.text.clone(), score: chunk.score, cache_distance: 0, // Would be computed from wiki-graph access_count: 1, last_accessed_slot: i as u32, }) .collect(); let slots = self.aligner.assign_slots(&cached); self.aligner.preload_hot_chunks( cached.iter().take(5).map(|c| (c.chunk_id.as_str(), c.text.as_str())).collect() )?; let cache_time = start_cache.elapsed().as_millis() as u64; self.profiler.record("cache_align", cache_time); // Step 8: Build optimized chunks with all metadata let mut optimized_chunks = Vec::new(); for (i, chunk) in selected_opt.iter().enumerate() { let slot = slots.iter().find(|(id, _)| id == &chunk.id).map(|(_, s)| *s).unwrap_or(0); let metadata = MetadataExtractor::extract(&chunk.id, &chunk.text); optimized_chunks.push(OptimizedChunk { id: chunk.id.clone(), text: chunk.text.clone(), tfidf_score: chunk.score * 0.4, // Approximate semantic_score: chunk.score * 0.6, metadata_boost: 0.0, // Already applied final_score: chunk.score, category: metadata.category, cache_distance: 0, cache_slot: slot, }); } let total_latency = start.elapsed().as_millis() as u64; Ok(QueryResult { query: query.to_string(), selected_chunks: optimized_chunks, selection_metrics, cache_metrics: self.aligner.get_metrics(), profiling: self.profiler.summary(), total_latency_ms: total_latency, query_intent, }) } } /// Memory projection for multi-project queries pub struct MemoryProjection { projects: HashMap>, } impl MemoryProjection { pub fn new() -> Self { Self { projects: HashMap::new(), } } pub fn register_project( &mut self, project: &str, orchestrator: Arc, ) { self.projects.insert(project.to_string(), orchestrator); } pub async fn query_project( &self, project: &str, query: &str, candidates: Vec<(String, String)>, context: &QueryContext, ) -> Result { let orchestrator = self .projects .get(project) .ok_or_else(|| anyhow::anyhow!("Project not found: {}", project))?; orchestrator.execute(query, candidates, context).await } pub fn projects(&self) -> Vec<&str> { self.projects.keys().map(|s| s.as_str()).collect() } } #[cfg(test)] mod tests { use super::*; use std::collections::BTreeMap; #[test] fn test_query_context_default() { let ctx = QueryContext::default(); assert_eq!(ctx.project, "default"); assert_eq!(ctx.budget_bytes, 8192); assert_eq!(ctx.max_wiki_hops, 3); } #[test] fn test_memory_projection_register() { let mut proj = MemoryProjection::new(); let vocab = Arc::new(BTreeMap::new()); let scorer = Arc::new(mem_core::GlobalTfIdfScorer::new(vocab)); let semantic = Arc::new(mem_core::SemanticScorer::new()); let orchestrator = Arc::new(QueryOrchestrator::new(scorer, semantic, &QueryContext::default())); proj.register_project("test", orchestrator); assert!(proj.projects().contains(&"test")); } #[test] fn test_memory_projection_unknown_project() { let proj = MemoryProjection::new(); let candidates = vec![("doc1".to_string(), "content".to_string())]; let ctx = QueryContext::default(); let result = tokio::runtime::Runtime::new() .unwrap() .block_on(proj.query_project("unknown", "test", candidates, &ctx)); assert!(result.is_err()); } #[test] fn test_optimized_chunk_creation() { let chunk = OptimizedChunk { id: "doc1".to_string(), text: "Test content".to_string(), tfidf_score: 0.5, semantic_score: 0.8, metadata_boost: 0.1, final_score: 0.9, category: crate::chunk_metadata::ChunkCategory::Solution, cache_distance: 2, cache_slot: 0, }; assert_eq!(chunk.id, "doc1"); assert_eq!(chunk.final_score, 0.9); assert!(chunk.final_score <= 1.0); } #[test] fn test_query_result_structure() { let result = QueryResult { query: "test".to_string(), selected_chunks: vec![], selection_metrics: SelectionMetrics { selected_count: 0, rejected_count: 0, total_bytes: 0, budget_used_pct: 0.0, avg_score: 0.0, dedup_removed: 0, }, cache_metrics: crate::cache_alignment::CacheMetrics::new(), profiling: vec![], total_latency_ms: 100, query_intent: QueryIntent::Unknown, }; assert_eq!(result.total_latency_ms, 100); assert_eq!(result.selected_chunks.len(), 0); } }