jasonwitty
6b50b23467
Furo renders its own page-local table of contents and reports an error when a page adds one with contents::. No other page in the tree uses that directive; these two now follow suit. The plugin is named "AI Summary" where the name of the plugin is meant, and "AI summary" where the generated text itself is meant. That also fixes the entry in the plugin navigation, which takes its label from the page title. Two references to the configuration page rendered as its title, the bare YAML key "ai_summary:", which reads as a typo in a sentence. They now carry explicit link text. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
329 lines
14 KiB
Python
329 lines
14 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
"""Implementation of the AI Summary plugin, which shows a generated answer above
|
|
the search results. The answer comes from an LLM server that implements the
|
|
`OpenAI chat completions API`_ (Ollama, vLLM, llama.cpp, LM Studio, Hugging Face
|
|
TGI, ...) and that the administrator runs.
|
|
|
|
- :ref:`ai_summary plugin` describes the design and the request flow.
|
|
- :ref:`settings ai_summary` describes how to configure it.
|
|
|
|
This module holds the plugin itself and the ``/ai_summary`` endpoint
|
|
(:py:obj:`ai_summary_view`, registered in :py:obj:`searx.webapp`). The endpoint
|
|
streams the answer to the browser, so that the result page is never delayed by
|
|
the LLM; :py:obj:`SXNGPlugin.post_search` only adds an empty
|
|
:py:obj:`searx.result_types.AiSummary` placeholder for the client to fill.
|
|
|
|
Settings of the ``ai_summary:`` section are defined in
|
|
:py:obj:`searx.ai_summary.SettingsAISummary`.
|
|
|
|
.. _OpenAI chat completions API: https://platform.openai.com/docs/api-reference/chat
|
|
"""
|
|
|
|
import typing as t
|
|
|
|
import json
|
|
import logging
|
|
import re
|
|
import time
|
|
from urllib.parse import urlparse
|
|
|
|
import flask
|
|
import httpx
|
|
from flask_babel import gettext
|
|
|
|
from searx import get_setting
|
|
from searx.ai_summary import SettingsAISummary, build_chat_messages
|
|
from searx.extended_types import sxng_request
|
|
from searx.result_types import EngineResults
|
|
import searx.ai_summary
|
|
|
|
from . import Plugin, PluginInfo
|
|
|
|
if t.TYPE_CHECKING:
|
|
from searx.search import SearchWithPlugins
|
|
from searx.extended_types import SXNG_Request
|
|
from . import PluginCfg
|
|
|
|
VALID_ROLES = ("user", "assistant")
|
|
MODEL_NAME_REGEXP = re.compile(r"[A-Za-z0-9._:/-]{1,128}")
|
|
|
|
|
|
def _get_client(base_url: str, cfg: SettingsAISummary, api_key: str = "") -> httpx.Client:
|
|
"""HTTP client for one request to the LLM server at ``base_url``. The
|
|
``api_key`` (if any) is sent in an ``Authorization: Bearer`` header, see
|
|
:py:obj:`_server_api_key`."""
|
|
# the OpenAI API paths are prefixed with /v1, unless the base URL already
|
|
# points into an API prefix
|
|
base_url = base_url.rstrip("/")
|
|
if not base_url.endswith("/v1"):
|
|
base_url += "/v1"
|
|
return httpx.Client(
|
|
base_url=base_url,
|
|
headers={"Authorization": f"Bearer {api_key}"} if api_key else None,
|
|
timeout=httpx.Timeout(connect=cfg.connect_timeout, read=cfg.read_timeout, write=10.0, pool=10.0),
|
|
)
|
|
|
|
|
|
def _valid_server(url: str) -> bool:
|
|
try:
|
|
parsed = urlparse(url)
|
|
except ValueError:
|
|
return False
|
|
return parsed.scheme in ("http", "https") and bool(parsed.netloc) and len(url) <= 256
|
|
|
|
|
|
def _server_id(url: str) -> tuple[str, str, int, str] | None:
|
|
"""Identity of an LLM server URL (scheme, host, port, path) for comparing
|
|
two URLs, or ``None`` if the URL is unusable. The ``/v1`` API prefix is
|
|
not part of the identity, :py:obj:`_get_client` appends it when missing."""
|
|
try:
|
|
parsed = urlparse(url)
|
|
except ValueError:
|
|
return None
|
|
# .hostname (not .netloc) drops the userinfo, so that a server URL like
|
|
# http://llm.example.org@untrusted.example.org/ is identified by the host
|
|
# the request is actually sent to (untrusted.example.org)
|
|
if parsed.scheme not in ("http", "https") or not parsed.hostname:
|
|
return None
|
|
port = parsed.port or (443 if parsed.scheme == "https" else 80)
|
|
path = parsed.path.rstrip("/")
|
|
if path.endswith("/v1"):
|
|
path = path[: -len("/v1")].rstrip("/")
|
|
return (parsed.scheme, parsed.hostname.lower(), port, path)
|
|
|
|
|
|
def _server_api_key(cfg: SettingsAISummary, server: str, user_api_key: str = "") -> str:
|
|
"""The API key to send to ``server``:
|
|
|
|
- the administrator's :py:obj:`cfg.api_key
|
|
<searx.ai_summary.SettingsAISummary.api_key>` if ``server`` *is* the
|
|
administrator's server (:py:obj:`cfg.base_url
|
|
<searx.ai_summary.SettingsAISummary.base_url>`),
|
|
- otherwise the user's own ``ai_summary_api_key`` preference, which belongs
|
|
to the server in the user's own ``ai_summary_server`` preference.
|
|
|
|
The administrator's key is never sent to a server a user configured --
|
|
that would hand every user of the instance a way to capture it."""
|
|
server_id = _server_id(server)
|
|
if server_id is not None and server_id == _server_id(cfg.base_url):
|
|
return cfg.api_key
|
|
return user_api_key
|
|
|
|
|
|
def _user_server(request: "SXNG_Request", cfg: SettingsAISummary) -> str:
|
|
"""The LLM server URL for this request: the user's ``ai_summary_server``
|
|
preference, or the administrator's default."""
|
|
server = str(request.preferences.get_value("ai_summary_server") or "").strip()
|
|
# Credentials are stripped from a user's server URL: httpx turns them into
|
|
# an Authorization header, and a user should not be able to make SearXNG
|
|
# send a header of their choosing to a host of their choosing. An
|
|
# administrator can still use credentials in the configured base_url (e.g.
|
|
# an LLM server behind basic auth).
|
|
if server and "@" in urlparse(server).netloc:
|
|
server = ""
|
|
return server or cfg.base_url
|
|
|
|
|
|
class SXNGPlugin(Plugin):
|
|
"""Plugin that adds the AI summary placeholder to the result page, the
|
|
``/ai_summary`` endpoint itself is registered in :py:obj:`searx.webapp`."""
|
|
|
|
id = "ai_summary"
|
|
|
|
def __init__(self, plg_cfg: "PluginCfg"):
|
|
super().__init__(plg_cfg)
|
|
|
|
self.info = PluginInfo(
|
|
id=self.id,
|
|
name=gettext("AI Summary"),
|
|
description=gettext(
|
|
"Show an AI generated summary of the search query on top of the"
|
|
" result page (uses a local LLM server, see the settings below)."
|
|
),
|
|
preference_section="ai",
|
|
)
|
|
|
|
def init(self, app: "flask.Flask") -> bool:
|
|
cfg: SettingsAISummary = get_setting("ai_summary")
|
|
|
|
if cfg.base_url:
|
|
searx.ai_summary.MODELS = list(cfg.models) or self._probe_models(cfg)
|
|
if not cfg.model and searx.ai_summary.MODELS:
|
|
cfg.model = searx.ai_summary.MODELS[0]
|
|
if cfg.model and cfg.model not in searx.ai_summary.MODELS:
|
|
searx.ai_summary.MODELS.insert(0, cfg.model)
|
|
|
|
return True
|
|
|
|
def _probe_models(self, cfg: SettingsAISummary) -> list[str]:
|
|
"""Request the list of models from the LLM server (``GET
|
|
/v1/models``). The server might not be up when SearXNG starts, a
|
|
failing probe only leaves the model suggestion list empty."""
|
|
try:
|
|
with _get_client(cfg.base_url, cfg, cfg.api_key) as client:
|
|
resp = client.get("/models")
|
|
resp.raise_for_status()
|
|
models = [model["id"] for model in resp.json().get("data", [])]
|
|
except (httpx.HTTPError, ValueError, KeyError) as exc:
|
|
self.log.warning("can't request model list from %s: %s", cfg.base_url, exc)
|
|
models = []
|
|
return models or ([cfg.model] if cfg.model else [])
|
|
|
|
def post_search(self, request: "SXNG_Request", search: "SearchWithPlugins") -> EngineResults | None:
|
|
results = EngineResults()
|
|
sq = search.search_query
|
|
cfg: SettingsAISummary = get_setting("ai_summary")
|
|
|
|
skip = (
|
|
sq.pageno > 1
|
|
# post_search is also called for the json, csv and rss formats,
|
|
# the placeholder is only useful on the HTML result page
|
|
or request.form.get("format", "html") != "html"
|
|
or "general" not in sq.categories
|
|
# an infobox (e.g. wikipedia / wikidata) or an instant answer
|
|
# (e.g. ddg definitions) most likely already answers the query
|
|
or bool(search.result_container.infoboxes)
|
|
or bool(search.result_container.answers)
|
|
or not sq.query.strip()
|
|
# without an LLM server (user preference or admin default)
|
|
# there is nothing to show
|
|
or not _user_server(request, cfg)
|
|
)
|
|
if skip:
|
|
return None
|
|
|
|
grounding = bool(request.preferences.get_value("ai_summary_grounding"))
|
|
results.add(results.types.AiSummary(query=sq.query, grounding=grounding))
|
|
return results
|
|
|
|
|
|
def _bad_request(msg: str) -> flask.Response:
|
|
return flask.Response(json.dumps({"error": msg}), status=400, mimetype="application/json")
|
|
|
|
|
|
def _validate_messages(messages: t.Any, cfg: SettingsAISummary) -> list[dict[str, str]]:
|
|
if not isinstance(messages, list) or not messages or len(messages) > cfg.max_history_messages:
|
|
raise ValueError("invalid messages")
|
|
for msg in messages:
|
|
if not isinstance(msg, dict) or msg.keys() != {"role", "content"}:
|
|
raise ValueError("invalid message")
|
|
if msg["role"] not in VALID_ROLES or not isinstance(msg["content"], str):
|
|
raise ValueError("invalid message")
|
|
if not msg["content"].strip() or len(msg["content"]) > cfg.max_message_length:
|
|
raise ValueError("invalid message")
|
|
if messages[-1]["role"] != "user":
|
|
raise ValueError("last message is not a user message")
|
|
return messages
|
|
|
|
|
|
def _validate_context(context: t.Any, cfg: SettingsAISummary) -> list[dict[str, str]]:
|
|
if not isinstance(context, list) or len(context) > cfg.max_context_items:
|
|
raise ValueError("invalid context")
|
|
for item in context:
|
|
if not isinstance(item, dict) or not item.keys() <= {"title", "url", "snippet"}:
|
|
raise ValueError("invalid context item")
|
|
for val in item.values():
|
|
if not isinstance(val, str) or len(val) > cfg.max_message_length:
|
|
raise ValueError("invalid context item")
|
|
return context
|
|
|
|
|
|
def _validate_payload(payload: t.Any, cfg: SettingsAISummary) -> tuple[list[dict[str, str]], list[dict[str, str]]]:
|
|
"""Validate the request body of the ``/ai_summary`` endpoint and return
|
|
the ``messages`` and ``context`` lists. Raises a :py:obj:`ValueError` for
|
|
any malformed payload."""
|
|
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("payload is not an object")
|
|
return _validate_messages(payload.get("messages"), cfg), _validate_context(payload.get("context", []), cfg)
|
|
|
|
|
|
def ai_summary_view() -> flask.Response:
|
|
"""Stream an AI generated answer for the messages in the request body,
|
|
response is NDJSON: ``{"delta": ..}`` lines followed by one final
|
|
``{"done": true, ..}`` line."""
|
|
|
|
cfg: SettingsAISummary = get_setting("ai_summary")
|
|
|
|
if SXNGPlugin.id not in sxng_request.user_plugins:
|
|
return flask.Response(json.dumps({"error": "plugin is not enabled"}), status=403, mimetype="application/json")
|
|
|
|
try:
|
|
messages, context = _validate_payload(sxng_request.get_json(force=True, silent=True), cfg)
|
|
except ValueError as exc:
|
|
return _bad_request(str(exc))
|
|
|
|
server = _user_server(sxng_request, cfg)
|
|
if not _valid_server(server):
|
|
return _bad_request("no valid LLM server configured")
|
|
|
|
model = str(sxng_request.preferences.get_value("ai_summary_model") or "").strip() or cfg.model
|
|
if not MODEL_NAME_REGEXP.fullmatch(model):
|
|
return _bad_request("no valid model configured")
|
|
|
|
chat_payload = {
|
|
"model": model,
|
|
"messages": build_chat_messages(cfg, messages, context),
|
|
"stream": True,
|
|
}
|
|
|
|
# open the upstream connection before streaming, a connection error is
|
|
# reported as HTTP 502 instead of a line in an already started stream
|
|
user_api_key = str(sxng_request.preferences.get_value("ai_summary_api_key") or "").strip()
|
|
client = _get_client(server, cfg, _server_api_key(cfg, server, user_api_key))
|
|
stream_ctx = client.stream("POST", "/chat/completions", json=chat_payload)
|
|
upstream = None
|
|
try:
|
|
upstream = stream_ctx.__enter__() # pylint: disable=unnecessary-dunder-call
|
|
if upstream.status_code != 200:
|
|
stream_ctx.__exit__(None, None, None)
|
|
upstream = None
|
|
except httpx.HTTPError:
|
|
upstream = None
|
|
if upstream is None:
|
|
client.close()
|
|
return flask.Response(json.dumps({"error": "upstream error"}), status=502, mimetype="application/json")
|
|
|
|
# from here on nothing must be read from the request context, the
|
|
# generator runs after the request context has been torn down
|
|
log = logging.getLogger("searx.plugins.ai_summary")
|
|
|
|
def ndjson(obj: dict[str, t.Any]) -> bytes:
|
|
# the generator bypasses flask's response encoding (direct_passthrough)
|
|
return (json.dumps(obj) + "\n").encode()
|
|
|
|
def generate():
|
|
start = time.monotonic()
|
|
try:
|
|
# the upstream is a SSE stream: "data: {..}" lines, terminated by
|
|
# a "data: [DONE]" line
|
|
for line in upstream.iter_lines():
|
|
if time.monotonic() - start > cfg.stream_timeout:
|
|
yield ndjson({"done": True, "error": "timeout"})
|
|
return
|
|
line = line.strip()
|
|
if not line or line.startswith(":") or not line.startswith("data:"):
|
|
continue
|
|
payload = line[len("data:") :].strip()
|
|
if payload == "[DONE]":
|
|
break
|
|
data = json.loads(payload)
|
|
choices = data.get("choices") or [{}]
|
|
delta = choices[0].get("delta", {}).get("content") or ""
|
|
if delta:
|
|
yield ndjson({"delta": delta})
|
|
yield ndjson({"done": True, "model": model})
|
|
except (httpx.HTTPError, ValueError) as exc:
|
|
log.warning("error while streaming from the LLM server: %s", exc)
|
|
yield ndjson({"done": True, "error": "upstream error"})
|
|
finally:
|
|
stream_ctx.__exit__(None, None, None)
|
|
client.close()
|
|
|
|
return flask.Response(
|
|
generate(),
|
|
mimetype="application/x-ndjson",
|
|
headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"},
|
|
direct_passthrough=True,
|
|
)
|