Add document embedding and retrieval via a dedicated embedding model

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

@@ -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 |

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<tracing::Level>,
+
+    /// 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<PathBuf>,
 }

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

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

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";

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

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<String> {
+pub(crate) async fn research(
+    topic: &str,
+    docs: &[PathBuf],
+    show_progress: bool,
+) -> anyhow::Result<String> {
     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<Review> = 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<Option<String>> {
+    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))
+}

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<GatheredFindings> {
@@ -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)

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<ollama::EmbeddingModel, Document>;
+
+/// 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<Document>,
+    show_progress: bool,
+) -> anyhow::Result<DocumentIndex> {
+    let spinner = Spinner::start(
+        show_progress,
+        format!("{EMBED_EMOJI} Embedding {} document(s)...", documents.len()),
+    );
+
+    let embedded = client
+        .embeddings::<Document>(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<String> {
+    let request = VectorSearchRequest::builder()
+        .query(query)
+        .samples(TOP_N_EXCERPTS)
+        .build();
+
+    let results = index.top_n::<Document>(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::<Vec<_>>()
+        .join("\n\n"))
+}