chemenu/tools/chemenu/search/ripgrep.py

"""The lexical backend: `rg` over `kb/`, parsed from its JSON output.

Why shell out instead of scanning in Python: `rg` is already the retrieval
layer the agent instructions point at, it handles large trees fast, and its
`--json` mode gives line numbers and matched text without reparsing files.

Three safety properties are load-bearing and must survive any edit here:

1. The query is passed as an *argv element*, never through a shell. There is
   no `shell=True` anywhere in this module, so a query containing `;`, `$(...)`
   or backticks is searched for literally rather than executed.
2. `--fixed-strings` is the default. A user-supplied regex is opt-in via
   `--regex`, so an accidental `.*` in a search term is a literal, and a
   pathological pattern cannot be introduced without asking for one.
3. A user-supplied pattern is evaluated **only** by `rg`, whose engine is
   linear in the input. Nothing here hands it to Python's `re`, which
   backtracks - see `_contains`.
"""
from __future__ import annotations

import json
import subprocess
from pathlib import Path
from typing import Iterable

from chemenu import config
from chemenu.errors import BackendError
from chemenu.page import Page
from chemenu.search.base import page_key
from chemenu.search.filters import collection_of
from chemenu.search.types import Match, SearchHit, SearchQuery

# Only the first few matching lines per page are kept. The hit is a pointer for
# deciding whether to open the page, not a substitute for reading it, and an
# unbounded excerpt list is exactly the token cost this command exists to avoid.
MAX_MATCHES_PER_PAGE = 3

# Ranking tiers. An *exact* title match is worth more than a title that merely
# contains the term, because "longhorn" should surface the page called Longhorn
# ahead of every page whose title mentions it - otherwise a well-connected
# source page outranks the subject it is about.
EXACT_TITLE_WEIGHT = 10.0
TITLE_WEIGHT = 5.0
SUMMARY_WEIGHT = 3.0
LINE_WEIGHT = 1.0

# A hang-breaker, not a performance budget. A fixed-string search over this
# corpus costs 6 ms and a deliberately broad regex 1.4 s, so nothing legitimate
# comes near this; it exists so that a pathological pattern, a corpus on a
# stalled network mount, or an `rg` that never returns fails as an error
# instead of holding the caller open forever while its output buffers into the
# heap. The caller sees the ordinary `RipgrepFailed` path.
RIPGREP_TIMEOUT_SECONDS = 30.0


class RipgrepMissing(BackendError):
    """Raised when the `rg` executable is not on PATH."""


class RipgrepFailed(BackendError):
    """Raised when `rg` exits with an error (exit code 2 or above), or had to
    be killed for running past `RIPGREP_TIMEOUT_SECONDS`."""


def build_argv(query: SearchQuery, root: Path) -> list[str]:
    """The exact command line. Split out so a test can assert the safety
    properties above without running anything."""
    argv = ["rg", "--json", "--smart-case", "--glob", "*.md"]
    if not query.regex:
        argv.append("--fixed-strings")
    # `--` terminates option parsing: a query starting with `-` is a search
    # term, not a flag.
    argv += ["--", query.text or "", str(root)]
    return argv


def _iter_match_records(stdout: str) -> Iterable[dict]:
    for line in stdout.splitlines():
        if not line.strip():
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if record.get("type") == "match":
            yield record.get("data", {})


def _record_path(data: dict) -> str | None:
    path = data.get("path") or {}
    return path.get("text")


