From 19c92ccd52aef9adc0060ec3d3a3af58471686d7 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Wed, 19 Aug 2026 15:13:07 +0200 Subject: [PATCH] refactor: cleanup pass on document embedding/retrieval - retrieval::retrieve_relevant now shows a spinner around the query- embedding call and gets its own tracing span, matching every other model-calling phase in the pipeline (it previously ran invisibly and untraced) - add progress::report() to dedupe the show_progress-gated eprintln! pattern shared by research.rs and retrieval.rs - documents::push_document builds its TextSplitter once per collect_documents call instead of once per file, skips the char-truncation walk entirely when a file is already under the limit, and pulls the source-string branch out of the map closure - researcher::gather_findings drops a doc_context emptiness check that build_doc_context already guarantees - documents.rs tests use tempfile::tempdir() instead of a hand-rolled TempDir type Verified with cargo build/test/clippy -D warnings/fmt --check. --- Cargo.lock | 1 + Cargo.toml | 1 + src/documents.rs | 70 +++++++++++++++++------------------------------ src/history.rs | 6 +++- src/progress.rs | 9 ++++++ src/research.rs | 11 ++++---- src/researcher.rs | 4 +-- src/retrieval.rs | 17 ++++++++---- 8 files changed, 61 insertions(+), 58 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 80d92de..64bf95e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2129,6 +2129,7 @@ dependencies = [ "scraper", "serde", "serde_json", + "tempfile", "text-splitter", "tokio", "tracing", diff --git a/Cargo.toml b/Cargo.toml index e13d433..d42ec1b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,3 +27,4 @@ tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } [dev-dependencies] serde_json = "1" +tempfile = "3" diff --git a/src/documents.rs b/src/documents.rs index c7a80a5..e3fe5f5 100644 --- a/src/documents.rs +++ b/src/documents.rs @@ -1,6 +1,6 @@ use rig::embeddings::{EmbedError, TextEmbedder}; use std::path::{Path, PathBuf}; -use text_splitter::TextSplitter; +use text_splitter::{Characters, TextSplitter}; /// Safety valve against accidentally pointing `--doc` at a huge binary or /// log file — not a content limit. Real documents are chunked in full (see @@ -40,6 +40,7 @@ impl rig::Embed for Document { /// (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 splitter = TextSplitter::new(CHUNK_CHARS); let mut documents = Vec::new(); for path in paths { @@ -47,30 +48,38 @@ pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result) { +fn push_document(path: &Path, splitter: &TextSplitter, documents: &mut Vec) { match std::fs::read_to_string(path) { - Ok(text) => { - let text: String = text.chars().take(MAX_FILE_CHARS).collect(); - let chunks: Vec<&str> = TextSplitter::new(CHUNK_CHARS).chunks(&text).collect(); + Ok(mut text) => { + // Cheap upper bound on char count, so the vast majority of + // documents (well under the limit) skip the char-by-char walk + // entirely. + if text.len() > MAX_FILE_CHARS && text.chars().count() > MAX_FILE_CHARS { + text = text.chars().take(MAX_FILE_CHARS).collect(); + } + let chunks: Vec<&str> = splitter.chunks(&text).collect(); let total = chunks.len(); - - documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { - source: if total > 1 { + let source = |i: usize| { + if total > 1 { format!("{} (part {}/{total})", path.display(), i + 1) } else { path.display().to_string() - }, + } + }; + + documents.extend(chunks.into_iter().enumerate().map(|(i, chunk)| Document { + source: source(i), text: chunk.to_string(), })); } @@ -83,40 +92,11 @@ fn push_document(path: &Path, documents: &mut Vec) { #[cfg(test)] mod tests { use super::*; - - /// A directory unique to this test run, under the OS temp dir, cleaned - /// up on drop so a panic mid-test can't leave stray files behind for the - /// next run. - struct TempDir(PathBuf); - - impl TempDir { - fn new(name: &str) -> Self { - let path = std::env::temp_dir().join(format!( - "doubleo7-test-{name}-{}-{}", - std::process::id(), - std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .expect("system clock is after the epoch") - .as_nanos() - )); - std::fs::create_dir_all(&path).expect("create temp dir"); - Self(path) - } - - fn path(&self) -> &Path { - &self.0 - } - } - - impl Drop for TempDir { - fn drop(&mut self) { - let _ = std::fs::remove_dir_all(&self.0); - } - } + use tempfile::tempdir; #[test] fn collect_documents_reads_an_explicit_file() { - let dir = TempDir::new("explicit-file"); + let dir = tempdir().unwrap(); let file = dir.path().join("notes.txt"); std::fs::write(&file, "hello from a file").unwrap(); @@ -129,7 +109,7 @@ mod tests { #[test] fn collect_documents_reads_every_file_in_a_directory_non_recursively() { - let dir = TempDir::new("directory"); + 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(); @@ -150,7 +130,7 @@ mod tests { #[test] fn collect_documents_skips_unreadable_paths_instead_of_failing() { - let dir = TempDir::new("missing"); + let dir = tempdir().unwrap(); let missing = dir.path().join("does-not-exist.txt"); assert!(collect_documents(&[missing]).unwrap().is_empty()); @@ -158,7 +138,7 @@ mod tests { #[test] fn collect_documents_splits_a_large_file_into_multiple_chunks() { - let dir = TempDir::new("large-file"); + 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. diff --git a/src/history.rs b/src/history.rs index d506fd6..09b6fa7 100644 --- a/src/history.rs +++ b/src/history.rs @@ -270,7 +270,11 @@ 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"), ]; diff --git a/src/progress.rs b/src/progress.rs index faef9e3..013a3cc 100644 --- a/src/progress.rs +++ b/src/progress.rs @@ -69,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) { + 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 diff --git a/src/research.rs b/src/research.rs index 98899a5..e07a610 100644 --- a/src/research.rs +++ b/src/research.rs @@ -1,5 +1,5 @@ use crate::documents; -use crate::progress::REJECTED_EMOJI; +use crate::progress::{self, REJECTED_EMOJI}; use crate::researcher::gather_findings; use crate::retrieval; use crate::review::{self, Review}; @@ -57,12 +57,13 @@ pub(crate) async fn research( break; } - if show_progress { - eprintln!( + progress::report( + show_progress, + format!( "{REJECTED_EMOJI} Findings rejected — revising for round {}...", round + 1 - ); - } + ), + ); feedback = Some(review); } diff --git a/src/researcher.rs b/src/researcher.rs index 95f9218..4c3c72c 100644 --- a/src/researcher.rs +++ b/src/researcher.rs @@ -70,14 +70,14 @@ pub(crate) async fn gather_findings( }; let task = match doc_context { - Some(context) if !context.is_empty() => format!( + 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}" ), - _ => task, + None => task, }; let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); diff --git a/src/retrieval.rs b/src/retrieval.rs index 684858d..c8bd396 100644 --- a/src/retrieval.rs +++ b/src/retrieval.rs @@ -1,6 +1,6 @@ use crate::documents::Document; use crate::models::EMBEDDING_MODEL; -use crate::progress::{EMBED_EMOJI, Spinner}; +use crate::progress::{self, EMBED_EMOJI, Spinner}; use rig::client::EmbeddingsClient; use rig::providers::ollama; use rig::vector_store::VectorStoreIndex; @@ -44,6 +44,7 @@ pub(crate) async fn build_index( /// 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, @@ -54,14 +55,20 @@ pub(crate) async fn retrieve_relevant( .samples(TOP_N_EXCERPTS) .build(); + let spinner = Spinner::start( + show_progress, + format!("{EMBED_EMOJI} Retrieving relevant excerpts..."), + ); let results = index.top_n::(request).await?; + drop(spinner); - if show_progress { - eprintln!( + 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