web_search/
merger.rs

1//! Search result merger with reranking support
2
3use std::collections::{HashMap, HashSet};
4
5use crate::providers::SearchResult;
6
7/// Merge strategy for combining results
8#[derive(Debug, Clone, Copy, Default)]
9pub enum MergeStrategy {
10    /// Reciprocal Rank Fusion (default)
11    #[default]
12    Rrf,
13    /// Weighted scoring
14    Weighted,
15    /// Interleaved (round-robin)
16    Interleave,
17}
18
19/// Options for merging search results
20#[derive(Debug, Clone, Default)]
21pub struct MergeOptions {
22    /// Merge strategy to use
23    pub strategy: MergeStrategy,
24    /// Weights for each provider (provider name -> weight)
25    pub weights: HashMap<String, f64>,
26    /// RRF k parameter (default: 60)
27    pub rrf_k: Option<f64>,
28    /// Whether to remove duplicate URLs (default: true)
29    pub remove_duplicates: bool,
30}
31
32impl MergeOptions {
33    /// Create new merge options with default values
34    pub fn new() -> Self {
35        Self {
36            strategy: MergeStrategy::Rrf,
37            weights: HashMap::new(),
38            rrf_k: None,
39            remove_duplicates: true,
40        }
41    }
42
43    /// Set the merge strategy
44    pub fn with_strategy(mut self, strategy: MergeStrategy) -> Self {
45        self.strategy = strategy;
46        self
47    }
48
49    /// Set provider weights
50    pub fn with_weights(mut self, weights: HashMap<String, f64>) -> Self {
51        self.weights = weights;
52        self
53    }
54
55    /// Set the RRF k parameter
56    pub fn with_rrf_k(mut self, k: f64) -> Self {
57        self.rrf_k = Some(k);
58        self
59    }
60}
61
62/// Normalize URL for deduplication
63fn normalize_url(url: &str) -> String {
64    match url::Url::parse(url) {
65        Ok(parsed) => {
66            let mut normalized = format!("{}{}", parsed.host_str().unwrap_or(""), parsed.path());
67            normalized = normalized.trim_end_matches('/').to_lowercase();
68            normalized
69        }
70        Err(_) => url.to_lowercase(),
71    }
72}
73
74/// Calculate RRF score
75fn rrf_score(rank: usize, k: f64) -> f64 {
76    1.0 / (k + rank as f64)
77}
78
79/// Merge results using Reciprocal Rank Fusion
80pub fn merge_with_rrf(
81    results_by_provider: &HashMap<String, Vec<SearchResult>>,
82    options: &MergeOptions,
83) -> Vec<SearchResult> {
84    let k = options.rrf_k.unwrap_or(60.0);
85    let mut scores_by_url: HashMap<String, f64> = HashMap::new();
86    let mut results_by_url: HashMap<String, SearchResult> = HashMap::new();
87    let mut sources_by_url: HashMap<String, HashSet<String>> = HashMap::new();
88
89    for (provider, results) in results_by_provider {
90        let weight = options.weights.get(provider).copied().unwrap_or(1.0);
91
92        for result in results {
93            let normalized_url = normalize_url(&result.url);
94            let score = rrf_score(result.rank, k) * weight;
95
96            *scores_by_url.entry(normalized_url.clone()).or_insert(0.0) += score;
97
98            sources_by_url
99                .entry(normalized_url.clone())
100                .or_default()
101                .insert(result.source.clone());
102
103            results_by_url
104                .entry(normalized_url)
105                .or_insert_with(|| result.clone());
106        }
107    }
108
109    let mut merged: Vec<_> = scores_by_url
110        .into_iter()
111        .map(|(url, score)| {
112            let mut result = results_by_url.remove(&url).unwrap();
113            result.score = Some(score);
114            let sources: Vec<_> = sources_by_url
115                .get(&url)
116                .map(|s| s.iter().cloned().collect())
117                .unwrap_or_default();
118            if sources.len() > 1 {
119                result.sources = Some(sources);
120            }
121            (score, result)
122        })
123        .collect();
124
125    merged.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
126
127    merged
128        .into_iter()
129        .enumerate()
130        .map(|(i, (_, mut result))| {
131            result.rank = i + 1;
132            result
133        })
134        .collect()
135}
136
137/// Merge results using weighted scoring
138pub fn merge_with_weights(
139    results_by_provider: &HashMap<String, Vec<SearchResult>>,
140    options: &MergeOptions,
141) -> Vec<SearchResult> {
142    let max_rank = 100.0;
143    let mut scores_by_url: HashMap<String, f64> = HashMap::new();
144    let mut results_by_url: HashMap<String, SearchResult> = HashMap::new();
145    let mut sources_by_url: HashMap<String, HashSet<String>> = HashMap::new();
146
147    for (provider, results) in results_by_provider {
148        let weight = options.weights.get(provider).copied().unwrap_or(1.0);
149
150        for result in results {
151            let normalized_url = normalize_url(&result.url);
152            let score = ((max_rank - result.rank as f64 + 1.0) / max_rank) * weight;
153
154            *scores_by_url.entry(normalized_url.clone()).or_insert(0.0) += score;
155
156            sources_by_url
157                .entry(normalized_url.clone())
158                .or_default()
159                .insert(result.source.clone());
160
161            results_by_url
162                .entry(normalized_url)
163                .or_insert_with(|| result.clone());
164        }
165    }
166
167    let mut merged: Vec<_> = scores_by_url
168        .into_iter()
169        .map(|(url, score)| {
170            let mut result = results_by_url.remove(&url).unwrap();
171            result.score = Some(score);
172            let sources: Vec<_> = sources_by_url
173                .get(&url)
174                .map(|s| s.iter().cloned().collect())
175                .unwrap_or_default();
176            if sources.len() > 1 {
177                result.sources = Some(sources);
178            }
179            (score, result)
180        })
181        .collect();
182
183    merged.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
184
185    merged
186        .into_iter()
187        .enumerate()
188        .map(|(i, (_, mut result))| {
189            result.rank = i + 1;
190            result
191        })
192        .collect()
193}
194
195/// Merge results using interleaving (round-robin)
196pub fn merge_with_interleave(
197    results_by_provider: &HashMap<String, Vec<SearchResult>>,
198    options: &MergeOptions,
199) -> Vec<SearchResult> {
200    let mut results = Vec::new();
201    let mut seen_urls: HashSet<String> = HashSet::new();
202
203    let providers: Vec<_> = results_by_provider.keys().collect();
204    let max_len = results_by_provider
205        .values()
206        .map(|v| v.len())
207        .max()
208        .unwrap_or(0);
209
210    for i in 0..max_len {
211        for provider in &providers {
212            if let Some(provider_results) = results_by_provider.get(*provider) {
213                if i < provider_results.len() {
214                    let result = &provider_results[i];
215
216                    if options.remove_duplicates {
217                        let normalized = normalize_url(&result.url);
218                        if seen_urls.contains(&normalized) {
219                            continue;
220                        }
221                        seen_urls.insert(normalized);
222                    }
223
224                    let mut new_result = result.clone();
225                    new_result.rank = results.len() + 1;
226                    results.push(new_result);
227                }
228            }
229        }
230    }
231
232    results
233}
234
235/// Merge search results using the specified strategy
236pub fn merge_results(
237    results_by_provider: &HashMap<String, Vec<SearchResult>>,
238    options: &MergeOptions,
239) -> Vec<SearchResult> {
240    match options.strategy {
241        MergeStrategy::Rrf => merge_with_rrf(results_by_provider, options),
242        MergeStrategy::Weighted => merge_with_weights(results_by_provider, options),
243        MergeStrategy::Interleave => merge_with_interleave(results_by_provider, options),
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    fn create_test_result(url: &str, title: &str, source: &str, rank: usize) -> SearchResult {
252        SearchResult {
253            title: title.to_string(),
254            url: url.to_string(),
255            snippet: String::new(),
256            source: source.to_string(),
257            rank,
258            score: None,
259            sources: None,
260        }
261    }
262
263    #[test]
264    fn test_rrf_merge() {
265        let mut results_by_provider = HashMap::new();
266
267        results_by_provider.insert(
268            "google".to_string(),
269            vec![
270                create_test_result("https://example.com/1", "Result 1", "google", 1),
271                create_test_result("https://example.com/2", "Result 2", "google", 2),
272            ],
273        );
274
275        results_by_provider.insert(
276            "bing".to_string(),
277            vec![
278                create_test_result("https://example.com/2", "Result 2", "bing", 1),
279                create_test_result("https://example.com/3", "Result 3", "bing", 2),
280            ],
281        );
282
283        let options = MergeOptions::new();
284        let merged = merge_with_rrf(&results_by_provider, &options);
285
286        assert_eq!(merged.len(), 3);
287        assert!(merged[0].url.contains("example.com/2"));
288    }
289
290    #[test]
291    fn test_url_normalization() {
292        assert_eq!(
293            normalize_url("https://example.com/path/"),
294            normalize_url("https://example.com/path")
295        );
296        assert_eq!(
297            normalize_url("https://Example.COM/Path"),
298            normalize_url("https://example.com/path")
299        );
300    }
301}
web_search/merger.rs

web_search/
merger.rs
