From f82524fd154a1bf4a7277db868147afdb0cd507a Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 27 Jun 2024 12:30:09 +0800
Subject: feat: implement native rag url loader (#660)

---
 src/config/input.rs | 13 ++++++++++---
 src/config/mod.rs   | 16 ++++++----------
 2 files changed, 16 insertions(+), 13 deletions(-)

(limited to 'src/config')

diff --git a/src/config/input.rs b/src/config/input.rs
index a41b1c7..3b15d1f 100644
--- a/src/config/input.rs
+++ b/src/config/input.rs
@@ -10,7 +10,6 @@ use crate::utils::{base64_encode, sha256, AbortSignal};
 use anyhow::{bail, Context, Result};
 use fancy_regex::Regex;
 use lazy_static::lazy_static;
-use mime_guess::from_path;
 use std::{
     collections::HashMap,
     fs::File,
@@ -407,8 +406,16 @@ fn is_image_ext(path: &Path) -> bool {
 
 fn read_media_to_data_url<P: AsRef<Path>>(image_path: P) -> Result<String> {
     let image_path = image_path.as_ref();
-
-    let mime_type = from_path(image_path).first_or_octet_stream().to_string();
+    let mime_type = match image_path.extension().and_then(|v| v.to_str()) {
+        Some(extension) => match extension {
+            "png" => "image/png",
+            "jpg" | "jpeg" => "image/jpeg",
+            "webp" => "image/webp",
+            "gif" => "image/gif",
+            _ => bail!("Unsupported media type"),
+        },
+        None => bail!("Unknown media type"),
+    };
     let mut file = File::open(image_path)?;
     let mut buffer = Vec::new();
     file.read_to_end(&mut buffer)?;
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 69609f4..ff337c0 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -1434,16 +1434,12 @@ impl Config {
     }
 
     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);
-        });
+        [("pdf", "pdftotext $1 -"), ("docx", "pandoc --to plain $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);
+            });
     }
 }
 
-- 
cgit v1.2.3