Source code for tools.prowlarr_search

"""Search torrent indexers via a local Prowlarr instance (Docker *arr stack)."""

from __future__ import annotations

import jsonutil as json
import logging
from typing import TYPE_CHECKING
from urllib.parse import quote

import aiohttp
from tools.alter_privileges import has_privilege, PRIVILEGES

if TYPE_CHECKING:
    from tool_context import ToolContext

logger = logging.getLogger(__name__)

TOOL_NAME = "search_prowlarr"
TOOL_DESCRIPTION = (
    "Search all enabled torrent indexers in Prowlarr at once (same as Prowlarr interactive search: "
    "no indexer filter is applied). Results are sorted by seeders (highest first), then truncated to "
    "max_results (default 5). Indexers must be enabled in Prowlarr; disabled indexers are never queried. "
    "Requires prowlarr_base_url and prowlarr_api_key in config.yaml."
)
TOOL_PARAMETERS = {
    "type": "object",
    "properties": {
        "query": {
            "type": "string",
            "description": "Search query (keywords, release title, etc.).",
        },
        "max_results": {
            "type": "integer",
            "default": 5,
            "description": (
                "Maximum number of releases to return after sorting by seeders (highest first). "
                "Default 5; capped at 200."
            ),
        },
    },
    "required": ["query"],
}


def _int_seeders(item: object) -> int:
    """Extract a Prowlarr release's seeder count as a safe sort key.

    Reads the ``"seeders"`` field from a single Prowlarr release dictionary
    and coerces it to an integer, defending against every shape the upstream
    API may return: non-dict items, a missing or ``None`` seeders field, and
    non-numeric values all collapse to ``0`` rather than raising. This makes
    the function safe to use directly as the ``key`` argument to
    :func:`sorted`. This is a pure function with no side effects.

    This is used by :func:`_sort_and_limit_releases` in this module as the
    sort key (passed as ``key=_int_seeders``, without an explicit call site);
    no other internal callers were found.

    Args:
        item: A single search-result element from Prowlarr, expected to be a
            ``dict`` but tolerated as any object.

    Returns:
        int: The release's seeder count, or ``0`` when the value is absent,
        ``None``, non-numeric, or the item is not a mapping.
    """
    if not isinstance(item, dict):
        return 0
    v = item.get("seeders")
    if v is None:
        return 0
    try:
        return int(v)
    except (TypeError, ValueError):
        return 0


def _sort_and_limit_releases(data: object, max_results: int) -> object:
    """Rank a list of Prowlarr releases by seeders and keep the top ``max_results``.

    Sorts the raw search response by seeder count (highest first, using
    :func:`_int_seeders` as the key) and truncates to the requested cap, clamped to
    the ``[1, 200]`` range. Non-list inputs are returned unchanged so that an error
    payload from Prowlarr passes through untouched. This is a pure function with no
    side effects.

    Called by :func:`run` on the parsed Prowlarr JSON before it is re-serialised for
    the caller; no other internal callers.

    Args:
        data: The decoded Prowlarr search response; expected to be a list of release
            dicts but tolerated as any object.
        max_results (int): Desired number of releases to keep; clamped to
            ``[1, 200]``.

    Returns:
        object: The top ``max_results`` releases sorted by seeders when ``data`` is a
        list, otherwise ``data`` unchanged.
    """
    if not isinstance(data, list):
        return data
    cap = max(1, min(max_results, 200))
    ranked = sorted(data, key=_int_seeders, reverse=True)
    return ranked[:cap]


def _fetch_limit_for_sort(max_results: int) -> int:
    """Compute how many rows to request from Prowlarr so seeder-ranking is meaningful.

    Because Prowlarr merges and dedupes results across every enabled indexer, asking
    for only the final ``max_results`` rows would rank a tiny, arbitrary slice. This
    over-fetches a larger pool — roughly 80x the (clamped) requested cap — bounded to
    ``[200, 3000]`` so the top-by-seeders selection in :func:`_sort_and_limit_releases`
    is drawn from a representative set without overloading the indexers. Pure
    arithmetic with no side effects.

    Called by :func:`run` to build the ``limit`` query parameter for the Prowlarr
    request; no other internal callers.

    Args:
        max_results (int): The caller's desired result count; internally clamped to
            ``[1, 200]`` before scaling.

    Returns:
        int: The fetch limit to send to Prowlarr, within ``[200, 3000]``.
    """
    cap = max(1, min(max_results, 200))
    # Pool large enough that top-by-seeders is meaningful after merge/dedupe.
    return min(3000, max(200, cap * 80))


def _normalize_base(url: str) -> str:
    """Strip trailing slashes from a Prowlarr base URL before path joining.

    Removes any trailing ``/`` characters so the configured
    ``prowlarr_base_url`` can be concatenated with ``/api/v1/search`` without
    producing a doubled slash. This is a pure string helper with no network
    or other side effects.

    This is called by :func:`run` in this module when assembling the search
    request URL; no other internal callers were found.

    Args:
        url: The raw Prowlarr base URL from ``config.yaml`` (e.g.
            ``"http://localhost:9696/"``).

    Returns:
        str: The same URL with all trailing slashes removed.
    """
    return url.rstrip("/")


