S121 · Multi-Agent Orchestration — LangGraph, CrewAI Patterns

"""multi_agent.py — supervisor + researcher + writer + reviewer via LangGraph.
 
Requires: pip install langgraph langchain-openai
"""
from __future__ import annotations
import os
from typing import Annotated, Literal, TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.messages import (
    HumanMessage, AIMessage, SystemMessage, BaseMessage,
)
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
 
MODEL_MAIN   = "gpt-4o-mini"      # cheap for supervisor + workers
MODEL_REV    = "gpt-4o-mini"      # separate LLM ideally, same here for cost
MAX_STEPS    = 12
 
 
# ---------- Shared state ----------
class State(TypedDict):
    messages: Annotated[list[BaseMessage], add_messages]  # running convo
    next_agent: str                                        # supervisor's routing decision
    research_notes: str
    draft: str
    review: str
    step: int
 
 
# ---------- Prompts per role ----------
def sys(role: str) -> str:
    return {
        "supervisor": (
            "You are the SUPERVISOR of a small writing team.\n"
            "Team: researcher, writer, reviewer.\n"
            "Read the state and pick the SINGLE next role. Respond with ONE word: "
            "researcher | writer | reviewer | FINISH.\n"
            "- If research_notes empty: pick researcher.\n"
            "- If draft empty: pick writer.\n"
            "- If review empty: pick reviewer.\n"
            "- If review says APPROVED: pick FINISH.\n"
            "- If review says REVISE: pick writer."
        ),
        "researcher": (
            "You are the RESEARCHER. Given the user's request, produce 3–5 "
            "bullet points of factual claims + sources. Keep it under 200 words. "
            "Return ONLY the bullet points."
        ),
        "writer": (
            "You are the WRITER. Given the research notes and (optionally) a prior "
            "review, produce a polished 200-word summary. Return ONLY the summary text."
        ),
        "reviewer": (
            "You are the REVIEWER. Read the draft. If it accurately reflects the "
            "research and is under 250 words, respond with exactly 'APPROVED'. "
            "Otherwise respond with 'REVISE: <2-sentence critique>'."
        ),
    }[role]
 
 
llm_main = ChatOpenAI(model=MODEL_MAIN, temperature=0)
llm_rev  = ChatOpenAI(model=MODEL_REV,  temperature=0)
 
 
# ---------- Node implementations ----------
def supervisor(state: State) -> State:
    prompt = (
        f"USER REQUEST: {state['messages'][0].content}\n"
        f"research_notes: {state['research_notes'][:200] or '(empty)'}\n"
        f"draft: {state['draft'][:200] or '(empty)'}\n"
        f"review: {state['review'][:200] or '(empty)'}\n"
    )
    r = llm_main.invoke([SystemMessage(sys("supervisor")),
                         HumanMessage(prompt)])
    choice = r.content.strip().split()[0].lower()
    print(f"[step {state['step']}] SUPERVISOR → {choice}")
    return {"next_agent": choice, "step": state["step"] + 1}
 
 
def researcher(state: State) -> State:
    r = llm_main.invoke([SystemMessage(sys("researcher")),
                         HumanMessage(state['messages'][0].content)])
    print(f"[step {state['step']}] RESEARCHER produced {len(r.content)} chars")
    return {"research_notes": r.content}
 
 
def writer(state: State) -> State:
    prompt = (f"RESEARCH NOTES:\n{state['research_notes']}\n\n"
              f"PRIOR REVIEW (if any): {state['review']}\n\n"
              f"Write the summary now.")
    r = llm_main.invoke([SystemMessage(sys("writer")), HumanMessage(prompt)])
    print(f"[step {state['step']}] WRITER produced {len(r.content)} chars")
    return {"draft": r.content, "review": ""}       # clear review so we re-review
 
 
def reviewer(state: State) -> State:
    prompt = (f"RESEARCH NOTES:\n{state['research_notes']}\n\n"
              f"DRAFT:\n{state['draft']}")
    r = llm_rev.invoke([SystemMessage(sys("reviewer")), HumanMessage(prompt)])
    print(f"[step {state['step']}] REVIEWER → {r.content[:80]}...")
    return {"review": r.content}
 
 
# ---------- Routing ----------
def route(state: State) -> Literal["researcher","writer","reviewer","finish"]:
    if state["step"] >= MAX_STEPS:
        print("!! MAX_STEPS hit — forcing FINISH")
        return "finish"
    return {
        "researcher": "researcher",
        "writer":     "writer",
        "reviewer":   "reviewer",
        "finish":     "finish",
    }.get(state["next_agent"], "finish")
 
 
# ---------- Build the graph ----------
g = StateGraph(State)
g.add_node("supervisor", supervisor)
g.add_node("researcher", researcher)
g.add_node("writer",     writer)
g.add_node("reviewer",   reviewer)
 
g.add_edge(START, "supervisor")
g.add_conditional_edges("supervisor", route, {
    "researcher": "researcher",
    "writer":     "writer",
    "reviewer":   "reviewer",
    "finish":     END,
})
# After each worker, hand back to the supervisor.
for worker in ("researcher", "writer", "reviewer"):
    g.add_edge(worker, "supervisor")
 
app = g.compile()
 
 
# ---------- Run ----------
if __name__ == "__main__":
    initial = {
        "messages": [HumanMessage("Write a 200-word summary of how HNSW graph indexes work.")],
        "next_agent": "",
        "research_notes": "", "draft": "", "review": "", "step": 0,
    }
    final = app.invoke(initial)
    print("\n===== FINAL DRAFT =====\n")
    print(final["draft"])