89 lines
3.4 KiB
Rust
89 lines
3.4 KiB
Rust
|
|
use rig::client::{AgentClientExt, Nothing};
|
||
|
|
use rig::completion::Prompt;
|
||
|
|
use rig::providers::ollama;
|
||
|
|
use crate::deep_research::tools::{FetchPage, SearchWeb};
|
||
|
|
|
||
|
|
/// 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;
|
||
|
|
|
||
|
|
const DEFAULT_TOPIC: &str =
|
||
|
|
"What are the latest advances in running large language models locally, on consumer hardware?";
|
||
|
|
|
||
|
|
fn initialize_observability() {
|
||
|
|
tracing_subscriber::fmt()
|
||
|
|
.with_env_filter(
|
||
|
|
tracing_subscriber::EnvFilter::try_from_default_env()
|
||
|
|
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
|
||
|
|
)
|
||
|
|
.with_span_events(tracing_subscriber::fmt::format::FmtSpan::CLOSE)
|
||
|
|
.with_writer(std::io::stderr)
|
||
|
|
.init();
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Performs deep research via a two-stage agentic flow: a tool-calling agent
|
||
|
|
/// gathers and cross-checks evidence from the web, then a second agent turns
|
||
|
|
/// those raw notes into a structured report.
|
||
|
|
pub(crate) async fn start() -> anyhow::Result<()> {
|
||
|
|
initialize_observability();
|
||
|
|
|
||
|
|
let topic = std::env::args().nth(1).unwrap_or_else(|| DEFAULT_TOPIC.to_string());
|
||
|
|
|
||
|
|
let report = research(&topic).await?;
|
||
|
|
|
||
|
|
println!("{report}");
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn research(topic: &str) -> anyhow::Result<String> {
|
||
|
|
let client = ollama::Client::new(Nothing)?;
|
||
|
|
|
||
|
|
let researcher = client
|
||
|
|
.agent(RESEARCHER_MODEL)
|
||
|
|
.preamble(
|
||
|
|
"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. Once you are confident you have enough evidence, stop calling tools \
|
||
|
|
and reply with a plain-text dump of every fact you gathered, the source URL it came \
|
||
|
|
from, and any open questions or contradictions between sources. This is raw research \
|
||
|
|
material for a writer, not a final report, so favor completeness over polish.",
|
||
|
|
)
|
||
|
|
.tool(SearchWeb)
|
||
|
|
.tool(FetchPage)
|
||
|
|
.build();
|
||
|
|
|
||
|
|
tracing::info!(topic, "starting research phase");
|
||
|
|
|
||
|
|
let findings = researcher
|
||
|
|
.runner(topic)
|
||
|
|
.max_turns(MAX_RESEARCH_TURNS)
|
||
|
|
.run()
|
||
|
|
.await?
|
||
|
|
.output;
|
||
|
|
|
||
|
|
tracing::info!(findings = %findings, "research phase complete, starting writing phase");
|
||
|
|
|
||
|
|
let writer = client
|
||
|
|
.agent(WRITER_MODEL)
|
||
|
|
.preamble(
|
||
|
|
"You turn raw research notes into a clear, well-organized report for the reader. \
|
||
|
|
Structure the report with headings, cite source URLs inline next to the claims they \
|
||
|
|
support, and call out any open questions or contradictions the research turned up. Do \
|
||
|
|
not invent facts beyond what the notes provide.",
|
||
|
|
)
|
||
|
|
.build();
|
||
|
|
|
||
|
|
let report = writer
|
||
|
|
.prompt(format!("Topic: {topic}\n\nResearch notes:\n{findings}"))
|
||
|
|
.await?;
|
||
|
|
|
||
|
|
Ok(report)
|
||
|
|
}
|