# ---------------------------------------------------------------------------
# Authorization
# ---------------------------------------------------------------------------


async def _check_web_search_access(ctx) -> str | None:
    """Authorize the caller for web search, returning an error JSON or ``None``.

    Pulls ``user_id``/``redis``/``config`` off the tool context and tests the
    ``WEB_SEARCH`` privilege via :func:`tools.alter_privileges.has_privilege`.
    Returns ``None`` when access is granted; otherwise returns a JSON string with
    ``success=False`` and an explanatory error so the caller can surface it
    verbatim. Read-only.

    Called at the top of :func:`run` (the ``search_prowlarr`` tool) to gate the
    search behind the privilege.

    Args:
        ctx: The tool context; its ``user_id``, ``redis`` and ``config``
            attributes drive the privilege lookup.

    Returns:
        str | None: ``None`` if the user may search, else a JSON error string.
    """
    user_id = getattr(ctx, "user_id", "") or ""
    redis = getattr(ctx, "redis", None)
    config = getattr(ctx, "config", None)
    if not await has_privilege(redis, user_id, PRIVILEGES["WEB_SEARCH"], config):
        return json.dumps(
            {
                "success": False,
                "error": "The user does not have the WEB_SEARCH privilege. Ask an admin to grant it.",
            }
        )
    return None


[docs] async def run( query: str, max_results: int = 5, ctx: ToolContext | None = None, ) -> str: """Search every enabled Prowlarr indexer for a query and return ranked JSON. Entry point for the ``search_prowlarr`` tool. It authorizes the caller, reads the Prowlarr endpoint from config, performs one unfiltered interactive search across all enabled indexers, then ranks the merged releases by seeders and truncates to ``max_results`` before returning them as a JSON string. It gates on :func:`_check_web_search_access` (the ``WEB_SEARCH`` privilege, resolved against Redis), reads ``prowlarr_base_url`` and ``prowlarr_api_key`` from ``ctx.config``, normalises the base URL via :func:`_normalize_base`, and over-fetches using :func:`_fetch_limit_for_sort`. The actual call is an async ``aiohttp`` GET to ``/api/v1/search`` (120s timeout, ``X-Api-Key`` header); the JSON response is ranked and trimmed by :func:`_sort_and_limit_releases`. Network and unexpected errors are caught, logged via ``logger.exception``, and folded into an error payload rather than propagated. No module state is mutated. Registered via the single-tool module contract (``TOOL_NAME`` / ``run``) and dispatched by ``tool_loader.load_tools``; no direct in-repo Python callers. Args: query (str): Search keywords / release title; must be non-empty. max_results (int): Releases to return after seeder-ranking; coerced to an int and capped at 200 (default 5). ctx (ToolContext | None): The tool context supplying ``user_id``, ``redis``, and ``config`` (Prowlarr URL and API key). Returns: str: A JSON array of the top releases on success, or a JSON ``{"error": ...}`` string when authorization fails, the query is empty, config is missing the URL/key, Prowlarr returns a non-200 or invalid JSON, or a network/unexpected error occurs. """ auth_err = await _check_web_search_access(ctx) if auth_err: return auth_err if not query or not query.strip(): return json.dumps({"error": "Search query cannot be empty"}) try: cap = int(max_results) except (TypeError, ValueError): cap = 5 cap = max(1, min(cap, 200)) cfg = ctx.config if ctx else None if not cfg: return json.dumps({"error": "Tool context not available (config missing)."}) base = getattr(cfg, "prowlarr_base_url", "") or "" api_key = getattr(cfg, "prowlarr_api_key", "") or "" if not base.strip(): return json.dumps( {"error": "prowlarr_base_url is not set in config.yaml."}, ) if not api_key.strip(): return json.dumps( {"error": "prowlarr_api_key is not set in config.yaml."}, ) q = quote(query.strip(), safe="") fetch_limit = _fetch_limit_for_sort(cap) url = f"{_normalize_base(base)}/api/v1/search" f"?query={q}&limit={fetch_limit}" headers = {"X-Api-Key": api_key.strip()} try: async with aiohttp.ClientSession() as session: async with session.get( url, headers=headers, timeout=aiohttp.ClientTimeout(total=120), ) as response: text = await response.text() if response.status == 200: try: data = json.loads(text) except json.JSONDecodeError: return json.dumps( { "error": "Invalid JSON from Prowlarr", "body": text[:2000], }, ) trimmed = _sort_and_limit_releases(data, cap) return json.dumps(trimmed, indent=2) return json.dumps( {"error": f"HTTP {response.status}", "body": text[:2000]}, ) except aiohttp.ClientError as e: logger.exception("Prowlarr search request failed") return json.dumps({"error": f"Network error: {e}"}) except Exception as e: logger.exception("Prowlarr search unexpected error") return json.dumps({"error": f"Unexpected error: {e}"})