owui-tools/vm/vmd/websearch.py

"""Search and page-reading for the sub-agent, done on the host.

WHY THIS IS NOT A SECOND SEARCH CLIENT

The api container already has `web_search` and `fetch_page`: a SearXNG
query with the engines this house trusts, a result formatter tuned for a
model to read, a cache, a rate limit, and - for fetch_page - readability
extraction, PDF handling and an SSRF guard that refuses private
addresses. Writing a second one here would mean two search behaviours to
keep in step, and the copy would be the one nobody remembers to fix.

So this calls the same endpoint the main model's tool calls, through the
same gateway, with the same bearer token. One implementation, two
callers.

WHY THE AGENT CANNOT DO THIS ITSELF

SearXNG lives inside netguard's network namespace, on loopback, with a
default-deny filter around it. The VM cannot reach it (that is the whole
point of the sandbox) and neither can the agent container, whose
allowlist is vmgw alone. vmd is the privileged half and already holds the
Gitea token, so it is where a credentialled call belongs.

WHAT COMES BACK IS UNTRUSTED TEXT

Search results and page text are written by strangers. They reach the
agent's context and, through its summary, the main model's. The existing
defences apply unchanged: vmd assembles the tool result header itself,
the agent's prose is delimited and labelled untrusted, and an unfreeze
request can only be produced by vmd with a nonce. Nothing here is
special-cased; it is simply the same exposure the main model already has
when it searches.
"""

from __future__ import annotations

import json
import os
import urllib.error
import urllib.request

# The gateway, not the api directly: it is the front door every other
# caller uses, so a change to auth or routing applies here too. Pinned in
# docker-compose.yml at 172.22.0.10, on a subnet that is also pinned.
# The tool routes live under /tools - the same prefix Open WebUI calls.
TOOLS_URL = os.environ.get("OWUITOOLS_TOOLS_URL", "http://172.22.0.10:8000").rstrip("/")
PREFIX = "/tools"
AUTH_TOKEN = os.environ.get("OWUITOOLS_AUTH_TOKEN", "")

TIMEOUT = 45.0


class SearchError(RuntimeError):
    """Phrased for the model to read."""


def configured() -> bool:
    return bool(AUTH_TOKEN)


def _post(path: str, payload: dict) -> str:
    if not AUTH_TOKEN:
        raise SearchError(
            "web search is not configured on this host, so there is nothing to search with. "
            "Work from what is installed in the machine instead."
        )
    data = json.dumps(payload).encode()
    request = urllib.request.Request(f"{TOOLS_URL}{PREFIX}{path}", data=data, method="POST")
    request.add_header("Authorization", f"Bearer {AUTH_TOKEN}")
    request.add_header("Content-Type", "application/json")
    request.add_header("Accept", "application/json")
    # Same reasoning as gitea.py: an inherited http_proxy would send a
    # credentialled request somewhere nobody chose.
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    try:
        with opener.open(request, timeout=TIMEOUT) as response:
            body = response.read().decode("utf-8", "replace")
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:300]
        if exc.code == 401:
            raise SearchError("the search service rejected this host's credentials.") from exc
        raise SearchError(f"the search service answered {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise SearchError(f"could not reach the search service: {exc.reason}") from exc

    # These endpoints declare `response_model=str`, so the body is a JSON
    # string. Anything else means the shape changed and guessing would
    # put a stray pair of quotes in front of the model.
    try:
        decoded = json.loads(body)
    except ValueError:
        return body
    return decoded if isinstance(decoded, str) else json.dumps(decoded)


# THE RESEARCH SUITE THE MAIN MODEL HAS, MINUS WHAT THE MACHINE DOES BETTER.
#
# Offered to every session. Reading a page as text costs one call where
# reading it off the screen costs a screenshot per scroll, and a session
# with no cheap way to look something up reasons it out instead.
#
# What is deliberately NOT here, and the rule behind it: anything the
# machine can answer about itself, in itself. `dns_lookup`, `http_headers`,
# `inspect_tls` and `lookup_ip_or_asn` are all one curl, dig or
# `openssl s_client` away inside a machine that has internet - and running
# them THERE is the entire point of having a machine. Answering from the
# host would also answer from the host's network position, which is a
# different question than the one being asked.
#
# Also not here: `matlab_run`, which has its own sandbox and whose work
# python in the machine can do; and `vm_agent`, for reasons that should
# not need stating.
RESEARCH_TOOLS = {
    "search_stackexchange": "/search_stackexchange",
    "search_papers": "/search_papers",
    "wikipedia_lookup": "/wikipedia_lookup",
    "read_man_page": "/read_man_page",
    "lookup_package": "/lookup_package",
    "lookup_rfc": "/lookup_rfc",
    "lookup_cve": "/lookup_cve",
    "find_datasheet": "/find_datasheet",
}


def research(tool: str, arguments: dict) -> str:
    """One of the main model's lookup tools, called the same way it calls it."""
    path = RESEARCH_TOOLS.get(tool)
    if path is None:
        raise SearchError(f"'{tool}' is not a tool this session has.")
    if not isinstance(arguments, dict):
        raise SearchError("that tool was called with something that is not a set of arguments.")
    # The endpoint validates its own arguments and says so in prose when
    # they are wrong - it is the same 422 the main model gets, and the
    # same wording. Nothing is re-validated here, because a second
    # opinion that disagrees is worse than no second opinion.
    return _post(path, arguments).replace("fetch_page", "read_page")


def search(query: str, max_results: int = 6, site: str = "",
           time_range: str = "") -> str:
    query = (query or "").strip()
    if not query:
        raise SearchError("no search terms were given.")
    if len(query) > 400:
        raise SearchError("that search is too long; say it in a sentence or less.")
    payload: dict = {"query": query, "max_results": max(1, min(int(max_results or 6), 15))}
    if site:
        payload["site"] = site
    if time_range:
        payload["time_range"] = time_range
    # The endpoint's own advice names the MAIN model's tool - "read the
    # most promising result with fetch_page". This agent has no such
    # tool; its name for the same thing is read_page. Telling a model to
    # call something it was not given is how a turn gets spent on a
    # refusal, so the one word is translated on the way through.
    return _post("/web_search", payload).replace("fetch_page", "read_page")


def read(url: str, max_chars: int = 6000, offset: int = 0, search_term: str = "") -> str:
    url = (url or "").strip()
    if not url:
        raise SearchError("no URL was given.")
    payload: dict = {"url": url, "max_chars": max(500, min(int(max_chars or 6000), 20000))}
    if offset:
        payload["offset"] = int(offset)
    if search_term:
        payload["search"] = search_term
    # Same translation as above: the continuation hint says "call
    # fetch_page again with offset=...", and this agent's tool is
    # read_page. The offsets and the advice are right; only the name is
    # from the other caller's vocabulary.
    return _post("/fetch_page", payload).replace("fetch_page", "read_page")