From 95bad975f4fb47fe86df1837660a1341009ea12a Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Wed, 26 Jun 2024 08:18:58 +0800
Subject: feat: custom rag document loaders (#650)

---
 src/config/mod.rs | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

(limited to 'src/config')

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 62ed28e..26c0bda 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -115,6 +115,8 @@ pub struct Config {
     pub rag_min_score_vector_search: f32,
     pub rag_min_score_keyword_search: f32,
     pub rag_min_score_rerank: f32,
+    #[serde(default)]
+    pub rag_document_loaders: HashMap<String, String>,
     pub rag_template: Option<String>,
 
     pub highlight: bool,
@@ -174,6 +176,7 @@ impl Default for Config {
             rag_min_score_vector_search: 0.0,
             rag_min_score_keyword_search: 0.0,
             rag_min_score_rerank: 0.0,
+            rag_document_loaders: Default::default(),
             rag_template: None,
 
             save_session: None,
@@ -229,6 +232,7 @@ impl Config {
         config.setup_model()?;
         config.setup_highlight();
         config.setup_light_theme()?;
+        config.setup_rag_document_loaders();
 
         Ok(config)
     }
@@ -1440,6 +1444,19 @@ impl Config {
         };
         Ok(())
     }
+
+    fn setup_rag_document_loaders(&mut self) {
+        [
+            ("pdf", "pdftotext $1 -"),
+            ("docx", "pandoc --to plain $1"),
+            ("url", "curl -fsSL $1"),
+        ]
+        .into_iter()
+        .for_each(|(k, v)| {
+            let (k, v) = (k.to_string(), v.to_string());
+            self.rag_document_loaders.entry(k).or_insert(v);
+        });
+    }
 }
 
 #[derive(Debug, Clone, Deserialize, Default)]
-- 
cgit v1.2.3