wiki-kb/wikikb/search.py

"""Hybrid retrieval: FTS5 (BM25) + embedding cosine, fused with Reciprocal Rank Fusion."""

from __future__ import annotations

import re
import sqlite3
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np

from . import config
from . import variant as variant_mod
from .embed import EmbeddingModel, embed_query

_WORD = re.compile(r"\w+", re.UNICODE)


@dataclass(frozen=True)
class Result:
    """A ranked retrieval hit."""

    title: str
    url: str
    source: str
    variant: str
    section: str
    text: str
    score: float


def _fts_match_query(query: str) -> str:
    """Build an FTS5 MATCH: OR of quoted terms so any term can hit.

    OR (not AND) keeps recall high on natural-language queries; BM25's IDF then
    ranks chunks that match the rarer, more meaningful terms above those that only
    share common words.
    """
    return " OR ".join(f'"{word}"' for word in _WORD.findall(query.lower()))


def _allowed_ids(con: sqlite3.Connection, scope: str) -> set[int] | None:
    """Chunk ids in scope, or None when unrestricted."""
    if scope == "any" or not config.VARIANTS:
        return None
    rows = con.execute("SELECT id, variant FROM chunks").fetchall()
    return {int(r[0]) for r in rows if variant_mod.matches(r[1], scope)}


def _fts_ids(
    con: sqlite3.Connection, query: str, limit: int, allowed: set[int] | None
) -> list[int]:
    match = _fts_match_query(query)
    if not match:
        return []
    # Weight matches in the title/section columns above body text, so a page whose
    # heading is literally about the query terms outranks one that merely mentions
    # them in passing. Weights map to the FTS columns in order: (text, title, section).
    # Filter BEFORE the LIMIT: truncating first would leave only a handful of hits
    # whenever the excluded edition dominates the lexical matches for a query.
    rows = con.execute(
        "SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? "
        "ORDER BY bm25(chunks_fts, 1.0, 5.0, 3.0)",
        (match,),
    ).fetchall()
    ids = [int(r[0]) for r in rows]
    if allowed is not None:
        ids = [i for i in ids if i in allowed]
    return ids[:limit]


def _vector_ids(
    vectors: np.ndarray, query_vec: np.ndarray, limit: int, allowed: set[int] | None
) -> list[int]:
    if vectors.shape[0] == 0:
        return []
    sims = vectors @ query_vec
    out: list[int] = []
    for index in np.argsort(-sims):
        rowid = int(index) + 1  # rowid = embedding index + 1
        if allowed is None or rowid in allowed:
            out.append(rowid)
            if len(out) == limit:
                break
    return out


def reciprocal_rank_fusion(rankings: list[list[int]], k: int = config.RRF_K) -> list[int]:
    """Fuse ranked id lists; higher fused score = better."""
    scores: dict[int, float] = {}
    for ranking in rankings:
        for rank, cid in enumerate(ranking):
            scores[cid] = scores.get(cid, 0.0) + 1.0 / (k + rank + 1)
    return sorted(scores, key=lambda cid: -scores[cid])


def search(
    query: str,
    *,
    k: int = config.DEFAULT_TOP_K,
    variant: str | None = None,
    db_path: Path | None = None,
    embeddings_path: Path | None = None,
    model: EmbeddingModel | None = None,
) -> list[Result]:
    """Run hybrid retrieval and return the top-k fused results.

    Paths default to the live ``config`` values, resolved at call time.
    """
    scope = variant or config.DEFAULT_VARIANT
    db_path = db_path or config.DB_PATH
    embeddings_path = embeddings_path or config.EMBEDDINGS_PATH
    con = sqlite3.connect(db_path)
    try:
        allowed = _allowed_ids(con, scope)
        fts = _fts_ids(con, query, config.CANDIDATE_POOL, allowed)
        vectors = np.load(embeddings_path)
        vec = _vector_ids(vectors, embed_query(query, model=model), config.CANDIDATE_POOL, allowed)
        fused = reciprocal_rank_fusion([fts, vec])[:k]
        if not fused:
            return []
        placeholders = ",".join("?" * len(fused))
        rows = {
            row[0]: row
            for row in con.execute(
                "SELECT id, title, url, source, variant, section, text FROM chunks "
                f"WHERE id IN ({placeholders})",
                fused,
            ).fetchall()
        }
    finally:
        con.close()

    results: list[Result] = []
    for rank, cid in enumerate(fused):
        row = rows.get(cid)
        if row is None:  # pragma: no cover - defensive; ids come from the same db
            continue
        results.append(
            Result(
                title=row[1],
                url=row[2],
                source=row[3],
                variant=row[4],
                section=row[5],
                text=row[6],
                score=1.0 / (rank + 1),
            )
        )
    return results


def _format(results: list[Result]) -> str:
    if not results:
        return "(no results)"
    blocks = []
    for i, r in enumerate(results, 1):
        snippet = r.text if len(r.text) <= 500 else r.text[:500] + "…"
        blocks.append(f"{i}. [{r.variant}] {r.title} > {r.section}\n   {r.url}\n   {snippet}")
    return "\n\n".join(blocks)


def main(argv: list[str] | None = None) -> int:
    args = sys.argv[1:] if argv is None else argv
    if not args:
        print('usage: run.sh search "<query>"', file=sys.stderr)
        return 2
    print(_format(search(" ".join(args))))
    return 0


if __name__ == "__main__":
    sys.exit(main())