src/documents.rs

@@ -20,6 +20,25 @@ pub(crate) struct Document {
     pub(crate) text: String,
 }
 
+impl Document {
+    /// Builds a `Document` for the `index`-th of `total` chunks split from
+    /// `path`, labeling multi-chunk files with a "(part N/total)" suffix so
+    /// retrieved excerpts can be traced back to their position in the
+    /// source file.
+    fn from_chunk(path: &Path, index: usize, chunk: &str, total: usize) -> Self {
+        let source = if total > 1 {
+            format!("{} (part {}/{total})", path.display(), index + 1)
+        } else {
+            path.display().to_string()
+        };
+
+        Self {
+            source,
+            text: chunk.to_string(),
+        }
+    }
+}
+
 impl rig::Embed for Document {
     fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
         embedder.embed(self.text.clone());
@@ -37,42 +56,48 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Documen
     let splitter = TextSplitter::new(CHUNK_CHARS);
     let mut documents = Vec::new();
 
-    for path in paths {
-        if path.is_dir() {
-            for entry in std::fs::read_dir(path)? {
+    let (directories, files): (Vec<&PathBuf>, Vec<&PathBuf>) =
+        paths.iter().partition(|path| path.is_dir());
+
+    for dir in directories {
+        std::fs::read_dir(dir)?
+            .map(|entry| -> anyhow::Result<Option<PathBuf>> {
                 let entry = entry?;
-                if entry.file_type()?.is_file() {
-                    push_document(&entry.path(), &splitter, &mut documents);
-                }
-            }
-        } else {
-            push_document(path, &splitter, &mut documents);
-        }
+                Ok(entry.file_type()?.is_file().then(|| entry.path()))
+            })
+            .filter_map(Result::transpose)
+            .collect::<anyhow::Result<Vec<PathBuf>>>()?
+            .iter()
+            .for_each(|path| documents.extend(chunk_document(path, &splitter)));
     }
 
+    files
+        .iter()
+        .for_each(|file| documents.extend(chunk_document(file, &splitter)));
+
     Ok(documents)
 }
 
-fn push_document(path: &Path, splitter: &TextSplitter<Characters>, documents: &mut Vec<Document>) {
+/// Reads `path` and splits it into chunked `Document`s via `splitter`. An
+/// unreadable path (permissions, non-UTF-8, gone by the time it's read)
+/// yields an empty `Vec` and a warning rather than failing the caller.
+fn chunk_document(path: &Path, splitter: &TextSplitter<Characters>) -> Vec<Document> {
     match std::fs::read_to_string(path) {
         Ok(text) => {
             let chunks: Vec<&str> = splitter.chunks(&text).collect();
-            let total = chunks.len();
-            let source = |i: usize| {
-                if total > 1 {
-                    format!("{} (part {}/{total})", path.display(), i + 1)
-                } else {
-                    path.display().to_string()
-                }
-            };
+            let chunk_count = chunks.len();
 
-            documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document {
-                source: source(i),
-                text: chunk.to_string(),
-            }));
+            chunks
+                .into_iter()
+                .enumerate()
+                .map(|(chunk_index, chunk)| {
+                    Document::from_chunk(path, chunk_index, chunk, chunk_count)
+                })
+                .collect()
         }
         Err(err) => {
-            tracing::warn!(path = %path.display(), %err, "skipping unreadable document")
+            tracing::warn!(path = %path.display(), %err, "skipping unreadable document");
+            Vec::new()
         }
     }
 }
@@ -160,4 +185,137 @@ mod tests {
             vec!["the text to embed".to_string()]
         );
     }
