Simple AI Workflow with Reflection, Structured Output and Validation

%%capture
%pip install -q "openai>=2.0.0" "python-dotenv>=1.0.0"

# OpenAI SDK: provides a clean interface for calling LLMs
# We'll use it with OpenRouter, which is OpenAI API-compatible
from openai import OpenAI

# Standard library
import json
import os

# # Uncomment the lines below if running in Google Colab
# from google.colab import userdata
# os.environ["OPENROUTER_API_KEY"] = userdata.get("OPENROUTER_API_KEY")
# print("API key loaded from Colab Secrets")

# Load a local .env file when one is present.
# This does not replace a key already loaded from Colab Secrets.
from dotenv import load_dotenv

env_file_found = load_dotenv(dotenv_path=".env")
print("Local .env file found:", env_file_found)

# Sanity check: verify that a non-empty API key is available.
api_key = os.getenv("OPENROUTER_API_KEY")
assert api_key, "OPENROUTER_API_KEY not found. Please check your setup."
print("OPENROUTER_API_KEY present:", bool(api_key))

# Initialize the OpenAI client pointing to OpenRouter.
client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=api_key,
)

# Use one model and one reasoning setting for all three roles.
MODEL = "openai/gpt-5.6-luna"
REASONING_EFFORT = "none"

print(f"Client initialized. Using model: {MODEL}")
print(f"Reasoning effort: {REASONING_EFFORT}")

# A small, predefined bound keeps retries predictable.
MAX_ATTEMPTS = 2

WRITER_SCHEMA = {
    "type": "object",
    "properties": {
        "title": {"type": "string"},
        "sections": {"type": "array", "items": {"type": "string"}},
        "conclusion": {"type": "string"},
    },
    "required": ["title", "sections", "conclusion"],
    "additionalProperties": False,
}

CRITIC_SCHEMA = {
    "type": "object",
    "properties": {
        "strengths": {"type": "array", "items": {"type": "string"}},
        "problems": {"type": "array", "items": {"type": "string"}},
        "suggestions": {"type": "array", "items": {"type": "string"}},
    },
    "required": ["strengths", "problems", "suggestions"],
    "additionalProperties": False,
}

REVISER_SCHEMA = {
    "type": "object",
    "properties": {"revised_report": {"type": "string"}},
    "required": ["revised_report"],
    "additionalProperties": False,
}

def validate_writer(data):
    errors = []
    for field in ("title", "sections", "conclusion"):
        if field not in data:
            errors.append(f"missing required field: {field}")
    return errors


def validate_critic(data):
    errors = []
    for field in ("strengths", "problems", "suggestions"):
        if field not in data:
            errors.append(f"missing required field: {field}")
        elif not isinstance(data[field], list) or not data[field]:
            errors.append(f"{field} must be a non-empty list")
    return errors


def validate_reviser(data):
    if "revised_report" not in data:
        return ["missing required field: revised_report"]
    if not isinstance(data["revised_report"], str) or not data["revised_report"].strip():
        return ["revised_report must contain non-empty text"]
    return []


def call_structured_step(messages, schema_name, schema, validator, model=MODEL):
    """Call one workflow step, validate it, and retry with specific feedback."""
    feedback = None

    for attempt in range(1, MAX_ATTEMPTS + 1):
        attempt_messages = list(messages)
        if feedback:
            attempt_messages.append({
                "role": "user",
                "content": f"Your previous output was invalid. Fix these issues: {feedback}",
            })

        response = client.chat.completions.create(
            model=model,
            messages=attempt_messages,
            reasoning_effort=REASONING_EFFORT,
            response_format={
                "type": "json_schema",
                "json_schema": {"name": schema_name, "strict": True, "schema": schema},
            },
        )

        raw_output = response.choices[0].message.content or ""
        try:
            data = json.loads(raw_output)
            errors = validator(data) if isinstance(data, dict) else ["output must be a JSON object"]
        except json.JSONDecodeError as error:
            data = None
            errors = [f"output was not valid JSON: {error.msg}"]

        if not errors:
            return data, response

        feedback = "; ".join(errors)
        print(f"{schema_name} attempt {attempt} failed validation: {feedback}")

    raise RuntimeError(
        f"{schema_name} failed validation after {MAX_ATTEMPTS} attempts: {feedback}"
    )


