summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--src/rag/loader.rs67
-rw-r--r--src/rag/mod.rs25
2 files changed, 56 insertions, 36 deletions
diff --git a/src/rag/loader.rs b/src/rag/loader.rs
index 9f70c8b..99bc489 100644
--- a/src/rag/loader.rs
+++ b/src/rag/loader.rs
@@ -9,6 +9,7 @@ use tokio::io::AsyncWriteExt;
pub const RECURSIVE_URL_LOADER: &str = "recursive_url";
pub const URL_LOADER: &str = "url";
+pub const EXTENSION_METADATA: &str = "__extension__";
lazy_static! {
static ref CLIENT: Result<reqwest::Client> = {
@@ -22,35 +23,35 @@ lazy_static! {
pub async fn load(
loaders: &HashMap<String, String>,
path: &str,
- loader_name: &str,
+ extension: &str,
) -> Result<Vec<RagDocument>> {
- if loader_name == RECURSIVE_URL_LOADER {
+ if extension == RECURSIVE_URL_LOADER {
let loader_command = loaders
- .get(loader_name)
- .with_context(|| format!("RAG document loader '{loader_name}' not configured"))?;
- let contents = run_loader_command(path, loader_name, loader_command)?;
+ .get(extension)
+ .with_context(|| format!("RAG document loader '{extension}' not configured"))?;
+ let contents = run_loader_command(path, extension, loader_command)?;
let output = match parse_json_documents(&contents) {
Some(v) => v,
None => vec![RagDocument::new(contents)],
};
Ok(output)
} else {
- match loaders.get(loader_name) {
- Some(loader_command) => load_with_command(path, loader_name, loader_command),
+ match loaders.get(extension) {
+ Some(loader_command) => load_with_command(path, extension, loader_command),
None => {
- if loader_name == URL_LOADER {
+ if extension == URL_LOADER {
load_url(loaders, path).await
} else {
- load_plain(path, loader_name).await
+ load_plain(path, extension).await
}
}
}
}
}
-async fn load_plain(path: &str, loader_name: &str) -> Result<Vec<RagDocument>> {
+async fn load_plain(path: &str, extension: &str) -> Result<Vec<RagDocument>> {
let contents = tokio::fs::read_to_string(path).await?;
- if loader_name == "json" {
+ if extension == "json" {
if let Some(documents) = parse_json_documents(&contents) {
return Ok(documents);
}
@@ -66,43 +67,55 @@ async fn load_url(loaders: &HashMap<String, String>, path: &str) -> Result<Vec<R
Err(ref err) => bail!("{err}"),
};
let mut res = client.get(path).send().await?;
- let loader_name = path
+
+ let mut metadata: RagMetadata = Default::default();
+ metadata.insert("path".into(), path.to_string());
+
+ let extension = path
.rsplit_once('/')
.and_then(|(_, pair)| pair.rsplit_once('.').map(|(_, ext)| ext))
.unwrap_or("txt");
- let contents = match loaders.get(loader_name) {
+ let extension = extension.to_lowercase();
+ let document = match loaders.get(&extension) {
Some(loader_command) => {
let save_path = env::temp_dir()
- .join(format!("aichat-download-{}.{loader_name}", sha256(path)))
+ .join(format!("aichat-download-{}.{extension}", sha256(path)))
.display()
.to_string();
let mut save_file = tokio::fs::File::create(&save_path).await?;
while let Some(chunk) = res.chunk().await? {
save_file.write_all(&chunk).await?;
}
- run_loader_command(&save_path, loader_name, loader_command)?
+ let contents = run_loader_command(&save_path, &extension, loader_command)?;
+ metadata.insert(EXTENSION_METADATA.into(), "txt".to_string());
+ RagDocument::new(contents).with_metadata(metadata)
+ }
+ None => {
+ let contents = res.text().await?;
+ metadata.insert(EXTENSION_METADATA.into(), extension);
+ RagDocument::new(contents).with_metadata(metadata)
}
- None => res.text().await?,
};
- let mut document = RagDocument::new(contents);
- document.metadata.insert("path".into(), path.to_string());
Ok(vec![document])
}
fn load_with_command(
path: &str,
- loader_name: &str,
+ extension: &str,
loader_command: &str,
) -> Result<Vec<RagDocument>> {
- let contents = run_loader_command(path, loader_name, loader_command)?;
+ let contents = run_loader_command(path, extension, loader_command)?;
let mut document = RagDocument::new(contents);
document.metadata.insert("path".into(), path.to_string());
+ document
+ .metadata
+ .insert(EXTENSION_METADATA.into(), "txt".to_string());
Ok(vec![document])
}
-fn run_loader_command(path: &str, loader_name: &str, loader_command: &str) -> Result<String> {
+fn run_loader_command(path: &str, extension: &str, loader_command: &str) -> Result<String> {
let cmd_args = shell_words::split(loader_command).with_context(|| {
- anyhow!("Invalid rag document loader '{loader_name}': `{loader_command}`")
+ anyhow!("Invalid rag document loader '{extension}': `{loader_command}`")
})?;
let mut use_stdout = true;
let outpath = env::temp_dir()
@@ -184,7 +197,6 @@ fn parse_json_documents(data: &str) -> Option<Vec<RagDocument>> {
"html",
"markdown",
"text",
- "data",
]
.into_iter()
.map(|v| v.to_string())
@@ -196,7 +208,7 @@ fn parse_json_documents(data: &str) -> Option<Vec<RagDocument>> {
if let Some(page_content) = obj.get(&key).and_then(|v| v.as_str()) {
let page_content = page_content.to_string();
obj.remove(&key);
- let metadata: IndexMap<_, _> = obj
+ let mut metadata: IndexMap<_, _> = obj
.into_iter()
.map(|(k, v)| {
if let Value::String(v) = v {
@@ -206,6 +218,11 @@ fn parse_json_documents(data: &str) -> Option<Vec<RagDocument>> {
}
})
.collect();
+ if key == "markdown" {
+ metadata.insert(EXTENSION_METADATA.into(), "md".into());
+ } else if key == "html" {
+ metadata.insert(EXTENSION_METADATA.into(), "html".into());
+ }
return Some(RagDocument {
page_content,
metadata,
@@ -337,7 +354,7 @@ mod tests {
let mut metadata = IndexMap::new();
metadata.insert("k1".into(), "1".into());
- let data = r#"[{"k1": 1, "data": "foo" }]"#;
+ let data = r#"[{"k1": 1, "text": "foo" }]"#;
assert_eq!(
parse_json_documents(data).unwrap(),
vec![RagDocument::new("foo").with_metadata(metadata.clone())]
diff --git a/src/rag/mod.rs b/src/rag/mod.rs
index 9e37e6f..536092f 100644
--- a/src/rag/mod.rs
+++ b/src/rag/mod.rs
@@ -272,20 +272,24 @@ impl Rag {
if let Some(spinner) = &spinner {
let _ = spinner.set_message(String::new());
}
- for (index, (path, loader_name)) in new_paths.into_iter().enumerate() {
+ for (index, (path, extension)) in new_paths.into_iter().enumerate() {
println!("Loading {path} [{}/{new_paths_len}]", index + 1);
- let documents = load(&loaders, &path, &loader_name)
+ let documents = load(&loaders, &path, &extension)
.await
.with_context(|| format!("Failed to load '{path}'"))?;
- let separator = get_separators(&loader_name);
- let splitter = RecursiveCharacterTextSplitter::new(
- self.data.chunk_size,
- self.data.chunk_overlap,
- &separator,
- );
let splitted_documents: Vec<_> = documents
.into_iter()
- .flat_map(|document| {
+ .flat_map(|mut document| {
+ let extension = document
+ .metadata
+ .swap_remove(EXTENSION_METADATA)
+ .unwrap_or_else(|| extension.clone());
+ let separator = get_separators(&extension);
+ let splitter = RecursiveCharacterTextSplitter::new(
+ self.data.chunk_size,
+ self.data.chunk_overlap,
+ &separator,
+ );
let metadata = document
.metadata
.iter()
@@ -299,7 +303,7 @@ impl Rag {
splitter.split_documents(&[document], &split_options)
})
.collect();
- let display_path = if loader_name == RECURSIVE_URL_LOADER {
+ let display_path = if extension == RECURSIVE_URL_LOADER {
format!("{path}**")
} else {
path
@@ -557,7 +561,6 @@ impl RagDocument {
}
}
- #[allow(unused)]
pub fn with_metadata(mut self, metadata: RagMetadata) -> Self {
self.metadata = metadata;
self
generated by cgit v1.2.3 (git 2.39.1) at 2026-10-03 01:22:47 +0000