1use serde::{Deserialize, Serialize};
4use std::fmt;
5
6pub const DEFAULT_RRF_K: usize = 60;
8
9#[derive(Debug, Clone, PartialEq)]
11pub enum FusionStrategyError {
12 RrfKZero,
14 WeightNaN,
16 WeightInfinite,
18 WeightNotPositive,
20 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#[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#[derive(Debug, Clone, PartialEq, Serialize)]
83#[serde(rename_all = "snake_case")]
84#[serde(try_from = "RawFusionStrategy")]
85pub enum FusionStrategy {
86 Rrf {
88 k: usize,
90 },
91
92 Weighted {
94 weights: Vec<f64>,
96 },
97
98 WeightedRrf {
100 k: usize,
102 weights: Vec<f64>,
104 },
105
106 Union,
108
109 VectorOnly,
111
112 KeywordOnly,
114
115 Custom {
117 name: String,
119 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 #[inline]
143 pub fn rrf() -> Self {
144 Self::Rrf { k: DEFAULT_RRF_K }
145 }
146
147 #[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 #[inline]
160 pub fn rrf_with_k(k: usize) -> Self {
161 Self::Rrf { k: k.max(1) }
162 }
163
164 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 #[inline]
181 pub fn weighted(weights: Vec<f64>) -> Self {
182 Self::try_weighted(weights).expect("weights must be finite")
183 }
184
185 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 #[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 #[inline]
217 pub fn union() -> Self {
218 Self::Union
219 }
220
221 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;