From 6cd2d6515893270883452dcb054aa9c3fa01aa60 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 12:59:03 +0200 Subject: [PATCH] Split core.rs into one file per concern MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit core.rs had grown into a 612-line grab-bag mixing six unrelated concerns: CLI arg parsing, logging setup, top-level orchestration, the researcher agent phase, chat-history reconstruction utilities, the summarizer agent phase, and the writer agent phase — while review.rs, tools.rs, stream.rs, and progress.rs already correctly isolated their own concerns. This splits core.rs to match that existing pattern instead of being the one file that doesn't follow it: - cli.rs — Cli struct + DEFAULT_TOPIC - observability.rs — initialize_observability - models.rs — RESEARCHER_MODEL / WRITER_MODEL (previously duplicated across call sites, now a single source of truth) - history.rs — pure chat-history parsing/reconstruction helpers (partial_findings_from_history, annotated_transcript_from_history, and their private helpers), plus their unit tests. Also dedupes MAX_TOOL_RESULT_CHARS, which was previously defined twice. - researcher.rs — gather_findings + GatheredFindings (the tool-calling research phase) - summarizer.rs — summarize_partial_history (the max-turns recovery agent) - writer.rs — write_report - research.rs — the top-level research() orchestration loop main.rs now only does argument parsing, logging setup, and the top-level call — no orchestration logic of its own. Unit tests stay co-located with the code they test per Rust convention (not pulled into separate files) rather than under "prefer new files" — that applies to production code organization here. No behavior changes; cargo test/clippy/fmt all clean. --- deep_research/src/cli.rs | 19 + deep_research/src/core.rs | 612 ----------------------------- deep_research/src/history.rs | 292 ++++++++++++++ deep_research/src/main.rs | 20 +- deep_research/src/models.rs | 8 + deep_research/src/observability.rs | 13 + deep_research/src/research.rs | 58 +++ deep_research/src/researcher.rs | 98 +++++ deep_research/src/summarizer.rs | 69 ++++ deep_research/src/writer.rs | 77 ++++ 10 files changed, 646 insertions(+), 620 deletions(-) create mode 100644 deep_research/src/cli.rs delete mode 100644 deep_research/src/core.rs create mode 100644 deep_research/src/history.rs create mode 100644 deep_research/src/models.rs create mode 100644 deep_research/src/observability.rs create mode 100644 deep_research/src/research.rs create mode 100644 deep_research/src/researcher.rs create mode 100644 deep_research/src/summarizer.rs create mode 100644 deep_research/src/writer.rs diff --git a/deep_research/src/cli.rs b/deep_research/src/cli.rs new file mode 100644 index 0000000..269a257 --- /dev/null +++ b/deep_research/src/cli.rs @@ -0,0 +1,19 @@ +use clap::Parser; + +pub(crate) const DEFAULT_TOPIC: &str = + "What are the latest advances in running large language models locally, on consumer hardware?"; + +/// Deep research agentic loop over local Gemma models: a tool-calling agent +/// gathers and cross-checks web evidence, a reviewer agent gates it, and a +/// writer agent turns approved findings into a structured report. +#[derive(Parser)] +#[command(name = "doubleo7-research", version, about)] +pub(crate) struct Cli { + /// Research topic to investigate + pub(crate) topic: Option, + + /// Emit logs at this level (off by default; passing this also enables a + /// progress spinner to switch off, since the logs already show progress) + #[arg(short = 'l', long, value_name = "LEVEL")] + pub(crate) log_level: Option, +} diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs deleted file mode 100644 index 3c14a91..0000000 --- a/deep_research/src/core.rs +++ /dev/null @@ -1,612 +0,0 @@ -use crate::progress::{REJECTED_EMOJI, REPORT_EMOJI, RESEARCH_EMOJI, SUMMARIZE_EMOJI, Spinner}; -use crate::review::{self, Review}; -use crate::stream::write_text_stream; -use crate::tools::{FetchPage, SearchWeb}; -use clap::Parser; -use rig::client::{AgentClientExt, Nothing}; -use rig::completion::message::{ToolResultContent, UserContent}; -use rig::completion::{AssistantContent, Message, Prompt, PromptError}; -use rig::providers::ollama; -use rig::streaming::StreamingPrompt; -use std::io::Write; - -/// The tool-calling research loop needs to reliably decide what to search -/// for, when a page is worth fetching, and when it has enough evidence — -/// that's a reasoning-heavy job best given to the largest local Gemma -/// variant. Turning the gathered notes into prose afterwards is comparatively -/// mechanical, so the smaller/faster variant handles that pass instead. -const RESEARCHER_MODEL: &str = "gemma4:26b"; -const WRITER_MODEL: &str = "gemma4-e4b:latest"; -const MAX_RESEARCH_TURNS: usize = 12; -/// Research/review rounds before giving up and writing the report from -/// whatever the last pass produced, rather than looping forever on a topic -/// the reviewer can never be satisfied with. -const MAX_RESEARCH_ROUNDS: usize = 3; - -/// Deep research agentic loop over local Gemma models: a tool-calling agent -/// gathers and cross-checks web evidence, a reviewer agent gates it, and a -/// writer agent turns approved findings into a structured report. -#[derive(Parser)] -#[command(name = "doubleo7-research", version, about)] -pub(crate) struct Cli { - /// Research topic to investigate - pub(crate) topic: Option, - - /// Emit logs at this level (off by default; passing this also enables a - /// progress spinner to switch off, since the logs already show progress) - #[arg(short = 'l', long, value_name = "LEVEL")] - pub(crate) log_level: Option, -} - -/// Only initializes a subscriber (and thus produces any log output at all) -/// when the caller opted in via `--log-level` — otherwise tracing's macros -/// are no-ops, leaving the terminal clean for the spinner. -pub(crate) fn initialize_observability(log_level: tracing::Level) { - tracing_subscriber::fmt() - .with_env_filter( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(log_level.to_string())), - ) - .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE) - .with_writer(std::io::stderr) - .init(); -} - -/// The least-agentic shape that fits: a plain Rust loop putting *this code*, -/// 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 { - let client = ollama::Client::new(Nothing)?; - - let mut findings = String::new(); - let mut feedback: Option = 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?; - 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. - 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"); - - if approved || round == MAX_RESEARCH_ROUNDS { - break; - } - - if show_progress { - eprintln!( - "{REJECTED_EMOJI} Findings rejected — revising for round {}...", - round + 1 - ); - } - - feedback = Some(review); - } - - write_report(&client, topic, &findings, incomplete, show_progress).await -} - -/// Findings gathered by a research pass, and whether the researcher was cut -/// off by the turn budget before it could conclude on its own. -struct GatheredFindings { - findings: String, - incomplete: bool, -} - -/// Best-effort reconstruction of research material from a chat history left -/// behind when the researcher hit its turn budget mid-investigation: whatever -/// prose the model wrote between tool calls, plus the raw text of every tool -/// result (search snippets, fetched pages), truncated so one huge page can't -/// crowd out everything else that was found. -fn partial_findings_from_history(chat_history: &[Message]) -> String { - const MAX_TOOL_RESULT_CHARS: usize = 2000; - - let sections: Vec = chat_history - .iter() - .flat_map(|message| { - extract_assistant_text(message) - .into_iter() - .chain(tool_result_text(message, MAX_TOOL_RESULT_CHARS)) - }) - .collect(); - - match sections.is_empty() { - true => "The researcher exhausted its turn budget before gathering any usable evidence." - .to_string(), - false => sections.join("\n"), - } -} - -/// Plain-text blocks from an assistant message, if any. -fn extract_assistant_text(message: &Message) -> Vec { - let Message::Assistant { content, .. } = message else { - return Vec::new(); - }; - content - .iter() - .filter_map(|item: &AssistantContent| match item { - AssistantContent::Text(text) => Some(text.text().to_string()), - _ => None, - }) - .collect() -} - -/// Plain-text tool results attached to a user message (that's where rig -/// places them), each truncated to `max_chars` so one huge fetched page -/// can't crowd out everything else that was found. -fn tool_result_text(message: &Message, max_chars: usize) -> Vec { - let Message::User { content } = message else { - return Vec::new(); - }; - - content - .iter() - .filter_map(|item: &UserContent| match item { - UserContent::ToolResult(tool_result) => Some(tool_result), - _ => None, - }) - .flat_map(|tool_result| tool_result.content.iter()) - .filter_map( - |tool_result_content: &ToolResultContent| match tool_result_content { - ToolResultContent::Text(text) => Some(truncate(text.text(), max_chars)), - _ => None, - }, - ) - .collect() -} - -fn truncate(text: &str, max_chars: usize) -> String { - let mut result = text.to_string(); - if result.len() > max_chars { - result.truncate(max_chars); - result.push_str(" ...[truncated]"); - } - result -} - -/// Chronological transcript of a partial research run, annotated with tool -/// calls (so a fetch's URL or a search's query stays attached to its -/// result) rather than just the bare result text `partial_findings_from_history` -/// collects — the summarizer agent below needs that context to attribute -/// facts to the right source. -fn annotated_transcript_from_history(chat_history: &[Message]) -> String { - chat_history - .iter() - .flat_map(transcript_lines) - .collect::>() - .join("\n") -} - -fn transcript_lines(message: &Message) -> Vec { - const MAX_TOOL_RESULT_CHARS: usize = 2000; - - match message { - Message::Assistant { content, .. } => content - .iter() - .filter_map(|item| match item { - AssistantContent::Text(text) => Some(format!("Note: {}", text.text())), - AssistantContent::ToolCall(call) => Some(format!( - "Called {}({})", - call.function.name, call.function.arguments - )), - _ => None, - }) - .collect(), - Message::User { content } => content - .iter() - .filter_map(|item| match item { - UserContent::ToolResult(result) => Some(result), - _ => None, - }) - .flat_map(|result| result.content.iter()) - .filter_map(|part| match part { - ToolResultContent::Text(text) => Some(format!( - "Result: {}", - truncate(text.text(), MAX_TOOL_RESULT_CHARS) - )), - _ => None, - }) - .collect(), - Message::System { .. } => Vec::new(), - } -} - -/// Agentic alternative to the plain programmatic extraction above: hands the -/// annotated transcript to a fresh model call and asks it to reconstruct the -/// same footnote-style findings dump the researcher would have written -/// itself, had it not run out of turns. This can dedupe repeated URLs and -/// restore correct `[n]` citation numbering in a way string concatenation -/// can't — but it's a model call like any other, so on an empty transcript -/// or a failure it falls back to `partial_findings_from_history` rather than -/// letting a second turn-budget problem take down the one recovery path -/// that's supposed to be bulletproof. -#[tracing::instrument(skip(client, chat_history), fields(gen_ai.agent.name = "history-summarizer"))] -async fn summarize_partial_history( - client: &ollama::Client, - topic: &str, - chat_history: &[Message], - show_progress: bool, -) -> String { - let transcript = annotated_transcript_from_history(chat_history); - - if transcript.is_empty() { - tracing::warn!("no usable transcript to summarize; skipping summarizer pass"); - return partial_findings_from_history(chat_history); - } - - let summarizer = client - .agent(WRITER_MODEL) - .name("history-summarizer") - .preamble( - "You are reconstructing research notes from a research session that was cut off \ - before the researcher could write its own summary. You'll be given a raw transcript \ - of tool calls (searches run, pages fetched), their results, and any interim comments \ - the researcher made. Turn this into a footnote-style findings dump: write each fact \ - the transcript actually supports, followed by a bracketed number like [1], then end \ - with a 'Sources' section mapping each number to its exact URL, one per line. Reuse \ - the same number for a URL that appears more than once. Do not invent facts beyond \ - what the transcript shows, and note explicitly where it looks thin or cuts off \ - mid-investigation.", - ) - .build(); - - let spinner = Spinner::start( - show_progress, - format!("{SUMMARIZE_EMOJI} Reconstructing partial findings..."), - ); - let result = summarizer - .prompt(format!( - "Topic: {topic}\n\nPartial research transcript:\n{transcript}" - )) - .await; - drop(spinner); - - match result { - Ok(findings) => { - tracing::info!(findings = %findings, "summarizer reconstructed partial findings"); - findings - } - Err(error) => { - tracing::warn!(%error, "summarizer agent failed; falling back to programmatic extraction"); - partial_findings_from_history(chat_history) - } - } -} - -/// 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"))] -async fn gather_findings( - client: &ollama::Client, - topic: &str, - feedback: Option<&Review>, - round: usize, - show_progress: bool, -) -> anyhow::Result { - let current_date = chrono::offset::Local::now().to_string(); - - let researcher = client - .agent(RESEARCHER_MODEL) - .name("researcher") - .preamble(format!( - "You are a meticulous research assistant. Use the search_web and fetch_page tools to \ - investigate the user's topic: run several searches with varied phrasing, fetch the \ - most promising pages, and cross-check claims across at least two sources before \ - trusting them. Ensure sources are up to date: the current date is {current_date}. \ - Once you are confident you have enough evidence, stop calling tools and reply with a \ - plain-text dump of every fact you gathered, using footnote-style citations: write each \ - fact followed by a bracketed number like [1], then at the end of your reply list a \ - 'Sources' section mapping each number to the exact URL it came from, one per line, e.g. \ - '[1] https://example.com/page'. Reuse the same number when multiple facts come from the \ - same URL — do not give one URL two different numbers. Also call out any open questions \ - or contradictions between sources, citing the footnotes involved. This is raw research \ - material for a writer, not a final report, so favor completeness over polish.") - .as_str(), - ) - .tool(SearchWeb) - .tool(FetchPage) - .build(); - - let task = match feedback { - None => topic.to_string(), - Some(review) => format!( - "Topic: {topic}\n\n\ - You already ran a research pass on this topic. A reviewer checked it against its \ - cited sources and found it insufficient. Do more research to address the reviewer's \ - feedback, then produce an updated findings dump: carry forward what's solid, and \ - add, correct, or better-source whatever the gaps call for. Gaps include out-of-date,\ - irrelevant, or clearly wrong information. The date is {current_date}.\n\n\ - Solid findings from the last pass — keep and build on these:\n{}\n\n\ - Gaps the reviewer found — conclusions not actually backed by their source, sources \ - that don't line up with the conclusion drawn from them, or parts of the topic still \ - uncovered:\n{}", - review.solid_findings, review.gaps - ), - }; - - let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); - let run_result = researcher - .runner(task) - .max_turns(MAX_RESEARCH_TURNS) - .run() - .await; - drop(spinner); - - let gathered = match run_result { - Ok(response) => GatheredFindings { - findings: response.output, - incomplete: false, - }, - Err(PromptError::MaxTurnsError { chat_history, .. }) => { - tracing::warn!(round, "researcher hit its turn budget before concluding"); - GatheredFindings { - findings: summarize_partial_history(client, topic, &chat_history, show_progress) - .await, - incomplete: true, - } - } - Err(err) => return Err(err.into()), - }; - - tracing::info!(round, findings = %gathered.findings, incomplete = gathered.incomplete, "research phase complete"); - - Ok(gathered) -} - -#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "writer"))] -async fn write_report( - client: &ollama::Client, - topic: &str, - findings: &str, - incomplete: bool, - show_progress: bool, -) -> anyhow::Result { - let writer = client - .agent(WRITER_MODEL) - .name("writer") - .preamble( - "You turn raw research notes into a clear, well-organized report for the reader. The \ - notes use footnote-style citations — a bracketed number like [1] after a fact, with a \ - Sources section mapping numbers to URLs. Preserve this scheme in your report: keep the \ - same [n] markers next to the claims they support (renumbering only if you drop unused \ - sources), and end the report with a 'Sources' section listing every footnote number \ - still in use next to its exact URL. Structure the body with headings, and call out any \ - open questions or contradictions the research turned up. Do not invent facts beyond \ - what the notes provide.", - ) - .build(); - - let task = if incomplete { - format!( - "Topic: {topic}\n\n\ - The researcher ran out of turns before it could finish investigating, so these notes \ - are an incomplete, best-effort reconstruction rather than the researcher's own \ - conclusions — treat any citations already in them as tentative. Write up what the \ - notes actually support, note explicitly where evidence is thin or missing, and open \ - the report with a short callout that this research is inconclusive and further \ - manual searching is needed.\n\n\ - Research notes:\n{findings}" - ) - } else { - format!("Topic: {topic}\n\nResearch notes:\n{findings}") - }; - - let spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report...")); - let response_stream = writer.stream_prompt(task).await; - - // Locked once for the whole stream rather than per chunk (as print! - // would do internally) — chunks arrive in a tight loop, so re-acquiring - // the lock on every one adds up. The spinner keeps running until the - // stream's first chunk arrives, so the terminal stays covered through - // the gap between sending the prompt and generation actually starting. - let stdout = std::io::stdout(); - let mut handle = stdout.lock(); - - // Guaranteed regardless of whether the writer model actually heeds the - // instruction above — the reader should never mistake a budget-exhausted - // run for a complete one just because the model forgot to say so. - let disclaimer = incomplete.then_some( - "> **Disclaimer:** incomplete research: the researcher exhausted its turn \ - budget before it could finish investigating. The findings below are partial and may be \ - incomplete or unbalanced; treat them as a starting point and verify further manually.\n\n", - ); - if let Some(disclaimer) = disclaimer { - write!(handle, "{disclaimer}")?; - } - - let report = write_text_stream(response_stream, &mut handle, spinner).await?; - writeln!(handle)?; - - Ok(match disclaimer { - Some(disclaimer) => format!("{disclaimer}{report}"), - None => report, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn truncate_leaves_short_text_untouched() { - assert_eq!(truncate("hello", 10), "hello"); - } - - #[test] - fn truncate_leaves_exact_length_text_untouched() { - assert_eq!(truncate("hello", 5), "hello"); - } - - #[test] - fn truncate_cuts_long_text_and_marks_it() { - assert_eq!(truncate("hello world", 5), "hello ...[truncated]"); - } - - #[test] - fn assistant_text_extracts_text_blocks() { - let message = Message::assistant("found it"); - assert_eq!( - extract_assistant_text(&message), - vec!["found it".to_string()] - ); - } - - #[test] - fn assistant_text_ignores_tool_calls() { - let message = Message::Assistant { - id: None, - content: rig::OneOrMany::one(AssistantContent::tool_call( - "call-1", - "search_web", - serde_json::json!({ "query": "test" }), - )), - }; - assert!(extract_assistant_text(&message).is_empty()); - } - - #[test] - fn assistant_text_ignores_non_assistant_messages() { - assert!(extract_assistant_text(&Message::user("hi")).is_empty()); - assert!(extract_assistant_text(&Message::system("be careful")).is_empty()); - } - - #[test] - fn tool_result_text_extracts_and_truncates() { - let short = Message::tool_result("call-1", "short result"); - assert_eq!( - tool_result_text(&short, 100), - vec!["short result".to_string()] - ); - - let long = Message::tool_result("call-2", "0123456789"); - assert_eq!( - tool_result_text(&long, 5), - vec!["01234 ...[truncated]".to_string()] - ); - } - - #[test] - fn tool_result_text_ignores_non_tool_result_content() { - assert!(tool_result_text(&Message::user("plain text, no tool result"), 100).is_empty()); - assert!(tool_result_text(&Message::assistant("also ignored"), 100).is_empty()); - } - - #[test] - fn partial_findings_from_history_falls_back_when_nothing_usable() { - let history = vec![ - Message::system("preamble"), - Message::Assistant { - id: None, - content: rig::OneOrMany::one(AssistantContent::tool_call( - "call-1", - "search_web", - serde_json::json!({ "query": "test" }), - )), - }, - ]; - - assert_eq!( - partial_findings_from_history(&history), - "The researcher exhausted its turn budget before gathering any usable evidence." - ); - } - - #[test] - 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::assistant("Cross-checking..."), - Message::tool_result("call-2", "Result B [2]"), - ]; - - assert_eq!( - partial_findings_from_history(&history), - "Checking sources...\nResult A [1]\nCross-checking...\nResult B [2]" - ); - } - - #[test] - fn transcript_lines_prefixes_assistant_notes() { - let message = Message::assistant("still checking this claim"); - assert_eq!( - transcript_lines(&message), - vec!["Note: still checking this claim".to_string()] - ); - } - - #[test] - fn transcript_lines_renders_tool_calls_with_their_arguments() { - let message = Message::Assistant { - id: None, - content: rig::OneOrMany::one(AssistantContent::tool_call( - "call-1", - "search_web", - serde_json::json!({ "query": "uruguay senior engineers" }), - )), - }; - assert_eq!( - transcript_lines(&message), - vec!["Called search_web({\"query\":\"uruguay senior engineers\"})".to_string()] - ); - } - - #[test] - fn transcript_lines_prefixes_and_truncates_tool_results() { - let short = Message::tool_result("call-1", "found via search [1]"); - assert_eq!( - transcript_lines(&short), - vec!["Result: found via search [1]".to_string()] - ); - } - - #[test] - fn transcript_lines_ignores_system_messages() { - assert!(transcript_lines(&Message::system("be thorough")).is_empty()); - } - - #[test] - fn annotated_transcript_from_history_joins_calls_results_and_notes_in_order() { - let history = vec![ - Message::Assistant { - id: None, - content: rig::OneOrMany::one(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::assistant("that source looks solid"), - ]; - - assert_eq!( - annotated_transcript_from_history(&history), - "Called search_web({\"query\":\"test\"})\n\ - Result: 1. Example\n https://example.com\n snippet\n\ - Note: that source looks solid" - ); - } - - #[test] - fn annotated_transcript_from_history_is_empty_with_no_usable_content() { - assert_eq!( - annotated_transcript_from_history(&[Message::system("preamble")]), - "" - ); - } -} diff --git a/deep_research/src/history.rs b/deep_research/src/history.rs new file mode 100644 index 0000000..5eb75f4 --- /dev/null +++ b/deep_research/src/history.rs @@ -0,0 +1,292 @@ +use rig::completion::message::{ToolResultContent, UserContent}; +use rig::completion::{AssistantContent, Message}; + +const MAX_TOOL_RESULT_CHARS: usize = 2000; + +/// Best-effort reconstruction of research material from a chat history left +/// behind when the researcher hit its turn budget mid-investigation: whatever +/// prose the model wrote between tool calls, plus the raw text of every tool +/// result (search snippets, fetched pages), truncated so one huge page can't +/// crowd out everything else that was found. +pub(crate) fn partial_findings_from_history(chat_history: &[Message]) -> String { + let sections: Vec = chat_history + .iter() + .flat_map(|message| { + extract_assistant_text(message) + .into_iter() + .chain(tool_result_text(message, MAX_TOOL_RESULT_CHARS)) + }) + .collect(); + + match sections.is_empty() { + true => "The researcher exhausted its turn budget before gathering any usable evidence." + .to_string(), + false => sections.join("\n"), + } +} + +/// Plain-text blocks from an assistant message, if any. +fn extract_assistant_text(message: &Message) -> Vec { + let Message::Assistant { content, .. } = message else { + return Vec::new(); + }; + content + .iter() + .filter_map(|item: &AssistantContent| match item { + AssistantContent::Text(text) => Some(text.text().to_string()), + _ => None, + }) + .collect() +} + +/// Plain-text tool results attached to a user message (that's where rig +/// places them), each truncated to `max_chars` so one huge fetched page +/// can't crowd out everything else that was found. +fn tool_result_text(message: &Message, max_chars: usize) -> Vec { + let Message::User { content } = message else { + return Vec::new(); + }; + + content + .iter() + .filter_map(|item: &UserContent| match item { + UserContent::ToolResult(tool_result) => Some(tool_result), + _ => None, + }) + .flat_map(|tool_result| tool_result.content.iter()) + .filter_map( + |tool_result_content: &ToolResultContent| match tool_result_content { + ToolResultContent::Text(text) => Some(truncate(text.text(), max_chars)), + _ => None, + }, + ) + .collect() +} + +fn truncate(text: &str, max_chars: usize) -> String { + let mut result = text.to_string(); + if result.len() > max_chars { + result.truncate(max_chars); + result.push_str(" ...[truncated]"); + } + result +} + +/// Chronological transcript of a partial research run, annotated with tool +/// calls (so a fetch's URL or a search's query stays attached to its +/// result) rather than just the bare result text `partial_findings_from_history` +/// collects — the summarizer agent needs that context to attribute facts to +/// the right source. +pub(crate) fn annotated_transcript_from_history(chat_history: &[Message]) -> String { + chat_history + .iter() + .flat_map(transcript_lines) + .collect::>() + .join("\n") +} + +fn transcript_lines(message: &Message) -> Vec { + match message { + Message::Assistant { content, .. } => content + .iter() + .filter_map(|item| match item { + AssistantContent::Text(text) => Some(format!("Note: {}", text.text())), + AssistantContent::ToolCall(call) => Some(format!( + "Called {}({})", + call.function.name, call.function.arguments + )), + _ => None, + }) + .collect(), + Message::User { content } => content + .iter() + .filter_map(|item| match item { + UserContent::ToolResult(result) => Some(result), + _ => None, + }) + .flat_map(|result| result.content.iter()) + .filter_map(|part| match part { + ToolResultContent::Text(text) => Some(format!( + "Result: {}", + truncate(text.text(), MAX_TOOL_RESULT_CHARS) + )), + _ => None, + }) + .collect(), + Message::System { .. } => Vec::new(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn truncate_leaves_short_text_untouched() { + assert_eq!(truncate("hello", 10), "hello"); + } + + #[test] + fn truncate_leaves_exact_length_text_untouched() { + assert_eq!(truncate("hello", 5), "hello"); + } + + #[test] + fn truncate_cuts_long_text_and_marks_it() { + assert_eq!(truncate("hello world", 5), "hello ...[truncated]"); + } + + #[test] + fn assistant_text_extracts_text_blocks() { + let message = Message::assistant("found it"); + assert_eq!( + extract_assistant_text(&message), + vec!["found it".to_string()] + ); + } + + #[test] + fn assistant_text_ignores_tool_calls() { + let message = Message::Assistant { + id: None, + content: rig::OneOrMany::one(AssistantContent::tool_call( + "call-1", + "search_web", + serde_json::json!({ "query": "test" }), + )), + }; + assert!(extract_assistant_text(&message).is_empty()); + } + + #[test] + fn assistant_text_ignores_non_assistant_messages() { + assert!(extract_assistant_text(&Message::user("hi")).is_empty()); + assert!(extract_assistant_text(&Message::system("be careful")).is_empty()); + } + + #[test] + fn tool_result_text_extracts_and_truncates() { + let short = Message::tool_result("call-1", "short result"); + assert_eq!( + tool_result_text(&short, 100), + vec!["short result".to_string()] + ); + + let long = Message::tool_result("call-2", "0123456789"); + assert_eq!( + tool_result_text(&long, 5), + vec!["01234 ...[truncated]".to_string()] + ); + } + + #[test] + fn tool_result_text_ignores_non_tool_result_content() { + assert!(tool_result_text(&Message::user("plain text, no tool result"), 100).is_empty()); + assert!(tool_result_text(&Message::assistant("also ignored"), 100).is_empty()); + } + + #[test] + fn partial_findings_from_history_falls_back_when_nothing_usable() { + let history = vec![ + Message::system("preamble"), + Message::Assistant { + id: None, + content: rig::OneOrMany::one(AssistantContent::tool_call( + "call-1", + "search_web", + serde_json::json!({ "query": "test" }), + )), + }, + ]; + + assert_eq!( + partial_findings_from_history(&history), + "The researcher exhausted its turn budget before gathering any usable evidence." + ); + } + + #[test] + 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::assistant("Cross-checking..."), + Message::tool_result("call-2", "Result B [2]"), + ]; + + assert_eq!( + partial_findings_from_history(&history), + "Checking sources...\nResult A [1]\nCross-checking...\nResult B [2]" + ); + } + + #[test] + fn transcript_lines_prefixes_assistant_notes() { + let message = Message::assistant("still checking this claim"); + assert_eq!( + transcript_lines(&message), + vec!["Note: still checking this claim".to_string()] + ); + } + + #[test] + fn transcript_lines_renders_tool_calls_with_their_arguments() { + let message = Message::Assistant { + id: None, + content: rig::OneOrMany::one(AssistantContent::tool_call( + "call-1", + "search_web", + serde_json::json!({ "query": "uruguay senior engineers" }), + )), + }; + assert_eq!( + transcript_lines(&message), + vec!["Called search_web({\"query\":\"uruguay senior engineers\"})".to_string()] + ); + } + + #[test] + fn transcript_lines_prefixes_and_truncates_tool_results() { + let short = Message::tool_result("call-1", "found via search [1]"); + assert_eq!( + transcript_lines(&short), + vec!["Result: found via search [1]".to_string()] + ); + } + + #[test] + fn transcript_lines_ignores_system_messages() { + assert!(transcript_lines(&Message::system("be thorough")).is_empty()); + } + + #[test] + fn annotated_transcript_from_history_joins_calls_results_and_notes_in_order() { + let history = vec![ + Message::Assistant { + id: None, + content: rig::OneOrMany::one(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::assistant("that source looks solid"), + ]; + + assert_eq!( + annotated_transcript_from_history(&history), + "Called search_web({\"query\":\"test\"})\n\ + Result: 1. Example\n https://example.com\n snippet\n\ + Note: that source looks solid" + ); + } + + #[test] + fn annotated_transcript_from_history_is_empty_with_no_usable_content() { + assert_eq!( + annotated_transcript_from_history(&[Message::system("preamble")]), + "" + ); + } +} diff --git a/deep_research/src/main.rs b/deep_research/src/main.rs index 087a844..1b71913 100644 --- a/deep_research/src/main.rs +++ b/deep_research/src/main.rs @@ -1,32 +1,36 @@ use clap::Parser; -mod core; +mod cli; +mod history; +mod models; +mod observability; mod progress; +mod research; +mod researcher; mod review; mod stream; +mod summarizer; mod tools; - -pub(crate) const DEFAULT_TOPIC: &str = - "What are the latest advances in running large language models locally, on consumer hardware?"; +mod writer; #[tokio::main] async fn main() -> anyhow::Result<()> { - let cli = core::Cli::parse(); + let cli = cli::Cli::parse(); let show_progress = match cli.log_level { Some(level) => { - core::initialize_observability(level); + observability::initialize_observability(level); false } None => true, }; - let topic = cli.topic.unwrap_or_else(|| DEFAULT_TOPIC.to_string()); + let topic = cli.topic.unwrap_or_else(|| cli::DEFAULT_TOPIC.to_string()); // 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. - core::research(&topic, show_progress).await?; + research::research(&topic, show_progress).await?; Ok(()) } diff --git a/deep_research/src/models.rs b/deep_research/src/models.rs new file mode 100644 index 0000000..b5f9120 --- /dev/null +++ b/deep_research/src/models.rs @@ -0,0 +1,8 @@ +/// The tool-calling research loop needs to reliably decide what to search +/// for, when a page is worth fetching, and when it has enough evidence — +/// that's a reasoning-heavy job best given to the largest local Gemma +/// variant. Turning gathered notes into prose (writing, summarizing) is +/// comparatively mechanical, so the smaller/faster variant handles those +/// passes instead. +pub(crate) const RESEARCHER_MODEL: &str = "gemma4:26b"; +pub(crate) const WRITER_MODEL: &str = "gemma4-e4b:latest"; diff --git a/deep_research/src/observability.rs b/deep_research/src/observability.rs new file mode 100644 index 0000000..4d54229 --- /dev/null +++ b/deep_research/src/observability.rs @@ -0,0 +1,13 @@ +/// Only initializes a subscriber (and thus produces any log output at all) +/// when the caller opted in via `--log-level` — otherwise tracing's macros +/// are no-ops, leaving the terminal clean for the spinner. +pub(crate) fn initialize_observability(log_level: tracing::Level) { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(log_level.to_string())), + ) + .with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE) + .with_writer(std::io::stderr) + .init(); +} diff --git a/deep_research/src/research.rs b/deep_research/src/research.rs new file mode 100644 index 0000000..2dc803e --- /dev/null +++ b/deep_research/src/research.rs @@ -0,0 +1,58 @@ +use crate::progress::REJECTED_EMOJI; +use crate::researcher::gather_findings; +use crate::review::{self, Review}; +use crate::writer::write_report; +use rig::client::Nothing; +use rig::providers::ollama; + +/// Research/review rounds before giving up and writing the report from +/// whatever the last pass produced, rather than looping forever on a topic +/// the reviewer can never be satisfied with. +const MAX_RESEARCH_ROUNDS: usize = 3; + +/// The least-agentic shape that fits: a plain Rust loop putting *this code*, +/// 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 { + let client = ollama::Client::new(Nothing)?; + + let mut findings = String::new(); + let mut feedback: Option = 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?; + 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. + 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"); + + if approved || round == MAX_RESEARCH_ROUNDS { + break; + } + + if show_progress { + eprintln!( + "{REJECTED_EMOJI} Findings rejected — revising for round {}...", + round + 1 + ); + } + + feedback = Some(review); + } + + write_report(&client, topic, &findings, incomplete, show_progress).await +} diff --git a/deep_research/src/researcher.rs b/deep_research/src/researcher.rs new file mode 100644 index 0000000..ad327a8 --- /dev/null +++ b/deep_research/src/researcher.rs @@ -0,0 +1,98 @@ +use crate::models::RESEARCHER_MODEL; +use crate::progress::{RESEARCH_EMOJI, Spinner}; +use crate::review::Review; +use crate::summarizer::summarize_partial_history; +use crate::tools::{FetchPage, SearchWeb}; +use rig::client::AgentClientExt; +use rig::completion::PromptError; +use rig::providers::ollama; + +const MAX_RESEARCH_TURNS: usize = 12; + +/// Findings gathered by a research pass, and whether the researcher was cut +/// off by the turn budget before it could conclude on its own. +pub(crate) struct GatheredFindings { + pub(crate) findings: String, + pub(crate) incomplete: bool, +} + +/// 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"))] +pub(crate) async fn gather_findings( + client: &ollama::Client, + topic: &str, + feedback: Option<&Review>, + round: usize, + show_progress: bool, +) -> anyhow::Result { + let current_date = chrono::offset::Local::now().to_string(); + + let researcher = client + .agent(RESEARCHER_MODEL) + .name("researcher") + .preamble(format!( + "You are a meticulous research assistant. Use the search_web and fetch_page tools to \ + investigate the user's topic: run several searches with varied phrasing, fetch the \ + most promising pages, and cross-check claims across at least two sources before \ + trusting them. Ensure sources are up to date: the current date is {current_date}. \ + Once you are confident you have enough evidence, stop calling tools and reply with a \ + plain-text dump of every fact you gathered, using footnote-style citations: write each \ + fact followed by a bracketed number like [1], then at the end of your reply list a \ + 'Sources' section mapping each number to the exact URL it came from, one per line, e.g. \ + '[1] https://example.com/page'. Reuse the same number when multiple facts come from the \ + same URL — do not give one URL two different numbers. Also call out any open questions \ + or contradictions between sources, citing the footnotes involved. This is raw research \ + material for a writer, not a final report, so favor completeness over polish.") + .as_str(), + ) + .tool(SearchWeb) + .tool(FetchPage) + .build(); + + let task = match feedback { + None => topic.to_string(), + Some(review) => format!( + "Topic: {topic}\n\n\ + You already ran a research pass on this topic. A reviewer checked it against its \ + cited sources and found it insufficient. Do more research to address the reviewer's \ + feedback, then produce an updated findings dump: carry forward what's solid, and \ + add, correct, or better-source whatever the gaps call for. Gaps include out-of-date,\ + irrelevant, or clearly wrong information. The date is {current_date}.\n\n\ + Solid findings from the last pass — keep and build on these:\n{}\n\n\ + Gaps the reviewer found — conclusions not actually backed by their source, sources \ + that don't line up with the conclusion drawn from them, or parts of the topic still \ + uncovered:\n{}", + review.solid_findings, review.gaps + ), + }; + + let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); + let run_result = researcher + .runner(task) + .max_turns(MAX_RESEARCH_TURNS) + .run() + .await; + drop(spinner); + + let gathered = match run_result { + Ok(response) => GatheredFindings { + findings: response.output, + incomplete: false, + }, + Err(PromptError::MaxTurnsError { chat_history, .. }) => { + tracing::warn!(round, "researcher hit its turn budget before concluding"); + GatheredFindings { + findings: summarize_partial_history(client, topic, &chat_history, show_progress) + .await, + incomplete: true, + } + } + Err(err) => return Err(err.into()), + }; + + tracing::info!(round, findings = %gathered.findings, incomplete = gathered.incomplete, "research phase complete"); + + Ok(gathered) +} diff --git a/deep_research/src/summarizer.rs b/deep_research/src/summarizer.rs new file mode 100644 index 0000000..c25e053 --- /dev/null +++ b/deep_research/src/summarizer.rs @@ -0,0 +1,69 @@ +use crate::history::{annotated_transcript_from_history, partial_findings_from_history}; +use crate::models::WRITER_MODEL; +use crate::progress::{SUMMARIZE_EMOJI, Spinner}; +use rig::client::AgentClientExt; +use rig::completion::{Message, Prompt}; +use rig::providers::ollama; + +/// Agentic alternative to the plain programmatic extraction in `history`: +/// hands the annotated transcript to a fresh model call and asks it to +/// reconstruct the same footnote-style findings dump the researcher would +/// have written itself, had it not run out of turns. This can dedupe +/// repeated URLs and restore correct `[n]` citation numbering in a way +/// string concatenation can't — but it's a model call like any other, so on +/// an empty transcript or a failure it falls back to +/// `partial_findings_from_history` rather than letting a second turn-budget +/// problem take down the one recovery path that's supposed to be +/// bulletproof. +#[tracing::instrument(skip(client, chat_history), fields(gen_ai.agent.name = "history-summarizer"))] +pub(crate) async fn summarize_partial_history( + client: &ollama::Client, + topic: &str, + chat_history: &[Message], + show_progress: bool, +) -> String { + let transcript = annotated_transcript_from_history(chat_history); + + if transcript.is_empty() { + tracing::warn!("no usable transcript to summarize; skipping summarizer pass"); + return partial_findings_from_history(chat_history); + } + + let summarizer = client + .agent(WRITER_MODEL) + .name("history-summarizer") + .preamble( + "You are reconstructing research notes from a research session that was cut off \ + before the researcher could write its own summary. You'll be given a raw transcript \ + of tool calls (searches run, pages fetched), their results, and any interim comments \ + the researcher made. Turn this into a footnote-style findings dump: write each fact \ + the transcript actually supports, followed by a bracketed number like [1], then end \ + with a 'Sources' section mapping each number to its exact URL, one per line. Reuse \ + the same number for a URL that appears more than once. Do not invent facts beyond \ + what the transcript shows, and note explicitly where it looks thin or cuts off \ + mid-investigation.", + ) + .build(); + + let spinner = Spinner::start( + show_progress, + format!("{SUMMARIZE_EMOJI} Reconstructing partial findings..."), + ); + let result = summarizer + .prompt(format!( + "Topic: {topic}\n\nPartial research transcript:\n{transcript}" + )) + .await; + drop(spinner); + + match result { + Ok(findings) => { + tracing::info!(findings = %findings, "summarizer reconstructed partial findings"); + findings + } + Err(error) => { + tracing::warn!(%error, "summarizer agent failed; falling back to programmatic extraction"); + partial_findings_from_history(chat_history) + } + } +} diff --git a/deep_research/src/writer.rs b/deep_research/src/writer.rs new file mode 100644 index 0000000..bc5d3c7 --- /dev/null +++ b/deep_research/src/writer.rs @@ -0,0 +1,77 @@ +use crate::models::WRITER_MODEL; +use crate::progress::{REPORT_EMOJI, Spinner}; +use crate::stream::write_text_stream; +use rig::client::AgentClientExt; +use rig::providers::ollama; +use rig::streaming::StreamingPrompt; +use std::io::Write; + +#[tracing::instrument(skip(client, findings), fields(gen_ai.agent.name = "writer"))] +pub(crate) async fn write_report( + client: &ollama::Client, + topic: &str, + findings: &str, + incomplete: bool, + show_progress: bool, +) -> anyhow::Result { + let writer = client + .agent(WRITER_MODEL) + .name("writer") + .preamble( + "You turn raw research notes into a clear, well-organized report for the reader. The \ + notes use footnote-style citations — a bracketed number like [1] after a fact, with a \ + Sources section mapping numbers to URLs. Preserve this scheme in your report: keep the \ + same [n] markers next to the claims they support (renumbering only if you drop unused \ + sources), and end the report with a 'Sources' section listing every footnote number \ + still in use next to its exact URL. Structure the body with headings, and call out any \ + open questions or contradictions the research turned up. Do not invent facts beyond \ + what the notes provide.", + ) + .build(); + + let task = if incomplete { + format!( + "Topic: {topic}\n\n\ + The researcher ran out of turns before it could finish investigating, so these notes \ + are an incomplete, best-effort reconstruction rather than the researcher's own \ + conclusions — treat any citations already in them as tentative. Write up what the \ + notes actually support, note explicitly where evidence is thin or missing, and open \ + the report with a short callout that this research is inconclusive and further \ + manual searching is needed.\n\n\ + Research notes:\n{findings}" + ) + } else { + format!("Topic: {topic}\n\nResearch notes:\n{findings}") + }; + + let spinner = Spinner::start(show_progress, format!("{REPORT_EMOJI} Writing report...")); + let response_stream = writer.stream_prompt(task).await; + + // Locked once for the whole stream rather than per chunk (as print! + // would do internally) — chunks arrive in a tight loop, so re-acquiring + // the lock on every one adds up. The spinner keeps running until the + // stream's first chunk arrives, so the terminal stays covered through + // the gap between sending the prompt and generation actually starting. + let stdout = std::io::stdout(); + let mut handle = stdout.lock(); + + // Guaranteed regardless of whether the writer model actually heeds the + // instruction above — the reader should never mistake a budget-exhausted + // run for a complete one just because the model forgot to say so. + let disclaimer = incomplete.then_some( + "> **Disclaimer:** incomplete research: the researcher exhausted its turn \ + budget before it could finish investigating. The findings below are partial and may be \ + incomplete or unbalanced; treat them as a starting point and verify further manually.\n\n", + ); + if let Some(disclaimer) = disclaimer { + write!(handle, "{disclaimer}")?; + } + + let report = write_text_stream(response_stream, &mut handle, spinner).await?; + writeln!(handle)?; + + Ok(match disclaimer { + Some(disclaimer) => format!("{disclaimer}{report}"), + None => report, + }) +}