Blame for tiny_gdn/web_search.py

"""Execute Nemotron Agentic-v2 `web-search` the way Tercet-R stage-3 SFT saw it.

`nvidia/Nemotron-SFT-Agentic-v2` `data/search.jsonl` (packed as
`agentic_v2_search`) declares a single tool named `web-search` whose
observations are compact Tavily Search JSON:

    {"query", "follow_up_questions", "answer", "images", "results",
     "response_time", "request_id"}

Each `results[]` item is `{url, title, content, score, raw_content}`
(about ten hits, with an LLM `answer` when Tavily supplies one). The chat
template then prefixes that JSON with `<|tool_response|>`.
"""

from __future__ import annotations

import json
import os
import time
from collections.abc import Mapping, Sequence
from html import unescape
from html.parser import HTMLParser
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import parse_qs, quote, unquote, urlencode, urlparse
from urllib.request import Request, urlopen


NEMOTRON_WEB_SEARCH_TOOL: dict[str, Any] = {
    "type": "function",
    "function": {
        "name": "web-search",
        "description": "Search the web for a query.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query.",
                }
            },
            "required": ["query"],
        },
    },
}

WEB_SEARCH_TOOL_NAME = "web-search"
DEFAULT_MAX_RESULTS = 10
TAVILY_SEARCH_URL = "https://api.tavily.com/search"
DUCKDUCKGO_HTML_URL = "https://html.duckduckgo.com/html/"
WIKIPEDIA_API_URL = "https://en.wikipedia.org/w/api.php"
USER_AGENT = (
    "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
)
REQUEST_TIMEOUT_SECONDS = 30.0


def is_web_search_tool_name(name: str) -> bool:
    return name.strip().lower().replace("_", "-") == WEB_SEARCH_TOOL_NAME


def query_from_arguments(arguments: Mapping[str, Any] | None) -> str:
    if not arguments:
        return ""
    raw = arguments.get("query")
    if isinstance(raw, str):
        return raw.strip()
    if raw is None:
        return ""
    return str(raw).strip()


def dumps_search_payload(payload: Mapping[str, Any]) -> str:
    return json.dumps(canonical_search_payload(payload), ensure_ascii=False)


def canonical_search_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
    results_raw = payload.get("results")
    results: list[dict[str, Any]] = []
    if isinstance(results_raw, Sequence) and not isinstance(results_raw, (str, bytes)):
        for item in results_raw:
            if isinstance(item, Mapping):
                results.append(_canonical_result(item))
    images_raw = payload.get("images")
    images: list[Any] = []
    if isinstance(images_raw, Sequence) and not isinstance(images_raw, (str, bytes)):
        images = list(images_raw)
    return {
        "query": str(payload.get("query") or ""),
        "follow_up_questions": payload.get("follow_up_questions"),
        "answer": payload.get("answer"),
        "images": images,
        "results": results,
        "response_time": payload.get("response_time"),
        "request_id": payload.get("request_id"),
    }


def _canonical_result(item: Mapping[str, Any]) -> dict[str, Any]:
    if item.get("error") and not item.get("url") and not item.get("title"):
        return {"error": str(item["error"])}
    result: dict[str, Any] = {
        "url": str(item.get("url") or ""),
        "title": str(item.get("title") or ""),
        "content": str(item.get("content") or ""),
        "score": _as_score(item.get("score")),
        "raw_content": item.get("raw_content"),
    }
    if item.get("error"):
        result["error"] = str(item["error"])
    return result


def _as_score(value: Any) -> float | None:
    if value is None or value == "":
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def search_web(
    query: str,
    *,
    max_results: int = DEFAULT_MAX_RESULTS,
    api_key: str | None = None,
) -> str:
    """Return a Tavily-shaped JSON string for one `web-search` observation."""

    cleaned = query.strip()
    started = time.perf_counter()
    if not cleaned:
        return dumps_search_payload(
            _empty_payload(
                query,
                error="Search query was empty",
                response_time=time.perf_counter() - started,
            )
        )
    key = (api_key if api_key is not None else os.environ.get("TAVILY_API_KEY", "")).strip()
    if key:
        try:
            return dumps_search_payload(
                _search_tavily(cleaned, api_key=key, max_results=max_results)
            )
        except Exception as error:
            fallback_error = f"Tavily search failed: {error}"
            try:
                return dumps_search_payload(
                    _search_without_tavily(
                        cleaned,
                        max_results=max_results,
                        fallback_error=fallback_error,
                    )
                )
            except Exception:
                return dumps_search_payload(
                    _empty_payload(
                        cleaned,
                        error=fallback_error,
                        response_time=time.perf_counter() - started,
                    )
                )
    try:
        return dumps_search_payload(
            _search_without_tavily(cleaned, max_results=max_results)
        )
    except Exception as error:
        return dumps_search_payload(
            _empty_payload(
                cleaned,
                error=str(error),
                response_time=time.perf_counter() - started,
            )
        )


