Compare commits
19 commits
chore/rebr
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| 8d7f8b6683 | |||
|
0790b2b6d0 | ||
| 90502a3eed | |||
|
|
ba9b61fcf3 | ||
|
|
f916787962 | ||
|
|
6632d63239 | ||
|
|
5e55861a3f | ||
|
|
19c92ccd52 | ||
|
|
53b4f572df | ||
| ad2ba1181a | |||
|
|
0558ac5c4f | ||
|
|
ce0bf8dafd | ||
|
|
9e7291214b | ||
|
|
cd580f3bec | ||
|
|
8f4af514a3 | ||
| a30c18ac44 | |||
|
|
b5f12a500c | ||
|
|
ff9f9455a6 | ||
|
|
ec15893c9c |
15 changed files with 751 additions and 94 deletions
151
Cargo.lock
generated
151
Cargo.lock
generated
|
|
@ -517,6 +517,18 @@ version = "1.1.2"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
|
||||
|
||||
[[package]]
|
||||
name = "auto_enums"
|
||||
version = "0.8.10"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3091d68264354f211516b91dce6f71046e444fab1867716035f736667243affb"
|
||||
dependencies = [
|
||||
"derive_utils",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "autocfg"
|
||||
version = "1.5.1"
|
||||
|
|
@ -1068,6 +1080,15 @@ version = "0.8.7"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
|
||||
|
||||
[[package]]
|
||||
name = "core_maths"
|
||||
version = "0.1.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "77745e017f5edba1a9c1d854f6f3a52dac8a12dd5af5d2f54aecf61e43d80d30"
|
||||
dependencies = [
|
||||
"libm",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cpufeatures"
|
||||
version = "0.2.17"
|
||||
|
|
@ -2039,6 +2060,17 @@ dependencies = [
|
|||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "derive_utils"
|
||||
version = "0.16.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "dc05a5d33db20c784f873e84934ad94bb209a090987ac5f62fede2c178234f23"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.3",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "digest"
|
||||
version = "0.10.7"
|
||||
|
|
@ -2097,6 +2129,8 @@ dependencies = [
|
|||
"scraper",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tempfile",
|
||||
"text-splitter",
|
||||
"tokio",
|
||||
"tracing",
|
||||
"tracing-subscriber",
|
||||
|
|
@ -2874,11 +2908,32 @@ checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb"
|
|||
dependencies = [
|
||||
"displaydoc",
|
||||
"litemap",
|
||||
"serde",
|
||||
"tinystr",
|
||||
"writeable",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_locale_fallback"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "251af8e57c9400e3eb58242fe5b8b1152b2a64fdf4cf632f923c38ccee6f2fa9"
|
||||
dependencies = [
|
||||
"icu_locale_core",
|
||||
"icu_locale_fallback_data",
|
||||
"icu_provider",
|
||||
"potential_utf",
|
||||
"tinystr",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_locale_fallback_data"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "decf2a22ec8fa68f1a0c1129a3f8583f8f8bc24e8b9ccbe98ead99f62a4dc3a8"
|
||||
|
||||
[[package]]
|
||||
name = "icu_normalizer"
|
||||
version = "2.3.0"
|
||||
|
|
@ -2928,6 +2983,8 @@ checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428"
|
|||
dependencies = [
|
||||
"displaydoc",
|
||||
"icu_locale_core",
|
||||
"serde",
|
||||
"stable_deref_trait",
|
||||
"writeable",
|
||||
"yoke",
|
||||
"zerofrom",
|
||||
|
|
@ -2935,6 +2992,29 @@ dependencies = [
|
|||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_segmenter"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "82d07aafccd67af15d02512a6adf5896fbc5ed00f2e99b471d2efa14016db3db"
|
||||
dependencies = [
|
||||
"core_maths",
|
||||
"icu_collections",
|
||||
"icu_locale_fallback",
|
||||
"icu_provider",
|
||||
"icu_segmenter_data",
|
||||
"potential_utf",
|
||||
"smallvec",
|
||||
"utf8_iter",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "icu_segmenter_data"
|
||||
version = "2.3.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ae293c039020f9ec10710af98d29ce6aa2051486638b49c9a6409f3b4a9e98ad"
|
||||
|
||||
[[package]]
|
||||
name = "ident_case"
|
||||
version = "1.0.1"
|
||||
|
|
@ -4763,6 +4843,8 @@ version = "0.1.6"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
"writeable",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
|
|
@ -5414,9 +5496,9 @@ checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4"
|
|||
|
||||
[[package]]
|
||||
name = "rig"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2ce03971e6115d30ef53fb3244d06718a4e62bbde82c103065600c09459b989a"
|
||||
checksum = "b989f4060e5f3c1ea4b8b6999f6759949d820704286cc6976b7a688cc0133029"
|
||||
dependencies = [
|
||||
"rig-agent",
|
||||
"rig-core",
|
||||
|
|
@ -5429,9 +5511,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-agent"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2b0796bbf47d7b76670401aac975bc619cf7fba3482b22dfe14992edaa9c2e04"
|
||||
checksum = "b66477c0d0a786358d078e7e3541a73089dd9009100d6e2b3fe3b0d730e8550f"
|
||||
dependencies = [
|
||||
"async-stream",
|
||||
"fastrand",
|
||||
|
|
@ -5451,9 +5533,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-core"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "35f5520515ae8f6851adcbc6fde9eea8e96f657418c062e16c82cd81cce44e8e"
|
||||
checksum = "432d83e0facf16749f91fe729cbffca84437e8062d2f4e92f4f12e903693922d"
|
||||
dependencies = [
|
||||
"as-any",
|
||||
"async-stream",
|
||||
|
|
@ -5475,6 +5557,7 @@ dependencies = [
|
|||
"schemars 1.2.2",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"sha2",
|
||||
"thiserror",
|
||||
"tokio",
|
||||
"tokio-tungstenite",
|
||||
|
|
@ -5485,9 +5568,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-derive"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "eb868fcebdf3ba425e3afad2e4926bb6d9e1188a856843b00bcee2e15c07424f"
|
||||
checksum = "de0a33f1bac45f16e50146c248bcbbfaa44518c7252d274e972c7f4ad71aaba7"
|
||||
dependencies = [
|
||||
"convert_case",
|
||||
"proc-macro-crate",
|
||||
|
|
@ -5498,9 +5581,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-fastembed"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "69ba93356242940918fac76a4b1cfe78b964f757e5cb65bb3feebd597635cfb0"
|
||||
checksum = "f1dff99f9bea13108bb14df5431b6fc0738beb87f1fe5ebfcf26ae9fe54a188d"
|
||||
dependencies = [
|
||||
"fastembed",
|
||||
"rig-core",
|
||||
|
|
@ -5512,9 +5595,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-helixdb"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d1d40ae8fa89af33793626b898ec527e246183ef8b4b5095afee5f6a65cbbb1a"
|
||||
checksum = "b5fa2b188a1e9e65d29c34738c97465f0822c57540d6d45fd4f04dad906c966a"
|
||||
dependencies = [
|
||||
"reqwest 0.13.4",
|
||||
"rig-core",
|
||||
|
|
@ -5525,12 +5608,11 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-lancedb"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f703b9e081ee4d77e60953598a5071a5cff4b8b702a2f06d251310950b0dae7c"
|
||||
checksum = "c5fe3555202c5bd4e30948414c2e234fddd9260d4520046733fcf0d58718f3a1"
|
||||
dependencies = [
|
||||
"arrow-array",
|
||||
"deranged",
|
||||
"futures",
|
||||
"lancedb",
|
||||
"rig-core",
|
||||
|
|
@ -5540,9 +5622,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "rig-milvus"
|
||||
version = "0.41.0"
|
||||
version = "0.42.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "76662eb7f75de4e54b5f71e616430dc7f75d96d32f109d9435ad426b511eb565"
|
||||
checksum = "af9c7b7ff7ed9c16da983a21f1e81f425c9805355427a92f4d2d97ba1fe2aa12"
|
||||
dependencies = [
|
||||
"reqwest 0.13.4",
|
||||
"rig-core",
|
||||
|
|
@ -6409,6 +6491,23 @@ dependencies = [
|
|||
"new_debug_unreachable",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "text-splitter"
|
||||
version = "0.32.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3f3eeec76988617ff1434d754d7e8e197be2cb8981ea77c17e32f3f7a2c6f95e"
|
||||
dependencies = [
|
||||
"ahash",
|
||||
"auto_enums",
|
||||
"either",
|
||||
"icu_provider",
|
||||
"icu_segmenter",
|
||||
"itertools 0.14.0",
|
||||
"memchr",
|
||||
"strum 0.28.0",
|
||||
"thiserror",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "thiserror"
|
||||
version = "2.0.20"
|
||||
|
|
@ -6507,6 +6606,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643"
|
||||
dependencies = [
|
||||
"displaydoc",
|
||||
"serde_core",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
|
|
@ -6610,9 +6710,9 @@ dependencies = [
|
|||
|
||||
[[package]]
|
||||
name = "tokio-tungstenite"
|
||||
version = "0.28.0"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857"
|
||||
checksum = "8f72a05e828585856dacd553fba484c242c46e391fb0e58917c942ee9202915c"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"log",
|
||||
|
|
@ -6799,9 +6899,9 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
|
|||
|
||||
[[package]]
|
||||
name = "tungstenite"
|
||||
version = "0.28.0"
|
||||
version = "0.29.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442"
|
||||
checksum = "6c01152af293afb9c7c2a57e4b559c5620b421f6d133261c60dd2d0cdb38e6b8"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"data-encoding",
|
||||
|
|
@ -6813,7 +6913,6 @@ dependencies = [
|
|||
"rustls-pki-types",
|
||||
"sha1",
|
||||
"thiserror",
|
||||
"utf-8",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -6934,12 +7033,6 @@ dependencies = [
|
|||
"serde",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "utf-8"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
|
||||
|
||||
[[package]]
|
||||
name = "utf8-ranges"
|
||||
version = "1.0.5"
|
||||
|
|
@ -7560,6 +7653,7 @@ dependencies = [
|
|||
"displaydoc",
|
||||
"yoke",
|
||||
"zerofrom",
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -7568,6 +7662,7 @@ version = "0.11.7"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8"
|
||||
dependencies = [
|
||||
"serde",
|
||||
"yoke",
|
||||
"zerofrom",
|
||||
"zerovec-derive",
|
||||
|
|
|
|||
|
|
@ -16,13 +16,15 @@ clap = { version = "4", features = ["derive"] }
|
|||
futures = "0.3"
|
||||
indicatif = "0.18.6"
|
||||
reqwest = { version = "0.13.4", features = ["query", "json"] }
|
||||
rig = "0.41.0"
|
||||
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"
|
||||
|
|
|
|||
50
README.md
50
README.md
|
|
@ -27,22 +27,27 @@ 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:
|
||||
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 ┌──────────┐
|
||||
│ researcher │ ───────────────► │ reviewer │
|
||||
│ (tool-using)│ ◄─────────────── │ │
|
||||
└──────┬──────┘ gaps/feedback └────┬─────┘
|
||||
┌────────────┐ excerpts ┌─────────────┐ approve/reject ┌──────────┐
|
||||
│ retriever │ ───────────► │ researcher │ ───────────────► │ reviewer │
|
||||
│(--doc only)│ │ (tool-using)│ ◄─────────────── │ │
|
||||
└────────────┘ └──────┬──────┘ gaps/feedback └────┬─────┘
|
||||
│ turn budget exhausted │ approved,
|
||||
│ mid-investigation │ or out of rounds
|
||||
▼ ▼
|
||||
┌──────────────┐ ┌──────────────┐
|
||||
│ summarizer │──findings───►│ writer │──► report
|
||||
│ (recovery) │ │ │
|
||||
└──────────────┘ └──────────────┘
|
||||
┌──────────────┐ ┌──────────────┐
|
||||
│ summarizer │──findings───►│ writer │──► report
|
||||
│ (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
|
||||
gathers and cross-checks evidence, capped at a fixed model-call budget so
|
||||
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:
|
||||
cargo run -- -l info "your research topic"
|
||||
|
||||
# give the researcher your own documents to draw on, alongside the web —
|
||||
# repeatable, and a directory contributes every file directly inside it
|
||||
# (one level deep, not recursive):
|
||||
cargo run -- --doc ./notes.txt --doc ./research-docs/ "your research topic"
|
||||
```
|
||||
|
||||
Uploaded documents are chunked (see `documents.rs`), embedded with a
|
||||
dedicated embedding model (see `EMBEDDING_MODEL` in
|
||||
[`src/models.rs`](./src/models.rs)) into an in-memory vector index, then the
|
||||
excerpts most relevant to the topic are retrieved and folded into the
|
||||
researcher's task alongside anything it finds on the web — the same
|
||||
footnote-citation scheme applies to both.
|
||||
|
||||
Known limitation: retrieval returns a fixed top-N chunks
|
||||
(`retrieval::TOP_N_EXCERPTS`). A document with several chunks that all read
|
||||
as similar to the query — several incident reports, several revisions of
|
||||
the same section — can crowd out the one chunk that actually answers it,
|
||||
since only the top N by similarity are ever returned regardless of how many
|
||||
plausible candidates exist. Reproduced deliberately (a 13.7 KB / 12-chunk
|
||||
document with 6 near-identical "incident report" sections, only one of
|
||||
which had the real answer, was built specifically to stress this — the top
|
||||
5 slots filled entirely with distractors and the answer chunk was
|
||||
excluded), so it's a real edge case, not a hypothetical. Not fixed for now
|
||||
since it takes a document engineered to trigger it, but worth knowing if a
|
||||
report seems to be missing something you know is in an uploaded document.
|
||||
|
||||
## Project layout
|
||||
|
||||
Split one concern per file rather than one large module:
|
||||
|
|
@ -94,6 +124,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 |
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
use clap::Parser;
|
||||
use std::path::PathBuf;
|
||||
|
||||
pub(crate) const DEFAULT_TOPIC: &str =
|
||||
"What are the latest advances in running large language models locally, on consumer hardware?";
|
||||
|
|
@ -16,4 +17,11 @@ pub(crate) struct Doubleo7 {
|
|||
/// progress spinner to switch off, since the logs already show progress)
|
||||
#[arg(short = 'l', long, value_name = "LEVEL")]
|
||||
pub(crate) log_level: Option<tracing::Level>,
|
||||
|
||||
/// A file, or a directory of files, to embed and make available to the
|
||||
/// researcher as retrieved excerpts (repeatable). Directories are
|
||||
/// scanned one level deep — non-directory entries only, subdirectories
|
||||
/// are skipped rather than walked.
|
||||
#[arg(short = 'd', long = "doc", value_name = "PATH")]
|
||||
pub(crate) docs: Vec<PathBuf>,
|
||||
}
|
||||
|
|
|
|||
345
src/documents.rs
Normal file
345
src/documents.rs
Normal 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)");
|
||||
}
|
||||
}
|
||||
|
|
@ -149,11 +149,11 @@ mod tests {
|
|||
fn assistant_text_ignores_tool_calls() {
|
||||
let message = Message::Assistant {
|
||||
id: None,
|
||||
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||
content: vec![AssistantContent::tool_call(
|
||||
"call-1",
|
||||
"search_web",
|
||||
serde_json::json!({ "query": "test" }),
|
||||
)),
|
||||
serde_json::json!({ "query": "test"}),
|
||||
)],
|
||||
};
|
||||
assert!(extract_assistant_text(&message).is_empty());
|
||||
}
|
||||
|
|
@ -166,13 +166,13 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn tool_result_text_extracts_and_truncates() {
|
||||
let short = Message::tool_result("call-1", "short result");
|
||||
let short = Message::tool_result("call-1", "call-1-name", "short result");
|
||||
assert_eq!(
|
||||
tool_result_text(&short, 100),
|
||||
vec!["short result".to_string()]
|
||||
);
|
||||
|
||||
let long = Message::tool_result("call-2", "0123456789");
|
||||
let long = Message::tool_result("call-2", "call-2-result", "0123456789");
|
||||
assert_eq!(
|
||||
tool_result_text(&long, 5),
|
||||
vec!["01234 ...[truncated]".to_string()]
|
||||
|
|
@ -191,11 +191,11 @@ mod tests {
|
|||
Message::system("preamble"),
|
||||
Message::Assistant {
|
||||
id: None,
|
||||
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||
content: vec![AssistantContent::tool_call(
|
||||
"call-1",
|
||||
"search_web",
|
||||
serde_json::json!({ "query": "test" }),
|
||||
)),
|
||||
)],
|
||||
},
|
||||
];
|
||||
|
||||
|
|
@ -209,9 +209,9 @@ mod tests {
|
|||
fn partial_findings_from_history_collects_assistant_text_and_tool_results_in_order() {
|
||||
let history = vec![
|
||||
Message::assistant("Checking sources..."),
|
||||
Message::tool_result("call-1", "Result A [1]"),
|
||||
Message::tool_result("call-1", "call-1-name", "Result A [1]"),
|
||||
Message::assistant("Cross-checking..."),
|
||||
Message::tool_result("call-2", "Result B [2]"),
|
||||
Message::tool_result("call-2", "call-2-name", "Result B [2]"),
|
||||
];
|
||||
|
||||
assert_eq!(
|
||||
|
|
@ -233,11 +233,11 @@ mod tests {
|
|||
fn transcript_lines_renders_tool_calls_with_their_arguments() {
|
||||
let message = Message::Assistant {
|
||||
id: None,
|
||||
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||
content: vec![AssistantContent::tool_call(
|
||||
"call-1",
|
||||
"search_web",
|
||||
serde_json::json!({ "query": "uruguay senior engineers" }),
|
||||
)),
|
||||
)],
|
||||
};
|
||||
assert_eq!(
|
||||
transcript_lines(&message),
|
||||
|
|
@ -247,7 +247,7 @@ mod tests {
|
|||
|
||||
#[test]
|
||||
fn transcript_lines_prefixes_and_truncates_tool_results() {
|
||||
let short = Message::tool_result("call-1", "found via search [1]");
|
||||
let short = Message::tool_result("call-1", "call-1-name", "found via search [1]");
|
||||
assert_eq!(
|
||||
transcript_lines(&short),
|
||||
vec!["Result: found via search [1]".to_string()]
|
||||
|
|
@ -264,13 +264,17 @@ mod tests {
|
|||
let history = vec![
|
||||
Message::Assistant {
|
||||
id: None,
|
||||
content: rig::OneOrMany::one(AssistantContent::tool_call(
|
||||
content: vec![AssistantContent::tool_call(
|
||||
"call-1",
|
||||
"search_web",
|
||||
serde_json::json!({ "query": "test" }),
|
||||
)),
|
||||
)],
|
||||
},
|
||||
Message::tool_result("call-1", "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,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(())
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ pub(crate) const REVIEW_EMOJI: &str = "🧐";
|
|||
pub(crate) const REJECTED_EMOJI: &str = "❌";
|
||||
pub(crate) const REPORT_EMOJI: &str = "✍️";
|
||||
pub(crate) const SUMMARIZE_EMOJI: &str = "🧩";
|
||||
pub(crate) const EMBED_EMOJI: &str = "📚";
|
||||
|
||||
/// The spinner currently on screen, if any — set by `Spinner::start` and
|
||||
/// cleared on drop. Tool implementations don't otherwise have a handle to
|
||||
|
|
@ -68,6 +69,15 @@ impl Drop for Spinner {
|
|||
}
|
||||
}
|
||||
|
||||
/// Prints a status line, but only when progress display is on — matching
|
||||
/// `Spinner`'s own `enabled` gate so a call site doesn't need to repeat the
|
||||
/// `if show_progress` check inline.
|
||||
pub(crate) fn report(enabled: bool, message: impl Into<String>) {
|
||||
if enabled {
|
||||
eprintln!("{}", message.into());
|
||||
}
|
||||
}
|
||||
|
||||
/// Updates the message of whatever spinner is currently running, if any.
|
||||
/// A no-op when progress display is off (no spinner was ever started, so
|
||||
/// `active()` stays empty) or between phases (the previous `Spinner` has
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
use crate::progress::REJECTED_EMOJI;
|
||||
use crate::documents;
|
||||
use crate::progress::{self, REJECTED_EMOJI};
|
||||
use crate::researcher::gather_findings;
|
||||
use crate::retrieval;
|
||||
use crate::review::{self, Review};
|
||||
use crate::writer::write_report;
|
||||
use rig::client::Nothing;
|
||||
use rig::providers::ollama;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Research/review rounds before giving up and writing the report from
|
||||
/// whatever the last pass produced, rather than looping forever on a topic
|
||||
|
|
@ -14,45 +17,82 @@ const MAX_RESEARCH_ROUNDS: usize = 3;
|
|||
/// not a model, in charge of when to stop — re-running research with the
|
||||
/// reviewer's feedback folded in until it approves or the round budget runs
|
||||
/// out, then writing the report from whatever the last pass produced.
|
||||
pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result<String> {
|
||||
pub(crate) async fn research(
|
||||
topic: &str,
|
||||
docs: &[PathBuf],
|
||||
show_progress: bool,
|
||||
) -> anyhow::Result<String> {
|
||||
let client = ollama::Client::new(Nothing)?;
|
||||
|
||||
let doc_context = build_doc_context(&client, topic, docs, show_progress).await?;
|
||||
|
||||
let mut findings = String::new();
|
||||
let mut feedback: Option<Review> = None;
|
||||
let mut incomplete = false;
|
||||
|
||||
for round in 1..=MAX_RESEARCH_ROUNDS {
|
||||
let gathered =
|
||||
gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?;
|
||||
let gathered = gather_findings(
|
||||
&client,
|
||||
topic,
|
||||
feedback.as_ref(),
|
||||
doc_context.as_deref(),
|
||||
round,
|
||||
show_progress,
|
||||
)
|
||||
.await?;
|
||||
findings = gathered.findings;
|
||||
incomplete = gathered.incomplete;
|
||||
|
||||
// The researcher ran out of turns mid-investigation rather than
|
||||
// concluding on its own — another round would just repeat the same
|
||||
// dead end, so stop and write up whatever was gathered.
|
||||
// Prevent repeat iterations where researcher hits same dead end.
|
||||
if incomplete {
|
||||
tracing::info!(round, "researcher exhausted its turn budget");
|
||||
break;
|
||||
}
|
||||
|
||||
let review = review::review_findings(&client, topic, &findings, show_progress).await?;
|
||||
let approved = review.approved;
|
||||
|
||||
tracing::info!(round, approved, "review verdict");
|
||||
tracing::info!(round, review.approved, "review verdict");
|
||||
|
||||
if approved || round == MAX_RESEARCH_ROUNDS {
|
||||
if review.approved || round == MAX_RESEARCH_ROUNDS {
|
||||
break;
|
||||
}
|
||||
|
||||
if show_progress {
|
||||
eprintln!(
|
||||
progress::report(
|
||||
show_progress,
|
||||
format!(
|
||||
"{REJECTED_EMOJI} Findings rejected — revising for round {}...",
|
||||
round + 1
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
feedback = Some(review);
|
||||
}
|
||||
|
||||
write_report(&client, topic, &findings, incomplete, show_progress).await
|
||||
}
|
||||
|
||||
/// Embeds any uploaded documents and retrieves the excerpts most relevant to
|
||||
/// `topic`, once, up front — nothing later in the loop changes which
|
||||
/// excerpts are relevant, so there's no reason to repeat this per round.
|
||||
/// `None` when no documents were given, or none of them yielded a usable
|
||||
/// excerpt (an empty `docs` list is deliberately not an error).
|
||||
async fn build_doc_context(
|
||||
client: &ollama::Client,
|
||||
topic: &str,
|
||||
docs: &[PathBuf],
|
||||
show_progress: bool,
|
||||
) -> anyhow::Result<Option<String>> {
|
||||
if docs.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let documents = documents::collect_documents(docs)?;
|
||||
if documents.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let index = retrieval::build_index(client, documents, show_progress).await?;
|
||||
let excerpts = retrieval::retrieve_relevant(&index, topic, show_progress).await?;
|
||||
|
||||
Ok((!excerpts.is_empty()).then_some(excerpts))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,11 +19,12 @@ pub(crate) struct GatheredFindings {
|
|||
/// Wraps the tool-calling research loop in its own span so it's visible as a
|
||||
/// single unit in traces, distinct from the writing and review phases and
|
||||
/// nesting rig's own per-turn `chat`/`execute_tool` spans underneath it.
|
||||
#[tracing::instrument(skip(client, feedback), fields(gen_ai.agent.name = "researcher"))]
|
||||
#[tracing::instrument(skip(client, feedback, doc_context), fields(gen_ai.agent.name = "researcher"))]
|
||||
pub(crate) async fn gather_findings(
|
||||
client: &ollama::Client,
|
||||
topic: &str,
|
||||
feedback: Option<&Review>,
|
||||
doc_context: Option<&str>,
|
||||
round: usize,
|
||||
show_progress: bool,
|
||||
) -> anyhow::Result<GatheredFindings> {
|
||||
|
|
@ -68,6 +69,17 @@ pub(crate) async fn gather_findings(
|
|||
),
|
||||
};
|
||||
|
||||
let task = match doc_context {
|
||||
Some(context) => format!(
|
||||
"{task}\n\n\
|
||||
Relevant excerpts from documents the user uploaded — treat these as trusted primary \
|
||||
sources alongside anything you find on the web, and cite them with the same \
|
||||
footnote scheme (their Sources entry can just be the document path shown below):\n\
|
||||
{context}"
|
||||
),
|
||||
None => task,
|
||||
};
|
||||
|
||||
let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching..."));
|
||||
let run_result = researcher
|
||||
.runner(task)
|
||||
|
|
|
|||
81
src/retrieval.rs
Normal file
81
src/retrieval.rs
Normal 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"))
|
||||
}
|
||||
|
|
@ -1,4 +1,5 @@
|
|||
use rig::client::AgentClientExt;
|
||||
use rig::completion::TypedPrompt;
|
||||
use rig::providers::ollama;
|
||||
use rig::schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
|
@ -25,10 +26,18 @@ pub(crate) struct Review {
|
|||
pub(crate) gaps: String,
|
||||
}
|
||||
|
||||
/// Uses rig's typed extractor — a forced tool call into a `submit(Review)`
|
||||
/// schema — rather than parsing free-text output, so the verdict and its
|
||||
/// two feedback fields always come back structured instead of relying on
|
||||
/// scanning prose for a trailing yes/no.
|
||||
/// Retries for a review that fails to come back as valid structured output
|
||||
/// (e.g. a transient network error), on top of the initial attempt.
|
||||
const REVIEW_RETRIES: usize = 2;
|
||||
|
||||
/// Uses rig's typed-prompt API — `Native` structured output constraining the
|
||||
/// model's own reply to the `Review` schema — rather than a forced tool call
|
||||
/// or parsing free-text output. A forced tool call needs `tool_choice`
|
||||
/// support, which Ollama doesn't have: the reviewer's small local model would
|
||||
/// just answer in prose and never call the tool, so extraction only ever
|
||||
/// exhausted its retries and errored out. `Native` mode instead uses Ollama's
|
||||
/// own `format` JSON-schema constraint, which is honored regardless of
|
||||
/// tool-calling ability.
|
||||
#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "reviewer"))]
|
||||
pub(crate) async fn review_findings(
|
||||
client: &ollama::Client,
|
||||
|
|
@ -37,7 +46,7 @@ pub(crate) async fn review_findings(
|
|||
show_progress: bool,
|
||||
) -> anyhow::Result<Review> {
|
||||
let reviewer = client
|
||||
.extractor::<Review>(REVIEWER_MODEL)
|
||||
.agent(REVIEWER_MODEL)
|
||||
.preamble(
|
||||
"You are a skeptical fact-checker reviewing another researcher's notes before they \
|
||||
get turned into a report. Approve only if every conclusion in the findings is \
|
||||
|
|
@ -47,20 +56,36 @@ pub(crate) async fn review_findings(
|
|||
Always separate the solid, well-supported findings from the gaps so a follow-up \
|
||||
research pass knows what to keep and what to dig into further.",
|
||||
)
|
||||
.retries(2)
|
||||
.build();
|
||||
|
||||
let spinner = Spinner::start(
|
||||
show_progress,
|
||||
format!("{REVIEW_EMOJI} Reviewing findings..."),
|
||||
);
|
||||
let review = reviewer
|
||||
.extract(format!(
|
||||
"Topic: {topic}\n\nResearch findings to review:\n{findings}"
|
||||
))
|
||||
.await?;
|
||||
|
||||
let prompt = format!("Topic: {topic}\n\nResearch findings to review:\n{findings}");
|
||||
|
||||
let mut last_error = None;
|
||||
let mut extracted_review = None;
|
||||
|
||||
for attempt in 0..=REVIEW_RETRIES {
|
||||
match reviewer.prompt_typed(prompt.clone()).await {
|
||||
Ok(r) => {
|
||||
extracted_review = Some(r);
|
||||
break;
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!("Attempt {attempt} to extract JSON failed: {e:?}.");
|
||||
last_error = Some(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
drop(spinner);
|
||||
|
||||
let review: Review = extracted_review
|
||||
.ok_or_else(|| last_error.expect("loop always sets last_error on failure"))?;
|
||||
|
||||
tracing::info!(approved = review.approved, gaps = %review.gaps, "review complete");
|
||||
|
||||
Ok(review)
|
||||
|
|
|
|||
|
|
@ -18,14 +18,11 @@ use std::io::Write;
|
|||
/// covering the gap between the prompt being sent and generation starting
|
||||
/// (otherwise the terminal would go blank for however long that takes)
|
||||
/// rather than being dropped by the caller before this is even called.
|
||||
pub(crate) async fn write_text_stream<R>(
|
||||
mut stream: impl Stream<Item = Result<MultiTurnStreamItem<R>, StreamingError>> + Unpin,
|
||||
pub(crate) async fn write_text_stream(
|
||||
mut stream: impl Stream<Item = Result<MultiTurnStreamItem, StreamingError>> + Unpin,
|
||||
writer: &mut impl Write,
|
||||
mut spinner: Spinner,
|
||||
) -> anyhow::Result<String>
|
||||
where
|
||||
R: Clone,
|
||||
{
|
||||
) -> anyhow::Result<String> {
|
||||
let mut text = String::new();
|
||||
|
||||
while let Some(chunk) = stream.next().await {
|
||||
|
|
@ -50,10 +47,7 @@ mod tests {
|
|||
use rig::completion::CompletionError;
|
||||
use rig::message::Text;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct DummyResponse;
|
||||
|
||||
fn text_item(text: &str) -> Result<MultiTurnStreamItem<DummyResponse>, StreamingError> {
|
||||
fn text_item(text: &str) -> Result<MultiTurnStreamItem, StreamingError> {
|
||||
Ok(MultiTurnStreamItem::StreamAssistantItem(
|
||||
StreamedAssistantContent::Text(Text::new(text)),
|
||||
))
|
||||
|
|
@ -76,7 +70,7 @@ mod tests {
|
|||
#[tokio::test]
|
||||
async fn ignores_non_text_items() {
|
||||
let final_item = Ok(MultiTurnStreamItem::final_response(
|
||||
rig::OneOrMany::one(rig::message::AssistantContent::text("ignored")),
|
||||
vec![rig::message::AssistantContent::text("ignored")],
|
||||
rig::completion::Usage::new(),
|
||||
));
|
||||
let items = vec![text_item("kept"), final_item];
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ pub(crate) async fn search_web(
|
|||
) -> Result<String, ToolExecutionError> {
|
||||
progress::set_activity(format!("{SEARCH_EMOJI} Searching: {query}"));
|
||||
|
||||
// TODO re-use same client object?
|
||||
let response = reqwest::Client::new()
|
||||
.get(format!("{}/search", searxng_base_url()))
|
||||
.query(&[("q", query.as_str()), ("format", "json")])
|
||||
|
|
|
|||
Loading…
Reference in a new issue