code.py

"""Reflection — let the model critique its own work and revise it.

Generate, critique, revise, repeat. The critic runs as a separate call with a
separate prompt, which matters: asking one call to "write it well" gets you
much less than asking a fresh call "what's wrong with this?".

The loop needs a stop condition or it revises forever. Here the critic emits a
sentinel when it has nothing left to say, and MAX_ROUNDS caps it regardless.
"""

from langchain_core.prompts import ChatPromptTemplate

from resources.agent import llm
from resources.helper import show_response

TASK = (
    "Write a Python function `parse_duration(text)` that turns strings like "
    "'2h30m' or '45s' into a total number of seconds."
)

draft_prompt = ChatPromptTemplate.from_template(
    "{task}\n\nReply with the code only."
)

# The critic is deliberately adversarial and deliberately separate. It never
# sees itself as the author, so it has no stake in defending the draft.
critique_prompt = ChatPromptTemplate.from_template(
    "You are a strict code reviewer. Find real bugs and missing edge cases in "
    "the code below. Be specific and terse.\n"
    f"If the code is genuinely correct and complete, reply with exactly {{done}}.\n\n"
    "Task: {task}\n\nCode:\n{draft}"
)

revise_prompt = ChatPromptTemplate.from_template(
    "Revise the code to address every point in the review.\n"
    "Reply with the corrected code only.\n\n"
    "Task: {task}\n\nCode:\n{draft}\n\nReview:\n{critique}"
)

DONE = "APPROVED"  # sentinel the critic emits when it has nothing left
MAX_ROUNDS = 3  # hard cap: reflection without one loops forever


def reflect(task):
    """Draft, then critique-and-revise until approved or out of rounds."""
    draft = llm.invoke(draft_prompt.format(task=task)).content
    print(f"\n{'=' * 60}\nDRAFT\n{'=' * 60}")
    print(draft)

    for round_number in range(1, MAX_ROUNDS + 1):
        critique = llm.invoke(
            critique_prompt.format(task=task, draft=draft, done=DONE)
        )
        print(f"\n{'=' * 60}\nCRITIQUE (round {round_number})\n{'=' * 60}")
        show_response(critique)

        # Stop condition one: the critic has nothing left to complain about.
        if DONE in critique.content:
            print(f"\n[reflect] approved after {round_number} round(s)")
            return draft

        revision = llm.invoke(
            revise_prompt.format(task=task, draft=draft, critique=critique.content)
        )
        draft = revision.content
        print(f"\n{'=' * 60}\nREVISION (round {round_number})\n{'=' * 60}")
        show_response(revision)

    # Stop condition two: out of rounds. Return the best draft we have rather
    # than looping - diminishing returns are real, and so is the token bill.
    print(f"\n[reflect] hit MAX_ROUNDS={MAX_ROUNDS}, returning latest draft")
    return draft


if __name__ == "__main__":
    final = reflect(TASK)
    print(f"\n{'=' * 60}\nFINAL\n{'=' * 60}")
    print(final)