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