Learnixo
Back to blog
AI Systemsintermediate

Creating a StateGraph

Walk through every step of building a LangGraph StateGraph ā€” initialization, registering nodes, connecting edges, and compiling to a runnable Pregel graph.

Asma Hafeez KhanMay 15, 20267 min read
LangGraphAI AgentsStateGraphPregelPython
Share:š•

Creating a StateGraph

The StateGraph class is the core of every LangGraph application. You use it to register node functions, declare how execution flows between them, and compile everything into a runnable graph. This lesson covers every step of that process with complete, working code.

StateGraph Overview

Python
from langgraph.graph import StateGraph, START, END

The full lifecycle of a StateGraph:

  1. Initialize ā€” StateGraph(MyState) binds the graph to a state schema
  2. Add nodes ā€” add_node(name, function) registers callable steps
  3. Add edges ā€” add_edge(A, B) or add_conditional_edges(A, router, map) declares transitions
  4. Set entry point ā€” add_edge(START, first_node) tells the graph where to begin
  5. Compile ā€” compile() validates the graph and returns a CompiledGraph
  6. Invoke ā€” app.invoke(input) runs the graph and returns the final state

Step 1: Initialize with State Schema

Python
from typing import TypedDict, Optional, Annotated
from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, START, END

class ArticleState(TypedDict):
    topic: str
    outline: Optional[str]
    draft: Optional[str]
    critique: Optional[str]
    final_article: Optional[str]
    revision_count: int
    messages: Annotated[list[BaseMessage], add_messages]

# Initialize ā€” no nodes or edges yet
builder = StateGraph(ArticleState)

The StateGraph constructor takes a single argument: the state type. All nodes registered on this graph must accept ArticleState as their first argument and return a dict with a subset of its keys.

Step 2: Write Node Functions

Python
from langchain_openai import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)

def create_outline(state: ArticleState) -> dict:
    """Generate a structured outline for the article."""
    response = llm.invoke(
        f"Create a detailed outline for an article about: {state['topic']}\n"
        f"Include 5-7 main sections with bullet points."
    )
    print(f"[outline] Created outline ({len(response.content)} chars)")
    return {
        "outline": response.content,
        "messages": [AIMessage(content=f"Outline created: {response.content[:100]}...")],
    }

def write_draft(state: ArticleState) -> dict:
    """Write a full draft based on the outline."""
    response = llm.invoke(
        f"Write a comprehensive article based on this outline:\n\n{state['outline']}\n\n"
        f"Topic: {state['topic']}\n"
        f"Write approximately 800 words."
    )
    print(f"[draft] Written ({len(response.content)} chars)")
    return {"draft": response.content}

def critique_draft(state: ArticleState) -> dict:
    """Critique the draft for quality, accuracy, and completeness."""
    response = llm.invoke(
        f"Critique this article draft. Be specific about what is missing, unclear, or weak.\n\n"
        f"Draft:\n{state['draft']}\n\n"
        f"Provide 3-5 specific improvement suggestions."
    )
    print(f"[critique] Generated critique")
    return {
        "critique": response.content,
        "revision_count": state.get("revision_count", 0) + 1,
    }

def revise_draft(state: ArticleState) -> dict:
    """Revise the draft based on critique."""
    response = llm.invoke(
        f"Revise this article based on the critique provided.\n\n"
        f"Original Draft:\n{state['draft']}\n\n"
        f"Critique:\n{state['critique']}\n\n"
        f"Produce an improved version."
    )
    print(f"[revise] Revision {state.get('revision_count', 0)} complete")
    return {"draft": response.content}

def finalize(state: ArticleState) -> dict:
    """Finalize and format the article for publication."""
    final = (
        f"# {state['topic'].title()}\n\n"
        f"{state['draft']}\n\n"
        f"---\n"
        f"*Article underwent {state.get('revision_count', 0)} revision(s)*"
    )
    return {"final_article": final}

Step 3: Register Nodes with add_node

