Initial Commit

vm/vmd/websearch.py

@@ -0,0 +1,172 @@
+"""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")