engine/

hybrid.rs

//! Hybrid search combining vector similarity and full-text search
//!
//! Provides a unified search experience by combining:
//! - Vector similarity scores (cosine, euclidean, dot product)
//! - Full-text BM25 scores
//!
//! Two fusion strategies are supported:
//! - **RRF** (default): Reciprocal Rank Fusion — `score(d) = Σ 1/(k + rank_r(d))`, k=60
//! - **MinMax**: weighted linear combination of min-max normalised scores

use std::collections::HashMap;

use common::FusionStrategy;

use crate::fulltext::FullTextResult;

/// RRF smoothing constant (Cormack et al., SIGIR 2009 — k=60 is the canonical default).
const RRF_K: f32 = 60.0;

/// A vector search result row: (id, score, optional metadata, optional vector).
type VectorResultRow = (String, f32, Option<serde_json::Value>, Option<Vec<f32>>);

/// Configuration for hybrid search
#[derive(Debug, Clone)]
pub struct HybridConfig {
    /// Weight for vector search (0.0 to 1.0) — used by MinMax strategy only.
    pub vector_weight: f32,
    /// Whether to require matches in both indices.
    pub require_both: bool,
    /// CE-14: Fusion strategy (default: RRF).
    pub fusion_strategy: FusionStrategy,
}

impl Default for HybridConfig {
    fn default() -> Self {
        Self {
            vector_weight: 0.5,
            require_both: false,
            fusion_strategy: FusionStrategy::Rrf,
        }
    }
}

/// Raw score from a single search type
#[derive(Debug, Clone)]
struct RawScore {
    /// Original score before normalization
    score: f32,
    /// Additional data (metadata, vector)
    metadata: Option<serde_json::Value>,
    vector: Option<Vec<f32>>,
}

/// Result of hybrid search
#[derive(Debug, Clone)]
pub struct HybridResult {
    /// Document/vector ID
    pub id: String,
    /// Combined score (weighted average of normalized scores)
    pub combined_score: f32,
    /// Normalized vector similarity score (0-1)
    pub vector_score: f32,
    /// Normalized text search score (0-1)
    pub text_score: f32,
    /// Optional metadata
    pub metadata: Option<serde_json::Value>,
    /// Optional vector values
    pub vector: Option<Vec<f32>>,
}

/// Hybrid search engine that combines vector and text search
pub struct HybridSearcher {
    config: HybridConfig,
}

impl HybridSearcher {
    pub fn new(config: HybridConfig) -> Self {
        Self { config }
    }

    pub fn with_vector_weight(mut self, weight: f32) -> Self {
        self.config.vector_weight = weight.clamp(0.0, 1.0);
        self
    }

    /// CE-14: Override the fusion strategy.
    pub fn with_fusion_strategy(mut self, strategy: FusionStrategy) -> Self {
        self.config.fusion_strategy = strategy;
        self
    }

    /// Combine vector search results with full-text search results.
    ///
    /// Dispatches to [`Self::rrf_search`] (RRF, default) or [`Self::minmax_search`]
    /// depending on `config.fusion_strategy`.
    ///
    /// # Arguments
    /// * `vector_results` - Results from vector similarity search (id, score, metadata, vector)
    /// * `text_results` - Results from full-text BM25 search
    /// * `top_k` - Number of results to return
    pub fn search(
        &self,
        vector_results: Vec<VectorResultRow>,
        text_results: Vec<FullTextResult>,
        top_k: usize,
    ) -> Vec<HybridResult> {
        match self.config.fusion_strategy {
            FusionStrategy::Rrf => self.rrf_search(vector_results, text_results, top_k),
            FusionStrategy::MinMax => self.minmax_search(vector_results, text_results, top_k),
        }
    }

