Agent state machines (LangGraph-shaped)

Open in Colab ▶️ Run this notebook in Colab

Agent state machines (LangGraph-shaped)#

Track 04 - Agents · Notebook 03 · Runtime: ≈30 s on CPU

Prerequisites: 04_agents/01 (ReAct from scratch).

Reference: LangGraph docs - StateGraph.

What#

LangGraph models an agent as a directed graph over a shared state dict. Nodes are pure functions (state) -> state_update; edges can be unconditional or routed by a predicate on the state. The engine runs a loop: execute the current node, apply the state update, route to the next node, repeat until an END node is reached.

This structure is what distinguishes a “real” agent framework from just a for-loop with a parser: you can declare the whole control flow up front, visualise it, and swap out nodes without touching the rest.

We don’t depend on langgraph. A 50-line StateGraph clone built from dict state + edge tables reproduces the core semantics, and makes the shape of the abstraction obvious.

from llm_systems_cookbook.nb import bootstrap

from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any

s = bootstrap("04_agents_03_langgraph_state_machines")

Minimal StateGraph#

Three operations:

  • add_node(name, fn) - register a pure state->update function.

  • add_edge(src, dest) - unconditional transition.

  • add_conditional_edges(src, router, mapping) - router(state) returns a key; mapping[key] is the next node.

run(state) drives the loop from "__start__" until it reaches "__end__" or hits the step limit.

END = "__end__"
START = "__start__"


@dataclass
class StateGraph:
    nodes: dict[str, Callable[[dict], dict]] = field(default_factory=dict)
    edges: dict[str, str] = field(default_factory=dict)
    conditional: dict[str, tuple[Callable[[dict], str], dict[str, str]]] = field(default_factory=dict)
    entry: str = ""

    def add_node(self, name: str, fn: Callable[[dict], dict]) -> "StateGraph":
        self.nodes[name] = fn
        return self

    def add_edge(self, src: str, dest: str) -> "StateGraph":
        self.edges[src] = dest
        return self

    def add_conditional_edges(
        self, src: str, router: Callable[[dict], str], mapping: dict[str, str]
    ) -> "StateGraph":
        self.conditional[src] = (router, mapping)
        return self

    def set_entry(self, name: str) -> "StateGraph":
        self.entry = name
        return self

    def run(self, state: dict, max_steps: int = 30) -> tuple[dict, list[str]]:
        trace: list[str] = []
        cur = self.entry
        for _ in range(max_steps):
            trace.append(cur)
            if cur == END:
                return state, trace
            fn = self.nodes[cur]
            update = fn(state)
            state = {**state, **update}
            if cur in self.conditional:
                router, mapping = self.conditional[cur]
                cur = mapping[router(state)]
            elif cur in self.edges:
                cur = self.edges[cur]
            else:
                cur = END
        return state, trace

An agent that plans, acts, checks, and loops#

Four nodes:

  • plan: generate a next action (here: pick the tool with the first unmet requirement).

  • act: dispatch the action to a tool stub.

  • check: inspect the result; decide whether we’re done.

  • END.

This is the canonical shape for a “reasoning” agent - ReAct expressed as an explicit graph rather than a string-format loop.

# Toy task: an agent must gather three facts about a country before
# answering. State tracks which facts are already in `known`.
FACTS: dict[str, dict[str, str]] = {
    "japan": {"capital": "Tokyo", "population": "125M", "currency": "Yen"},
    "brazil": {"capital": "Brasilia", "population": "215M", "currency": "Real"},
    "egypt": {"capital": "Cairo", "population": "110M", "currency": "Pound"},
}

REQUIRED = ["capital", "population", "currency"]


def plan_node(state: dict) -> dict:
    known = state.get("known", {})
    next_req = next((r for r in REQUIRED if r not in known), None)
    return {"planned_tool": next_req}


def act_node(state: dict) -> dict:
    tool = state["planned_tool"]
    country = state["country"]
    result = FACTS[country].get(tool, "UNKNOWN")
    known = {**state.get("known", {}), tool: result}
    return {"known": known, "last_result": result}


def check_node(state: dict) -> dict:
    done = all(r in state.get("known", {}) for r in REQUIRED)
    answer = None
    if done:
        answer = " / ".join(f"{k}: {state['known'][k]}" for k in REQUIRED)
    return {"done": done, "answer": answer}


