doubleo7/src/documents.rs

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

/// Safety valve against accidentally pointing `--doc` at a huge binary or
/// log file — not a content limit. Real documents are chunked in full (see
/// `CHUNK_CHARS` below), so nothing meaningful gets silently dropped short
/// of this.
const MAX_FILE_CHARS: usize = 2_000_000;

/// 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 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();

    for path in paths {
        if path.is_dir() {
            for entry in std::fs::read_dir(path)? {
                let entry = entry?;
                if entry.file_type()?.is_file() {
                    push_document(&entry.path(), &splitter, &mut documents);
                }
            }
        } else {
            push_document(path, &splitter, &mut documents);
        }
    }

    Ok(documents)
}

fn push_document(path: &Path, splitter: &TextSplitter<Characters>, documents: &mut Vec<Document>) {
    match std::fs::read_to_string(path) {
        Ok(mut text) => {
            // Cheap upper bound on char count, so the vast majority of
            // documents (well under the limit) skip the char-by-char walk
            // entirely.
            if text.len() > MAX_FILE_CHARS && text.chars().count() > MAX_FILE_CHARS {
                text = text.chars().take(MAX_FILE_CHARS).collect();
            }
            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()
                }
            };

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

#[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()]
        );
    }
}