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
//!
//! The final score is computed as:
//! `score = vector_weight * vector_score + (1 - vector_weight) * text_score`

use std::collections::HashMap;

use crate::fulltext::FullTextResult;

/// 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)
    pub vector_weight: f32,
    /// Whether to require matches in both indices
    pub require_both: bool,
}

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

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

    /// Combine vector search results with full-text search results
    ///
    /// # 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
    ///
    /// # Returns
    /// Combined and re-ranked results
    pub fn search(
        &self,
        vector_results: Vec<VectorResultRow>,
        text_results: Vec<FullTextResult>,
        top_k: usize,
    ) -> Vec<HybridResult> {
        // Collect all unique IDs and their raw scores
        let mut vector_scores: HashMap<String, RawScore> = HashMap::new();
        let mut text_scores: HashMap<String, f32> = HashMap::new();

        // Track min/max for normalization
        let mut vector_min = f32::MAX;
        let mut vector_max = f32::MIN;
        let mut text_min = f32::MAX;
        let mut text_max = f32::MIN;

        // Collect vector scores
        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,
                },
            );
        }

        // Collect text scores
        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);
        }

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

        // Compute combined scores
        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);

            // Skip if require_both and missing one
            if self.config.require_both && (vector_raw.is_none() || text_raw.is_none()) {
                continue;
            }

            // Normalize scores to 0-1 range
            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
            };

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

            // Get metadata and vector from vector results if available
            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,
            });
        }

        // Sort by combined score (descending)
        results.sort_by(|a, b| {
            b.combined_score
                .partial_cmp(&a.combined_score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Return top-k
        results.truncate(top_k);
        results
    }
}

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

/// 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() {
        let searcher = HybridSearcher::new(HybridConfig {
            vector_weight: 1.0,
            require_both: false,
        });

        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() {
        let searcher = HybridSearcher::new(HybridConfig {
            vector_weight: 0.0,
            require_both: false,
        });

        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,
        });

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