Module 4: Middleware and Customization

The AgentMiddleware Class: Composed Middleware

Capsule overview

In the previous capsules you learned to customize agents with individual hooks: @before_model to intercept before the model call, @after_model to act after the response, @wrap_model_call to wrap the whole call, and @wrap_tool_call to control tool execution. You also saw how to use dynamic models, dynamic tools and dynamic prompts to change the agent's behavior at runtime.

But there's a problem that shows up fast in real projects: the hooks are scattered. Your file has a @before_model for logging over here, a @wrap_tool_call for auth over there, a @wrap_model_call for model routing further down. When a second agent needs the same logging, you copy and paste. When a third one needs logging + auth but not routing, you copy selectively. Within a few weeks you have duplicated, inconsistent hooks that are painful to maintain.

The AgentMiddleware class solves this. It lets you package several hooks, custom state and tools into a single reusable unit — like a plugin you can install in any agent. You define a class that inherits from AgentMiddleware, implement the hooks you need as methods, and hand it to the agent. When you have several middleware, you compose them in a list: middleware=[LoggingMiddleware(), AuthMiddleware(), CachingMiddleware()]. LangChain runs them in order, creating layers of behavior that stack without stepping on each other.

The problem: individual hooks don't scale

Imagine you have three agents in your system and each one needs a different combination of behaviors:

Agent	Logging	Auth	Rate limiting
Support agent	✅	✅	✅
Internal agent	✅	❌	❌
Analytics agent	✅	✅	❌

With individual hooks, your code looks like this:

from dotenv import load_dotenv
load_dotenv()

from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

model = init_chat_model("openai:gpt-4.1-mini")

# Support agent: needs logging + auth + rate limiting
support_agent = create_agent(
    model, [search],
    prompt="You are a support agent.",
    before_model=log_before,       # ← loose hook
    after_model=log_after,         # ← loose hook
    wrap_tool_call=auth_tool,      # ← loose hook
    wrap_model_call=rate_limit,    # ← loose hook
)

# Internal agent: logging only
internal_agent = create_agent(
    model, [search],
    prompt="You are an internal agent.",
    before_model=log_before,       # ← copied
    after_model=log_after,         # ← copied
)

# Analytics agent: logging + auth
analysis_agent = create_agent(
    model, [search],
    prompt="You are an analytics agent.",
    before_model=log_before,       # ← copied again
    after_model=log_after,         # ← copied again
    wrap_tool_call=auth_tool,      # ← copied
)

Three agents, and you've already copied hooks 6 times. If you want to change how logging works, you have to update every agent. If you add a fourth agent with logging + rate limiting, you copy again. This pattern doesn't scale.

The AgentMiddleware class: the solution

AgentMiddleware is a base class that combines hooks, state and tools into a single unit. Instead of passing individual hooks, you create classes that encapsulate behavior and pass them to the agent as a list.

The basic structure

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

class LoggingMiddleware(AgentMiddleware):
    """Log every operation the agent performs."""

    def before_model(self, messages, config):
        print(f"[LOG] Model call with {len(messages)} messages")

    def after_model(self, response, config):
        preview = response.content[:80] if response.content else "(no content)"
        print(f"[LOG] Model answered: {preview}")

    def wrap_tool_call(self, tool_call, config, call_next):
        print(f"[LOG] Running tool: {tool_call['name']}")
        result = call_next(tool_call, config)
        preview = str(result)[:80]
        print(f"[LOG] Result: {preview}")
        return result

@tool
def search(query: str) -> str:
    """Search for information about a topic."""
    return f"Results about {query}: Python is a versatile language..."

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are a research assistant.",
    middleware=[LoggingMiddleware()],
)

result = agent.invoke(
    {"messages": [("user", "What is Python?")]}
)
print(result["messages"][-1].content)
# Expected output:
# [LOG] Model call with 2 messages
# [LOG] Running tool: search
# [LOG] Result: Results about Python: Python is a versatile language...
# [LOG] Model call with 4 messages
# [LOG] Model answered: Python is a high-level programming language...
# Python is a high-level programming language...

