langgraph_supervisor_orchestrating_researcher_coder_worker_agents.py

1import operator
2from typing import Annotated, List, TypedDict, Union
3
4from langchain_openai import ChatOpenAI
5from langchain_core.messages import BaseMessage, HumanMessage
6from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
7from langgraph.graph import END, StateGraph, START
8from langgraph.prebuilt import create_react_agent
9
10# 1. Define the tools and agents
11def researcher(state):
12    # This is a placeholder for a real research tool/agent
13    return {"messages": [HumanMessage(content="Researcher: Found info on the topic.", name="Researcher")]}
14
15def coder(state):
16    # This is a placeholder for a real coding tool/agent
17    return {"messages": [HumanMessage(content="Coder: Wrote the script.", name="Coder")]}
18
19# 2. Define the Supervisor State
20class AgentState(TypedDict):
21    messages: Annotated[List[BaseMessage], operator.add]
22    next: str
23
24# 3. Create the Supervisor Node
25members = ["Researcher", "Coder"]
26system_prompt = (
27    "You are a supervisor tasked with managing a conversation between the"
28    " following workers: {members}. Given the following user request,"
29    " respond with the worker to act next. Each worker will perform a"
30    " task and respond with their results and status. When finished,"
31    " respond with FINISH."
32)
33
34options = ["FINISH"] + members
35llm = ChatOpenAI(model="gpt-4o")
36
37def supervisor_node(state: AgentState):
38    prompt = ChatPromptTemplate.from_messages([
39        ("system", system_prompt),
40        MessagesPlaceholder(variable_name="messages"),
41        ("system", "Given the conversation above, who should act next?"
42                   " Or should we FINISH? Select one of: {options}"),
43    ]).partial(options=str(options), members=", ".join(members))
44    
45    chain = prompt | llm.with_structured_output(TypedDict("Route", {"next": str}))
46    return chain.invoke(state)
47
48# 4. Build the Graph
49workflow = StateGraph(AgentState)
50
51# Add nodes
52workflow.add_node("Researcher", researcher)
53workflow.add_node("Coder", coder)
54workflow.add_node("supervisor", supervisor_node)
55
56# Define edges: Workers always report back to supervisor
57for member in members:
58    workflow.add_edge(member, "supervisor")
59
60# The supervisor decides the next step via conditional edges
61conditional_map = {k: k for k in members}
62conditional_map["FINISH"] = END
63workflow.add_conditional_edges("supervisor", lambda x: x["next"], conditional_map)
64
65workflow.add_edge(START, "supervisor")
66
67graph = workflow.compile()
68
69# 5. Execute
70for s in graph.stream({"messages": [HumanMessage(content="Write a report on AI and then code a summary script.")]}):
71    if "__end__" not in s:
72        print(s)
73        print("---")