xz_rag/types/
rag.rs

1use serde::{Deserialize, Serialize};
2
3use super::retrieval::{RetrieveRequest, RetrieveResult};
4
5// === RAG Request ===
6
7/// End-to-end RAG request including retrieval, context, and generation config.
8#[derive(Debug, Clone, Serialize, Deserialize)]
9pub struct RagRequest {
10    /// User query string.
11    pub query: String,
12    /// Optional system prompt override.
13    pub system_prompt: Option<String>,
14    /// Chat history for conversational context.
15    pub history: Vec<ChatMessage>,
16    /// Retrieval configuration.
17    pub retrieve_config: RetrieveRequest,
18    /// Generation parameters.
19    pub generation: RagGenerationConfig,
20    /// Optional context assembly configuration.
21    pub context_config: Option<ContextConfig>,
22    /// Optional prompt template name.
23    pub prompt_template: Option<String>,
24    /// Additional request options.
25    pub options: RequestOptions,
26}
27
28impl RagRequest {
29    /// Create a new builder for `RagRequest`.
30    pub fn builder(query: impl Into<String>) -> RagRequestBuilder {
31        RagRequestBuilder {
32            query: query.into(),
33            system_prompt: None,
34            history: vec![],
35            retrieve_config: RetrieveRequest::builder("").build(),
36            generation: RagGenerationConfig::default(),
37            context_config: None,
38            prompt_template: None,
39            options: RequestOptions::default(),
40        }
41    }
42}
43
44/// Builder for `RagRequest`.
45pub struct RagRequestBuilder {
46    query: String,
47    system_prompt: Option<String>,
48    history: Vec<ChatMessage>,
49    retrieve_config: RetrieveRequest,
50    generation: RagGenerationConfig,
51    context_config: Option<ContextConfig>,
52    prompt_template: Option<String>,
53    options: RequestOptions,
54}
55
56impl RagRequestBuilder {
57    /// Set the system prompt.
58    pub fn system_prompt(mut self, sp: impl Into<String>) -> Self {
59        self.system_prompt = Some(sp.into());
60        self
61    }
62
63    /// Set the chat history.
64    pub fn history(mut self, h: Vec<ChatMessage>) -> Self {
65        self.history = h;
66        self
67    }
68
69    /// Set the retrieval configuration.
70    pub fn retrieve_config(mut self, rc: RetrieveRequest) -> Self {
71        self.retrieve_config = rc;
72        self
73    }
74
75    /// Set the generation parameters.
76    pub fn generation(mut self, g: RagGenerationConfig) -> Self {
77        self.generation = g;
78        self
79    }
80
81    /// Set the context assembly configuration.
82    pub fn context_config(mut self, cc: ContextConfig) -> Self {
83        self.context_config = Some(cc);
84        self
85    }
86
87    /// Set the prompt template name.
88    pub fn prompt_template(mut self, pt: impl Into<String>) -> Self {
89        self.prompt_template = Some(pt.into());
90        self
91    }
92
93    /// Set additional request options.
94    pub fn options(mut self, opts: RequestOptions) -> Self {
95        self.options = opts;
96        self
97    }
98
99    /// Build the `RagRequest`.
100    pub fn build(self) -> RagRequest {
101        RagRequest {
102            query: self.query,
103            system_prompt: self.system_prompt,
104            history: self.history,
105            retrieve_config: self.retrieve_config,
106            generation: self.generation,
107            context_config: self.context_config,
108            prompt_template: self.prompt_template,
109            options: self.options,
110        }
111    }
112}
113
114// === Chat Message ===
115
116/// A single chat message in the conversation history.
117#[derive(Debug, Clone, Serialize, Deserialize)]
118pub struct ChatMessage {
119    /// Role of the message sender.
120    pub role: ChatRole,
121    /// Content of the message.
122    pub content: String,
123}
124
125/// Role of a chat message participant.
126#[derive(Debug, Clone, Serialize, Deserialize)]
127pub enum ChatRole {
128    /// System instruction message.
129    System,
130    /// User message.
131    User,
132    /// Assistant response.
133    Assistant,
134}
135
136// === Request Options ===
137
138/// Additional options for a RAG request.
139#[derive(Debug, Clone, Serialize, Deserialize, Default)]
140pub struct RequestOptions {
141    /// Optional namespace for scoping the retrieval.
142    pub namespace: Option<String>,
143    /// Optional timeout in milliseconds.
144    pub timeout_ms: Option<u64>,
145    /// Optional retry count for transient failures.
146    pub retry_count: Option<u32>,
147    /// Whether to stream the response.
148    pub stream: bool,
149}
150
151// === RAG Generation Config ===
152
153/// Generation parameters for the LLM.
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub struct RagGenerationConfig {
156    /// Maximum context tokens for generation.
157    pub max_context_tokens: usize,
158    /// Model identifier (e.g. "gpt-4o", "claude-3").
159    pub model: Option<String>,
160    /// Temperature for sampling (0.0-1.0).
161    pub temperature: Option<f32>,
162    /// Maximum output tokens.
163    pub max_output_tokens: Option<usize>,
164    /// Whether to stream the generation.
165    pub stream: bool,
166}
167
168impl Default for RagGenerationConfig {
169    fn default() -> Self {
170        Self {
171            max_context_tokens: 4096,
172            model: None,
173            temperature: Some(0.7),
174            max_output_tokens: Some(1024),
175            stream: false,
176        }
177    }
178}
179
180// === Context Config ===
181
182/// Configuration for context assembly from retrieved chunks.
183#[derive(Debug, Clone, Serialize, Deserialize)]
184pub struct ContextConfig {
185    /// Maximum context window size in tokens.
186    pub max_context_tokens: usize,
187    /// Tokens reserved for the system prompt.
188    pub system_prompt_reserve: usize,
189    /// Tokens reserved for the user query.
190    pub query_reserve: usize,
191    /// Tokens reserved for the model output.
192    pub output_reserve: usize,
193    /// Minimum overlap between chunks in the context.
194    pub min_chunk_overlap: usize,
195    /// Citation format to use.
196    pub citation_format: CitationFormat,
197}
198
199impl Default for ContextConfig {
200    fn default() -> Self {
201        Self {
202            max_context_tokens: 4096,
203            system_prompt_reserve: 256,
204            query_reserve: 128,
205            output_reserve: 512,
206            min_chunk_overlap: 50,
207            citation_format: CitationFormat::Numeric,
208        }
209    }
210}
211
212/// Citation format style.
213#[derive(Debug, Clone, Serialize, Deserialize)]
214pub enum CitationFormat {
215    /// Numeric citations like `[1]`, `[2]`, etc.
216    Numeric,
217    /// Use chunk IDs as citation markers.
218    ChunkId,
219    /// Use source document names as citation markers.
220    SourceName,
221}
222
223// === RAG Response ===
224
225/// Full RAG response including answer, citations, and usage statistics.
226#[derive(Debug, Clone, Serialize, Deserialize)]
227pub struct RagResponse {
228    /// Generated answer text.
229    pub answer: String,
230    /// List of citations referenced in the answer.
231    pub citations: Vec<Citation>,
232    /// Token usage statistics.
233    pub usage: RagTokenUsage,
234    /// Retrieval statistics from the pipeline.
235    pub retrieve_stats: RetrieveResult,
236    /// Total latency in milliseconds.
237    pub total_latency_ms: u64,
238    /// Model used for generation.
239    pub model: Option<String>,
240}
241
242// === Citation ===
243
244/// A single citation referencing a retrieved chunk.
245#[derive(Debug, Clone, Serialize, Deserialize)]
246pub struct Citation {
247    /// Citation index number.
248    pub index: usize,
249    /// ID of the referenced chunk.
250    pub chunk_id: String,
251    /// Content of the referenced chunk.
252    pub content: String,
253    /// Optional title of the source document.
254    pub document_title: Option<String>,
255    /// Relevance score of the chunk.
256    pub score: f32,
257    /// Channel that produced this chunk.
258    pub channel: String,
259}
260
261// === Token Usage ===
262
263/// Token usage statistics for a RAG operation.
264#[derive(Debug, Clone, Serialize, Deserialize, Default)]
265pub struct RagTokenUsage {
266    /// Tokens consumed by the context.
267    pub context_tokens: usize,
268    /// Tokens consumed by the prompt.
269    pub prompt_tokens: usize,
270    /// Tokens consumed by the completion.
271    pub completion_tokens: usize,
272    /// Total tokens consumed.
273    pub total_tokens: usize,
274    /// Number of chunks included in the context.
275    pub chunks_used: usize,
276    /// Number of chunks dropped due to budget limits.
277    pub chunks_dropped: usize,
278}
279
280// === Prompt Template ===
281
282/// Prompt template for formatting the LLM prompt.
283#[derive(Debug, Clone, Serialize, Deserialize)]
284pub struct PromptTemplate {
285    /// Template name.
286    pub name: String,
287    /// System prompt content.
288    pub system: String,
289    /// User message template with `{query}` and `{context}` placeholders.
290    pub user_template: String,
291    /// Prefix prepended to the context block.
292    pub context_prefix: String,
293    /// Format string for each chunk in the context.
294    pub chunk_format: String,
295    /// Suffix appended to the context block.
296    pub context_suffix: String,
297    /// Citation instruction appended to the prompt.
298    pub citation_instruction: String,
299}
300
301impl PromptTemplate {
302    /// Create a default QA prompt template.
303    pub fn default_qa() -> Self {
304        Self {
305            name: "default_qa".into(),
306            system: "You are a helpful assistant. Answer the user's question based on the provided context.".into(),
307            user_template: "Question: {query}\n\nContext:\n{context}\n\nAnswer:".into(),
308            context_prefix: "Relevant information:\n".into(),
309            chunk_format: "[{index}] {content}\n".into(),
310            context_suffix: "\n".into(),
311            citation_instruction: "Please cite sources using [N] notation when referencing the context.".into(),
312        }
313    }
314
315    /// Render the user template with the given query and context.
316    pub fn render(&self, query: &str, context: &str) -> String {
317        self.user_template.replace("{query}", query).replace("{context}", context)
318    }
319}
320
321// === Streaming Event ===
322
323/// Events emitted during a streaming RAG operation.
324#[derive(Debug, Clone, Serialize, Deserialize)]
325#[serde(tag = "type")]
326pub enum RagStreamEvent {
327    /// Retrieval phase has started.
328    RetrievalStarted {
329        /// Number of channels being queried.
330        channel_count: usize,
331    },
332    /// A single channel has completed retrieval.
333    ChannelDone {
334        /// Channel identifier.
335        channel: String,
336        /// Number of hits from this channel.
337        hits: usize,
338        /// Latency for this channel in milliseconds.
339        latency_ms: u64,
340    },
341    /// Generation phase has started with the assembled context.
342    GenerationStarted {
343        /// Number of chunks included in the context.
344        context_chunks: usize,
345        /// Tokens consumed by the context.
346        context_tokens: usize,
347    },
348    /// A delta of generated content.
349    ContentDelta {
350        /// Text delta chunk.
351        delta: String,
352    },
353    /// A citation has been identified in the generated content.
354    Citation {
355        /// Referenced chunk ID.
356        chunk_id: String,
357        /// Citation index.
358        index: usize,
359    },
360    /// Generation is complete.
361    Done {
362        /// Total latency in milliseconds.
363        total_latency_ms: u64,
364        /// Full list of citations.
365        citations: Vec<Citation>,
366        /// Token usage statistics.
367        usage: RagTokenUsage,
368    },
369}
370
371// === Built Context ===
372
373/// Assembled context from retrieved chunks, ready for generation.
374#[derive(Debug, Clone, Serialize, Deserialize)]
375pub struct BuiltContext {
376    /// Formatted context text.
377    pub context_text: String,
378    /// Citations extracted from the context.
379    pub citations: Vec<Citation>,
380    /// Number of chunks included in the context.
381    pub chunks_used: usize,
382    /// Number of chunks dropped due to budget limits.
383    pub chunks_dropped: usize,
384    /// Tokens consumed by the context.
385    pub tokens_used: usize,
386}
xz_rag/types/rag.rs

xz_rag/types/
rag.rs
