aivar/rules.py

"""Stage 4 — IFAB rule retrieval. Fully local/offline.

Hybrid retrieval over the IFAB Laws of the Game 2025/26:
HNSW vector search (fastembed int8 ONNX, no torch) + SQLite FTS5 BM25,
fused with Reciprocal Rank Fusion. Index is prebuilt by
scripts/build_rules_index.py and committed to the repo.
"""
import json
import re
import sqlite3

import numpy as np

from .config import (EMBED_MODEL, INDEX_META_PATH, INDEX_PATH, RULES_DB,
                     TOP_K_BM25, TOP_K_FINAL, TOP_K_VECTOR)
from .schemas import IncidentType, RuleChunk

RRF_K = 60
_embedder = None
_index = None

SCHEMA = """
CREATE TABLE IF NOT EXISTS chunks(
    id INTEGER PRIMARY KEY,
    law TEXT NOT NULL,
    section TEXT NOT NULL,
    text TEXT NOT NULL,
    page INTEGER NOT NULL
);
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
    text, content='chunks', content_rowid='id'
);
"""

# Which laws matter per incident type — used to boost retrieval queries.
INCIDENT_LAW_HINTS: dict[IncidentType, str] = {
    IncidentType.SLIDING_TACKLE: "Law 12 fouls tackles challenges careless reckless excessive force",
    IncidentType.STANDING_TACKLE: "Law 12 fouls tackles challenges careless reckless",
    IncidentType.HANDBALL: "Law 12 handball offence hand arm unnaturally bigger",
    IncidentType.PUSH: "Law 12 pushes an opponent direct free kick",
    IncidentType.HOLDING: "Law 12 holds an opponent holding offence",
    IncidentType.SHIRT_PULL: "Law 12 holds an opponent holding shirt",
    IncidentType.DANGEROUS_PLAY: "Law 12 playing in a dangerous manner indirect free kick",
    IncidentType.SIMULATION: "Law 12 simulation attempts to deceive the referee yellow card caution",
    IncidentType.OFFSIDE: "Law 11 offside position interfering with play gaining advantage",
    IncidentType.DOGSO: "Law 12 denying an obvious goal-scoring opportunity DOGSO sending-off",
    IncidentType.VIOLENT_CONDUCT: "Law 12 violent conduct excessive force brutality not challenging for the ball",
    IncidentType.SERIOUS_FOUL_PLAY: "Law 12 serious foul play excessive force endangers the safety of an opponent",
    IncidentType.NO_INCIDENT: "Law 12 fouls and misconduct fair challenge",
}


def connect() -> sqlite3.Connection:
    con = sqlite3.connect(RULES_DB)
    con.row_factory = sqlite3.Row
    con.executescript(SCHEMA)
    return con


def embedder():
    global _embedder
    if _embedder is None:
        from fastembed import TextEmbedding
        _embedder = TextEmbedding(model_name=EMBED_MODEL)
    return _embedder


def embed_texts(texts: list[str]) -> np.ndarray:
    vecs = np.array(list(embedder().embed(texts)), dtype=np.float32)
    vecs /= np.linalg.norm(vecs, axis=1, keepdims=True) + 1e-12
    return vecs


def vector_index(dim: int):
    global _index
    if _index is None:
        import hnswlib
        meta = json.loads(INDEX_META_PATH.read_text())
        idx = hnswlib.Index(space="cosine", dim=meta["dim"])
        idx.load_index(str(INDEX_PATH), max_elements=meta["max_elements"])
        idx.set_ef(128)
        _index = idx
    return _index


def _fts_query(text: str) -> str:
    words = re.findall(r"[a-zA-Z]{3,}", text)
    return " OR ".join(dict.fromkeys(w.lower() for w in words)) or "foul"


def retrieve(incident: IncidentType, description: str, k: int = TOP_K_FINAL) -> list[RuleChunk]:
    """Hybrid RRF retrieval keyed off incident type + Gemini's description."""
    query = f"{INCIDENT_LAW_HINTS.get(incident, '')} {description}".strip()
    con = connect()
    try:
        qvec = embed_texts([query])[0]
        idx = vector_index(qvec.shape[0])
        n = min(TOP_K_VECTOR, idx.element_count)
        labels, _ = idx.knn_query(qvec, k=n)
        vec_ids = [int(l) for l in labels[0]]

        bm_rows = con.execute(
            "SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?",
            (_fts_query(query), TOP_K_BM25),
        ).fetchall()
        bm_ids = [r["rowid"] for r in bm_rows]

        scores: dict[int, float] = {}
        for ranking in (vec_ids, bm_ids):
            for rank, cid in enumerate(ranking):
                scores[cid] = scores.get(cid, 0.0) + 1.0 / (RRF_K + rank + 1)

        top = sorted(scores, key=lambda c: -scores[c])[:k]
        if not top:
            return []
        marks = ",".join("?" * len(top))
        rows = {r["id"]: r for r in con.execute(
            f"SELECT * FROM chunks WHERE id IN ({marks})", top)}
        return [RuleChunk(law=rows[c]["law"], section=rows[c]["section"],
                          text=rows[c]["text"], page=rows[c]["page"],
                          score=scores[c]) for c in top if c in rows]
    finally:
        con.close()


def format_context(chunks: list[RuleChunk]) -> str:
    parts = []
    for c in chunks:
        parts.append(f"[{c.law} — {c.section} (p.{c.page})]\n{c.text}")
    return "\n\n".join(parts)