zeph_memory/document/loader/
text.rs1use std::collections::HashMap;
5use std::path::Path;
6use std::pin::Pin;
7
8use super::super::{
9 DEFAULT_MAX_FILE_SIZE, Document, DocumentError, DocumentLoader, DocumentMetadata,
10};
11
12pub struct TextLoader {
13 pub max_file_size: u64,
14}
15
16impl Default for TextLoader {
17 fn default() -> Self {
18 Self {
19 max_file_size: DEFAULT_MAX_FILE_SIZE,
20 }
21 }
22}
23
24impl DocumentLoader for TextLoader {
25 fn load(
26 &self,
27 path: &Path,
28 ) -> Pin<Box<dyn std::future::Future<Output = Result<Vec<Document>, DocumentError>> + Send + '_>>
29 {
30 let path = path.to_path_buf();
31 let max_size = self.max_file_size;
32 Box::pin(async move {
33 let path = tokio::fs::canonicalize(&path).await?;
34
35 let meta = tokio::fs::metadata(&path).await?;
36 if meta.len() > max_size {
37 return Err(DocumentError::FileTooLarge(meta.len()));
38 }
39
40 let ext = path.extension().and_then(|e| e.to_str()).unwrap_or("");
41
42 let content_type = match ext {
43 "md" | "markdown" => "text/markdown",
44 _ => "text/plain",
45 };
46
47 let content = tokio::fs::read_to_string(&path).await?;
48
49 Ok(vec![Document {
50 content,
51 metadata: DocumentMetadata {
52 source: path.display().to_string(),
53 content_type: content_type.to_owned(),
54 extra: HashMap::new(),
55 },
56 }])
57 })
58 }
59
60 fn supported_extensions(&self) -> &[&str] {
61 &["txt", "md", "markdown"]
62 }
63}
64
65#[cfg(test)]
66mod tests {
67 use super::*;
68 use std::assert_matches;
69
70 #[tokio::test]
71 async fn load_text_file() {
72 let dir = tempfile::tempdir().unwrap();
73 let file = dir.path().join("test.txt");
74 std::fs::write(&file, "hello world").unwrap();
75
76 let docs = TextLoader::default().load(&file).await.unwrap();
77 assert_eq!(docs.len(), 1);
78 assert_eq!(docs[0].content, "hello world");
79 assert_eq!(docs[0].metadata.content_type, "text/plain");
80 }
81
82 #[tokio::test]
83 async fn load_markdown_file() {
84 let dir = tempfile::tempdir().unwrap();
85 let file = dir.path().join("readme.md");
86 std::fs::write(&file, "# Title").unwrap();
87
88 let docs = TextLoader::default().load(&file).await.unwrap();
89 assert_eq!(docs[0].metadata.content_type, "text/markdown");
90 }
91
92 #[tokio::test]
93 async fn load_nonexistent_file() {
94 let result = TextLoader::default()
95 .load(Path::new("/nonexistent/file.txt"))
96 .await;
97 assert!(result.is_err());
98 }
99
100 #[tokio::test]
101 async fn load_empty_file() {
102 let dir = tempfile::tempdir().unwrap();
103 let file = dir.path().join("empty.txt");
104 std::fs::write(&file, "").unwrap();
105
106 let docs = TextLoader::default().load(&file).await.unwrap();
107 assert_eq!(docs.len(), 1);
108 assert!(docs[0].content.is_empty());
109 }
110
111 #[tokio::test]
112 async fn load_markdown_extension_variant() {
113 let dir = tempfile::tempdir().unwrap();
114 let file = dir.path().join("doc.markdown");
115 std::fs::write(&file, "content").unwrap();
116
117 let docs = TextLoader::default().load(&file).await.unwrap();
118 assert_eq!(docs[0].metadata.content_type, "text/markdown");
119 }
120
121 #[tokio::test]
122 async fn unknown_extension_treated_as_plain_text() {
123 let dir = tempfile::tempdir().unwrap();
124 let file = dir.path().join("data.csv");
125 std::fs::write(&file, "a,b,c").unwrap();
126
127 let docs = TextLoader::default().load(&file).await.unwrap();
128 assert_eq!(docs[0].metadata.content_type, "text/plain");
129 }
130
131 #[test]
132 fn supported_extensions_list() {
133 let loader = TextLoader::default();
134 let exts = loader.supported_extensions();
135 assert!(exts.contains(&"txt"));
136 assert!(exts.contains(&"md"));
137 assert!(exts.contains(&"markdown"));
138 }
139
140 #[tokio::test]
141 async fn metadata_source_is_canonical() {
142 let dir = tempfile::tempdir().unwrap();
143 let file = dir.path().join("test.txt");
144 std::fs::write(&file, "data").unwrap();
145
146 let docs = TextLoader::default().load(&file).await.unwrap();
147 let canonical = std::fs::canonicalize(&file).unwrap();
148 assert_eq!(docs[0].metadata.source, canonical.display().to_string());
149 }
150
151 #[tokio::test]
152 async fn file_too_large_rejected() {
153 let dir = tempfile::tempdir().unwrap();
154 let file = dir.path().join("big.txt");
155 std::fs::write(&file, "x").unwrap();
156
157 let loader = TextLoader { max_file_size: 0 };
158 let result = loader.load(&file).await;
159 assert_matches!(result, Err(DocumentError::FileTooLarge(_)));
160 }
161}