Skip to main content

lc_rag/
hybrid.rs

1// src/retrieval/hybrid.rs
2//! Hybrid retrieval module
3//!
4//! Combines BM25 keyword retrieval + vector semantic retrieval
5
6use lc_vector_stores::Document;
7use std::collections::HashMap;
8
9/// The k parameter in the RRF fusion algorithm (default 60).
10pub const RRF_K: usize = 60;
11
12/// Generate a stable document ID from content hash to avoid collisions (H46).
13///
14/// P2-3: Replaces `DefaultHasher` with FNV-1a 64-bit. `DefaultHasher`'s algorithm is an
15/// internal std implementation detail, not guaranteed stable across processes/versions;
16/// FNV-1a is a fully-specified deterministic hash, so when `doc.id` is missing, fusion
17/// dedup does not drift across processes/versions.
18fn 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/// Retrieval result (with score)
26#[derive(Debug, Clone)]
27pub struct RetrievedDocument {
28    /// Document content
29    pub document: Document,
30    /// Fused score
31    pub score: f64,
32    /// Retrieval source (BM25 / vector / hybrid)
33    pub source: RetrievalSource,
34}
35
36/// Retrieval source
37#[derive(Debug, Clone, Copy, PartialEq)]
38pub enum RetrievalSource {
39    /// From BM25 keyword retrieval
40    BM25,
41    /// From vector semantic retrieval
42    Vector,
43    /// From RRF fusion
44    Hybrid,
45}
46
47/// Filters retrieval results by a minimum score (P1-2).
48///
49/// Eliminates the duplicated `score > 0.0` ghost-threshold implementation across
50/// `unified_hybrid` / `graph_rag::matcher`. `min_score` is compared on the **raw score scale**:
51/// the default 0.0 keeps the old behavior (only positive similarities are kept). Cosine
52/// similarity ranges over [-1, 1]; with non-normalized embedding models the cosine of a
53/// relevant document can be negative, so different models may lower the threshold.
54pub 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
58/// RRF fusion algorithm
59///
60/// Formula: RRF_score(d) = Σ 1/(k + rank(d))
61///
62/// Arguments:
63/// - bm25_results: BM25 retrieval results, sorted by score descending
64/// - vector_results: vector retrieval results, sorted by similarity descending
65/// - k: the RRF parameter, usually 60
66///
67/// Returns:
68/// - the fused document list, sorted by RRF score descending
69pub 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    // Process BM25 results
77    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    // Process vector results
90    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    // Sort by RRF score
103    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        // doc1 appears in both lists, so its score should be highest
152        let first_doc_id = results[0].document.id.clone().unwrap_or_default();
153        println!("最高分文档: {}", first_doc_id);
154    }
155
156    /// P1-2: Shares the filter_by_score utility — the default 0.0 keeps only positive
157    /// similarities; lowering the threshold keeps negative-similarity documents (with
158    /// non-normalized embedding models a relevant document's cosine can be negative).
159    #[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        // Default threshold 0.0: keep only strictly-greater scores (matches the old `score > 0.0` behavior)
164        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        // Lowering the threshold keeps negative similarities
169        let relaxed = filter_by_score(scored.clone(), -0.5);
170        assert_eq!(relaxed.len(), 4);
171
172        // Raising the threshold is stricter
173        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    /// P1-2: filter_by_score also works for f64 scores.
179    #[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    /// P2-3: `doc_content_hash` is a deterministic hash — the same content yields the same
188    /// result across calls, different content yields different results. FNV-1a is fully
189    /// specified and does not drift across processes/versions.
190    #[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}