Split core.rs into one file per concern
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.
This commit is contained in:
parent
c150c67f1e
commit
6cd2d65158
10 changed files with 646 additions and 620 deletions
19
deep_research/src/cli.rs
Normal file
19
deep_research/src/cli.rs
Normal file
|
|
@ -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<String>,
|
||||
|
||||
/// 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<tracing::Level>,
|
||||
}
|
||||
|
|
@ -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<String>,
|
||||
|
||||
/// 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<tracing::Level>,
|
||||
}
|
||||
|
||||
/// 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<String> {
|
||||
let client = ollama::Client::new(Nothing)?;
|
||||
|
||||
let mut findings = String::new();
|
||||
let mut feedback: Option<Review> = None;
|
||||
let mut incomplete = false;
|
||||
|
||||
for round in 1..=MAX_RESEARCH_ROUNDS {
|
||||
let gathered =
|
||||
gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?;
|
||||
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<String> = 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<String> {
|
||||
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<String> {
|
||||
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::<Vec<String>>()
|
||||
.join("\n")
|
||||
}
|
||||
|
||||
fn transcript_lines(message: &Message) -> Vec<String> {
|
||||
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<GatheredFindings> {
|
||||
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<String> {
|
||||
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")]),
|
||||
""
|
||||
);
|
||||
}
|
||||
}
|
||||
292
deep_research/src/history.rs
Normal file
292
deep_research/src/history.rs
Normal file
|
|
@ -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<String> = 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<String> {
|
||||
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<String> {
|
||||
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::<Vec<String>>()
|
||||
.join("\n")
|
||||
}
|
||||
|
||||
fn transcript_lines(message: &Message) -> Vec<String> {
|
||||
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")]),
|
||||
""
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
@ -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(())
|
||||
}
|
||||
|
|
|
|||
8
deep_research/src/models.rs
Normal file
8
deep_research/src/models.rs
Normal file
|
|
@ -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";
|
||||
13
deep_research/src/observability.rs
Normal file
13
deep_research/src/observability.rs
Normal file
|
|
@ -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();
|
||||
}
|
||||
58
deep_research/src/research.rs
Normal file
58
deep_research/src/research.rs
Normal file
|
|
@ -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<String> {
|
||||
let client = ollama::Client::new(Nothing)?;
|
||||
|
||||
let mut findings = String::new();
|
||||
let mut feedback: Option<Review> = None;
|
||||
let mut incomplete = false;
|
||||
|
||||
for round in 1..=MAX_RESEARCH_ROUNDS {
|
||||
let gathered =
|
||||
gather_findings(&client, topic, feedback.as_ref(), round, show_progress).await?;
|
||||
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
|
||||
}
|
||||
98
deep_research/src/researcher.rs
Normal file
98
deep_research/src/researcher.rs
Normal file
|
|
@ -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<GatheredFindings> {
|
||||
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)
|
||||
}
|
||||
69
deep_research/src/summarizer.rs
Normal file
69
deep_research/src/summarizer.rs
Normal file
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
77
deep_research/src/writer.rs
Normal file
77
deep_research/src/writer.rs
Normal file
|
|
@ -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<String> {
|
||||
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,
|
||||
})
|
||||
}
|
||||
Loading…
Reference in a new issue