From document ingestion to cited RAG answers

"""Actual parsing, SQL filtering, vector search, BM25, RRF and contract evaluation.

Embeddings are replayed from an actual local-model inference snapshot. The generation
stage uses explicitly authored fixtures. No network or inference is needed for replay.
"""
from collections import Counter
import json
import math
import re
import sqlite3

from application_data import FIXTURES, corpus_hash, load_chunks, read_cases
from prompt_iteration import CANDIDATES, build_prompt, grade


def load_snapshot(chunks, path=None):
    snapshot = json.loads((path or FIXTURES / "embeddings.json").read_text(encoding="utf-8"))
    if snapshot.get("schemaVersion") != 1 or snapshot.get("corpusSha256") != corpus_hash(chunks):
        raise ValueError("embedding snapshot is stale; rebuild after changing documents")
    if set(snapshot["documents"]) != {chunk["id"] for chunk in chunks}:
        raise ValueError("embedding snapshot has different chunk identities")
    dimension = snapshot["dimension"]
    if type(dimension) is not int or dimension < 1 or not snapshot.get("model"):
        raise ValueError("invalid embedding contract")
    for vector in [*snapshot["documents"].values(), *snapshot["queries"].values()]:
        if (len(vector) != dimension or not all(type(x) in (int, float) and math.isfinite(x) for x in vector)
                or not any(vector)):
            raise ValueError("invalid embedding vector")
    return snapshot


def tokenize(text):
    return re.findall(r"[a-z0-9]+", text.lower())


def cosine(left, right):
    return sum(a * b for a, b in zip(left, right)) / (math.sqrt(sum(a * a for a in left)) * math.sqrt(sum(b * b for b in right)))


def keyword_scores(query, chunks):
    terms = set(tokenize(query))
    bags = [Counter(tokenize(chunk["text"])) for chunk in chunks]
    average = sum(sum(bag.values()) for bag in bags) / len(bags)
    scores = []
    for chunk, bag in zip(chunks, bags):
        score = 0.0
        for term in terms:
            frequency = bag[term]
            if frequency:
                df = sum(term in other for other in bags)
                idf = math.log(1 + (len(bags) - df + 0.5) / (df + 0.5))
                score += idf * frequency * 2.2 / (frequency + 1.2 * (0.25 + 0.75 * sum(bag.values()) / average))
        scores.append((chunk["id"], score))
    return sorted((row for row in scores if row[1] > 0), key=lambda row: (-row[1], row[0]))


def index_chunks(chunks):
    db = sqlite3.connect(":memory:")
    db.execute("CREATE TABLE chunks(id TEXT PRIMARY KEY, tenant TEXT NOT NULL, current INTEGER NOT NULL, body TEXT NOT NULL)")
    db.executemany("INSERT INTO chunks VALUES(?,?,?,?)", [(chunk["id"], chunk["tenant"], int(chunk["current"]), json.dumps(chunk)) for chunk in chunks])
    return db


def retrieve(query, chunks, snapshot, tenant="shop-a", top_k=3):
    if type(top_k) is not int or top_k < 1:
        raise ValueError("top_k must be a positive integer")
    if query not in snapshot["queries"]:
        raise ValueError("query was not recorded; regenerate embeddings for the changed query")
    db = index_chunks(chunks)
    try:
        allowed = [json.loads(row[0]) for row in db.execute("SELECT body FROM chunks WHERE tenant=? AND current=1 ORDER BY id", (tenant,))]
    finally:
        db.close()
    if not allowed:
        return {"allowed": 0, "dense": [], "keyword": [], "context": []}
    query_vector = snapshot["queries"][query]
    dense = sorted([(chunk["id"], cosine(query_vector, snapshot["documents"][chunk["id"]])) for chunk in allowed], key=lambda row: (-row[1], row[0]))
    keyword = keyword_scores(query, allowed)
    fused = Counter()
    for ranking in (dense, keyword):
        for rank, (chunk_id, _) in enumerate(ranking, 1):
            fused[chunk_id] += 1 / (60 + rank)
    selected = [cid for cid, _ in sorted(fused.items(), key=lambda row: (-row[1], row[0]))[:top_k]]
    by_id = {chunk["id"]: chunk for chunk in allowed}
    return {"allowed": len(allowed), "dense": [cid for cid, _ in dense], "keyword": [cid for cid, _ in keyword], "context": [by_id[cid] for cid in selected]}


def run_pipeline(top_k=3):
    chunks = load_chunks()
    snapshot = load_snapshot(chunks)
    rows = []
    for case in read_cases():
        retrieved = retrieve(case["query"], chunks, snapshot, top_k=top_k)
        request = build_prompt(case, retrieved["context"])
        answer = CANDIDATES["v2"][case["id"]]
        context_ids = [chunk["id"] for chunk in retrieved["context"]]
        required = set(case["requiredEvidence"])
        rows.append({"case": case["id"], "contextIds": context_ids,
                     "request": request, "candidate": answer,
                     "evidenceRecall": len(required.intersection(context_ids)) / len(required) if required else None,
                     **grade(case, answer, retrieved["context"])})
    return {"chunks": chunks, "snapshot": snapshot, "rows": rows}


def demo():
    report = run_pipeline()
    answerable = [row["evidenceRecall"] for row in report["rows"] if row["evidenceRecall"] is not None]
    return {"chunks": len(report["chunks"]), "dimension": report["snapshot"]["dimension"],
            "embeddingModel": report["snapshot"]["model"], "queries": len(report["rows"]),
            "contextIds": {row["case"]: row["contextIds"] for row in report["rows"]},
            "contractPasses": sum(row["contractPassed"] for row in report["rows"]),
            "meanNecessaryEvidenceRecallAt3": sum(answerable) / len(answerable),
            "semanticQuality": "not_scored", "generation": "authored_fixtures",
            "embedding": "real_local_inference_snapshot"}


if __name__ == "__main__":
    print(json.dumps(demo(), ensure_ascii=False, sort_keys=True))