diff --git a/Cargo.lock b/Cargo.lock index 0d434f0..64bf95e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -517,6 +517,18 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "auto_enums" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3091d68264354f211516b91dce6f71046e444fab1867716035f736667243affb" +dependencies = [ + "derive_utils", + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -1068,6 +1080,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_maths" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77745e017f5edba1a9c1d854f6f3a52dac8a12dd5af5d2f54aecf61e43d80d30" +dependencies = [ + "libm", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -2039,6 +2060,17 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "derive_utils" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc05a5d33db20c784f873e84934ad94bb209a090987ac5f62fede2c178234f23" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "digest" version = "0.10.7" @@ -2097,6 +2129,8 @@ dependencies = [ "scraper", "serde", "serde_json", + "tempfile", + "text-splitter", "tokio", "tracing", "tracing-subscriber", @@ -2874,11 +2908,32 @@ checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb" dependencies = [ "displaydoc", "litemap", + "serde", "tinystr", "writeable", "zerovec", ] +[[package]] +name = "icu_locale_fallback" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "251af8e57c9400e3eb58242fe5b8b1152b2a64fdf4cf632f923c38ccee6f2fa9" +dependencies = [ + "icu_locale_core", + "icu_locale_fallback_data", + "icu_provider", + "potential_utf", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_locale_fallback_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "decf2a22ec8fa68f1a0c1129a3f8583f8f8bc24e8b9ccbe98ead99f62a4dc3a8" + [[package]] name = "icu_normalizer" version = "2.3.0" @@ -2928,6 +2983,8 @@ checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428" dependencies = [ "displaydoc", "icu_locale_core", + "serde", + "stable_deref_trait", "writeable", "yoke", "zerofrom", @@ -2935,6 +2992,29 @@ dependencies = [ "zerovec", ] +[[package]] +name = "icu_segmenter" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82d07aafccd67af15d02512a6adf5896fbc5ed00f2e99b471d2efa14016db3db" +dependencies = [ + "core_maths", + "icu_collections", + "icu_locale_fallback", + "icu_provider", + "icu_segmenter_data", + "potential_utf", + "smallvec", + "utf8_iter", + "zerovec", +] + +[[package]] +name = "icu_segmenter_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae293c039020f9ec10710af98d29ce6aa2051486638b49c9a6409f3b4a9e98ad" + [[package]] name = "ident_case" version = "1.0.1" @@ -4763,6 +4843,8 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" dependencies = [ + "serde_core", + "writeable", "zerovec", ] @@ -6409,6 +6491,23 @@ dependencies = [ "new_debug_unreachable", ] +[[package]] +name = "text-splitter" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f3eeec76988617ff1434d754d7e8e197be2cb8981ea77c17e32f3f7a2c6f95e" +dependencies = [ + "ahash", + "auto_enums", + "either", + "icu_provider", + "icu_segmenter", + "itertools 0.14.0", + "memchr", + "strum 0.28.0", + "thiserror", +] + [[package]] name = "thiserror" version = "2.0.20" @@ -6507,6 +6606,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" dependencies = [ "displaydoc", + "serde_core", "zerovec", ] @@ -7553,6 +7653,7 @@ dependencies = [ "displaydoc", "yoke", "zerofrom", + "zerovec", ] [[package]] @@ -7561,6 +7662,7 @@ version = "0.11.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8" dependencies = [ + "serde", "yoke", "zerofrom", "zerovec-derive", diff --git a/Cargo.toml b/Cargo.toml index 47af285..d42ec1b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,9 +20,11 @@ rig = "0.42.0" schemars = "1" scraper = "0.27" serde = { version = "1.0.229", features = ["derive"] } +text-splitter = "0.32" tokio = { version = "1.53.1", features = ["full"] } tracing = "0.1" tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } [dev-dependencies] serde_json = "1" +tempfile = "3" diff --git a/README.md b/README.md index f5ac686..814d92a 100644 --- a/README.md +++ b/README.md @@ -79,8 +79,33 @@ cargo run -- "your research topic" # or, with tracing spans on stderr instead of the progress spinner: cargo run -- -l info "your research topic" + +# give the researcher your own documents to draw on, alongside the web — +# repeatable, and a directory contributes every file directly inside it +# (one level deep, not recursive): +cargo run -- --doc ./notes.txt --doc ./research-docs/ "your research topic" ``` +Uploaded documents are chunked (see `documents.rs`), embedded with a +dedicated embedding model (see `EMBEDDING_MODEL` in +[`src/models.rs`](./src/models.rs)) into an in-memory vector index, then the +excerpts most relevant to the topic are retrieved and folded into the +researcher's task alongside anything it finds on the web — the same +footnote-citation scheme applies to both. + +Known limitation: retrieval returns a fixed top-N chunks +(`retrieval::TOP_N_EXCERPTS`). A document with several chunks that all read +as similar to the query — several incident reports, several revisions of +the same section — can crowd out the one chunk that actually answers it, +since only the top N by similarity are ever returned regardless of how many +plausible candidates exist. Reproduced deliberately (a 13.7 KB / 12-chunk +document with 6 near-identical "incident report" sections, only one of +which had the real answer, was built specifically to stress this — the top +5 slots filled entirely with distractors and the answer chunk was +excluded), so it's a real edge case, not a hypothetical. Not fixed for now +since it takes a document engineered to trigger it, but worth knowing if a +report seems to be missing something you know is in an uploaded document. + ## Project layout Split one concern per file rather than one large module: @@ -94,6 +119,8 @@ Split one concern per file rather than one large module: | `summarizer.rs` | Max-turns recovery: reconstructs findings via a model call | | `writer.rs` | Turns findings into the final streamed report | | `history.rs` | Pure, unit-tested helpers for parsing a rig chat history into usable text | +| `documents.rs` | Resolves `--doc` paths into embeddable documents | +| `retrieval.rs` | Embeds documents into an in-memory vector index and retrieves relevant excerpts | | `tools.rs` | `search_web` (SearXNG) and `fetch_page` tool implementations | | `stream.rs` | Drains a streaming prompt response to the terminal | | `progress.rs` | The terminal spinner and per-phase emoji | diff --git a/src/cli.rs b/src/cli.rs index 6d3d9f2..495e30e 100644 --- a/src/cli.rs +++ b/src/cli.rs @@ -1,4 +1,5 @@ use clap::Parser; +use std::path::PathBuf; pub(crate) const DEFAULT_TOPIC: &str = "What are the latest advances in running large language models locally, on consumer hardware?"; @@ -16,4 +17,11 @@ pub(crate) struct Doubleo7 { /// progress spinner to switch off, since the logs already show progress) #[arg(short = 'l', long, value_name = "LEVEL")] pub(crate) log_level: Option, + + /// A file, or a directory of files, to embed and make available to the + /// researcher as retrieved excerpts (repeatable). Directories are + /// scanned one level deep — non-directory entries only, subdirectories + /// are skipped rather than walked. + #[arg(short = 'd', long = "doc", value_name = "PATH")] + pub(crate) docs: Vec, } diff --git a/src/documents.rs b/src/documents.rs new file mode 100644 index 0000000..24c5090 --- /dev/null +++ b/src/documents.rs @@ -0,0 +1,345 @@ +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> { + 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> { + let entry = entry?; + Ok(entry.file_type()?.is_file().then(|| entry.path())) + }) + .filter_map(Result::transpose) + .collect::>>()? + .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) -> Vec { + 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)"); + } +} diff --git a/src/history.rs b/src/history.rs index d506fd6..09b6fa7 100644 --- a/src/history.rs +++ b/src/history.rs @@ -270,7 +270,11 @@ mod tests { serde_json::json!({ "query": "test" }), )], }, - Message::tool_result("call-1", "call-1-name", "1. Example\n https://example.com\n snippet"), + Message::tool_result( + "call-1", + "call-1-name", + "1. Example\n https://example.com\n snippet", + ), Message::assistant("that source looks solid"), ]; diff --git a/src/main.rs b/src/main.rs index c1eb207..f1c4a8e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,14 @@ use clap::Parser; mod cli; +mod documents; mod history; mod models; mod observability; mod progress; mod research; mod researcher; +mod retrieval; mod review; mod stream; mod summarizer; @@ -30,7 +32,7 @@ async fn main() -> anyhow::Result<()> { // The report streams to stdout as the writer generates it, so nothing // left to print here — the return value only matters to callers that // embed `research` rather than running it as this binary. - research::research(&topic, show_progress).await?; + research::research(&topic, &cli.docs, show_progress).await?; Ok(()) } diff --git a/src/models.rs b/src/models.rs index b5f9120..4f421b1 100644 --- a/src/models.rs +++ b/src/models.rs @@ -6,3 +6,9 @@ /// passes instead. pub(crate) const RESEARCHER_MODEL: &str = "gemma4:26b"; pub(crate) const WRITER_MODEL: &str = "gemma4-e4b:latest"; + +/// Deliberately a dedicated embedding model rather than reusing a chat model +/// for embeddings — it's trained for semantic similarity, not chat, and +/// Ollama's `nomic-embed-text` is a well-known identifier Rig already knows +/// the output dimensionality for. +pub(crate) const EMBEDDING_MODEL: &str = "nomic-embed-text"; diff --git a/src/progress.rs b/src/progress.rs index d5117af..013a3cc 100644 --- a/src/progress.rs +++ b/src/progress.rs @@ -11,6 +11,7 @@ pub(crate) const REVIEW_EMOJI: &str = "🧐"; pub(crate) const REJECTED_EMOJI: &str = "❌"; pub(crate) const REPORT_EMOJI: &str = "✍️"; pub(crate) const SUMMARIZE_EMOJI: &str = "🧩"; +pub(crate) const EMBED_EMOJI: &str = "📚"; /// The spinner currently on screen, if any — set by `Spinner::start` and /// cleared on drop. Tool implementations don't otherwise have a handle to @@ -68,6 +69,15 @@ impl Drop for Spinner { } } +/// Prints a status line, but only when progress display is on — matching +/// `Spinner`'s own `enabled` gate so a call site doesn't need to repeat the +/// `if show_progress` check inline. +pub(crate) fn report(enabled: bool, message: impl Into) { + if enabled { + eprintln!("{}", message.into()); + } +} + /// Updates the message of whatever spinner is currently running, if any. /// A no-op when progress display is off (no spinner was ever started, so /// `active()` stays empty) or between phases (the previous `Spinner` has diff --git a/src/research.rs b/src/research.rs index ffd68c7..e07a610 100644 --- a/src/research.rs +++ b/src/research.rs @@ -1,9 +1,12 @@ -use crate::progress::REJECTED_EMOJI; +use crate::documents; +use crate::progress::{self, REJECTED_EMOJI}; use crate::researcher::gather_findings; +use crate::retrieval; use crate::review::{self, Review}; use crate::writer::write_report; use rig::client::Nothing; use rig::providers::ollama; +use std::path::PathBuf; /// Research/review rounds before giving up and writing the report from /// whatever the last pass produced, rather than looping forever on a topic @@ -14,16 +17,29 @@ const MAX_RESEARCH_ROUNDS: usize = 3; /// not a model, in charge of when to stop — re-running research with the /// reviewer's feedback folded in until it approves or the round budget runs /// out, then writing the report from whatever the last pass produced. -pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result { +pub(crate) async fn research( + topic: &str, + docs: &[PathBuf], + show_progress: bool, +) -> anyhow::Result { let client = ollama::Client::new(Nothing)?; + let doc_context = build_doc_context(&client, topic, docs, show_progress).await?; + let mut findings = String::new(); let mut feedback: Option = None; let mut incomplete = false; for round in 1..=MAX_RESEARCH_ROUNDS { - let gathered = - gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?; + let gathered = gather_findings( + &client, + topic, + feedback.as_ref(), + doc_context.as_deref(), + round, + show_progress, + ) + .await?; findings = gathered.findings; incomplete = gathered.incomplete; @@ -41,15 +57,42 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result break; } - if show_progress { - eprintln!( + progress::report( + show_progress, + format!( "{REJECTED_EMOJI} Findings rejected — revising for round {}...", round + 1 - ); - } + ), + ); feedback = Some(review); } write_report(&client, topic, &findings, incomplete, show_progress).await } + +/// Embeds any uploaded documents and retrieves the excerpts most relevant to +/// `topic`, once, up front — nothing later in the loop changes which +/// excerpts are relevant, so there's no reason to repeat this per round. +/// `None` when no documents were given, or none of them yielded a usable +/// excerpt (an empty `docs` list is deliberately not an error). +async fn build_doc_context( + client: &ollama::Client, + topic: &str, + docs: &[PathBuf], + show_progress: bool, +) -> anyhow::Result> { + if docs.is_empty() { + return Ok(None); + } + + let documents = documents::collect_documents(docs)?; + if documents.is_empty() { + return Ok(None); + } + + let index = retrieval::build_index(client, documents, show_progress).await?; + let excerpts = retrieval::retrieve_relevant(&index, topic, show_progress).await?; + + Ok((!excerpts.is_empty()).then_some(excerpts)) +} diff --git a/src/researcher.rs b/src/researcher.rs index ad327a8..4c3c72c 100644 --- a/src/researcher.rs +++ b/src/researcher.rs @@ -19,11 +19,12 @@ pub(crate) struct GatheredFindings { /// Wraps the tool-calling research loop in its own span so it's visible as a /// single unit in traces, distinct from the writing and review phases and /// nesting rig's own per-turn `chat`/`execute_tool` spans underneath it. -#[tracing::instrument(skip(client, feedback), fields(gen_ai.agent.name = "researcher"))] +#[tracing::instrument(skip(client, feedback, doc_context), fields(gen_ai.agent.name = "researcher"))] pub(crate) async fn gather_findings( client: &ollama::Client, topic: &str, feedback: Option<&Review>, + doc_context: Option<&str>, round: usize, show_progress: bool, ) -> anyhow::Result { @@ -68,6 +69,17 @@ pub(crate) async fn gather_findings( ), }; + let task = match doc_context { + Some(context) => format!( + "{task}\n\n\ + Relevant excerpts from documents the user uploaded — treat these as trusted primary \ + sources alongside anything you find on the web, and cite them with the same \ + footnote scheme (their Sources entry can just be the document path shown below):\n\ + {context}" + ), + None => task, + }; + let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); let run_result = researcher .runner(task) diff --git a/src/retrieval.rs b/src/retrieval.rs new file mode 100644 index 0000000..c8bd396 --- /dev/null +++ b/src/retrieval.rs @@ -0,0 +1,81 @@ +use crate::documents::Document; +use crate::models::EMBEDDING_MODEL; +use crate::progress::{self, EMBED_EMOJI, Spinner}; +use rig::client::EmbeddingsClient; +use rig::providers::ollama; +use rig::vector_store::VectorStoreIndex; +use rig::vector_store::in_memory_store::{InMemoryVectorIndex, InMemoryVectorStore}; +use rig::vector_store::request::VectorSearchRequest; + +/// Excerpts returned per retrieval query — enough to give the researcher +/// useful grounding without one query result dominating its context. +const TOP_N_EXCERPTS: u64 = 5; + +pub(crate) type DocumentIndex = InMemoryVectorIndex; + +/// Embeds every collected document with the dedicated embedding model (see +/// `models::EMBEDDING_MODEL`) and builds a queryable in-memory vector index +/// from the result. Runs once per research invocation, up front — document +/// relevance to the topic doesn't change between research/review rounds, so +/// there's nothing to gain from re-embedding per round. +#[tracing::instrument(skip(client, documents), fields(gen_ai.agent.name = "embedder"))] +pub(crate) async fn build_index( + client: &ollama::Client, + documents: Vec, + show_progress: bool, +) -> anyhow::Result { + let spinner = Spinner::start( + show_progress, + format!("{EMBED_EMOJI} Embedding {} document(s)...", documents.len()), + ); + + let embedded = client + .embeddings::(EMBEDDING_MODEL) + .documents(documents)? + .build() + .await?; + + drop(spinner); + tracing::info!(count = embedded.len(), "documents embedded"); + + let model = client.embedding_model(EMBEDDING_MODEL); + Ok(InMemoryVectorStore::from_documents(embedded).index(model)) +} + +/// Retrieves the excerpts most semantically relevant to `query` from the +/// index, formatted for inclusion in the researcher's task. +#[tracing::instrument(skip(index), fields(gen_ai.agent.name = "retriever"))] +pub(crate) async fn retrieve_relevant( + index: &DocumentIndex, + query: &str, + show_progress: bool, +) -> anyhow::Result { + let request = VectorSearchRequest::builder() + .query(query) + .samples(TOP_N_EXCERPTS) + .build(); + + let spinner = Spinner::start( + show_progress, + format!("{EMBED_EMOJI} Retrieving relevant excerpts..."), + ); + let results = index.top_n::(request).await?; + drop(spinner); + + progress::report( + show_progress, + format!( + "{EMBED_EMOJI} Retrieved {} relevant excerpt(s) from uploaded documents", + results.len() + ), + ); + tracing::info!(count = results.len(), "retrieved document excerpts"); + + Ok(results + .into_iter() + .map(|(score, _id, doc): (f64, String, Document)| { + format!("[{score:.2}] {}\n{}", doc.source, doc.text) + }) + .collect::>() + .join("\n\n")) +}