From ec15893c9c423e2c1d8e3c02e8a324998eff8ee7 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 14:27:54 +0200 Subject: [PATCH 1/8] 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")) +} -- 2.45.2 From ff9f9455a65846f0f223877f39db32e9d240053a Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 14:43:14 +0200 Subject: [PATCH 2/8] Chunk uploaded documents instead of truncating them at 20K chars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A flat 20K-char cutoff (copied from the fetched-web-page limit) silently dropped everything past the first ~20KB of a larger file, and even within the cutoff, embedding a whole multi-page document as one vector made retrieval coarse — the vector just averages out whatever topics the document covers. Split each file into ~1500-char chunks via text-splitter (recursive semantic-boundary splitting: paragraph > sentence > word, never mid-word) and embed each chunk as its own document, tagged with its source and part number. This removes the practical size ceiling — a large file chunks the same way a short one does — and sharpens retrieval by letting it surface the specific passage relevant to a query. It also incidentally caps the worst-case retrieval payload: 5 chunks now tops out around 7500 chars versus the old worst case of 5 full 20K-char documents. Verified against a live Ollama nomic-embed-text pull: a short document still embeds as a single chunk, unchanged from before. --- Cargo.lock | 101 +++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/documents.rs | 77 +++++++++++++++++++++++++++++------- 3 files changed, 164 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 12247c2..8c28f42 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" @@ -1942,6 +1963,7 @@ dependencies = [ "scraper", "serde", "serde_json", + "text-splitter", "tokio", "tracing", "tracing-subscriber", @@ -2059,6 +2081,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" @@ -2874,11 +2907,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 +2982,8 @@ checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428" dependencies = [ "displaydoc", "icu_locale_core", + "serde", + "stable_deref_trait", "writeable", "yoke", "zerofrom", @@ -2935,6 +2991,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 +4842,8 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" dependencies = [ + "serde_core", + "writeable", "zerovec", ] @@ -6409,6 +6490,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 +6605,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" dependencies = [ "displaydoc", + "serde_core", "zerovec", ] @@ -7560,6 +7659,7 @@ dependencies = [ "displaydoc", "yoke", "zerofrom", + "zerovec", ] [[package]] @@ -7568,6 +7668,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 e6e0ec6..45bba54 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ rig = "0.41.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"] } diff --git a/src/documents.rs b/src/documents.rs index 8a11c2c..f0ba97d 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -1,13 +1,25 @@ use rig::embeddings::{EmbedError, TextEmbedder}; use std::path::{Path, PathBuf}; +use text_splitter::TextSplitter; -/// 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; +/// Safety valve against accidentally pointing `--doc` at a huge binary or +/// log file — not a content limit. Real documents are chunked in full (see +/// `CHUNK_CHARS` below), so nothing meaningful gets silently dropped short +/// of this. +const MAX_FILE_CHARS: usize = 2_000_000; -/// A single uploaded document, embedded via its full text so the researcher -/// can later retrieve semantically relevant excerpts by topic. +/// 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, @@ -21,11 +33,12 @@ impl rig::Embed for Document { } } -/// 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. +/// 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 mut documents = Vec::new(); @@ -47,10 +60,20 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result) { 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(), - }), + Ok(text) => { + let text: String = text.chars().take(MAX_FILE_CHARS).collect(); + let chunks: Vec<&str> = TextSplitter::new(CHUNK_CHARS).chunks(&text).collect(); + let total = chunks.len(); + + documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { + source: if total > 1 { + format!("{} (part {}/{total})", path.display(), i + 1) + } else { + path.display().to_string() + }, + text: chunk.to_string(), + })); + } Err(err) => { tracing::warn!(path = %path.display(), %err, "skipping unreadable document") } @@ -133,6 +156,30 @@ mod tests { assert!(collect_documents(&[missing]).unwrap().is_empty()); } + #[test] + fn collect_documents_splits_a_large_file_into_multiple_chunks() { + let dir = TempDir::new("large-file"); + 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 { -- 2.45.2 From b5f12a500cedae33ff37c79b95698d919c53c419 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 14:49:53 +0200 Subject: [PATCH 3/8] Document a real retrieval limitation found via manual large-document testing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Manually tested chunking against a 13.7KB/12-chunk document (previous verification only used a 424-char single-chunk file, which never exercised multi-chunk retrieval). Chunking itself held up. Retrieval didn't: built a document with 6 near-identical distractor sections and only one true answer, and the fixed top-5 slots filled entirely with distractors, excluding the chunk that actually answered the query. Documented as a known limitation rather than fixed now — it takes a document engineered to trigger it (several chunks that all read as similar to the query), not a typical upload. --- README.md | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 516748e..6bfbb9f 100644 --- a/README.md +++ b/README.md @@ -86,11 +86,25 @@ cargo run -- -l info "your research topic" 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. +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 -- 2.45.2 From 19c92ccd52aef9adc0060ec3d3a3af58471686d7 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 15:13:07 +0200 Subject: [PATCH 4/8] refactor: cleanup pass on document embedding/retrieval - retrieval::retrieve_relevant now shows a spinner around the query- embedding call and gets its own tracing span, matching every other model-calling phase in the pipeline (it previously ran invisibly and untraced) - add progress::report() to dedupe the show_progress-gated eprintln! pattern shared by research.rs and retrieval.rs - documents::push_document builds its TextSplitter once per collect_documents call instead of once per file, skips the char-truncation walk entirely when a file is already under the limit, and pulls the source-string branch out of the map closure - researcher::gather_findings drops a doc_context emptiness check that build_doc_context already guarantees - documents.rs tests use tempfile::tempdir() instead of a hand-rolled TempDir type Verified with cargo build/test/clippy -D warnings/fmt --check. --- Cargo.lock | 1 + Cargo.toml | 1 + src/documents.rs | 70 +++++++++++++++++------------------------------ src/history.rs | 6 +++- src/progress.rs | 9 ++++++ src/research.rs | 11 ++++---- src/researcher.rs | 4 +-- src/retrieval.rs | 17 ++++++++---- 8 files changed, 61 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 80d92de..64bf95e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2129,6 +2129,7 @@ dependencies = [ "scraper", "serde", "serde_json", + "tempfile", "text-splitter", "tokio", "tracing", diff --git a/Cargo.toml b/Cargo.toml index e13d433..d42ec1b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,3 +27,4 @@ tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } [dev-dependencies] serde_json = "1" +tempfile = "3" diff --git a/src/documents.rs b/src/documents.rs index c7a80a5..e3fe5f5 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -1,6 +1,6 @@ use rig::embeddings::{EmbedError, TextEmbedder}; use std::path::{Path, PathBuf}; -use text_splitter::TextSplitter; +use text_splitter::{Characters, TextSplitter}; /// Safety valve against accidentally pointing `--doc` at a huge binary or /// log file — not a content limit. Real documents are chunked in full (see @@ -40,6 +40,7 @@ impl rig::Embed for Document { /// (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(); for path in paths { @@ -47,30 +48,38 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result) { +fn push_document(path: &Path, splitter: &TextSplitter, documents: &mut Vec) { match std::fs::read_to_string(path) { - Ok(text) => { - let text: String = text.chars().take(MAX_FILE_CHARS).collect(); - let chunks: Vec<&str> = TextSplitter::new(CHUNK_CHARS).chunks(&text).collect(); + Ok(mut text) => { + // Cheap upper bound on char count, so the vast majority of + // documents (well under the limit) skip the char-by-char walk + // entirely. + if text.len() > MAX_FILE_CHARS && text.chars().count() > MAX_FILE_CHARS { + text = text.chars().take(MAX_FILE_CHARS).collect(); + } + let chunks: Vec<&str> = splitter.chunks(&text).collect(); let total = chunks.len(); - - documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { - source: if total > 1 { + let source = |i: usize| { + if total > 1 { format!("{} (part {}/{total})", path.display(), i + 1) } else { path.display().to_string() - }, + } + }; + + documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { + source: source(i), text: chunk.to_string(), })); } @@ -83,40 +92,11 @@ fn push_document(path: &Path, documents: &mut Vec) { #[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!( - "doubleo7-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); - } - } + use tempfile::tempdir; #[test] fn collect_documents_reads_an_explicit_file() { - let dir = TempDir::new("explicit-file"); + let dir = tempdir().unwrap(); let file = dir.path().join("notes.txt"); std::fs::write(&file, "hello from a file").unwrap(); @@ -129,7 +109,7 @@ mod tests { #[test] fn collect_documents_reads_every_file_in_a_directory_non_recursively() { - let dir = TempDir::new("directory"); + 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(); @@ -150,7 +130,7 @@ mod tests { #[test] fn collect_documents_skips_unreadable_paths_instead_of_failing() { - let dir = TempDir::new("missing"); + let dir = tempdir().unwrap(); let missing = dir.path().join("does-not-exist.txt"); assert!(collect_documents(&[missing]).unwrap().is_empty()); @@ -158,7 +138,7 @@ mod tests { #[test] fn collect_documents_splits_a_large_file_into_multiple_chunks() { - let dir = TempDir::new("large-file"); + 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. 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/progress.rs b/src/progress.rs index faef9e3..013a3cc 100644 --- a/src/progress.rs +++ b/src/progress.rs @@ -69,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 98899a5..e07a610 100644 --- a/src/research.rs +++ b/src/research.rs @@ -1,5 +1,5 @@ use crate::documents; -use crate::progress::REJECTED_EMOJI; +use crate::progress::{self, REJECTED_EMOJI}; use crate::researcher::gather_findings; use crate::retrieval; use crate::review::{self, Review}; @@ -57,12 +57,13 @@ pub(crate) async fn research( break; } - if show_progress { - eprintln!( + progress::report( + show_progress, + format!( "{REJECTED_EMOJI} Findings rejected — revising for round {}...", round + 1 - ); - } + ), + ); feedback = Some(review); } diff --git a/src/researcher.rs b/src/researcher.rs index 95f9218..4c3c72c 100644 --- a/src/researcher.rs +++ b/src/researcher.rs @@ -70,14 +70,14 @@ pub(crate) async fn gather_findings( }; let task = match doc_context { - Some(context) if !context.is_empty() => format!( + 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}" ), - _ => task, + None => task, }; let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); diff --git a/src/retrieval.rs b/src/retrieval.rs index 684858d..c8bd396 100644 --- a/src/retrieval.rs +++ b/src/retrieval.rs @@ -1,6 +1,6 @@ use crate::documents::Document; use crate::models::EMBEDDING_MODEL; -use crate::progress::{EMBED_EMOJI, Spinner}; +use crate::progress::{self, EMBED_EMOJI, Spinner}; use rig::client::EmbeddingsClient; use rig::providers::ollama; use rig::vector_store::VectorStoreIndex; @@ -44,6 +44,7 @@ pub(crate) async fn build_index( /// 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, @@ -54,14 +55,20 @@ pub(crate) async fn retrieve_relevant( .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); - if show_progress { - eprintln!( + 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 -- 2.45.2 From 5e55861a3f33068209ad43aed7631b15975418e4 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 15:23:52 +0200 Subject: [PATCH 5/8] refactor: drop MAX_FILE_CHARS cap on --doc file size MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It didn't actually guard what it claimed to: std::fs::read_to_string loads the whole file into memory before truncation ever ran, so it never prevented the OOM risk it was documented as protecting against — it only capped how much of the already-fully-read text got chunked afterward. Directory scanning is non-recursive and --doc is an explicit opt-in, so an oversized file is on the caller; chunking already handles arbitrarily long documents correctly. Can add a pre-read fs::metadata size check back later if real usage shows a need for it. --- src/documents.rs | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/src/documents.rs b/src/documents.rs index e3fe5f5..a167330 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -2,12 +2,6 @@ use rig::embeddings::{EmbedError, TextEmbedder}; use std::path::{Path, PathBuf}; use text_splitter::{Characters, TextSplitter}; -/// Safety valve against accidentally pointing `--doc` at a huge binary or -/// log file — not a content limit. Real documents are chunked in full (see -/// `CHUNK_CHARS` below), so nothing meaningful gets silently dropped short -/// of this. -const MAX_FILE_CHARS: usize = 2_000_000; - /// 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. @@ -61,13 +55,7 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result, documents: &mut Vec) { match std::fs::read_to_string(path) { - Ok(mut text) => { - // Cheap upper bound on char count, so the vast majority of - // documents (well under the limit) skip the char-by-char walk - // entirely. - if text.len() > MAX_FILE_CHARS && text.chars().count() > MAX_FILE_CHARS { - text = text.chars().take(MAX_FILE_CHARS).collect(); - } + Ok(text) => { let chunks: Vec<&str> = splitter.chunks(&text).collect(); let total = chunks.len(); let source = |i: usize| { -- 2.45.2 From 6632d632398442f15ab4800a7161a158b945c198 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 16:47:37 +0200 Subject: [PATCH 6/8] refactor: rework document/directory resolution and add its unit tests - collect_documents now partitions --doc paths into files vs. directories up front and handles each with its own loop, instead of a single branching if/else per path; the directory branch resolves entries via a fallible iterator chain (? inside a Result-returning map, collected with Result::transpose) rather than a nested for loop - extracted chunk_document (path, splitter) -> Vec as a pure function, and Document::from_chunk(path, index, chunk, total) to own the chunk's source-string formatting, both previously inlined in a function that mutated a shared Vec in place - renamed push_document -> chunk_document to match what it now does (reads and chunks a file into owned Documents) instead of what it used to do (push one Document into a caller-supplied Vec) - added unit tests for the reworked logic: multiple files/directories in one call, a mix of both in one call, an unreadable file inside a scanned directory (skipped, siblings kept) vs. an unreadable directory itself (fails the whole call, unlike a file), and direct tests of chunk_document and Document::from_chunk in isolation Verified with cargo build/test (34 passed)/clippy -D warnings/fmt --check. --- src/documents.rs | 206 +++++++++++++++++++++++++++++++++++++++++------ 1 file changed, 182 insertions(+), 24 deletions(-) diff --git a/src/documents.rs b/src/documents.rs index a167330..3b34919 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -20,6 +20,25 @@ pub(crate) struct Document { 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 = if total > 1 { + format!("{} (part {}/{total})", path.display(), index + 1) + } else { + 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()); @@ -37,42 +56,48 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result, 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?; - if entry.file_type()?.is_file() { - push_document(&entry.path(), &splitter, &mut documents); - } - } - } else { - push_document(path, &splitter, &mut documents); - } + 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) } -fn push_document(path: &Path, splitter: &TextSplitter, documents: &mut Vec) { +/// 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 total = chunks.len(); - let source = |i: usize| { - if total > 1 { - format!("{} (part {}/{total})", path.display(), i + 1) - } else { - path.display().to_string() - } - }; + let chunk_count = chunks.len(); - documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { - source: source(i), - text: chunk.to_string(), - })); + 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") + tracing::warn!(path = %path.display(), %err, "skipping unreadable document"); + Vec::new() } } } @@ -160,4 +185,137 @@ mod tests { 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"); + } + + #[cfg(unix)] + #[test] + fn collect_documents_skips_an_unreadable_file_inside_a_directory_but_keeps_the_rest() { + use std::os::unix::fs::PermissionsExt; + + 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; + + 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)"); + } } -- 2.45.2 From f916787962dc2b9f3d4ab9795675a97e7be6cf5a Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 17:04:10 +0200 Subject: [PATCH 7/8] chore: refactor some code. --- src/documents.rs | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/documents.rs b/src/documents.rs index 3b34919..dc51a9e 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -26,10 +26,11 @@ impl Document { /// 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 = if total > 1 { - format!("{} (part {}/{total})", path.display(), index + 1) - } else { - path.display().to_string() + let source = match total { + total_chunks if total_chunks > 1 => { + format!("{} (part {}/{total})", path.display(), index + 1) + } + _ => path.display().to_string(), }; Self { -- 2.45.2 From ba9b61fcf3b2d613be8ce064be4c67817a08dd66 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 17:34:29 +0200 Subject: [PATCH 8/8] fix: skip chmod-0o000 permission tests when running as root MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI runs the test suite as root inside an unmodified Docker base image (data.forgejo.org/oci/node:20-bookworm), where chmod 0o000 doesn't actually block reads — root bypasses Unix permission bits entirely. The two new permission-based tests from the previous commit passed locally (non-root) but failed in CI for exactly that reason. Skip them under root via a raw geteuid() FFI check instead of asserting behavior the OS isn't enforcing. Verified with cargo build/test (34 passed)/clippy -D warnings/fmt --check. --- src/documents.rs | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/src/documents.rs b/src/documents.rs index dc51a9e..24c5090 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -237,11 +237,29 @@ mod tests { 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"); @@ -263,6 +281,11 @@ mod tests { 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(); -- 2.45.2