Every method is optional. You only implement the hooks you need:

✅ before_model(messages, config) — runs before each model call
✅ after_model(response, config) — runs after each model response
✅ wrap_model_call(messages, config, call_next) — wraps the whole model call
✅ wrap_tool_call(tool_call, config, call_next) — wraps each tool execution
✅ get_tools(state) — returns dynamic tools based on the state
✅ get_prompt(state) — returns a dynamic prompt based on the state

Building a middleware class step by step

Let's build an authentication middleware that validates permissions before running tools and filters which tools are available based on the user's role.

Step 1: define the middleware

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for general information."""
    return f"Results for: {query}"

@tool
def calculator(expression: str) -> str:
    """Evaluate math expressions."""
    try:
        return str(eval(expression))
    except Exception as e:
        return f"Error: {e}"

@tool
def admin_panel(action: str) -> str:
    """Run administrative actions. Admins only."""
    return f"Administrative action executed: {action}"

class AuthMiddleware(AgentMiddleware):
    """Middleware that controls access by user role."""

    def get_tools(self, state):
        user_role = state.get("user_role", "basic")
        if user_role == "admin":
            return [search, calculator, admin_panel]
        return [search, calculator]

    def wrap_tool_call(self, tool_call, config, call_next):
        state = config.get("configurable", {})
        user_role = state.get("user_role", "basic")
        tool_name = tool_call["name"]

        if tool_name == "admin_panel" and user_role != "admin":
            return "Error: You don't have permission to run administrative actions."

        print(f"[AUTH] User ({user_role}) running: {tool_name}")
        return call_next(tool_call, config)

Step 2: use the middleware in an agent

from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    user_role: str

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model,
    tools=[search, calculator, admin_panel],
    prompt="You are an assistant. Use the available tools to help the user.",
    state_schema=AgentState,
    middleware=[AuthMiddleware()],
)

# Basic user
result_basic = agent.invoke({
    "messages": [("user", "Run an administrative action")],
    "user_role": "basic",
})
print(result_basic["messages"][-1].content)
# Expected output:
# [AUTH] User (basic) running: admin_panel
# I don't have permission to run administrative actions...

# Admin user
result_admin = agent.invoke({
    "messages": [("user", "Run an administrative backup action")],
    "user_role": "admin",
})
print(result_admin["messages"][-1].content)
# Expected output:
# [AUTH] User (admin) running: admin_panel
# The administrative backup action was executed successfully.

The middleware filters tools in get_tools (the basic user never even sees admin_panel) and also validates in wrap_tool_call as a second layer of security.

Composing several middleware

Composition is where AgentMiddleware shines. You pass a list of middleware and LangChain runs them in order — the first middleware in the list is the outermost layer.

from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class LoggingMiddleware(AgentMiddleware):
    """Layer 1: Log every operation."""

    def wrap_model_call(self, messages, config, call_next):
        start = time.time()
        print(f"[LOG] → Model call ({len(messages)} msgs)")
        response = call_next(messages, config)
        elapsed = time.time() - start
        print(f"[LOG] ← Response in {elapsed:.2f}s")
        return response

    def wrap_tool_call(self, tool_call, config, call_next):
        print(f"[LOG] 🔧 Tool: {tool_call['name']}")
        result = call_next(tool_call, config)
        print(f"[LOG] 📥 Result: {str(result)[:60]}")
        return result

class CachingMiddleware(AgentMiddleware):
    """Layer 2: Cache tool responses."""

    def __init__(self):
        self._cache = {}

    def wrap_tool_call(self, tool_call, config, call_next):
        cache_key = f"{tool_call['name']}:{tool_call['args']}"
        if cache_key in self._cache:
            print(f"[CACHE] Hit for {tool_call['name']}")
            return self._cache[cache_key]

        result = call_next(tool_call, config)
        self._cache[cache_key] = result
        print(f"[CACHE] Stored: {tool_call['name']}")
        return result

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are an assistant. Search for information when asked.",
    middleware=[LoggingMiddleware(), CachingMiddleware()],
)

result = agent.invoke(
    {"messages": [("user", "Search about Python and then about Python again")]}
)
print(result["messages"][-1].content)
# Expected output:
# [LOG] → Model call (2 msgs)
# [LOG] ← Response in 1.23s
# [LOG] 🔧 Tool: search
# [CACHE] Stored: search
# [LOG] 📥 Result: Results for: Python
# [LOG] → Model call (4 msgs)
# [LOG] ← Response in 0.89s
# [LOG] 🔧 Tool: search
# [CACHE] Hit for search
# [LOG] 📥 Result: Results for: Python
# ...

Order matters: how middleware stack

When you pass middleware=[A(), B(), C()], LangChain runs them like this:

A.wrap_model_call →
    B.wrap_model_call →
        C.wrap_model_call →
            the real model
        ← C returns
    ← B returns
← A returns

The first middleware in the list is the outermost layer (the first to run and the last to receive the result). That matters when order changes the behavior:

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class TimerMiddleware(AgentMiddleware):
    """Measure total time (including the time spent in inner middleware)."""

    def wrap_model_call(self, messages, config, call_next):
        import time
        start = time.time()
        response = call_next(messages, config)
        elapsed = time.time() - start
        print(f"[TIMER] Total time (model + inner middleware): {elapsed:.2f}s")
        return response

class RetryMiddleware(AgentMiddleware):
    """Retry when the model fails."""

    def wrap_model_call(self, messages, config, call_next):
        max_retries = 3
        for attempt in range(max_retries):
            try:
                return call_next(messages, config)
            except Exception as e:
                if attempt == max_retries - 1:
                    raise
                print(f"[RETRY] Attempt {attempt + 1} failed: {e}. Retrying...")

model = init_chat_model("openai:gpt-4.1-mini")

# Timer is the outer layer → it measures total time, retries included
agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[TimerMiddleware(), RetryMiddleware()],
)

result = agent.invoke(
    {"messages": [("user", "Search about Python")]}
)
# Expected output:
# [TIMER] Total time (model + inner middleware): 1.45s

If you flip the order ([RetryMiddleware(), TimerMiddleware()]), the timer would measure only the model's time, without the retries. The order you pick depends on what you want to measure or control.

A practical ordering rule:

Position	Middleware type	Why
First (outer)	Logging / Metrics	Captures everything, including retries and errors
Middle	Auth / Validation	Blocks before anything reaches the model
Last (inner)	Retry / Cache	Acts directly on the model call

Middleware as reusable modules

The real advantage is that an AgentMiddleware is just a regular Python class — you can parameterize it, inherit from it, and share it across projects.

Parameterizable middleware

from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class ConfigurableLoggingMiddleware(AgentMiddleware):
    """Logging with a configurable level."""

    def __init__(self, level: str = "info", log_tools: bool = True, log_model: bool = True):
        self.level = level
        self.log_tools = log_tools
        self.log_model = log_model

    def before_model(self, messages, config):
        if self.log_model:
            print(f"[{self.level.upper()}] Model: {len(messages)} messages")

    def wrap_tool_call(self, tool_call, config, call_next):
        if self.log_tools:
            print(f"[{self.level.upper()}] Tool: {tool_call['name']}")
        return call_next(tool_call, config)

model = init_chat_model("openai:gpt-4.1-mini")

# Development: full logging
dev_agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[ConfigurableLoggingMiddleware(level="debug", log_tools=True, log_model=True)],
)

# Production: tools only
prod_agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[ConfigurableLoggingMiddleware(level="warn", log_tools=True, log_model=False)],
)

result = dev_agent.invoke({"messages": [("user", "Search Python")]})
# Expected output:
# [DEBUG] Model: 2 messages
# [DEBUG] Tool: search
# [DEBUG] Model: 4 messages

Middleware with internal state

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class MetricsMiddleware(AgentMiddleware):
    """Collect usage metrics."""

    def __init__(self):
        self.model_calls = 0
        self.tool_calls = 0
        self.total_input_messages = 0

    def before_model(self, messages, config):
        self.model_calls += 1
        self.total_input_messages += len(messages)

    def wrap_tool_call(self, tool_call, config, call_next):
        self.tool_calls += 1
        return call_next(tool_call, config)

    def get_metrics(self) -> dict:
        return {
            "model_calls": self.model_calls,
            "tool_calls": self.tool_calls,
            "total_input_messages": self.total_input_messages,
        }

metrics = MetricsMiddleware()

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are a research assistant.",
    middleware=[metrics],
)

result = agent.invoke({"messages": [("user", "Search about Python and TypeScript")]})

print(metrics.get_metrics())
# Expected output:
# {'model_calls': 2, 'tool_calls': 2, 'total_input_messages': 6}

Because metrics is an instance that keeps state, you can query the metrics after every invocation. That's especially handy for monitoring dashboards.

Common built-in patterns

These are the middleware patterns that show up in nearly every production project. Use them as a base and adapt them.

Rate limiting middleware

from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class RateLimitMiddleware(AgentMiddleware):
    """Limit how often the model is called."""

    def __init__(self, max_calls_per_minute: int = 20):
        self.max_calls = max_calls_per_minute
        self._timestamps: list[float] = []

    def wrap_model_call(self, messages, config, call_next):
        now = time.time()
        self._timestamps = [t for t in self._timestamps if now - t < 60]

        if len(self._timestamps) >= self.max_calls:
            wait_time = 60 - (now - self._timestamps[0])
            print(f"[RATE] Limit reached. Waiting {wait_time:.1f}s...")
            time.sleep(wait_time)

        self._timestamps.append(time.time())
        return call_next(messages, config)

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[RateLimitMiddleware(max_calls_per_minute=10)],
)

result = agent.invoke({"messages": [("user", "Search about Python")]})
print(result["messages"][-1].content)
# Expected output (no rate limit hit):
# Python is a programming language...

Error recovery middleware

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def unreliable_api(query: str) -> str:
    """An API that can fail."""
    import random
    if random.random() < 0.3:
        raise ConnectionError("API temporarily unavailable")
    return f"Data about: {query}"

class ErrorRecoveryMiddleware(AgentMiddleware):
    """Handle tool errors with retry and fallback."""

    def __init__(self, max_retries: int = 2, fallback_message: str = "Service unavailable"):
        self.max_retries = max_retries
        self.fallback_message = fallback_message

    def wrap_tool_call(self, tool_call, config, call_next):
        for attempt in range(self.max_retries + 1):
            try:
                return call_next(tool_call, config)
            except Exception as e:
                if attempt < self.max_retries:
                    print(f"[RECOVERY] {tool_call['name']} failed (attempt {attempt + 1}): {e}")
                    continue
                print(f"[RECOVERY] {tool_call['name']} failed for good. Falling back.")
                return self.fallback_message

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [unreliable_api],
    prompt="You are an assistant. Search for information when asked.",
    middleware=[ErrorRecoveryMiddleware(max_retries=2)],
)

result = agent.invoke({"messages": [("user", "Find data about ML")]})
print(result["messages"][-1].content)
# Expected output (with a successful retry):
# [RECOVERY] unreliable_api failed (attempt 1): API temporarily unavailable
# ML (Machine Learning) is a branch of artificial intelligence...

Building a middleware library for your organization

When you have stable middleware used across several projects, the next step is packaging them as a Python module any team can import.

Recommended structure

my_middleware/
├── __init__.py
├── logging.py
├── auth.py
├── rate_limiting.py
├── caching.py
└── metrics.py

Example of a shared module

# my_middleware/__init__.py
from .logging import LoggingMiddleware
from .auth import AuthMiddleware
from .rate_limiting import RateLimitMiddleware
from .caching import CachingMiddleware
from .metrics import MetricsMiddleware

__all__ = [
    "LoggingMiddleware",
    "AuthMiddleware",
    "RateLimitMiddleware",
    "CachingMiddleware",
    "MetricsMiddleware",
]

# my_middleware/logging.py
from langchain.agents import AgentMiddleware

class LoggingMiddleware(AgentMiddleware):
    """Standardized logging for every agent in the organization."""

    def __init__(self, service_name: str = "default", log_level: str = "info"):
        self.service_name = service_name
        self.log_level = log_level

    def before_model(self, messages, config):
        print(f"[{self.service_name}] [{self.log_level.upper()}] "
              f"Model call: {len(messages)} messages")

    def after_model(self, response, config):
        token_count = getattr(response, "usage_metadata", {})
        print(f"[{self.service_name}] [{self.log_level.upper()}] "
              f"Model response: {len(response.content)} chars")

    def wrap_tool_call(self, tool_call, config, call_next):
        print(f"[{self.service_name}] Tool: {tool_call['name']}")
        result = call_next(tool_call, config)
        print(f"[{self.service_name}] Tool done: {tool_call['name']}")
        return result

# Usage in any project across the organization
from dotenv import load_dotenv
load_dotenv()

from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain_core.tools import tool

from my_middleware import LoggingMiddleware, RateLimitMiddleware

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[
        LoggingMiddleware(service_name="support-bot", log_level="debug"),
        RateLimitMiddleware(max_calls_per_minute=30),
    ],
)

result = agent.invoke({"messages": [("user", "Search about Python")]})
# Expected output:
# [support-bot] [DEBUG] Model call: 2 messages
# [support-bot] [DEBUG] Model response: 156 chars
# [support-bot] Tool: search
# [support-bot] Tool done: search

Troubleshooting

1. The middleware hooks never run

Cause: The method name doesn't match exactly. The valid names are before_model, after_model, wrap_model_call, wrap_tool_call, get_tools, get_prompt.

Fix: Double-check that the method names are exact. Python won't warn you if you define before_Model (capital M) — it simply won't run.

2. `wrap_model_call` and `before_model` both run

Cause: That's correct behavior. before_model runs first, then wrap_model_call. They aren't mutually exclusive.

Fix: If you only need one, implement only that one. Use before_model / after_model for simple side effects (logging, metrics). Use wrap_model_call when you need to modify the call or its response.

3. Stateful middleware shares state across invocations

Cause: If you use the same middleware instance for several agents or invocations, the state (self._cache, self.model_calls, etc.) accumulates.

Fix: That's intentional for middleware like MetricsMiddleware (you want the metrics to accumulate). For middleware where you want fresh state, create a new instance per agent or implement a reset() method.

4. `call_next` isn't available in `before_model`

Cause: before_model and after_model are side-effect hooks — they don't wrap the call. Only wrap_model_call and wrap_tool_call receive call_next.

Fix: If you need to modify the model call (not just observe it), use wrap_model_call instead of before_model.

5. The middleware order doesn't produce the result you expected

Cause: Remember that the first middleware in the list is the outermost layer. If A comes before B, A.wrap_model_call wraps B.wrap_model_call.

Fix: Draw the flow as onion layers: [outer → inner]. Logging goes first (it captures everything), retry goes last (right on top of the model).

Exercises

Exercise 1: A basic timestamp middleware (Basic)

Build a TimestampMiddleware that prints the date and time of every model call and every tool execution.

See solution

from dotenv import load_dotenv
load_dotenv()

from datetime import datetime
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class TimestampMiddleware(AgentMiddleware):
    """Add timestamps to every operation."""

    def before_model(self, messages, config):
        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{now}] Model call ({len(messages)} msgs)")

    def wrap_tool_call(self, tool_call, config, call_next):
        now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{now}] Tool: {tool_call['name']}")
        result = call_next(tool_call, config)
        after = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        print(f"[{after}] Tool finished: {tool_call['name']}")
        return result

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[TimestampMiddleware()],
)

result = agent.invoke({"messages": [("user", "Search about Python")]})
print(result["messages"][-1].content)
# Expected output:
# [2026-02-28 15:30:00] Model call (2 msgs)
# [2026-02-28 15:30:01] Tool: search
# [2026-02-28 15:30:01] Tool finished: search
# [2026-02-28 15:30:01] Model call (4 msgs)
# Python is a programming language...

Exercise 2: A token-counting middleware (Basic)

Build a TokenCounterMiddleware that counts how many model calls were made and how many tools were run. The middleware should have a report() method that prints a summary.

See solution

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

@tool
def calculator(expression: str) -> str:
    """Evaluate a math expression."""
    return str(eval(expression))

class TokenCounterMiddleware(AgentMiddleware):
    """Count model calls and tool executions."""

    def __init__(self):
        self.model_calls = 0
        self.tool_calls = 0
        self.tools_used: dict[str, int] = {}

    def before_model(self, messages, config):
        self.model_calls += 1

    def wrap_tool_call(self, tool_call, config, call_next):
        self.tool_calls += 1
        name = tool_call["name"]
        self.tools_used[name] = self.tools_used.get(name, 0) + 1
        return call_next(tool_call, config)

    def report(self):
        print(f"\n📊 Usage report:")
        print(f"   Model calls: {self.model_calls}")
        print(f"   Tool executions: {self.tool_calls}")
        print(f"   Tools used: {self.tools_used}")

counter = TokenCounterMiddleware()
model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search, calculator],
    prompt="You are a research assistant.",
    middleware=[counter],
)

result = agent.invoke(
    {"messages": [("user", "Search about Python and compute 2**10")]}
)
counter.report()
# Expected output:
# 📊 Usage report:
#    Model calls: 2
#    Tool executions: 2
#    Tools used: {'search': 1, 'calculator': 1}

Exercise 3: Compose logging + rate limiting (Intermediate)

Build two middleware (SimpleLogger and SimpleRateLimiter) and compose them in an agent. The rate limiter should wait 1 second between model calls. The logger should show that the wait happens.

See solution

from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class SimpleLogger(AgentMiddleware):
    """Logging with timestamps."""

    def wrap_model_call(self, messages, config, call_next):
        start = time.time()
        print(f"[LOG] → Model call started")
        response = call_next(messages, config)
        elapsed = time.time() - start
        print(f"[LOG] ← Model call finished ({elapsed:.2f}s)")
        return response

class SimpleRateLimiter(AgentMiddleware):
    """Wait between model calls."""

    def __init__(self, min_interval: float = 1.0):
        self.min_interval = min_interval
        self._last_call = 0.0

    def wrap_model_call(self, messages, config, call_next):
        now = time.time()
        elapsed = now - self._last_call
        if elapsed < self.min_interval and self._last_call > 0:
            wait = self.min_interval - elapsed
            print(f"[RATE] Waiting {wait:.2f}s...")
            time.sleep(wait)
        self._last_call = time.time()
        return call_next(messages, config)

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[SimpleLogger(), SimpleRateLimiter()],
)

result = agent.invoke({"messages": [("user", "Search about Python")]})
print(result["messages"][-1].content)
# Expected output:
# [LOG] → Model call started
# [LOG] ← Model call finished (1.12s)
# [LOG] → Model call started
# [RATE] Waiting 0.88s...
# [LOG] ← Model call finished (2.01s)

Exercise 4: An audit middleware with history (Intermediate)

Build an AuditMiddleware that stores a history of every operation (model and tools) with a timestamp, a type and details. It should have a get_audit_log() method that returns the full list.

See solution

from dotenv import load_dotenv
load_dotenv()

from datetime import datetime
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class AuditMiddleware(AgentMiddleware):
    """Keep an audit log of every operation."""

    def __init__(self):
        self._log: list[dict] = []

    def before_model(self, messages, config):
        self._log.append({
            "timestamp": datetime.now().isoformat(),
            "type": "model_call",
            "details": {"message_count": len(messages)},
        })

    def after_model(self, response, config):
        self._log.append({
            "timestamp": datetime.now().isoformat(),
            "type": "model_response",
            "details": {
                "content_length": len(response.content) if response.content else 0,
                "has_tool_calls": bool(getattr(response, "tool_calls", [])),
            },
        })

    def wrap_tool_call(self, tool_call, config, call_next):
        self._log.append({
            "timestamp": datetime.now().isoformat(),
            "type": "tool_call",
            "details": {
                "tool": tool_call["name"],
                "args": tool_call["args"],
            },
        })
        result = call_next(tool_call, config)
        self._log.append({
            "timestamp": datetime.now().isoformat(),
            "type": "tool_result",
            "details": {
                "tool": tool_call["name"],
                "result_length": len(str(result)),
            },
        })
        return result

    def get_audit_log(self) -> list[dict]:
        return self._log

audit = AuditMiddleware()
model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[audit],
)

result = agent.invoke({"messages": [("user", "Search about Rust")]})

for entry in audit.get_audit_log():
    print(f"  [{entry['type']}] {entry['details']}")
# Expected output:
#   [model_call] {'message_count': 2}
#   [model_response] {'content_length': 0, 'has_tool_calls': True}
#   [tool_call] {'tool': 'search', 'args': {'query': 'Rust'}}
#   [tool_result] {'tool': 'search', 'result_length': 24}
#   [model_call] {'message_count': 4}
#   [model_response] {'content_length': 185, 'has_tool_calls': False}

Exercise 5: Middleware that swaps the model by cost (Advanced)

Build a CostAwareMiddleware that uses gpt-4.1-mini when the input has fewer than 100 words and gpt-4.1 when it has more. It should record which model was used on each call.

See solution

from dotenv import load_dotenv
load_dotenv()

from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"

class CostAwareMiddleware(AgentMiddleware):
    """Pick the model based on the length of the input."""

    def __init__(self):
        self.model_log: list[dict] = []
        self._mini = init_chat_model("openai:gpt-4.1-mini")
        self._full = init_chat_model("openai:gpt-4.1")

    def wrap_model_call(self, messages, config, call_next):
        total_words = sum(
            len(str(m.content).split())
            for m in messages
            if hasattr(m, "content") and m.content
        )

        if total_words < 100:
            model_name = "gpt-4.1-mini"
            selected_model = self._mini
        else:
            model_name = "gpt-4.1"
            selected_model = self._full

        self.model_log.append({
            "model": model_name,
            "word_count": total_words,
            "reason": "short input" if total_words < 100 else "long input",
        })
        print(f"[COST] Using {model_name} ({total_words} words)")

        return selected_model.invoke(messages)

model = init_chat_model("openai:gpt-4.1-mini")
cost_middleware = CostAwareMiddleware()

agent = create_agent(
    model, [search],
    prompt="You are an assistant.",
    middleware=[cost_middleware],
)

result = agent.invoke({"messages": [("user", "Hi")]})
print(result["messages"][-1].content)
print(f"\nModels used: {cost_middleware.model_log}")
# Expected output:
# [COST] Using gpt-4.1-mini (3 words)
# Hi! How can I help you?
# Models used: [{'model': 'gpt-4.1-mini', 'word_count': 3, 'reason': 'short input'}]

Exercise 6: A complete library with 3 composed middleware (Challenge)

Build three middleware (LoggerMW, AuthMW, MetricsMW), compose them in an agent whose state_schema includes user_role, and show that:

The logger records everything
Auth filters tools by role
The metrics accumulate correctly

See solution

from dotenv import load_dotenv
load_dotenv()

from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain.agents import AgentMiddleware, create_agent
from langchain.chat_models import init_chat_model
from langchain_core.tools import tool

@tool
def search(query: str) -> str:
    """Search for general information."""
    return f"Results: {query}"

@tool
def admin_action(action: str) -> str:
    """Run an administrative action (admins only)."""
    return f"Admin: {action} executed"

class LoggerMW(AgentMiddleware):
    def __init__(self):
        self.logs: list[str] = []

    def before_model(self, messages, config):
        entry = f"MODEL_CALL({len(messages)} msgs)"
        self.logs.append(entry)
        print(f"[LOG] {entry}")

    def wrap_tool_call(self, tool_call, config, call_next):
        entry = f"TOOL({tool_call['name']})"
        self.logs.append(entry)
        print(f"[LOG] {entry}")
        return call_next(tool_call, config)

class AuthMW(AgentMiddleware):
    def get_tools(self, state):
        role = state.get("user_role", "basic")
        if role == "admin":
            return [search, admin_action]
        return [search]

class MetricsMW(AgentMiddleware):
    def __init__(self):
        self.calls = {"model": 0, "tools": 0}

    def before_model(self, messages, config):
        self.calls["model"] += 1

    def wrap_tool_call(self, tool_call, config, call_next):
        self.calls["tools"] += 1
        return call_next(tool_call, config)

class AgentState(TypedDict):
    messages: Annotated[list, add_messages]
    user_role: str

logger = LoggerMW()
metrics = MetricsMW()

model = init_chat_model("openai:gpt-4.1-mini")

agent = create_agent(
    model,
    tools=[search, admin_action],
    prompt="You are an assistant. Use the available tools.",
    state_schema=AgentState,
    middleware=[logger, AuthMW(), metrics],
)

result = agent.invoke({
    "messages": [("user", "Search for information about Python")],
    "user_role": "basic",
})

print(f"\n📊 Metrics: {metrics.calls}")
print(f"📝 Logs: {logger.logs}")
# Expected output:
# [LOG] MODEL_CALL(2 msgs)
# [LOG] TOOL(search)
# [LOG] MODEL_CALL(4 msgs)
#
# 📊 Metrics: {'model': 2, 'tools': 1}
# 📝 Logs: ['MODEL_CALL(2 msgs)', 'TOOL(search)', 'MODEL_CALL(4 msgs)']

Summary

In this capsule you learned:

The AgentMiddleware class packages hooks, state and tools into a reusable unit — like a plugin for your agents
Structure: you define methods (before_model, after_model, wrap_model_call, wrap_tool_call, get_tools, get_prompt) inside a class
Composition: you pass middleware=[A(), B(), C()] to the agent — LangChain runs them as layers, with A as the outermost
Order matters: logging/metrics go first (they capture everything), retry/cache go last (right on top of the model)
Internal state: middleware can keep state (self.model_calls, self._cache) that persists across invocations
Parameterizable: as regular Python classes, they take an __init__ with parameters to configure their behavior (log level, rate limits, etc.)
An organization-wide library: you can build a middleware package (my_middleware/) that every team imports and composes to fit their needs
Common patterns: logging, auth/permissions, rate limiting, caching, error recovery, metrics — all implemented as reusable middleware

Next capsule: The project — you'll build an agent with dynamic model routing that combines everything from this module: composed middleware, dynamic models, dynamic tools and logging.

The AgentMiddleware Class: Composed Middleware

Capsule overview

The problem: individual hooks don't scale

The AgentMiddleware class: the solution

The basic structure

Building a middleware class step by step

Step 1: define the middleware

Step 2: use the middleware in an agent

Composing several middleware

Order matters: how middleware stack

Middleware as reusable modules

Parameterizable middleware

Middleware with internal state

Common built-in patterns

Rate limiting middleware

Error recovery middleware

Building a middleware library for your organization

Recommended structure

Example of a shared module

Troubleshooting

1. The middleware hooks never run

2. wrap_model_call and before_model both run

3. Stateful middleware shares state across invocations

4. call_next isn't available in before_model

5. The middleware order doesn't produce the result you expected

Exercises

Exercise 1: A basic timestamp middleware (Basic)

Exercise 2: A token-counting middleware (Basic)

Exercise 3: Compose logging + rate limiting (Intermediate)

Exercise 4: An audit middleware with history (Intermediate)

Exercise 5: Middleware that swaps the model by cost (Advanced)

Exercise 6: A complete library with 3 composed middleware (Challenge)

Summary

Further reading

2. `wrap_model_call` and `before_model` both run

4. `call_next` isn't available in `before_model`