xz_rag/types/
retrieval.rs

1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use super::chunk::ChunkMetadata;
6
7use crate::pipeline::channel::ChannelConfig;
8
9// === Retrieval Request ===
10
11/// Multi-channel retrieval request.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct RetrieveRequest {
14    /// Search query string.
15    pub query: String,
16    /// Channel configurations for this request.
17    pub channels: Vec<ChannelConfig>,
18    /// Global filters applied across all channels.
19    ///
20    /// **Note on channel support**: `StructuredFilter` values are currently only applied by the
21    /// `metadata` channel. The `semantic`, `bm25`, and `graph` channels silently ignore
22    /// `global_filters` in their current implementation.
23    pub global_filters: Vec<StructuredFilter>,
24    /// Maximum number of results after fusion.
25    pub top_k: usize,
26    /// Optional namespace for scoping.
27    pub namespace: Option<String>,
28    /// Whether to include embedding vectors in results.
29    pub include_embeddings: bool,
30    /// Optional query preprocessing strategy.
31    pub query_preprocessing: Option<QueryPreprocessing>,
32}
33
34impl RetrieveRequest {
35    /// Create a new builder for `RetrieveRequest`.
36    pub fn builder(query: impl Into<String>) -> RetrieveRequestBuilder {
37        RetrieveRequestBuilder {
38            query: query.into(),
39            channels: vec![ChannelConfig::semantic(0.5, 10)],
40            global_filters: vec![],
41            top_k: 10,
42            namespace: None,
43            include_embeddings: false,
44            query_preprocessing: None,
45        }
46    }
47}
48
49/// Builder for `RetrieveRequest`.
50pub struct RetrieveRequestBuilder {
51    query: String,
52    channels: Vec<ChannelConfig>,
53    global_filters: Vec<StructuredFilter>,
54    top_k: usize,
55    namespace: Option<String>,
56    include_embeddings: bool,
57    query_preprocessing: Option<QueryPreprocessing>,
58}
59
60impl RetrieveRequestBuilder {
61    /// Set the channels to search.
62    pub fn channels(mut self, channels: Vec<ChannelConfig>) -> Self {
63        self.channels = channels;
64        self
65    }
66
67    /// Set the maximum number of results.
68    pub fn top_k(mut self, top_k: usize) -> Self {
69        self.top_k = top_k;
70        self
71    }
72
73    /// Set the namespace for scoping.
74    pub fn namespace(mut self, ns: impl Into<String>) -> Self {
75        self.namespace = Some(ns.into());
76        self
77    }
78
79    /// Build the `RetrieveRequest`.
80    pub fn build(self) -> RetrieveRequest {
81        RetrieveRequest {
82            query: self.query,
83            channels: self.channels,
84            global_filters: self.global_filters,
85            top_k: self.top_k,
86            namespace: self.namespace,
87            include_embeddings: self.include_embeddings,
88            query_preprocessing: self.query_preprocessing,
89        }
90    }
91}
92
93// === Query Preprocessing ===
94
95/// Query preprocessing strategy.
96#[derive(Debug, Clone, Serialize, Deserialize)]
97pub enum QueryPreprocessing {
98    /// No preprocessing.
99    None,
100    /// HYDE: generate a hypothetical answer passage.
101    Hyde,
102    /// Generate N query variations for better recall.
103    QueryExpansion {
104        /// Number of variations to generate.
105        count: usize,
106    },
107    /// Translate query to English before retrieval.
108    TranslateToEnglish,
109}
110
111// === Structured Filter ===
112
113/// Structured filter for metadata-based retrieval.
114#[derive(Debug, Clone, Serialize, Deserialize)]
115pub enum StructuredFilter {
116    /// Field equals value.
117    MetadataEq {
118        /// Metadata key.
119        key: String,
120        /// Expected value.
121        value: String,
122    },
123    /// Field is one of the given values.
124    MetadataIn {
125        /// Metadata key.
126        key: String,
127        /// Allowed values.
128        values: Vec<String>,
129    },
130    /// Field does not equal value.
131    MetadataNe {
132        /// Metadata key.
133        key: String,
134        /// Excluded value.
135        value: String,
136    },
137    /// Field exists (has any value).
138    MetadataExists {
139        /// Metadata key.
140        key: String,
141    },
142    /// Raw SQL filter expression.
143    SqlFilter(String),
144    /// Logical AND of two filters.
145    And(Box<StructuredFilter>, Box<StructuredFilter>),
146    /// Logical OR of two filters.
147    Or(Box<StructuredFilter>, Box<StructuredFilter>),
148}
149
150// === Retrieve Result ===
151
152/// Result of a multi-channel retrieval operation.
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub struct RetrieveResult {
155    /// Ranked list of retrieved chunks.
156    pub hits: Vec<RetrievedChunk>,
157    /// Per-channel statistics.
158    pub channel_report: HashMap<String, ChannelStats>,
159    /// Total latency in milliseconds.
160    pub latency_ms: u64,
161    /// The query text actually used (after preprocessing).
162    pub effective_query: String,
163}
164
165/// A single retrieved chunk with score and metadata.
166#[derive(Debug, Clone, Serialize, Deserialize)]
167pub struct RetrievedChunk {
168    /// Chunk identifier.
169    pub chunk_id: String,
170    /// Parent document identifier.
171    pub document_id: String,
172    /// Chunk text content.
173    pub content: String,
174    /// Relevance score (post-fusion).
175    pub score: f32,
176    /// Channel that produced this chunk.
177    pub channel: String,
178    /// Original score from the channel (pre-fusion).
179    pub channel_score: f32,
180    /// Chunk metadata.
181    pub metadata: ChunkMetadata,
182    /// Optional embedding vector.
183    pub embedding: Option<Vec<f32>>,
184}
185
186/// Per-channel retrieval statistics.
187#[derive(Debug, Clone, Serialize, Deserialize)]
188pub struct ChannelStats {
189    /// Channel type identifier.
190    pub channel_type: String,
191    /// Number of hits from this channel.
192    pub hits: usize,
193    /// Latency in milliseconds.
194    pub latency_ms: u64,
195    /// Minimum score among hits.
196    pub min_score: f32,
197    /// Maximum score among hits.
198    pub max_score: f32,
199}
xz_rag/types/retrieval.rs

xz_rag/types/
retrieval.rs
