src/documents.rs

@@ -20,26 +20,6 @@ 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 = match total {
-            total_chunks if total_chunks > 1 => {
-                format!("{} (part {}/{total})", path.display(), index + 1)
-            }
-            _ => 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());
@@ -57,48 +37,42 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Documen
     let splitter = TextSplitter::new(CHUNK_CHARS);
     let mut documents = Vec::new();
 
-    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>> {
+    for path in paths {
+        if path.is_dir() {
+            for entry in std::fs::read_dir(path)? {
                 let entry = entry?;
-                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)));
+                if entry.file_type()?.is_file() {
+                    push_document(&entry.path(), &splitter, &mut documents);
+                }
+            }
+        } else {
+            push_document(path, &splitter, &mut documents);
+        }
     }
 
-    files
-        .iter()
-        .for_each(|file| documents.extend(chunk_document(file, &splitter)));
-
     Ok(documents)
 }
 
-/// 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> {
+fn push_document(path: &Path, splitter: &TextSplitter<Characters>, documents: &mut Vec<Document>) {
     match std::fs::read_to_string(path) {
         Ok(text) => {
             let chunks: Vec<&str> = splitter.chunks(&text).collect();
-            let chunk_count = chunks.len();
+            let total = chunks.len();
+            let source = |i: usize| {
+                if total > 1 {
+                    format!("{} (part {}/{total})", path.display(), i + 1)
+                } else {
+                    path.display().to_string()
+                }
+            };
 
-            chunks
-                .into_iter()
-                .enumerate()
-                .map(|(chunk_index, chunk)| {
-                    Document::from_chunk(path, chunk_index, chunk, chunk_count)
-                })
-                .collect()
+            documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document {
+                source: source(i),
+                text: chunk.to_string(),
+            }));
         }
         Err(err) => {
-            tracing::warn!(path = %path.display(), %err, "skipping unreadable document");
-            Vec::new()
+            tracing::warn!(path = %path.display(), %err, "skipping unreadable document")
         }
     }
 }
@@ -186,137 +160,4 @@ 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)");
-    }
 }