path: root/crates/asist-tools/src/search.rs

//! Web search.
//!
//! The result will be read aloud, so what matters is not a list of links but
//! an answer. That is why a search engine that summarizes is preferred
//! (Tavily returns an already written paragraph), and the headlines only come
//! along as a fallback when there is no summary.

use std::path::PathBuf;
use std::time::{Duration, Instant};

use asist_core::error::{Error, Result};
use asist_core::proc;
use asist_core::tools::Tool;
use serde_json::{json, Value};

/// Where the results come from.
#[derive(Debug, Clone)]
pub enum SearchBackend {
    /// Tavily API. It returns an already written answer besides the
    /// results, which is exactly what is needed to speak it.
    Tavily { api_key: String },
    /// SearXNG instance, your own or a trusted one. No key, but it only
    /// returns results: the model has to do the summary.
    SearxNG { base_url: String },
    /// An external program that prints the results as JSON.
    ///
    /// It is the keyless path: `scripts/search-ddgs.sh` queries DuckDuckGo and
    /// friends through the `ddgs` library. It works for anything else that writes
    /// JSON to standard output (a bridge to an MCP server, an internal search
    /// engine), so it is also the extension point of the search feature.
    ///
    /// In `args`, `{query}` and `{max}` are substituted before running.
    Command { program: PathBuf, args: Vec<String> },
}

impl SearchBackend {
    /// Default keyless backend: the script that wraps ddgs.
    pub fn ddgs(script: impl Into<PathBuf>) -> Self {
        SearchBackend::Command {
            program: script.into(),
            args: vec!["{query}".into(), "{max}".into()],
        }
    }
}

impl SearchBackend {
    pub fn label(&self) -> &'static str {
        match self {
            SearchBackend::Tavily { .. } => "Tavily",
            SearchBackend::SearxNG { .. } => "SearXNG",
            SearchBackend::Command { .. } => "command",
        }
    }
}

pub struct WebSearch {
    backend: SearchBackend,
    max_results: usize,
    timeout: Duration,
}

impl WebSearch {
    pub fn new(backend: SearchBackend, max_results: usize, timeout: Duration) -> Self {
        Self {
            backend,
            max_results: max_results.clamp(1, 10),
            timeout,
        }
    }

    pub fn backend(&self) -> &SearchBackend {
        &self.backend
    }

    fn agent(&self) -> ureq::Agent {
        ureq::Agent::config_builder()
            .timeout_global(Some(self.timeout))
            .build()
            .into()
    }

    fn fail(message: impl Into<String>) -> Error {
        Error::Tool {
            tool: "buscar_en_internet".into(),
            message: message.into(),
        }
    }

    fn tavily(&self, key: &str, query: &str) -> Result<String> {
        let body = json!({
            "query": query,
            "max_results": self.max_results,
            "search_depth": "basic",
            // The written answer is the reason to use this engine: without it
            // another model round would be spent summarizing.
            "include_answer": true,
        });
        let mut response = self
            .agent()
            .post("https://api.tavily.com/search")
            .header("Authorization", &format!("Bearer {key}"))
            .send_json(&body)
            .map_err(|e| Self::fail(describe_ureq(&e)))?;

        let parsed: Value = response
            .body_mut()
            .read_json()
            .map_err(|e| Self::fail(format!("respuesta ilegible: {e}")))?;

        let (answer, results) = extract(&parsed, self.max_results);
        compose(&answer, &results)
    }

    /// Runs the configured program and translates whatever it prints.
    fn command(&self, program: &std::path::Path, args: &[String], query: &str) -> Result<String> {
        let rendered: Vec<String> = args
            .iter()
            .map(|arg| {
                arg.replace("{query}", query)
                    .replace("{max}", &self.max_results.to_string())
            })
            .collect();

        // The arguments go to `execve` as they are: the query comes from what
        // the microphone heard, and it must never end up interpreted by a
        // shell.
        let output = proc::run(program, &rendered, self.timeout, None)
            .map_err(|e| Self::fail(e.to_string()))?;
        if !output.success() {
            return Err(Self::fail(format!(
                "el buscador falló: {}",
                output.last_error_line()
            )));
        }

        let parsed: Value = serde_json::from_slice(&output.stdout)
            .map_err(|e| Self::fail(format!("el buscador no devolvió JSON válido: {e}")))?;
        let (answer, results) = extract(&parsed, self.max_results);
        compose(&answer, &results)
    }

    fn searxng(&self, base_url: &str, query: &str) -> Result<String> {
        let url = format!(
            "{}/search?q={}&format=json&language=es",
            base_url.trim_end_matches('/'),
            urlencode(query)
        );
        let mut response = self
            .agent()
            .get(&url)
            .call()
            .map_err(|e| Self::fail(describe_ureq(&e)))?;

        let parsed: Value = response
            .body_mut()
            .read_json()
            .map_err(|e| Self::fail(format!("respuesta ilegible: {e}")))?;

        let (answer, results) = extract(&parsed, self.max_results);
        compose(&answer, &results)
    }
}

impl Tool for WebSearch {
    fn name(&self) -> &str {
        "buscar_en_internet"
    }

    fn description(&self) -> &str {
        "Busca información actual en internet y devuelve un resumen. Úsala para \
         noticias, precios, resultados, el tiempo o cualquier cosa posterior a tu \
         entrenamiento, en vez de responder de memoria."
    }

    fn parameters(&self) -> Value {
        json!({
            "type": "object",
            "properties": {
                "consulta": {
                    "type": "string",
                    "description": "Qué buscar, en lenguaje natural. Por ejemplo: tiempo en Buenos Aires mañana"
                }
            },
            "required": ["consulta"]
        })
    }

    fn acknowledgement(&self) -> Option<&str> {
        Some("Déjame que lo busque.")
    }

    fn call(&self, args: &Value) -> Result<String> {
        let query = args
            .get("consulta")
            .and_then(Value::as_str)
            .map(str::trim)
            .filter(|q| !q.is_empty())
            .ok_or_else(|| Self::fail("falta «consulta»"))?;

        let started = Instant::now();
        let result = match &self.backend {
            SearchBackend::Tavily { api_key } => self.tavily(api_key, query),
            SearchBackend::SearxNG { base_url } => self.searxng(base_url, query),
            SearchBackend::Command { program, args } => self.command(program, args, query),
        };
        tracing::info!(
            target: "tools",
            backend = self.backend.label(),
            query = query,
            ms = started.elapsed().as_millis(),
            ok = result.is_ok(),
            "search"
        );
        result
    }
}

/// Pulls the answer and results out of a JSON without requiring a specific
/// shape.
///
/// Each engine names the fields its own way (`href` or `url`, `body` or
/// `content`, a bare list or inside `results`), and what matters here is that
/// a new script works without touching Rust.
fn extract(parsed: &Value, max: usize) -> (String, Vec<String>) {
    let answer = parsed
        .get("answer")
        .and_then(Value::as_str)
        .unwrap_or("")
        .trim()
        .to_string();

    let rows = parsed
        .as_array()
        .or_else(|| parsed.get("results").and_then(Value::as_array));

    let results = rows
        .map(|items| {
            items
                .iter()
                .take(max)
                .filter_map