troupe-standup/wizard-work/null/raglite/raglite.py

#!/usr/bin/env python3
"""raglite -- minimal retrieval for agents: inverted index + BM25 over passages.

Stdlib only. Deterministic. No network. No dependencies.

A "passage" is a blank-line-delimited block of a text/markdown file, addressed
as "<path>::<paragraph-index>". The index is plain JSON so it can be diffed,
cached and inspected by hand.

Library:
    idx = raglite.build_index(["./docs"])     # path may be a file or a directory
    raglite.save_index(idx, "index.json")
    idx = raglite.load_index("index.json")
    hits = raglite.search(idx, "bm25 ranking", k=3)

CLI:
    python3 raglite.py index ./docs --out index.json
    python3 raglite.py query index.json "bm25 ranking" -k 3 --json
"""

from __future__ import annotations

import argparse
import json
import math
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

TOKEN_RE = re.compile(r"[A-Za-z0-9_']+")
PASSAGE_SPLIT = re.compile(r"\n[ \t]*\n")
DEFAULT_EXTENSIONS = (".md", ".txt")


def tokenize(text: str) -> list[str]:
    """Lowercased alphanumeric tokens; punctuation is a separator."""
    return TOKEN_RE.findall(text.lower())


def split_passages(text: str) -> list[str]:
    """Split on blank lines, drop empties, keep original order."""
    return [p.strip() for p in PASSAGE_SPLIT.split(text) if p.strip()]


def build_index(paths, k1: float = 1.5, b: float = 0.75,
                extensions=DEFAULT_EXTENSIONS) -> dict:
    """Build a BM25 index over every passage of every file under `paths`."""
    postings: dict[str, Counter] = defaultdict(Counter)
    lengths: dict[str, int] = {}
    texts: dict[str, str] = {}
    n = 0

    for raw in paths:
        p = Path(raw)
        if p.is_dir():
            files = sorted(q for q in p.rglob("*")
                           if q.suffix.lower() in extensions)
        elif p.is_file():
            files = [p]
        else:
            raise FileNotFoundError(f"no such path: {raw}")

        for f in files:
            body = f.read_text(encoding="utf-8", errors="replace")
            for i, passage in enumerate(split_passages(body)):
                toks = tokenize(passage)
                if not toks:
                    continue
                pid = f"{f}::{i}"
                for term, tf in Counter(toks).items():
                    postings[term][pid] = tf
                lengths[pid] = len(toks)
                texts[pid] = passage
                n += 1

    total = sum(lengths.values())
    return {
        "k1": k1,
        "b": b,
        "n_docs": n,
        "avgdl": (total / n) if n else 0.0,
        "lengths": lengths,
        "texts": texts,
        "postings": {t: dict(c) for t, c in postings.items()},
    }


def _idf(df: int, n: int) -> float:
    """BM25 idf, Robertson's non-negative form."""
    return math.log(1.0 + (n - df + 0.5) / (df + 0.5))


def search(idx: dict, query: str, k: int = 5) -> list[dict]:
    """Rank passages for `query`; returns <= k hits, best first.

    Each hit: {"id", "score", "text"}. Score is the raw BM25 sum (unrounded).
    """
    k1, b = idx["k1"], idx["b"]
    n, avgdl = idx["n_docs"], idx["avgdl"]
    if not avgdl or k <= 0:
        return []

    scores: dict[str, float] = defaultdict(float)
    for term in set(tokenize(query)):
        post = idx["postings"].get(term)
        if not post:
            continue
        idf = _idf(len(post), n)
        for pid, tf in post.items():
            dl = idx["lengths"][pid]
            denom = tf + k1 * (1.0 - b + b * dl / avgdl)
            scores[pid] += idf * (tf * (k1 + 1.0)) / denom

    ranked = sorted(scores.items(), key=lambda kv: (-kv[1], kv[0]))
    return [{"id": pid, "score": s, "text": idx["texts"][pid]}
            for pid, s in ranked[:k]]


def save_index(idx: dict, path) -> None:
    Path(path).write_text(json.dumps(idx, sort_keys=True), encoding="utf-8")


def load_index(path) -> dict:
    return json.loads(Path(path).read_text(encoding="utf-8"))


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(prog="raglite", description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    pi = sub.add_parser("index", help="build an index over a file or directory")
    pi.add_argument("path")
    pi.add_argument("--out", required=True)
    pi.add_argument("--k1", type=float, default=1.5)
    pi.add_argument("--b", type=float, default=0.75)

    pq = sub.add_parser("query", help="rank passages in an index")
    pq.add_argument("index")
    pq.add_argument("query")
    pq.add_argument("-k", type=int, default=5)
    pq.add_argument("--json", action="store_true", help="emit JSON hits")

    args = ap.parse_args(argv)

    if args.cmd == "index":
        idx = build_index([args.path], k1=args.k1, b=args.b)
        save_index(idx, args.out)
        print(f"indexed {idx['n_docs']} passages from {args.path} -> {args.out}")
        return 0

    idx = load_index(args.index)
    hits = search(idx, args.query, k=args.k)
    if args.json:
        print(json.dumps(hits, indent=2))
    else:
        for i, h in enumerate(hits, 1):
            print(f"{i}. [{h['score']:.6f}] {h['id']}")
            print(f"   {h['text'][:160]}")
    return 0


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