    /// Reciprocal Rank Fusion (Cormack et al., SIGIR 2009).
    ///
    /// Each document receives `score(d) = Σ_r 1 / (k + rank_r(d))` where k=60.
    /// Documents appearing in only one result list receive 0 from the missing retriever.
    fn rrf_search(
        &self,
        vector_results: Vec<VectorResultRow>,
        text_results: Vec<FullTextResult>,
        top_k: usize,
    ) -> Vec<HybridResult> {
        let mut vector_map: HashMap<String, RawScore> = HashMap::new();
        let mut vector_ranks: HashMap<String, usize> = HashMap::new();
        let mut text_ranks: HashMap<String, usize> = HashMap::new();

        // Sort vector results by score descending, assign 1-based ranks.
        let mut sorted_vec = vector_results;
        sorted_vec.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        for (i, (id, score, metadata, vector)) in sorted_vec.into_iter().enumerate() {
            vector_ranks.insert(id.clone(), i + 1);
            vector_map.insert(
                id,
                RawScore {
                    score,
                    metadata,
                    vector,
                },
            );
        }

        // Sort text results by score descending, assign 1-based ranks.
        let mut sorted_text = text_results;
        sorted_text.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        for (i, result) in sorted_text.into_iter().enumerate() {
            text_ranks.insert(result.doc_id, i + 1);
        }

        // Union of all document IDs.
        let mut all_ids: Vec<String> = vector_map
            .keys()
            .chain(text_ranks.keys())
            .cloned()
            .collect();
        all_ids.sort();
        all_ids.dedup();

        let total = all_ids.len().max(1) as f32;
        let mut results: Vec<HybridResult> = Vec::with_capacity(all_ids.len());

        for id in all_ids {
            let vec_rank = vector_ranks.get(&id).copied().unwrap_or(0);
            let txt_rank = text_ranks.get(&id).copied().unwrap_or(0);

            if self.config.require_both && (vec_rank == 0 || txt_rank == 0) {
                continue;
            }

            let vec_rrf = if vec_rank > 0 {
                1.0 / (RRF_K + vec_rank as f32)
            } else {
                0.0
            };
            let txt_rrf = if txt_rank > 0 {
                1.0 / (RRF_K + txt_rank as f32)
            } else {
                0.0
            };
            let combined = vec_rrf + txt_rrf;

            // Rank-normalised display scores (0=absent, 1=top-ranked).
            let vector_score = if vec_rank > 0 {
                1.0 - (vec_rank as f32 - 1.0) / total
            } else {
                0.0
            };
            let text_score = if txt_rank > 0 {
                1.0 - (txt_rank as f32 - 1.0) / total
            } else {
                0.0
            };

            let raw = vector_map.get(&id);
            results.push(HybridResult {
                id,
                combined_score: combined,
                vector_score,
                text_score,
                metadata: raw.and_then(|r| r.metadata.clone()),
                vector: raw.and_then(|r| r.vector.clone()),
            });
        }

        results.sort_by(|a, b| {
            b.combined_score
                .partial_cmp(&a.combined_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(top_k);
        results
    }

    /// Weighted min-max normalization (legacy fusion strategy).
    fn minmax_search(
        &self,
        vector_results: Vec<VectorResultRow>,
        text_results: Vec<FullTextResult>,
        top_k: usize,
    ) -> Vec<HybridResult> {
        let mut vector_scores: HashMap<String, RawScore> = HashMap::new();
        let mut text_scores: HashMap<String, f32> = HashMap::new();

        let mut vector_min = f32::MAX;
        let mut vector_max = f32::MIN;
        let mut text_min = f32::MAX;
        let mut text_max = f32::MIN;

        for (id, score, metadata, vector) in vector_results {
            vector_min = vector_min.min(score);
            vector_max = vector_max.max(score);
            vector_scores.insert(
                id,
                RawScore {
                    score,
                    metadata,
                    vector,
                },
            );
        }

        for result in text_results {
            text_min = text_min.min(result.score);
            text_max = text_max.max(result.score);
            text_scores.insert(result.doc_id, result.score);
        }

        let mut all_ids: Vec<String> = vector_scores
            .keys()
            .chain(text_scores.keys())
            .cloned()
            .collect();
        all_ids.sort();
        all_ids.dedup();

        let mut results: Vec<HybridResult> = Vec::new();

        for id in all_ids {
            let vector_raw = vector_scores.get(&id);
            let text_raw = text_scores.get(&id);

            if self.config.require_both && (vector_raw.is_none() || text_raw.is_none()) {
                continue;
            }

            let vector_normalized = if let Some(raw) = vector_raw {
                normalize_score(raw.score, vector_min, vector_max)
            } else {
                0.0
            };

            let text_normalized = if let Some(&score) = text_raw {
                normalize_score(score, text_min, text_max)
            } else {
                0.0
            };

            let combined = self.config.vector_weight * vector_normalized
                + (1.0 - self.config.vector_weight) * text_normalized;

            let (metadata, vector) = if let Some(raw) = vector_raw {
                (raw.metadata.clone(), raw.vector.clone())
            } else {
                (None, None)
            };

            results.push(HybridResult {
                id,
                combined_score: combined,
                vector_score: vector_normalized,
                text_score: text_normalized,
                metadata,
                vector,
            });
        }

        results.sort_by(|a, b| {
            b.combined_score
                .partial_cmp(&a.combined_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(top_k);
        results
    }
}

impl Default for HybridSearcher {
    fn default() -> Self {
        Self::new(HybridConfig::default())
    }
}

// ============================================================================
// CE-12c: Adaptive hybrid weighting
// ============================================================================

/// Return an adaptive `vector_weight` (0.0–1.0) for a [`HybridSearcher`]
/// based on the inferred [`QueryKind`].
///
/// | QueryKind | vector_weight | Rationale                          |
/// |-----------|---------------|------------------------------------|
/// | Keyword   | 0.25          | Exact-term signals dominate         |
/// | Hybrid    | 0.50          | Balanced blend                      |
/// | Semantic  | 0.75          | Embedding captures intent better    |
pub fn adaptive_vector_weight(kind: crate::routing::QueryKind) -> f32 {
    match kind {
        crate::routing::QueryKind::Keyword => 0.25,
        crate::routing::QueryKind::Hybrid => 0.50,
        crate::routing::QueryKind::Semantic => 0.75,
    }
}

/// Normalize a score to 0-1 range using min-max normalization
fn normalize_score(score: f32, min: f32, max: f32) -> f32 {
    if (max - min).abs() < f32::EPSILON {
        // All scores are the same, return 1.0
        1.0
    } else {
        (score - min) / (max - min)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_hybrid_search_basic() {
        let searcher = HybridSearcher::default();

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.7, None, None),
            ("doc3".to_string(), 0.5, None, None),
        ];

        let text_results = vec![
            FullTextResult {
                doc_id: "doc1".to_string(),
                score: 3.0,
            },
            FullTextResult {
                doc_id: "doc2".to_string(),
                score: 4.0,
            },
            FullTextResult {
                doc_id: "doc4".to_string(),
                score: 2.0,
            },
        ];

        let results = searcher.search(vector_results, text_results, 10);

        // All 4 documents should be in results
        assert_eq!(results.len(), 4);

        // Check that doc1 and doc2 have both scores >= 0
        // (normalized scores, min becomes 0.0)
        let doc1 = results.iter().find(|r| r.id == "doc1").unwrap();
        assert!(doc1.vector_score > 0.0);
        assert!(doc1.text_score >= 0.0);
        assert!(doc1.combined_score > 0.0);

        let doc2 = results.iter().find(|r| r.id == "doc2").unwrap();
        assert!(doc2.vector_score > 0.0);
        assert!(doc2.text_score > 0.0); // doc2 has highest text score, should be 1.0
        assert!(doc2.combined_score > 0.0);

        // doc2 should have the highest text score (normalized to 1.0)
        assert_eq!(doc2.text_score, 1.0);
    }

    #[test]
    fn test_hybrid_search_vector_only() {
        // MinMax with vector_weight=1.0 → combined_score == vector_score
        let searcher = HybridSearcher::new(HybridConfig {
            vector_weight: 1.0,
            require_both: false,
            fusion_strategy: FusionStrategy::MinMax,
        });

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.5, None, None),
        ];

        let text_results = vec![FullTextResult {
            doc_id: "doc1".to_string(),
            score: 1.0,
        }];

        let results = searcher.search(vector_results, text_results, 10);

        // doc1 should be first (highest vector score)
        assert_eq!(results[0].id, "doc1");
        assert_eq!(results[0].combined_score, results[0].vector_score);
    }

    #[test]
    fn test_hybrid_search_text_only() {
        // MinMax with vector_weight=0.0 → combined_score == text_score
        let searcher = HybridSearcher::new(HybridConfig {
            vector_weight: 0.0,
            require_both: false,
            fusion_strategy: FusionStrategy::MinMax,
        });

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.5, None, None),
        ];

        let text_results = vec![
            FullTextResult {
                doc_id: "doc1".to_string(),
                score: 1.0,
            },
            FullTextResult {
                doc_id: "doc2".to_string(),
                score: 3.0,
            },
        ];

        let results = searcher.search(vector_results, text_results, 10);

        // doc2 should be first (highest text score)
        assert_eq!(results[0].id, "doc2");
        assert_eq!(results[0].combined_score, results[0].text_score);
    }

    #[test]
    fn test_hybrid_search_require_both() {
        let searcher = HybridSearcher::new(HybridConfig {
            vector_weight: 0.5,
            require_both: true,
            ..Default::default()
        });

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.7, None, None),
        ];

