Skip to main content

docling_rag/source/
folder.rs

1//! Local folder document source. Recursively lists files under a root directory;
2//! works over any mounted filesystem (local disk, NFS, FUSE, …).
3
4use super::{DocumentSource, SourceRef};
5use crate::Result;
6use async_trait::async_trait;
7use std::path::{Path, PathBuf};
8
9/// Reads documents from a directory tree.
10#[derive(Debug, Clone)]
11pub struct FolderSource {
12    root: PathBuf,
13}
14
15impl FolderSource {
16    /// Create a source rooted at `root`.
17    pub fn new(root: impl AsRef<Path>) -> Self {
18        FolderSource {
19            root: root.as_ref().to_path_buf(),
20        }
21    }
22}
23
24/// Recursively collect regular files under `dir` (depth-first, sorted for
25/// determinism). Unreadable directories are skipped rather than failing the walk.
26fn walk(dir: &Path, out: &mut Vec<PathBuf>) {
27    let Ok(entries) = std::fs::read_dir(dir) else {
28        return;
29    };
30    let mut paths: Vec<PathBuf> = entries.filter_map(|e| e.ok().map(|e| e.path())).collect();
31    paths.sort();
32    for path in paths {
33        if path.is_dir() {
34            walk(&path, out);
35        } else if path.is_file() {
36            out.push(path);
37        }
38    }
39}
40
41#[async_trait]
42impl DocumentSource for FolderSource {
43    async fn list(&self) -> Result<Vec<SourceRef>> {
44        let root = self.root.clone();
45        // Directory walking is blocking I/O; keep it off the async reactor.
46        let files = tokio::task::spawn_blocking(move || {
47            let mut out = Vec::new();
48            walk(&root, &mut out);
49            out
50        })
51        .await
52        .map_err(|e| crate::RagError::Source(format!("walk join: {e}")))?;
53
54        Ok(files
55            .into_iter()
56            .map(|p| SourceRef {
57                name: p
58                    .file_name()
59                    .map(|n| n.to_string_lossy().into_owned())
60                    .unwrap_or_default(),
61                rel_path: p
62                    .strip_prefix(&self.root)
63                    .unwrap_or(&p)
64                    .to_string_lossy()
65                    .into_owned(),
66                uri: format!("file://{}", p.display()),
67            })
68            .collect())
69    }
70
71    async fn fetch(&self, r: &SourceRef) -> Result<Vec<u8>> {
72        let path = r.uri.strip_prefix("file://").unwrap_or(&r.uri);
73        Ok(tokio::fs::read(path).await?)
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    #[tokio::test]
82    async fn lists_and_fetches_files_recursively() {
83        let dir = std::env::temp_dir().join(format!("rag-folder-{}", crate::model::new_id()));
84        std::fs::create_dir_all(dir.join("sub")).unwrap();
85        std::fs::write(dir.join("a.md"), b"# A\n\nalpha").unwrap();
86        std::fs::write(dir.join("sub/b.md"), b"# B\n\nbeta").unwrap();
87
88        let src = FolderSource::new(&dir);
89        let refs = src.list().await.unwrap();
90        assert_eq!(refs.len(), 2);
91        let bytes = src.fetch(&refs[0]).await.unwrap();
92        assert!(bytes.starts_with(b"# "));
93
94        std::fs::remove_dir_all(&dir).ok();
95    }
96}