lc_rag/loaders/
text.rs

1// src/retrieval/loaders/text.rs
2//! Text document loader implementation
3//!
4//! Provides loading content from plain-text files.
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use std::path::PathBuf;
9
10/// Text document loader
11///
12/// Supports loading plain-text files (.txt), treating the entire file content as one document.
13pub struct TextLoader {
14    /// Text file path
15    pub path: PathBuf,
16
17    /// Whether to split by line (optional)
18    /// If true, each line is returned as a separate document
19    pub split_by_line: bool,
20}
21
22impl TextLoader {
23    /// Creates a new Text loader
24    ///
25    /// # Arguments
26    /// * `path` - the text file path
27    pub fn new(path: impl Into<PathBuf>) -> Self {
28        Self {
29            path: path.into(),
30            split_by_line: false,
31        }
32    }
33
34    /// Creates a line-splitting Text loader
35    ///
36    /// Each line is returned as a separate document.
37    ///
38    /// # Arguments
39    /// * `path` - the text file path
40    pub fn new_with_line_split(path: impl Into<PathBuf>) -> Self {
41        Self {
42            path: path.into(),
43            split_by_line: true,
44        }
45    }
46
47    /// Sets whether to split by line
48    pub fn with_split_by_line(mut self, split: bool) -> Self {
49        self.split_by_line = split;
50        self
51    }
52}
53
54#[async_trait]
55impl DocumentLoader for TextLoader {
56    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
57        // Verify the file exists
58        if !self.path.exists() {
59            return Err(LoaderError::Other(format!(
60                "text file does not exist: {}",
61                self.path.display()
62            )));
63        }
64
65        // Read the file content
66        let content = std::fs::read_to_string(&self.path)?;
67
68        if self.split_by_line {
69            // Split by line
70            let lines: Vec<&str> = content.lines().filter(|l| !l.trim().is_empty()).collect();
71            let documents = lines
72                .iter()
73                .enumerate()
74                .map(|(idx, line)| {
75                    let mut doc = Document::new(line.to_string());
76                    doc = doc.with_metadata("source".to_string(), self.path.display().to_string());
77                    doc = doc.with_metadata("format".to_string(), "text".to_string());
78                    doc = doc.with_metadata("line_number".to_string(), (idx + 1).to_string());
79                    doc
80                })
81                .collect();
82
83            Ok(documents)
84        } else {
85            // Treat the entire file as one document
86            let mut document = Document::new(content);
87            document =
88                document.with_metadata("source".to_string(), self.path.display().to_string());
89            document = document.with_metadata("format".to_string(), "text".to_string());
90
91            Ok(vec![document])
92        }
93    }
94}
95
96#[cfg(test)]
97mod tests {
98    use super::*;
99    use std::io::Write;
100    use tempfile::NamedTempFile;
101
102    #[tokio::test]
103    async fn test_text_loader_nonexistent() {
104        let loader = TextLoader::new("./nonexistent.txt");
105        let result = loader.load().await;
106
107        assert!(result.is_err());
108        match result.unwrap_err() {
109            LoaderError::Other(msg) => assert!(msg.contains("does not exist")),
110            _ => panic!("Expected Other error"),
111        }
112    }
113
114    #[tokio::test]
115    async fn test_text_loader_single_document() {
116        let mut temp_file = NamedTempFile::new().unwrap();
117        write!(temp_file, "Hello, World!\nThis is a test.").unwrap();
118
119        let loader = TextLoader::new(temp_file.path());
120        let result = loader.load().await;
121
122        assert!(result.is_ok());
123        let docs = result.unwrap();
124        assert_eq!(docs.len(), 1);
125        assert!(docs[0].content.contains("Hello, World!"));
126        assert_eq!(
127            docs[0].metadata.get("format"),
128            Some(&serde_json::Value::String("text".to_string()))
129        );
130    }
131
132    #[tokio::test]
133    async fn test_text_loader_split_by_line() {
134        let mut temp_file = NamedTempFile::new().unwrap();
135        writeln!(temp_file, "Line 1").unwrap();
136        writeln!(temp_file, "Line 2").unwrap();
137        writeln!(temp_file, "Line 3").unwrap();
138
139        let loader = TextLoader::new_with_line_split(temp_file.path());
140        let result = loader.load().await;
141
142        assert!(result.is_ok());
143        let docs = result.unwrap();
144        assert_eq!(docs.len(), 3);
145        assert_eq!(docs[0].content, "Line 1");
146        assert_eq!(
147            docs[0].metadata.get("line_number"),
148            Some(&serde_json::Value::String("1".to_string()))
149        );
150    }
151
152    #[tokio::test]
153    async fn test_text_loader_skip_empty_lines() {
154        let mut temp_file = NamedTempFile::new().unwrap();
155        writeln!(temp_file, "Line 1").unwrap();
156        writeln!(temp_file).unwrap();
157        writeln!(temp_file, "   ").unwrap();
158        writeln!(temp_file, "Line 2").unwrap();
159
160        let loader = TextLoader::new_with_line_split(temp_file.path());
161        let result = loader.load().await;
162
163        assert!(result.is_ok());
164        let docs = result.unwrap();
165        assert_eq!(docs.len(), 2); // empty lines are skipped
166    }
167}
lc_rag/loaders/text.rs

lc_rag/loaders/
text.rs
