Multi-Agent Systems

1from typing import TypedDict, Annotated, Literal
2from langgraph.graph import StateGraph, END
3from langchain_openai import ChatOpenAI
4from langchain_core.prompts import ChatPromptTemplate
5import operator
6
7llm = ChatOpenAI(model="gpt-4o-mini")
8
9class TeamState(TypedDict):
10    messages: Annotated[list, operator.add]
11    task: str
12    research: str
13    draft: str
14    review: str
15    next_agent: str
16
17# Supervisor decides which agent to call
18supervisor_prompt = ChatPromptTemplate.from_template("""
19You are a team supervisor managing: researcher, writer, reviewer.
20
21Task: {task}
22Current progress:
23- Research: {research}
24- Draft: {draft}
25- Review: {review}
26
27Decide the next agent to call. Respond with ONLY one of:
28- "researcher" (if more research needed)
29- "writer" (if ready to write/rewrite)
30- "reviewer" (if draft needs review)
31- "FINISH" (if task is complete)
32""")
33
34def supervisor_node(state):
35    """Supervisor routes to appropriate agent."""
36    chain = supervisor_prompt | llm
37    result = chain.invoke({
38        "task": state["task"],
39        "research": state.get("research", "Not done"),
40        "draft": state.get("draft", "Not done"),
41        "review": state.get("review", "Not done")
42    })
43    
44    next_agent = result.content.strip().lower()
45    return {"next_agent": next_agent}
46
47def researcher_node(state):
48    """Research agent gathers information."""
49    result = llm.invoke(
50        f"Research this topic thoroughly:\n{state['task']}\n"
51        f"Provide key findings, data points, and sources."
52    ).content
53    return {
54        "research": result,
55        "messages": [{"role": "researcher", "content": result}]
56    }
57
58def writer_node(state):
59    """Writer agent creates content."""
60    review_feedback = state.get("review", "")
61    
62    prompt = f"Write a comprehensive article about:\n{state['task']}\n"
63    prompt += f"Based on research:\n{state.get('research', '')}\n"
64    if review_feedback and review_feedback != "Not done":
65        prompt += f"Address this feedback:\n{review_feedback}"
66    
67    result = llm.invoke(prompt).content
68    return {
69        "draft": result,
70        "messages": [{"role": "writer", "content": result}]
71    }
72
73def reviewer_node(state):
74    """Reviewer agent checks quality."""
75    result = llm.invoke(
76        f"Review this article critically:\n{state.get('draft', '')}\n"
77        f"Check: accuracy, completeness, clarity, engagement.\n"
78        f"Provide specific feedback."
79    ).content
80    return {
81        "review": result,
82        "messages": [{"role": "reviewer", "content": result}]
83    }
84
85def route_supervisor(state):
86    """Route based on supervisor decision."""
87    next_agent = state.get("next_agent", "researcher")
88    if next_agent == "finish":
89        return END
90    return next_agent
91
92# Build multi-agent graph
93graph = StateGraph(TeamState)
94graph.add_node("supervisor", supervisor_node)
95graph.add_node("researcher", researcher_node)
96graph.add_node("writer", writer_node)
97graph.add_node("reviewer", reviewer_node)
98
99graph.set_entry_point("supervisor")
100graph.add_conditional_edges("supervisor", route_supervisor, {
101    "researcher": "researcher",
102    "writer": "writer",
103    "reviewer": "reviewer",
104    END: END
105})
106
107# All agents report back to supervisor
108graph.add_edge("researcher", "supervisor")
109graph.add_edge("writer", "supervisor")
110graph.add_edge("reviewer", "supervisor")
111
112app = graph.compile()
113
114# Run
115result = app.invoke({
116    "task": "Write an article about AI adoption in Vietnamese businesses",
117    "messages": []
118})

