doubleo7/src/documents.rs

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

/// Mirrors `tools::MAX_PAGE_CHARS` — one huge uploaded file shouldn't blow
/// out the embedding model's input any more than one huge fetched page
/// should blow out the researcher's context.
const MAX_DOCUMENT_CHARS: usize = 20_000;

/// A single uploaded document, embedded via its full text so the researcher
/// can later retrieve semantically relevant excerpts by topic.
#[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 documents to embed: a file is read
/// directly, 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 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(), &mut documents);
                }
            }
        } else {
            push_document(path, &mut documents);
        }
    }

    Ok(documents)
}

fn push_document(path: &Path, documents: &mut Vec<Document>) {
    match std::fs::read_to_string(path) {
        Ok(text) => documents.push(Document {
            source: path.display().to_string(),
            text: text.chars().take(MAX_DOCUMENT_CHARS).collect(),
        }),
        Err(err) => {
            tracing::warn!(path = %path.display(), %err, "skipping unreadable document")
        }
    }
}

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

    /// A directory unique to this test run, under the OS temp dir, cleaned
    /// up on drop so a panic mid-test can't leave stray files behind for the
    /// next run.
    struct TempDir(PathBuf);

    impl TempDir {
        fn new(name: &str) -> Self {
            let path = std::env::temp_dir().join(format!(
                "deep_research-test-{name}-{}-{}",
                std::process::id(),
                std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .expect("system clock is after the epoch")
                    .as_nanos()
            ));
            std::fs::create_dir_all(&path).expect("create temp dir");
            Self(path)
        }

        fn path(&self) -> &Path {
            &self.0
        }
    }

    impl Drop for TempDir {
        fn drop(&mut self) {
            let _ = std::fs::remove_dir_all(&self.0);
        }
    }

    #[test]
    fn collect_documents_reads_an_explicit_file() {
        let dir = TempDir::new("explicit-file");
        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::new("directory");
        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::new("missing");
        let missing = dir.path().join("does-not-exist.txt");

        assert!(collect_documents(&[missing]).unwrap().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()]
        );
    }
}