docling_rag/source/
folder.rs1use super::{DocumentSource, SourceRef};
5use crate::Result;
6use async_trait::async_trait;
7use std::path::{Path, PathBuf};
8
9#[derive(Debug, Clone)]
11pub struct FolderSource {
12 root: PathBuf,
13}
14
15impl FolderSource {
16 pub fn new(root: impl AsRef<Path>) -> Self {
18 FolderSource {
19 root: root.as_ref().to_path_buf(),
20 }
21 }
22}
23
24fn 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 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}