Reflection: draft, critique, revise

"""
Reflection: draft, critique, revise

The cheapest reliability win available. One model drafts, a second pass
critiques that draft against a stated rubric, and a third rewrites using the
critique. Loop until the score clears the bar or the budget runs out.

    Draft -> Critique (scored) -> Revise -> ... -> Accept or exhaust

The critic is a separate call with a separate instruction, not the same call
asked to "check your work". A model grading its own output inside one
completion tends to agree with itself.

Two things worth knowing before you reach for this:

  - It costs you a multiple of the latency and tokens of a single call. Spend
    it where quality matters more than speed, not everywhere.
  - Scores from a model cluster hard around the middle of whatever range you
    give it. A 1 to 10 scale mostly returns 6, 7 and 8. Ask for specific
    boolean checks instead when you need a real gate.

Run: python reflection.py
"""

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

load_dotenv()

client = OpenAI()

MODEL = "gpt-5.4-mini"
MAX_ROUNDS = 3

#: What "good" means here, stated once so the critic and the reviser cannot
#: drift apart. Editing this changes the gate for both.
RUBRIC = """
1. Every claim is supported by something in the brief. No invented specifics.
2. States what is NOT covered, rather than implying full coverage.
3. Concrete nouns and verbs. No filler adjectives.
4. Under 120 words.
""".strip()


class Critique(BaseModel):
    """Boolean checks, not a 1-to-10 score.

    A model asked for a score returns the middle of the range almost
    regardless of quality. Specific checks force it to commit to something
    falsifiable, and they give the reviser an actionable list.
    """

    passes_rubric: bool = Field(description="True only if every numbered rule holds.")
    failing_rules: list[int] = Field(
        description="Rule numbers that fail. Empty when passes_rubric is true."
    )
    what_to_change: str = Field(
        description="Concrete instruction for the rewrite. Empty if nothing to change."
    )


def draft(brief: str) -> str:
    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": "Write a short internal summary for the given brief.",
            },
            {"role": "user", "content": brief},
        ],
    )
    return completion.choices[0].message.content or ""


def critique(brief: str, text: str) -> Critique:
    """A separate call with an adversarial instruction, not self-review."""
    completion = client.beta.chat.completions.parse(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "You are reviewing someone else's draft against a rubric. "
                    "You did not write it and you gain nothing by approving it. "
                    "Mark a rule as failing whenever you are unsure it holds.\n\n"
                    "Rubric:\n" + RUBRIC
                ),
            },
            {"role": "user", "content": f"Brief:\n{brief}\n\nDraft:\n{text}"},
        ],
        response_format=Critique,
    )
    return completion.choices[0].message.parsed


def revise(brief: str, text: str, verdict: Critique) -> str:
    completion = client.chat.completions.create(
        model=MODEL,
        messages=[
            {
                "role": "system",
                "content": (
                    "Rewrite the draft to fix exactly the listed problems. Change "
                    "nothing else. Do not add new claims.\n\nRubric:\n" + RUBRIC
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Brief:\n{brief}\n\nDraft:\n{text}\n\n"
                    f"Failing rules: {verdict.failing_rules}\n"
                    f"Required change: {verdict.what_to_change}"
                ),
            },
        ],
    )
    return completion.choices[0].message.content or ""


def run(brief: str) -> str:
    print("  drafting...")
    text = draft(brief)
    print(f"      draft: {text[:90]}...")

    for round_number in range(1, MAX_ROUNDS + 1):
        print(f"  [{round_number}/{MAX_ROUNDS}] critiquing...")
        verdict = critique(brief, text)

        if verdict.passes_rubric:
            print(f"      ACCEPTED after {round_number} round(s)")
            print(f"      final: {text}")
            return text

        print(f"      failing rules: {verdict.failing_rules}")
        print(f"      change       : {verdict.what_to_change}")
        text = revise(brief, text, verdict)
        print(f"      revised: {text[:90]}...")

    # Budget exhausted is a real outcome. Return the best text with the
    # honest caveat rather than presenting it as if it passed.
    print(f"      NOT ACCEPTED after {MAX_ROUNDS} rounds, returning last revision")
    return text


if __name__ == "__main__":
    brief = (
        "Our nightly sync job failed twice this week. Both times it was the "
        "same downstream timeout. We have not yet found why the timeout "
        "started happening. Write the update for the team channel."
    )
    print(f"\n> {brief[:70]}...")
    run(brief)