Skip to main content

adk_rag/
config.rs

1//! Configuration for the RAG pipeline.
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{RagError, Result};
6
7/// Configuration parameters for the RAG pipeline.
8#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
9pub struct RagConfig {
10    /// Maximum chunk size in characters.
11    pub chunk_size: usize,
12    /// Number of overlapping characters between consecutive chunks.
13    pub chunk_overlap: usize,
14    /// Number of top results to return from vector search.
15    pub top_k: usize,
16    /// Minimum similarity score for results (results below this are filtered out).
17    pub similarity_threshold: f32,
18}
19
20impl Default for RagConfig {
21    fn default() -> Self {
22        Self { chunk_size: 512, chunk_overlap: 100, top_k: 10, similarity_threshold: 0.0 }
23    }
24}
25
26impl RagConfig {
27    /// Create a new builder for constructing a [`RagConfig`].
28    pub fn builder() -> RagConfigBuilder {
29        RagConfigBuilder::default()
30    }
31}
32
33/// Builder for constructing a validated [`RagConfig`].
34#[derive(Debug, Clone, Default)]
35pub struct RagConfigBuilder {
36    config: RagConfig,
37}
38
39impl RagConfigBuilder {
40    /// Set the maximum chunk size in characters.
41    pub fn chunk_size(mut self, size: usize) -> Self {
42        self.config.chunk_size = size;
43        self
44    }
45
46    /// Set the overlap between consecutive chunks in characters.
47    pub fn chunk_overlap(mut self, overlap: usize) -> Self {
48        self.config.chunk_overlap = overlap;
49        self
50    }
51
52    /// Set the number of top results to return from vector search.
53    pub fn top_k(mut self, k: usize) -> Self {
54        self.config.top_k = k;
55        self
56    }
57
58    /// Set the minimum similarity threshold for filtering results.
59    pub fn similarity_threshold(mut self, threshold: f32) -> Self {
60        self.config.similarity_threshold = threshold;
61        self
62    }
63
64    /// Build the [`RagConfig`], validating that parameters are consistent.
65    ///
66    /// # Errors
67    ///
68    /// Returns [`RagError::ConfigError`] if:
69    /// - `chunk_overlap >= chunk_size`
70    /// - `top_k == 0`
71    pub fn build(self) -> Result<RagConfig> {
72        if self.config.chunk_overlap >= self.config.chunk_size {
73            return Err(RagError::ConfigError(format!(
74                "chunk_overlap ({}) must be less than chunk_size ({})",
75                self.config.chunk_overlap, self.config.chunk_size
76            )));
77        }
78        if self.config.top_k == 0 {
79            return Err(RagError::ConfigError("top_k must be greater than zero".to_string()));
80        }
81        Ok(self.config)
82    }
83}