langchain_rust/text_splitter/
plain_text_splitter.rs

1use async_trait::async_trait;
2
3use super::{TextSplitter, TextSplitterError};
4
5// Options is a struct that contains options for a plain text splitter.
6#[derive(Debug, Clone)]
7pub struct PlainTextSplitterOptions {
8    pub chunk_size: usize,
9    pub chunk_overlap: usize,
10    pub trim_chunks: bool,
11}
12
13impl Default for PlainTextSplitterOptions {
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl PlainTextSplitterOptions {
20    pub fn new() -> Self {
21        PlainTextSplitterOptions {
22            chunk_size: 512,
23            chunk_overlap: 0,
24            trim_chunks: false,
25        }
26    }
27
28    pub fn with_chunk_size(mut self, chunk_size: usize) -> Self {
29        self.chunk_size = chunk_size;
30        self
31    }
32
33    pub fn with_chunk_overlap(mut self, chunk_overlap: usize) -> Self {
34        self.chunk_overlap = chunk_overlap;
35        self
36    }
37
38    pub fn with_trim_chunks(mut self, trim_chunks: bool) -> Self {
39        self.trim_chunks = trim_chunks;
40        self
41    }
42
43    pub fn chunk_size(&self) -> usize {
44        self.chunk_size
45    }
46
47    pub fn chunk_overlap(&self) -> usize {
48        self.chunk_overlap
49    }
50
51    pub fn trim_chunks(&self) -> bool {
52        self.trim_chunks
53    }
54}
55
56pub struct PlainTextSplitter {
57    splitter_options: PlainTextSplitterOptions,
58}
59
60impl Default for PlainTextSplitter {
61    fn default() -> Self {
62        PlainTextSplitter::new(PlainTextSplitterOptions::default())
63    }
64}
65
66impl PlainTextSplitter {
67    pub fn new(options: PlainTextSplitterOptions) -> PlainTextSplitter {
68        PlainTextSplitter {
69            splitter_options: options,
70        }
71    }
72}
73
74#[async_trait]
75impl TextSplitter for PlainTextSplitter {
76    async fn split_text(&self, text: &str) -> Result<Vec<String>, TextSplitterError> {
77        let splitter = text_splitter::TextSplitter::new(
78            text_splitter::ChunkConfig::new(self.splitter_options.chunk_size)
79                .with_trim(self.splitter_options.trim_chunks)
80                .with_overlap(self.splitter_options.chunk_overlap)?,
81        );
82
83        Ok(splitter.chunks(text).map(|x| x.to_string()).collect())
84    }
85}
langchain_rust/text_splitter/plain_text_splitter.rs

langchain_rust/text_splitter/
plain_text_splitter.rs
