searxng/searx/ai_summary.py

# 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]