docling_rag/
config.rs

1//! Configuration, loaded from the process environment (and a `.env` file).
2//!
3//! Every knob has a documented default so the crate runs out of the box with the
4//! offline-friendly stack (bundled SQLite + Ollama). See `.env.example` at the
5//! repo root for the full list.
6
7use crate::model::RetrievalMode;
8use crate::{RagError, Result};
9use std::str::FromStr;
10
11/// Which database backend backs the [`crate::store::VectorStore`].
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum DbBackend {
14    /// Bundled SQLite (default, zero external services).
15    Sqlite,
16    /// PostgreSQL (+ pgvector). Requires the `postgres` cargo feature.
17    Postgres,
18    /// Pure in-memory store — never persisted; used by tests and quick evals.
19    Memory,
20}
21
22/// Which embedding provider produces vectors.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum EmbedProvider {
25    /// Ollama HTTP API (default; e.g. `bge-m3`, 1024-dim).
26    Ollama,
27    /// Google Gemini embeddings (`gemini-embedding-001`, truncated to `dim`).
28    Gemini,
29    /// Local ONNX model. Requires the `onnx-embed` cargo feature.
30    Onnx,
31    /// Deterministic hashing embedder — no network, used for offline tests/eval.
32    Hash,
33}
34
35/// Which document source feeds the ingestion pipeline.
36#[derive(Debug, Clone, Copy, PartialEq, Eq)]
37pub enum SourceKind {
38    /// A local directory (works over any FUSE / network mount).
39    Folder,
40    /// FTP. Requires the `remote-sources` cargo feature.
41    Ftp,
42    /// SFTP. Requires the `remote-sources` cargo feature.
43    Sftp,
44}
45
46/// Which message queue drives async ingestion.
47#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum QueueKind {
49    /// In-process tokio channel (default).
50    Memory,
51    /// RabbitMQ (AMQP). Requires the `rabbitmq` cargo feature.
52    RabbitMq,
53    /// Redis pub/sub. Requires the `redis` cargo feature.
54    Redis,
55}
56
57/// The fully-resolved configuration for a RAG session.
58#[derive(Debug, Clone)]
59pub struct RagConfig {
60    // --- database ---
61    pub db_backend: DbBackend,
62    pub database_url: String,
63
64    // --- embedding ---
65    pub embed_provider: EmbedProvider,
66    pub embed_model: String,
67    pub embed_dim: usize,
68    pub ollama_base_url: String,
69    pub gemini_api_key: Option<String>,
70    pub gemini_model: String,
71    pub embed_onnx_path: String,
72    pub embed_tokenizer_path: String,
73
74    // --- llm (OpenRouter) ---
75    pub openrouter_api_key: Option<String>,
76    pub openrouter_base_url: String,
77    pub llm_model: String,
78
79    // --- OCR ---
80    /// OCR recognition language (`RAG_OCR_LANG`): `en` (default — English
81    /// PP-OCRv3, docling-pdf's own default) or `ch` (the multilingual
82    /// conformance-validated model; weak Latin word spacing). Maps onto
83    /// docling-pdf's `DOCLING_RS_OCR_LANG` — explicit `DOCLING_OCR_*` env
84    /// overrides always win.
85    pub ocr_lang: OcrLang,
86
87    // --- chunking ---
88    pub chunker: ChunkerKind,
89    pub chunk_size: usize,
90    pub chunk_overlap: f32,
91    pub chunk_unit: ChunkUnit,
92    /// Path to a HuggingFace `tokenizer.json` for the `hybrid` chunker's token
93    /// counts (`RAG_CHUNK_TOKENIZER`). Required when `chunker = hybrid`.
94    pub chunk_tokenizer: Option<String>,
95
96    // --- retrieval ---
97    pub retrieval_mode: RetrievalMode,
98    pub top_k: usize,
99    pub rrf_k: f32,
100    pub multiquery_n: usize,
101    /// BM25 term-frequency saturation (`RAG_BM25_K1`, default 1.2).
102    pub bm25_k1: f32,
103    /// BM25 length normalization (`RAG_BM25_B`, default 0.75).
104    pub bm25_b: f32,
105
106    // --- sources ---
107    pub source: SourceKind,
108    pub source_path: String,
109    pub source_url: Option<String>,
110    pub source_user: Option<String>,
111    pub source_password: Option<String>,
112    /// Optional folder to dump each ingested document's converted Markdown into
113    /// (for debugging / re-ingestion). `None` disables the dump.
114    pub documents_output: Option<String>,
115
116    // --- queue ---
117    pub queue: QueueKind,
118    pub rabbitmq_url: Option<String>,
119    pub redis_url: Option<String>,
120
121    // --- REST API ---
122    /// Bind address for `serve` (default `127.0.0.1:8080`).
123    pub http_addr: String,
124    /// Accepted API keys (comma-separated in `RAG_API_KEYS`). The server refuses
125    /// to start with an empty list — auth is fail-closed.
126    pub api_keys: Vec<String>,
127}
128
129/// OCR recognition language (`RAG_OCR_LANG`).
130#[derive(Debug, Clone, Copy, PartialEq, Eq)]
131pub enum OcrLang {
132    /// ch_PP-OCRv3 — multilingual; what docling conformance is measured with.
133    Ch,
134    /// en_PP-OCRv3 (default) — English-only, proper Latin word spacing.
135    En,
136}
137
138fn parse_ocr_lang(s: &str) -> Result<OcrLang> {
139    // The engine's own spellings (#388): `en`/`ch` plus BCP-47 tags for
140    // either language (`en-US`, `zh-Hans`, `iso:zh-CN`).
141    if s.trim().is_empty() {
142        return Ok(OcrLang::En);
143    }
144    match docling::OcrLang::parse(s) {
145        Some(docling::OcrLang::Ch) => Ok(OcrLang::Ch),
146        Some(docling::OcrLang::En) => Ok(OcrLang::En),
147        None => Err(RagError::config(format!(
148            "RAG_OCR_LANG={:?} is not supported ({})",
149            s.trim(),
150            docling::OcrLang::ACCEPTED
151        ))),
152    }
153}
154
155/// Which chunker ingestion runs (`RAG_CHUNKER`).
156#[derive(Debug, Clone, Copy, PartialEq, Eq)]
157pub enum ChunkerKind {
158    /// The Markdown sliding-window chunker (`chunk_size`/`chunk_overlap`/
159    /// `chunk_unit`), streaming — chunks pages while later pages still convert.
160    /// Default.
161    Window,
162    /// docling's structure-driven `HierarchicalChunker`: one chunk per document
163    /// item with its heading path (buffered — needs the whole document).
164    Hierarchical,
165    /// docling's `HybridChunker`: hierarchical chunks refined against a real
166    /// tokenizer budget (`chunk_size` tokens, `chunk_tokenizer` counts them).
167    Hybrid,
168}
169
170/// The unit used to measure chunk size / overlap.
171#[derive(Debug, Clone, Copy, PartialEq, Eq)]
172pub enum ChunkUnit {
173    /// Whitespace-delimited words (fast, no tokenizer). Default.
174    Word,
175    /// Approximate subword tokens (`chars/4` heuristic — no tokenizer dependency).
176    Token,
177}
178
179impl Default for RagConfig {
180    fn default() -> Self {
181        RagConfig {
182            db_backend: DbBackend::Sqlite,
183            database_url: "sqlite://data/rag.db".to_string(),
184            embed_provider: EmbedProvider::Ollama,
185            embed_model: "bge-m3".to_string(),
186            embed_dim: 1024,
187            ollama_base_url: "http://localhost:11434".to_string(),
188            gemini_api_key: None,
189            gemini_model: "gemini-embedding-001".to_string(),
190            embed_onnx_path: ".models/embed/bge-m3.onnx".to_string(),
191            embed_tokenizer_path: ".models/embed/tokenizer.json".to_string(),
192            openrouter_api_key: None,
193            openrouter_base_url: "https://openrouter.ai/api/v1".to_string(),
194            llm_model: "deepseek/deepseek-chat".to_string(),
195            ocr_lang: OcrLang::En,
196            chunker: ChunkerKind::Window,
197            chunk_size: 300,
198            chunk_overlap: 0.05,
199            chunk_unit: ChunkUnit::Word,
200            chunk_tokenizer: None,
201            retrieval_mode: RetrievalMode::Hybrid,
202            top_k: 5,
203            rrf_k: 60.0,
204            multiquery_n: 4,
205            bm25_k1: 1.2,
206            bm25_b: 0.75,
207            source: SourceKind::Folder,
208            source_path: "./input".to_string(),
209            source_url: None,
210            source_user: None,
211            source_password: None,
212            documents_output: None,
213            queue: QueueKind::Memory,
214            rabbitmq_url: None,
215            redis_url: None,
216            http_addr: "127.0.0.1:8080".to_string(),
217            api_keys: Vec::new(),
218        }
219    }
220}
221
222impl RagConfig {
223    /// Load a `.env` file (if present) then resolve config from the environment.
224    ///
225    /// Missing keys fall back to the [`Default`] values, so this never fails on an
226    /// empty environment; it only errors on an *invalid* value (bad number, unknown
227    /// backend name).
228    pub fn from_env() -> Result<Self> {
229        // Best-effort: a missing .env is not an error.
230        let _ = dotenvy::dotenv();
231        Self::from_env_inner()
232    }
233
234    fn from_env_inner() -> Result<Self> {
235        let d = RagConfig::default();
236        let cfg = RagConfig {
237            db_backend: match env_str("RAG_DB_BACKEND") {
238                Some(s) => parse_db_backend(&s)?,
239                None => d.db_backend,
240            },
241            database_url: env_str("RAG_DATABASE_URL").unwrap_or(d.database_url),
242            embed_provider: match env_str("RAG_EMBED_PROVIDER") {
243                Some(s) => parse_embed_provider(&s)?,
244                None => d.embed_provider,
245            },
246            embed_model: env_str("RAG_EMBED_MODEL").unwrap_or(d.embed_model),
247            embed_dim: env_parse("RAG_EMBED_DIM", d.embed_dim)?,
248            ollama_base_url: env_str("OLLAMA_BASE_URL").unwrap_or(d.ollama_base_url),
249            gemini_api_key: env_str("GEMINI_API_KEY"),
250            gemini_model: env_str("RAG_GEMINI_MODEL").unwrap_or(d.gemini_model),
251            embed_onnx_path: env_str("RAG_EMBED_ONNX_PATH").unwrap_or(d.embed_onnx_path),
252            embed_tokenizer_path: env_str("RAG_EMBED_TOKENIZER").unwrap_or(d.embed_tokenizer_path),
253            openrouter_api_key: env_str("OPENROUTER_API_KEY"),
254            openrouter_base_url: env_str("OPENROUTER_BASE_URL").unwrap_or(d.openrouter_base_url),
255            llm_model: env_str("RAG_LLM_MODEL").unwrap_or(d.llm_model),
256            ocr_lang: match env_str("RAG_OCR_LANG") {
257                Some(s) => parse_ocr_lang(&s)?,
258                None => d.ocr_lang,
259            },
260            chunker: match env_str("RAG_CHUNKER") {
261                Some(s) => parse_chunker_kind(&s)?,
262                None => d.chunker,
263            },
264            chunk_size: env_parse("RAG_CHUNK_SIZE", d.chunk_size)?,
265            chunk_overlap: env_parse("RAG_CHUNK_OVERLAP", d.chunk_overlap)?,
266            chunk_unit: match env_str("RAG_CHUNK_UNIT") {
267                Some(s) => parse_chunk_unit(&s)?,
268                None => d.chunk_unit,
269            },
270            chunk_tokenizer: env_str("RAG_CHUNK_TOKENIZER"),
271            retrieval_mode: match env_str("RAG_RETRIEVAL_MODE") {
272                Some(s) => RetrievalMode::from_str(&s)?,
273                None => d.retrieval_mode,
274            },
275            top_k: env_parse("RAG_TOP_K", d.top_k)?,
276            rrf_k: env_parse("RAG_RRF_K", d.rrf_k)?,
277            multiquery_n: env_parse("RAG_MULTIQUERY_N", d.multiquery_n)?,
278            bm25_k1: env_parse("RAG_BM25_K1", d.bm25_k1)?,
279            bm25_b: env_parse("RAG_BM25_B", d.bm25_b)?,
280            source: match env_str("RAG_SOURCE") {
281                Some(s) => parse_source_kind(&s)?,
282                None => d.source,
283            },
284            source_path: env_str("RAG_SOURCE_PATH").unwrap_or(d.source_path),
285            source_url: env_str("RAG_SOURCE_URL"),
286            source_user: env_str("RAG_SOURCE_USER"),
287            source_password: env_str("RAG_SOURCE_PASSWORD"),
288            documents_output: env_str("RAG_DOCUMENTS_OUTPUT"),
289            queue: match env_str("RAG_QUEUE") {
290                Some(s) => parse_queue_kind(&s)?,
291                None => d.queue,
292            },
293            rabbitmq_url: env_str("RABBITMQ_URL"),
294            redis_url: env_str("REDIS_URL"),
295            http_addr: env_str("RAG_HTTP_ADDR").unwrap_or(d.http_addr),
296            api_keys: env_str("RAG_API_KEYS")
297                .map(|s| {
298                    s.split(',')
299                        .map(|k| k.trim().to_string())
300                        .filter(|k| !k.is_empty())
301                        .collect()
302                })
303                .unwrap_or_default(),
304        };
305        cfg.validate()?;
306        Ok(cfg)
307    }
308
309    /// Sanity-check numeric ranges. Returns the first violation.
310    pub fn validate(&self) -> Result<()> {
311        if self.embed_dim == 0 {
312            return Err(RagError::config("RAG_EMBED_DIM must be > 0"));
313        }
314        if self.chunk_size == 0 {
315            return Err(RagError::config("RAG_CHUNK_SIZE must be > 0"));
316        }
317        if !(0.0..0.95).contains(&self.chunk_overlap) {
318            return Err(RagError::config("RAG_CHUNK_OVERLAP must be in [0.0, 0.95)"));
319        }
320        if self.chunker == ChunkerKind::Hybrid
321            && self.chunk_tokenizer.is_none()
322            && !std::path::Path::new(docling::chunker::DEFAULT_TOKENIZER_PATH).exists()
323        {
324            return Err(RagError::config(format!(
325                "RAG_CHUNKER=hybrid needs a HuggingFace tokenizer.json: set RAG_CHUNK_TOKENIZER \
326                 or run scripts/install/download_dependencies.sh (populates {})",
327                docling::chunker::DEFAULT_TOKENIZER_PATH
328            )));
329        }
330        if self.top_k == 0 {
331            return Err(RagError::config("RAG_TOP_K must be > 0"));
332        }
333        Ok(())
334    }
335}
336
337use docling_core::env::nonempty as env_str;
338
339fn env_parse<T>(key: &str, default: T) -> Result<T>
340where
341    T: FromStr,
342    T::Err: std::fmt::Display,
343{
344    match env_str(key) {
345        Some(s) => s
346            .parse::<T>()
347            .map_err(|e| RagError::config(format!("{key}: invalid value '{s}': {e}"))),
348        None => Ok(default),
349    }
350}
351
352fn parse_db_backend(s: &str) -> Result<DbBackend> {
353    match s.to_ascii_lowercase().as_str() {
354        "sqlite" => Ok(DbBackend::Sqlite),
355        "postgres" | "postgresql" | "pg" => Ok(DbBackend::Postgres),
356        "memory" | "mem" | "inmemory" => Ok(DbBackend::Memory),
357        other => Err(RagError::config(format!(
358            "unknown RAG_DB_BACKEND '{other}'"
359        ))),
360    }
361}
362
363fn parse_embed_provider(s: &str) -> Result<EmbedProvider> {
364    match s.to_ascii_lowercase().as_str() {
365        "ollama" => Ok(EmbedProvider::Ollama),
366        "gemini" | "google" => Ok(EmbedProvider::Gemini),
367        "onnx" | "local" => Ok(EmbedProvider::Onnx),
368        "hash" | "test" | "fake" => Ok(EmbedProvider::Hash),
369        other => Err(RagError::config(format!(
370            "unknown RAG_EMBED_PROVIDER '{other}'"
371        ))),
372    }
373}
374
375fn parse_chunker_kind(s: &str) -> Result<ChunkerKind> {
376    match s.to_ascii_lowercase().as_str() {
377        "window" => Ok(ChunkerKind::Window),
378        "hierarchical" => Ok(ChunkerKind::Hierarchical),
379        "hybrid" => Ok(ChunkerKind::Hybrid),
380        other => Err(RagError::config(format!(
381            "unknown RAG_CHUNKER '{other}' (expected: window, hierarchical, hybrid)"
382        ))),
383    }
384}
385
386fn parse_chunk_unit(s: &str) -> Result<ChunkUnit> {
387    match s.to_ascii_lowercase().as_str() {
388        "word" | "words" => Ok(ChunkUnit::Word),
389        "token" | "tokens" => Ok(ChunkUnit::Token),
390        other => Err(RagError::config(format!(
391            "unknown RAG_CHUNK_UNIT '{other}'"
392        ))),
393    }
394}
395
396fn parse_source_kind(s: &str) -> Result<SourceKind> {
397    match s.to_ascii_lowercase().as_str() {
398        "folder" | "dir" | "directory" | "local" => Ok(SourceKind::Folder),
399        "ftp" => Ok(SourceKind::Ftp),
400        "sftp" => Ok(SourceKind::Sftp),
401        other => Err(RagError::config(format!("unknown RAG_SOURCE '{other}'"))),
402    }
403}
404
405fn parse_queue_kind(s: &str) -> Result<QueueKind> {
406    match s.to_ascii_lowercase().as_str() {
407        "memory" | "mem" | "inproc" => Ok(QueueKind::Memory),
408        "rabbitmq" | "amqp" | "rabbit" => Ok(QueueKind::RabbitMq),
409        "redis" => Ok(QueueKind::Redis),
410        other => Err(RagError::config(format!("unknown RAG_QUEUE '{other}'"))),
411    }
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    #[test]
419    fn ocr_lang_parses_and_rejects_unknown() {
420        assert_eq!(parse_ocr_lang("").unwrap(), OcrLang::En);
421        assert_eq!(parse_ocr_lang("ch").unwrap(), OcrLang::Ch);
422        assert_eq!(parse_ocr_lang(" EN ").unwrap(), OcrLang::En);
423        assert_eq!(parse_ocr_lang("zh-Hans").unwrap(), OcrLang::Ch);
424        assert_eq!(parse_ocr_lang("en-US").unwrap(), OcrLang::En);
425        assert!(parse_ocr_lang("de").is_err());
426        assert!(parse_ocr_lang("de").is_err());
427    }
428
429    #[test]
430    fn defaults_are_valid_and_match_spec() {
431        let c = RagConfig::default();
432        c.validate().unwrap();
433        assert_eq!(c.chunk_size, 300);
434        assert!((c.chunk_overlap - 0.05).abs() < 1e-6);
435        assert_eq!(c.embed_dim, 1024);
436        assert_eq!(c.embed_model, "bge-m3");
437        assert_eq!(c.llm_model, "deepseek/deepseek-chat");
438        assert_eq!(c.retrieval_mode, RetrievalMode::Hybrid);
439    }
440
441    #[test]
442    fn validate_rejects_bad_overlap() {
443        let c = RagConfig {
444            chunk_overlap: 0.99,
445            ..Default::default()
446        };
447        assert!(c.validate().is_err());
448        let c = RagConfig {
449            chunk_size: 0,
450            ..Default::default()
451        };
452        assert!(c.validate().is_err());
453    }
454
455    #[test]
456    fn backend_parsers() {
457        assert_eq!(parse_db_backend("Postgres").unwrap(), DbBackend::Postgres);
458        assert_eq!(parse_embed_provider("HASH").unwrap(), EmbedProvider::Hash);
459        assert_eq!(parse_queue_kind("amqp").unwrap(), QueueKind::RabbitMq);
460        assert!(parse_source_kind("carrier-pigeon").is_err());
461    }
462}
docling_rag/config.rs

docling_rag/
config.rs
