Add document embedding and retrieval via a dedicated embedding model #10

Merged
schaefera merged 10 commits from worktree-deep-research-max-turns-report into master 2026-08-19 15:37:33 +00:00
3 changed files with 164 additions and 15 deletions
Showing only changes of commit ff9f9455a6 - Show all commits

101
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"
@ -1942,6 +1963,7 @@ dependencies = [
"scraper", "scraper",
"serde", "serde",
"serde_json", "serde_json",
"text-splitter",
"tokio", "tokio",
"tracing", "tracing",
"tracing-subscriber", "tracing-subscriber",
@ -2059,6 +2081,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"
@ -2874,11 +2907,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 +2982,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 +2991,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 +4842,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 +6490,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 +6605,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643"
dependencies = [ dependencies = [
"displaydoc", "displaydoc",
"serde_core",
"zerovec", "zerovec",
] ]
@ -7560,6 +7659,7 @@ dependencies = [
"displaydoc", "displaydoc",
"yoke", "yoke",
"zerofrom", "zerofrom",
"zerovec",
] ]
[[package]] [[package]]
@ -7568,6 +7668,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,6 +20,7 @@ rig = "0.41.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"] }

View file

@ -1,13 +1,25 @@
use rig::embeddings::{EmbedError, TextEmbedder}; use rig::embeddings::{EmbedError, TextEmbedder};
use std::path::{Path, PathBuf}; use std::path::{Path, PathBuf};
use text_splitter::TextSplitter;
/// Mirrors `tools::MAX_PAGE_CHARS` — one huge uploaded file shouldn't blow /// Safety valve against accidentally pointing `--doc` at a huge binary or
/// out the embedding model's input any more than one huge fetched page /// log file — not a content limit. Real documents are chunked in full (see
/// should blow out the researcher's context. /// `CHUNK_CHARS` below), so nothing meaningful gets silently dropped short
const MAX_DOCUMENT_CHARS: usize = 20_000; /// of this.
const MAX_FILE_CHARS: usize = 2_000_000;
/// A single uploaded document, embedded via its full text so the researcher /// Target chunk size handed to the embedding model: small enough that a
/// can later retrieve semantically relevant excerpts by topic. /// 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)] #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub(crate) struct Document { pub(crate) struct Document {
pub(crate) source: String, pub(crate) source: String,
@ -21,11 +33,12 @@ impl rig::Embed for Document {
} }
} }
/// Resolves CLI-provided paths into documents to embed: a file is read /// Resolves CLI-provided paths into chunked documents to embed: a file is
/// directly, a directory contributes every non-directory entry inside it — /// read and split into chunks, a directory contributes every non-directory
/// one level deep, not recursive, so a stray nested folder can't silently /// entry inside it — one level deep, not recursive, so a stray nested
/// pull in unrelated files. Unreadable entries (permissions, non-UTF-8) are /// folder can't silently pull in unrelated files. Unreadable entries
/// skipped with a warning rather than failing the whole run. /// (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>> { pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Document>> {
let mut documents = Vec::new(); let mut documents = Vec::new();
@ -47,10 +60,20 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Documen
fn push_document(path: &Path, documents: &mut Vec<Document>) { fn push_document(path: &Path, documents: &mut Vec<Document>) {
match std::fs::read_to_string(path) { match std::fs::read_to_string(path) {
Ok(text) => documents.push(Document { Ok(text) => {
source: path.display().to_string(), let text: String = text.chars().take(MAX_FILE_CHARS).collect();
text: text.chars().take(MAX_DOCUMENT_CHARS).collect(), let chunks: Vec<&str> = TextSplitter::new(CHUNK_CHARS).chunks(&text).collect();
}), let total = chunks.len();
documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document {
source: if total > 1 {
format!("{} (part {}/{total})", path.display(), i + 1)
} else {
path.display().to_string()
},
text: chunk.to_string(),
}));
}
Err(err) => { Err(err) => {
tracing::warn!(path = %path.display(), %err, "skipping unreadable document") tracing::warn!(path = %path.display(), %err, "skipping unreadable document")
} }
@ -133,6 +156,30 @@ mod tests {
assert!(collect_documents(&[missing]).unwrap().is_empty()); assert!(collect_documents(&[missing]).unwrap().is_empty());
} }
#[test]
fn collect_documents_splits_a_large_file_into_multiple_chunks() {
let dir = TempDir::new("large-file");
let file = dir.path().join("big.txt");
// Well over CHUNK_CHARS, and with paragraph breaks so the splitter
// has real semantic boundaries to chunk on.
let paragraph = "word ".repeat(100);
std::fs::write(&file, vec![paragraph; 10].join("\n\n")).unwrap();
let documents = collect_documents(std::slice::from_ref(&file)).unwrap();
assert!(
documents.len() > 1,
"expected a large file to produce multiple chunks"
);
for (i, doc) in documents.iter().enumerate() {
assert_eq!(
doc.source,
format!("{} (part {}/{})", file.display(), i + 1, documents.len())
);
assert!(!doc.text.is_empty());
}
}
#[test] #[test]
fn document_embed_yields_its_full_text() { fn document_embed_yields_its_full_text() {
let document = Document { let document = Document {