Evaluate a RAG Pipeline: Retrieval Metrics, Faithfulness and a Regression Gate

"""rag.py: the pipeline under evaluation. Retrieval + generation over kb/*.md.

    python3 rag.py index --size 500 --overlap 60          # chunk + embed, writes index.json
    python3 rag.py ask "How long is a password reset link valid?" --k 5
    python3 rag.py run --qa qa.jsonl --out runs/mine.jsonl --size 500 --k 5 [--no-answers] [--model ...]

Chunks carry provenance: the paragraph ids they cover, which is what qa.jsonl labels.
"""
import argparse
import json
import math
import os
import re
import sys
import time
import urllib.request

BASE_URL = os.environ.get("RAG_BASE_URL", "http://nim-proxy.labs.svc:8080/v1")
EMBED_MODEL = os.environ.get("RAG_EMBED_MODEL", "nvidia/llama-nemotron-embed-vl-1b-v2")
EMBED_DIMENSIONS = int(os.environ.get("RAG_EMBED_DIMENSIONS", "1024"))
CHAT_MODEL = os.environ.get("RAG_CHAT_MODEL", "meta/llama-3.1-8b-instruct")
KB_DIR = os.environ.get("RAG_KB", "kb")
INDEX_PATH = "index.json"

SYSTEM = ("You are the Northwind Cloud support assistant. Answer the question using ONLY the numbered "
          "context passages. Be specific: quote the numbers, durations and menu paths the passages give. "
          "If the passages do not contain the answer, reply exactly: "
          "\"I do not have that information in the knowledge base.\" Answer in at most three sentences, "
          "as a direct reply: never mention the passages, their numbers or the word context.")


# ------------------------------------------------------------------ corpus
def load_corpus(kb_dir=KB_DIR):
    """[{doc, title, paragraphs}] with the title line removed from the paragraphs."""
    docs = []
    for name in sorted(os.listdir(kb_dir)):
        if not name.endswith(".md"):
            continue
        raw = open(os.path.join(kb_dir, name), encoding="utf-8").read()
        title, body = name[:-3], raw
        m = re.match(r"^---\n(.*?)\n---\n", raw, re.S)
        if m:
            body = raw[m.end():]
            t = re.search(r"^source:\s*(.+)$", m.group(1), re.M)
            title = t.group(1).strip() if t else title
        paragraphs = [p.strip() for p in re.split(r"\n\s*\n", body) if p.strip() and not p.strip().startswith("# ")]
        docs.append({"doc": name[:-3], "title": title, "paragraphs": paragraphs})
    return docs


def chunk_doc(doc_id, paragraphs, size=500, overlap=60):
    """Split one document's paragraphs into chunks of at most `size` characters.

    Paragraph boundaries come first: whole paragraphs are packed into a chunk until the next one
    would not fit. A paragraph longer than `size` on its own is hard-split into windows of `size`
    characters that overlap by `overlap`. Every chunk records the paragraph ids it covers
    (`<doc>:p<n>`, 1-based) so retrieval can be scored against paragraph-level labels."""
    # TODO (Step 01): pack whole paragraphs into chunks of at most `size` characters; a paragraph
    #   longer than `size` is hard-split into windows of `size` with `overlap`; every chunk is
    #   {"doc": doc_id, "paragraphs": [<doc>:p<n>, ...], "text": ...}; number them "<doc>:c<i>" at the end.
    raise NotImplementedError("Step 01: implement chunk_doc()")


def chunk_corpus(docs, size, overlap):
    out = []
    for d in docs:
        for c in chunk_doc(d["doc"], d["paragraphs"], size, overlap):
            c["title"] = d["title"]
            out.append(c)
    return out


# ------------------------------------------------------------------ embeddings + search
def embed(texts, input_type):
    """One batched request per 50 texts through the lab proxy (rate-limited, so never per chunk)."""
    out, texts = [], list(texts)
    for i in range(0, len(texts), 50):
        body = json.dumps({"model": EMBED_MODEL, "input": texts[i:i + 50], "input_type": input_type,
                           "dimensions": EMBED_DIMENSIONS, "encoding_format": "float"}).encode()
        req = urllib.request.Request(BASE_URL.rstrip("/") + "/embeddings", data=body, method="POST",
                                     headers={"Content-Type": "application/json"})
        data = json.loads(urllib.request.urlopen(req, timeout=90).read())
        out.extend(item["embedding"] for item in sorted(data["data"], key=lambda x: x["index"]))
    return out


def _norm(v):
    n = math.sqrt(sum(x * x for x in v)) or 1.0
    return [x / n for x in v]