Python
builder.add_node("create_outline", create_outline)
builder.add_node("write_draft", write_draft)
builder.add_node("critique_draft", critique_draft)
builder.add_node("revise_draft", revise_draft)
builder.add_node("finalize", finalize)

add_node takes:

  • name (str) ā€” the identifier used in add_edge calls
  • function (callable) ā€” the node function with signature (state) -> dict

Node names must be unique within a graph. They can contain letters, digits, hyphens, and underscores.

Step 4: Connect Edges

Unconditional Edges

Python
# Fixed flow
builder.add_edge(START, "create_outline")
builder.add_edge("create_outline", "write_draft")
builder.add_edge("write_draft", "critique_draft")

Conditional Edge for the Revision Loop

After critique, we either revise again or finalize:

Python
def should_revise(state: ArticleState) -> str:
    """Decide whether to revise again or finalize."""
    max_revisions = 2
    current_count = state.get("revision_count", 0)

    if current_count < max_revisions:
        # Check if critique suggests major problems
        critique = state.get("critique", "")
        major_issues = any(
            phrase in critique.lower()
            for phrase in ["completely wrong", "major issues", "rewrite", "inaccurate"]
        )
        if major_issues:
            return "revise_draft"

    # No more revisions needed
    return "finalize"

builder.add_conditional_edges(
    "critique_draft",           # from this node
    should_revise,              # call this router
    {
        "revise_draft": "revise_draft",
        "finalize": "finalize",
    }
)
builder.add_edge("revise_draft", "critique_draft")  # loop back for another critique
builder.add_edge("finalize", END)

Step 5: Compile

Python
app = builder.compile()

compile() does the following:

  • Validates that all nodes referenced in edges exist
  • Validates that the graph has an entry point
  • Validates that all paths eventually reach END
  • Returns a CompiledGraph (also called a Pregel graph internally)

If validation fails, you get a descriptive error:

ValueError: Node 'write_draft' referenced in add_edge but not registered with add_node

Compile with Checkpointing

Python
from langgraph.checkpoint.memory import MemorySaver

checkpointer = MemorySaver()
app = builder.compile(checkpointer=checkpointer)

With a checkpointer, state is saved after every node. This enables resumability, time travel, and human-in-the-loop patterns.

Step 6: Invoke, Stream, or astream

invoke ā€” Synchronous, Returns Final State

Python
import os
os.environ["OPENAI_API_KEY"] = "sk-..."

initial_state = {
    "topic": "The future of AI agents in healthcare",
    "outline": None,
    "draft": None,
    "critique": None,
    "final_article": None,
    "revision_count": 0,
    "messages": [],
}

result = app.invoke(initial_state)
print(result["final_article"])

stream ā€” Synchronous, Yields After Each Node

Python
for chunk in app.stream(initial_state):
    node_name = list(chunk.keys())[0]
    state_update = chunk[node_name]
    print(f"\n--- After node: {node_name} ---")
    for key, value in state_update.items():
        if isinstance(value, str):
            print(f"  {key}: {value[:80]}...")
        else:
            print(f"  {key}: {value}")

stream() yields one dict per node execution. The dict key is the node name; the value is the partial state update that node returned. This is ideal for showing progress to the user in real time.

astream ā€” Asynchronous Streaming (for FastAPI, etc.)

Python
import asyncio

async def run_agent():
    async for chunk in app.astream(initial_state):
        node_name = list(chunk.keys())[0]
        print(f"Completed node: {node_name}")

asyncio.run(run_agent())

Getting the Final State After Execution

invoke() returns the complete final state as a dict:

Python
result = app.invoke(initial_state)

# Access any field
print(result["final_article"])
print(f"Total revisions: {result['revision_count']}")
print(f"Messages: {len(result['messages'])}")

When using checkpointing, you can retrieve the state for a specific thread at any time:

Python
config = {"configurable": {"thread_id": "article-123"}}
result = app.invoke(initial_state, config=config)

# Get current state without running the graph
current_state = app.get_state(config)
print(current_state.values["final_article"])

Complete Working Example

Here is the full three-node graph from this lesson, self-contained and runnable:

