From ff9f9455a65846f0f223877f39db32e9d240053a Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 14:43:14 +0200 Subject: [PATCH] Chunk uploaded documents instead of truncating them at 20K chars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A flat 20K-char cutoff (copied from the fetched-web-page limit) silently dropped everything past the first ~20KB of a larger file, and even within the cutoff, embedding a whole multi-page document as one vector made retrieval coarse — the vector just averages out whatever topics the document covers. Split each file into ~1500-char chunks via text-splitter (recursive semantic-boundary splitting: paragraph > sentence > word, never mid-word) and embed each chunk as its own document, tagged with its source and part number. This removes the practical size ceiling — a large file chunks the same way a short one does — and sharpens retrieval by letting it surface the specific passage relevant to a query. It also incidentally caps the worst-case retrieval payload: 5 chunks now tops out around 7500 chars versus the old worst case of 5 full 20K-char documents. Verified against a live Ollama nomic-embed-text pull: a short document still embeds as a single chunk, unchanged from before. --- Cargo.lock | 101 +++++++++++++++++++++++++++++++++++++++++++++++ Cargo.toml | 1 + src/documents.rs | 77 +++++++++++++++++++++++++++++------- 3 files changed, 164 insertions(+), 15 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 12247c2..8c28f42 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -517,6 +517,18 @@ version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" +[[package]] +name = "auto_enums" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3091d68264354f211516b91dce6f71046e444fab1867716035f736667243affb" +dependencies = [ + "derive_utils", + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "autocfg" version = "1.5.1" @@ -1068,6 +1080,15 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "core_maths" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77745e017f5edba1a9c1d854f6f3a52dac8a12dd5af5d2f54aecf61e43d80d30" +dependencies = [ + "libm", +] + [[package]] name = "cpufeatures" version = "0.2.17" @@ -1942,6 +1963,7 @@ dependencies = [ "scraper", "serde", "serde_json", + "text-splitter", "tokio", "tracing", "tracing-subscriber", @@ -2059,6 +2081,17 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "derive_utils" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc05a5d33db20c784f873e84934ad94bb209a090987ac5f62fede2c178234f23" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + [[package]] name = "digest" version = "0.10.7" @@ -2874,11 +2907,32 @@ checksum = "d56e28588da92eee5c3201a6eff33fabdd49b62269c8938d4ff050ce4d900deb" dependencies = [ "displaydoc", "litemap", + "serde", "tinystr", "writeable", "zerovec", ] +[[package]] +name = "icu_locale_fallback" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "251af8e57c9400e3eb58242fe5b8b1152b2a64fdf4cf632f923c38ccee6f2fa9" +dependencies = [ + "icu_locale_core", + "icu_locale_fallback_data", + "icu_provider", + "potential_utf", + "tinystr", + "zerovec", +] + +[[package]] +name = "icu_locale_fallback_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "decf2a22ec8fa68f1a0c1129a3f8583f8f8bc24e8b9ccbe98ead99f62a4dc3a8" + [[package]] name = "icu_normalizer" version = "2.3.0" @@ -2928,6 +2982,8 @@ checksum = "92a7ed671a6aad807a8651a2e1782a6598fda9ce5185dd8158549e95a91c6428" dependencies = [ "displaydoc", "icu_locale_core", + "serde", + "stable_deref_trait", "writeable", "yoke", "zerofrom", @@ -2935,6 +2991,29 @@ dependencies = [ "zerovec", ] +[[package]] +name = "icu_segmenter" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82d07aafccd67af15d02512a6adf5896fbc5ed00f2e99b471d2efa14016db3db" +dependencies = [ + "core_maths", + "icu_collections", + "icu_locale_fallback", + "icu_provider", + "icu_segmenter_data", + "potential_utf", + "smallvec", + "utf8_iter", + "zerovec", +] + +[[package]] +name = "icu_segmenter_data" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae293c039020f9ec10710af98d29ce6aa2051486638b49c9a6409f3b4a9e98ad" + [[package]] name = "ident_case" version = "1.0.1" @@ -4763,6 +4842,8 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d83eb9bc6d8e5cf568e7a1101d60ee05e81ed50ea106026f3d18deeb046d7661" dependencies = [ + "serde_core", + "writeable", "zerovec", ] @@ -6409,6 +6490,23 @@ dependencies = [ "new_debug_unreachable", ] +[[package]] +name = "text-splitter" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f3eeec76988617ff1434d754d7e8e197be2cb8981ea77c17e32f3f7a2c6f95e" +dependencies = [ + "ahash", + "auto_enums", + "either", + "icu_provider", + "icu_segmenter", + "itertools 0.14.0", + "memchr", + "strum 0.28.0", + "thiserror", +] + [[package]] name = "thiserror" version = "2.0.20" @@ -6507,6 +6605,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b1e27c91459209c2986af3dcf603a5a74a4368754ce37414f59acc971167f643" dependencies = [ "displaydoc", + "serde_core", "zerovec", ] @@ -7560,6 +7659,7 @@ dependencies = [ "displaydoc", "yoke", "zerofrom", + "zerovec", ] [[package]] @@ -7568,6 +7668,7 @@ version = "0.11.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94b5c6b5976d66c1d703c4fd17d3f5e43c8cedaacf604961b171adc7130896d8" dependencies = [ + "serde", "yoke", "zerofrom", "zerovec-derive", diff --git a/Cargo.toml b/Cargo.toml index e6e0ec6..45bba54 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ rig = "0.41.0" schemars = "1" scraper = "0.27" serde = { version = "1.0.229", features = ["derive"] } +text-splitter = "0.32" tokio = { version = "1.53.1", features = ["full"] } tracing = "0.1" tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } diff --git a/src/documents.rs b/src/documents.rs index 8a11c2c..f0ba97d 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -1,13 +1,25 @@ use rig::embeddings::{EmbedError, TextEmbedder}; use std::path::{Path, PathBuf}; +use text_splitter::TextSplitter; -/// Mirrors `tools::MAX_PAGE_CHARS` — one huge uploaded file shouldn't blow -/// out the embedding model's input any more than one huge fetched page -/// should blow out the researcher's context. -const MAX_DOCUMENT_CHARS: usize = 20_000; +/// Safety valve against accidentally pointing `--doc` at a huge binary or +/// log file — not a content limit. Real documents are chunked in full (see +/// `CHUNK_CHARS` below), so nothing meaningful gets silently dropped short +/// of this. +const MAX_FILE_CHARS: usize = 2_000_000; -/// A single uploaded document, embedded via its full text so the researcher -/// can later retrieve semantically relevant excerpts by topic. +/// Target chunk size handed to the embedding model: small enough that a +/// handful of retrieved chunks stays well within a local model's context +/// window, large enough to keep a paragraph or two of context in each one. +/// `TextSplitter` treats this as an upper bound, not a fixed size — it +/// recursively splits on the largest semantic boundary (paragraph, +/// sentence, word, ...) that still fits, so a chunk never cuts a sentence +/// mid-word just to hit the target exactly. +const CHUNK_CHARS: usize = 1_500; + +/// A single chunk of an uploaded document, embedded via its own text so +/// retrieval can surface just the passage relevant to a query rather than +/// an entire (possibly very long) file at once. #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub(crate) struct Document { pub(crate) source: String, @@ -21,11 +33,12 @@ impl rig::Embed for Document { } } -/// Resolves CLI-provided paths into documents to embed: a file is read -/// directly, a directory contributes every non-directory entry inside it — -/// one level deep, not recursive, so a stray nested folder can't silently -/// pull in unrelated files. Unreadable entries (permissions, non-UTF-8) are -/// skipped with a warning rather than failing the whole run. +/// Resolves CLI-provided paths into chunked documents to embed: a file is +/// read and split into chunks, a directory contributes every non-directory +/// entry inside it — one level deep, not recursive, so a stray nested +/// folder can't silently pull in unrelated files. Unreadable entries +/// (permissions, non-UTF-8) are skipped with a warning rather than failing +/// the whole run. pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result> { let mut documents = Vec::new(); @@ -47,10 +60,20 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result) { match std::fs::read_to_string(path) { - Ok(text) => documents.push(Document { - source: path.display().to_string(), - text: text.chars().take(MAX_DOCUMENT_CHARS).collect(), - }), + Ok(text) => { + let text: String = text.chars().take(MAX_FILE_CHARS).collect(); + let chunks: Vec<&str> = TextSplitter::new(CHUNK_CHARS).chunks(&text).collect(); + let total = chunks.len(); + + documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { + source: if total > 1 { + format!("{} (part {}/{total})", path.display(), i + 1) + } else { + path.display().to_string() + }, + text: chunk.to_string(), + })); + } Err(err) => { tracing::warn!(path = %path.display(), %err, "skipping unreadable document") } @@ -133,6 +156,30 @@ mod tests { assert!(collect_documents(&[missing]).unwrap().is_empty()); } + #[test] + fn collect_documents_splits_a_large_file_into_multiple_chunks() { + let dir = TempDir::new("large-file"); + let file = dir.path().join("big.txt"); + // Well over CHUNK_CHARS, and with paragraph breaks so the splitter + // has real semantic boundaries to chunk on. + let paragraph = "word ".repeat(100); + std::fs::write(&file, vec![paragraph; 10].join("\n\n")).unwrap(); + + let documents = collect_documents(std::slice::from_ref(&file)).unwrap(); + + assert!( + documents.len() > 1, + "expected a large file to produce multiple chunks" + ); + for (i, doc) in documents.iter().enumerate() { + assert_eq!( + doc.source, + format!("{} (part {}/{})", file.display(), i + 1, documents.len()) + ); + assert!(!doc.text.is_empty()); + } + } + #[test] fn document_embed_yields_its_full_text() { let document = Document {