1use serde::{Deserialize, Serialize};
4
5use crate::error::{RagError, Result};
6
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
9pub struct RagConfig {
10 pub chunk_size: usize,
12 pub chunk_overlap: usize,
14 pub top_k: usize,
16 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 pub fn builder() -> RagConfigBuilder {
29 RagConfigBuilder::default()
30 }
31}
32
33#[derive(Debug, Clone, Default)]
35pub struct RagConfigBuilder {
36 config: RagConfig,
37}
38
39impl RagConfigBuilder {
40 pub fn chunk_size(mut self, size: usize) -> Self {
42 self.config.chunk_size = size;
43 self
44 }
45
46 pub fn chunk_overlap(mut self, overlap: usize) -> Self {
48 self.config.chunk_overlap = overlap;
49 self
50 }
51
52 pub fn top_k(mut self, k: usize) -> Self {
54 self.config.top_k = k;
55 self
56 }
57
58 pub fn similarity_threshold(mut self, threshold: f32) -> Self {
60 self.config.similarity_threshold = threshold;
61 self
62 }
63
64 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}