Agent Design Patterns: ReAct, Reflexion, Plan-and-Execute, and Supervisor-Worker

from typing import TypedDict, Annotated, Literal
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
import operator

# --- State Definition ---
class AnalysisState(TypedDict):
    task: str                          # Original task description
    data_analysis: str                 # Output from data analyst worker
    report_draft: str                  # Output from report writer worker
    quality_feedback: str              # Output from quality checker worker
    final_report: str                  # Final assembled report
    next_step: str                     # Supervisor routing decision
    iterations: int                    # Guard against infinite loops
    status: str                        # "in_progress" | "complete" | "failed"

llm = ChatOpenAI(model="gpt-4o", temperature=0.1)

# --- Worker Node Functions ---
def data_analyst_node(state: AnalysisState) -> dict:
    """Specialized agent for quantitative data analysis."""
    response = llm.invoke([
        SystemMessage(content=(
            "You are a senior data analyst specializing in enterprise financial and operational data. "
            "Analyze the task provided and produce a structured quantitative analysis with key metrics, "
            "trends, and data-driven findings. Be specific and cite figures."
        )),
        HumanMessage(content=f"Task: {state['task']}\n\nProvide your data analysis.")
    ])
    return {"data_analysis": response.content}

def report_writer_node(state: AnalysisState) -> dict:
    """Specialized agent for structured business writing."""
    response = llm.invoke([
        SystemMessage(content=(
            "You are a senior business analyst who writes clear, executive-ready reports. "
            "Transform the provided data analysis into a well-structured report with an executive summary, "
            "key findings, and recommendations. Write for a CFO/CTO audience."
        )),
        HumanMessage(content=(
            f"Task: {state['task']}\n\n"
            f"Data Analysis:\n{state['data_analysis']}\n\n"
            "Write the report."
        ))
    ])
    return {"report_draft": response.content}

def quality_checker_node(state: AnalysisState) -> dict:
    """Specialized agent for report quality assurance."""
    response = llm.invoke([
        SystemMessage(content=(
            "You are a quality assurance specialist for business reports. "
            "Review the provided report draft for: factual consistency with the source data analysis, "
            "logical coherence of recommendations, completeness of required sections, and executive clarity. "
            "Return either APPROVED with brief rationale, or REVISE with specific feedback."
        )),
        HumanMessage(content=(
            f"Source Data Analysis:\n{state['data_analysis']}\n\n"
            f"Report Draft:\n{state['report_draft']}\n\n"
            "Quality assessment:"
        ))
    ])
    return {"quality_feedback": response.content}

def supervisor_node(state: AnalysisState) -> dict:
    """Routes task to the appropriate next worker based on current state."""
    iterations = state.get("iterations", 0) + 1

    # Hard stop: prevent infinite loops
    if iterations > 6:
        return {
            "next_step": "end",
            "status": "failed",
            "final_report": state.get("report_draft", "Max iterations reached without completion."),
            "iterations": iterations
        }

    # Routing logic based on state
    if not state.get("data_analysis"):
        return {"next_step": "data_analyst", "iterations": iterations, "status": "in_progress"}

    if not state.get("report_draft"):
        return {"next_step": "report_writer", "iterations": iterations, "status": "in_progress"}

    if not state.get("quality_feedback"):
        return {"next_step": "quality_checker", "iterations": iterations, "status": "in_progress"}

    # Check quality checker decision
    feedback = state.get("quality_feedback", "")
    if "APPROVED" in feedback.upper():
        return {
            "next_step": "end",
            "status": "complete",
            "final_report": state["report_draft"],
            "iterations": iterations
        }
    elif "REVISE" in feedback.upper() and iterations <= 4:
        # Request revision — clear draft to trigger re-write with feedback injected into task
        revised_task = (
            f"{state['task']}\n\n"
            f"QUALITY FEEDBACK REQUIRING REVISION:\n{feedback}"
        )
        return {
            "next_step": "report_writer",
            "task": revised_task,
            "report_draft": "",
            "quality_feedback": "",
            "iterations": iterations,
            "status": "in_progress"
        }
    else:
        # Accept current draft after max revision attempts
        return {
            "next_step": "end",
            "status": "complete",
            "final_report": state["report_draft"],
            "iterations": iterations
        }

# --- Router Function for Conditional Edges ---
def route_from_supervisor(state: AnalysisState) -> Literal["data_analyst", "report_writer", "quality_checker", "__end__"]:
    next_step = state.get("next_step", "data_analyst")
    if next_step == "end":
        return END
    return next_step

# --- Graph Construction ---
def build_analysis_graph() -> StateGraph:
    graph = StateGraph(AnalysisState)

    graph.add_node("supervisor", supervisor_node)
    graph.add_node("data_analyst", data_analyst_node)
    graph.add_node("report_writer", report_writer_node)
    graph.add_node("quality_checker", quality_checker_node)

    graph.set_entry_point("supervisor")

    # All workers route back to supervisor after completion
    graph.add_edge("data_analyst", "supervisor")
    graph.add_edge("report_writer", "supervisor")
    graph.add_edge("quality_checker", "supervisor")

    # Supervisor uses conditional edges to route to next worker or end
    graph.add_conditional_edges("supervisor", route_from_supervisor)

    return graph.compile(checkpointer=MemorySaver())

def main():
    graph = build_analysis_graph()
    config = {"configurable": {"thread_id": "procurement-analysis-001"}}

    initial_state = AnalysisState(
        task="Analyze Q1 2026 procurement spend across the top 5 vendor categories. Identify cost reduction opportunities and flag any vendors with deteriorating performance metrics.",
        data_analysis="",
        report_draft="",
        quality_feedback="",
        final_report="",
        next_step="",
        iterations=0,
        status="in_progress"
    )

    result = graph.invoke(initial_state, config)
    print(f"Status: {result['status']} (completed in {result['iterations']} supervisor iterations)")
    print("\n=== FINAL REPORT ===")
    print(result["final_report"])

if __name__ == "__main__":
    main()