Compare commits

...

14 commits

Author SHA1 Message Date
8d7f8b6683 Merge pull request 'docs: add retriever to the architecture diagram' (#13) from docs/update-architecture-diagram into master
All checks were successful
CI / test (push) Successful in 1m1s
Reviewed-on: #13
Reviewed-by: Austin Schaefer <austin.schaefer@mailo.eu>
2026-08-19 15:54:57 +00:00
Austin Schaefer
0790b2b6d0 docs: add retriever to the architecture diagram
All checks were successful
CI / test (pull_request) Successful in 1m3s
The diagram still showed the original four-agent pipeline
(researcher/reviewer/writer/summarizer) with no mention of the
document embedding/retrieval step added in #10. Added a retriever
box feeding excerpts into the researcher, plus a bullet describing
it, matching the retrieval pipeline already documented in prose
further down the README.
2026-08-19 17:48:32 +02:00
90502a3eed Merge pull request 'Add document embedding and retrieval via a dedicated embedding model' (#10) from worktree-deep-research-max-turns-report into master
All checks were successful
CI / test (push) Successful in 5m37s
Reviewed-on: #10
Reviewed-by: Austin Schaefer <austin.schaefer@mailo.eu>
2026-08-19 15:37:31 +00:00
Austin Schaefer
ba9b61fcf3 fix: skip chmod-0o000 permission tests when running as root
All checks were successful
CI / test (pull_request) Successful in 1m22s
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.
2026-08-19 17:34:29 +02:00
Austin Schaefer
f916787962 chore: refactor some code.
Some checks failed
CI / test (pull_request) Failing after 1m50s
2026-08-19 17:04:10 +02:00
Austin Schaefer
6632d63239 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<Document> 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.
2026-08-19 16:47:37 +02:00
Austin Schaefer
5e55861a3f refactor: drop MAX_FILE_CHARS cap on --doc file size
All checks were successful
CI / test (pull_request) Successful in 1m23s
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.
2026-08-19 15:23:52 +02:00
Austin Schaefer
19c92ccd52 refactor: cleanup pass on document embedding/retrieval
All checks were successful
CI / test (pull_request) Successful in 6m56s
- 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.
2026-08-19 15:13:07 +02:00
Austin Schaefer
53b4f572df Merge master into worktree-deep-research-max-turns-report
Some checks failed
CI / test (pull_request) Failing after 54s
2026-08-19 15:03:24 +02:00
ad2ba1181a Merge pull request 'worktree-fix-ollama-json-extraction' (#12) from worktree-fix-ollama-json-extraction into master
Some checks failed
CI / test (push) Failing after 46s
Reviewed-on: #12
2026-08-19 12:19:44 +00:00
Austin Schaefer
8f4af514a3 Merge master (rebrand to doubleo7) into this branch
All checks were successful
CI / test (pull_request) Successful in 3m38s
Resolves the Cargo.lock conflict by regenerating it, and picks up
master's rename (deep_research -> doubleo7 crate/binary name, Cli ->
Doubleo7 struct) cleanly through everything this branch added
(--doc/documents.rs/retrieval.rs). Also renamed a leftover
"deep_research-test-" temp-dir prefix in documents.rs's tests for
consistency with the rebrand.
2026-08-18 16:19:39 +02:00
Austin Schaefer
b5f12a500c Document a real retrieval limitation found via manual large-document testing
All checks were successful
CI / test (pull_request) Successful in 11m11s
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.
2026-08-18 14:49:53 +02:00
Austin Schaefer
ff9f9455a6 Chunk uploaded documents instead of truncating them at 20K chars
All checks were successful
CI / test (pull_request) Successful in 11m5s
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.
2026-08-18 14:43:14 +02:00
Austin Schaefer
ec15893c9c Add document embedding and retrieval via a dedicated embedding model
All checks were successful
CI / test (pull_request) Successful in 1m52s
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.
2026-08-18 14:27:54 +02:00
12 changed files with 670 additions and 23 deletions

102
Cargo.lock generated
View file

@ -517,6 +517,18 @@ version = "1.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" 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]] [[package]]
name = "autocfg" name = "autocfg"
version = "1.5.1" version = "1.5.1"
@ -1068,6 +1080,15 @@ version = "0.8.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" 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]] [[package]]
name = "cpufeatures" name = "cpufeatures"
version = "0.2.17" version = "0.2.17"
@ -2039,6 +2060,17 @@ dependencies = [
"syn 2.0.119", "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]] [[package]]
name = "digest" name = "digest"
version = "0.10.7" version = "0.10.7"
@ -2097,6 +2129,8 @@ dependencies = [
"scraper", "scraper",
"serde", "serde",
"serde_json", "serde_json",
"tempfile",
"text-splitter",
"tokio", "tokio",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
@ -2874,11 +2908,32 @@ checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb"
dependencies = [ dependencies = [
"displaydoc", "displaydoc",
"litemap", "litemap",
"serde",
"tinystr", "tinystr",
"writeable", "writeable",
"zerovec", "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]] [[package]]
name = "icu_normalizer" name = "icu_normalizer"
version = "2.3.0" version = "2.3.0"
@ -2928,6 +2983,8 @@ checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428"
dependencies = [ dependencies = [
"displaydoc", "displaydoc",
"icu_locale_core", "icu_locale_core",
"serde",
"stable_deref_trait",
"writeable", "writeable",
"yoke", "yoke",
"zerofrom", "zerofrom",
@ -2935,6 +2992,29 @@ dependencies = [
"zerovec", "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]] [[package]]
name = "ident_case" name = "ident_case"
version = "1.0.1" version = "1.0.1"
@ -4763,6 +4843,8 @@ version = "0.1.6"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661"
dependencies = [ dependencies = [
"serde_core",
"writeable",
"zerovec", "zerovec",
] ]
@ -6409,6 +6491,23 @@ dependencies = [
"new_debug_unreachable", "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]] [[package]]
name = "thiserror" name = "thiserror"
version = "2.0.20" version = "2.0.20"
@ -6507,6 +6606,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643"
dependencies = [ dependencies = [
"displaydoc", "displaydoc",
"serde_core",
"zerovec", "zerovec",
] ]
@ -7553,6 +7653,7 @@ dependencies = [
"displaydoc", "displaydoc",
"yoke", "yoke",
"zerofrom", "zerofrom",
"zerovec",
] ]
[[package]] [[package]]
@ -7561,6 +7662,7 @@ version = "0.11.7"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8" checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8"
dependencies = [ dependencies = [
"serde",
"yoke", "yoke",
"zerofrom", "zerofrom",
"zerovec-derive", "zerovec-derive",

View file

@ -20,9 +20,11 @@ rig = "0.42.0"
schemars = "1" schemars = "1"
scraper = "0.27" scraper = "0.27"
serde = { version = "1.0.229", features = ["derive"] } serde = { version = "1.0.229", features = ["derive"] }
text-splitter = "0.32"
tokio = { version = "1.53.1", features = ["full"] } tokio = { version = "1.53.1", features = ["full"] }
tracing = "0.1" tracing = "0.1"
tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } tracing-subscriber = { version = "0.3.23", features = ["env-filter"] }
[dev-dependencies] [dev-dependencies]
serde_json = "1" serde_json = "1"
tempfile = "3"

View file

@ -27,22 +27,27 @@ failure was fixed, not just papered over.
## Architecture ## Architecture
Four small agents, each with one job, coordinated by plain Rust control Four small agents, each with one job, coordinated by plain Rust control
flow — not a framework's agent graph, not an LLM deciding when to stop: flow — not a framework's agent graph, not an LLM deciding when to stop —
plus an optional retrieval step when `--doc` documents are supplied:
``` ```
┌─────────────┐ approve/reject ┌──────────┐ ┌────────────┐ excerpts ┌─────────────┐ approve/reject ┌──────────┐
│ researcher │ ───────────────► │ reviewer │ │ retriever │ ───────────► │ researcher │ ───────────────► │ reviewer │
│ (tool-using)│ ◄─────────────── │ │ │(--doc only)│ │ (tool-using)│ ◄─────────────── │ │
└──────┬──────┘ gaps/feedback └────┬─────┘ └────────────┘ └──────┬──────┘ gaps/feedback └────┬─────┘
│ turn budget exhausted │ approved, │ turn budget exhausted │ approved,
│ mid-investigation │ or out of rounds │ mid-investigation │ or out of rounds
▼ ▼ ▼ ▼
┌──────────────┐ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│ summarizer │──findings───►│ writer │──► report │ summarizer │──findings───►│ writer │──► report
│ (recovery) │ │ │ │ (recovery) │ │ │
└──────────────┘ └──────────────┘ └──────────────┘ └──────────────┘
``` ```
- **retriever** (only when `--doc` paths are given) — chunks and embeds
user-supplied documents with a dedicated embedding model, then retrieves
the excerpts most relevant to the topic once up front and folds them into
the researcher's task, cited with the same footnote scheme as web sources.
- **researcher** — a tool-calling agent (`search_web`, `fetch_page`) that - **researcher** — a tool-calling agent (`search_web`, `fetch_page`) that
gathers and cross-checks evidence, capped at a fixed model-call budget so gathers and cross-checks evidence, capped at a fixed model-call budget so
a confused model can't loop forever. a confused model can't loop forever.
@ -79,8 +84,33 @@ cargo run -- "your research topic"
# or, with tracing spans on stderr instead of the progress spinner: # or, with tracing spans on stderr instead of the progress spinner:
cargo run -- -l info "your research topic" 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 ## Project layout
Split one concern per file rather than one large module: Split one concern per file rather than one large module:
@ -94,6 +124,8 @@ Split one concern per file rather than one large module:
| `summarizer.rs` | Max-turns recovery: reconstructs findings via a model call | | `summarizer.rs` | Max-turns recovery: reconstructs findings via a model call |
| `writer.rs` | Turns findings into the final streamed report | | `writer.rs` | Turns findings into the final streamed report |
| `history.rs` | Pure, unit-tested helpers for parsing a rig chat history into usable text | | `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 | | `tools.rs` | `search_web` (SearXNG) and `fetch_page` tool implementations |
| `stream.rs` | Drains a streaming prompt response to the terminal | | `stream.rs` | Drains a streaming prompt response to the terminal |
| `progress.rs` | The terminal spinner and per-phase emoji | | `progress.rs` | The terminal spinner and per-phase emoji |

View file

@ -1,4 +1,5 @@
use clap::Parser; use clap::Parser;
use std::path::PathBuf;
pub(crate) const DEFAULT_TOPIC: &str = pub(crate) const DEFAULT_TOPIC: &str =
"What are the latest advances in running large language models locally, on consumer hardware?"; "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) /// progress spinner to switch off, since the logs already show progress)
#[arg(short = 'l', long, value_name = "LEVEL")] #[arg(short = 'l', long, value_name = "LEVEL")]
pub(crate) log_level: Option<tracing::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>,
} }

345
src/documents.rs Normal file
View file

@ -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)");
}
}

View file

@ -270,7 +270,11 @@ mod tests {
serde_json::json!({ "query": "test" }), 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"), Message::assistant("that source looks solid"),
]; ];

View file

@ -1,12 +1,14 @@
use clap::Parser; use clap::Parser;
mod cli; mod cli;
mod documents;
mod history; mod history;
mod models; mod models;
mod observability; mod observability;
mod progress; mod progress;
mod research; mod research;
mod researcher; mod researcher;
mod retrieval;
mod review; mod review;
mod stream; mod stream;
mod summarizer; mod summarizer;
@ -30,7 +32,7 @@ async fn main() -> anyhow::Result<()> {
// The report streams to stdout as the writer generates it, so nothing // The report streams to stdout as the writer generates it, so nothing
// left to print here — the return value only matters to callers that // left to print here — the return value only matters to callers that
// embed `research` rather than running it as this binary. // embed `research` rather than running it as this binary.
research::research(&topic, show_progress).await?; research::research(&topic, &cli.docs, show_progress).await?;
Ok(()) Ok(())
} }

View file

@ -6,3 +6,9 @@
/// passes instead. /// passes instead.
pub(crate) const RESEARCHER_MODEL: &str = "gemma4:26b"; pub(crate) const RESEARCHER_MODEL: &str = "gemma4:26b";
pub(crate) const WRITER_MODEL: &str = "gemma4-e4b:latest"; 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";

View file

@ -11,6 +11,7 @@ pub(crate) const REVIEW_EMOJI: &str = "🧐";
pub(crate) const REJECTED_EMOJI: &str = ""; pub(crate) const REJECTED_EMOJI: &str = "";
pub(crate) const REPORT_EMOJI: &str = "✍️"; pub(crate) const REPORT_EMOJI: &str = "✍️";
pub(crate) const SUMMARIZE_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 /// The spinner currently on screen, if any — set by `Spinner::start` and
/// cleared on drop. Tool implementations don't otherwise have a handle to /// 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. /// 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 /// A no-op when progress display is off (no spinner was ever started, so
/// `active()` stays empty) or between phases (the previous `Spinner` has /// `active()` stays empty) or between phases (the previous `Spinner` has

View file

@ -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::researcher::gather_findings;
use crate::retrieval;
use crate::review::{self, Review}; use crate::review::{self, Review};
use crate::writer::write_report; use crate::writer::write_report;
use rig::client::Nothing; use rig::client::Nothing;
use rig::providers::ollama; use rig::providers::ollama;
use std::path::PathBuf;
/// Research/review rounds before giving up and writing the report from /// Research/review rounds before giving up and writing the report from
/// whatever the last pass produced, rather than looping forever on a topic /// 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 /// 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 /// reviewer's feedback folded in until it approves or the round budget runs
/// out, then writing the report from whatever the last pass produced. /// 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 client = ollama::Client::new(Nothing)?;
let doc_context = build_doc_context(&client, topic, docs, show_progress).await?;
let mut findings = String::new(); let mut findings = String::new();
let mut feedback: Option<Review> = None; let mut feedback: Option<Review> = None;
let mut incomplete = false; let mut incomplete = false;
for round in 1..=MAX_RESEARCH_ROUNDS { 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; findings = gathered.findings;
incomplete = gathered.incomplete; incomplete = gathered.incomplete;
@ -41,15 +57,42 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result
break; break;
} }
if show_progress { progress::report(
eprintln!( show_progress,
format!(
"{REJECTED_EMOJI} Findings rejected — revising for round {}...", "{REJECTED_EMOJI} Findings rejected — revising for round {}...",
round + 1 round + 1
),
); );
}
feedback = Some(review); feedback = Some(review);
} }
write_report(&client, topic, &findings, incomplete, show_progress).await 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))
}

View file

@ -19,11 +19,12 @@ pub(crate) struct GatheredFindings {
/// Wraps the tool-calling research loop in its own span so it's visible as a /// 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 /// single unit in traces, distinct from the writing and review phases and
/// nesting rig's own per-turn `chat`/`execute_tool` spans underneath it. /// 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( pub(crate) async fn gather_findings(
client: &ollama::Client, client: &ollama::Client,
topic: &str, topic: &str,
feedback: Option<&Review>, feedback: Option<&Review>,
doc_context: Option<&str>,
round: usize, round: usize,
show_progress: bool, show_progress: bool,
) -> anyhow::Result<GatheredFindings> { ) -> 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 spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching..."));
let run_result = researcher let run_result = researcher
.runner(task) .runner(task)

81
src/retrieval.rs Normal file
View file

@ -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"))
}