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)