lc_rag/loaders/
text.rs

1// src/retrieval/loaders/text.rs
2//! Text 文档加载器实现
3//!
4//! 提供从纯文本文件加载内容的功能。
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use std::path::PathBuf;
9
10/// Text 文档加载器
11///
12/// 支持加载纯文本文件（.txt），将整个文件内容作为一个文档。
13pub struct TextLoader {
14    /// 文本文件路径
15    pub path: PathBuf,
16
17    /// 是否按行分割（可选）
18    /// 如果为 true，每行作为一个独立文档
19    pub split_by_line: bool,
20}
21
22impl TextLoader {
23    /// 创建新的 Text 加载器
24    ///
25    /// # 参数
26    /// * `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    /// 创建按行分割的 Text 加载器
35    ///
36    /// 每行文本将作为独立文档返回。
37    ///
38    /// # 参数
39    /// * `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    /// 设置是否按行分割
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        // 验证文件存在
58        if !self.path.exists() {
59            return Err(LoaderError::Other(format!(
60                "文本文件不存在: {}",
61                self.path.display()
62            )));
63        }
64
65        // 读取文件内容
66        let content = std::fs::read_to_string(&self.path)?;
67
68        if self.split_by_line {
69            // 按行分割
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            // 整个文件作为一个文档
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("不存在")),
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!(docs[0].metadata.get("format"), Some(&"text".to_string()));
127    }
128
129    #[tokio::test]
130    async fn test_text_loader_split_by_line() {
131        let mut temp_file = NamedTempFile::new().unwrap();
132        writeln!(temp_file, "Line 1").unwrap();
133        writeln!(temp_file, "Line 2").unwrap();
134        writeln!(temp_file, "Line 3").unwrap();
135
136        let loader = TextLoader::new_with_line_split(temp_file.path());
137        let result = loader.load().await;
138
139        assert!(result.is_ok());
140        let docs = result.unwrap();
141        assert_eq!(docs.len(), 3);
142        assert_eq!(docs[0].content, "Line 1");
143        assert_eq!(docs[0].metadata.get("line_number"), Some(&"1".to_string()));
144    }
145
146    #[tokio::test]
147    async fn test_text_loader_skip_empty_lines() {
148        let mut temp_file = NamedTempFile::new().unwrap();
149        writeln!(temp_file, "Line 1").unwrap();
150        writeln!(temp_file).unwrap();
151        writeln!(temp_file, "   ").unwrap();
152        writeln!(temp_file, "Line 2").unwrap();
153
154        let loader = TextLoader::new_with_line_split(temp_file.path());
155        let result = loader.load().await;
156
157        assert!(result.is_ok());
158        let docs = result.unwrap();
159        assert_eq!(docs.len(), 2); // 空行被跳过
160    }
161}
lc_rag/loaders/text.rs

lc_rag/loaders/
text.rs
