lc_rag/
hybrid.rs

1// src/retrieval/hybrid.rs
2//! 混合检索模块
3//!
4//! 结合 BM25 关键词检索 + 向量语义检索
5
6use lc_vector_stores::Document;
7use std::collections::HashMap;
8
9pub const RRF_K: usize = 60;
10
11/// Generate a stable document ID from content hash to avoid collisions (H46).
12///
13/// P2-3: 用 FNV-1a 64 替代 `DefaultHasher`。`DefaultHasher` 的算法是 std 内部
14/// 实现细节,不保证跨进程/跨版本稳定;FNV-1a 是完全指定的确定性哈希,同一
15/// 内容的 `doc.id` 缺失时融合去重不会漂移。
16fn doc_content_hash(doc: &Document) -> String {
17    use std::hash::{Hash, Hasher};
18    let mut hasher = fnv::FnvHasher::default();
19    doc.content.hash(&mut hasher);
20    format!("{:016x}", hasher.finish())
21}
22
23/// 检索结果（带分数）
24#[derive(Debug, Clone)]
25pub struct RetrievedDocument {
26    pub document: Document,
27    pub score: f64,
28    pub source: RetrievalSource,
29}
30
31/// 检索来源
32#[derive(Debug, Clone, Copy, PartialEq)]
33pub enum RetrievalSource {
34    BM25,
35    Vector,
36    Hybrid,
37}
38
39/// 按最小分数过滤检索结果(P1-2)。
40///
41/// 消除 `score > 0.0` 幽灵阈值在 `unified_hybrid` / `graph_rag::matcher`
42/// 两处的重复实现。`min_score` 在**原始分数尺度**上比较:
43/// 默认 0.0 保持旧行为(只保留正相似度)。余弦相似度范围 [-1,1],
44/// 非归一化嵌入模型下相关文档的余弦可能为负,不同模型可自行调低阈值。
45pub fn filter_by_score<T, S: PartialOrd>(scored: Vec<(T, S)>, min_score: S) -> Vec<(T, S)> {
46    scored.into_iter().filter(|(_, s)| *s > min_score).collect()
47}
48
49/// RRF 融合算法
50///
51/// 公式: RRF_score(d) = Σ 1/(k + rank(d))
52///
53/// 参数:
54/// - bm25_results: BM25 检索结果，按分数降序排列
55/// - vector_results: 向量检索结果，按相似度降序排列
56/// - k: RRF 参数，通常为 60
57///
58/// 返回:
59/// - 融合后的文档列表，按 RRF 分数降序排列
60pub fn reciprocal_rank_fusion(
61    bm25_results: Vec<Document>,
62    vector_results: Vec<Document>,
63    k: usize,
64) -> Vec<RetrievedDocument> {
65    let mut rrf_scores: HashMap<String, (f64, Document)> = HashMap::new();
66
67    // BM25 结果处理
68    for (rank, doc) in bm25_results.iter().enumerate() {
69        let doc_id = doc.id.clone().unwrap_or_else(|| doc_content_hash(doc));
70        let rrf_contribution = 1.0 / (k as f64 + (rank + 1) as f64);
71
72        rrf_scores
73            .entry(doc_id.clone())
74            .and_modify(|(score, _existing_doc)| {
75                *score += rrf_contribution;
76            })
77            .or_insert((rrf_contribution, doc.clone()));
78    }
79
80    // 向量结果处理
81    for (rank, doc) in vector_results.iter().enumerate() {
82        let doc_id = doc.id.clone().unwrap_or_else(|| doc_content_hash(doc));
83        let rrf_contribution = 1.0 / (k as f64 + (rank + 1) as f64);
84
85        rrf_scores
86            .entry(doc_id.clone())
87            .and_modify(|(score, _)| {
88                *score += rrf_contribution;
89            })
90            .or_insert((rrf_contribution, doc.clone()));
91    }
92
93    // 按 RRF 分数排序
94    let mut results: Vec<RetrievedDocument> = rrf_scores
95        .into_iter()
96        .map(|(_, (score, doc))| RetrievedDocument {
97            document: doc,
98            score,
99            source: RetrievalSource::Hybrid,
100        })
101        .collect();
102
103    results.sort_by(|a, b| {
104        b.score
105            .partial_cmp(&a.score)
106            .unwrap_or(std::cmp::Ordering::Equal)
107    });
108
109    results
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn test_rrf_basic() {
118        let bm25_docs = vec![
119            Document::new("Rust系统编程").with_id("doc1"),
120            Document::new("Python数据科学").with_id("doc2"),
121            Document::new("Go并发编程").with_id("doc3"),
122        ];
123
124        let vector_docs = vec![
125            Document::new("Rust系统编程").with_id("doc1"),
126            Document::new("JavaScript前端").with_id("doc4"),
127            Document::new("Python数据科学").with_id("doc2"),
128        ];
129
130        let results = reciprocal_rank_fusion(bm25_docs, vector_docs, 60);
131
132        println!("RRF 融合结果:");
133        for (i, r) in results.iter().enumerate() {
134            println!(
135                "  [{}] doc_id={}, score={:.4}",
136                i,
137                r.document.id.clone().unwrap_or_default(),
138                r.score
139            );
140        }
141
142        // doc1 在两个列表都出现，分数应该最高
143        let first_doc_id = results[0].document.id.clone().unwrap_or_default();
144        println!("最高分文档: {}", first_doc_id);
145    }
146
147    /// P1-2: 共享 filter_by_score 工具函数——默认 0.0 只保留正相似度,
148    /// 调低阈值可保留负相似度文档(非归一化嵌入模型下相关文档余弦可为负)。
149    #[test]
150    fn test_filter_by_score() {
151        let scored = vec![("a", 0.9_f32), ("b", 0.2), ("c", -0.3), ("d", 0.0)];
152
153        // 默认阈值 0.0: 严格大于才保留(与旧 `score > 0.0` 行为一致)
154        let filtered = filter_by_score(scored.clone(), 0.0);
155        let ids: Vec<&str> = filtered.iter().map(|(id, _)| *id).collect();
156        assert_eq!(ids, vec!["a", "b"]);
157
158        // 调低阈值可保留负相似度
159        let relaxed = filter_by_score(scored.clone(), -0.5);
160        assert_eq!(relaxed.len(), 4);
161
162        // 调高阈值更严格
163        let strict = filter_by_score(scored.clone(), 0.5);
164        let ids: Vec<&str> = strict.iter().map(|(id, _)| *id).collect();
165        assert_eq!(ids, vec!["a"]);
166    }
167
168    /// P1-2: filter_by_score 对 f64 分数同样适用。
169    #[test]
170    fn test_filter_by_score_f64() {
171        let scored = vec![("x", 0.8_f64), ("y", 0.0), ("z", -0.5)];
172        let filtered = filter_by_score(scored, 0.0);
173        let ids: Vec<&str> = filtered.iter().map(|(id, _)| *id).collect();
174        assert_eq!(ids, vec!["x"]);
175    }
176
177    /// P2-3: `doc_content_hash` 为确定性哈希——同内容多次调用结果一致,
178    /// 不同内容结果不同。FNV-1a 完全指定,跨进程/跨版本不漂移。
179    #[test]
180    fn test_doc_content_hash_stable() {
181        let content = "Rust 系统编程与并发";
182        let doc_a = Document::new(content.to_string());
183        let doc_b = Document::new(content.to_string());
184        let doc_c = Document::new("Python 数据科学");
185
186        let hash_a1 = doc_content_hash(&doc_a);
187        let hash_a2 = doc_content_hash(&doc_b);
188        assert_eq!(hash_a1, hash_a2, "相同内容应产生相同哈希");
189
190        let hash_c = doc_content_hash(&doc_c);
191        assert_ne!(hash_a1, hash_c, "不同内容应产生不同哈希");
192        assert_eq!(hash_a1.len(), 16, "应为 64 位哈希的 16 位十六进制表示");
193    }
194}
lc_rag/hybrid.rs

lc_rag/
hybrid.rs
