doubleo7/deep_research/src/tools.rs

use crate::progress::{self, FETCH_EMOJI, SEARCH_EMOJI};
use rig::tool::ToolExecutionError;
use scraper::{Html, Selector};

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.
#[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,
) -> Result<String, ToolExecutionError> {
    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)")
        .send()
        .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() {
        return Ok("No results found.".to_string());
    }

    Ok(results
        .into_iter()
        .enumerate()
        .map(|(i, r)| format!("{}. {}\n   {}\n   {}", i + 1, r.title, r.url, r.snippet))
        .collect::<Vec<_>>()
        .join("\n\n"))
}

/// 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)
)]
pub(crate) async fn fetch_page(
    /// The URL to fetch
    url: String,
) -> Result<String, ToolExecutionError> {
    progress::set_activity(format!("{FETCH_EMOJI} Fetching: {url}"));

    let response = reqwest::Client::new()
        .get(&url)
        .header("User-Agent", "Mozilla/5.0 (research-agent)")
        .send()
        .await
        .map_err(ToolExecutionError::from_error)?;

    let body = response
        .text()
        .await
        .map_err(ToolExecutionError::from_error)?;

    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=<real-url>`, 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<SearchResult> {
    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::<String>().trim().to_string();
            let snippet = result
                .select(&snippet_selector)
                .next()
                .map(|el| el.text().collect::<String>().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 =
        Selector::parse("p, h1, h2, h3, h4, h5, li, td").expect("valid selector");

    let mut text: String = document
        .select(&content_selector)
        .map(|el| el.text().collect::<Vec<_>>().join(" "))
        .collect::<Vec<_>>()
        .join("\n");

    if text.trim().is_empty() {
        text = document.root_element().text().collect::<Vec<_>>().join(" ");
    }

    let collapsed = text.split_whitespace().collect::<Vec<_>>().join(" ");
    collapsed.chars().take(MAX_PAGE_CHARS).collect()
}