RAG with citations you can check

import hashlib
import re
from dataclasses import dataclass

@dataclass(frozen=True)
class Chunk:
    id: str        # stable across re-ingestion, because citations outlive runs
    doc: str       # which file it came from
    heading: str   # heading path: shown in the citation, embedded with the text
    start: int     # character offset in the file, so a citation is checkable by hand
    text: str

HEADING = re.compile(r"^(#{1,6})\s+(.*)$")

def chunk_markdown(doc: str, text: str, max_chars: int = 1200) -> list[Chunk]:
    """Group paragraphs under their heading. A chunk never crosses a heading."""
    chunks: list[Chunk] = []
    heading, buf, start, pos = "", [], 0, 0

    def flush() -> None:
        nonlocal buf
        if not buf:
            return
        body = "\n\n".join(buf)
        # Content-addressed: an unchanged document re-ingests to the same ids,
        # so a citation printed last week still resolves to the same text.
        cid = hashlib.sha1(f"{doc}|{start}|{body}".encode()).hexdigest()[:12]
        chunks.append(Chunk(cid, doc, heading, start, body))
        buf = []

    for para in text.split("\n\n"):
        stripped = para.strip()
        match = HEADING.match(stripped)
        if match:
            flush()
            heading = match.group(2).strip()
        elif stripped:
            if sum(len(p) for p in buf) + len(stripped) > max_chars:
                flush()
            if not buf:
                start = pos
            buf.append(stripped)
        pos += len(para) + 2   # the blank line the split consumed

    flush()
    return chunks

import numpy as np

def embed(texts: list[str]) -> np.ndarray:
    """Your embedding model, batched. One request per batch, not per chunk."""
    out = [embed_batch(texts[i : i + 64]) for i in range(0, len(texts), 64)]
    return np.vstack(out).astype("float32")

# Embed the heading with the body. The chunk will be read alone, so it is
# indexed the way it will be read.
vectors = embed([f"{c.heading}\n\n{c.text}" for c in chunks])
vectors /= np.linalg.norm(vectors, axis=1, keepdims=True)

def search(query: str, k: int = 20) -> list[tuple[Chunk, float]]:
    q = embed([query])[0]
    q /= np.linalg.norm(q)
    scores = vectors @ q                      # every similarity in one multiply
    top = np.argsort(-scores)[:k]
    return [(chunks[i], float(scores[i])) for i in top]

# The cheapest correctness check there is: text lifted out of a chunk must
# retrieve that chunk.
probe = chunks[100]
assert search(probe.text[:200], k=1)[0][0].id == probe.id

def inspect(query: str, k: int = 5) -> None:
    for rank, (chunk, score) in enumerate(search(query, k), 1):
        print(f"{rank}. {score:.3f}  {chunk.doc} > {chunk.heading}  [{chunk.id}]")
        print("   " + chunk.text[:300].replace("\n", " "))
        print()

for q in QUESTIONS:
    print("=" * 72)
    print(q)
    print("=" * 72)
    inspect(q)

question                                    answerable  note
Can I expense a train ticket?               yes         rank 1 is the travel policy
What is the notice period for interns?      no          all five are the full-time policy
Which regions is plan Pro available in?     partial     right page, the table came out as prose
What does error ERR_4013 mean?              no          matched every other error code
How do I rotate an API key?                 yes         rank 3, ranks 1-2 are release notes

SYSTEM = """Answer the question using only the sources below.

Every sentence of your answer must end with the ids of the sources it came from,
written as [id]. Do not use anything you know that is not in the sources.
If the sources do not answer the question, reply with exactly: NOT_IN_SOURCES."""

def build_context(hits: list[Chunk]) -> str:
    return "\n\n".join(f"[{c.id}] {c.doc} > {c.heading}\n{c.text}" for c in hits)

CITE = re.compile(r"\[([0-9a-f]{12})\]")

def bad_citations(answer: str, hits: list[Chunk]) -> list[str]:
    """Ids the answer cites that were never sent. Non-empty means the model
    invented a source, which is a bug, not a wording problem."""
    allowed = {c.id for c in hits}
    return [cid for cid in CITE.findall(answer) if cid not in allowed]

def uncited_sentences(answer: str) -> list[str]:
    sentences = re.split(r"(?<=[.!?])\s+", answer.strip())
    return [s for s in sentences if s and not CITE.search(s)]

from rank_bm25 import BM25Okapi

tokenised = [f"{c.heading} {c.text}".lower().split() for c in chunks]
bm25 = BM25Okapi(tokenised)

def candidates(query: str, k: int = 40) -> list[Chunk]:
    """Union of both retrievers. No score blending: the two scores are not
    on the same scale, and the reranker is about to decide the order anyway."""
    dense = [c for c, _ in search(query, k)]
    lexical_ranks = np.argsort(-bm25.get_scores(query.lower().split()))[:k]
    seen, out = set(), []
    for c in dense + [chunks[i] for i in lexical_ranks]:
        if c.id not in seen:
            seen.add(c.id)
            out.append(c)
    return out

def rerank(query: str, pool: list[Chunk], k: int = 5) -> list[Chunk]:
    """A cross-encoder reads query and chunk together, so it can separate
    'refunds on annual plans' from 'refunds on monthly plans'. Far too slow
    to run over the corpus, which is exactly why it runs second."""
    scores = cross_encoder.predict([(query, f"{c.heading}\n{c.text}") for c in pool])
    ranked = sorted(zip(scores, pool), key=lambda pair: -pair[0])
    return [c for _, c in ranked[:k]]

# cases.jsonl, one per line:
# {"q": "How do I rotate an API key?", "gold": ["4c1f9ab207de", "9930ba7715c1"]}

def recall_at_k(cases: list[dict], k: int) -> float:
    """Did at least one chunk that contains the answer survive into the top k?"""
    hits = 0
    for case in cases:
        got = {c.id for c in rerank(case["q"], candidates(case["q"]), k)}
        hits += bool(got & set(case["gold"]))
    return hits / len(cases)

for k in (1, 5, 20):
    print(f"recall@{k}: {recall_at_k(cases, k):.2f}")

FAITHFUL = """Source:
{source}

Sentence:
{sentence}

Does the source state the sentence? Reply with one word: YES or NO."""

def faithfulness(answer: str, by_id: dict[str, Chunk]) -> float:
    verdicts = []
    for sentence in re.split(r"(?<=[.!?])\s+", answer.strip()):
        ids = CITE.findall(sentence)
        if not ids:
            continue
        source = "\n\n".join(by_id[i].text for i in ids)
        reply = ask(FAITHFUL.format(source=source, sentence=sentence))
        verdicts.append(reply.strip().upper().startswith("YES"))
    return sum(verdicts) / len(verdicts) if verdicts else 0.0

def answer(query: str, min_score: float = 0.35, min_hits: int = 2) -> dict:
    pool = candidates(query)
    top = rerank(query, pool, k=5)
    scores = cross_encoder.predict([(query, f"{c.heading}\n{c.text}") for c in top])

    # One weak hit is the shape of a question the corpus does not cover.
    strong = [c for c, s in zip(top, scores) if s >= min_score]
    if len(strong) < min_hits:
        return {"abstained": True, "reason": "no strong match", "sources": []}

    text = generate(SYSTEM, build_context(strong), query)
    if text.strip() == "NOT_IN_SOURCES":
        return {"abstained": True, "reason": "model found no answer", "sources": strong}

    assert not bad_citations(text, strong)
    return {"abstained": False, "answer": text, "sources": strong}