class RipgrepBackend:
    """Lexical search over page bodies and frontmatter text."""

    name = "rg"

    def __init__(self, search_root: Path | None = None, repo_root: Path | None = None):
        # Two roots, because they answer different questions: `search_root` is
        # what rg walks, `repo_root` is what the resulting paths are made
        # relative to so they match the keys in `pages`. They differ only in
        # tests, but conflating them makes the backend untestable outside the
        # real repo.
        self.search_root = search_root or config.KB_DIR
        self.repo_root = repo_root or config.ROOT

    def search(self, query: SearchQuery, pages: dict[str, Page]) -> list[SearchHit]:
        if not query.text:
            return []

        argv = build_argv(query, self.search_root)
        try:
            proc = subprocess.run(
                argv,
                capture_output=True,
                text=True,
                check=False,
                timeout=RIPGREP_TIMEOUT_SECONDS,
            )
        except subprocess.TimeoutExpired as exc:
            raise RipgrepFailed(
                f"rg did not finish within {RIPGREP_TIMEOUT_SECONDS:g}s and was killed. "
                "A search that takes this long is a pathological pattern or an "
                "unresponsive corpus directory, not a slow answer - narrow the query "
                "or use a fixed string instead of --regex."
            ) from exc
        except FileNotFoundError as exc:  # pragma: no cover - depends on host
            raise RipgrepMissing(
                "ripgrep (rg) is not installed or not on PATH. It is the search "
                "backend; install it (e.g. `pacman -S ripgrep`, `apt install ripgrep`) "
                "and retry. wikitool deliberately has no Python fallback: a fallback "
                "would answer differently from the documented backend without saying so."
            ) from exc

        # rg exits 1 for "no matches found", which is an answer, not a failure.
        if proc.returncode >= 2:
            raise RipgrepFailed(f"rg exited {proc.returncode}: {proc.stderr.strip()}")

        by_path: dict[str, list[Match]] = {}
        for data in _iter_match_records(proc.stdout):
            raw_path = _record_path(data)
            if raw_path is None:
                continue
            key = page_key(Path(raw_path), self.repo_root)
            if key not in pages:
                # Not a page: a generated index, a contract, or a file outside
                # the corpus. The pages dict is the authority on what exists.
                continue
            found = by_path.setdefault(key, [])
            if len(found) >= MAX_MATCHES_PER_PAGE:
                continue
            found.append(
                Match(
                    line=data.get("line_number", 0),
                    text=(data.get("lines", {}).get("text") or "").rstrip("\n"),
                )
            )

        hits = [
            build_hit(pages[key], key, matches, query, backend=self.name, kb_dir=self.search_root)
            for key, matches in by_path.items()
        ]
        hits.sort(key=lambda h: (-h.score, h.title.lower()))
        return hits


def _contains(haystack: str, query: SearchQuery) -> bool:
    """Literal containment, used only for the title and summary ranking boosts.

    It never evaluates `query.text` as a regex, even when `query.regex` is set.
    It used to, via `re.search`, and Python's engine backtracks: `(\\w+\\s?)+$`
    against 114 characters of ordinary page text does not terminate in eight
    seconds, while a pattern that fails deterministically takes 0.2 ms - the
    difference is the pattern, not the haystack. `build_hit` calls this twice
    per hit, and a pattern as cheap as `\\w` matches every page, so one request
    bought two unbounded searches per page in the corpus.

    Deleting the branch rather than bounding it is the right trade: `rg` has
    already applied the pattern with a linear engine by the time we get here,
    and the page is a hit *because* of that. What is lost is only the extra
    weight a regex hit in the title would have scored - and since a summary and
    an H1 are themselves lines in the file, `rg` still counts them. A pattern
    that is mostly literal (`longhorn`) still earns its boost through the test
    below; one that is not gets ranked by match count alone.
    """
    if not query.text:
        return False
    return query.text.lower() in haystack.lower()


def build_hit(
    page: Page,
    key: str,
    matches: list[Match],
    query: SearchQuery,
    backend: str,
    kb_dir: Path | None = None,
) -> SearchHit:
    """Turn a page plus its matching lines into an enriched, scored hit.

    Ranking is deliberately crude and explainable: an exact title match
    outweighs a partial one, which outweighs a summary match, which outweighs
    body matches. An agent scanning results should be able to predict the
    order, which a tuned scorer would not give.
    """
    summary = str(page.frontmatter.get("summary") or "")
    score = LINE_WEIGHT * len(matches)
    if query.text and page.title.lower() == query.text.lower():
        score += EXACT_TITLE_WEIGHT
    elif _contains(page.title, query):
        score += TITLE_WEIGHT
    if _contains(summary, query):
        score += SUMMARY_WEIGHT

    tags = page.frontmatter.get("tags") or []
    modified = page.frontmatter.get("modified") or page.frontmatter.get("date")

    return SearchHit(
        title=page.title,
        path=key,
        collection=collection_of(page.path, kb_dir),
        kind=page.kind,
        subtype=page.subtype,
        summary=summary,
        tags=[str(t) for t in tags] if isinstance(tags, (list, tuple)) else [str(tags)],
        modified=str(modified) if modified else None,
        score=score,
        backend=backend,
        matches=matches,
    )