xz-rag 0.1.1

Multi-channel Retrieval-Augmented Generation engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

use std::collections::HashMap;

use crate::types::retrieval::RetrievedChunk;

/// RRF (Reciprocal Rank Fusion) for merging multi-channel results.
#[derive(Debug, Clone)]
pub struct RrfFusion {
    /// RRF k constant (higher values smooth out rank differences).
    pub k: usize,
}

impl RrfFusion {
    /// Create a new RRF fusion instance with the given k constant.
    pub fn new(k: usize) -> Self {
        Self { k }
    }

    /// Fuse results from multiple channels using RRF.
    ///
    /// `channel_results` maps channel name -> ranked hits.
    /// Returns fused hits sorted by RRF score descending.
    pub fn fuse(
        &self,
        channel_results: HashMap<String, Vec<RetrievedChunk>>,
    ) -> Vec<RetrievedChunk> {
        let mut rrf_scores: HashMap<String, f32> = HashMap::new();
        let mut best_per_chunk: HashMap<String, RetrievedChunk> = HashMap::new();

        for hits in channel_results.values() {
            for (rank, hit) in hits.iter().enumerate() {
                let rrf = 1.0 / (self.k as f32 + (rank + 1) as f32);
                *rrf_scores.entry(hit.chunk_id.clone()).or_default() += rrf;

                // Keep highest original score per chunk
                best_per_chunk
                    .entry(hit.chunk_id.clone())
                    .and_modify(|existing| {
                        if hit.score > existing.score {
                            *existing = hit.clone();
                        }
                    })
                    .or_insert_with(|| hit.clone());
            }
        }

        let mut fused: Vec<RetrievedChunk> = best_per_chunk
            .into_iter()
            .map(|(chunk_id, mut chunk)| {
                chunk.score = rrf_scores.get(&chunk_id).copied().unwrap_or(0.0);
                chunk
            })
            .collect();

        fused.sort_by(|a, b| b.score.total_cmp(&a.score));
        fused
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_hit(id: &str, score: f32, channel: &str) -> RetrievedChunk {
        RetrievedChunk {
            chunk_id: id.to_string(),
            document_id: format!("doc_{}", id),
            content: format!("content_{}", id),
            score,
            channel: channel.to_string(),
            channel_score: score,
            metadata: crate::types::chunk::ChunkMetadata::default(),
            embedding: None,
        }
    }

    #[test]
    fn test_rrf_fusion_with_nan_scores() {
        let fusion = RrfFusion::new(60);

        let mut channel_results = HashMap::new();
        channel_results.insert(
            "semantic".to_string(),
            vec![make_hit("c1", f32::NAN, "semantic"), make_hit("c2", f32::NAN, "semantic")],
        );

        // Must not panic when sorting NaN scores
        let fused = fusion.fuse(channel_results);
        assert_eq!(fused.len(), 2);
    }

    #[test]
    fn test_rrf_fusion() {
        let fusion = RrfFusion::new(60);

        let mut channel_results = HashMap::new();
        channel_results.insert(
            "semantic".to_string(),
            vec![
                make_hit("c1", 0.9, "semantic"),
                make_hit("c2", 0.8, "semantic"),
                make_hit("c3", 0.5, "semantic"),
            ],
        );
        channel_results.insert(
            "bm25".to_string(),
            vec![
                make_hit("c3", 0.7, "bm25"),
                make_hit("c2", 0.4, "bm25"),
                make_hit("c4", 0.6, "bm25"),
            ],
        );

        let fused = fusion.fuse(channel_results);
        assert!(!fused.is_empty());

        // c2 appeared in both channels -> higher RRF score
        let c2 = fused.iter().find(|h| h.chunk_id == "c2").unwrap();
        let c4 = fused.iter().find(|h| h.chunk_id == "c4").unwrap();
        assert!(c2.score > c4.score);
    }
}