Skip to main content

zeph_memory/document/loader/
text.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use 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}