use anyhow::Result;
use crate::DocumentSplitter;
use crate::Documents;
use crate::TextSplitter;
pub struct MarkdownSplitter {
pub splitter_chunk_size: usize,
}
impl MarkdownSplitter {
pub fn create() -> Self {
MarkdownSplitter {
splitter_chunk_size: 400,
}
}
pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
self.splitter_chunk_size = chunk_size;
self
}
}
impl DocumentSplitter for MarkdownSplitter {
fn separators(&self) -> Vec<String> {
vec![
"\n## ".to_string(),
"\n### ".to_string(),
"\n#### ".to_string(),
"\n##### ".to_string(),
"\n###### ".to_string(),
]
}
fn split_documents(&self, documents: &Documents) -> Result<Documents> {
let text_splitter = TextSplitter::create()
.with_chunk_size(self.splitter_chunk_size)
.with_separators(self.separators());
text_splitter.split_documents(documents)
}
}