langchain_rust/text_splitter/
text_splitter.rs

1use std::collections::HashMap;
2
3use async_trait::async_trait;
4use serde_json::Value;
5
6use crate::schemas::Document;
7
8use super::TextSplitterError;
9
10#[async_trait]
11pub trait TextSplitter: Send + Sync {
12    async fn split_text(&self, text: &str) -> Result<Vec<String>, TextSplitterError>;
13
14    async fn split_documents(
15        &self,
16        documents: &[Document],
17    ) -> Result<Vec<Document>, TextSplitterError> {
18        let mut texts: Vec<String> = Vec::new();
19        let mut metadatas: Vec<HashMap<String, Value>> = Vec::new();
20        documents.iter().for_each(|d| {
21            texts.push(d.page_content.clone());
22            metadatas.push(d.metadata.clone());
23        });
24
25        self.create_documents(&texts, &metadatas).await
26    }
27
28    async fn create_documents(
29        &self,
30        text: &[String],
31        metadatas: &[HashMap<String, Value>],
32    ) -> Result<Vec<Document>, TextSplitterError> {
33        let mut metadatas = metadatas.to_vec();
34        if metadatas.is_empty() {
35            metadatas = vec![HashMap::new(); text.len()];
36        }
37
38        if text.len() != metadatas.len() {
39            return Err(TextSplitterError::MetadataTextMismatch);
40        }
41
42        let mut documents: Vec<Document> = Vec::new();
43        for i in 0..text.len() {
44            let chunks = self.split_text(&text[i]).await?;
45            for chunk in chunks {
46                let document = Document::new(chunk).with_metadata(metadatas[i].clone());
47                documents.push(document);
48            }
49        }
50
51        Ok(documents)
52    }
53}
langchain_rust/text_splitter/text_splitter.rs

langchain_rust/text_splitter/
text_splitter.rs
