Archives/aichat
2
0
Fork 0
mirror of https://github.com/sigoden/aichat synced 2024-11-08 13:10:28 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: rag chunk contains path metadata (#647)

Browse Source
This commit is contained in:
sigoden 2024-06-25 12:55:38 +08:00 committed by GitHub
parent 2bc9607b00
commit 9f19108a71
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 9 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
src/rag/mod.rs
View File

@ -262,8 +262,10 @@ impl Rag {
);
let documents = load(&path, &extension)
.with_context(|| format!("Failed to load file at '{path}'"))?;
let documents =
splitter.split_documents(&documents, &SplitterChunkHeaderOptions::default());
let split_options = SplitterChunkHeaderOptions::default().with_chunk_header(&format!(
"<document_metadata>\npath: {path}\n</document_metadata>\n\n"
));
let documents = splitter.split_documents(&documents, &split_options);
rag_files.push(RagFile { path, documents });
progress(
&progress_tx,

20
src/rag/splitter/mod.rs
View File

@ -106,7 +106,6 @@ impl RecursiveCharacterTextSplitter {
let SplitterChunkHeaderOptions {
chunk_header,
chunk_overlap_header,
append_chunk_overlap_header,
} = chunk_header_options;
let mut documents = Vec::new();
@ -144,7 +143,7 @@ impl RecursiveCharacterTextSplitter {
Ordering::Equal => {}
}
if *append_chunk_overlap_header {
if let Some(chunk_overlap_header) = chunk_overlap_header {
page_content += chunk_overlap_header;
}
}
@ -288,16 +287,14 @@ impl RecursiveCharacterTextSplitter {
pub struct SplitterChunkHeaderOptions {
pub chunk_header: String,
pub chunk_overlap_header: String,
pub append_chunk_overlap_header: bool,
pub chunk_overlap_header: Option<String>,
}
impl Default for SplitterChunkHeaderOptions {
fn default() -> Self {
Self {
chunk_header: "".into(),
chunk_overlap_header: "(cont'd) ".into(),
append_chunk_overlap_header: false,
chunk_overlap_header: None,
}
}
}
@ -313,14 +310,7 @@ impl SplitterChunkHeaderOptions {
// Set the value of chunk_overlap_header
#[allow(unused)]
pub fn with_chunk_overlap_header(mut self, overlap_header: &str) -> Self {
self.chunk_overlap_header = overlap_header.to_string();
self
}
// Set the value of append_chunk_overlap_header
#[allow(unused)]
pub fn with_append_chunk_overlap_header(mut self, value: bool) -> Self {
self.append_chunk_overlap_header = value;
self.chunk_overlap_header = Some(overlap_header.to_string());
self
}
}
@ -414,7 +404,7 @@ mod tests {
let splitter = RecursiveCharacterTextSplitter::new(3, 0, &[" "]);
let chunk_header_options = SplitterChunkHeaderOptions::default()
.with_chunk_header("SOURCE NAME: testing\n-----\n")
.with_append_chunk_overlap_header(true);
.with_chunk_overlap_header("(cont'd) ");
let mut metadata1 = IndexMap::new();
metadata1.insert("source".into(), "1".into());
let mut metadata2 = IndexMap::new();