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

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

/// Channel type identifier.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum ChannelType {
    /// Semantic vector search channel.
    Semantic,
    /// Metadata-based search channel.
    Metadata,
    /// BM25 full-text search channel.
    Bm25,
    /// Full-text search (generic).
    FullText,
    /// Knowledge graph search channel.
    Graph,
    /// Custom channel type.
    Custom(String),
}

impl ChannelType {
    /// Return the string representation of this channel type.
    pub fn as_str(&self) -> &str {
        match self {
            Self::Semantic => "semantic",
            Self::Metadata => "metadata",
            Self::Bm25 => "bm25",
            Self::FullText => "fulltext",
            Self::Graph => "graph",
            Self::Custom(s) => s.as_str(),
        }
    }
}

/// Per-channel configuration.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelConfig {
    /// Channel type.
    pub channel_type: ChannelType,
    /// Weight for RRF fusion.
    pub weight: f32,
    /// Maximum number of results from this channel.
    pub top_k: usize,
    /// Minimum score threshold.
    pub min_score: Option<f32>,
    /// Additional channel-specific parameters.
    pub params: HashMap<String, serde_json::Value>,
}

impl ChannelConfig {
    /// Create a new channel config.
    pub fn new(channel_type: ChannelType, weight: f32, top_k: usize) -> Self {
        Self { channel_type, weight, top_k, min_score: None, params: HashMap::new() }
    }

    /// Create a semantic channel config.
    pub fn semantic(weight: f32, top_k: usize) -> Self {
        Self::new(ChannelType::Semantic, weight, top_k)
    }

    /// Create a metadata channel config.
    pub fn metadata(weight: f32, top_k: usize) -> Self {
        Self::new(ChannelType::Metadata, weight, top_k)
    }

    /// Set a minimum score threshold.
    pub fn with_min_score(mut self, min_score: f32) -> Self {
        self.min_score = Some(min_score);
        self
    }

    /// Add a channel-specific parameter.
    pub fn with_param(
        mut self,
        key: impl Into<String>,
        value: impl Into<serde_json::Value>,
    ) -> Self {
        self.params.insert(key.into(), value.into());
        self
    }
}

/// Multi-channel pipeline orchestrating retrieval and fusion.
#[derive(Debug, Clone)]
pub struct ChannelPipeline {
    /// Ordered list of channel configs.
    pub channels: Vec<ChannelConfig>,
    /// RRF k constant for score smoothing.
    pub rrf_k: usize,
    /// Whether to normalize scores before fusion.
    pub normalize_scores: bool,
}

impl ChannelPipeline {
    /// Create a new pipeline with the given channels.
    pub fn new(channels: Vec<ChannelConfig>) -> Self {
        Self { channels, rrf_k: 60, normalize_scores: true }
    }

    /// Set the RRF k constant.
    pub fn with_rrf_k(mut self, k: usize) -> Self {
        self.rrf_k = k;
        self
    }

    /// Set whether to normalize scores before fusion.
    pub fn with_normalize(mut self, normalize: bool) -> Self {
        self.normalize_scores = normalize;
        self
    }
}