ragrig_gpui/rag/
config.rs

1//! Editable Ragrig configuration for the settings page and profiles.
2
3use std::{
4    fs,
5    path::{Path, PathBuf},
6};
7
8use ragrig::{
9    ChatConfig, EmbedConfig, GenerationParams, MemoryConfig, MemoryStrategyKind, ParseConfig,
10    PdfParserBackend, RagrigConfig,
11    types::{EmbeddingProvider, Provider},
12};
13
14/// Chat backend offered by the settings page.
15#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
16pub enum ChatProvider {
17    /// Local Ollama server.
18    #[default]
19    Ollama,
20    /// DeepSeek cloud API; the key comes from `DEEPSEEK_API_KEY`.
21    DeepSeek,
22}
23
24impl ChatProvider {
25    /// Human-readable name.
26    pub fn label(self) -> &'static str {
27        match self {
28            Self::Ollama => "Ollama",
29            Self::DeepSeek => "DeepSeek",
30        }
31    }
32
33    /// Options for a dropdown.
34    pub fn options() -> &'static [(Self, &'static str)] {
35        &[(Self::Ollama, "Ollama"), (Self::DeepSeek, "DeepSeek")]
36    }
37}
38
39/// PDF backend offered by the settings page.
40#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
41pub enum PdfBackend {
42    /// Legacy flat-text extraction.
43    #[default]
44    Extract,
45    /// Markdown-native extraction.
46    Unpdf,
47    /// Structured, layout-aware extraction.
48    Sink,
49    /// Binary scavenger that never panics.
50    Internal,
51}
52
53impl PdfBackend {
54    /// Human-readable name.
55    pub fn label(self) -> &'static str {
56        match self {
57            Self::Extract => "pdf-extract",
58            Self::Unpdf => "unpdf",
59            Self::Sink => "pdfsink",
60            Self::Internal => "sloppy-pdf",
61        }
62    }
63
64    /// Options for a dropdown.
65    pub fn options() -> &'static [(Self, &'static str)] {
66        &[
67            (Self::Extract, "pdf-extract (flat text)"),
68            (Self::Unpdf, "unpdf (Markdown)"),
69            (Self::Sink, "pdfsink (layout)"),
70            (Self::Internal, "sloppy-pdf (scavenger)"),
71        ]
72    }
73}
74
75/// Memory strategy offered by the settings page.
76#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
77pub enum MemoryChoice {
78    /// Rewrite the query with a small model before retrieval.
79    #[default]
80    Rewrite,
81    /// Pass the transcript through without rewriting.
82    Transcript,
83    /// Blend the most recent saved session into the query.
84    Log,
85    /// Summarise past sessions with an LLM.
86    Summary,
87    /// No transcript, no rewriting.
88    Off,
89}
90
91impl MemoryChoice {
92    /// Human-readable name.
93    pub fn label(self) -> &'static str {
94        match self {
95            Self::Rewrite => "Rewrite",
96            Self::Transcript => "Transcript",
97            Self::Log => "Log",
98            Self::Summary => "Summary",
99            Self::Off => "Off",
100        }
101    }
102
103    /// Options for a dropdown.
104    pub fn options() -> &'static [(Self, &'static str)] {
105        &[
106            (Self::Rewrite, "Rewrite (small model)"),
107            (Self::Transcript, "Transcript only"),
108            (Self::Log, "Log (last session)"),
109            (Self::Summary, "Summary (LLM)"),
110            (Self::Off, "Off"),
111        ]
112    }
113
114    fn uses_model(self) -> bool {
115        matches!(self, Self::Rewrite | Self::Log | Self::Summary)
116    }
117}
118
119/// Everything needed to build one [`ragrig::RagAgent`].
120#[derive(Clone, Debug)]
121pub struct RagSetup {
122    /// Ragrig's own configuration tree.
123    pub config: RagrigConfig,
124    /// Prompt text applied after building; takes precedence over
125    /// `config.chat.system_prompt_path`.
126    pub system_prompt: Option<String>,
127    /// Whether answers stream token by token.
128    pub stream: bool,
129    /// Whether `@key` mentions in a question attach the paper as one-shot
130    /// context. On by default.
131    pub resolve_citations: bool,
132}
133
134impl RagSetup {
135    /// Wrap a config without a prompt override.
136    pub fn new(config: RagrigConfig) -> Self {
137        Self {
138            config,
139            system_prompt: None,
140            stream: true,
141            resolve_citations: true,
142        }
143    }
144
145    /// Wrap a config and apply `prompt` before the agent is built.
146    pub fn with_system_prompt(config: RagrigConfig, prompt: impl Into<String>) -> Self {
147        Self {
148            config,
149            system_prompt: Some(prompt.into()),
150            stream: true,
151            resolve_citations: true,
152        }
153    }
154}
155
156impl From<RagrigConfig> for RagSetup {
157    fn from(config: RagrigConfig) -> Self {
158        Self::new(config)
159    }
160}
161
162/// Editable configuration shown by the settings page.
163///
164/// This mirrors the Ragrig fields a chat surface exposes. Secrets never enter
165/// this struct: the DeepSeek key is read from the environment by the worker.
166#[derive(Clone, Debug, PartialEq)]
167pub struct RagSettings {
168    /// Workspace folder holding the store and sessions.
169    pub workspace: String,
170    /// Active chat backend.
171    pub provider: ChatProvider,
172    /// Ollama chat model.
173    pub ollama_model: String,
174    /// DeepSeek chat model.
175    pub deepseek_model: String,
176    /// Sampling temperature.
177    pub temperature: f64,
178    /// Maximum generated tokens.
179    pub max_tokens: usize,
180    /// Chat context window in tokens.
181    pub context_tokens: usize,
182    /// System prompt sent before the conversation.
183    pub system_prompt: String,
184    /// Stream tokens into the transcript.
185    pub stream: bool,
186    /// Ollama embedding model.
187    pub embed_model: String,
188    /// Chunks retrieved per query.
189    pub top_k: usize,
190    /// Minimum similarity for a chunk to enter the prompt.
191    pub similarity_threshold: f64,
192    /// Target chunk size in tokens.
193    pub chunk_size: usize,
194    /// Token overlap between chunks.
195    pub chunk_overlap: usize,
196    /// PDF parser selected as primary.
197    pub pdf_backend: PdfBackend,
198    /// Keep the sloppy PDF fallback registered.
199    pub sloppy_pdf: bool,
200    /// Conversation memory strategy.
201    pub memory: MemoryChoice,
202    /// Small model used for rewriting and summaries.
203    pub memory_model: String,
204    /// Base URL of the GROBID server.
205    pub grobid_url: String,
206    /// Rename PDFs to their metadata name during reference extraction.
207    pub rename_pdfs: bool,
208    /// Complete extracted references against OpenAlex.
209    pub openalex: bool,
210    /// Attach `@key` papers mentioned in a question as one-shot context.
211    pub resolve_citations: bool,
212    /// Named folder corpora in `name=path` form.
213    pub corpus_dirs: Vec<String>,
214    /// Profile name typed by the user.
215    pub profile_name: String,
216    /// Profiles found in the workspace.
217    pub profiles: Vec<String>,
218    /// Last settings action result, shown under the Apply row.
219    pub status: Option<String>,
220}
221
222impl Default for RagSettings {
223    fn default() -> Self {
224        let config = RagrigConfig::default();
225        let mut settings = Self::from_config(&config);
226        settings.workspace = default_workspace().to_string_lossy().into_owned();
227        settings
228    }
229}
230
231impl RagSettings {
232    /// Read the editable fields from a Ragrig configuration.
233    pub fn from_config(config: &RagrigConfig) -> Self {
234        let provider = match config.chat.provider {
235            Provider::Deepseek => ChatProvider::DeepSeek,
236            _ => ChatProvider::Ollama,
237        };
238        let pdf_backend = match config.parse.pdf_parser {
239            PdfParserBackend::Unpdf => PdfBackend::Unpdf,
240            PdfParserBackend::Sink => PdfBackend::Sink,
241            PdfParserBackend::Internal => PdfBackend::Internal,
242            _ => PdfBackend::Extract,
243        };
244        let memory = match config.memory.strategy {
245            MemoryStrategyKind::Transcript => MemoryChoice::Transcript,
246            MemoryStrategyKind::Log => MemoryChoice::Log,
247            MemoryStrategyKind::Summary => MemoryChoice::Summary,
248            MemoryStrategyKind::Off => MemoryChoice::Off,
249            _ => MemoryChoice::Rewrite,
250        };
251        let system_prompt = config
252            .chat
253            .system_prompt_path
254            .as_ref()
255            .map(|path| {
256                if path.is_absolute() {
257                    path.clone()
258                } else {
259                    config.workspace.join(path)
260                }
261            })
262            .and_then(|path| fs::read_to_string(path).ok())
263            .filter(|text| !text.trim().is_empty())
264            .unwrap_or_else(default_prompt);
265
266        Self {
267            workspace: config.workspace.to_string_lossy().into_owned(),
268            provider,
269            ollama_model: config.chat.model.clone(),
270            deepseek_model: config.chat.deepseek_model.clone(),
271            temperature: config.chat.params.temperature.unwrap_or(0.2),
272            max_tokens: config.chat.params.max_tokens.unwrap_or(2048),
273            context_tokens: config.chat.context_tokens,
274            system_prompt,
275            stream: true,
276            embed_model: config.embed.model.clone(),
277            top_k: config.embed.top_k,
278            similarity_threshold: config.embed.similarity_threshold,
279            chunk_size: config.parse.chunk_size,
280            chunk_overlap: config.parse.chunk_overlap,
281            pdf_backend,
282            sloppy_pdf: config.parse.sloppy_pdf,
283            memory,
284            memory_model: config.memory.model.clone(),
285            grobid_url: "http://localhost:8070".to_owned(),
286            rename_pdfs: false,
287            openalex: false,
288            resolve_citations: true,
289            corpus_dirs: config.corpus_dirs.clone(),
290            profile_name: String::new(),
291            profiles: Vec::new(),
292            status: None,
293        }
294    }
295
296    /// Validate and convert to a Ragrig configuration.
297    pub fn to_config(&self) -> Result<RagrigConfig, String> {
298        self.validate()?;
299
300        let workspace = PathBuf::from(self.workspace.trim());
301        let chat = ChatConfig {
302            provider: match self.provider {
303                ChatProvider::Ollama => Provider::Ollama,
304                ChatProvider::DeepSeek => Provider::Deepseek,
305            },
306            model: self.ollama_model.trim().to_owned(),
307            deepseek_model: self.deepseek_model.trim().to_owned(),
308            // Secrets only ever come from the environment.
309            deepseek_api_key: None,
310            params: GenerationParams {
311                temperature: Some(self.temperature),
312                top_p: None,
313                max_tokens: Some(self.max_tokens),
314                seed: None,
315            },
316            context_tokens: self.context_tokens,
317            system_prompt_path: None,
318            ..ChatConfig::default()
319        };
320
321        let embed = EmbedConfig {
322            provider: EmbeddingProvider::Ollama,
323            model: self.embed_model.trim().to_owned(),
324            top_k: self.top_k,
325            similarity_threshold: self.similarity_threshold,
326            ..EmbedConfig::default()
327        };
328
329        let parse = ParseConfig {
330            pdf_parser: match self.pdf_backend {
331                PdfBackend::Extract => PdfParserBackend::Extract,
332                PdfBackend::Unpdf => PdfParserBackend::Unpdf,
333                PdfBackend::Sink => PdfParserBackend::Sink,
334                PdfBackend::Internal => PdfParserBackend::Internal,
335            },
336            sloppy_pdf: self.sloppy_pdf,
337            chunk_size: self.chunk_size,
338            chunk_overlap: self.chunk_overlap,
339        };
340
341        let memory = MemoryConfig {
342            model: self.memory_model.trim().to_owned(),
343            rewrite_prompt_path: None,
344            strategy: match self.memory {
345                MemoryChoice::Rewrite => MemoryStrategyKind::Rewrite,
346                MemoryChoice::Transcript => MemoryStrategyKind::Transcript,
347                MemoryChoice::Log => MemoryStrategyKind::Log,
348                MemoryChoice::Summary => MemoryStrategyKind::Summary,
349                MemoryChoice::Off => MemoryStrategyKind::Off,
350            },
351        };
352
353        Ok(RagrigConfig {
354            workspace,
355            chat,
356            embed,
357            parse,
358            memory,
359            corpus_dirs: self
360                .corpus_dirs
361                .iter()
362                .map(|entry| entry.trim().to_owned())
363                .filter(|entry| !entry.is_empty())
364                .collect(),
365            corpus_urls: Vec::new(),
366            semantic_scholar_api_key: None,
367        })
368    }
369
370    /// Validate and produce a build setup, including the prompt override.
371    pub fn to_setup(&self) -> Result<RagSetup, String> {
372        let config = self.to_config()?;
373        let prompt = self.system_prompt.trim();
374        let mut setup = if prompt.is_empty() {
375            RagSetup::new(config)
376        } else {
377            RagSetup::with_system_prompt(config, prompt)
378        };
379        setup.stream = self.stream;
380        setup.resolve_citations = self.resolve_citations;
381        Ok(setup)
382    }
383
384    /// Check the editable values without touching the filesystem.
385    pub fn validate(&self) -> Result<(), String> {
386        if self.workspace.trim().is_empty() {
387            return Err("choose a workspace folder".to_owned());
388        }
389        match self.provider {
390            ChatProvider::Ollama => {
391                if self.ollama_model.trim().is_empty() {
392                    return Err("choose an Ollama chat model".to_owned());
393                }
394            }
395            ChatProvider::DeepSeek => {
396                if self.deepseek_model.trim().is_empty() {
397                    return Err("choose a DeepSeek chat model".to_owned());
398                }
399                if std::env::var_os("DEEPSEEK_API_KEY").is_none() {
400                    return Err(
401                        "DEEPSEEK_API_KEY is not set; export it before building the agent"
402                            .to_owned(),
403                    );
404                }
405            }
406        }
407        if !self.temperature.is_finite() || !(0.0..=2.0).contains(&self.temperature) {
408            return Err("temperature must be between 0 and 2".to_owned());
409        }
410        if self.max_tokens == 0 {
411            return Err("max tokens must be at least 1".to_owned());
412        }
413        if self.context_tokens < 256 {
414            return Err("context tokens must be at least 256".to_owned());
415        }
416        if self.embed_model.trim().is_empty() {
417            return Err("choose an embedding model".to_owned());
418        }
419        if self.top_k == 0 {
420            return Err("top-k must be at least 1".to_owned());
421        }
422        if !self.similarity_threshold.is_finite()
423            || !(0.0..=1.0).contains(&self.similarity_threshold)
424        {
425            return Err("similarity threshold must be between 0 and 1".to_owned());
426        }
427        if self.chunk_size == 0 {
428            return Err("chunk size must be at least 1".to_owned());
429        }
430        if self.chunk_overlap >= self.chunk_size {
431            return Err("chunk overlap must be smaller than the chunk size".to_owned());
432        }
433        if self.memory.uses_model() && self.memory_model.trim().is_empty() {
434            return Err("choose a memory model or turn memory off".to_owned());
435        }
436        if self.grobid_url.trim().is_empty() {
437            return Err("enter the GROBID server URL".to_owned());
438        }
439        for entry in &self.corpus_dirs {
440            let entry = entry.trim();
441            if entry.is_empty() {
442                continue;
443            }
444            let Some((name, path)) = entry.split_once('=') else {
445                return Err(format!("corpus '{entry}' must use name=path"));
446            };
447            if name.trim().is_empty() || path.trim().is_empty() {
448                return Err(format!("corpus '{entry}' needs a name and a path"));
449            }
450        }
451        Ok(())
452    }
453
454    /// Re-read the profile list from the workspace.
455    pub fn refresh_profiles(&mut self) {
456        let workspace = Path::new(self.workspace.trim());
457        match RagrigConfig::list_profiles(workspace) {
458            Ok(profiles) => {
459                self.profiles = profiles;
460                if !self.profile_name.trim().is_empty()
461                    && !self
462                        .profiles
463                        .iter()
464                        .any(|name| name == self.profile_name.trim())
465                {
466                    self.profiles.insert(0, self.profile_name.trim().to_owned());
467                }
468            }
469            Err(error) => {
470                self.profiles.clear();
471                self.status = Some(format!("cannot list profiles: {error:#}"));
472            }
473        }
474    }
475
476    /// Save these settings as a named profile in the workspace.
477    pub fn save_profile(&mut self) -> Result<(), String> {
478        let name = self.profile_name.trim().to_owned();
479        validate_profile_name(&name)?;
480        let mut config = self.to_config()?;
481        let workspace = Path::new(self.workspace.trim());
482        let prompt = self.system_prompt.trim();
483        if prompt.is_empty() {
484            config.chat.system_prompt_path = None;
485        } else {
486            let dir = RagrigConfig::profiles_dir(workspace);
487            fs::create_dir_all(&dir)
488                .map_err(|error| format!("cannot create profiles folder: {error}"))?;
489            let relative = PathBuf::from(format!(".ragrig/profiles/{name}.prompt.md"));
490            fs::write(workspace.join(&relative), prompt).map_err(|error| {
491                format!("cannot write the prompt for profile '{name}': {error}")
492            })?;
493            config.chat.system_prompt_path = Some(relative);
494        }
495        config
496            .save_to_profile(workspace, &name)
497            .map_err(|error| format!("cannot save profile '{name}': {error:#}"))?;
498        self.refresh_profiles();
499        self.status = Some(format!("Saved profile '{name}'"));
500        Ok(())
501    }
502
503    /// Load a named profile from the workspace into the editable fields.
504    pub fn load_profile(&mut self, name: &str) -> Result<(), String> {
505        validate_profile_name(name)?;
506        let workspace = Path::new(self.workspace.trim());
507        let config = RagrigConfig::load_from_profile(workspace, name)
508            .map_err(|error| format!("cannot load profile '{name}': {error:#}"))?;
509        let profiles = self.profiles.clone();
510        let mut loaded = Self::from_config(&config);
511        loaded.profiles = profiles;
512        loaded.profile_name = name.to_owned();
513        loaded.status = Some(format!("Loaded profile '{name}' — press Apply"));
514        *self = loaded;
515        self.refresh_profiles();
516        Ok(())
517    }
518
519    /// Remember a folder corpus so later applies keep it.
520    pub fn add_corpus(&mut self, name: &str, folder: &Path) {
521        let entry = format!("{name}={}", folder.display());
522        self.corpus_dirs.retain(|existing| {
523            existing
524                .split_once('=')
525                .is_none_or(|(existing_name, _)| existing_name != name)
526        });
527        self.corpus_dirs.push(entry);
528    }
529}
530
531fn default_workspace() -> PathBuf {
532    std::env::var_os("RAGRIG_WORKSPACE")
533        .map(PathBuf::from)
534        .or_else(|| std::env::current_dir().ok())
535        .unwrap_or_else(|| PathBuf::from("."))
536}
537
538fn default_prompt() -> String {
539    "Answer from the attached context and the indexed corpus. Cite each paper by its @Key. \
540     If the context does not support an answer, say what is missing instead of guessing."
541        .to_owned()
542}
543
544fn validate_profile_name(name: &str) -> Result<(), String> {
545    if name.is_empty() {
546        return Err("enter a profile name first".to_owned());
547    }
548    if !name
549        .chars()
550        .all(|character| character.is_ascii_alphanumeric() || character == '-' || character == '_')
551    {
552        return Err("profile names may use letters, digits, '-' and '_' only".to_owned());
553    }
554    Ok(())
555}
556
557#[cfg(test)]
558mod tests {
559    use super::*;
560
561    fn fixture() -> RagSettings {
562        RagSettings {
563            workspace: std::env::temp_dir().display().to_string(),
564            ..RagSettings::default()
565        }
566    }
567
568    #[test]
569    fn defaults_convert_to_a_valid_setup() {
570        let setup = fixture().to_setup().expect("valid defaults");
571        assert_eq!(setup.config.parse.chunk_size, 1024);
572        assert_eq!(setup.config.parse.chunk_overlap, 128);
573        assert!(setup.system_prompt.is_some());
574        assert!(setup.resolve_citations);
575
576        let mut settings = fixture();
577        settings.resolve_citations = false;
578        assert!(!settings.to_setup().expect("valid setup").resolve_citations);
579    }
580
581    #[test]
582    fn invalid_values_are_rejected() {
583        let mut settings = fixture();
584        settings.chunk_overlap = settings.chunk_size;
585        assert_eq!(
586            settings.validate(),
587            Err("chunk overlap must be smaller than the chunk size".to_owned())
588        );
589
590        settings = fixture();
591        settings.top_k = 0;
592        assert_eq!(
593            settings.validate(),
594            Err("top-k must be at least 1".to_owned())
595        );
596
597        settings = fixture();
598        settings.similarity_threshold = 1.5;
599        assert_eq!(
600            settings.validate(),
601            Err("similarity threshold must be between 0 and 1".to_owned())
602        );
603
604        settings = fixture();
605        settings.corpus_dirs = vec!["broken".to_owned()];
606        assert_eq!(
607            settings.validate(),
608            Err("corpus 'broken' must use name=path".to_owned())
609        );
610    }
611
612    #[test]
613    fn profiles_round_trip_prompt_and_api_key_slot() {
614        let workspace = unique_temp_dir("ragrig-gpui-profile");
615        let mut settings = fixture();
616        settings.workspace = workspace.display().to_string();
617        settings.ollama_model = "qwen3.5:9b".to_owned();
618        settings.system_prompt = "Cite primary sources only.".to_owned();
619        settings.profile_name = "lab".to_owned();
620        settings.status = None;
621
622        settings.save_profile().expect("save profile");
623        let path = workspace.join(".ragrig/profiles/lab.json");
624        let json = std::fs::read_to_string(&path).expect("profile file");
625        assert!(!json.contains("sk-"));
626        assert!(json.contains("\"deepseek_api_key\": null"));
627        assert!(json.contains("lab.prompt.md"));
628        assert_eq!(
629            std::fs::read_to_string(workspace.join(".ragrig/profiles/lab.prompt.md"))
630                .expect("prompt file"),
631            "Cite primary sources only."
632        );
633
634        settings.workspace = workspace.display().to_string();
635        settings.refresh_profiles();
636        settings.profile_name = "lab".to_owned();
637        settings.system_prompt = "overwritten".to_owned();
638        settings.load_profile("lab").expect("load profile");
639        assert_eq!(settings.ollama_model, "qwen3.5:9b");
640        assert_eq!(settings.system_prompt, "Cite primary sources only.");
641
642        settings.profile_name = "../escape".to_owned();
643        assert!(settings.save_profile().is_err());
644
645        let _ = std::fs::remove_dir_all(workspace);
646    }
647
648    #[test]
649    fn adding_a_corpus_replaces_the_same_name() {
650        let mut settings = fixture();
651        settings.add_corpus("papers", Path::new("/docs/a"));
652        settings.add_corpus("papers", Path::new("/docs/b"));
653        assert_eq!(settings.corpus_dirs, vec!["papers=/docs/b".to_owned()]);
654    }
655
656    fn unique_temp_dir(prefix: &str) -> PathBuf {
657        static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
658        let counter = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
659        let dir = std::env::temp_dir().join(format!("{prefix}-{}-{counter}", std::process::id()));
660        std::fs::create_dir_all(&dir).expect("create temp dir");
661        dir
662    }
663}
ragrig_gpui/rag/config.rs

ragrig_gpui/rag/
config.rs
