lc_rag/loaders/
csv.rs

1// src/retrieval/loaders/csv.rs
2//! CSV 文档加载器实现
3//!
4//! 提供给定列作为文档内容的方式来加载 CSV 文件。
5
6use super::{Document, DocumentLoader, LoaderError};
7use async_trait::async_trait;
8use csv::Reader;
9use std::fs::File;
10use std::io::BufReader;
11use std::path::PathBuf;
12
13/// CSV 文档加载器
14pub struct CSVLoader {
15    /// CSV 文件路径
16    pub path: PathBuf,
17
18    /// 作为文档内容的列名
19    pub content_column: String,
20}
21
22impl CSVLoader {
23    /// 创建新的 CSV 加载器
24    ///
25    /// # 参数
26    /// * `path` - CSV 文件路径
27    /// * `content_column` - 作为文档内容的列名
28    pub fn new(path: impl Into<PathBuf>, content_column: impl Into<String>) -> Self {
29        Self {
30            path: path.into(),
31            content_column: content_column.into(),
32        }
33    }
34}
35
36#[async_trait]
37impl DocumentLoader for CSVLoader {
38    async fn load(&self) -> Result<Vec<Document>, LoaderError> {
39        // 验证文件存在
40        if !self.path.exists() {
41            return Err(LoaderError::Other(format!(
42                "CSV 文件不存在: {}",
43                self.path.display()
44            )));
45        }
46
47        // 创建 CSV reader
48        let file = File::open(&self.path)?;
49        let buf_reader = BufReader::new(file);
50        let mut reader = Reader::from_reader(buf_reader);
51
52        let mut documents = Vec::new();
53
54        // 直接处理可能的错误
55        let headers_result = reader.headers();
56        let headers = match headers_result {
57            Ok(headers) => headers.clone(),
58            Err(e) => return Err(LoaderError::CsvError(e.to_string())),
59        };
60
61        for result in reader.records() {
62            let record = match result {
63                Ok(record) => record,
64                Err(e) => return Err(LoaderError::CsvError(e.to_string())),
65            };
66
67            // 查找内容列的索引
68            let content_idx = headers
69                .iter()
70                .position(|h| h == self.content_column.as_str());
71
72            if let Some(idx) = content_idx {
73                // 获取内容
74                let content = record.get(idx).unwrap_or_default().to_string();
75
76                // 如果内容為空則跳過此行
77                if content.is_empty() {
78                    continue;
79                }
80
81                // 创建文档内容，包含 CSV 行的所有列
82                let mut document = Document::new(content);
83
84                // 添加所有列的值作为元数据
85                for (i, header) in headers.iter().enumerate() {
86                    let value = record.get(i).unwrap_or_default().to_string();
87                    document = document.with_metadata(header.to_string(), value);
88                }
89
90                // 添加文件源信息
91                document =
92                    document.with_metadata("source".to_string(), self.path.display().to_string());
93                document = document.with_metadata("format".to_string(), "csv".to_string());
94                document = document
95                    .with_metadata("content_column".to_string(), self.content_column.clone());
96
97                documents.push(document);
98            } else {
99                // 如果内容列不存在，返回错误
100                return Err(LoaderError::CsvError(format!(
101                    "内容列 '{}' 在 CSV 文件中不存在",
102                    self.content_column
103                )));
104            }
105        }
106
107        Ok(documents)
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use std::io::Write;
115    use tempfile::NamedTempFile;
116
117    #[tokio::test]
118    async fn test_csv_loader_nonexistent() {
119        let loader = CSVLoader::new("./nonexistent.csv", "content");
120        let result = loader.load().await;
121
122        assert!(result.is_err());
123    }
124
125    #[tokio::test]
126    async fn test_csv_loader_content_column_not_found() {
127        let mut temp_file = NamedTempFile::new().unwrap();
128        writeln!(temp_file, "col1,col2").unwrap();
129        writeln!(temp_file, "val1,val2").unwrap();
130
131        let loader = CSVLoader::new(temp_file.path(), "content");
132        let result = loader.load().await;
133
134        assert!(result.is_err());
135        match result.unwrap_err() {
136            LoaderError::CsvError(msg) => assert!(msg.contains("不存在")),
137            _ => panic!("Expected CsvError"),
138        }
139    }
140
141    #[tokio::test]
142    async fn test_csv_loader_valid_data() {
143        let mut temp_file = NamedTempFile::new().unwrap();
144        writeln!(temp_file, "title,content,author").unwrap();
145        writeln!(temp_file, "Example Title,\"This is the content\",John Doe").unwrap();
146        writeln!(temp_file, "Another Title,\"More content\",Jane Smith").unwrap();
147
148        let loader = CSVLoader::new(temp_file.path(), "content");
149        let result = loader.load().await;
150
151        assert!(result.is_ok());
152        let docs = result.unwrap();
153        assert_eq!(docs.len(), 2);
154
155        // 检查第一个文档的内容和元数据
156        if !docs.is_empty() {
157            let doc = &docs[0];
158            assert!(doc.content.contains("This is the content"));
159            assert_eq!(
160                doc.metadata.get("title"),
161                Some(&"Example Title".to_string())
162            );
163            assert_eq!(doc.metadata.get("author"), Some(&"John Doe".to_string()));
164            assert_eq!(
165                doc.metadata.get("content"),
166                Some(&"This is the content".to_string())
167            );
168            assert_eq!(doc.metadata.get("format"), Some(&"csv".to_string()));
169            assert_eq!(
170                doc.metadata.get("content_column"),
171                Some(&"content".to_string())
172            );
173        }
174    }
175}
lc_rag/loaders/csv.rs

lc_rag/loaders/
csv.rs