+
+    #[test]
+    fn collect_documents_reads_multiple_explicit_files() {
+        let dir = tempdir().unwrap();
+        let a = dir.path().join("a.txt");
+        let b = dir.path().join("b.txt");
+        std::fs::write(&a, "a").unwrap();
+        std::fs::write(&b, "b").unwrap();
+
+        let mut documents = collect_documents(&[a, b]).unwrap();
+        documents.sort_by(|x, y| x.text.cmp(&y.text));
+
+        assert_eq!(documents.len(), 2);
+        assert_eq!(documents[0].text, "a");
+        assert_eq!(documents[1].text, "b");
+    }
+
+    #[test]
+    fn collect_documents_reads_multiple_directories() {
+        let first = tempdir().unwrap();
+        let second = tempdir().unwrap();
+        std::fs::write(first.path().join("a.txt"), "a").unwrap();
+        std::fs::write(second.path().join("b.txt"), "b").unwrap();
+
+        let mut documents =
+            collect_documents(&[first.path().to_path_buf(), second.path().to_path_buf()]).unwrap();
+        documents.sort_by(|x, y| x.text.cmp(&y.text));
+
+        assert_eq!(documents.len(), 2);
+        assert_eq!(documents[0].text, "a");
+        assert_eq!(documents[1].text, "b");
+    }
+
+    #[test]
+    fn collect_documents_handles_a_mix_of_files_and_directories_in_one_call() {
+        let explicit = tempdir().unwrap();
+        let explicit_file = explicit.path().join("explicit.txt");
+        std::fs::write(&explicit_file, "explicit").unwrap();
+
+        let scanned = tempdir().unwrap();
+        std::fs::write(scanned.path().join("scanned.txt"), "scanned").unwrap();
+
+        let mut documents =
+            collect_documents(&[explicit_file, scanned.path().to_path_buf()]).unwrap();
+        documents.sort_by(|x, y| x.text.cmp(&y.text));
+
+        assert_eq!(documents.len(), 2);
+        assert_eq!(documents[0].text, "explicit");
+        assert_eq!(documents[1].text, "scanned");
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn collect_documents_skips_an_unreadable_file_inside_a_directory_but_keeps_the_rest() {
+        use std::os::unix::fs::PermissionsExt;
+
+        let dir = tempdir().unwrap();
+        let readable = dir.path().join("readable.txt");
+        let unreadable = dir.path().join("unreadable.txt");
+        std::fs::write(&readable, "readable").unwrap();
+        std::fs::write(&unreadable, "unreadable").unwrap();
+        std::fs::set_permissions(&unreadable, std::fs::Permissions::from_mode(0o000)).unwrap();
+
+        let documents = collect_documents(&[dir.path().to_path_buf()]);
+
+        std::fs::set_permissions(&unreadable, std::fs::Permissions::from_mode(0o644)).unwrap();
+
+        let documents = documents.unwrap();
+        assert_eq!(documents.len(), 1);
+        assert_eq!(documents[0].text, "readable");
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn collect_documents_fails_outright_on_an_unreadable_directory() {
+        use std::os::unix::fs::PermissionsExt;
+
+        let dir = tempdir().unwrap();
+        std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o000)).unwrap();
+
+        let result = collect_documents(&[dir.path().to_path_buf()]);
+
+        std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755)).unwrap();
+
+        assert!(
+            result.is_err(),
+            "an unreadable directory itself should fail the whole call, unlike an unreadable file"
+        );
+    }
+
+    #[test]
+    fn chunk_document_splits_a_readable_file_by_the_given_splitter() {
+        let dir = tempdir().unwrap();
+        let file = dir.path().join("small.txt");
+        std::fs::write(&file, "hello").unwrap();
+        let splitter = TextSplitter::new(CHUNK_CHARS);
+
+        let chunks = chunk_document(&file, &splitter);
+
+        assert_eq!(chunks.len(), 1);
+        assert_eq!(chunks[0].text, "hello");
+        assert_eq!(chunks[0].source, file.display().to_string());
+    }
+
+    #[test]
+    fn chunk_document_returns_empty_for_an_unreadable_path() {
+        let dir = tempdir().unwrap();
+        let missing = dir.path().join("does-not-exist.txt");
+        let splitter = TextSplitter::new(CHUNK_CHARS);
+
+        assert!(chunk_document(&missing, &splitter).is_empty());
+    }
+
+    #[test]
+    fn document_from_chunk_uses_the_bare_path_when_there_is_only_one_chunk() {
+        let path = Path::new("/tmp/notes.txt");
+
+        let document = Document::from_chunk(path, 0, "text", 1);
+
+        assert_eq!(document.source, "/tmp/notes.txt");
+        assert_eq!(document.text, "text");
+    }
+
+    #[test]
+    fn document_from_chunk_appends_a_part_suffix_when_there_are_multiple_chunks() {
+        let path = Path::new("/tmp/notes.txt");
+
+        let first = Document::from_chunk(path, 0, "a", 3);
+        let last = Document::from_chunk(path, 2, "c", 3);
+
+        assert_eq!(first.source, "/tmp/notes.txt (part 1/3)");
+        assert_eq!(last.source, "/tmp/notes.txt (part 3/3)");
+    }
 }