Module 4: Middleware and Customization
@wrap_model_call: Intercepting Model Calls
Capsule overview
@wrap_model_call is the most powerful piece of LangChain's middleware system. Where @before_model and @after_model only observe — they see what goes into the model and what comes out, but can't change anything — @wrap_model_call can modify both the request on the way in and the response on the way out. It wraps the entire model call, giving you full control over what it receives and what it returns.
In the previous capsule you learned to use @before_model and @after_model for logging and monitoring. They were like security cameras: they see everything but never step in. @wrap_model_call is like a security guard: it sees everything, and it can stop, modify, or redirect what passes through. By the end of this capsule you'll know how to intercept model calls to inject system messages, filter long conversations, add metadata to responses, and apply conditional logic based on the content or the agent's state.
The handler pattern
@wrap_model_call receives three parameters:
messages— The list of messages that will be sent to the model (system, human, AI, tool messages).config— The execution configuration (recursion_limit, configurable, metadata).call_next— The function that makes the actual model call. You must always call it.
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
@tool
def get_weather(city: str) -> str:
"""Get the current weather for a city."""
return f"The weather in {city} is sunny, 22°C"
def wrap_model_call(messages, config, call_next):
"""Intercept the model call — sees everything, can change everything."""
print(f"[WRAP] Intercepting {len(messages)} messages")
response = call_next(messages, config)
print(f"[WRAP] Response received: {response.content[:80]}")
return response
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[get_weather],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "What's the weather in Madrid?")]})
print(result["messages"][-1].content)
# Expected output:
# [WRAP] Intercepting 1 messages
# [WRAP] Response received: ...
# [WRAP] Intercepting 3 messages
# [WRAP] Response received: The weather in Madrid is sunny...
# The weather in Madrid is sunny, at 22°C.
The middleware ran twice: once when the model decided to call get_weather, and once when it produced the final answer. That's the ReAct loop in action — every trip through the model node passes through your middleware.
If you don't call call_next, the model never runs and the agent throws an error. call_next is the bridge to the model — without it, there's no response.
The ModelRequest concept
When @wrap_model_call intercepts a call, messages is the complete package headed for the model: the system message (if you configured prompt), the conversation history, and the latest user message or tool result.
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
@tool
def calculator(expression: str) -> str:
"""Evaluate a math expression."""
return str(eval(expression))
call_number = 0
def wrap_model_call(messages, config, call_next):
"""Inspect what the model receives on each call."""
global call_number
call_number += 1
print(f"\n=== Call #{call_number} to the model ===")
for i, msg in enumerate(messages):
msg_type = type(msg).__name__
if hasattr(msg, "tool_calls") and msg.tool_calls:
calls = [(tc["name"], tc["args"]) for tc in msg.tool_calls]
print(f" [{i}] {msg_type}: tool_calls={calls}")
else:
content = msg.content[:60] if msg.content else "(no content)"
print(f" [{i}] {msg_type}: {content}")
return call_next(messages, config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[calculator],
prompt="You are a math assistant. Show your work step by step.",
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "What is 15 * 8?")]})
print(f"\nAnswer: {result['messages'][-1].content}")
# Expected output:
# === Call #1 to the model ===
# [0] SystemMessage: You are a math assistant. Show your work step by step.
# [1] HumanMessage: What is 15 * 8?
#
# === Call #2 to the model ===
# [0] SystemMessage: You are a math assistant. Show your work step by step.
# [1] HumanMessage: What is 15 * 8?
# [2] AIMessage: tool_calls=[('calculator', {'expression': '15 * 8'})]
# [3] ToolMessage: 120
#
# Answer: 15 × 8 = 120.
On call #1, the model gets the system prompt + the question. On call #2, it gets all of that plus the tool call it made and the result.
Modifying the request: injecting system messages
The most common use case — adding a dynamic system message on every call:
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage
from datetime import datetime
@tool
def search(query: str) -> str:
"""Search the internet for information."""
return f"Result: LangChain v1.2 ships create_agent with middleware."
def wrap_model_call(messages, config, call_next):
"""Inject a system message carrying today's date."""
date_msg = SystemMessage(
content=f"Current date: {datetime.now().strftime('%Y-%m-%d %H:%M')}. "
f"Always mention the date when you give time-sensitive information."
)
modified_messages = [date_msg] + list(messages)
return call_next(modified_messages, config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "What's new in LangChain?")]})
print(result["messages"][-1].content)
# Expected output:
# As of 2026-02-28, LangChain v1.2 ships create_agent with middleware support.
Unlike prompt (which you set when you create the agent), this injection happens on every call to the model and can change dynamically.
Modifying the request: filtering sensitive messages
You can redact information before it ever reaches the model:
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage
import re
@tool
def process_order(order_id: str) -> str:
"""Process an order by its ID."""
return f"Order {order_id} processed successfully."
def wrap_model_call(messages, config, call_next):
"""Redact credit card numbers."""
redacted = []
for msg in messages:
if isinstance(msg, HumanMessage):
cleaned = re.sub(
r'\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b',
'[CARD REDACTED]',
msg.content
)
redacted.append(HumanMessage(content=cleaned))
else:
redacted.append(msg)
return call_next(redacted, config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[process_order],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({
"messages": [("user", "Process my order ORD-123. My card is 4532 1234 5678 9012")]
})
print(result["messages"][-1].content)
# Expected output:
# Order ORD-123 has been processed successfully.
The model never saw the card number — that's security at the infrastructure level, not something you're trusting the prompt to handle.
Reading state inside the middleware
If your agent uses a custom state_schema, you can reach the state through config:
from dotenv import load_dotenv
load_dotenv()
from typing import TypedDict, Annotated
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import AnyMessage, SystemMessage
import operator
class SupportState(TypedDict):
messages: Annotated[list[AnyMessage], operator.add]
user_tier: str
@tool
def search_kb(query: str) -> str:
"""Search the knowledge base."""
return f"Article: '{query}' — Restart the device and check the connection."
def wrap_model_call(messages, config, call_next):
"""Adapt the behavior to the user's tier."""
user_tier = config.get("configurable", {}).get("user_tier", "free")
tier_instruction = (
"This is a premium user. Be extra thorough and offer a follow-up."
if user_tier == "premium"
else "Free user. Be concise. Suggest upgrading for advanced support."
)
tier_msg = SystemMessage(content=tier_instruction)
return call_next([tier_msg] + list(messages), config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search_kb],
state_schema=SupportState,
wrap_model_call=wrap_model_call,
)
result = agent.invoke(
{"messages": [("user", "My internet isn't working")], "user_tier": "premium"},
config={"configurable": {"user_tier": "premium"}},
)
print(result["messages"][-1].content)
# Expected output:
# Sorry you're having trouble with your connection. As a premium user,
# here are the detailed steps I'd suggest...
Modifying the model's responses
@wrap_model_call can also change what the model gives back:
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import AIMessage
@tool
def get_price(product: str) -> str:
"""Get the price of a product."""
prices = {"laptop": "999", "mouse": "29", "keyboard": "79"}
return prices.get(product.lower(), "Product not found")
def wrap_model_call(messages, config, call_next):
"""Append a disclaimer to the final answer."""
response = call_next(messages, config)
if response.content and not response.tool_calls:
response = AIMessage(
content=response.content + "\n\n_AI-generated answer. Verify before acting on it._",
tool_calls=response.tool_calls,
response_metadata=response.response_metadata,
)
return response
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[get_price],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "How much is a mouse?")]})
print(result["messages"][-1].content)
# Expected output:
# The mouse costs $29.
#
# _AI-generated answer. Verify before acting on it._
It only appends the disclaimer to final answers (no tool_calls), not to the intermediate ones where the model is deciding which tool to call.
Conditional logic: dynamic behavior
Apply different logic depending on what's in the messages:
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage, HumanMessage
@tool
def search(query: str) -> str:
"""Search the internet for information."""
return f"Result: {query} — relevant information found."
def wrap_model_call(messages, config, call_next):
"""Apply instructions based on the type of question."""
last_human = None
for msg in reversed(messages):
if isinstance(msg, HumanMessage):
last_human = msg.content.lower()
break
if last_human and any(w in last_human for w in ["code", "program", "script"]):
extra = SystemMessage(
content="The user is asking for code. Give complete, working code."
)
messages = [extra] + list(messages)
elif last_human and any(w in last_human for w in ["explain", "what is", "how does"]):
extra = SystemMessage(
content="The user is asking for an explanation. Use simple analogies."
)
messages = [extra] + list(messages)
return call_next(messages, config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "Explain what a decorator is in Python")]})
print(result["messages"][-1].content[:200])
# Expected output:
# A decorator in Python is like gift wrapping. Imagine you have a function that
# does something, and you want to add extra behavior without changing it...
Comparison: @before/@after vs @wrap_model_call
| Aspect | @before_model / @after_model | @wrap_model_call |
|---|---|---|
| Can see the messages | ✅ | ✅ |
| Can see the response | ✅ (@after_model only) | ✅ |
| Can modify messages | ❌ | ✅ |
| Can modify the response | ❌ | ✅ |
| Can block the call | ❌ | ✅ (just don't call call_next) |
| Can retry the model | ❌ | ✅ |
| Main use case | Logging, monitoring | Transformation, control |
Rule of thumb: use @before_model / @after_model when you only need to observe. Use @wrap_model_call when you need to intervene.
Troubleshooting
Problem 1: "AttributeError: 'NoneType' object has no attribute..."
Symptom: Error when running the agent with wrap_model_call.
Cause: Your function doesn't return anything (you forgot return call_next(...) or return response).
Fix: Make sure you always return the result:
def wrap_model_call(messages, config, call_next):
response = call_next(messages, config)
return response # ← never forget the return
Problem 2: The middleware runs more times than expected
Symptom: The middleware's prints show up 3, 4, or more times. Cause: It runs on every trip through the model node in the ReAct loop. If the agent goes through 3 rounds of tools, the middleware runs 4 times. Fix: That's the expected behavior. To act only on the final answer:
def wrap_model_call(messages, config, call_next):
response = call_next(messages, config)
if not response.tool_calls:
print("[WRAP] This is the final answer")
return response
Problem 3: Modifying the response loses the tool_calls
Symptom: The agent breaks after you modify the AIMessage.
Cause: You didn't carry over the original's tool_calls when you built the new AIMessage.
Fix: Always preserve the tool_calls:
from langchain_core.messages import AIMessage
def wrap_model_call(messages, config, call_next):
response = call_next(messages, config)
return AIMessage(
content=response.content + " [modified]",
tool_calls=response.tool_calls,
response_metadata=response.response_metadata,
)
Exercises
Exercise 1: Basic logging middleware (Easy)
Write a @wrap_model_call that prints how many messages the model receives and whether the response carries tool calls or text. Try it on an agent with a calculator tool.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
@tool
def calculator(expression: str) -> str:
"""Evaluate a math expression."""
return str(eval(expression))
def wrap_model_call(messages, config, call_next):
print(f"[LOG] Input: {len(messages)} messages")
response = call_next(messages, config)
has_tools = bool(response.tool_calls)
print(f"[LOG] Output: {'tool_calls' if has_tools else 'text'}")
return response
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[calculator],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "What is 42 * 58?")]})
print(f"\nAnswer: {result['messages'][-1].content}")
# Expected output:
# [LOG] Input: 1 messages
# [LOG] Output: tool_calls
# [LOG] Input: 3 messages
# [LOG] Output: text
#
# Answer: 42 × 58 = 2,436.
Explanation: On the first call, the model decides to call calculator (tool_calls). On the second, it produces the final answer (text).
Exercise 2: Forcing a language (Easy)
Write middleware that forces the model to always answer in English, no matter what language the question is in. Inject a system message carrying that instruction.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import SystemMessage
@tool
def search(query: str) -> str:
"""Search the internet for information."""
return f"Resultado para '{query}': Python fue creado por Guido van Rossum en 1991."
def wrap_model_call(messages, config, call_next):
language_msg = SystemMessage(
content="ABSOLUTE RULE: Always respond in English, "
"regardless of the language of the question or of the tool results."
)
return call_next([language_msg] + list(messages), config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "¿Quién creó Python?")]})
print(result["messages"][-1].content)
# Expected output:
# Python was created by Guido van Rossum in 1991.
Explanation: Even though the question is in Spanish and the tool returns Spanish, the injected system message forces the answer into English.
Exercise 3: Cumulative cost tracker (Medium)
Write middleware that tracks tokens and estimated cost. At the end, print a summary with total calls, tokens, and cost in USD.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
@tool
def search(query: str) -> str:
"""Search the internet for information."""
return f"Result: {query} is a broad topic."
costs = {"calls": 0, "input_tokens": 0, "output_tokens": 0}
def wrap_model_call(messages, config, call_next):
costs["calls"] += 1
response = call_next(messages, config)
usage = response.response_metadata.get("token_usage", {})
costs["input_tokens"] += usage.get("prompt_tokens", 0)
costs["output_tokens"] += usage.get("completion_tokens", 0)
return response
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "What is Python?")]})
print(f"Answer: {result['messages'][-1].content}")
total = (costs["input_tokens"] / 1000) * 0.00015 + (costs["output_tokens"] / 1000) * 0.0006
print(f"\nCalls: {costs['calls']}")
print(f"Tokens: {costs['input_tokens']} in / {costs['output_tokens']} out")
print(f"Cost: ${total:.6f} USD")
# Expected output:
# Answer: Python is a programming language...
#
# Calls: 2
# Tokens: ~170 in / ~50 out
# Cost: $0.000056 USD
Explanation: The costs dictionary accumulates data across calls. In production you'd swap the dict for a database.
Exercise 4: Content filter with a whitelist (Advanced)
Write middleware that only allows certain topics. If the user asks about anything outside the whitelist, return a rejection message without calling the model (so you burn no tokens).
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.messages import AIMessage, HumanMessage
@tool
def search(query: str) -> str:
"""Search the internet for information."""
return f"Result: {query} — information found."
ALLOWED = ["python", "langchain", "programming", "artificial intelligence"]
def wrap_model_call(messages, config, call_next):
last_human = None
for msg in reversed(messages):
if isinstance(msg, HumanMessage):
last_human = msg.content.lower()
break
if last_human and not any(topic in last_human for topic in ALLOWED):
return AIMessage(
content="I can only help with programming, "
"Python, LangChain, and AI. Any questions on those topics?"
)
return call_next(messages, config)
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
wrap_model_call=wrap_model_call,
)
for q in ["What is Python?", "Best pasta recipe?", "How does LangChain work?"]:
result = agent.invoke({"messages": [("user", q)]})
print(f"Q: {q}\nA: {result['messages'][-1].content[:100]}\n")
# Expected output:
# Q: What is Python?
# A: Python is a programming language...
#
# Q: Best pasta recipe?
# A: I can only help with programming, Python, LangChain, and AI...
#
# Q: How does LangChain work?
# A: LangChain is a framework for building applications...
Explanation: For the pasta question, the middleware returns a rejection AIMessage without ever calling call_next — the model never runs and not a single token gets spent.
Exercise 5: Composed middleware — logging + cost tracking (Challenge)
Combine input/output logging and cost tracking in a single @wrap_model_call. Print a report at the end.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
@tool
def search(query: str) -> str:
"""Search the internet for information."""
return f"Result for '{query}': public information."
@tool
def calculator(expression: str) -> str:
"""Evaluate a math expression."""
return str(eval(expression))
metrics = {"calls": 0, "input_tokens": 0, "output_tokens": 0}
def wrap_model_call(messages, config, call_next):
metrics["calls"] += 1
call_num = metrics["calls"]
print(f"[#{call_num}] INPUT: {len(messages)} messages")
response = call_next(messages, config)
usage = response.response_metadata.get("token_usage", {})
metrics["input_tokens"] += usage.get("prompt_tokens", 0)
metrics["output_tokens"] += usage.get("completion_tokens", 0)
output_type = "tool_calls" if response.tool_calls else "text"
print(f"[#{call_num}] OUTPUT: {output_type}")
return response
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search, calculator],
wrap_model_call=wrap_model_call,
)
result = agent.invoke({"messages": [("user", "What is Python, and what is 100 * 3.14?")]})
print(f"\nAnswer: {result['messages'][-1].content}")
total_cost = (
(metrics["input_tokens"] / 1000) * 0.00015
+ (metrics["output_tokens"] / 1000) * 0.0006
)
print(f"\n=== REPORT ===")
print(f"Calls: {metrics['calls']}")
print(f"Tokens: {metrics['input_tokens']} in / {metrics['output_tokens']} out")
print(f"Cost: ${total_cost:.6f} USD")
# Expected output:
# [#1] INPUT: 1 messages
# [#1] OUTPUT: tool_calls
# [#2] INPUT: 4 messages
# [#2] OUTPUT: text
#
# Answer: Python is a language... 100 × 3.14 = 314.
#
# === REPORT ===
# Calls: 2
# Tokens: ~200 in / ~60 out
# Cost: $0.000066 USD
Explanation: A single middleware handles both logging and cost tracking. In production, you'd probably split those responsibilities into separate middleware using the AgentMiddleware class (Capsule 07).
Summary
In this capsule you learned:
@wrap_model_callreceives (messages, config, call_next) and wraps the entire model call- Unlike
@before_model/@after_model, it can modify both the request and the response - You must always call
call_nextfor the model to run (unless you're deliberately blocking it) - The middleware runs on every step of the ReAct loop — once per trip through the model node
- You can inject dynamic system messages that change with the context
- You can filter messages before they go out (redacting sensitive information)
- You can modify responses after they come back (disclaimers, metadata)
- Conditional logic: analyze the content to apply different behavior
- Key use cases: cost tracking, content filtering, request transformation, response enrichment
- Execution order with the other middleware:
@before_model→@wrap_model_call→@after_model
Next capsule: @wrap_tool_call — you'll learn to intercept tool execution to add retry logic, custom error handling, and timing.
Further reading
- create_agent API Reference — Full documentation of the middleware parameters
- LangGraph Agents — Middleware — The official guide to middleware in agents
- Custom Agent Middleware — How-to for custom middleware
- Token Usage Tracking — Tracking token usage through metadata
- Content Moderation with LangChain — Content moderation techniques
Module 4 — LangChain & LangGraph: From Chains to Agents