doubleo7/src/researcher.rs

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, doc_context), fields(gen_ai.agent.name = "researcher"))]
pub(crate) async fn gather_findings(
    client: &ollama::Client,
    topic: &str,
    feedback: Option<&Review>,
    doc_context: Option<&str>,
    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 task = match doc_context {
        Some(context) => format!(
            "{task}\n\n\
             Relevant excerpts from documents the user uploaded — treat these as trusted primary \
             sources alongside anything you find on the web, and cite them with the same \
             footnote scheme (their Sources entry can just be the document path shown below):\n\
             {context}"
        ),
        None => task,
    };

    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)
}