1use lc_vector_stores::Document;
7use std::collections::HashMap;
8
9pub const RRF_K: usize = 60;
11
12fn doc_content_hash(doc: &Document) -> String {
19 use std::hash::{Hash, Hasher};
20 let mut hasher = fnv::FnvHasher::default();
21 doc.content.hash(&mut hasher);
22 format!("{:016x}", hasher.finish())
23}
24
25#[derive(Debug, Clone)]
27pub struct RetrievedDocument {
28 pub document: Document,
30 pub score: f64,
32 pub source: RetrievalSource,
34}
35
36#[derive(Debug, Clone, Copy, PartialEq)]
38pub enum RetrievalSource {
39 BM25,
41 Vector,
43 Hybrid,
45}
46
47pub fn filter_by_score<T, S: PartialOrd>(scored: Vec<(T, S)>, min_score: S) -> Vec<(T, S)> {
55 scored.into_iter().filter(|(_, s)| *s > min_score).collect()
56}
57
58pub fn reciprocal_rank_fusion(
70 bm25_results: Vec<Document>,
71 vector_results: Vec<Document>,
72 k: usize,
73) -> Vec<RetrievedDocument> {
74 let mut rrf_scores: HashMap<String, (f64, Document)> = HashMap::new();
75
76 for (rank, doc) in bm25_results.iter().enumerate() {
78 let doc_id = doc.id.clone().unwrap_or_else(|| doc_content_hash(doc));
79 let rrf_contribution = 1.0 / (k as f64 + (rank + 1) as f64);
80
81 rrf_scores
82 .entry(doc_id.clone())
83 .and_modify(|(score, _existing_doc)| {
84 *score += rrf_contribution;
85 })
86 .or_insert((rrf_contribution, doc.clone()));
87 }
88
89 for (rank, doc) in vector_results.iter().enumerate() {
91 let doc_id = doc.id.clone().unwrap_or_else(|| doc_content_hash(doc));
92 let rrf_contribution = 1.0 / (k as f64 + (rank + 1) as f64);
93
94 rrf_scores
95 .entry(doc_id.clone())
96 .and_modify(|(score, _)| {
97 *score += rrf_contribution;
98 })
99 .or_insert((rrf_contribution, doc.clone()));
100 }
101
102 let mut results: Vec<RetrievedDocument> = rrf_scores
104 .into_iter()
105 .map(|(_, (score, doc))| RetrievedDocument {
106 document: doc,
107 score,
108 source: RetrievalSource::Hybrid,
109 })
110 .collect();
111
112 results.sort_by(|a, b| {
113 b.score
114 .partial_cmp(&a.score)
115 .unwrap_or(std::cmp::Ordering::Equal)
116 });
117
118 results
119}
120
121#[cfg(test)]
122mod tests {
123 use super::*;
124
125 #[test]
126 fn test_rrf_basic() {
127 let bm25_docs = vec![
128 Document::new("Rust系统编程").with_id("doc1"),
129 Document::new("Python数据科学").with_id("doc2"),
130 Document::new("Go并发编程").with_id("doc3"),
131 ];
132
133 let vector_docs = vec![
134 Document::new("Rust系统编程").with_id("doc1"),
135 Document::new("JavaScript前端").with_id("doc4"),
136 Document::new("Python数据科学").with_id("doc2"),
137 ];
138
139 let results = reciprocal_rank_fusion(bm25_docs, vector_docs, 60);
140
141 println!("RRF 融合结果:");
142 for (i, r) in results.iter().enumerate() {
143 println!(
144 " [{}] doc_id={}, score={:.4}",
145 i,
146 r.document.id.clone().unwrap_or_default(),
147 r.score
148 );
149 }
150
151 let first_doc_id = results[0].document.id.clone().unwrap_or_default();
153 println!("最高分文档: {}", first_doc_id);
154 }
155
156 #[test]
160 fn test_filter_by_score() {
161 let scored = vec![("a", 0.9_f32), ("b", 0.2), ("c", -0.3), ("d", 0.0)];
162
163 let filtered = filter_by_score(scored.clone(), 0.0);
165 let ids: Vec<&str> = filtered.iter().map(|(id, _)| *id).collect();
166 assert_eq!(ids, vec!["a", "b"]);
167
168 let relaxed = filter_by_score(scored.clone(), -0.5);
170 assert_eq!(relaxed.len(), 4);
171
172 let strict = filter_by_score(scored.clone(), 0.5);
174 let ids: Vec<&str> = strict.iter().map(|(id, _)| *id).collect();
175 assert_eq!(ids, vec!["a"]);
176 }
177
178 #[test]
180 fn test_filter_by_score_f64() {
181 let scored = vec![("x", 0.8_f64), ("y", 0.0), ("z", -0.5)];
182 let filtered = filter_by_score(scored, 0.0);
183 let ids: Vec<&str> = filtered.iter().map(|(id, _)| *id).collect();
184 assert_eq!(ids, vec!["x"]);
185 }
186
187 #[test]
191 fn test_doc_content_hash_stable() {
192 let content = "Rust 系统编程与并发";
193 let doc_a = Document::new(content.to_string());
194 let doc_b = Document::new(content.to_string());
195 let doc_c = Document::new("Python 数据科学");
196
197 let hash_a1 = doc_content_hash(&doc_a);
198 let hash_a2 = doc_content_hash(&doc_b);
199 assert_eq!(hash_a1, hash_a2, "相同内容应产生相同哈希");
200
201 let hash_c = doc_content_hash(&doc_c);
202 assert_ne!(hash_a1, hash_c, "不同内容应产生不同哈希");
203 assert_eq!(hash_a1.len(), 16, "应为 64 位哈希的 16 位十六进制表示");
204 }
205}