From 187d3271446edb918fb64f566c0312285b325eb0 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 11:26:53 +0200 Subject: [PATCH 1/6] Write a partial report instead of erroring out when research hits max turns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the researcher agent exhausts its turn budget mid-investigation, rig now surfaces PromptError::MaxTurnsError with the chat history intact rather than nothing at all. Catch it, reconstruct a findings dump from whatever assistant text and tool results the run produced, and still write a report from that — clearly flagged as incomplete — instead of propagating the error and losing all the work. --- deep_research/src/core.rs | 134 +++++++++++++++++++++++++++++++++----- 1 file changed, 119 insertions(+), 15 deletions(-) diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index 0a55425..dd177a7 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -4,6 +4,8 @@ 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, PromptError}; use rig::providers::ollama; use rig::streaming::StreamingPrompt; use std::io::Write; @@ -60,9 +62,21 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result let mut findings = String::new(); let mut feedback: Option = None; + let mut incomplete = false; for round in 1..=MAX_RESEARCH_ROUNDS { - findings = gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?; + 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; @@ -80,7 +94,61 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result feedback = Some(review); } - write_report(&client, topic, &findings, show_progress).await + 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 mut sections = Vec::new(); + + for message in chat_history { + match message { + Message::Assistant { content, .. } => { + for item in content.iter() { + if let AssistantContent::Text(text) = item { + sections.push(text.text().to_string()); + } + } + } + Message::User { content } => { + for item in content.iter() { + if let UserContent::ToolResult(result) = item { + for part in result.content.iter() { + if let ToolResultContent::Text(text) = part { + let mut snippet = text.text().to_string(); + if snippet.len() > MAX_TOOL_RESULT_CHARS { + snippet.truncate(MAX_TOOL_RESULT_CHARS); + snippet.push_str(" [...truncated]"); + } + sections.push(snippet); + } + } + } + } + } + Message::System { .. } => {} + } + } + + if sections.is_empty() { + "The researcher exhausted its turn budget before gathering any usable evidence." + .to_string() + } else { + sections.join("\n\n") + } } /// Wraps the tool-calling research loop in its own span so it's visible as a @@ -93,7 +161,7 @@ async fn gather_findings( feedback: Option<&Review>, round: usize, show_progress: bool, -) -> anyhow::Result { +) -> anyhow::Result { let current_date = chrono::offset::Local::now().to_string(); let researcher = client @@ -136,17 +204,24 @@ async fn gather_findings( }; let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); - let findings = researcher - .runner(task) - .max_turns(MAX_RESEARCH_TURNS) - .run() - .await? - .output; + let run_result = researcher.runner(task).max_turns(MAX_RESEARCH_TURNS).run().await; drop(spinner); - tracing::info!(round, findings = %findings, "research phase complete"); + 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: partial_findings_from_history(&chat_history), + incomplete: true, + } + } + Err(err) => return Err(err.into()), + }; - Ok(findings) + 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"))] @@ -154,6 +229,7 @@ async fn write_report( client: &ollama::Client, topic: &str, findings: &str, + incomplete: bool, show_progress: bool, ) -> anyhow::Result { let writer = client @@ -171,10 +247,22 @@ async fn write_report( ) .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 incomplete and may be a raw, unsorted mix of tool output rather than clean \ + footnoted findings. 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(format!("Topic: {topic}\n\nResearch notes:\n{findings}")) - .await; + 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 @@ -183,8 +271,24 @@ async fn write_report( // 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:** this research did not complete — 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(report) + Ok(match disclaimer { + Some(disclaimer) => format!("{disclaimer}{report}"), + None => report, + }) } -- 2.45.2 From 3ddc48fe18500b88483c6851f7f0c1ac747c08c0 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 11:30:45 +0200 Subject: [PATCH 2/6] Flatten partial_findings_from_history's nested for/match/if-let pyramid Split into two filter_map-based helpers (assistant_text, tool_result_text) and a shared truncate() so the extraction reads as a flat iterator chain instead of four levels of nesting. --- deep_research/src/core.rs | 81 ++++++++++++++++++++++++--------------- 1 file changed, 51 insertions(+), 30 deletions(-) diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index dd177a7..300ecaf 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -112,36 +112,14 @@ struct GatheredFindings { fn partial_findings_from_history(chat_history: &[Message]) -> String { const MAX_TOOL_RESULT_CHARS: usize = 2000; - let mut sections = Vec::new(); - - for message in chat_history { - match message { - Message::Assistant { content, .. } => { - for item in content.iter() { - if let AssistantContent::Text(text) = item { - sections.push(text.text().to_string()); - } - } - } - Message::User { content } => { - for item in content.iter() { - if let UserContent::ToolResult(result) = item { - for part in result.content.iter() { - if let ToolResultContent::Text(text) = part { - let mut snippet = text.text().to_string(); - if snippet.len() > MAX_TOOL_RESULT_CHARS { - snippet.truncate(MAX_TOOL_RESULT_CHARS); - snippet.push_str(" [...truncated]"); - } - sections.push(snippet); - } - } - } - } - } - Message::System { .. } => {} - } - } + let sections: Vec = chat_history + .iter() + .flat_map(|message| { + assistant_text(message) + .into_iter() + .chain(tool_result_text(message, MAX_TOOL_RESULT_CHARS)) + }) + .collect(); if sections.is_empty() { "The researcher exhausted its turn budget before gathering any usable evidence." @@ -151,6 +129,49 @@ fn partial_findings_from_history(chat_history: &[Message]) -> String { } } +/// Plain-text blocks from an assistant message, if any. +fn assistant_text(message: &Message) -> Vec { + let Message::Assistant { content, .. } = message else { + return Vec::new(); + }; + content + .iter() + .filter_map(|item| 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| match item { + UserContent::ToolResult(result) => Some(result), + _ => None, + }) + .flat_map(|result| result.content.iter()) + .filter_map(|part| match part { + ToolResultContent::Text(text) => Some(truncate(text.text(), max_chars)), + _ => None, + }) + .collect() +} + +fn truncate(text: &str, max_chars: usize) -> String { + if text.len() <= max_chars { + text.to_string() + } else { + format!("{} [...truncated]", &text[..max_chars]) + } +} + /// 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. -- 2.45.2 From e4460467443194b8281e9dea606c905e53a83e25 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 11:43:14 +0200 Subject: [PATCH 3/6] Add unit tests for the decomposed max-turns recovery helpers Covers truncate, assistant_text, tool_result_text, and partial_findings_from_history in isolation (fallback text, ordering, truncation, and filtering out non-text content). gather_findings/ write_report/research still need a live ollama client and aren't covered here. --- deep_research/Cargo.toml | 3 ++ deep_research/src/core.rs | 104 ++++++++++++++++++++++++++++++++++++-- 2 files changed, 102 insertions(+), 5 deletions(-) diff --git a/deep_research/Cargo.toml b/deep_research/Cargo.toml index 9c4e5b8..179c84c 100644 --- a/deep_research/Cargo.toml +++ b/deep_research/Cargo.toml @@ -17,3 +17,6 @@ serde = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } + +[dev-dependencies] +serde_json = "1" diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index 300ecaf..53ff559 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -121,11 +121,10 @@ fn partial_findings_from_history(chat_history: &[Message]) -> String { }) .collect(); - if sections.is_empty() { - "The researcher exhausted its turn budget before gathering any usable evidence." - .to_string() - } else { - sections.join("\n\n") + match sections.is_empty() { + true => "The researcher exhausted its turn budget before gathering any usable evidence." + .to_string(), + false => sections.join("\n"), } } @@ -313,3 +312,98 @@ async fn write_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!(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!(assistant_text(&message).is_empty()); + } + + #[test] + fn assistant_text_ignores_non_assistant_messages() { + assert!(assistant_text(&Message::user("hi")).is_empty()); + assert!(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]" + ); + } +} -- 2.45.2 From 21030462b15d3f4cc98687b5e3135a2e181275e4 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 11:54:23 +0200 Subject: [PATCH 4/6] Prototype an agentic summarizer for max-turns recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add summarize_partial_history: a one-shot writer-model pass that turns an annotated transcript (tool calls with their args, so a fetch's URL or a search's query stays attached to its result, plus results and interim notes) into a proper footnote-style findings dump, instead of the flat concatenation partial_findings_from_history produces on its own. It's wired in as the MaxTurnsError recovery path in gather_findings, but partial_findings_from_history stays as the fallback for an empty transcript or if the summarizer call itself fails — the one guaranteed recovery path shouldn't have a second turn-budget/model failure as a single point of failure. Observability: instrumented with the same #[tracing::instrument(fields( gen_ai.agent.name = ...))] + spinner pattern as the researcher/reviewer/ writer phases, with info!/warn! events on success, empty-transcript skip, and summarizer failure. Also: cargo fmt across the crate (unrelated formatting drift had accumulated), and adds unit tests for the new transcript_lines / annotated_transcript_from_history helpers. --- Cargo.lock | 1 + deep_research/src/core.rs | 226 +++++++++++++++++++++++++++++++--- deep_research/src/main.rs | 2 +- deep_research/src/progress.rs | 1 + deep_research/src/review.rs | 9 +- deep_research/src/stream.rs | 17 +-- deep_research/src/tools.rs | 29 ++++- 7 files changed, 254 insertions(+), 31 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 15dba5d..b2a255e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1941,6 +1941,7 @@ dependencies = [ "schemars 1.2.2", "scraper", "serde", + "serde_json", "tokio", "tracing", "tracing-subscriber", diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index 53ff559..1a31483 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -1,11 +1,11 @@ -use crate::progress::{REJECTED_EMOJI, REPORT_EMOJI, RESEARCH_EMOJI, Spinner}; +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, PromptError}; +use rig::completion::{AssistantContent, Message, Prompt, PromptError}; use rig::providers::ollama; use rig::streaming::StreamingPrompt; use std::io::Write; @@ -23,7 +23,6 @@ const MAX_RESEARCH_TURNS: usize = 12; /// 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. @@ -88,7 +87,10 @@ pub(crate) async fn research(topic: &str, show_progress: bool) -> anyhow::Result } if show_progress { - eprintln!("{REJECTED_EMOJI} Findings rejected — revising for round {}...", round + 1); + eprintln!( + "{REJECTED_EMOJI} Findings rejected — revising for round {}...", + round + 1 + ); } feedback = Some(review); @@ -123,7 +125,7 @@ fn partial_findings_from_history(chat_history: &[Message]) -> String { match sections.is_empty() { true => "The researcher exhausted its turn budget before gathering any usable evidence." - .to_string(), + .to_string(), false => sections.join("\n"), } } @@ -171,6 +173,115 @@ fn truncate(text: &str, max_chars: usize) -> String { } } +/// 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. @@ -224,15 +335,23 @@ async fn gather_findings( }; let spinner = Spinner::start(show_progress, format!("{RESEARCH_EMOJI} Researching...")); - let run_result = researcher.runner(task).max_turns(MAX_RESEARCH_TURNS).run().await; + 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 }, + 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: partial_findings_from_history(&chat_history), + findings: summarize_partial_history(client, topic, &chat_history, show_progress) + .await, incomplete: true, } } @@ -271,10 +390,11 @@ async fn write_report( format!( "Topic: {topic}\n\n\ The researcher ran out of turns before it could finish investigating, so these notes \ - are incomplete and may be a raw, unsorted mix of tool output rather than clean \ - footnoted findings. 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\ + 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 { @@ -296,7 +416,7 @@ async fn write_report( // 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:** this research did not complete — the researcher exhausted its turn \ + "> **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", ); @@ -360,10 +480,16 @@ mod tests { #[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()]); + 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()]); + assert_eq!( + tool_result_text(&long, 5), + vec!["01234 [...truncated]".to_string()] + ); } #[test] @@ -406,4 +532,74 @@ mod tests { "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 1d1b553..087a844 100644 --- a/deep_research/src/main.rs +++ b/deep_research/src/main.rs @@ -1,8 +1,8 @@ use clap::Parser; +mod core; mod progress; mod review; -mod core; mod stream; mod tools; diff --git a/deep_research/src/progress.rs b/deep_research/src/progress.rs index ed0312f..d5117af 100644 --- a/deep_research/src/progress.rs +++ b/deep_research/src/progress.rs @@ -10,6 +10,7 @@ pub(crate) const FETCH_EMOJI: &str = "📄"; pub(crate) const REVIEW_EMOJI: &str = "🧐"; pub(crate) const REJECTED_EMOJI: &str = "❌"; pub(crate) const REPORT_EMOJI: &str = "✍️"; +pub(crate) const SUMMARIZE_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 diff --git a/deep_research/src/review.rs b/deep_research/src/review.rs index 0c6bb2c..8b43b59 100644 --- a/deep_research/src/review.rs +++ b/deep_research/src/review.rs @@ -50,9 +50,14 @@ pub(crate) async fn review_findings( .retries(2) .build(); - let spinner = Spinner::start(show_progress, format!("{REVIEW_EMOJI} Reviewing findings...")); + let spinner = Spinner::start( + show_progress, + format!("{REVIEW_EMOJI} Reviewing findings..."), + ); let review = reviewer - .extract(format!("Topic: {topic}\n\nResearch findings to review:\n{findings}")) + .extract(format!( + "Topic: {topic}\n\nResearch findings to review:\n{findings}" + )) .await?; drop(spinner); diff --git a/deep_research/src/stream.rs b/deep_research/src/stream.rs index bb3622b..9e24127 100644 --- a/deep_research/src/stream.rs +++ b/deep_research/src/stream.rs @@ -64,9 +64,10 @@ mod tests { let items = vec![text_item("Hello, "), text_item("world!")]; let mut written = Vec::new(); - let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) - .await - .unwrap(); + let accumulated = + write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) + .await + .unwrap(); assert_eq!(accumulated, "Hello, world!"); assert_eq!(String::from_utf8(written).unwrap(), "Hello, world!"); @@ -81,9 +82,10 @@ mod tests { let items = vec![text_item("kept"), final_item]; let mut written = Vec::new(); - let accumulated = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) - .await - .unwrap(); + let accumulated = + write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")) + .await + .unwrap(); assert_eq!(accumulated, "kept"); assert_eq!(String::from_utf8(written).unwrap(), "kept"); @@ -97,7 +99,8 @@ mod tests { let items = vec![text_item("kept"), Err(error)]; let mut written = Vec::new(); - let result = write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")).await; + let result = + write_text_stream(stream::iter(items), &mut written, Spinner::start(false, "")).await; assert!(result.is_err()); assert_eq!(String::from_utf8(written).unwrap(), "kept"); diff --git a/deep_research/src/tools.rs b/deep_research/src/tools.rs index 85c5f91..5a8d8c0 100644 --- a/deep_research/src/tools.rs +++ b/deep_research/src/tools.rs @@ -8,7 +8,10 @@ const MAX_PAGE_CHARS: usize = 6000; /// Searches the web via DuckDuckGo's HTML endpoint (no API key required) and /// returns each hit's title, URL, and snippet so the caller can decide which /// pages are worth fetching in full. -#[rig::tool_macro(description = "Search the web for pages related to a query", required(query))] +#[rig::tool_macro( + description = "Search the web for pages related to a query", + required(query) +)] pub(crate) async fn search_web( /// The search query query: String, @@ -23,7 +26,10 @@ pub(crate) async fn search_web( .await .map_err(ToolExecutionError::from_error)?; - let body = response.text().await.map_err(ToolExecutionError::from_error)?; + let body = response + .text() + .await + .map_err(ToolExecutionError::from_error)?; let results = parse_search_results(&body); if results.is_empty() { @@ -40,7 +46,10 @@ pub(crate) async fn search_web( /// Fetches a page and returns its main text content, stripped of markup and /// truncated so a single fetch can't blow out the model's context window. -#[rig::tool_macro(description = "Fetch a web page and return its readable text content", required(url))] +#[rig::tool_macro( + description = "Fetch a web page and return its readable text content", + required(url) +)] pub(crate) async fn fetch_page( /// The URL to fetch url: String, @@ -54,7 +63,10 @@ pub(crate) async fn fetch_page( .await .map_err(ToolExecutionError::from_error)?; - let body = response.text().await.map_err(ToolExecutionError::from_error)?; + let body = response + .text() + .await + .map_err(ToolExecutionError::from_error)?; Ok(extract_readable_text(&body)) } @@ -91,7 +103,11 @@ fn parse_search_results(body: &str) -> Vec { if title.is_empty() || url.is_empty() { None } else { - Some(SearchResult { title, url, snippet }) + Some(SearchResult { + title, + url, + snippet, + }) } }) .take(MAX_SEARCH_RESULTS) @@ -118,7 +134,8 @@ fn resolve_ddg_redirect(href: &str) -> String { fn extract_readable_text(html: &str) -> String { let document = Html::parse_document(html); - let content_selector = Selector::parse("p, h1, h2, h3, h4, h5, li, td").expect("valid selector"); + let content_selector = + Selector::parse("p, h1, h2, h3, h4, h5, li, td").expect("valid selector"); let mut text: String = document .select(&content_selector) -- 2.45.2 From c150c67f1e505df2b695b5afba9d6df38994def7 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 12:25:29 +0200 Subject: [PATCH 5/6] Switch search_web from DuckDuckGo HTML scraping to local SearXNG MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DuckDuckGo's HTML endpoint rate-limits after enough requests, and a rate-limited response is indistinguishable from a genuine empty result — which is exactly what burned a full 12-turn research run on 13 consecutive "No results found" responses. Swapping to a local SearXNG instance's JSON API (no HTML scraping needed) fixes both problems: SearXNG spreads queries across multiple upstream engines instead of hammering one, and this machine already runs an instance. This tool is explicitly local-only and never released, so the base URL is a plain default (localhost:8080) overridable via SEARXNG_URL, not a general-purpose config surface. Evaluated the two third-party SearXNG crates on crates.io first (searxng, searxng-client) — both are single-maintainer v0.1.0 packages with no adoption signal and no official alternative exists, so a hand-rolled reqwest + serde call was the better bet for something this small. Drops the DuckDuckGo-specific HTML parsing (parse_search_results, resolve_ddg_redirect, the .result/.result__a/.result__snippet scraper selectors) entirely — fetch_page's extract_readable_text still needs scraper for arbitrary fetched pages, so that dependency stays. Adds an #[ignore]'d live smoke test (search_web_returns_real_results_from_local_searxng) for manually verifying against a running instance; not run by default since there's no CI environment with SearXNG available. --- deep_research/Cargo.toml | 2 +- deep_research/src/core.rs | 49 ++++++++------ deep_research/src/tools.rs | 131 ++++++++++++++++--------------------- 3 files changed, 87 insertions(+), 95 deletions(-) diff --git a/deep_research/Cargo.toml b/deep_research/Cargo.toml index 179c84c..18f0f38 100644 --- a/deep_research/Cargo.toml +++ b/deep_research/Cargo.toml @@ -9,7 +9,7 @@ chrono = "0.4.45" clap = { version = "4", features = ["derive"] } futures = { workspace = true } indicatif = "0.18.6" -reqwest = { workspace = true, features = ["query"] } +reqwest = { workspace = true, features = ["query", "json"] } rig = { workspace = true } schemars = "1" scraper = "0.27" diff --git a/deep_research/src/core.rs b/deep_research/src/core.rs index 1a31483..3c14a91 100644 --- a/deep_research/src/core.rs +++ b/deep_research/src/core.rs @@ -117,7 +117,7 @@ fn partial_findings_from_history(chat_history: &[Message]) -> String { let sections: Vec = chat_history .iter() .flat_map(|message| { - assistant_text(message) + extract_assistant_text(message) .into_iter() .chain(tool_result_text(message, MAX_TOOL_RESULT_CHARS)) }) @@ -131,13 +131,13 @@ fn partial_findings_from_history(chat_history: &[Message]) -> String { } /// Plain-text blocks from an assistant message, if any. -fn assistant_text(message: &Message) -> Vec { +fn extract_assistant_text(message: &Message) -> Vec { let Message::Assistant { content, .. } = message else { return Vec::new(); }; content .iter() - .filter_map(|item| match item { + .filter_map(|item: &AssistantContent| match item { AssistantContent::Text(text) => Some(text.text().to_string()), _ => None, }) @@ -151,26 +151,30 @@ fn tool_result_text(message: &Message, max_chars: usize) -> Vec { let Message::User { content } = message else { return Vec::new(); }; + 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(truncate(text.text(), max_chars)), + .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 { - if text.len() <= max_chars { - text.to_string() - } else { - format!("{} [...truncated]", &text[..max_chars]) + 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 @@ -182,7 +186,7 @@ fn annotated_transcript_from_history(chat_history: &[Message]) -> String { chat_history .iter() .flat_map(transcript_lines) - .collect::>() + .collect::>() .join("\n") } @@ -449,13 +453,16 @@ mod tests { #[test] fn truncate_cuts_long_text_and_marks_it() { - assert_eq!(truncate("hello world", 5), "hello [...truncated]"); + assert_eq!(truncate("hello world", 5), "hello ...[truncated]"); } #[test] fn assistant_text_extracts_text_blocks() { let message = Message::assistant("found it"); - assert_eq!(assistant_text(&message), vec!["found it".to_string()]); + assert_eq!( + extract_assistant_text(&message), + vec!["found it".to_string()] + ); } #[test] @@ -468,13 +475,13 @@ mod tests { serde_json::json!({ "query": "test" }), )), }; - assert!(assistant_text(&message).is_empty()); + assert!(extract_assistant_text(&message).is_empty()); } #[test] fn assistant_text_ignores_non_assistant_messages() { - assert!(assistant_text(&Message::user("hi")).is_empty()); - assert!(assistant_text(&Message::system("be careful")).is_empty()); + assert!(extract_assistant_text(&Message::user("hi")).is_empty()); + assert!(extract_assistant_text(&Message::system("be careful")).is_empty()); } #[test] @@ -488,7 +495,7 @@ mod tests { let long = Message::tool_result("call-2", "0123456789"); assert_eq!( tool_result_text(&long, 5), - vec!["01234 [...truncated]".to_string()] + vec!["01234 ...[truncated]".to_string()] ); } diff --git a/deep_research/src/tools.rs b/deep_research/src/tools.rs index 5a8d8c0..51237c2 100644 --- a/deep_research/src/tools.rs +++ b/deep_research/src/tools.rs @@ -1,13 +1,39 @@ use crate::progress::{self, FETCH_EMOJI, SEARCH_EMOJI}; use rig::tool::ToolExecutionError; use scraper::{Html, Selector}; +use serde::Deserialize; const MAX_SEARCH_RESULTS: usize = 6; const MAX_PAGE_CHARS: usize = 6000; -/// Searches the web via DuckDuckGo's HTML endpoint (no API key required) and -/// returns each hit's title, URL, and snippet so the caller can decide which -/// pages are worth fetching in full. +/// Local-only tool: scraping DuckDuckGo directly shares rate-limit fate with +/// every other bot hitting it from this IP, and a rate-limited response +/// looks identical to a genuine "no results" — which is exactly what took +/// down a research run over a dozen turns without ever surfacing as an +/// error. A self-hosted SearXNG instance has its own JSON API (so no HTML +/// scraping) and spreads queries across multiple upstream engines instead +/// of hammering one. This is deliberately not configurable beyond the env +/// var below — this tool is never meant to run anywhere but this machine. +fn searxng_base_url() -> String { + std::env::var("SEARXNG_URL").unwrap_or_else(|_| "http://localhost:8080".to_string()) +} + +#[derive(Deserialize)] +struct SearxngResponse { + results: Vec, +} + +#[derive(Deserialize)] +struct SearxngResult { + title: String, + url: String, + #[serde(default)] + content: String, +} + +/// Searches the web via a local SearXNG instance's JSON API and returns each +/// hit's title, URL, and snippet so the caller can decide which pages are +/// worth fetching in full. #[rig::tool_macro( description = "Search the web for pages related to a query", required(query) @@ -19,27 +45,27 @@ pub(crate) async fn search_web( progress::set_activity(format!("{SEARCH_EMOJI} Searching: {query}")); let response = reqwest::Client::new() - .get("https://html.duckduckgo.com/html/") - .query(&[("q", query.as_str())]) - .header("User-Agent", "Mozilla/5.0 (research-agent)") + .get(format!("{}/search", searxng_base_url())) + .query(&[("q", query.as_str()), ("format", "json")]) .send() .await .map_err(ToolExecutionError::from_error)?; - let body = response - .text() + let parsed: SearxngResponse = response + .json() .await .map_err(ToolExecutionError::from_error)?; - let results = parse_search_results(&body); - if results.is_empty() { + if parsed.results.is_empty() { return Ok("No results found.".to_string()); } - Ok(results + Ok(parsed + .results .into_iter() + .take(MAX_SEARCH_RESULTS) .enumerate() - .map(|(i, r)| format!("{}. {}\n {}\n {}", i + 1, r.title, r.url, r.snippet)) + .map(|(i, r)| format!("{}. {}\n {}\n {}", i + 1, r.title, r.url, r.content)) .collect::>() .join("\n\n")) } @@ -71,67 +97,6 @@ pub(crate) async fn fetch_page( Ok(extract_readable_text(&body)) } -struct SearchResult { - title: String, - url: String, - snippet: String, -} - -/// DuckDuckGo's HTML results page wraps each hit in a `.result` block; the -/// title/link lives in `.result__a` and links are redirected through -/// `duckduckgo.com/l/?uddg=`, so the real URL has to be pulled back -/// out of that query parameter rather than used as-is. -fn parse_search_results(body: &str) -> Vec { - let document = Html::parse_document(body); - let result_selector = Selector::parse(".result").expect("valid selector"); - let title_selector = Selector::parse(".result__a").expect("valid selector"); - let snippet_selector = Selector::parse(".result__snippet").expect("valid selector"); - - document - .select(&result_selector) - .filter_map(|result| { - let title_el = result.select(&title_selector).next()?; - let href = title_el.value().attr("href")?; - let url = resolve_ddg_redirect(href); - let title = title_el.text().collect::().trim().to_string(); - let snippet = result - .select(&snippet_selector) - .next() - .map(|el| el.text().collect::().trim().to_string()) - .unwrap_or_default(); - - if title.is_empty() || url.is_empty() { - None - } else { - Some(SearchResult { - title, - url, - snippet, - }) - } - }) - .take(MAX_SEARCH_RESULTS) - .collect() -} - -fn resolve_ddg_redirect(href: &str) -> String { - let full = if href.starts_with("//") { - format!("https:{href}") - } else { - href.to_string() - }; - - reqwest::Url::parse(&full) - .ok() - .and_then(|parsed| { - parsed - .query_pairs() - .find(|(k, _)| k == "uddg") - .map(|(_, v)| v.into_owned()) - }) - .unwrap_or(full) -} - fn extract_readable_text(html: &str) -> String { let document = Html::parse_document(html); let content_selector = @@ -150,3 +115,23 @@ fn extract_readable_text(html: &str) -> String { let collapsed = text.split_whitespace().collect::>().join(" "); collapsed.chars().take(MAX_PAGE_CHARS).collect() } + +#[cfg(test)] +mod live_smoke_test { + use super::*; + + /// Not run by default — this tool is local-only by design, so there's no + /// CI environment where a SearXNG instance would exist to test against. + /// Run manually with `cargo test -- --ignored` when SEARXNG_URL (or the + /// localhost:8080 default) points at a running instance. + #[tokio::test] + #[ignore = "hits a real local SearXNG instance; run manually with --ignored"] + async fn search_web_returns_real_results_from_local_searxng() { + let output = search_web("uruguay senior software engineer hiring 2026".to_string()) + .await + .expect("search_web should succeed against a live local SearXNG instance"); + + println!("{output}"); + assert_ne!(output, "No results found."); + } +} -- 2.45.2 From 6cd2d6515893270883452dcb054aa9c3fa01aa60 Mon Sep 17 00:00:00 2001 From: Austin Schaefer Date: Tue, 18 Aug 2026 12:59:03 +0200 Subject: [PATCH 6/6] 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, + }) +} -- 2.45.2