Building AI Agents with LangGraph: State Machines for LLM Workflows

1from typing import TypedDict, Annotated, List
2from langgraph.graph import StateGraph, END
3import operator
4
5class AgentState(TypedDict):
6    messages: Annotated[List[dict], operator.add]  # accumulates
7    tool_results: Annotated[List[dict], operator.add]
8    iteration: int
9    max_iterations: int
10    task_complete: bool
11    final_answer: str

1from langchain_openai import ChatOpenAI
2from langchain_core.tools import tool
3
4llm = ChatOpenAI(model="gpt-4o", temperature=0)
5
6# Define tools
7@tool
8def search_web(query: str) -> str:
9    """Search the web for current information."""
10    # ... implementation
11    return f"Search results for: {query}"
12
13@tool
14def run_code(code: str) -> str:
15    """Execute Python code and return output."""
16    try:
17        exec_globals = {}
18        exec(code, exec_globals)
19        return str(exec_globals.get("_result", "No output"))
20    except Exception as e:
21        return f"Error: {e}"
22
23tools = [search_web, run_code]
24llm_with_tools = llm.bind_tools(toals)
25
26def reason_node(state: AgentState) -> dict:
27    """LLM decides what to do next."""
28    response = llm_with_tools.invoke(state["messages"])
29    return {"messages": [response]}
30
31def tool_node(state: AgentState) -> dict:
32    """Execute tool calls from the last message."""
33    last_msg = state["messages"][-1]
34    results = []
35
36    for tool_call in last_msg.tool_calls:
37        tool_name = tool_call["name"]
38        tool_args = tool_call["args"]
39        result = {"name": tool_name, "args": tool_args,
40                  "output": globals()[tool_name].invoke(tool_args)}
41        results.append(result)
42
43    return {
44        "tool_results": results,
45        "messages": [{"role": "tool", "content": str(results)}],
46    }
47
48def check_completion(state: AgentState) -> dict:
49    """Check if the task is done."""
50    last_msg = state["messages"][-1]
51    has_tool_calls = bool(getattr(last_msg, "tool_calls", None))
52
53    return {
54        "task_complete": not has_tool_calls,
55        "iteration": state["iteration"] + 1,
56    }

1def should_continue(state: AgentState) -> str:
2    """Route to next node based on state."""
3    if state["task_complete"]:
4        return "finalize"
5    if state["iteration"] >= state["max_iterations"]:
6        return "finalize"
7    return "tools"
8
9def finalize_node(state: AgentState) -> dict:
10    """Extract the final answer."""
11    last_msg = state["messages"][-1]
12    answer = last_msg.content if hasattr(last_msg, "content") else str(last_msg)
13
14    if not state["task_complete"]:
15        answer = f"Task incomplete after {state['iteration']} iterations. Best answer: {answer}"
16
17    return {"final_answer": answer}

1# Build the state machine
2graph = StateGraph(AgentState)
3
4# Add nodes
5graph.add_node("reason", reason_node)
6graph.add_node("tools", tool_node)
7graph.add_node("check", check_completion_node)
8graph.add_node("finalize", finalize_node)
9
10# Set entry point
11graph.set_entry_point("reason")
12
13# Add edges
14graph.add_edge("reason", "tools")      # reason → always try tools
15graph.add_edge("tools", "check")       # tools → check completion
16graph.add_conditional_edges(
17    "check",
18    should_continue,                   # conditional routing
19    {
20        "tools": "tools",              # loop back to tools
21        "finalize": "finalize",        # or finish
22    },
23)
24graph.add_edge("finalize", END)
25
26# Compile
27app = graph.compile()

1START → reason → tools → check ──→ tools (loop)
2                              │
3                              └──→ finalize → END

1initial_state = {
2    "messages": [{"role": "user", "content": "Research the latest ESP32-S3 specs and write a summary"}],
3    "tool_results": [],
4    "iteration": 0,
5    "max_iterations": 5,
6    "task_complete": False,
7    "final_answer": "",
8}
9
10result = app.invoke(initial_state)
11print(result["final_answer"])

1for event in app.stream(initial_state, stream_mode="values"):
2    print(f"Step {event['iteration']}: {event['messages'][-1]}")

1from langgraph.checkpoint.memory import MemorySaver
2
3# Add checkpointing
4checkpointer = MemorySaver()
5app = graph.compile(
6    checkpointer=checkpointer,
7    interrupt_before=["tools"],  # pause before tool execution
8)
9
10# Run — pauses before tools
11config = {"configurable": {"thread_id": "session-1"}}
12result = app.invoke(initial_state, config)
13
14# Human reviews the planned tool calls
15last_msg = result["messages"][-1]
16print(f"Agent wants to call: {last_msg.tool_calls}")
17
18# Human approves → resume
19approved_result = app.invoke(None, config)  # None = continue from checkpoint

1from langgraph.checkpoint.postgres import PostgresSaver
2
3# Persistent checkpointing with Postgres
4checkpointer = PostgresSaver.from_conn_string(
5    "postgresql://user:pass@localhost/agent_db"
6)
7app = graph.compile(checkpointer=checkpointer)
8
9# Session 1
10config = {"configurable": {"thread_id": "user-123-session-1"}}
11app.invoke({"messages": [{"role": "user", "content": "I'm working on an ESP32 project"}]}, config)
12
13# Session 2 — agent remembers
14config2 = {"configurable": {"thread_id": "user-123-session-2"}}
15app.invoke({"messages": [{"role": "user", "content": "What was I working on?"}]}, config2)
16# Agent recalls: "You were working on an ESP32 project"