poimen-memory/crates/mem-cli/src/query_router.rs

/// 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,
    pub confidence_score: f32,          // Multi-signal confidence (0-1)
    pub is_valid: bool,                 // Passes validation gate
}

/// 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;

        // Phase 8: Answer Validation (confidence scoring)
        use crate::query::answer_validator::{AnswerValidator, AnswerValidationConfig, ConfidenceSignals};
        let validator = AnswerValidator::new(AnswerValidationConfig::default());
        let avg_score = selected_chunks.iter().map(|c| c.final_score).sum::<f32>() 
            / (selected_chunks.len() as f32).max(1.0);
        let signals = ConfidenceSignals {
            search_score: avg_score,
            evidence_count: selected_chunks.len(),
            evidence_confidence: avg_score,
            temporal_score: 0.9,  // Assume recent chunks
            entity_coverage: 0.85,
            contradiction_score: 1.0,  // No contradictions by default
        };
        let validated = validator.validate("", &signals);

        Ok(RoutedResult {
            selected_chunks,
            route,
            wiki_scope_size,
            prefilter_size,
            metrics,
            latency_ms,
            confidence_score: validated.overall_confidence,
            is_valid: validated.is_valid,
        })
    }

    /// 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,
            confidence_score: 1.0,
            is_valid: true,
        })
    }

    /// 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)
    }
}