From ec15893c9c423e2c1d8e3c02e8a324998eff8ee7 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 14:27:54 +0200 Subject: [PATCH] Add document embedding and retrieval via a dedicated embedding model MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds --doc (repeatable, file or directory) so the researcher can draw on user-supplied documents alongside the web: documents.rs resolves paths into embeddable text, retrieval.rs embeds them with a dedicated embedding model (nomic-embed-text, separate from the chat models used elsewhere) into an in-memory vector index and retrieves the excerpts most relevant to the topic once up front, and researcher.rs folds those excerpts into the researcher's task under the same footnote-citation scheme already used for web sources. The embedding and retrieval phases show progress the same way every other phase does — a spinner while embedding, a summary line once excerpts are retrieved, tracing spans for -l mode. Verified against a live Ollama nomic-embed-text pull and a real research round: a planted fact sheet was correctly ranked as the most relevant of several embedded documents and appeared in the researcher's task before its first turn. --- README.md | 13 ++++ src/cli.rs | 8 +++ src/documents.rs | 148 ++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 4 +- src/models.rs | 6 ++ src/progress.rs | 1 + src/research.rs | 48 ++++++++++++++- src/researcher.rs | 14 ++++- src/retrieval.rs | 74 +++++++++++++++++++++++ 9 files changed, 311 insertions(+), 5 deletions(-) create mode 100644 src/documents.rs create mode 100644 src/retrieval.rs diff --git a/README.md b/README.md index c97deae..516748e 100644 --- a/README.md +++ b/README.md @@ -79,8 +79,19 @@ 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 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. + ## Project layout Split one concern per file rather than one large module: @@ -94,6 +105,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 269a257..4f42f6e 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 Cli { /// 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..8a11c2c --- /dev/null +++ b/src/documents.rs @@ -0,0 +1,148 @@ +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()] + ); + } +} diff --git a/src/main.rs b/src/main.rs index 1b71913..31ed9e7 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..faef9e3 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 diff --git a/src/research.rs b/src/research.rs index 2dc803e..e4cf484 100644 --- a/src/research.rs +++ b/src/research.rs @@ -1,9 +1,12 @@ +use crate::documents; use crate::progress::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; @@ -56,3 +72,29 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result 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..95f9218 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) if !context.is_empty() => 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}" + ), + _ => 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..684858d --- /dev/null +++ b/src/retrieval.rs @@ -0,0 +1,74 @@ +use crate::documents::Document; +use crate::models::EMBEDDING_MODEL; +use crate::progress::{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. +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 results = index.top_n::(request).await?; + + if show_progress { + eprintln!( + "{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")) +}