Merge pull request 'Add document embedding and retrieval via a dedicated embedding model' (#10) from worktree-deep-research-max-turns-report into master

Reviewed-on: #10
Reviewed-by: Austin Schaefer <austin.schaefer@mailo.eu>

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

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"

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 |

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<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,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<Vec<Document>> {
+    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<Option<PathBuf>> {
+                let entry = entry?;
+                Ok(entry.file_type()?.is_file().then(|| entry.path()))
+            })
+            .filter_map(Result::transpose)
+            .collect::<anyhow::Result<Vec<PathBuf>>>()?
+            .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<Characters>) -> Vec<Document> {
+    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)");
+    }
+}

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"),
         ];
 

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
@@ -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<String>) {
+    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

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<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 =
+        let gathered = gather_findings(
-            gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?;
+            &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 {
+        progress::report(
-            eprintln!(
+            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<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) => 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)

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<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.
+#[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<String> {
+    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::<Document>(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::<Vec<_>>()
+        .join("\n\n"))
+}