LangGraph Section Flow Middleware

Section-based flow control for LangGraph React agents. Divide a conversational agent into discrete, self-contained phases — each with its own prompt, tools, and transition rules — without writing a custom graph.

Why section flow?

A typical React agent sees all tools and the full system prompt on every call. As workflows grow (qualification → recommendation → booking → payment), this creates two problems:

1. Context pollution — the agent is distracted by tools and instructions that are irrelevant to the current step.
2. Unclear guardrails — it's hard to restrict what the agent can do at each stage without building a bespoke multi-node graph.

SectionFlowMiddleware solves both by layering a lightweight state machine on top of create_react_agent. You describe sections in plain Python; the middleware handles everything else.

Features

Installation

Install the latest release from PyPI:

pip install langgraph-state-machine

You'll also need a model provider package for whichever LLM you use, e.g.:

pip install langchain-openai      # OpenAI
pip install langchain-anthropic   # Anthropic

Requirements: Python ≥ 3.10, langchain ≥ 1.0.0, langgraph ≥ 0.2.0

To hack on the library itself:

git clone https://github.com/mahmoud661/langgraph-state-machine
cd langgraph-state-machine
pip install -e ".[dev]"
pytest

How to use it

Using the middleware is four steps: define your tools, group them into sections, wrap the sections in a SectionFlowMiddleware, and pass that middleware to your agent. The package installs as section_flow:

import os

from langchain.agents import create_react_agent
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

from section_flow import SectionFlowMiddleware, SectionConfig

# --- 1. Define your tools ------------------------------------------------
@tool
def save_user_info(name: str, budget: int, category: str) -> str:
    """Save the user's name, budget, and product category."""
    return f"Saved: name={name}, budget={budget}, category={category}"

@tool
def search_catalog(query: str, max_price: int) -> str:
    """Search the product catalog and return matching items."""
    return f"Top results for '{query}' under ${max_price}: ..."

@tool
def process_payment(item_name: str, amount: float) -> str:
    """Process payment for the selected item."""
    return f"Payment of ${amount:.2f} for '{item_name}' processed."

# --- 2. Describe your workflow as sections -------------------------------
sections = {
    "gather": SectionConfig(
        name="gather",
        prompt=(
            "Collect the user's name, budget, and product category. "
            "Once you have all three, call save_user_info and transition to 'recommend'."
        ),
        tools=[save_user_info],
        allowed_transitions=["recommend"],
    ),
    "recommend": SectionConfig(
        name="recommend",
        prompt="Present exactly three options, then move to 'checkout' when the user picks one.",
        tools=[search_catalog],
        allowed_transitions=["checkout"],
    ),
    "checkout": SectionConfig(
        name="checkout",
        prompt="Confirm the item and total price, then process payment.",
        tools=[process_payment],
        allowed_transitions=[],  # terminal section
    ),
}

# --- 3. Create the middleware ---------------------------------------------
middleware = SectionFlowMiddleware(
    sections=sections,
    initial_section="gather",
    fallback_section="gather",
)

# --- 4. Attach it to your agent -------------------------------------------
agent = create_react_agent(
    model=ChatOpenAI(model="gpt-4o", api_key=os.environ["OPENAI_API_KEY"]),
    system_prompt="You are a friendly shopping assistant.",
    middleware=[middleware],
)

result = agent.invoke(
    {"messages": [{"role": "user", "content": "Hi, I need a new laptop."}]}
)
print(result["current_section"])   # section the agent ended in
print(result["messages"][-1].content)

The agent starts in gather and only sees save_user_info plus the built-in change_section tool. When the model decides it has what it needs, it calls change_section(target_section="recommend") — from then on it sees the recommend prompt and tools instead. You never write graph nodes or edges; the middleware rewrites the prompt and tool list before every model call based on current_section in the agent state.

Runnable end-to-end scripts live in examples/.

Core concepts

SectionConfig

Each section is a pydantic.BaseModel:

SectionConfig(
    name="gather",                        # unique identifier
    prompt="...",                         # injected system message fragment
    tools=[my_tool],                      # tool objects available in this section
    allowed_transitions=["recommend"],    # sections this one may transition to
    required_state_fields={"budget": int},# fields that must exist in section_data before entering
    auto_transition_conditions=lambda state: (
        "recommend" if state.get("section_data", {}).get("budget") else None
    ),
    strict_validation=True,               # enforce allowed_transitions and required_state_fields
    on_enter=lambda state: None,          # lifecycle hook (called on entry)
    on_exit=lambda state: None,           # lifecycle hook (called on exit)
    llm=ChatOpenAI(model="gpt-4o-mini"), # optional per-section model override
    allowed_subagents=["research_agent"], # limit task-tool subagents for this section
)

SectionFlowState

The middleware extends your agent state with three fields:

class MyState(SectionFlowState):         # merge with your existing state
    messages: Annotated[list, add_messages]

# Fields added by the middleware:
# current_section  – name of the active section
# section_data     – shared dict for cross-section data (e.g. {"budget": 1000})
# visited_sections – ordered list of activated sections

Transitions

There are three ways to advance the flow:

Auto-transition (callable)

SectionConfig(
    name="gather",
    auto_transition_conditions=lambda state: (
        "recommend"
        if state.get("section_data", {}).get("budget")
        else None
    ),
    ...
)

Auto-transition (priority list)

from section_flow import TransitionCondition

SectionConfig(
    name="gather",
    auto_transition_conditions=[
        TransitionCondition(
            target="vip_recommend",
            condition=lambda s: s.get("section_data", {}).get("budget", 0) > 5000,
            priority=10,
        ),
        TransitionCondition(
            target="recommend",
            condition=lambda s: bool(s.get("section_data", {}).get("budget")),
            priority=0,
        ),
    ],
    ...
)

Per-section LLM

from langchain_openai import ChatOpenAI

SectionConfig(
    name="checkout",
    prompt="Process payment carefully.",
    tools=[process_payment],
    llm=ChatOpenAI(model="gpt-4o"),  # override the graph's default model here
)

Advanced usage

Pre-built SectionManager

Reuse the same manager across multiple agents:

from section_flow import SectionManager, SectionFlowMiddleware

manager = SectionManager(
    sections=sections,
    initial_section="gather",
    fallback_section="gather",
)

agent1 = create_react_agent(..., middleware=[SectionFlowMiddleware(section_manager=manager)])
agent2 = create_react_agent(..., middleware=[SectionFlowMiddleware(section_manager=manager)])

Global tools

Tools that should always be available regardless of section:

SectionFlowMiddleware(
    sections=sections,
    initial_section="gather",
    global_tools=[escalate_to_human],   # overrides section tools with same name
)

Disable the transition tool

If you prefer purely automatic or state-driven transitions:

SectionFlowMiddleware(
    sections=sections,
    initial_section="gather",
    include_transition_tool=False,
)

Runtime cache invalidation

If you modify sections at runtime (hot-reload), clear the tool cache:

middleware.clear_tool_cache()

API reference

SectionFlowMiddleware

SectionConfig

Examples

See the examples/ directory:

How it works

User message
     │
     ▼
before_model()          ← initialise state, resolve fallbacks, fire auto-transitions
     │
     ▼
wrap_model_call()       ← prepend section prompt, filter tools, swap model if needed
     │
     ▼
LLM call
     │
     ▼
[agent calls change_section tool]  ← updates current_section in state
     │
     ▼
next before_model() ...

Contributing

See CONTRIBUTING.md.

License

Apache 2.0 — see LICENSE.