Build a RAG pipeline in Python

import re
from dataclasses import dataclass, field
from html import escape

from fake_openai import fake_client
from mini_rag import HybridIndex
from site_docs import COURSE_NAMES, load_sections

INSTRUCTIONS = """You answer questions about the course lessons in the <source> tags.
- Use only the sources. Cite every claim inline as [1], [2].
- If the sources do not answer the question, reply exactly: I don't know.
- At most 4 sentences."""
PROMPT_VERSION = "answer-v3"

@dataclass
class Answer:
    text: str
    sources: list = field(default_factory=list)
    usage: dict = field(default_factory=dict)
    grounded: bool = True

class RagPipeline:
    def __init__(self, client, index, model="gpt-5.6-luna", budget_words=700):
        self.client, self.index, self.model, self.budget = client, index, model, budget_words

    def retrieve(self, question):
        chosen, seen, used = [], set(), 0
        for hit in self.index.search(question, k=20):
            lesson = hit.doc["url"].split("#")[0]
            words = len(hit.doc["text"].split())
            if lesson in seen or (chosen and used + words > self.budget):
                continue                                   # one chunk per lesson, within budget
            chosen.append(hit.doc)
            seen.add(lesson)
            used += words
        return chosen

    @staticmethod
    def build_input(question, docs):
        blocks = [f'<source id="{i}" title="{escape(COURSE_NAMES[d["course"]])} › {escape(d["title"])}">\n'
                  f'{escape(d["text"], quote=False)}\n</source>' for i, d in enumerate(docs, 1)]
        return "\n".join(blocks) + f"\n<question>{escape(question, quote=False)}</question>"

    def answer(self, question):
        docs = self.retrieve(question)
        if not docs:
            return Answer("I don't know.", grounded=False)
        r = self.client.responses.create(model=self.model, instructions=INSTRUCTIONS,
                                         input=self.build_input(question, docs),
                                         metadata={"prompt_version": PROMPT_VERSION})
        cited = sorted({int(n) for n in re.findall(r"\[(\d+)\]", r.output_text)})
        valid = [n for n in cited if 1 <= n <= len(docs)]
        usage = {"input": r.usage.input_tokens, "output": r.usage.output_tokens}
        if not valid or len(valid) != len(cited):            # no citation, or a made-up one
            return Answer("I don't know.", usage=usage, grounded=False)
        return Answer(r.output_text, [docs[n - 1] for n in valid], usage)

def stand_in_model(body):
    """Offline stand-in for the LLM: quotes the first prose sentence of source 1, cited."""
    m = re.search(r'<source id="1"[^>]*>\n(.*?)\n</source>', body["input"], re.S)
    for line in (m.group(1).splitlines() if m else []):
        sentence = re.match(r"([A-Z][^.!?]{20,}[.!?])(\s|$)", line)     # skips code and diagrams
        if sentence:
            return f"{sentence.group(1)} [1]"
    return "I don't know."

rag = RagPipeline(fake_client(responder=stand_in_model), HybridIndex(load_sections()))
for q in ["What is a covering index?", "What does the GIL do?", "When should I use a pandas UDF?",
          "What does the KV cache store?"]:
    a = rag.answer(q)
    print(f"Q: {q}\nA: {a.text}")
    print("   sources:", [s["url"] for s in a.sources], "| tokens:", a.usage)

import re

from fake_openai import fake_client

docs = [{"url": "sql/26-indexes.html#covering", "text": "A covering index contains every column the query needs."}]
client = fake_client(script=["Covering indexes avoid table lookups [1]. They also compress data by 90% [4]."])

reply = client.responses.create(model="gpt-5.6-luna", input="...").output_text
cited = sorted({int(n) for n in re.findall(r"\[(\d+)\]", reply)})
invalid = [n for n in cited if not 1 <= n <= len(docs)]
print("reply:  ", reply)
print("cited:  ", cited, "| invalid:", invalid)
print("verdict:", "reject, not grounded" if invalid else "ok")