1from langgraph.graph import StateGraph, END
2from typing import TypedDict, Annotated, list
3import operator
4
5class State(TypedDict):
6 goal: str
7 research: str
8 code: str
9 review: str
10 messages: Annotated[list, operator.add]
11
12def supervisor(state: State) -> dict:
13 # 1. Decide next specialist
14 if not state.get("research"): return {"messages": ["-> researcher"], "next": "researcher"}
15 if not state.get("code"): return {"messages": ["-> coder"], "next": "coder"}
16 if not state.get("review"): return {"messages": ["-> critic"], "next": "critic"}
17 return {"next": END}
18
19def researcher(state: State) -> dict:
20 findings = call_llm(f"Research: {state['goal']}", role="researcher")
21 return {"research": findings, "messages": [f"researcher: {findings[:60]}..."]}
22
23def coder(state: State) -> dict:
24 impl = call_llm(f"Code based on: {state['research']}", role="coder")
25 return {"code": impl, "messages": [f"coder: wrote {len(impl)} chars"]}
26
27def critic(state: State) -> dict:
28 review = call_llm(f"Review:\n{state['code']}", role="critic")
29 return {"review": review, "messages": [f"critic: {review[:60]}..."]}
30
31# 2. Build graph
32g = StateGraph(State)
33g.add_node("supervisor", supervisor); g.add_node("researcher", researcher)
34g.add_node("coder", coder); g.add_node("critic", critic)
35g.set_entry_point("supervisor")
36g.add_conditional_edges("supervisor", lambda s: s["next"], {"researcher":"researcher","coder":"coder","critic":"critic", END: END})
37for n in ["researcher","coder","critic"]: g.add_edge(n, "supervisor")
38
39# 3. Compile + run
40app = g.compile()
41result = app.invoke({"goal": "Build a Python rate limiter"})
42print(result["messages"])
