Corrective RAG: grade what you retrieved before you answer

"""
Corrective RAG: grade what you retrieved before you answer from it

Plain RAG trusts the retriever. Corrective RAG does not: every retrieved
passage is graded for relevance, the ones that fail are discarded, and the
answer is written only from what survived.

    Retrieve -> Grade each passage -> Discard failures -> Answer from survivors

The point is not better retrieval. The point is that a keyword-ish retriever
returns plausible-looking passages that are about the wrong thing, and those
passages will otherwise be summarised into a confident, wrong answer.

If nothing survives grading, this refuses to answer. That refusal is the
feature. An agent that answers anyway is the failure mode the pattern exists
to prevent.

Run: python corrective_rag.py
"""

from typing import Literal

from dotenv import load_dotenv
from openai import OpenAI
from pydantic import BaseModel, Field

load_dotenv()

client = OpenAI()

MODEL = "gpt-5.4-mini"
TOP_K = 4

# --------------------------------------------------------------
# Corpus. Deliberately seeded with near-miss distractors: several
# passages share vocabulary with the questions but answer something else.
# --------------------------------------------------------------

CORPUS: dict[str, str] = {
    "retry-policy": (
        "Failed jobs retry three times with exponential backoff starting at 2s. "
        "After the third failure the job moves to the dead-letter queue."
    ),
    "retry-ui-copy": (
        "The retry button in the dashboard is labelled 'Try again' and appears "
        "only after a job has finished failing. Copy is owned by the design team."
    ),
    "queue-limits": (
        "The dead-letter queue holds 10,000 messages. Beyond that, the oldest "
        "messages are dropped and a PagerDuty alert fires."
    ),
    "auth-tokens": (
        "Service tokens expire after 24 hours. Refresh happens automatically on "
        "the next authenticated call; there is no background refresh loop."
    ),
    "token-budget": (
        "Each tenant has a monthly token budget. Exceeding it does not block "
        "requests; it raises a soft warning on the billing page."
    ),
}


def retrieve(question: str, k: int = TOP_K) -> list[str]:
    """Deliberately naive term-overlap retrieval, so grading has real work."""
    terms = {t for t in question.lower().replace("?", "").split() if len(t) > 3}
    scored = [
        (len(terms & set(f"{key} {text}".lower().split())), key)
        for key, text in CORPUS.items()
    ]
    return [key for score, key in sorted(scored, reverse=True) if score > 0][:k]


# --------------------------------------------------------------
# Grade: one call per passage, so a bad passage cannot hide in a batch
# --------------------------------------------------------------


class Grade(BaseModel):
    verdict: Literal["relevant", "irrelevant"]
    why: str = Field(description="One sentence. Required for both verdicts.")


def grade(question: str, passage_id: str, passage: str) -> Grade:
    completion = client.beta.chat.completions.parse(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Decide whether the passage contains information that helps "
                    "answer the question. Sharing vocabulary is not relevance. "
                    "If it is about a neighbouring topic, mark it irrelevant."
                ),
            },
            {
                "role": "user",
                "content": f"Question: {question}\n\n[{passage_id}] {passage}",
            },
        ],
        response_format=Grade,
    )
    return completion.choices[0].message.parsed


# --------------------------------------------------------------
# Answer: only from what survived
# --------------------------------------------------------------


def answer_from(question: str, kept: list[tuple[str, str]]) -> str:
    context = "\n".join(f"[{pid}] {text}" for pid, text in kept)
    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Answer using ONLY the passages provided. Cite the passage id "
                    "in square brackets for each claim. If the passages do not "
                    "contain the answer, say so plainly instead of inferring."
                ),
            },
            {"role": "user", "content": f"Question: {question}\n\n{context}"},
        ],
    )
    return completion.choices[0].message.content or ""


def run(question: str) -> str:
    retrieved = retrieve(question)
    print(f"  retrieved: {retrieved or '(nothing)'}")

    if not retrieved:
        return "No passage matched the question. Not answering from memory."

    kept: list[tuple[str, str]] = []
    for passage_id in retrieved:
        # Printed before the call so a slow grading pass reads as work in
        # progress rather than a stall.
        print(f"  grading  : {passage_id}...")
        verdict = grade(question, passage_id, CORPUS[passage_id])
        print(f"      {verdict.verdict:<10} {verdict.why}")
        if verdict.verdict == "relevant":
            kept.append((passage_id, CORPUS[passage_id]))

    if not kept:
        # The whole point of the pattern. Answering here would be the bug.
        message = (
            f"All {len(retrieved)} retrieved passages were graded irrelevant. "
            "Refusing to answer rather than inferring from unrelated context."
        )
        print(f"  REFUSED  : {message}")
        return message

    print(f"  kept     : {[pid for pid, _ in kept]}")
    result = answer_from(question, kept)
    print(f"  answer   : {result}")
    return result


if __name__ == "__main__":
    for question in [
        "How many times does a failed job retry before it is dead-lettered?",
        "What is our policy on refunding annual subscriptions?",
    ]:
        print(f"\n> {question}")
        run(question)