Add document embedding and retrieval via a dedicated embedding model #10
9 changed files with 311 additions and 5 deletions
13
README.md
13
README.md
|
|
@ -79,8 +79,19 @@ 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 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.
|
||||
|
||||
## Project layout
|
||||
|
||||
Split one concern per file rather than one large module:
|
||||
|
|
@ -94,6 +105,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 Cli {
|
|||
/// 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>,
|
||||
}
|
||||
|
|
|
|||
148
src/documents.rs
Normal file
148
src/documents.rs
Normal file
|
|
@ -0,0 +1,148 @@
|
|||
use rig::embeddings::{EmbedError, TextEmbedder};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// 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;
|
||||
|
||||
/// A single uploaded document, embedded via its full text so the researcher
|
||||
/// can later retrieve semantically relevant excerpts by topic.
|
||||
#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub(crate) struct Document {
|
||||
pub(crate) source: String,
|
||||
pub(crate) text: 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 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.
|
||||
pub(crate) fn collect_documents(paths: &[PathBuf]) -> anyhow::Result<Vec<Document>> {
|
||||
let mut documents = Vec::new();
|
||||
|
||||
for path in paths {
|
||||
if path.is_dir() {
|
||||
for entry in std::fs::read_dir(path)? {
|
||||
let entry = entry?;
|
||||
if entry.file_type()?.is_file() {
|
||||
push_document(&entry.path(), &mut documents);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
push_document(path, &mut documents);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(documents)
|
||||
}
|
||||
|
||||
fn push_document(path: &Path, documents: &mut Vec<Document>) {
|
||||
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(),
|
||||
}),
|
||||
Err(err) => {
|
||||
tracing::warn!(path = %path.display(), %err, "skipping unreadable document")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[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!(
|
||||
"deep_research-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);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn collect_documents_reads_an_explicit_file() {
|
||||
let dir = TempDir::new("explicit-file");
|
||||
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::new("directory");
|
||||
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::new("missing");
|
||||
let missing = dir.path().join("does-not-exist.txt");
|
||||
|
||||
assert!(collect_documents(&[missing]).unwrap().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()]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,9 +1,12 @@
|
|||
use crate::documents;
|
||||
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
|
||||
|
|
@ -14,16 +17,29 @@ 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;
|
||||
|
||||
|
|
@ -56,3 +72,29 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result
|
|||
|
||||
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) if !context.is_empty() => 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,
|
||||
};
|
||||
|
||||
let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching..."));
|
||||
let run_result = researcher
|
||||
.runner(task)
|
||||
|
|
|
|||
74
src/retrieval.rs
Normal file
74
src/retrieval.rs
Normal file
|
|
@ -0,0 +1,74 @@
|
|||
use crate::documents::Document;
|
||||
use crate::models::EMBEDDING_MODEL;
|
||||
use crate::progress::{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.
|
||||
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 results = index.top_n::<Document>(request).await?;
|
||||
|
||||
if show_progress {
|
||||
eprintln!(
|
||||
"{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