Skip to main content

khive_fusion/
strategy.rs

1//! Fusion strategy types with invariant validation.
2
3use serde::{Deserialize, Serialize};
4use std::fmt;
5
6/// Default RRF constant k=60, standard in literature (Cormack et al., 2009).
7pub const DEFAULT_RRF_K: usize = 60;
8
9/// Error returned when a [`FusionStrategy`] fails invariant validation.
10#[derive(Debug, Clone, PartialEq)]
11pub enum FusionStrategyError {
12    /// RRF k must be >= 1 to avoid division by zero.
13    RrfKZero,
14    /// Weighted strategy weights contain NaN.
15    WeightNaN,
16    /// Weighted strategy weights contain infinity.
17    WeightInfinite,
18    /// Weighted RRF requires every weight to be strictly positive.
19    WeightNotPositive,
20    /// Custom strategy name must not be empty.
21    CustomNameEmpty,
22}
23
24impl fmt::Display for FusionStrategyError {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        match self {
27            Self::RrfKZero => write!(f, "Rrf k must be >= 1"),
28            Self::WeightNaN => write!(f, "Weighted weights must not contain NaN"),
29            Self::WeightInfinite => write!(f, "Weighted weights must not contain infinity"),
30            Self::WeightNotPositive => write!(f, "WeightedRrf weights must be strictly positive"),
31            Self::CustomNameEmpty => write!(f, "Custom strategy name must not be empty"),
32        }
33    }
34}
35
36impl std::error::Error for FusionStrategyError {}
37
38/// Raw serde form used for deserialization before validation.
39#[derive(Deserialize)]
40#[serde(rename_all = "snake_case")]
41enum RawFusionStrategy {
42    #[serde(alias = "Rrf")]
43    Rrf { k: usize },
44    #[serde(alias = "Weighted")]
45    Weighted { weights: Vec<f64> },
46    #[serde(alias = "WeightedRrf")]
47    WeightedRrf { k: usize, weights: Vec<f64> },
48    #[serde(alias = "Union")]
49    Union,
50    #[serde(alias = "VectorOnly")]
51    VectorOnly,
52    #[serde(alias = "KeywordOnly")]
53    KeywordOnly,
54    #[serde(alias = "Custom")]
55    Custom {
56        name: String,
57        params: serde_json::Value,
58    },
59}
60
61impl TryFrom<RawFusionStrategy> for FusionStrategy {
62    type Error = FusionStrategyError;
63
64    fn try_from(raw: RawFusionStrategy) -> Result<Self, Self::Error> {
65        match raw {
66            RawFusionStrategy::Rrf { k } => FusionStrategy::try_rrf(k),
67            RawFusionStrategy::Weighted { weights } => FusionStrategy::try_weighted(weights),
68            RawFusionStrategy::WeightedRrf { k, weights } => {
69                FusionStrategy::try_weighted_rrf(k, weights)
70            }
71            RawFusionStrategy::Union => Ok(FusionStrategy::Union),
72            RawFusionStrategy::VectorOnly => Ok(FusionStrategy::VectorOnly),
73            RawFusionStrategy::KeywordOnly => Ok(FusionStrategy::KeywordOnly),
74            RawFusionStrategy::Custom { name, params } => FusionStrategy::try_custom(name, params),
75        }
76    }
77}
78
79/// Validated selection for combining ranked result lists.
80///
81/// See `crates/khive-fusion/docs/api/strategy-validation.md`.
82#[derive(Debug, Clone, PartialEq, Serialize)]
83#[serde(rename_all = "snake_case")]
84#[serde(try_from = "RawFusionStrategy")]
85pub enum FusionStrategy {
86    /// Reciprocal Rank Fusion (default, recommended). Rank-based, distribution-agnostic.
87    Rrf {
88        /// Smoothing constant (>= 1). Default: 60.
89        k: usize,
90    },
91
92    /// Weighted linear combination of scores. Weights normalized to 1.0; must be finite.
93    Weighted {
94        /// Weights for each source (will be normalized). Must be finite.
95        weights: Vec<f64>,
96    },
97
98    /// Weighted Reciprocal Rank Fusion over ordered sources. Weights are not normalized.
99    WeightedRrf {
100        /// Smoothing constant (>= 1).
101        k: usize,
102        /// One finite, strictly positive weight for each ordered source.
103        weights: Vec<f64>,
104    },
105
106    /// Take union with max score per ID.
107    Union,
108
109    /// Skip BM25 entirely -- return only vector (HNSW) results.
110    VectorOnly,
111
112    /// Skip HNSW entirely -- return only BM25 keyword results.
113    KeywordOnly,
114
115    /// Pack-defined or user-defined custom strategy dispatched by name at runtime.
116    Custom {
117        /// Strategy identifier registered with the fusion executor registry.
118        name: String,
119        /// Opaque parameters consumed by the executor.
120        params: serde_json::Value,
121    },
122}
123
124impl<'de> Deserialize<'de> for FusionStrategy {
125    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
126    where
127        D: serde::Deserializer<'de>,
128    {
129        let raw = RawFusionStrategy::deserialize(deserializer)?;
130        FusionStrategy::try_from(raw).map_err(serde::de::Error::custom)
131    }
132}
133
134impl Default for FusionStrategy {
135    fn default() -> Self {
136        Self::Rrf { k: DEFAULT_RRF_K }
137    }
138}
139
140impl FusionStrategy {
141    /// Create an RRF strategy with default k=60.
142    #[inline]
143    pub fn rrf() -> Self {
144        Self::Rrf { k: DEFAULT_RRF_K }
145    }
146
147    /// Create an RRF strategy, returning [`FusionStrategyError::RrfKZero`] for zero.
148    #[inline]
149    pub fn try_rrf(k: usize) -> Result<Self, FusionStrategyError> {
150        if k == 0 {
151            return Err(FusionStrategyError::RrfKZero);
152        }
153        Ok(Self::Rrf { k })
154    }
155
156    /// Create an RRF strategy, clamping k to at least 1.
157    ///
158    /// Prefer [`try_rrf`](Self::try_rrf) at public API boundaries.
159    #[inline]
160    pub fn rrf_with_k(k: usize) -> Self {
161        Self::Rrf { k: k.max(1) }
162    }
163
164    /// Create weighted fusion, returning an error for NaN or infinite weights.
165    pub fn try_weighted(weights: Vec<f64>) -> Result<Self, FusionStrategyError> {
166        for w in &weights {
167            if w.is_nan() {
168                return Err(FusionStrategyError::WeightNaN);
169            }
170            if w.is_infinite() {
171                return Err(FusionStrategyError::WeightInfinite);
172            }
173        }
174        Ok(Self::Weighted { weights })
175    }
176
177    /// Create a weighted strategy. Panics on NaN/infinity.
178    ///
179    /// Prefer [`try_weighted`](Self::try_weighted) at public API boundaries.
180    #[inline]
181    pub fn weighted(weights: Vec<f64>) -> Self {
182        Self::try_weighted(weights).expect("weights must be finite")
183    }
184
185    /// Create weighted reciprocal rank fusion, rejecting zero `k` and weights that are
186    /// non-finite or not strictly positive. Source-count validation happens when the strategy is
187    /// fused because a strategy does not carry its source lists.
188    pub fn try_weighted_rrf(k: usize, weights: Vec<f64>) -> Result<Self, FusionStrategyError> {
189        if k == 0 {
190            return Err(FusionStrategyError::RrfKZero);
191        }
192        for weight in &weights {
193            if weight.is_nan() {
194                return Err(FusionStrategyError::WeightNaN);
195            }
196            if weight.is_infinite() {
197                return Err(FusionStrategyError::WeightInfinite);
198            }
199            if *weight <= 0.0 {
200                return Err(FusionStrategyError::WeightNotPositive);
201            }
202        }
203        Ok(Self::WeightedRrf { k, weights })
204    }
205
206    /// Create weighted reciprocal rank fusion. Panics on zero `k` or invalid weights.
207    ///
208    /// Prefer [`try_weighted_rrf`](Self::try_weighted_rrf) at public API boundaries.
209    #[inline]
210    pub fn weighted_rrf(k: usize, weights: Vec<f64>) -> Self {
211        Self::try_weighted_rrf(k, weights)
212            .expect("weighted RRF requires k >= 1 and positive finite weights")
213    }
214
215    /// Create a union strategy.
216    #[inline]
217    pub fn union() -> Self {
218        Self::Union
219    }
220
221    /// Create a runtime-dispatched custom strategy, rejecting an empty name.
222    pub fn try_custom(
223        name: String,
224        params: serde_json::Value,
225    ) -> Result<Self, FusionStrategyError> {
226        if name.is_empty() {
227            return Err(FusionStrategyError::CustomNameEmpty);
228        }
229        Ok(Self::Custom { name, params })
230    }
231}
232
233#[cfg(test)]
234#[path = "strategy_tests.rs"]
235mod tests;