        let text_results = vec![FullTextResult {
            doc_id: "doc1".to_string(),
            score: 2.0,
        }];

        let results = searcher.search(vector_results, text_results, 10);

        // Only doc1 should be in results (only one with both scores)
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "doc1");
    }

    #[test]
    fn test_hybrid_search_top_k() {
        let searcher = HybridSearcher::default();

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.8, None, None),
            ("doc3".to_string(), 0.7, None, None),
            ("doc4".to_string(), 0.6, None, None),
            ("doc5".to_string(), 0.5, None, None),
        ];

        let text_results = vec![];

        let results = searcher.search(vector_results, text_results, 3);

        assert_eq!(results.len(), 3);
    }

    #[test]
    fn test_hybrid_search_with_metadata() {
        let searcher = HybridSearcher::default();

        let metadata = serde_json::json!({"title": "Test Document"});
        let vector = vec![1.0, 0.0, 0.0];

        let vector_results = vec![(
            "doc1".to_string(),
            0.9,
            Some(metadata.clone()),
            Some(vector.clone()),
        )];

        let text_results = vec![FullTextResult {
            doc_id: "doc1".to_string(),
            score: 2.0,
        }];

        let results = searcher.search(vector_results, text_results, 10);

        assert_eq!(results.len(), 1);
        assert_eq!(results[0].metadata, Some(metadata));
        assert_eq!(results[0].vector, Some(vector));
    }

    #[test]
    fn test_normalize_score() {
        // Normal case
        assert_eq!(normalize_score(5.0, 0.0, 10.0), 0.5);
        assert_eq!(normalize_score(0.0, 0.0, 10.0), 0.0);
        assert_eq!(normalize_score(10.0, 0.0, 10.0), 1.0);

        // All same scores
        assert_eq!(normalize_score(5.0, 5.0, 5.0), 1.0);
    }

    #[test]
    fn test_hybrid_searcher_builder() {
        let searcher = HybridSearcher::default().with_vector_weight(0.7);

        assert_eq!(searcher.config.vector_weight, 0.7);
    }

    #[test]
    fn test_vector_weight_clamping() {
        let searcher1 = HybridSearcher::default().with_vector_weight(1.5);
        assert_eq!(searcher1.config.vector_weight, 1.0);

        let searcher2 = HybridSearcher::default().with_vector_weight(-0.5);
        assert_eq!(searcher2.config.vector_weight, 0.0);
    }

    // --- CE-14: RRF tests ---

    #[test]
    fn test_rrf_default_strategy() {
        // Default HybridSearcher uses RRF
        let searcher = HybridSearcher::default();
        assert_eq!(searcher.config.fusion_strategy, FusionStrategy::Rrf);
    }

    #[test]
    fn test_rrf_ranks_correctly() {
        // doc1: rank 1 in vector, rank 2 in text → 1/(60+1) + 1/(60+2) ≈ 0.03254
        // doc2: rank 2 in vector, rank 1 in text → 1/(60+2) + 1/(60+1) ≈ 0.03254  (equal)
        // doc3: rank 3 in vector, not in text    → 1/(60+3) ≈ 0.01587
        let searcher = HybridSearcher::default(); // RRF

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.7, None, None),
            ("doc3".to_string(), 0.5, None, None),
        ];

        let text_results = vec![
            FullTextResult {
                doc_id: "doc2".to_string(),
                score: 5.0,
            },
            FullTextResult {
                doc_id: "doc1".to_string(),
                score: 3.0,
            },
        ];

        let results = searcher.search(vector_results, text_results, 10);

        assert_eq!(results.len(), 3);

        // doc1 and doc2 both appear in both lists — they score higher than doc3 (vector-only)
        let doc3 = results.iter().find(|r| r.id == "doc3").unwrap();
        let doc1 = results.iter().find(|r| r.id == "doc1").unwrap();
        assert!(doc1.combined_score > doc3.combined_score);

        // Scores are positive
        for r in &results {
            assert!(r.combined_score > 0.0);
        }
    }

    #[test]
    fn test_rrf_require_both() {
        let searcher = HybridSearcher::new(HybridConfig {
            require_both: true,
            ..Default::default() // fusion_strategy: Rrf
        });

        let vector_results = vec![
            ("doc1".to_string(), 0.9, None, None),
            ("doc2".to_string(), 0.7, None, None),
        ];

        let text_results = vec![FullTextResult {
            doc_id: "doc1".to_string(),
            score: 2.0,
        }];

        let results = searcher.search(vector_results, text_results, 10);

        // Only doc1 appears in both lists
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "doc1");
    }

    #[test]
    fn test_rrf_formula_k60() {
        // Verify RRF formula: score = 1/(60 + rank), k=60
        let searcher = HybridSearcher::default();

        let vector_results = vec![("doc1".to_string(), 1.0, None, None)];
        let text_results = vec![FullTextResult {
            doc_id: "doc1".to_string(),
            score: 1.0,
        }];

        let results = searcher.search(vector_results, text_results, 10);

        assert_eq!(results.len(), 1);
        // doc1 is rank 1 in both → score = 1/(60+1) + 1/(60+1) = 2/61 ≈ 0.032787
        let expected = 2.0 / (RRF_K + 1.0);
        assert!((results[0].combined_score - expected).abs() < 1e-5);
    }

    #[test]
    fn test_with_fusion_strategy_builder() {
        let searcher = HybridSearcher::default().with_fusion_strategy(FusionStrategy::MinMax);
        assert_eq!(searcher.config.fusion_strategy, FusionStrategy::MinMax);
    }
}