def format_writer_output(data):
    """Turn the validated Writer object into readable report text."""
    return "\n\n".join([data["title"], *data["sections"], data["conclusion"]])


def format_critique(data):
    """Turn the validated Critic object into readable feedback."""
    blocks = []
    for label, field in (("Strengths", "strengths"), ("Problems", "problems"), ("Suggestions", "suggestions")):
        blocks.append(label + ":\n- " + "\n- ".join(data[field]))
    return "\n\n".join(blocks)

# System prompt that defines the Report Writer persona
WRITER_SYSTEM_PROMPT = """You are a journalist who writes clear, informative reports.

When given a topic, write a short report that:
- Provides a balanced overview of the topic
- Includes relevant context and background
- Is written in a professional, accessible style

Return the report using the required structured output schema."""

def generate_initial_report(topic: str, system_prompt: str = WRITER_SYSTEM_PROMPT, model: str = MODEL, return_response=False):
    """Generate and validate the Writer's structured output."""
    data, response = call_structured_step(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Write a concise report on: {topic}"},
        ],
        schema_name="writer_output",
        schema=WRITER_SCHEMA,
        validator=validate_writer,
        model=model,
    )
    return (data, response) if return_response else data

# Make one validated Writer call and keep the complete response object.
topic = "The impact of artificial intelligence on healthcare"

writer_output, response = generate_initial_report(topic, return_response=True)

response.model_dump()

# Convert the validated Writer object to text for display and for the Critic.
initial_report = format_writer_output(writer_output)
initial_report

# System prompt that defines the Critic persona
CRITIC_SYSTEM_PROMPT = """You are an editor who provides constructive feedback on reports.

Identify strengths, problems, and specific suggestions for improvement. Do not rewrite the report.
Return the critique using the required structured output schema."""

def generate_critique(report: str, system_prompt: str = CRITIC_SYSTEM_PROMPT, model: str = MODEL):
    """Generate and validate the Critic's structured output."""
    data, _ = call_structured_step(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"Please critique this report:\n\n{report}"},
        ],
        schema_name="critic_output",
        schema=CRITIC_SCHEMA,
        validator=validate_critic,
        model=model,
    )
    return data

# Generate and validate the critique. Critic failure stops the workflow.
try:
    critic_output = generate_critique(initial_report)
    critique = format_critique(critic_output)
except RuntimeError as error:
    raise RuntimeError("Workflow stopped: the draft was not evaluated.") from error

critique

print("Critique word count:", len(critique.split()))

# System prompt that defines the Reviser persona
REVISER_SYSTEM_PROMPT = """You are a senior reporter who produces polished final reports.

Revise the original report by addressing the critique while preserving its useful content and voice.
Return the revision using the required structured output schema."""

def generate_revised_report(report: str, critique: str, system_prompt: str = REVISER_SYSTEM_PROMPT, model: str = MODEL):
    """Generate and validate the Reviser's structured output."""
    user_message = f"""Revise the report using the editorial feedback.

ORIGINAL REPORT:
{report}

EDITORIAL FEEDBACK:
{critique}"""

    data, _ = call_structured_step(
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message},
        ],
        schema_name="reviser_output",
        schema=REVISER_SCHEMA,
        validator=validate_reviser,
        model=model,
    )
    return data

# Generate and validate the revision. On failure, preserve prior valid work.
try:
    reviser_output = generate_revised_report(initial_report, critique)
    final_report = reviser_output["revised_report"]
except RuntimeError as error:
    final_report = initial_report
    print("Revision failed; preserving the original draft and critique.")
    print(error)

final_report

print("Initial report word count:", len(initial_report.split()))
print("Revised report word count:", len(final_report.split()))

from IPython.display import Markdown

comparison = f"""
## Initial Report

{initial_report}

---

## Critique

{critique}

---

## Final Report (After Reflection)

{final_report}
"""

Markdown(comparison)