Compare commits
No commits in common. "master" and "worktree-fix-ollama-json-extraction" have entirely different histories.
master
...
worktree-f
12 changed files with 23 additions and 670 deletions
102
Cargo.lock
generated
102
Cargo.lock
generated
|
|
@ -517,18 +517,6 @@ 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"
|
||||
|
|
@ -1080,15 +1068,6 @@ 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"
|
||||
|
|
@ -2060,17 +2039,6 @@ 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"
|
||||
|
|
@ -2129,8 +2097,6 @@ dependencies = [
|
|||
"scraper",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"text-splitter",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
|
|
@ -2908,32 +2874,11 @@ 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"
|
||||
|
|
@ -2983,8 +2928,6 @@ checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428"
|
|||
dependencies = [
|
||||
"displaydoc",
|
||||
"icu_locale_core",
|
||||
"serde",
|
||||
"stable_deref_trait",
|
||||
"writeable",
|
||||
"yoke",
|
||||
"zerofrom",
|
||||
|
|
@ -2992,29 +2935,6 @@ 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"
|
||||
|
|
@ -4843,8 +4763,6 @@ version = "0.1.6"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
"writeable",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
|
|
@ -6491,23 +6409,6 @@ 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"
|
||||
|
|
@ -6606,7 +6507,6 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643"
|
||||
dependencies = [
|
||||
"displaydoc",
|
||||
"serde_core",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
|
|
@ -7653,7 +7553,6 @@ dependencies = [
|
|||
"displaydoc",
|
||||
"yoke",
|
||||
"zerofrom",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -7662,7 +7561,6 @@ version = "0.11.7"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"yoke",
|
||||
"zerofrom",
|
||||
"zerovec-derive",
|
||||
|
|
|
|||
|
|
@ -20,11 +20,9 @@ rig = "0.42.0"
|
|||
schemars = "1"
|
||||
scraper = "0.27"
|
||||
serde = { version = "1.0.229", features = ["derive"] }
|
||||
text-splitter = "0.32"
|
||||
tokio = { version = "1.53.1", features = ["full"] }
|
||||
tracing = "0.1"
|
||||
tracing-subscriber = { version = "0.3.23", features = ["env-filter"] }
|
||||
|
||||
[dev-dependencies]
|
||||
serde_json = "1"
|
||||
tempfile = "3"
|
||||
|
|
|
|||
42
README.md
42
README.md
|
|
@ -27,14 +27,13 @@ failure was fixed, not just papered over.
|
|||
## Architecture
|
||||
|
||||
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 —
|
||||
plus an optional retrieval step when `--doc` documents are supplied:
|
||||
flow — not a framework's agent graph, not an LLM deciding when to stop:
|
||||
|
||||
```
|
||||
┌────────────┐ excerpts ┌─────────────┐ approve/reject ┌──────────┐
|
||||
│ retriever │ ───────────► │ researcher │ ───────────────► │ reviewer │
|
||||
│(--doc only)│ │ (tool-using)│ ◄─────────────── │ │
|
||||
└────────────┘ └──────┬──────┘ gaps/feedback └────┬─────┘
|
||||
┌─────────────┐ approve/reject ┌──────────┐
|
||||
│ researcher │ ───────────────► │ reviewer │
|
||||
│ (tool-using)│ ◄─────────────── │ │
|
||||
└──────┬──────┘ gaps/feedback └────┬─────┘
|
||||
│ turn budget exhausted │ approved,
|
||||
│ mid-investigation │ or out of rounds
|
||||
▼ ▼
|
||||
|
|
@ -44,10 +43,6 @@ plus an optional retrieval step when `--doc` documents are supplied:
|
|||
└──────────────┘ └──────────────┘
|
||||
```
|
||||
|
||||
- **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
|
||||
gathers and cross-checks evidence, capped at a fixed model-call budget so
|
||||
a confused model can't loop forever.
|
||||
|
|
@ -84,33 +79,8 @@ cargo run -- "your research topic"
|
|||
|
||||
# or, with tracing spans on stderr instead of the progress spinner:
|
||||
cargo run -- -l info "your research topic"
|
||||
|
||||
# give the researcher your own documents to draw on, alongside the web —
|
||||
# repeatable, and a directory contributes every file directly inside it
|
||||
# (one level deep, not recursive):
|
||||
cargo run -- --doc ./notes.txt --doc ./research-docs/ "your research topic"
|
||||
```
|
||||
|
||||
Uploaded documents are chunked (see `documents.rs`), embedded with a
|
||||
dedicated embedding model (see `EMBEDDING_MODEL` in
|
||||
[`src/models.rs`](./src/models.rs)) into an in-memory vector index, then the
|
||||
excerpts most relevant to the topic are retrieved and folded into the
|
||||
researcher's task alongside anything it finds on the web — the same
|
||||
footnote-citation scheme applies to both.
|
||||
|
||||
Known limitation: retrieval returns a fixed top-N chunks
|
||||
(`retrieval::TOP_N_EXCERPTS`). A document with several chunks that all read
|
||||
as similar to the query — several incident reports, several revisions of
|
||||
the same section — can crowd out the one chunk that actually answers it,
|
||||
since only the top N by similarity are ever returned regardless of how many
|
||||
plausible candidates exist. Reproduced deliberately (a 13.7 KB / 12-chunk
|
||||
document with 6 near-identical "incident report" sections, only one of
|
||||
which had the real answer, was built specifically to stress this — the top
|
||||
5 slots filled entirely with distractors and the answer chunk was
|
||||
excluded), so it's a real edge case, not a hypothetical. Not fixed for now
|
||||
since it takes a document engineered to trigger it, but worth knowing if a
|
||||
report seems to be missing something you know is in an uploaded document.
|
||||
|
||||
## Project layout
|
||||
|
||||
Split one concern per file rather than one large module:
|
||||
|
|
@ -124,8 +94,6 @@ 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 |
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
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?";
|
||||
|
|
@ -17,11 +16,4 @@ pub(crate) struct Doubleo7 {
|
|||
/// progress spinner to switch off, since the logs already show progress)
|
||||
#[arg(short = 'l', long, value_name = "LEVEL")]
|
||||
pub(crate) log_level: Option<tracing::Level>,
|
||||
|
||||
/// A file, or a directory of files, to embed and make available to the
|
||||
/// researcher as retrieved excerpts (repeatable). Directories are
|
||||
/// scanned one level deep — non-directory entries only, subdirectories
|
||||
/// are skipped rather than walked.
|
||||
#[arg(short = 'd', long = "doc", value_name = "PATH")]
|
||||
pub(crate) docs: Vec<PathBuf>,
|
||||
}
|
||||
|
|
|
|||
345
src/documents.rs
345
src/documents.rs
|
|
@ -1,345 +0,0 @@
|
|||
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)");
|
||||
}
|
||||
}
|
||||
|
|
@ -270,11 +270,7 @@ 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"),
|
||||
];
|
||||
|
||||
|
|
|
|||
|
|
@ -1,14 +1,12 @@
|
|||
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;
|
||||
|
|
@ -32,7 +30,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, &cli.docs, show_progress).await?;
|
||||
research::research(&topic, show_progress).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -6,9 +6,3 @@
|
|||
/// 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";
|
||||
|
|
|
|||
|
|
@ -11,7 +11,6 @@ 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
|
||||
|
|
@ -69,15 +68,6 @@ impl Drop for Spinner {
|
|||
}
|
||||
}
|
||||
|
||||
/// Prints a status line, but only when progress display is on — matching
|
||||
/// `Spinner`'s own `enabled` gate so a call site doesn't need to repeat the
|
||||
/// `if show_progress` check inline.
|
||||
pub(crate) fn report(enabled: bool, message: impl Into<String>) {
|
||||
if enabled {
|
||||
eprintln!("{}", message.into());
|
||||
}
|
||||
}
|
||||
|
||||
/// Updates the message of whatever spinner is currently running, if any.
|
||||
/// A no-op when progress display is off (no spinner was ever started, so
|
||||
/// `active()` stays empty) or between phases (the previous `Spinner` has
|
||||
|
|
|
|||
|
|
@ -1,12 +1,9 @@
|
|||
use crate::documents;
|
||||
use crate::progress::{self, REJECTED_EMOJI};
|
||||
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
|
||||
|
|
@ -17,29 +14,16 @@ 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,
|
||||
docs: &[PathBuf],
|
||||
show_progress: bool,
|
||||
) -> anyhow::Result<String> {
|
||||
pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result<String> {
|
||||
let client = ollama::Client::new(Nothing)?;
|
||||
|
||||
let doc_context = build_doc_context(&client, topic, docs, show_progress).await?;
|
||||
|
||||
let mut findings = String::new();
|
||||
let mut feedback: Option<Review> = None;
|
||||
let mut incomplete = false;
|
||||
|
||||
for round in 1..=MAX_RESEARCH_ROUNDS {
|
||||
let gathered = gather_findings(
|
||||
&client,
|
||||
topic,
|
||||
feedback.as_ref(),
|
||||
doc_context.as_deref(),
|
||||
round,
|
||||
show_progress,
|
||||
)
|
||||
.await?;
|
||||
let gathered =
|
||||
gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?;
|
||||
findings = gathered.findings;
|
||||
incomplete = gathered.incomplete;
|
||||
|
||||
|
|
@ -57,42 +41,15 @@ pub(crate) async fn research(
|
|||
break;
|
||||
}
|
||||
|
||||
progress::report(
|
||||
show_progress,
|
||||
format!(
|
||||
if show_progress {
|
||||
eprintln!(
|
||||
"{REJECTED_EMOJI} Findings rejected — revising for round {}...",
|
||||
round + 1
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
feedback = Some(review);
|
||||
}
|
||||
|
||||
write_report(&client, topic, &findings, incomplete, show_progress).await
|
||||
}
|
||||
|
||||
/// Embeds any uploaded documents and retrieves the excerpts most relevant to
|
||||
/// `topic`, once, up front — nothing later in the loop changes which
|
||||
/// excerpts are relevant, so there's no reason to repeat this per round.
|
||||
/// `None` when no documents were given, or none of them yielded a usable
|
||||
/// excerpt (an empty `docs` list is deliberately not an error).
|
||||
async fn build_doc_context(
|
||||
client: &ollama::Client,
|
||||
topic: &str,
|
||||
docs: &[PathBuf],
|
||||
show_progress: bool,
|
||||
) -> anyhow::Result<Option<String>> {
|
||||
if docs.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let documents = documents::collect_documents(docs)?;
|
||||
if documents.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let index = retrieval::build_index(client, documents, show_progress).await?;
|
||||
let excerpts = retrieval::retrieve_relevant(&index, topic, show_progress).await?;
|
||||
|
||||
Ok((!excerpts.is_empty()).then_some(excerpts))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,12 +19,11 @@ 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, doc_context), fields(gen_ai.agent.name = "researcher"))]
|
||||
#[tracing::instrument(skip(client, feedback), fields(gen_ai.agent.name = "researcher"))]
|
||||
pub(crate) async fn gather_findings(
|
||||
client: &ollama::Client,
|
||||
topic: &str,
|
||||
feedback: Option<&Review>,
|
||||
doc_context: Option<&str>,
|
||||
round: usize,
|
||||
show_progress: bool,
|
||||
) -> anyhow::Result<GatheredFindings> {
|
||||
|
|
@ -69,17 +68,6 @@ pub(crate) async fn gather_findings(
|
|||
),
|
||||
};
|
||||
|
||||
let task = match doc_context {
|
||||
Some(context) => format!(
|
||||
"{task}\n\n\
|
||||
Relevant excerpts from documents the user uploaded — treat these as trusted primary \
|
||||
sources alongside anything you find on the web, and cite them with the same \
|
||||
footnote scheme (their Sources entry can just be the document path shown below):\n\
|
||||
{context}"
|
||||
),
|
||||
None => task,
|
||||
};
|
||||
|
||||
let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching..."));
|
||||
let run_result = researcher
|
||||
.runner(task)
|
||||
|
|
|
|||
|
|
@ -1,81 +0,0 @@
|
|||
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"))
|
||||
}
|
||||
Loading…
Reference in a new issue