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