def _empty_payload(
    query: str,
    *,
    error: str,
    response_time: float | None,
) -> dict[str, Any]:
    return {
        "query": query,
        "follow_up_questions": None,
        "answer": None,
        "images": [],
        "results": [{"error": error}],
        "response_time": response_time,
        "request_id": None,
    }


def _search_tavily(
    query: str,
    *,
    api_key: str,
    max_results: int,
) -> dict[str, Any]:
    body = json.dumps(
        {
            "query": query,
            "max_results": max(1, min(int(max_results), 20)),
            "include_answer": True,
            "include_raw_content": False,
            "include_images": False,
            "search_depth": "basic",
        }
    ).encode("utf-8")
    payload = _read_json(
        TAVILY_SEARCH_URL,
        data=body,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    if not isinstance(payload, dict):
        raise RuntimeError("Tavily returned a non-object payload")
    payload.setdefault("query", query)
    return payload


def _search_without_tavily(
    query: str,
    *,
    max_results: int,
    fallback_error: str | None = None,
) -> dict[str, Any]:
    started = time.perf_counter()
    results: list[dict[str, Any]] = []
    errors: list[str] = []
    if fallback_error:
        errors.append(fallback_error)
    try:
        results.extend(_search_duckduckgo_html(query, max_results=max_results))
    except Exception as error:
        errors.append(f"DuckDuckGo search failed: {error}")
    if len(results) < max_results:
        try:
            wiki = _search_wikipedia(query, max_results=max_results)
            seen = {item["url"] for item in results}
            for item in wiki:
                if item["url"] in seen:
                    continue
                results.append(item)
                if len(results) >= max_results:
                    break
        except Exception as error:
            errors.append(f"Wikipedia search failed: {error}")
    elapsed = time.perf_counter() - started
    if not results:
        message = "; ".join(errors) if errors else "Search returned no results"
        return _empty_payload(query, error=message, response_time=elapsed)
    clipped = results[: max(1, min(int(max_results), 20))]
    for index, item in enumerate(clipped):
        if item.get("score") is None:
            item["score"] = round(max(0.05, 1.0 - (index * 0.07)), 8)
        item.setdefault("raw_content", None)
    answer = clipped[0].get("content") if clipped else None
    if isinstance(answer, str) and not answer.strip():
        answer = None
    return {
        "query": query,
        "follow_up_questions": None,
        "answer": answer,
        "images": [],
        "results": clipped,
        "response_time": elapsed,
        "request_id": None,
    }


def _search_duckduckgo_html(query: str, *, max_results: int) -> list[dict[str, Any]]:
    encoded = urlencode({"q": query, "kl": "us-en"})
    html = _read_text(
        DUCKDUCKGO_HTML_URL,
        data=encoded.encode("utf-8"),
        headers={
            "Content-Type": "application/x-www-form-urlencoded",
            "Accept": "text/html",
        },
        method="POST",
    )
    parser = _DuckDuckGoParser()
    parser.feed(html)
    parser.close()
    results: list[dict[str, Any]] = []
    seen: set[str] = set()
    for item in parser.results:
        url = item["url"]
        if not url or url in seen:
            continue
        seen.add(url)
        results.append(
            {
                "url": url,
                "title": item["title"],
                "content": item["content"],
                "score": None,
                "raw_content": None,
            }
        )
        if len(results) >= max_results:
            break
    return results


def _search_wikipedia(query: str, *, max_results: int) -> list[dict[str, Any]]:
    params = urlencode(
        {
            "action": "query",
            "list": "search",
            "srsearch": query,
            "srlimit": max(1, min(int(max_results), 10)),
            "srprop": "snippet",
            "format": "json",
            "utf8": "1",
        }
    )
    payload = _read_json(
        f"{WIKIPEDIA_API_URL}?{params}",
        headers={"Accept": "application/json"},
        method="GET",
    )
    if not isinstance(payload, dict):
        return []
    query_block = payload.get("query")
    hits = query_block.get("search") if isinstance(query_block, dict) else None
    if not isinstance(hits, list):
        return []
    results: list[dict[str, Any]] = []
    for hit in hits:
        if not isinstance(hit, dict):
            continue
        title = str(hit.get("title") or "").strip()
        if not title:
            continue
        snippet = unescape(_strip_tags(str(hit.get("snippet") or ""))).strip()
        results.append(
            {
                "url": f"https://en.wikipedia.org/wiki/{quote(title.replace(' ', '_'))}",
                "title": title,
                "content": snippet,
                "score": None,
                "raw_content": None,
            }
        )
    return results


def unwrap_ddg_url(href: str) -> str:
    raw = href.strip()
    if raw.startswith("//"):
        raw = f"https:{raw}"
    parsed = urlparse(raw)
    host = parsed.netloc.lower()
    if "duckduckgo.com" in host:
        values = parse_qs(parsed.query).get("uddg") or []
        if values:
            return unquote(values[0])
    return raw


def _strip_tags(value: str) -> str:
    return _TagStripper.strip(value)


class _TagStripper(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.chunks: list[str] = []

    def handle_data(self, data: str) -> None:
        self.chunks.append(data)

    @classmethod
    def strip(cls, value: str) -> str:
        parser = cls()
        parser.feed(value)
        parser.close()
        return "".join(parser.chunks)


class _DuckDuckGoParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.results: list[dict[str, str]] = []
        self._in_title = False
        self._in_snippet = False
        self._href = ""
        self._buffer: list[str] = []
        self._pending: dict[str, str] | None = None

    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
        classes = {item for item in (dict(attrs).get("class") or "").split() if item}
        if tag == "a" and "result__a" in classes:
            self._flush_pending()
            self._in_title = True
            self._href = dict(attrs).get("href") or ""
            self._buffer = []
            return
        if tag in {"a", "div"} and "result__snippet" in classes:
            self._in_snippet = True
            self._buffer = []

    def handle_endtag(self, tag: str) -> None:
        if tag == "a" and self._in_title:
            title = unescape("".join(self._buffer)).strip()
            url = unwrap_ddg_url(self._href)
            self._pending = {"title": title, "url": url, "content": ""}
            self._in_title = False
            self._buffer = []
            return
        if self._in_snippet and tag in {"a", "div"}:
            snippet = unescape("".join(self._buffer)).strip()
            if self._pending is not None:
                self._pending["content"] = snippet
                self.results.append(self._pending)
                self._pending = None
            self._in_snippet = False
            self._buffer = []

    def handle_data(self, data: str) -> None:
        if self._in_title or self._in_snippet:
            self._buffer.append(data)

    def close(self) -> None:
        self._flush_pending()
        super().close()

    def _flush_pending(self) -> None:
        if self._pending is None:
            return
        self.results.append(self._pending)
        self._pending = None


def _read_json(
    url: str,
    *,
    data: bytes | None = None,
    headers: Mapping[str, str],
    method: str,
) -> Any:
    text = _read_text(url, data=data, headers=headers, method=method)
    return json.loads(text)


def _read_text(
    url: str,
    *,
    data: bytes | None = None,
    headers: Mapping[str, str],
    method: str,
) -> str:
    merged = {"User-Agent": USER_AGENT, **dict(headers)}
    request = Request(url, data=data, headers=merged, method=method)
    try:
        with urlopen(request, timeout=REQUEST_TIMEOUT_SECONDS) as response:
            raw = response.read()
            charset = response.headers.get_content_charset() or "utf-8"
    except HTTPError as error:
        detail = error.read().decode("utf-8", errors="replace")
        raise RuntimeError(
            f"{error.code} {error.reason} from {url}: {detail[:500]}"
        ) from error
    except URLError as error:
        raise RuntimeError(f"Could not reach {url}: {error.reason}") from error
    return raw.decode(charset, errors="replace")