[feat] plugin: AI summary of search results via local Ollama server

Add an optional, disabled-by-default plugin that shows an AI generated
summary at the top of the result page, generated by a (local) Ollama
server:

- async: the result page is never delayed; a client plugin streams the
  answer (NDJSON over a new /ai_summary endpoint) into a placeholder
  answer with a typing indicator, collapsed behind a More button, with
  an inline follow-up chat
- trigger: first page of general searches only, skipped when an infobox
  or instant answer already answers the query
- grounding (per-user preference): send the top result snippets as
  context, the model answers from them instead of its own knowledge
- configuration: new AI Summary preferences tab (server URL, model,
  grounding) with instance defaults in a new ai_summary: settings
  section; all three preferences can be locked for public instances

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
jasonwitty
2026-07-28 10:47:18 -07:00
parent b060c780d0
commit b47fd08cb7
22 changed files with 1227 additions and 3 deletions
+218
View File
@@ -0,0 +1,218 @@
// SPDX-License-Identifier: AGPL-3.0-or-later
import { Plugin } from "../Plugin.ts";
import { settings } from "../toolkit.ts";
import { assertElement } from "../util/assertElement.ts";
type Message = { role: "user" | "assistant"; content: string };
type ContextItem = { title: string; url: string; snippet: string };
type StreamLine = { delta?: string; done?: boolean; model?: string; error?: string };
// keep in sync with the server side defaults (searx/ai_summary.py)
const MAX_CONTEXT_ITEMS = 5;
const MAX_HISTORY_MESSAGES = 12;
const MAX_TITLE_LENGTH = 300;
const MAX_SNIPPET_LENGTH = 1000;
/**
* Fills the AI summary placeholder (rendered by the ai_summary plugin of the
* server) by streaming an answer from the /ai_summary endpoint, with an
* expand button and a follow-up question chat.
*/
export default class AiSummary extends Plugin {
private readonly messages: Message[] = [];
private context: ContextItem[] = [];
private controller: AbortController | undefined;
public constructor() {
super("ai_summary");
}
protected async run(): Promise<void> {
const box = document.getElementById("ai_summary");
if (!box) return;
try {
const { query } = box.dataset;
if (!query) return;
if (box.dataset.grounding === "1") {
this.context = AiSummary.collectContext();
}
this.wireControls(box);
this.messages.push({ role: "user", content: query });
await this.exchange(box);
} catch (error) {
// never fail silently, always leave a message in the summary box
AiSummary.showError(box, error);
}
}
protected async post(): Promise<void> {
// noop
}
/**
* Scrape the top search results from the DOM as grounding context.
*/
private static collectContext(): ContextItem[] {
const items: ContextItem[] = [];
for (const article of document.querySelectorAll<HTMLElement>("#urls article.result")) {
if (items.length >= MAX_CONTEXT_ITEMS) break;
const link = article.querySelector<HTMLAnchorElement>("h3 a");
if (!link) continue;
items.push({
title: (link.textContent ?? "").trim().slice(0, MAX_TITLE_LENGTH),
url: link.href,
snippet: (article.querySelector(".content")?.textContent ?? "").trim().slice(0, MAX_SNIPPET_LENGTH)
});
}
return items;
}
private wireControls(box: HTMLElement): void {
const body = box.querySelector<HTMLElement>(".ai-summary-body");
const moreButton = box.querySelector<HTMLElement>(".ai-summary-more");
const followupForm = box.querySelector<HTMLFormElement>(".ai-summary-followup");
assertElement(body);
assertElement(moreButton);
assertElement(followupForm);
moreButton.addEventListener("click", () => {
const collapsed = body.classList.toggle("collapsed");
moreButton.textContent = collapsed
? (moreButton.dataset.btnTextCollapsed ?? "")
: (moreButton.dataset.btnTextNotCollapsed ?? "");
followupForm.classList.toggle("invisible", collapsed);
});
followupForm.addEventListener("submit", (event: Event) => {
event.preventDefault();
const input = followupForm.querySelector<HTMLInputElement>("input");
assertElement(input);
const question = input.value.trim();
if (!question || this.controller) return;
input.value = "";
this.messages.push({ role: "user", content: question });
// the server rejects too long histories, drop the oldest messages
this.messages.splice(0, this.messages.length - (MAX_HISTORY_MESSAGES - 1));
const questionElement = Object.assign(document.createElement("p"), {
textContent: question,
className: "ai-summary-question"
});
box.querySelector(".ai-summary-answers")?.append(questionElement);
void this.exchange(box);
});
}
/**
* Send the message history to /ai_summary and stream the answer into a new
* block in the answer area.
*/
private async exchange(box: HTMLElement): Promise<void> {
const answers = box.querySelector<HTMLElement>(".ai-summary-answers");
assertElement(answers);
const block = Object.assign(document.createElement("p"), {
className: "ai-summary-content typing"
});
answers.append(block);
const controller = new AbortController();
this.controller = controller;
let text = "";
const handleLine = (line: string): void => {
if (!line.trim()) return;
const data = JSON.parse(line) as StreamLine;
if (data.error) {
throw new Error(data.error);
}
if (data.delta) {
text += data.delta;
block.textContent = text;
this.updateMoreButton(box);
}
if (data.done && data.model) {
const modelElement = box.querySelector<HTMLElement>(".ai-summary-model");
if (modelElement) modelElement.textContent = data.model;
}
};
try {
const res = await fetch("./ai_summary", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ messages: this.messages, context: this.context }),
signal: controller.signal
});
if (!res.ok) {
throw new Error(res.statusText);
}
if (res.body) {
const reader = res.body.getReader();
const decoder = new TextDecoder();
let buffer = "";
for (;;) {
// biome-ignore lint/performance/noAwaitInLoops: chunks of a stream are read sequentially
const { done, value } = await reader.read();
if (done) break;
buffer += decoder.decode(value, { stream: true });
const lines = buffer.split("\n");
buffer = lines.pop() ?? "";
for (const line of lines) handleLine(line);
}
handleLine(buffer);
} else {
// some webviews (e.g. Brave iOS) wrap fetch and don't expose a
// streaming body: no typing effect, render the answer in one go
const body = await res.text();
for (const line of body.split("\n")) handleLine(line);
}
this.messages.push({ role: "assistant", content: text });
} catch (error) {
AiSummary.showError(box, error);
} finally {
block.classList.remove("typing");
this.controller = undefined;
this.updateMoreButton(box);
}
}
private static showError(box: HTMLElement, error: unknown): void {
console.error("Error loading AI summary:", error);
const errorElement = Object.assign(document.createElement("div"), {
textContent: settings.translations?.error_loading_ai_summary ?? "Error loading the AI summary",
className: "dialog-error"
});
errorElement.setAttribute("role", "alert");
(box.querySelector<HTMLElement>(".ai-summary-answers") ?? box).append(errorElement);
}
/**
* Show the expand button as soon as the (collapsed) body overflows.
*/
private updateMoreButton(box: HTMLElement): void {
const body = box.querySelector<HTMLElement>(".ai-summary-body");
const moreButton = box.querySelector<HTMLElement>(".ai-summary-more");
if (!(body && moreButton)) return;
if (!body.classList.contains("collapsed") || body.scrollHeight > body.clientHeight + 4) {
moreButton.classList.remove("invisible");
}
}
}
+7
View File
@@ -34,6 +34,13 @@ ready(() => {
where: [Endpoints.results] where: [Endpoints.results]
}); });
} }
if (settings.plugins?.includes("ai_summary")) {
load(() => import("./plugin/AiSummary.ts").then(({ default: Plugin }) => new Plugin()), {
on: "endpoint",
where: [Endpoints.results]
});
}
}); });
ready( ready(
+90
View File
@@ -0,0 +1,90 @@
// SPDX-License-Identifier: AGPL-3.0-or-later
@keyframes ai-summary-caret {
50% {
opacity: 0;
}
}
.ai-summary {
display: flex;
flex-direction: column;
gap: 0.5rem;
.ai-summary-header {
display: flex;
align-items: baseline;
gap: 0.5rem;
.ai-summary-title {
font-size: 0.9em;
font-weight: bold;
text-transform: uppercase;
letter-spacing: 0.05em;
}
.ai-summary-model {
font-size: 0.8em;
opacity: 0.7;
}
}
.ai-summary-body.collapsed {
max-height: 8em;
overflow: hidden;
mask-image: linear-gradient(to bottom, #000 60%, transparent 100%);
}
.ai-summary-answers {
display: flex;
flex-direction: column;
gap: 0.5rem;
}
.ai-summary-content {
margin: 0;
white-space: pre-line;
overflow-wrap: anywhere;
&.typing::after {
content: "▍";
animation: ai-summary-caret 1s step-end infinite;
}
}
.ai-summary-question {
margin: 0;
font-weight: bold;
}
.ai-summary-more {
align-self: center;
border: none;
.show-content-button();
}
.ai-summary-followup {
display: flex;
gap: 0.5rem;
input {
flex: 1;
padding: 5px 10px;
border: 1px solid var(--color-sidebar-border);
background: var(--color-answer-background);
color: var(--color-answer-font);
.rounded-corners-tiny;
}
button {
border: none;
.show-content-button();
}
}
.ai-summary-disclaimer {
margin: 0;
font-size: 0.8em;
opacity: 0.7;
}
}
+1
View File
@@ -13,6 +13,7 @@
@import "stats.less"; @import "stats.less";
@import "result_templates.less"; @import "result_templates.less";
@import "weather.less"; @import "weather.less";
@import "ai_summary.less";
// for index.html template // for index.html template
@import "index.less"; @import "index.less";
+1
View File
@@ -25,3 +25,4 @@ Settings
settings_outgoing settings_outgoing
settings_categories_as_tabs settings_categories_as_tabs
settings_plugins settings_plugins
settings_ai_summary
@@ -0,0 +1,36 @@
.. _settings ai_summary:
===============
``ai_summary:``
===============
Default configuration of the :ref:`AI summary plugin <ai_summary plugin>`.
Users configure the Ollama server URL and the model in the *AI Summary* tab of
their preferences; the values below only act as instance wide defaults.
.. code:: yaml
ai_summary:
base_url: "http://127.0.0.1:11434"
model: "llama3.2:3b"
.. attention::
A user configurable server URL allows any user of the instance to make the
SearXNG server send requests to a URL of their choice (`SSRF`_), and each
summary is real LLM work. This plugin is intended for private instances --
on a public instance lock the related preferences (:ref:`settings
preferences`):
.. code:: yaml
preferences:
lock:
- ai_summary_server
- ai_summary_model
- ai_summary_grounding
.. _SSRF: https://owasp.org/www-community/attacks/Server_Side_Request_Forgery
.. autoclass:: searx.ai_summary.SettingsAISummary
:members:
+8
View File
@@ -0,0 +1,8 @@
.. _ai_summary plugin:
==========
AI summary
==========
.. automodule:: searx.plugins.ai_summary
:members:
+1
View File
@@ -7,6 +7,7 @@ Built-in Plugins
.. toctree:: .. toctree::
:maxdepth: 1 :maxdepth: 1
ai_summary
calculator calculator
hash_plugin hash_plugin
hostnames hostnames
+3
View File
@@ -33,6 +33,9 @@ class SettingsPref(msgspec.Struct, kw_only=True, forbid_unknown_fields=True):
"theme", "theme",
"results_on_new_tab", "results_on_new_tab",
"doi_resolver", "doi_resolver",
"ai_summary_server",
"ai_summary_model",
"ai_summary_grounding",
"simple_style", "simple_style",
"center_alignment", "center_alignment",
"query_in_title", "query_in_title",
+121
View File
@@ -0,0 +1,121 @@
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Implementations needed for the AI summary plugin
(:py:obj:`searx.plugins.ai_summary`)."""
# pylint: disable=too-few-public-methods
# Struct fields aren't discovered in Python 3.14
# - https://github.com/searxng/searxng/issues/5284
from __future__ import annotations
__all__ = ["SettingsAISummary", "MODELS", "model_choices", "build_ollama_messages"]
import msgspec
DEFAULT_SYSTEM_PROMPT = (
"You are a search assistant. Answer the user's search query concisely in"
" a few short paragraphs of plain text. If you are unsure or don't know"
" the answer, say so."
)
DEFAULT_SYSTEM_PROMPT_GROUNDED = (
"You are a search assistant. Answer the user's search query concisely in"
" a few short paragraphs of plain text, using the following search results"
" as context when they are relevant. If you are unsure or don't know the"
" answer, say so.\n\nSearch results:\n\n{context}"
)
MODELS: list[str] = []
"""List of model names a user can select from. Populated once at
application setup by :py:obj:`searx.plugins.ai_summary.SXNGPlugin.init`."""
class SettingsAISummary(msgspec.Struct, kw_only=True, forbid_unknown_fields=True):
"""Options for configuring the AI summary plugin.
.. code:: yaml
ai_summary:
base_url: "http://127.0.0.1:11434"
model: "llama3.2:3b"
models:
- "llama3.2:3b"
- "gemma3:4b"
"""
base_url: str = ""
"""Default base URL of the Ollama server (e.g. ``http://127.0.0.1:11434``).
Users can set their own server URL in the preferences
(``ai_summary_server``) unless that preference is locked."""
model: str = ""
"""Name of the default model (e.g. ``llama3.2:3b``). If empty, the first
entry of :py:obj:`SettingsAISummary.models` is used. Users can set their
own model in the preferences (``ai_summary_model``) unless that preference
is locked."""
models: list[str] = []
"""List of model names suggested to the user in the preferences. If empty
and :py:obj:`SettingsAISummary.base_url` is set, the list is requested
once at application setup from the Ollama server (``GET /api/tags``)."""
grounding: bool = False
"""Default of the ``ai_summary_grounding`` user preference: ground the
summary on the search results. Users can still opt in/out in the
preferences unless that preference is locked."""
connect_timeout: float = 5.0
"""Timeout (seconds) to establish a TCP connection to the Ollama server."""
read_timeout: float = 30.0
"""Maximum gap (seconds) between two chunks of the token stream."""
stream_timeout: float = 120.0
"""Wall clock limit (seconds) for one completion."""
keep_alive: str = "5m"
"""How long the model stays loaded in memory after the request (passed
through to Ollama's ``keep_alive`` option)."""
max_context_items: int = 5
"""Maximum number of search results accepted as grounding context."""
max_history_messages: int = 12
"""Maximum number of messages (follow-up chat history) per request."""
max_message_length: int = 4000
"""Maximum length (characters) of a single message or context snippet."""
system_prompt: str = DEFAULT_SYSTEM_PROMPT
"""System prompt used when the *grounding* preference is off."""
system_prompt_grounded: str = DEFAULT_SYSTEM_PROMPT_GROUNDED
"""System prompt used when the *grounding* preference is on. The
placeholder ``{context}`` is replaced by an enumeration of the search
results sent along with the query."""
def model_choices() -> list[str]:
"""Model names a user can select from in the preferences."""
return list(MODELS)
def build_ollama_messages(
cfg: SettingsAISummary,
messages: list[dict[str, str]],
context: list[dict[str, str]] | None = None,
) -> list[dict[str, str]]:
"""Build the message list for Ollama's ``/api/chat`` from the (already
validated) request ``messages``, prepending a system prompt. When
``context`` items are given, the grounded system prompt is used and the
context items are serialized into its ``{context}`` placeholder."""
if context:
ctx_lines = [
f"[{no}] {item.get('title', '')}{item.get('snippet', '')} ({item.get('url', '')})"
for no, item in enumerate(context[: cfg.max_context_items], start=1)
]
system_prompt = cfg.system_prompt_grounded.replace("{context}", "\n".join(ctx_lines))
else:
system_prompt = cfg.system_prompt
return [{"role": "system", "content": system_prompt}, *messages]
+1 -1
View File
@@ -42,7 +42,7 @@ class PluginInfo:
description: str description: str
"""Short description of the *answerer*.""" """Short description of the *answerer*."""
preference_section: t.Literal["general", "ui", "privacy", "query"] | None = "general" preference_section: t.Literal["general", "ui", "privacy", "query", "ai"] | None = "general"
"""Section (tab/group) in the preferences where this plugin is shown to the """Section (tab/group) in the preferences where this plugin is shown to the
user. user.
+302
View File
@@ -0,0 +1,302 @@
# SPDX-License-Identifier: AGPL-3.0-or-later
"""Plugin that displays an AI generated summary of the search query at the top
of the result page. The summary is generated by a (local) `Ollama`_ server.
The Ollama server URL and the model are configured by the user in the *AI
Summary* tab of the preferences (``ai_summary_server``, ``ai_summary_model``);
the administrator can configure instance wide defaults in the ``ai_summary:``
section and lock the preferences via :ref:`settings preferences`.
.. attention::
A user configurable server URL allows any user of the instance to make the
SearXNG server send requests to a URL of their choice (`SSRF`_), and each
summary is real LLM work. This plugin is intended for private instances --
on a public instance, lock the ``ai_summary_server``, ``ai_summary_model``
and ``ai_summary_grounding`` preferences and configure the ``ai_summary:``
section instead.
The result page is never delayed by this plugin: it only places an empty
placeholder (:py:obj:`searx.result_types.AiSummary`) in the answer area, which
is filled asynchronously by the client (``client/simple/src/js/plugin/
AiSummary.ts``) from the ``/ai_summary`` endpoint (registered in
:py:obj:`searx.webapp`). The endpoint streams the tokens from Ollama's
``/api/chat`` to the client as `NDJSON`_.
A summary is only generated on the first page of a *general* search and only
if no engine has contributed an infobox (e.g. wikipedia / wikidata) or an
instant answer (e.g. ddg definitions) -- in these cases the query is most
likely a lookup of a well known term that is already answered.
The requests to the Ollama server are sent directly (not via
:py:obj:`searx.network`), an outgoing proxy configuration is deliberately not
applied to reach an Ollama server in the local network.
Configuration of the defaults (:py:obj:`searx.ai_summary.SettingsAISummary`):
.. code:: yaml
ai_summary:
base_url: "http://127.0.0.1:11434"
model: "llama3.2:3b"
.. code:: yaml
plugins:
searx.plugins.ai_summary.SXNGPlugin:
active: false
.. _Ollama: https://ollama.com/
.. _NDJSON: https://github.com/ndjson/ndjson-spec
.. _SSRF: https://owasp.org/www-community/attacks/Server_Side_Request_Forgery
"""
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_ollama_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) -> httpx.Client:
"""HTTP client for one request to the Ollama server at ``base_url``."""
return httpx.Client(
base_url=base_url,
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 _user_server(request: "SXNG_Request", cfg: SettingsAISummary) -> str:
"""The Ollama server URL for this request: the user's ``ai_summary_server``
preference, or the administrator's default."""
return str(request.preferences.get_value("ai_summary_server") or "").strip() 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 an Ollama 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 Ollama server (``GET
/api/tags``). 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) as client:
resp = client.get("/api/tags")
resp.raise_for_status()
models = [model["name"] for model in resp.json().get("models", [])]
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 Ollama 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 Ollama 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")
ollama_payload = {
"model": model,
"messages": build_ollama_messages(cfg, messages, context),
"stream": True,
"keep_alive": cfg.keep_alive,
}
# open the upstream connection before streaming, a connection error is
# reported as HTTP 502 instead of a line in an already started stream
client = _get_client(server, cfg)
stream_ctx = client.stream("POST", "/api/chat", json=ollama_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:
for line in upstream.iter_lines():
if time.monotonic() - start > cfg.stream_timeout:
yield ndjson({"done": True, "error": "timeout"})
return
if not line.strip():
continue
data = json.loads(line)
if data.get("done"):
yield ndjson({"done": True, "model": model})
return
delta = data.get("message", {}).get("content", "")
if delta:
yield ndjson({"delta": delta})
yield ndjson({"done": True, "model": model})
except (httpx.HTTPError, ValueError) as exc:
log.warning("error while streaming from Ollama: %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,
)
+14
View File
@@ -462,6 +462,20 @@ class Preferences:
locked="doi_resolver" in self.cfg.lock, locked="doi_resolver" in self.cfg.lock,
choices=DOI_RESOLVERS, choices=DOI_RESOLVERS,
), ),
# empty values fall back to the administrator's defaults in the
# ai_summary: section (searx.ai_summary.SettingsAISummary)
'ai_summary_server': StringSetting(
"",
locked="ai_summary_server" in self.cfg.lock,
),
'ai_summary_model': StringSetting(
"",
locked="ai_summary_model" in self.cfg.lock,
),
'ai_summary_grounding': BooleanSetting(
get_setting("ai_summary").grounding,
locked="ai_summary_grounding" in self.cfg.lock,
),
'simple_style': EnumStringSetting( 'simple_style': EnumStringSetting(
get_setting("ui.theme_args.simple_style"), get_setting("ui.theme_args.simple_style"),
locked="simple_style" in self.cfg.lock, locked="simple_style" in self.cfg.lock,
+3 -1
View File
@@ -19,6 +19,7 @@ __all__ = [
"EngineResults", "EngineResults",
"AnswerSet", "AnswerSet",
"Answer", "Answer",
"AiSummary",
"Translations", "Translations",
"WeatherAnswer", "WeatherAnswer",
"Code", "Code",
@@ -32,7 +33,7 @@ import typing as t
import abc import abc
from ._base import Result, MainResult, LegacyResult from ._base import Result, MainResult, LegacyResult
from .answer import AnswerSet, Answer, Translations, WeatherAnswer from .answer import AnswerSet, Answer, AiSummary, Translations, WeatherAnswer
from .keyvalue import KeyValue from .keyvalue import KeyValue
from .code import Code from .code import Code
from .paper import Paper from .paper import Paper
@@ -49,6 +50,7 @@ class ResultList(list[Result | LegacyResult], abc.ABC):
implemented).""" implemented)."""
Answer = Answer Answer = Answer
AiSummary = AiSummary
KeyValue = KeyValue KeyValue = KeyValue
Code = Code Code = Code
Paper = Paper Paper = Paper
+27 -1
View File
@@ -14,6 +14,10 @@ template.
:members: :members:
:show-inheritance: :show-inheritance:
.. autoclass:: AiSummary
:members:
:show-inheritance:
.. autoclass:: Translations .. autoclass:: Translations
:members: :members:
:show-inheritance: :show-inheritance:
@@ -29,7 +33,7 @@ template.
# pylint: disable=too-few-public-methods # pylint: disable=too-few-public-methods
__all__ = ["AnswerSet", "Answer", "Translations", "WeatherAnswer"] __all__ = ["AnswerSet", "Answer", "AiSummary", "Translations", "WeatherAnswer"]
from flask_babel import gettext from flask_babel import gettext
import msgspec import msgspec
@@ -92,6 +96,28 @@ class Answer(BaseAnswer, kw_only=True):
return hash(self.answer) return hash(self.answer)
class AiSummary(BaseAnswer, kw_only=True):
"""Placeholder for an AI generated summary of the search query
(:py:obj:`searx.plugins.ai_summary`). The summary itself is fetched
asynchronously by the client after the result page has been rendered."""
# The answers of an AnswerSet are sorted by their template name; this
# template sorts before the other answer templates, the AI summary is
# therefore rendered first in the answer area.
template: str = "answer/ai_summary.html"
query: str
"""The search query the summary is generated for."""
grounding: bool = False
"""Ground the summary on the search results (user's preference)."""
def __hash__(self):
"""The hash value of field *query* is the hash value of the
:py:obj:`AiSummary` object."""
return hash(self.query)
class Translations(BaseAnswer, kw_only=True): class Translations(BaseAnswer, kw_only=True):
"""Answer type with a list of translations. """Answer type with a list of translations.
+24
View File
@@ -257,6 +257,9 @@ plugins:
searx.plugins.tracker_url_remover.SXNGPlugin: searx.plugins.tracker_url_remover.SXNGPlugin:
active: true active: true
searx.plugins.ai_summary.SXNGPlugin:
active: false
# Configuration of the "Hostnames plugin": # Configuration of the "Hostnames plugin":
# #
@@ -284,6 +287,27 @@ plugins:
# '(.*\.)?youtu\.be$': 'yt.example.com' # '(.*\.)?youtu\.be$': 'yt.example.com'
# #
# Configuration of the "AI summary plugin", for more details see
# https://docs.searxng.org/admin/settings/settings_ai_summary.html
#
# ai_summary:
#
# # Base URL of the Ollama server; without this URL the plugin is inactive.
# base_url: "http://127.0.0.1:11434"
#
# # Default model; if empty, the first entry of models: is used.
# model: "llama3.2:3b"
#
# # Models the user can select from in the preferences; if empty, the list
# # is requested from the Ollama server (GET /api/tags) at startup.
# models:
# - "llama3.2:3b"
# - "gemma3:4b"
#
# # Ground summaries on the search results by default (users can still opt
# # in/out in their preferences).
# grounding: false
categories_as_tabs: categories_as_tabs:
general: general:
+2
View File
@@ -13,6 +13,7 @@ from os.path import dirname, abspath
import msgspec import msgspec
from typing_extensions import override from typing_extensions import override
from .ai_summary import SettingsAISummary
from .brand import SettingsBrand from .brand import SettingsBrand
from .sxng_locales import sxng_locales from .sxng_locales import sxng_locales
from ._settings import SettingsPref from ._settings import SettingsPref
@@ -267,6 +268,7 @@ SCHEMA: dict[str, t.Any] = {
'networks': {}, 'networks': {},
}, },
'plugins': SettingsValue(dict, {}), 'plugins': SettingsValue(dict, {}),
'ai_summary': SettingsAISummary,
'categories_as_tabs': SettingsValue(dict, CATEGORIES_AS_TABS), 'categories_as_tabs': SettingsValue(dict, CATEGORIES_AS_TABS),
'engines': SettingsValue(list, []), 'engines': SettingsValue(list, []),
'doi_resolvers': {}, 'doi_resolvers': {},
@@ -0,0 +1,19 @@
<div id="ai_summary" class="ai-summary hide_if_nojs" {{- ' ' -}}
data-query="{{ answer.query }}" {{- ' ' -}}
data-grounding="{{ '1' if answer.grounding else '0' }}">{{- '' -}}
<div class="ai-summary-header">{{- '' -}}
<span class="ai-summary-title">{{ _('AI Summary') }}</span>{{- '' -}}
<span class="ai-summary-model"></span>{{- '' -}}
</div>{{- '' -}}
<div class="ai-summary-body collapsed">{{- '' -}}
<div class="ai-summary-answers" aria-live="polite"></div>{{- '' -}}
</div>{{- '' -}}
<button type="button" class="ai-summary-more invisible" {{- ' ' -}}
data-btn-text-collapsed="{{ _('More') }}" {{- ' ' -}}
data-btn-text-not-collapsed="{{ _('Less') }}">{{ _('More') }}</button>{{- '' -}}
<form class="ai-summary-followup invisible">{{- '' -}}
<input type="text" placeholder="{{ _('Ask a follow-up question') }}" autocomplete="off" maxlength="2048">{{- '' -}}
<button type="submit">{{ _('Ask') }}</button>{{- '' -}}
</form>{{- '' -}}
<p class="ai-summary-disclaimer">{{ _('Generated by AI — may contain mistakes.') }}</p>{{- '' -}}
</div>
+9
View File
@@ -249,6 +249,15 @@
{%- endif -%} {%- endif -%}
{{- tab_footer() -}} {{- tab_footer() -}}
{# tab: ai #}
{%- if plugins_storage | selectattr('preference_section', 'equalto', 'ai') | list -%}
{{- tab_header('maintab', 'ai', _('AI Summary')) -}}
{{- plugin_preferences('ai') -}}
{%- include 'simple/preferences/ai_summary.html' -%}
{{- tab_footer() -}}
{%- endif -%}
{# tab: cookies #} {# tab: cookies #}
{{- tab_header('maintab', 'cookies', _('Cookies')) -}} {{- tab_header('maintab', 'cookies', _('Cookies')) -}}
@@ -0,0 +1,59 @@
{%- if 'ai_summary_server' not in locked_preferences -%}
<fieldset>{{- '' -}}
<legend id="pref_ai_summary_server">{{- _('Ollama server URL') -}}</legend>{{- '' -}}
<div class="value">{{- '' -}}
<input name="ai_summary_server" aria-labelledby="pref_ai_summary_server" type="text"
autocomplete="off" spellcheck="false" autocorrect="off"
placeholder="{{ ai_summary_default_server or 'http://127.0.0.1:11434' }}"
value="{{ preferences.get_value('ai_summary_server') }}">{{- '' -}}
</div>{{- '' -}}
<div class="description">
{{- _('URL of the Ollama server that generates the summaries, e.g. http://192.168.1.10:11434.') -}}
{{- ' ' -}}
{%- if ai_summary_default_server -%}
{{- _('Leave empty to use the default of this instance.') -}}
{%- endif -%}
</div>{{- '' -}}
</fieldset>{{- '' -}}
{%- endif -%}
{%- if 'ai_summary_model' not in locked_preferences -%}
<fieldset>{{- '' -}}
<legend id="pref_ai_summary_model">{{- _('AI summary model') -}}</legend>{{- '' -}}
<div class="value">{{- '' -}}
<input name="ai_summary_model" aria-labelledby="pref_ai_summary_model" type="text"
autocomplete="off" spellcheck="false" autocorrect="off" list="ai_summary_model_list"
placeholder="{{ ai_summary_default_model or 'llama3.2:3b' }}"
value="{{ preferences.get_value('ai_summary_model') }}">{{- '' -}}
<datalist id="ai_summary_model_list">
{%- for model in ai_summary_models -%}
<option value="{{ model }}"></option>
{%- endfor -%}
</datalist>{{- '' -}}
</div>{{- '' -}}
<div class="description">
{{- _('Name of the model used to generate the summaries, e.g. llama3.2:3b or gemma3:4b.') -}}
{{- ' ' -}}
{%- if ai_summary_default_model -%}
{{- _('Leave empty to use the default of this instance.') -}}
{%- endif -%}
</div>{{- '' -}}
</fieldset>{{- '' -}}
{%- endif -%}
{%- if 'ai_summary_grounding' not in locked_preferences -%}
<fieldset>{{- '' -}}
<legend id="pref_ai_summary_grounding">{{ _('Ground AI summary on search results') }}</legend>{{- '' -}}
<p class="value">{{- '' -}}
<input type="checkbox" {{- ' ' -}}
name="ai_summary_grounding" {{- ' ' -}}
aria-labelledby="pref_ai_summary_grounding" {{- ' ' -}}
class="checkbox-onoff" {{- ' ' -}}
{%- if preferences.get_value('ai_summary_grounding') -%}
checked
{%- endif -%}{{- ' ' -}}
>{{- '' -}}
</p>{{- '' -}}
<div class="description">
{{- _('Send the top search results along with the query, the model answers from this context instead of its own knowledge. Answers are more accurate and more current, but generating them is slower and needs more memory (VRAM) on the Ollama server.') -}}
</div>{{- '' -}}
</fieldset>{{- '' -}}
{%- endif -%}
+10
View File
@@ -93,8 +93,10 @@ from searx.preferences import (
ClientPref, ClientPref,
ValidationException, ValidationException,
) )
import searx.ai_summary
import searx.answerers import searx.answerers
import searx.plugins import searx.plugins
import searx.plugins.ai_summary
from searx.metrics import get_engines_stats, get_engine_errors, get_reliabilities, histogram, counter, openmetrics from searx.metrics import get_engines_stats, get_engine_errors, get_reliabilities, histogram, counter, openmetrics
@@ -328,6 +330,8 @@ def get_translations():
'Source': gettext('Source'), 'Source': gettext('Source'),
# infinite scroll # infinite scroll
'error_loading_next_page': gettext('Error loading the next page'), 'error_loading_next_page': gettext('Error loading the next page'),
# AI summary
'error_loading_ai_summary': gettext('Error loading the AI summary'),
} }
@@ -981,11 +985,17 @@ def preferences():
preferences_url_params = sxng_request.preferences.get_as_url_params(), preferences_url_params = sxng_request.preferences.get_as_url_params(),
locked_preferences = get_setting("preferences").lock, locked_preferences = get_setting("preferences").lock,
doi_resolvers = get_setting("doi_resolvers", {}), doi_resolvers = get_setting("doi_resolvers", {}),
ai_summary_models = searx.ai_summary.model_choices(),
ai_summary_default_model = get_setting("ai_summary").model,
ai_summary_default_server = get_setting("ai_summary").base_url,
# fmt: on # fmt: on
) )
app.add_url_rule('/favicon_proxy', methods=['GET'], endpoint="favicon_proxy", view_func=favicons.favicon_proxy) app.add_url_rule('/favicon_proxy', methods=['GET'], endpoint="favicon_proxy", view_func=favicons.favicon_proxy)
app.add_url_rule(
'/ai_summary', methods=['POST'], endpoint="ai_summary", view_func=searx.plugins.ai_summary.ai_summary_view
)
@app.route('/image_proxy', methods=['GET']) @app.route('/image_proxy', methods=['GET'])
+271
View File
@@ -0,0 +1,271 @@
# SPDX-License-Identifier: AGPL-3.0-or-later
# pylint: disable=missing-module-docstring,missing-class-docstring,invalid-name,protected-access
# pylint: disable=too-many-public-methods
import json
from contextlib import contextmanager
from mock import Mock
import searx.ai_summary
import searx.plugins
import searx.plugins.ai_summary
import searx.preferences
from searx.extended_types import sxng_request
from searx.result_types import AiSummary
from tests import SearxTestCase
from .test_plugins import get_search_mock
PLUGIN_FQN = "searx.plugins.ai_summary.SXNGPlugin"
BASE_URL = "http://127.0.0.1:11434"
MODEL = "test-model"
def ollama_stream_mock(lines: list[dict], status_code: int = 200) -> Mock:
"""A mock httpx client whose ``stream()`` context manager yields the given
(Ollama) NDJSON lines."""
upstream = Mock(status_code=status_code)
upstream.iter_lines.return_value = iter([json.dumps(line) for line in lines])
@contextmanager
def stream(*_args, **_kwargs):
yield upstream
client = Mock()
client.stream = stream
return client
class PluginAISummaryInit(SearxTestCase):
def test_active_without_base_url(self):
# the Ollama server can be configured by the user in the preferences,
# the plugin stays active without an admin configured default
self.setattr4test(searx.ai_summary, "MODELS", [])
storage = searx.plugins.PluginStorage()
storage.load_settings({PLUGIN_FQN: {"active": True}})
storage.init(self.app)
self.assertEqual(1, len(storage))
self.assertEqual([], searx.ai_summary.MODELS)
def test_model_list_from_settings(self):
cfg = searx.get_setting("ai_summary")
self.setattr4test(cfg, "base_url", BASE_URL)
self.setattr4test(cfg, "model", MODEL)
self.setattr4test(cfg, "models", [MODEL, "other-model"])
self.setattr4test(searx.ai_summary, "MODELS", [])
storage = searx.plugins.PluginStorage()
storage.load_settings({PLUGIN_FQN: {"active": True}})
storage.init(self.app)
self.assertEqual(1, len(storage))
self.assertEqual([MODEL, "other-model"], searx.ai_summary.MODELS)
class PluginAISummary(SearxTestCase):
def setUp(self):
super().setUp()
cfg = searx.get_setting("ai_summary")
self.setattr4test(cfg, "base_url", BASE_URL)
self.setattr4test(cfg, "model", MODEL)
self.setattr4test(cfg, "models", [MODEL])
self.setattr4test(searx.ai_summary, "MODELS", [])
self.storage = searx.plugins.PluginStorage()
self.storage.load_settings({PLUGIN_FQN: {"active": True}})
self.storage.init(self.app)
# the endpoint checks request.user_plugins (built in webapp.pre_request
# from the global plugin storage) -- enable the plugin like a browser
# with saved preferences does
self.client.set_cookie("disabled_plugins", "")
self.client.set_cookie("enabled_plugins", "ai_summary")
self.pref = searx.preferences.Preferences(["simple"], ["general"], {}, self.storage)
self.pref.parse_dict({"locale": "en"})
def mock_upstream(self, client_mock: Mock):
self.setattr4test(searx.plugins.ai_summary, "_get_client", lambda *_args, **_kwargs: client_mock)
def do_post_search(self, query, **kwargs) -> Mock:
kwargs.setdefault("categories", ["general"])
search = get_search_mock(query, user_plugins=["ai_summary"], **kwargs)
self.storage.post_search(sxng_request, search)
return search
def test_placeholder_answer_is_added(self):
with self.app.test_request_context():
sxng_request.preferences = self.pref
search = self.do_post_search("what is the best searx fork")
answer = AiSummary(query="what is the best searx fork", grounding=False)
self.assertIn(answer, search.result_container.answers)
def test_placeholder_carries_grounding_pref(self):
with self.app.test_request_context():
sxng_request.preferences = self.pref
self.pref.parse_dict({"ai_summary_grounding": "1"})
search = self.do_post_search("lorem ipsum")
answer = list(search.result_container.answers)[0]
self.assertTrue(answer.grounding)
def test_grounding_default_from_settings(self):
self.setattr4test(searx.get_setting("ai_summary"), "grounding", True)
pref = searx.preferences.Preferences(["simple"], ["general"], {}, self.storage)
with self.app.test_request_context():
sxng_request.preferences = pref
search = self.do_post_search("lorem ipsum")
answer = list(search.result_container.answers)[0]
self.assertTrue(answer.grounding)
def test_skip_pageno(self):
with self.app.test_request_context():
sxng_request.preferences = self.pref
search = self.do_post_search("lorem ipsum", pageno=2)
self.assertEqual(list(search.result_container.answers), [])
def test_skip_non_html_format(self):
with self.app.test_request_context("/search", method="POST", data={"format": "json"}):
sxng_request.preferences = self.pref
search = self.do_post_search("lorem ipsum")
self.assertEqual(list(search.result_container.answers), [])
def test_skip_non_general_category(self):
with self.app.test_request_context():
sxng_request.preferences = self.pref
search = self.do_post_search("lorem ipsum", categories=["images"])
self.assertEqual(list(search.result_container.answers), [])
def test_skip_infobox(self):
with self.app.test_request_context():
sxng_request.preferences = self.pref
search = get_search_mock("lorem ipsum", user_plugins=["ai_summary"], categories=["general"])
search.result_container.infoboxes.append(Mock())
self.storage.post_search(sxng_request, search)
self.assertEqual(list(search.result_container.answers), [])
def test_skip_instant_answer(self):
# e.g. the "ddg definitions" engine adds wikipedia abstracts as Answer
from searx.result_types import Answer # pylint: disable=import-outside-toplevel
with self.app.test_request_context():
sxng_request.preferences = self.pref
search = get_search_mock("lorem ipsum", user_plugins=["ai_summary"], categories=["general"])
engine_answer = Answer(answer="Lorem ipsum is placeholder text. More at Wikipedia")
search.result_container.answers.add(engine_answer)
self.storage.post_search(sxng_request, search)
self.assertEqual(list(search.result_container.answers), [engine_answer])
def test_skip_without_any_server(self):
self.setattr4test(searx.get_setting("ai_summary"), "base_url", "")
with self.app.test_request_context():
sxng_request.preferences = self.pref
search = self.do_post_search("lorem ipsum")
self.assertEqual(list(search.result_container.answers), [])
def test_placeholder_with_user_server_only(self):
self.setattr4test(searx.get_setting("ai_summary"), "base_url", "")
with self.app.test_request_context():
sxng_request.preferences = self.pref
self.pref.parse_dict({"ai_summary_server": "http://192.168.1.10:11434"})
search = self.do_post_search("lorem ipsum")
self.assertEqual(len(list(search.result_container.answers)), 1)
def test_endpoint_forbidden_when_disabled(self):
self.client.set_cookie("disabled_plugins", "ai_summary")
self.client.set_cookie("enabled_plugins", "")
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 403)
def test_endpoint_bad_request(self):
for body in [
None,
{},
{"messages": []},
{"messages": [{"role": "system", "content": "hi"}]},
{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "ho"}]},
{"messages": [{"role": "user", "content": "x" * 5000}]},
{"messages": [{"role": "user", "content": "hi"}], "context": [{"unknown_key": "x"}]},
{"messages": [{"role": "user", "content": "hi"}] * 13},
]:
res = self.client.post("/ai_summary", json=body)
self.assertEqual(res.status_code, 400, body)
def test_endpoint_invalid_server_pref(self):
self.client.set_cookie("ai_summary_server", "ftp://example.org")
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 400)
def test_endpoint_invalid_model_pref(self):
self.client.set_cookie("ai_summary_model", "bad model name!")
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 400)
def test_endpoint_no_server_configured(self):
self.setattr4test(searx.get_setting("ai_summary"), "base_url", "")
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 400)
def test_endpoint_streams_ndjson(self):
self.mock_upstream(
ollama_stream_mock(
[
{"message": {"content": "Hello "}, "done": False},
{"message": {"content": "world"}, "done": False},
{"done": True},
]
)
)
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 200)
self.assertEqual(res.headers["Content-Type"], "application/x-ndjson")
lines = [json.loads(line) for line in res.data.decode().splitlines() if line]
self.assertEqual(lines[0], {"delta": "Hello "})
self.assertEqual(lines[1], {"delta": "world"})
self.assertEqual(lines[2], {"done": True, "model": MODEL})
def test_endpoint_model_pref_wins(self):
self.mock_upstream(ollama_stream_mock([{"done": True}]))
self.client.set_cookie("ai_summary_model", "my-own-model:7b")
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
lines = [json.loads(line) for line in res.data.decode().splitlines() if line]
self.assertEqual(lines[-1], {"done": True, "model": "my-own-model:7b"})
def test_endpoint_upstream_error(self):
self.mock_upstream(ollama_stream_mock([], status_code=500))
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 502)
def test_endpoint_error_while_streaming(self):
upstream = Mock(status_code=200)
upstream.iter_lines.return_value = iter(
[json.dumps({"message": {"content": "Hello"}, "done": False}), "this is not json"]
)
@contextmanager
def stream(*_args, **_kwargs):
yield upstream
client_mock = Mock()
client_mock.stream = stream
self.mock_upstream(client_mock)
res = self.client.post("/ai_summary", json={"messages": [{"role": "user", "content": "hi"}]})
self.assertEqual(res.status_code, 200)
lines = [json.loads(line) for line in res.data.decode().splitlines() if line]
self.assertEqual(lines[0], {"delta": "Hello"})
self.assertEqual(lines[1], {"done": True, "error": "upstream error"})