Module 4: Middleware and Customization
Your First Middleware: Logging and Monitoring
Capsule overview
Before you intercept and modify an agent's behavior, you need to observe what it does. How many times does it call the model? What messages does it send? How long does each call take? How many tokens does it burn? Without visibility, you're working blind.
In this capsule you'll learn the two simplest hooks in the middleware system: before_model and after_model. The first runs right before each LLM call; the second, right after. Together they give you full visibility into the agent-model interaction. You'll build basic logging middleware, timing, token counting, and you'll see how it compares to LangChain's callbacks system.
By the end, you'll be able to add observability to any existing agent without changing a single line of its code.
before_model: intercepting before the call
before_model is a function that runs every time the agent is about to send messages to the model. It receives the list of messages that's about to go out.
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 web for information."""
return f"Results for '{query}': LangChain is a framework for LLMs."
def before_model(messages):
print(f"[BEFORE] Sending {len(messages)} messages to the model")
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
before_model=before_model
)
result = agent.invoke({"messages": [("user", "What is LangChain?")]})
print(result["messages"][-1].content)
# Expected output:
# [BEFORE] Sending 2 messages to the model
# [BEFORE] Sending 4 messages to the model
# LangChain is a framework for building applications with LLMs.
Why does [BEFORE] show up twice? Because the agent went through two rounds of the ReAct loop:
- First round: The agent sends 2 messages (system prompt + human message). The model decides to call
search. - Second round: The agent sends 4 messages (system prompt + human message + AI message with the tool_call + ToolMessage with the result). The model produces the final answer.
Every time the agent is about to call the model, before_model runs first.
Inspecting the messages
You can also iterate over the messages to see exactly what the model receives:
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 before_model(messages):
print(f"\n[BEFORE] {len(messages)} messages for the model:")
for msg in messages:
msg_type = type(msg).__name__
if hasattr(msg, "tool_calls") and msg.tool_calls:
print(f" {msg_type}: tool_calls={[tc['name'] for tc in msg.tool_calls]}")
else:
print(f" {msg_type}: {str(msg.content)[:80]}")
agent = create_agent("openai:gpt-4.1-mini", tools=[calculator], before_model=before_model)
result = agent.invoke({"messages": [("user", "What is 15 * 37?")]})
print(f"\nAnswer: {result['messages'][-1].content}")
# Expected output:
# [BEFORE] 2 messages for the model:
# SystemMessage: You are a helpful assistant.
# HumanMessage: What is 15 * 37?
#
# [BEFORE] 4 messages for the model:
# SystemMessage: You are a helpful assistant.
# HumanMessage: What is 15 * 37?
# AIMessage: tool_calls=['calculator']
# ToolMessage: 555
#
# Answer: 15 × 37 = 555.
This is invaluable for debugging: if the agent behaves unexpectedly, before_model shows you exactly what information it had to work with.
after_model: intercepting after the response
after_model runs every time the model returns a response, before the agent processes it (running tools or handing back the final result).
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 web for information."""
return f"Results for '{query}': Python was created by Guido van Rossum in 1991."
def after_model(response):
print(f"[AFTER] Response type: {type(response).__name__}")
if response.tool_calls:
for tc in response.tool_calls:
print(f" → Tool call: {tc['name']}({tc['args']})")
else:
print(f" → Final answer: {str(response.content)[:100]}")
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search],
after_model=after_model
)
result = agent.invoke({"messages": [("user", "Who created Python?")]})
# Expected output:
# [AFTER] Response type: AIMessage
# → Tool call: search({'query': 'who created Python'})
# [AFTER] Response type: AIMessage
# → Final answer: Python was created by Guido van Rossum in 1991.
after_model lets you tell two situations apart:
- The model decided to call a tool →
response.tool_callshas content - The model decided to answer directly →
response.contenthas the final answer
Combining before_model and after_model
Together, before_model and after_model give you full visibility into the request-response cycle between the agent and the model:
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 web for information."""
return f"Results for '{query}': LangChain was created by Harrison Chase in 2022."
@tool
def calculator(expression: str) -> str:
"""Evaluate a math expression."""
return str(eval(expression))
call_number = 0
def before_model(messages):
global call_number
call_number += 1
last_msg = messages[-1]
print(f"\n--- Call #{call_number}: {len(messages)} msgs, "
f"last={type(last_msg).__name__}: {str(last_msg.content)[:60]} ---")
def after_model(response):
if response.tool_calls:
print(f" → tool_calls: {[tc['name'] for tc in response.tool_calls]}")
else:
print(f" → final answer ({len(response.content)} chars)")
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[search, calculator],
before_model=before_model,
after_model=after_model
)
result = agent.invoke({
"messages": [("user", "Who created LangChain, and what is 2**16?")]
})
print(f"\nAnswer: {result['messages'][-1].content}")
# Expected output:
# --- Call #1: 2 msgs, last=HumanMessage: Who created LangChain... ---
# → tool_calls: ['search', 'calculator']
# --- Call #2: 5 msgs, last=ToolMessage: 65536 ---
# → final answer (78 chars)
# Answer: LangChain was created by Harrison Chase in 2022. And 2^16 = 65,536.
Now you can trace every interaction: how many calls the agent made, what it got each time, and what it decided to do.
Timing middleware: measuring latency
One of the most useful metrics in production is how long each model call takes. With before_model and after_model, you can measure it:
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
import time
@tool
def search(query: str) -> str:
"""Search the web for information."""
return f"Results for '{query}': FastAPI is a modern web framework for Python."
call_times = {}
def before_model(messages):
call_times["start"] = time.time()
def after_model(response):
elapsed = time.time() - call_times["start"]
action = "tool_calls" if response.tool_calls else "final answer"
print(f"[TIMER] {elapsed:.2f}s → {action}")
agent = create_agent(
"openai:gpt-4.1-mini", tools=[search],
before_model=before_model, after_model=after_model
)
result = agent.invoke({"messages": [("user", "What is FastAPI?")]})
print(f"Answer: {result['messages'][-1].content}")
# Expected output:
# [TIMER] 0.83s → tool_calls
# [TIMER] 0.65s → final answer
# Answer: FastAPI is a modern web framework for Python...
This pattern is the foundation of production monitoring. The metrics you're collecting here are the same ones platforms like LangSmith or LangFuse track automatically. In the exercises at the end you'll see a version with an accumulator that reports totals and averages.
Token counting middleware
If you need to keep costs under control, after_model can read usage_metadata from the response to accumulate tokens:
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 web for information."""
return f"Results for '{query}': LangGraph lets you build workflows as graphs."
token_usage = {"total_input": 0, "total_output": 0, "calls": 0}
def after_model(response):
token_usage["calls"] += 1
if hasattr(response, "usage_metadata") and response.usage_metadata:
usage = response.usage_metadata
token_usage["total_input"] += usage.get("input_tokens", 0)
token_usage["total_output"] += usage.get("output_tokens", 0)
print(f"[TOKENS] Call {token_usage['calls']}: "
f"input={usage.get('input_tokens', 0)}, output={usage.get('output_tokens', 0)}")
agent = create_agent("openai:gpt-4.1-mini", tools=[search], after_model=after_model)
result = agent.invoke({"messages": [("user", "What is LangGraph?")]})
total = token_usage["total_input"] + token_usage["total_output"]
print(f"\nTotal: {total} tokens across {token_usage['calls']} calls")
print(f"Answer: {result['messages'][-1].content}")
# Expected output:
# [TOKENS] Call 1: input=85, output=22
# [TOKENS] Call 2: input=130, output=45
# Total: 282 tokens across 2 calls
⚠️ Note: usage_metadata depends on the provider. OpenAI always includes it. Anthropic includes it if you enable tracking. Other providers may not support it at all.
The callbacks system: an alternative approach
LangChain has another mechanism for observing what happens internally: callbacks. Callbacks predate the middleware system and are still useful for certain cases.
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.callbacks import BaseCallbackHandler
@tool
def search(query: str) -> str:
"""Search the web for information."""
return f"Results for '{query}': Python is the most popular language of 2025."
class LoggingCallback(BaseCallbackHandler):
def on_llm_start(self, serialized, prompts, **kwargs):
print("[CALLBACK] LLM starting...")
def on_llm_end(self, response, **kwargs):
print("[CALLBACK] LLM finished")
def on_tool_start(self, serialized, input_str, **kwargs):
print(f"[CALLBACK] Tool: {serialized.get('name', '?')}")
def on_tool_end(self, output, **kwargs):
print(f"[CALLBACK] Tool done")
agent = create_agent("openai:gpt-4.1-mini", tools=[search])
result = agent.invoke(
{"messages": [("user", "What is the most popular language?")]},
config={"callbacks": [LoggingCallback()]}
)
print(f"Answer: {result['messages'][-1].content}")
# Output: [CALLBACK] LLM starting... → finished → Tool: search → done → LLM starting... → finished
Middleware vs callbacks: when do you use each?
| Aspect | Middleware (before_model/after_model) | Callbacks (BaseCallbackHandler) |
|---|---|---|
| Scope | Specific to the agent | Any LangChain component |
| Data access | Full messages and a typed response | Serialized, generic data |
| Modification | Can modify messages and responses | Observation only (read-only) |
| Configuration | A create_agent parameter | Passed in config at invoke time |
| Best for | Active middleware (routing, retry, filtering) | Passive logging, tracing, platform integrations |
Rule of thumb: Use middleware when you need to modify behavior. Use callbacks when you only need to observe from the outside. For simple logging, either works — middleware is the more direct route.
How the hooks get passed to create_agent
Middleware hooks go in as keyword arguments to create_agent. You can use one, the other, or both — they're optional and independent:
# before_model only
agent = create_agent("openai:gpt-4.1-mini", tools=tools, before_model=my_before)
# after_model only
agent = create_agent("openai:gpt-4.1-mini", tools=tools, after_model=my_after)
# Both
agent = create_agent("openai:gpt-4.1-mini", tools=tools, before_model=my_before, after_model=my_after)
Execution order
When you use both hooks, the sequence for each model call is:
1. before_model(messages) ← you intercept before
2. model.invoke(messages) ← the model does its thing
3. after_model(response) ← you intercept after
This cycle repeats on every iteration of the ReAct loop. If the agent makes 3 model calls, before_model and after_model each run 3 times.
Troubleshooting
Problem 1: "TypeError: before_model() takes 0 positional arguments but 1 was given"
Symptom: Error when invoking the agent.
Cause: Your before_model function doesn't accept the messages parameter.
Fix: before_model always receives one argument — the message list:
# Wrong — takes no arguments
def before_model():
print("Before the model")
# Right — takes messages
def before_model(messages):
print(f"Before the model: {len(messages)} messages")
Problem 2: "TypeError: after_model() takes 0 positional arguments but 1 was given"
Symptom: Error when invoking the agent.
Cause: Your after_model function doesn't accept the response parameter.
Fix: after_model always receives one argument — the model's response:
# Wrong
def after_model():
print("After the model")
# Right
def after_model(response):
print(f"After the model: {len(response.content)} chars")
Problem 3: The middleware prints nothing
Symptom: The agent works, but the middleware's print() calls never show up.
Cause: The middleware isn't wired into the agent. You probably passed the hooks under the wrong name, or didn't pass them at all.
Fix: Check that you're using the exact names as keyword arguments:
# Wrong — wrong names
agent = create_agent(model, tools, pre_model=fn, post_model=fn)
# Right — correct names
agent = create_agent(model, tools, before_model=fn, after_model=fn)
Problem 4: The timer shows 0.00s
Symptom: Your time measurements always come out as 0.00 seconds.
Cause: You're measuring the wrong thing — probably creating a fresh timestamp inside after_model instead of using the one from before_model.
Fix: Store the timestamp in a variable shared by both hooks:
import time
state = {"start": None}
def before_model(messages):
state["start"] = time.time()
def after_model(response):
elapsed = time.time() - state["start"]
print(f"Duration: {elapsed:.2f}s")
Problem 5: usage_metadata is None
Symptom: response.usage_metadata returns None, or doesn't exist.
Cause: The provider doesn't include token metadata, or the attribute doesn't exist in your version.
Fix: Guard with hasattr(response, "usage_metadata") and response.usage_metadata before you touch it. OpenAI always includes it; other providers may not.
Exercises
Exercise 1: Basic logging with before_model (Easy)
Create an agent with a get_weather tool and a before_model that prints how many messages go to the model on each call. Invoke the agent with "What's the weather in Tokyo?" and check that the log shows up before each model response.
See solution
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 rainy, 15°C, 80% humidity"
def before_model(messages):
print(f"[LOG] Sending {len(messages)} messages to the model")
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[get_weather],
before_model=before_model
)
result = agent.invoke({"messages": [("user", "What's the weather in Tokyo?")]})
print(f"Answer: {result['messages'][-1].content}")
# Expected output:
# [LOG] Sending 2 messages to the model
# [LOG] Sending 4 messages to the model
# Answer: The weather in Tokyo is rainy, 15°C with 80% humidity.
Explanation: before_model runs twice: once before the call that produces the tool call, and once before the call that produces the final answer using the tool's result.
Exercise 2: Detecting tool calls with after_model (Easy)
Write an after_model that prints whether the model decided to call tools or answer directly. Try it with two questions: one that needs a tool and one that doesn't.
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 after_model(response):
if response.tool_calls:
tools = [tc["name"] for tc in response.tool_calls]
print(f"[DECISION] Model called tools: {tools}")
else:
print(f"[DECISION] Model answered directly ({len(response.content)} chars)")
agent = create_agent(
"openai:gpt-4.1-mini",
tools=[calculator],
after_model=after_model
)
print("--- Question that needs a tool ---")
r1 = agent.invoke({"messages": [("user", "What is 99 * 77?")]})
print(f"A: {r1['messages'][-1].content}\n")
print("--- Question that does NOT need a tool ---")
r2 = agent.invoke({"messages": [("user", "What is a variable?")]})
print(f"A: {r2['messages'][-1].content}")
# Expected output:
# --- Question that needs a tool ---
# [DECISION] Model called tools: ['calculator']
# [DECISION] Model answered directly (18 chars)
# A: 99 × 77 = 7,623.
#
# --- Question that does NOT need a tool ---
# [DECISION] Model answered directly (85 chars)
# A: A variable is a named space in memory that stores a value...
Explanation: With a tool, after_model runs twice: first it detects the tool call, then it detects the final answer. Without a tool, it runs just once with the direct answer.
Exercise 3: Complete timing middleware (Medium)
Build middleware that times every model call and, at the end, prints a summary with the total time and the average per call.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
import time
@tool
def search(query: str) -> str:
"""Search the web for information."""
return f"Results for '{query}': Docker is a container platform."
timing = {"start": None, "durations": []}
def before_model(messages):
timing["start"] = time.time()
def after_model(response):
elapsed = time.time() - timing["start"]
timing["durations"].append(elapsed)
action = "tools" if response.tool_calls else "answer"
print(f"[TIMER] Call {len(timing['durations'])}: {elapsed:.3f}s ({action})")
agent = create_agent(
"openai:gpt-4.1-mini", tools=[search],
before_model=before_model, after_model=after_model
)
result = agent.invoke({"messages": [("user", "What is Docker?")]})
total = sum(timing["durations"])
avg = total / len(timing["durations"])
print(f"\nTotal: {total:.3f}s | Average: {avg:.3f}s | Calls: {len(timing['durations'])}")
print(f"Answer: {result['messages'][-1].content}")
# Expected output:
# [TIMER] Call 1: 0.872s (tools)
# [TIMER] Call 2: 0.641s (answer)
# Total: 1.513s | Average: 0.757s | Calls: 2
Explanation: before_model stores the start timestamp. after_model computes the delta and accumulates it. At the end, we work out the totals and averages.
Exercise 4: Accumulating token counts (Medium)
Write an after_model that accumulates input and output tokens across the whole conversation. Print the running total after each call and the grand total 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 web for information."""
return f"Results for '{query}': Kubernetes orchestrates containers across clusters."
tokens = {"input": 0, "output": 0, "calls": 0}
def after_model(response):
tokens["calls"] += 1
if hasattr(response, "usage_metadata") and response.usage_metadata:
inp = response.usage_metadata.get("input_tokens", 0)
out = response.usage_metadata.get("output_tokens", 0)
tokens["input"] += inp
tokens["output"] += out
print(f"[TOKENS] Call {tokens['calls']}: +{inp} in, +{out} out "
f"(total: {tokens['input'] + tokens['output']})")
agent = create_agent("openai:gpt-4.1-mini", tools=[search], after_model=after_model)
result = agent.invoke({"messages": [("user", "What is Kubernetes?")]})
print(f"\nTotal: {tokens['input'] + tokens['output']} tokens across {tokens['calls']} calls")
print(f"Answer: {result['messages'][-1].content}")
# Expected output:
# [TOKENS] Call 1: +85 in, +20 out (total: 105)
# [TOKENS] Call 2: +128 in, +52 out (total: 285)
# Total: 285 tokens across 2 calls
Explanation: The input grows on the second call because it now includes every previous message (system prompt + user + AI tool_call + ToolMessage).
Exercise 5: Logging with callbacks (Medium)
Implement logging using the callbacks system with BaseCallbackHandler. Create a handler that logs the start and end of every LLM call and every tool execution. Pass the handler through config={"callbacks": [handler]}.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
from langchain_core.callbacks import BaseCallbackHandler
@tool
def search(query: str) -> str:
"""Search the web for information."""
return f"Results for '{query}': Redis is an in-memory database."
class DetailedLogger(BaseCallbackHandler):
def __init__(self):
self.llm_calls = 0
self.tool_calls = 0
def on_llm_start(self, serialized, prompts, **kwargs):
self.llm_calls += 1
print(f"[CB] LLM call #{self.llm_calls} started")
def on_llm_end(self, response, **kwargs):
print(f"[CB] LLM call #{self.llm_calls} finished")
def on_tool_start(self, serialized, input_str, **kwargs):
self.tool_calls += 1
print(f"[CB] Tool '{serialized.get('name', '?')}' running...")
def on_tool_end(self, output, **kwargs):
print(f"[CB] Tool done: {str(output)[:60]}")
logger = DetailedLogger()
agent = create_agent("openai:gpt-4.1-mini", tools=[search])
result = agent.invoke(
{"messages": [("user", "What is Redis?")]},
config={"callbacks": [logger]}
)
print(f"\nSummary: {logger.llm_calls} LLM calls, {logger.tool_calls} tool calls")
print(f"Answer: {result['messages'][-1].content}")
# Expected output:
# [CB] LLM call #1 started → finished → Tool 'search' → done
# [CB] LLM call #2 started → finished
# Summary: 2 LLM calls, 1 tool calls
Explanation: Callbacks go in config, not as create_agent parameters. They capture events from the whole pipeline (LLM + tools).
Exercise 6: A factory for configurable middleware (Advanced)
Write a function create_monitor(name, verbose) that returns a (before_fn, after_fn) pair. With verbose=True, it prints the last 2 messages and the tool calls. With verbose=False, it prints only the call number and the duration. Create two agents with different levels of detail.
See solution
from dotenv import load_dotenv
load_dotenv()
from langchain.agents import create_agent
from langchain_core.tools import tool
import time
@tool
def search(query: str) -> str:
"""Search the web for information."""
return f"Results for '{query}': GraphQL is a query language for APIs."
def create_monitor(name, verbose=False):
state = {"start": None, "num": 0}
def before_fn(messages):
state["num"] += 1
state["start"] = time.time()
if verbose:
print(f"[{name}] Call #{state['num']}: {len(messages)} messages")
for msg in messages[-2:]:
print(f" └─ {type(msg).__name__}: {str(msg.content)[:60]}")
else:
print(f"[{name}] #{state['num']}...", end=" ")
def after_fn(response):
elapsed = time.time() - state["start"]
action = "tools" if response.tool_calls else "answer"
if verbose:
print(f"[{name}] → {action} ({elapsed:.2f}s)")
else:
print(f"{'T' if response.tool_calls else 'A'} {elapsed:.2f}s")
return before_fn, after_fn
before_v, after_v = create_monitor("VERBOSE", verbose=True)
agent_v = create_agent("openai:gpt-4.1-mini", tools=[search], before_model=before_v, after_model=after_v)
before_q, after_q = create_monitor("QUIET", verbose=False)
agent_q = create_agent("openai:gpt-4.1-mini", tools=[search], before_model=before_q, after_model=after_q)
print("=== Verbose ===")
r1 = agent_v.invoke({"messages": [("user", "What is GraphQL?")]})
print(f"A: {r1['messages'][-1].content}\n")
print("=== Quiet ===")
r2 = agent_q.invoke({"messages": [("user", "What is GraphQL?")]})
print(f"A: {r2['messages'][-1].content}")
# Expected output: Verbose shows detailed messages; Quiet shows only "#1... T 0.80s"
Explanation: The factory function wraps state and logic in a closure. Each instance gets its own counter and configuration, which lets you reuse the pattern across multiple agents.
Summary
In this capsule you learned:
before_model(messages)runs before every LLM call — it receives the full list of messages the model is about to processafter_model(response)runs after every LLM response — it receives theAIMessagewith the answer or the tool calls- Together they give you full visibility into the request-response cycle between the agent and the model
- They're passed as keyword arguments to
create_agent:before_model=fn,after_model=fn— optional and independent - Timing middleware calls
time.time()inbefore_modeland computes the delta inafter_model, keeping the timestamp in a shared variable - Token counting middleware reads
response.usage_metadatainafter_modelto accumulate input and output tokens (provider-dependent) - The callbacks system (
BaseCallbackHandler) is an alternative for passive logging — it captures events from the whole pipeline, not just the model - The key difference: middleware can modify behavior; callbacks only observe
- Factory functions (
create_monitor(name, verbose)) let you build configurable, reusable middleware
Next capsule: @wrap_model_call: Intercepting Model Calls — you'll learn the most powerful hook in the middleware system, the one that lets you modify requests before they reach the model and transform the responses on the way back.
Further reading
- create_agent API Reference — Full docs for
before_modelandafter_modelas parameters - LangChain Agents Overview — Conceptual guide to agents and middleware
- LangChain Callbacks Documentation — The callbacks system for logging and tracing
- BaseCallbackHandler API — Reference for every event you can capture with callbacks
- Token Usage Tracking — LangChain — How to track token usage per provider
- LangSmith Overview — The observability platform that automates the logging you just built by hand
- Python Closures — Real Python — Reference for the factory pattern used in reusable middleware
- time.time() vs time.perf_counter() — For more precise timing, consider
perf_counter
Module 4 — LangChain & LangGraph: From Chains to Agents