khive-retrieval 0.2.3

Hybrid retrieval composer (HNSW + BM25 + fusion + graph + cross-encoder) with deterministic scoring
Documentation
//! Tunable hybrid search configuration — Brain Phase 7 substrate.
//!
//! `SearchConfig` is a per-call configuration that controls how vector and
//! keyword results are retrieved and fused. It is the public API surface for
//! `recall()` and compose's internal search phase.
//!
//! # Defaults (backward-compatible)
//!
//! `SearchConfig::default()` is designed to produce **identical results** to
//! the pre-Phase-7 hardcoded search behavior: RRF with k=60, top_k=10,
//! no min_score filter. Existing callers that do not supply a `SearchConfig`
//! get the same behavior as before.
//!
//! # Presets
//!
//! | Preset | Strategy | vector_weight |
//! |--------|----------|---------------|
//! | `default()` / `hybrid_balanced()` | RRF (k=60) | 0.5 |
//! | `vector_only()` | VectorOnly | 1.0 |
//! | `keyword_only()` | KeywordOnly | 0.0 |
//!
//! # Usage in recall
//!
//! ```rust,ignore
//! let opts = RecallOptions {
//!     query: "metal inference kernel".to_string(),
//!     search: Some(SearchConfig::vector_only()),
//!     ..Default::default()
//! };
//! service.recall(opts).await?;
//! ```

use serde::{Deserialize, Serialize};

use khive_fusion::{FusionStrategy, DEFAULT_RRF_K};

/// Per-call configuration for hybrid search retrieval and fusion.
///
/// Added to `RecallOptions` and `ComposeOptions` as `search: Option<SearchConfig>`.
/// When `None`, callers receive identical behavior to pre-Phase-7 code.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct SearchConfig {
    /// Maximum number of results to return.
    ///
    /// Default: 10.
    #[serde(default = "default_top_k")]
    pub top_k: usize,

    /// Candidate pool multiplier over `top_k`.
    ///
    /// The retriever fetches `top_k * candidate_pool_multiplier` candidates
    /// before fusion and reranking. Higher values improve recall quality at
    /// the cost of more computation.
    ///
    /// Default: 3.
    #[serde(default = "default_multiplier")]
    pub candidate_pool_multiplier: usize,

    /// Fusion strategy for combining vector and keyword result lists.
    ///
    /// Default: RRF with k=60.
    #[serde(default = "default_fusion")]
    pub fusion_strategy: FusionStrategy,

    /// Weight for vector search in weighted fusion (0.0 to 1.0).
    ///
    /// Only used when `fusion_strategy` is `Weighted`. Keyword weight is
    /// implicitly `1.0 - vector_weight`.
    ///
    /// Default: 0.5 (balanced).
    #[serde(default = "default_vector_weight")]
    pub vector_weight: f64,

    /// Minimum score threshold.
    ///
    /// Results with a final score below this value are filtered out.
    /// When `None`, no threshold is applied.
    ///
    /// Default: None.
    #[serde(default)]
    pub min_score: Option<f64>,
}

fn default_top_k() -> usize {
    10
}

fn default_multiplier() -> usize {
    3
}

fn default_fusion() -> FusionStrategy {
    FusionStrategy::Rrf { k: DEFAULT_RRF_K }
}

fn default_vector_weight() -> f64 {
    0.5
}

impl Default for SearchConfig {
    fn default() -> Self {
        Self {
            top_k: default_top_k(),
            candidate_pool_multiplier: default_multiplier(),
            fusion_strategy: default_fusion(),
            vector_weight: default_vector_weight(),
            min_score: None,
        }
    }
}

impl SearchConfig {
    /// Preset: skip BM25 entirely, return only vector search results.
    ///
    /// Use when keyword search degrades quality (e.g., short queries, code search).
    pub fn vector_only() -> Self {
        Self {
            top_k: default_top_k(),
            candidate_pool_multiplier: 1,
            fusion_strategy: FusionStrategy::VectorOnly,
            vector_weight: 1.0,
            min_score: None,
        }
    }