def route_after_check(state: dict) -> str:
    return "end" if state["done"] else "loop"


graph = (
    StateGraph()
    .add_node("plan", plan_node)
    .add_node("act", act_node)
    .add_node("check", check_node)
    .add_edge("plan", "act")
    .add_edge("act", "check")
    .add_conditional_edges("check", route_after_check, {"loop": "plan", "end": END})
    .set_entry("plan")
)

final_state, trace = graph.run({"country": "japan", "known": {}})
print(f"trace: {' -> '.join(trace)}")
print(f"answer: {final_state.get('answer')}")

s.check(
    "graph_terminates_within_budget",
    lambda: trace[-1] == END and len(trace) < 30,
    msg=f"trace length = {len(trace)}",
)
s.check(
    "graph_visited_plan_three_times",
    lambda: trace.count("plan") == len(REQUIRED),
    msg=f"plan count = {trace.count('plan')}",
)
s.check(
    "final_answer_contains_all_required_facts",
    lambda: all(k in (final_state.get("answer") or "") for k in ("capital", "population", "currency")),
    msg=f"answer = {final_state.get('answer')}",
)
s.check(
    "first_visited_node_is_entry",
    lambda: trace[0] == "plan",
    msg=f"trace[0] = {trace[0]}",
)

# Run the graph on a second country to check state isolation.
final2, _ = graph.run({"country": "brazil", "known": {}})
s.check(
    "state_resets_per_run",
    lambda: "Brasilia" in (final2.get("answer") or ""),
    msg=f"brazil answer = {final2.get('answer')}",
)

Visit counts and state progress#

The left bar chart counts how many times each node ran for the Japan query; plan/act/check each fire once per required fact, so the counts read off the conditional edge’s loop structure directly. The right panel tracks the size of state["known"] after each act step, showing the monotone climb to the termination condition. Nothing here is a new benchmark - both come from the single graph.run(...) already performed above.

import matplotlib.pyplot as plt
from collections import Counter

counts = Counter(n for n in trace if n != END)
nodes = ["plan", "act", "check"]
visits = [counts.get(n, 0) for n in nodes]

# Per-iteration: how many facts are known after each 'act' step.
known_over_time: list[int] = []
st = {"country": "japan", "known": {}}
cur = graph.entry
for _ in range(30):
    if cur == END:
        break
    st = {**st, **graph.nodes[cur](st)}
    if cur == "act":
        known_over_time.append(len(st.get("known", {})))
    if cur in graph.conditional:
        r, m = graph.conditional[cur]; cur = m[r(st)]
    else:
        cur = graph.edges.get(cur, END)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(9, 3.2))
ax1.bar(nodes, visits, color=["tab:blue", "tab:orange", "tab:green"])
ax1.set_ylabel("visits per run"); ax1.set_title(f"node visit counts (trace len {len(trace)})")
for i, v in enumerate(visits):
    ax1.text(i, v + 0.05, str(v), ha="center")

ax2.plot(range(1, len(known_over_time) + 1), known_over_time, marker="o", color="tab:purple")
ax2.axhline(len(REQUIRED), color="tab:gray", linestyle=":", label=f"target {len(REQUIRED)}")
ax2.set_xlabel("act iteration"); ax2.set_ylabel("facts known")
ax2.set_title("state progress over iterations"); ax2.legend(); ax2.grid(True, alpha=0.3)
fig.tight_layout(); plt.show()

Exercises#

  1. Add a retry node. If a tool returns “UNKNOWN”, the check router should go to a retry node that falls back to a more expensive search tool, then back to check. Draw the graph.

  2. Parallel branches. Real LangGraph supports parallel fan-out and gather. Extend StateGraph with add_parallel(src, [n1, n2, n3], merge_fn) and redo the fact-gathering with three concurrent tool calls.

  3. Real LangGraph. pip install langgraph and express the same agent with langgraph.graph.StateGraph. Compare the shape.

References#

  • LangGraph documentation, Conceptual Guide and State management sections.

  • langgraph-prebuilt for common agent templates (create_react_agent, etc.).

s.summary()
s.save()