06. Advanced: LangGraph Supervisor

from langgraph.graph import StateGraph, END
from typing import TypedDict, Literal
 
class AgentState(TypedDict):
    messages: list
    next: str
 
def supervisor_node(state: AgentState) -> AgentState:
    """Supervisor decides who acts next"""
    response = llm.invoke([
        SystemMessage(content="""You are a supervisor managing a team.
        Based on the conversation, decide who should act next:
        - 'researcher' for information gathering
        - 'coder' for implementation
        - 'FINISH' if task is complete"""),
        *state["messages"]
    ])
    return {"next": response.content}
 
def researcher_node(state: AgentState) -> AgentState:
    """Research agent gathers information"""
    # ... research logic
    return {"messages": state["messages"] + [result]}
 
def coder_node(state: AgentState) -> AgentState:
    """Coder agent writes code"""
    # ... coding logic
    return {"messages": state["messages"] + [result]}

workflow = StateGraph(AgentState)
 
# Add nodes
workflow.add_node("supervisor", supervisor_node)
workflow.add_node("researcher", researcher_node)
workflow.add_node("coder", coder_node)
 
# Add edges
workflow.add_conditional_edges(
    "supervisor",
    lambda x: x["next"],
    {
        "researcher": "researcher",
        "coder": "coder",
        "FINISH": END
    }
)
workflow.add_edge("researcher", "supervisor")
workflow.add_edge("coder", "supervisor")
 
workflow.set_entry_point("supervisor")
app = workflow.compile()