src/gaia_agent/tools.py

import re
import tempfile
from pathlib import Path
from typing import Any, cast
from urllib.parse import urlparse

import httpx
from bs4 import BeautifulSoup
from ddgs import DDGS
from smolagents import tool
from yt_dlp import YoutubeDL

from gaia_agent.attachments import transcribe_media

MAX_QUERY_CHARACTERS = 500
MAX_PAGE_CHARACTERS = 12_000
MAX_SEARCH_RESULTS = 8
WEB_TIMEOUT_SECONDS = 20.0


@tool  # type: ignore[untyped-decorator]
def web_search(query: str) -> str:
    """Search the public web for evidence relevant to a research question.

    Args:
        query: A concise search-engine query of at most 500 characters.
    """
    query = query.strip()
    if not query or len(query) > MAX_QUERY_CHARACTERS:
        raise ValueError("query must contain between 1 and 500 characters")
    results = DDGS().text(query, max_results=MAX_SEARCH_RESULTS)
    return "\n\n".join(
        "\n".join(
            (
                f"Title: {item.get('title', '')}",
                f"URL: {item.get('href', '')}",
                f"Snippet: {item.get('body', '')}",
            )
        )
        for item in results
    )


@tool  # type: ignore[untyped-decorator]
def visit_webpage(url: str) -> str:
    """Download a public HTTP(S) webpage and return its readable text.

    Args:
        url: The public HTTP or HTTPS URL to inspect.
    """
    parsed = urlparse(url)
    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
        raise ValueError("only absolute HTTP(S) URLs are supported")
    if parsed.hostname in {"localhost", "127.0.0.1", "::1"}:
        raise ValueError("local network destinations are not allowed")
    response = httpx.get(
        url,
        timeout=WEB_TIMEOUT_SECONDS,
        follow_redirects=True,
        headers={"User-Agent": "GAIA-Course-Agent/1.0"},
    )
    response.raise_for_status()
    content_type = response.headers.get("content-type", "")
    if "text/html" not in content_type and "text/plain" not in content_type:
        return f"Non-text resource ({content_type}) at {response.url}"
    soup = BeautifulSoup(response.text, "html.parser")
    for element in soup(["script", "style", "noscript"]):
        element.decompose()
    text = re.sub(r"\n{3,}", "\n\n", soup.get_text("\n", strip=True))
    return f"Source URL: {response.url}\n\n{text[:MAX_PAGE_CHARACTERS]}"


@tool  # type: ignore[untyped-decorator]
def transcribe_youtube(url: str) -> str:
    """Download and transcribe the spoken audio from a public YouTube video.

    Args:
        url: A complete youtube.com or youtu.be video URL.
    """
    parsed = urlparse(url)
    hostname = (parsed.hostname or "").lower()
    if hostname not in {"youtube.com", "www.youtube.com", "youtu.be"}:
        raise ValueError("only YouTube video URLs are supported")
    with tempfile.TemporaryDirectory(prefix="gaia-youtube-") as directory:
        output = str(Path(directory) / "audio.%(ext)s")
        options = {
            "format": "bestaudio/best",
            "outtmpl": output,
            "quiet": True,
            "noplaylist": True,
        }
        with YoutubeDL(cast(Any, options)) as downloader:
            info = downloader.extract_info(url, download=True)
            downloaded = Path(downloader.prepare_filename(info))
        return transcribe_media(downloaded)