1class SupportState(TypedDict):
2    messages: Annotated[list, operator.add]
3    customer_query: str
4    category: str
5    resolution: str
6
7def triage_agent(state):
8    """Classify and route customer query."""
9    result = llm.invoke(
10        f"Classify this customer query into one category:\n"
11        f"Query: {state['customer_query']}\n\n"
12        f"Categories: billing, technical, sales, general\n"
13        f"Respond with ONLY the category name."
14    ).content.strip().lower()
15    
16    return {
17        "category": result,
18        "messages": [{"role": "triage", "content": f"Category: {result}"}]
19    }
20
21def billing_agent(state):
22    """Handle billing queries."""
23    result = llm.invoke(
24        f"You are a billing specialist. Help with:\n{state['customer_query']}\n"
25        f"Be specific about amounts, deadlines, and next steps."
26    ).content
27    return {
28        "resolution": result,
29        "messages": [{"role": "billing_agent", "content": result}]
30    }
31
32def technical_agent(state):
33    """Handle technical queries."""
34    result = llm.invoke(
35        f"You are a technical support specialist. Help with:\n{state['customer_query']}\n"
36        f"Provide step-by-step troubleshooting."
37    ).content
38    return {
39        "resolution": result,
40        "messages": [{"role": "tech_agent", "content": result}]
41    }
42
43def sales_agent(state):
44    """Handle sales queries."""
45    result = llm.invoke(
46        f"You are a sales specialist. Help with:\n{state['customer_query']}\n"
47        f"Recommend products and pricing."
48    ).content
49    return {
50        "resolution": result,
51        "messages": [{"role": "sales_agent", "content": result}]
52    }
53
54def general_agent(state):
55    """Handle general queries."""
56    result = llm.invoke(
57        f"Help with this general query:\n{state['customer_query']}"
58    ).content
59    return {
60        "resolution": result,
61        "messages": [{"role": "general_agent", "content": result}]
62    }
63
64def route_to_specialist(state):
65    """Route to specialist based on category."""
66    category = state.get("category", "general")
67    return category
68
69# Build support graph
70graph = StateGraph(SupportState)
71graph.add_node("triage", triage_agent)
72graph.add_node("billing", billing_agent)
73graph.add_node("technical", technical_agent)
74graph.add_node("sales", sales_agent)
75graph.add_node("general", general_agent)
76
77graph.set_entry_point("triage")
78graph.add_conditional_edges("triage", route_to_specialist, {
79    "billing": "billing",
80    "technical": "technical",
81    "sales": "sales",
82    "general": "general"
83})
84
85for node in ["billing", "technical", "sales", "general"]:
86    graph.add_edge(node, END)
87
88support_app = graph.compile()
89
90# Test
91result = support_app.invoke({
92    "customer_query": "Tôi muốn hủy gói subscription và hoàn tiền",
93    "messages": []
94})
95print(f"Category: {result['category']}")
96print(f"Resolution: {result['resolution']}")

1"""
2Level 1: Project Manager
3Level 2: Team Leads (Research Lead, Engineering Lead)
4Level 3: Worker Agents (Researchers, Engineers)
5
6Project Manager
7├── Research Lead
8│   ├── Web Researcher
9│   └── Data Analyst
10└── Engineering Lead
11    ├── Backend Developer
12    └── Frontend Developer
13"""
14
15def project_manager(state):
16    """Top-level coordinator."""
17    result = llm.invoke(
18        f"As project manager, decide next team to work:\n"
19        f"Task: {state['task']}\n"
20        f"Progress: {state.get('progress', 'Starting')}\n"
21        f"Choose: 'research_team' or 'engineering_team' or 'DONE'"
22    ).content.strip()
23    return {"next_team": result}
24
25def research_lead(state):
26    """Research team lead coordinates researchers."""
27    result = llm.invoke(
28        f"As research lead, coordinate research for:\n{state['task']}\n"
29        f"Synthesize findings from web research and data analysis."
30    ).content
31    return {
32        "research_output": result,
33        "progress": "Research complete",
34        "messages": [{"role": "research_lead", "content": result}]
35    }
36
37def engineering_lead(state):
38    """Engineering team lead coordinates developers."""
39    result = llm.invoke(
40        f"As engineering lead, plan implementation for:\n{state['task']}\n"
41        f"Based on research:\n{state.get('research_output', 'N/A')}"
42    ).content
43    return {
44        "engineering_output": result,
45        "progress": "Engineering complete",
46        "messages": [{"role": "engineering_lead", "content": result}]
47    }