Supervisor Pattern: Multi-Agent Coordination

User Query
    ↓
[Supervisor] ─→ "needs drug research" ─→ [Pharmacology Agent]
    ↑                                              ↓
    └─────── output reviewed ←──────────────────── ┘
    ↓
[Supervisor] ─→ "needs safety check" ─→ [Safety Agent]
    ↑                                              ↓
    └─────── output reviewed ←──────────────────── ┘
    ↓
[Supervisor] ─→ "complete" ─→ [END]

from langgraph.graph import StateGraph, END
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from typing import TypedDict, Annotated, Literal
import operator
import json

# State shared between supervisor and all agents
class SupervisorState(TypedDict):
    query: str
    messages: Annotated[list[dict], operator.add]
    next_agent: str
    pharmacology_report: str
    safety_report: str
    regulatory_report: str
    final_response: str

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

AGENTS = ["pharmacology_agent", "safety_agent", "regulatory_agent", "FINISH"]

def supervisor(state: SupervisorState) -> dict:
    """Supervisor decides which agent to invoke next."""
    system_prompt = f"""You are a supervisor coordinating a drug information team.

Available agents: {', '.join(AGENTS[:-1])}

Current query: {state['query']}

Information gathered so far:
- Pharmacology report: {'Completed' if state.get('pharmacology_report') else 'Not yet done'}
- Safety report: {'Completed' if state.get('safety_report') else 'Not yet done'}
- Regulatory report: {'Completed' if state.get('regulatory_report') else 'Not yet done'}

Decide which agent to invoke next, or FINISH if all needed information has been gathered.
Return JSON only: {{"next": "agent_name_or_FINISH", "reason": "one sentence"}}"""

    response = llm.invoke([SystemMessage(content=system_prompt)])
    parsed = json.loads(response.content)

    return {
        "next_agent": parsed["next"],
        "messages": [{"role": "supervisor", "content": f"Routing to {parsed['next']}: {parsed['reason']}"}],
    }

def pharmacology_agent(state: SupervisorState) -> dict:
    """Specialist in pharmacological mechanisms."""
    response = llm.invoke([
        SystemMessage(content="You are a clinical pharmacologist. Provide accurate pharmacological data."),
        HumanMessage(content=f"Research the pharmacological profile of the drug mentioned in: {state['query']}"),
    ])
    return {
        "pharmacology_report": response.content,
        "messages": [{"role": "pharmacology_agent", "content": response.content[:100] + "..."}],
    }

def safety_agent(state: SupervisorState) -> dict:
    """Specialist in drug safety and interactions."""
    context = f"Query: {state['query']}\nPharmacology data: {state.get('pharmacology_report', 'Not available')}"
    response = llm.invoke([
        SystemMessage(content="You are a drug safety specialist. Identify safety concerns and interactions."),
        HumanMessage(content=f"Analyze safety for: {context}"),
    ])
    return {
        "safety_report": response.content,
        "messages": [{"role": "safety_agent", "content": response.content[:100] + "..."}],
    }

def regulatory_agent(state: SupervisorState) -> dict:
    """Specialist in regulatory status."""
    response = llm.invoke([
        SystemMessage(content="You are a regulatory affairs specialist with FDA and EMA expertise."),
        HumanMessage(content=f"Provide regulatory information for: {state['query']}"),
    ])
    return {
        "regulatory_report": response.content,
        "messages": [{"role": "regulatory_agent", "content": response.content[:100] + "..."}],
    }

def synthesizer(state: SupervisorState) -> dict:
    """Combine all reports into a final response."""
    context = f"""Query: {state['query']}

Pharmacology: {state.get('pharmacology_report', 'Not gathered')}
Safety: {state.get('safety_report', 'Not gathered')}
Regulatory: {state.get('regulatory_report', 'Not gathered')}"""

    response = llm.invoke([
        SystemMessage(content="Synthesize the gathered information into a comprehensive drug information summary."),
        HumanMessage(content=context),
    ])
    return {"final_response": response.content}

def route_supervisor(state: SupervisorState) -> str:
    """Route based on supervisor's decision."""
    next_agent = state.get("next_agent", "FINISH")
    if next_agent == "FINISH":
        return "synthesizer"
    return next_agent

builder = StateGraph(SupervisorState)

# Add all nodes
builder.add_node("supervisor", supervisor)
builder.add_node("pharmacology_agent", pharmacology_agent)
builder.add_node("safety_agent", safety_agent)
builder.add_node("regulatory_agent", regulatory_agent)
builder.add_node("synthesizer", synthesizer)

# Entry: always start with supervisor
builder.set_entry_point("supervisor")

# Supervisor routes to agents dynamically
builder.add_conditional_edges(
    "supervisor",
    route_supervisor,
    {
        "pharmacology_agent": "pharmacology_agent",
        "safety_agent": "safety_agent",
        "regulatory_agent": "regulatory_agent",
        "synthesizer": "synthesizer",
    },
)

# After each agent, always return to supervisor for next decision
builder.add_edge("pharmacology_agent", "supervisor")
builder.add_edge("safety_agent", "supervisor")
builder.add_edge("regulatory_agent", "supervisor")

# Synthesizer is the final step
builder.add_edge("synthesizer", END)

app = builder.compile()

# Run
result = app.invoke({
    "query": "What are the key clinical considerations for warfarin therapy?",
    "messages": [],
    "next_agent": "",
    "pharmacology_report": "",
    "safety_report": "",
    "regulatory_report": "",
    "final_response": "",
})

print(result["final_response"])

class SupervisorStateWithLimit(SupervisorState):
    routing_steps: int
    max_routing_steps: int

def supervisor_with_limit(state: SupervisorStateWithLimit) -> dict:
    if state["routing_steps"] >= state["max_routing_steps"]:
        return {"next_agent": "FINISH", "routing_steps": state["routing_steps"] + 1}

    # Normal supervisor logic...
    result = supervisor(state)
    return {**result, "routing_steps": state["routing_steps"] + 1}