Python
import os
from typing import TypedDict, Optional, Annotated
from langchain_openai import ChatOpenAI
from langchain_core.messages import BaseMessage, AIMessage
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

os.environ["OPENAI_API_KEY"] = "sk-..."

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)

class ArticleState(TypedDict):
    topic: str
    outline: Optional[str]
    draft: Optional[str]
    critique: Optional[str]
    final_article: Optional[str]
    revision_count: int
    messages: Annotated[list[BaseMessage], add_messages]

def create_outline(state: ArticleState) -> dict:
    response = llm.invoke(f"Create a 5-section outline for an article about: {state['topic']}")
    return {"outline": response.content}

def write_draft(state: ArticleState) -> dict:
    response = llm.invoke(
        f"Write a 500-word article based on this outline:\n{state['outline']}\n\nTopic: {state['topic']}"
    )
    return {"draft": response.content}

def critique_draft(state: ArticleState) -> dict:
    response = llm.invoke(
        f"Provide 3 specific improvements for this article:\n\n{state['draft']}"
    )
    return {
        "critique": response.content,
        "revision_count": state.get("revision_count", 0) + 1,
    }

def revise_draft(state: ArticleState) -> dict:
    response = llm.invoke(
        f"Revise this article based on feedback:\n\nArticle:\n{state['draft']}\n\nFeedback:\n{state['critique']}"
    )
    return {"draft": response.content}

def finalize(state: ArticleState) -> dict:
    return {
        "final_article": f"# {state['topic']}\n\n{state['draft']}",
        "messages": [AIMessage(content="Article finalized.")],
    }

def should_revise(state: ArticleState) -> str:
    return "revise_draft" if state.get("revision_count", 0) < 2 else "finalize"

# Build
builder = StateGraph(ArticleState)
builder.add_node("create_outline", create_outline)
builder.add_node("write_draft", write_draft)
builder.add_node("critique_draft", critique_draft)
builder.add_node("revise_draft", revise_draft)
builder.add_node("finalize", finalize)

builder.add_edge(START, "create_outline")
builder.add_edge("create_outline", "write_draft")
builder.add_edge("write_draft", "critique_draft")
builder.add_conditional_edges("critique_draft", should_revise,
    {"revise_draft": "revise_draft", "finalize": "finalize"})
builder.add_edge("revise_draft", "critique_draft")
builder.add_edge("finalize", END)

app = builder.compile()

# Run
result = app.invoke({
    "topic": "Vector databases in production AI systems",
    "outline": None, "draft": None, "critique": None,
    "final_article": None, "revision_count": 0, "messages": [],
})

print(result["final_article"])
print(f"\nRevisions: {result['revision_count']}")

Common Errors and Fixes

Error: "Node X referenced in edge but not registered"

Python
# Wrong ā€” typo in node name
builder.add_edge("retreive", "generate")  # should be "retrieve"

# Fix ā€” names must match exactly
builder.add_node("retrieve", retrieve_fn)
builder.add_edge("retrieve", "generate")

Error: "Graph has no entry point"

Python
# Missing START edge
builder.add_edge(START, "first_node")  # add this

Error: "Graph has unreachable nodes"

Python
# A node was added but no edge leads to it
# Either add an edge to it or remove the node

Summary

Building a StateGraph follows five deterministic steps:

  1. StateGraph(StateSchema) ā€” bind to state type
  2. add_node(name, fn) ā€” register each node function
  3. add_edge(A, B) or add_conditional_edges(A, router, map) ā€” wire the graph
  4. add_edge(START, first_node) ā€” set entry point
  5. compile() ā€” validate and get runnable graph

The compiled graph exposes invoke(), stream(), and astream() for synchronous, streaming, and async execution respectively. Every invocation returns the complete final state, making it easy to extract results and debug execution.

Enjoyed this article?

Explore the AI Systems learning path for more.

Subscribe for weekly tutorials

Found this helpful?

Share:š•

Leave a comment

Have a question, correction, or just found this helpful? Leave a note below.