StateGraph for Agents: A Practical Starter (with code)

Most “agent” projects fail from workflow chaos, not model choice. This primer shows how to use LangGraph’s StateGraph to build reliable, stateful multi-step flows—with typed state, conditional edges, and copy-paste code you can ship today.

LangGraph’s StateGraph is a core building block for designing stateful, multi-step agent workflows. Each node represents an operation; edges define the flow; a shared State carries data across nodes. Below is a concise, hands-on guide with runnable examples.


What is StateGraph?

StateGraph lets you define a directed graph:

  • Nodes: functions that read/update a shared state.
  • Edges: connections (including conditional branches) between nodes.
  • State: a typed schema describing what data flows through the graph.

You can set an entry point, add conditional edges, compile the graph, and invoke it with an initial state.


Quick Start (Single-schema)

1) Define the state schema

from typing import TypedDict
from langgraph.graph import StateGraph, START, END

class GraphState(TypedDict):
    input: str
    output: str
    step_count: int        
2) Create a graph
workflow = StateGraph(GraphState)
        

3) Define node functions

def node_1(state: GraphState) -> dict:
    # Return partial updates that merge into the shared state
    return {
        "output": f"Processed: {state['input']}",
        "step_count": state.get("step_count", 0) + 1
    }

def node_2(state: GraphState) -> dict:
    return {
        "output": state["output"] + " -> further processing",
        "step_count": state["step_count"] + 1
    }
        

4) Add nodes and edges

workflow.add_node("step1", node_1)
workflow.add_node("step2", node_2)

workflow.add_edge(START, "step1")
workflow.add_edge("step1", "step2")
workflow.add_edge("step2", END)
        

5) (Optional) Conditional edges

def should_continue(state: GraphState) -> str:
    return "continue" if state["step_count"] < 3 else "end"

workflow.add_conditional_edges(
    "step1",
    should_continue,
    {
        "continue": "step2",
        "end": END
    }
)
        

6) Compile & run

app = workflow.compile()

initial_state = {"input": "Hello World", "output": "", "step_count": 0}
result = app.invoke(initial_state)
print(result)  # {'input': 'Hello World', 'output': 'Processed: Hello World -> further processing', 'step_count': 2}
        

What is State in LangGraph?

State defines:

  • The schema of your graph’s data (types and keys)
  • How updates are reduced/merged as nodes run
  • The contract for inputs/outputs vs. internal fields

Core principles

  • Type safety: Use TypedDict or Pydantic for clear contracts.
  • Separation of concerns: Distinguish input, output, and internal state.
  • Flexibility: Any node can read/write declared state channels.
  • Encapsulation: Keep private/intermediate fields internal.

Three related schemas

  • state_schema – full internal state of the graph (required)
  • input_schema – what the graph accepts as input (optional; subset of state_schema)
  • output_schema – what the graph returns (optional; subset of state_schema)

Flow: Input → [input_schema] → [state_schema] → [output_schema] → Output (Filtered In)     (Full Internal)   (Filtered Out)

This gives you:

  • Encapsulation: internal fields aren’t exposed
  • Clean interfaces: clear input/output contracts
  • Flexibility: add internal fields for node-to-node communication


Example: Split Input/Output from Internal State

Below is a realistic multi-schema setup. We keep “private” fields inside the overall state (prefixed with _) so all nodes operate on a single schema while still distinguishing external interfaces.

from typing import TypedDict, Optional
from langgraph.graph import StateGraph, START, END

# --- Public I/O schemas ---
class InputState(TypedDict):
    user_input: str

class OutputState(TypedDict):
    graph_output: str
    step_count: int

# --- Full internal state (used by the graph) ---
class OverallState(TypedDict, total=False):
    # Publicly relevant
    user_input: str
    intermediate_result: str
    processed_data: str
    graph_output: str
    step_count: int

    # “Private”/internal fields
    _private_data: str
    _internal_flag: bool

# --- Node functions (return partial updates) ---
def input_node(state: OverallState) -> dict:
    return {
        "user_input": state["user_input"],
        "intermediate_result": f"Start: {state['user_input']}",
        "step_count": 1
    }

def processing_node(state: OverallState) -> dict:
    processed = state["intermediate_result"].upper()
    return {
        "_private_data": f"Private: {processed}",
        "_internal_flag": len(state["user_input"]) > 5
    }

def intermediate_node(state: OverallState) -> dict:
    suffix = " (long text)" if state.get("_internal_flag") else " (short text)"
    return {
        "processed_data": state["_private_data"] + suffix
    }

def output_node(state: OverallState) -> dict:
    final_result = f"Final: {state['processed_data']} | steps: {state.get('step_count', 0)}"
    return {
        "graph_output": final_result
    }

# --- Build graph with I/O separation ---
builder = StateGraph(
    state_schema=OverallState,
    input_schema=InputState,
    output_schema=OutputState,
)

builder.add_node("input_processor", input_node)
builder.add_node("data_processor", processing_node)
builder.add_node("intermediate_processor", intermediate_node)
builder.add_node("output_processor", output_node)

builder.add_edge(START, "input_processor")
builder.add_edge("input_processor", "data_processor")
builder.add_edge("data_processor", "intermediate_processor")
builder.add_edge("intermediate_processor", "output_processor")
builder.add_edge("output_processor", END)

graph = builder.compile()

# --- Test run ---
input_data = {"user_input": "Hello LangGraph", "extra_input": "ignored"}
print("Input:", input_data)

result = graph.invoke(input_data)
print("Output:", result)
# Example:
# Input: {'user_input': 'Hello LangGraph', 'extra_input': 'ignored'}
# Output: {'graph_output': 'Final: Private: START: HELLO LANGGRAPH (long text) | steps: 1',
#          'step_count': 1}
        
Tip: You can render the graph if your environment supports it:

Why this design works

  • Encapsulation: Public APIs stay small and clean; internal plumbing stays internal.
  • Composability: Add nodes, edges, or private fields without breaking callers.
  • Safety: Typed schemas make large multi-agent systems more reliable.
  • Flexibility: Conditional edges let you branch based on state at runtime.


TL;DR

  • Use StateGraph for multi-step, stateful agent flows.
  • Keep a strict state_schema for internals, and lighter input_schema/output_schema for clean I/O.
  • Nodes return partial state updates; LangGraph reduces them into the shared state.
  • Add conditional edges for dynamic routing.

To view or add a comment, sign in

More articles by Bei-Bei Wang

Others also viewed

Explore content categories