doubleo7/src/documents.rs

use rig::embeddings::{EmbedError, TextEmbedder};
use std::path::{Path, PathBuf};
use text_splitter::{Characters, TextSplitter};

/// Target chunk size handed to the embedding model: small enough that a
/// handful of retrieved chunks stays well within a local model's context
/// window, large enough to keep a paragraph or two of context in each one.
/// `TextSplitter` treats this as an upper bound, not a fixed size — it
/// recursively splits on the largest semantic boundary (paragraph,
/// sentence, word, ...) that still fits, so a chunk never cuts a sentence
/// mid-word just to hit the target exactly.
const CHUNK_CHARS: usize = 1_500;

/// A single chunk of an uploaded document, embedded via its own text so
/// retrieval can surface just the passage relevant to a query rather than
/// an entire (possibly very long) file at once.
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(crate) struct Document {
    pub(crate) source: String,
    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());
        Ok(())
    }
}

/// Resolves CLI-provided paths into chunked documents to embed: a file is
/// read and split into chunks, a directory contributes every non-directory
/// entry inside it — one level deep, not recursive, so a stray nested
/// folder can't silently pull in unrelated files. Unreadable entries
/// (permissions, non-UTF-8) are skipped with a warning rather than failing
/// the whole run.
pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Document>> {
    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>> {
                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)));
    }

    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> {
    match std::fs::read_to_string(path) {
        Ok(text) => {
            let chunks: Vec<&str> = splitter.chunks(&text).collect();
            let chunk_count = chunks.len();

            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");
            Vec::new()
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    #[test]
    fn collect_documents_reads_an_explicit_file() {
        let dir = tempdir().unwrap();
        let file = dir.path().join("notes.txt");
        std::fs::write(&file, "hello from a file").unwrap();

        let documents = collect_documents(std::slice::from_ref(&file)).unwrap();

        assert_eq!(documents.len(), 1);
        assert_eq!(documents[0].source, file.display().to_string());
        assert_eq!(documents[0].text, "hello from a file");
    }

    #[test]
    fn collect_documents_reads_every_file_in_a_directory_non_recursively() {
        let dir = tempdir().unwrap();
        std::fs::write(dir.path().join("a.txt"), "a").unwrap();
        std::fs::write(dir.path().join("b.txt"), "b").unwrap();
        std::fs::create_dir(dir.path().join("nested")).unwrap();
        std::fs::write(dir.path().join("nested/c.txt"), "c").unwrap();

        let mut documents = collect_documents(&[dir.path().to_path_buf()]).unwrap();
        documents.sort_by(|a, b| a.source.cmp(&b.source));

        assert_eq!(documents.len(), 2);
        assert_eq!(documents[0].text, "a");
        assert_eq!(documents[1].text, "b");
    }

    #[test]
    fn collect_documents_is_empty_for_no_paths() {
        assert!(collect_documents(&[]).unwrap().is_empty());
    }

    #[test]
    fn collect_documents_skips_unreadable_paths_instead_of_failing() {
        let dir = tempdir().unwrap();
        let missing = dir.path().join("does-not-exist.txt");

        assert!(collect_documents(&[missing]).unwrap().is_empty());
    }

    #[test]
    fn collect_documents_splits_a_large_file_into_multiple_chunks() {
        let dir = tempdir().unwrap();
        let file = dir.path().join("big.txt");
        // Well over CHUNK_CHARS, and with paragraph breaks so the splitter
        // has real semantic boundaries to chunk on.
        let paragraph = "word ".repeat(100);
        std::fs::write(&file, vec![paragraph; 10].join("\n\n")).unwrap();

        let documents = collect_documents(std::slice::from_ref(&file)).unwrap();

        assert!(
            documents.len() > 1,
            "expected a large file to produce multiple chunks"
        );
        for (i, doc) in documents.iter().enumerate() {
            assert_eq!(
                doc.source,
                format!("{} (part {}/{})", file.display(), i + 1, documents.len())
            );
            assert!(!doc.text.is_empty());
        }
    }

    #[test]
    fn document_embed_yields_its_full_text() {
        let document = Document {
            source: "s".to_string(),
            text: "the text to embed".to_string(),
        };

        assert_eq!(
            rig::embeddings::to_texts(document).unwrap(),
            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");
    }

    /// Unix permission bits are meaningless to a root process — it can read
    /// anything regardless of mode — so the two `chmod 0o000` tests below
    /// would fail under a root-run CI container (e.g. an unmodified Docker
    /// base image) despite the code being correct. Skip rather than assert
    /// behavior the OS isn't actually enforcing.
    #[cfg(unix)]
    fn running_as_root() -> bool {
        unsafe extern "C" {
            fn geteuid() -> u32;
        }
        unsafe { geteuid() == 0 }
    }

    #[cfg(unix)]
    #[test]
    fn collect_documents_skips_an_unreadable_file_inside_a_directory_but_keeps_the_rest() {
        use std::os::unix::fs::PermissionsExt;

        if running_as_root() {
            eprintln!("skipping: running as root, chmod 0o000 has no effect");
            return;
        }

        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;

        if running_as_root() {
            eprintln!("skipping: running as root, chmod 0o000 has no effect");
            return;
        }

        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)");
    }
}