Five Multi-Agent Coordination Patterns That Actually Work in Enterprise

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


# --- State definition ---

class WorkflowState(TypedDict):
    goal: str
    plan: list[dict]           # Supervisor's task decomposition
    completed_tasks: list[dict]
    failed_tasks: list[dict]
    worker_results: Annotated[list, operator.add]
    final_output: str
    retry_count: int


# --- LLM setup ---

supervisor_llm = ChatOpenAI(model="gpt-4o", temperature=0)
worker_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.2)


# --- Supervisor node: decompose goal into tasks ---

def supervisor_plan(state: WorkflowState) -> WorkflowState:
    response = supervisor_llm.invoke([
        SystemMessage(content="""You are a task decomposition supervisor.
        Break the user's goal into 3-5 specialist tasks.
        Respond with JSON: {"tasks": [{"id": "t1", "specialist": "researcher|analyst|writer", "instruction": "..."}]}"""),
        HumanMessage(content=f"Goal: {state['goal']}")
    ])
    try:
        plan = json.loads(response.content)
        return {"plan": plan["tasks"], "retry_count": 0}
    except json.JSONDecodeError:
        return {"plan": [], "retry_count": state.get("retry_count", 0) + 1}


# --- Worker nodes: execute assigned tasks ---

SPECIALIST_PROMPTS = {
    "researcher": "You are a research specialist. Gather and summarise relevant information.",
    "analyst": "You are an analyst. Analyse the provided information and identify key insights.",
    "writer": "You are a writer. Produce clear, structured written output based on provided inputs.",
}

def execute_worker_tasks(state: WorkflowState) -> WorkflowState:
    results = []
    failed = []
    for task in state["plan"]:
        specialist = task.get("specialist", "analyst")
        system_prompt = SPECIALIST_PROMPTS.get(specialist, SPECIALIST_PROMPTS["analyst"])
        context = ""
        if results:
            context = "\nPrevious results:\n" + "\n".join(
                f"- {r['task_id']}: {r['result'][:200]}" for r in results
            )
        try:
            response = worker_llm.invoke([
                SystemMessage(content=system_prompt),
                HumanMessage(content=f"Task: {task['instruction']}{context}")
            ])
            results.append({"task_id": task["id"], "specialist": specialist,
                           "result": response.content, "status": "success"})
        except Exception as e:
            failed.append({"task_id": task["id"], "error": str(e)})
    return {"worker_results": results, "failed_tasks": failed}


# --- Supervisor aggregation node ---

def supervisor_aggregate(state: WorkflowState) -> WorkflowState:
    results_text = "\n\n".join(
        f"[{r['specialist'].upper()}]\n{r['result']}"
        for r in state.get("worker_results", [])
    )
    response = supervisor_llm.invoke([
        SystemMessage(content="Synthesise the worker outputs into a single coherent response."),
        HumanMessage(content=f"Goal: {state['goal']}\n\nWorker results:\n{results_text}")
    ])
    return {"final_output": response.content}


# --- Routing logic ---

def route_after_planning(state: WorkflowState) -> Literal["execute", "abort"]:
    if not state["plan"] and state.get("retry_count", 0) >= 2:
        return "abort"
    return "execute"

def route_after_execution(state: WorkflowState) -> Literal["aggregate", "retry"]:
    success_count = len(state.get("worker_results", []))
    total_tasks = len(state.get("plan", []))
    if success_count < total_tasks * 0.5:  # >50% failed → retry
        return "retry"
    return "aggregate"


# --- Build graph ---

def build_supervisor_worker_graph() -> any:
    builder = StateGraph(WorkflowState)
    builder.add_node("plan", supervisor_plan)
    builder.add_node("execute", execute_worker_tasks)
    builder.add_node("aggregate", supervisor_aggregate)
    builder.set_entry_point("plan")
    builder.add_conditional_edges("plan", route_after_planning,
                                  {"execute": "execute", "abort": END})
    builder.add_conditional_edges("execute", route_after_execution,
                                  {"aggregate": "aggregate", "retry": "plan"})
    builder.add_edge("aggregate", END)
    return builder.compile(checkpointer=MemorySaver())


if __name__ == "__main__":
    graph = build_supervisor_worker_graph()
    config = {"configurable": {"thread_id": "run-001"}}
    result = graph.invoke(
        {"goal": "Analyse Q3 2025 SaaS market trends and produce an executive summary.",
         "plan": [], "completed_tasks": [], "failed_tasks": [],
         "worker_results": [], "final_output": "", "retry_count": 0},
        config=config
    )
    print(result["final_output"])