def build_index(size, overlap, kb_dir=KB_DIR):
    chunks = chunk_corpus(load_corpus(kb_dir), size, overlap)
    vecs = embed([f"{c['title']}\n{c['text']}" for c in chunks], "passage")
    return {"size": size, "overlap": overlap, "chunks": chunks, "vectors": [_norm(v) for v in vecs]}


def search(index, qvec, k):
    q = _norm(qvec)
    scored = [(sum(a * b for a, b in zip(q, v)), i) for i, v in enumerate(index["vectors"])]
    scored.sort(reverse=True)
    return [dict(index["chunks"][i], score=round(s, 4)) for s, i in scored[:k]]


# ------------------------------------------------------------------ generation
def generate(question, contexts, model=CHAT_MODEL):
    ctx = "\n\n".join(f"[{i + 1}] ({c['title']}) {c['text']}" for i, c in enumerate(contexts))
    body = json.dumps({"model": model, "temperature": 0, "max_tokens": 220,
                       "messages": [{"role": "system", "content": SYSTEM},
                                    {"role": "user", "content": f"CONTEXT:\n{ctx}\n\nQUESTION: {question}"}]}).encode()
    req = urllib.request.Request(BASE_URL.rstrip("/") + "/chat/completions", data=body, method="POST",
                                 headers={"Content-Type": "application/json"})
    data = json.loads(urllib.request.urlopen(req, timeout=120).read())
    return data["choices"][0]["message"]["content"].strip(), data.get("model", model)


# ------------------------------------------------------------------ CLI
def read_jsonl(path):
    with open(path, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


def cmd_index(a):
    idx = build_index(a.size, a.overlap)
    json.dump(idx, open(INDEX_PATH, "w"))
    docs = len({c["doc"] for c in idx["chunks"]})
    multi = sum(1 for c in idx["chunks"] if len(c["paragraphs"]) > 1)
    print(f"indexed {len(idx['chunks'])} chunks from {docs} documents (size {a.size}, overlap {a.overlap}); "
          f"{multi} chunks span more than one paragraph")


def cmd_ask(a):
    if not os.path.exists(INDEX_PATH):
        sys.exit("no index.json: run `python3 rag.py index` first")
    idx = json.load(open(INDEX_PATH))
    hits = search(idx, embed([a.question], "query")[0], a.k)
    for i, h in enumerate(hits, start=1):
        print(f"{i}. {h['score']:.3f}  {h['id']:26} paragraphs={','.join(p.split(':')[1] for p in h['paragraphs'])}  {h['text'][:70]!r}")
    if not a.no_answer:
        answer, served = generate(a.question, hits, a.model)
        print(f"\n[{served}] {answer}")


def cmd_run(a):
    qa = read_jsonl(a.qa)
    idx = build_index(a.size, a.overlap)
    qvecs = embed([q["question"] for q in qa], "query")
    os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
    with open(a.out, "w", encoding="utf-8") as f:
        for q, qv in zip(qa, qvecs):
            hits = search(idx, qv, a.k)
            rec = {"id": q["id"], "question": q["question"], "config": {"size": a.size, "overlap": a.overlap, "k": a.k, "chunks": len(idx["chunks"])},
                   "retrieved": [{"id": h["id"], "doc": h["doc"], "paragraphs": h["paragraphs"], "score": h["score"],
                                  "title": h["title"], "text": h["text"]} for h in hits]}
            if not a.no_answers:
                rec["answer"], rec["model"] = generate(q["question"], hits, a.model)
                print(f"{q['id']}: {rec['answer'][:90]!r}")
                if a.pace:
                    time.sleep(a.pace)
            f.write(json.dumps(rec) + "\n")
    print(f"wrote {len(qa)} records to {a.out}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("index"); p.add_argument("--size", type=int, default=500); p.add_argument("--overlap", type=int, default=60); p.set_defaults(fn=cmd_index)
    p = sub.add_parser("ask"); p.add_argument("question"); p.add_argument("--k", type=int, default=5); p.add_argument("--model", default=CHAT_MODEL); p.add_argument("--no-answer", action="store_true"); p.set_defaults(fn=cmd_ask)
    p = sub.add_parser("run"); p.add_argument("--qa", default="qa.jsonl"); p.add_argument("--out", required=True); p.add_argument("--size", type=int, default=500)
    p.add_argument("--overlap", type=int, default=60); p.add_argument("--k", type=int, default=5); p.add_argument("--model", default=CHAT_MODEL); p.add_argument("--no-answers", action="store_true"); p.add_argument("--pace", type=float, default=0.0); p.set_defaults(fn=cmd_run)
    a = ap.parse_args(); a.fn(a)