    /// Preset: skip HNSW entirely, return only BM25 keyword results.
    ///
    /// Use for exact-match retrieval (e.g., medication names, identifiers).
    pub fn keyword_only() -> Self {
        Self {
            top_k: default_top_k(),
            candidate_pool_multiplier: 1,
            fusion_strategy: FusionStrategy::KeywordOnly,
            vector_weight: 0.0,
            min_score: None,
        }
    }

    /// Preset: balanced hybrid search using RRF with k=60.
    ///
    /// Equivalent to `SearchConfig::default()`. Combines vector and keyword
    /// results with equal weight using Reciprocal Rank Fusion.
    pub fn hybrid_balanced() -> Self {
        Self::default()
    }

    /// Set a custom top_k.
    #[must_use]
    pub fn with_top_k(mut self, top_k: usize) -> Self {
        self.top_k = top_k;
        self
    }

    /// Set the candidate pool multiplier.
    #[must_use]
    pub fn with_candidate_pool_multiplier(mut self, multiplier: usize) -> Self {
        self.candidate_pool_multiplier = multiplier;
        self
    }

    /// Set a minimum score filter.
    #[must_use]
    pub fn with_min_score(mut self, min: f64) -> Self {
        self.min_score = Some(min);
        self
    }

    /// Compute the candidate pool size from `top_k * candidate_pool_multiplier`.
    pub fn candidate_pool_size(&self) -> usize {
        self.top_k * self.candidate_pool_multiplier.max(1)
    }
}

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

    #[test]
    fn test_default_config() {
        let cfg = SearchConfig::default();
        assert_eq!(cfg.top_k, 10);
        assert_eq!(cfg.candidate_pool_multiplier, 3);
        assert!((cfg.vector_weight - 0.5).abs() < f64::EPSILON);
        assert!(cfg.min_score.is_none());
        assert!(matches!(cfg.fusion_strategy, FusionStrategy::Rrf { k: 60 }));
    }

    #[test]
    fn test_vector_only_preset() {
        let cfg = SearchConfig::vector_only();
        assert!(matches!(cfg.fusion_strategy, FusionStrategy::VectorOnly));
        assert!((cfg.vector_weight - 1.0).abs() < f64::EPSILON);
        assert_eq!(cfg.candidate_pool_multiplier, 1);
    }

    #[test]
    fn test_keyword_only_preset() {
        let cfg = SearchConfig::keyword_only();
        assert!(matches!(cfg.fusion_strategy, FusionStrategy::KeywordOnly));
        assert!((cfg.vector_weight - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_hybrid_balanced_is_default() {
        let balanced = SearchConfig::hybrid_balanced();
        let default = SearchConfig::default();
        assert_eq!(balanced.top_k, default.top_k);
        assert_eq!(
            balanced.candidate_pool_multiplier,
            default.candidate_pool_multiplier
        );
        assert!((balanced.vector_weight - default.vector_weight).abs() < f64::EPSILON);
    }

    #[test]
    fn test_candidate_pool_size() {
        let cfg = SearchConfig::default();
        assert_eq!(cfg.candidate_pool_size(), 30); // 10 * 3

        let cfg = SearchConfig::vector_only().with_top_k(5);
        assert_eq!(cfg.candidate_pool_size(), 5); // 5 * 1
    }

    #[test]
    fn test_builder_methods() {
        let cfg = SearchConfig::default()
            .with_top_k(20)
            .with_candidate_pool_multiplier(5)
            .with_min_score(0.3);
        assert_eq!(cfg.top_k, 20);
        assert_eq!(cfg.candidate_pool_multiplier, 5);
        assert_eq!(cfg.min_score, Some(0.3));
        assert_eq!(cfg.candidate_pool_size(), 100);
    }

    #[test]
    fn test_serde_roundtrip() {
        let cfg = SearchConfig {
            top_k: 15,
            candidate_pool_multiplier: 4,
            fusion_strategy: FusionStrategy::Weighted {
                weights: vec![0.7, 0.3],
            },
            vector_weight: 0.7,
            min_score: Some(0.1),
        };
        let json = serde_json::to_string(&cfg).unwrap();
        let back: SearchConfig = serde_json::from_str(&json).unwrap();
        assert_eq!(back.top_k, 15);
        assert_eq!(back.candidate_pool_multiplier, 4);
        assert_eq!(back.min_score, Some(0.1));
    }
}