Module 4: Middleware and Customization

@wrap_tool_call: Customizing Tool Execution

Capsule overview

Just as @wrap_model_call wraps model calls, @wrap_tool_call wraps tool execution. Every time the agent decides to call a tool, your middleware intercepts that execution before it happens and after it finishes. That lets you add retry logic, custom error handling, detailed logging, timing, and any logic you need around your tools — without changing the tools' own code.

In the previous capsule you learned to intercept calls to the model with @wrap_model_call. Now you complete the other side: the tools. With both middleware in place, you have full visibility and control over every operation the agent runs. By the end of this capsule you'll know how to build agents that handle tool errors intelligently, measure each tool's performance, and keep a detailed log of every execution.


The handler pattern

@wrap_tool_call receives three parameters:

  1. tool_call — A dictionary with name (the tool's name), args (the arguments), and id (a unique identifier for the call).
  2. config — The execution configuration (same as in @wrap_model_call).
  3. call_next — The function that runs the actual tool. You must always call it (unless you're deliberately blocking the execution).
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"

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

def wrap_tool_call(tool_call, config, call_next):
    """Intercept every tool execution."""
    print(f"[TOOL] Running: {tool_call['name']}({tool_call['args']})")

    result = call_next(tool_call, config)

    print(f"[TOOL] Result: {result}")
    return result

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[get_weather, calculator],
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({
    "messages": [("user", "What's the weather in Madrid, and what is 15 * 8?")]
})
print(f"\nAnswer: {result['messages'][-1].content}")

# Expected output:
# [TOOL] Running: get_weather({'city': 'Madrid'})
# [TOOL] Result: The weather in Madrid is sunny, 22°C
# [TOOL] Running: calculator({'expression': '15 * 8'})
# [TOOL] Result: 120
#
# Answer: The weather in Madrid is sunny, 22°C. And 15 × 8 = 120.

The middleware ran once per tool call. When the model asks for two tools in parallel, your middleware intercepts each of them individually.

The id inside tool_call is used internally to pair each ToolMessage with its corresponding tool call — you don't need to touch it.


Retry logic: retrying tools that fail

Tools fail — APIs go down, requests time out, networks drop. @wrap_tool_call lets you add automatic retries:

from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import create_agent
from langchain_core.tools import tool

attempt_counter = 0

@tool
def flaky_api(query: str) -> str:
    """Query an external API that sometimes fails."""
    global attempt_counter
    attempt_counter += 1
    if attempt_counter < 3:
        raise ConnectionError(f"API timeout (attempt {attempt_counter})")
    return f"Data for '{query}': temperature 22°C, humidity 65%"

def wrap_tool_call(tool_call, config, call_next):
    """Retry tools with exponential backoff."""
    tool_name = tool_call["name"]
    max_retries = 3

    for attempt in range(max_retries):
        try:
            result = call_next(tool_call, config)
            if attempt > 0:
                print(f"[RETRY] {tool_name} succeeded on attempt {attempt + 1}")
            return result
        except Exception as e:
            wait_time = 2 ** attempt
            if attempt < max_retries - 1:
                print(f"[RETRY] {tool_name} failed ({attempt + 1}/{max_retries}): {e}")
                time.sleep(wait_time)
            else:
                print(f"[RETRY] {tool_name} failed after {max_retries} attempts")
                raise

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[flaky_api],
    wrap_tool_call=wrap_tool_call,
)

attempt_counter = 0
result = agent.invoke({"messages": [("user", "What's the temperature?")]})
print(f"\nAnswer: {result['messages'][-1].content}")

# Expected output:
# [RETRY] flaky_api failed (1/3): API timeout (attempt 1)
# [RETRY] flaky_api failed (2/3): API timeout (attempt 2)
# [RETRY] flaky_api succeeded on attempt 3
#
# Answer: The current temperature is 22°C with 65% humidity.

Exponential backoff (1s, 2s, 4s…) keeps you from hammering an API that's already struggling.


Custom error handling: friendly errors

Instead of letting errors break the agent, catch the exceptions and return informative messages:

from dotenv import load_dotenv
load_dotenv()

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

@tool
def get_stock_price(symbol: str) -> str:
    """Get the current price of a stock."""
    prices = {"AAPL": "$178.50", "GOOGL": "$141.20"}
    if symbol not in prices:
        raise ValueError(f"Symbol '{symbol}' not found")
    return f"{symbol}: {prices[symbol]}"

@tool
def get_exchange_rate(from_currency: str, to_currency: str) -> str:
    """Get the exchange rate between two currencies."""
    rates = {("USD", "MXN"): "17.15"}
    key = (from_currency.upper(), to_currency.upper())
    if key not in rates:
        raise ConnectionError("Currency service unavailable")
    return f"1 {from_currency} = {rates[key]} {to_currency}"

def wrap_tool_call(tool_call, config, call_next):
    """Handle errors by type, with useful messages."""
    tool_name = tool_call["name"]
    try:
        return call_next(tool_call, config)
    except ValueError as e:
        print(f"[ERROR] {tool_name}: data not found — {e}")
        return f"Error in {tool_name}: {e}. Try a different value."
    except ConnectionError as e:
        print(f"[ERROR] {tool_name}: service down — {e}")
        return f"Error in {tool_name}: service unavailable. Try again later."
    except Exception as e:
        print(f"[ERROR] {tool_name}: unexpected error — {type(e).__name__}: {e}")
        return f"Unexpected error in {tool_name}."

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[get_stock_price, get_exchange_rate],
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({
    "messages": [("user", "How much is TSLA stock in Mexican pesos?")]
})
print(result["messages"][-1].content)

# Expected output:
# [ERROR] get_stock_price: data not found — Symbol 'TSLA' not found
# I couldn't get the price for TSLA. The symbol wasn't found...

The model gets a clear error message and works it into its answer. The agent keeps running.


Tool-level logging and timing

Record every tool call in full detail and measure how long it takes:

from dotenv import load_dotenv
load_dotenv()

import time
import json
from datetime import datetime
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}': Python is an interpreted language."

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

tool_log = []

def wrap_tool_call(tool_call, config, call_next):
    """Record every tool call with a timestamp, duration, and status."""
    entry = {
        "timestamp": datetime.now().isoformat(),
        "tool": tool_call["name"],
        "args": tool_call["args"],
        "status": "pending",
    }
    start = time.time()

    try:
        result = call_next(tool_call, config)
        entry["status"] = "success"
        entry["result"] = str(result)[:200]
        return result
    except Exception as e:
        entry["status"] = "error"
        entry["error"] = f"{type(e).__name__}: {e}"
        raise
    finally:
        entry["duration_ms"] = round((time.time() - start) * 1000, 1)
        tool_log.append(entry)

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[search, calculator],
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({
    "messages": [("user", "What is Python, and what is 2**10?")]
})
print(f"Answer: {result['messages'][-1].content}\n")

print("=== Tool Log ===")
for entry in tool_log:
    print(f"  {entry['tool']} ({entry['duration_ms']}ms) → {entry['status']}")
    print(f"    Args: {entry['args']}")

# Expected output:
# Answer: Python is an interpreted language. 2^10 = 1,024.
#
# === Tool Log ===
#   search (1.2ms) → success
#     Args: {'query': 'what is Python'}
#   calculator (0.3ms) → success
#     Args: {'expression': '2**10'}

The try/except/finally block guarantees that every tool call gets logged, successful or not. In production, you'd ship this log to CloudWatch, Datadog, and the like.


Combining @wrap_model_call and @wrap_tool_call

For complete observability, use both middleware together:

from dotenv import load_dotenv
load_dotenv()

import time
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: Python 3.12 ships with performance improvements."

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

session_log = []

def wrap_model_call(messages, config, call_next):
    start = time.time()
    response = call_next(messages, config)
    elapsed = round(time.time() - start, 3)
    tool_info = "→ tool calls" if response.tool_calls else "→ final answer"
    session_log.append(f"MODEL ({elapsed}s) {len(messages)} msgs {tool_info}")
    return response

def wrap_tool_call(tool_call, config, call_next):
    start = time.time()
    try:
        result = call_next(tool_call, config)
        elapsed = round(time.time() - start, 3)
        session_log.append(f"TOOL  ({elapsed}s) {tool_call['name']}({tool_call['args']}) → OK")
        return result
    except Exception as e:
        elapsed = round(time.time() - start, 3)
        session_log.append(f"TOOL  ({elapsed}s) {tool_call['name']} → ERROR: {e}")
        raise

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[search, calculator],
    wrap_model_call=wrap_model_call,
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({
    "messages": [("user", "What's the newest version of Python, and what is 7**4?")]
})
print(f"Answer: {result['messages'][-1].content}\n")

print("=== Session Log ===")
for entry in session_log:
    print(f"  {entry}")

# Expected output:
# Answer: The most recent version of Python is 3.12... 7^4 = 2,401.
#
# === Session Log ===
#   MODEL (0.8s) 1 msgs → tool calls
#   TOOL  (0.001s) search({'query': '...'}) → OK
#   TOOL  (0.001s) calculator({'expression': '7**4'}) → OK
#   MODEL (0.6s) 4 msgs → final answer

The log shows the whole sequence: model decides → tools run → model answers. Total visibility into the pipeline.


Comparison: @wrap_model_call vs @wrap_tool_call

Aspect@wrap_model_call@wrap_tool_call
InterceptsCalls to the model (LLM)Tool executions
Receives(messages, config, call_next)(tool_call, config, call_next)
RunsOn every trip through the model nodeOn every tool execution
Can modifyMessages and the model's responseTool arguments and results
Typical use caseInjecting prompts, filtering, cost trackingRetry, error handling, timing, logging

The two middleware are complementary. Together they cover both kinds of operation an agent performs.


Troubleshooting

Problem 1: "TypeError: wrap_tool_call() takes 2 positional arguments but 3 were given"

Symptom: Error when creating or running the agent. Cause: Your function has the wrong signature. Fix: The signature must be exactly (tool_call, config, call_next):

def wrap_tool_call(tool_call, config, call_next):
    return call_next(tool_call, config)

Problem 2: The retry loop never ends

Symptom: The agent hangs indefinitely. Cause: Your retry logic has no upper bound. Fix: Always set a max_retries and a fallback:

def wrap_tool_call(tool_call, config, call_next):
    for attempt in range(3):
        try:
            return call_next(tool_call, config)
        except Exception:
            if attempt == 2:
                return f"Error: {tool_call['name']} unavailable"
            time.sleep(2 ** attempt)

Problem 3: Error when modifying tool_call args

Symptom: Error when you try to modify tool_call["args"] directly. Cause: The dictionary may be immutable. Fix: Build a new dictionary:

def wrap_tool_call(tool_call, config, call_next):
    modified = {**tool_call, "args": {**tool_call["args"], "limit": 10}}
    return call_next(modified, config)

Problem 4: The middleware doesn't run for one of the tools

Symptom: The middleware works for some tools but not others. Cause: @wrap_tool_call runs for every tool. If it looks like it doesn't, the model simply isn't calling that tool. Fix: Check whether the model is calling the tool at all by reviewing its docstring. The model decides which tool to use based on the description.


Exercises

Exercise 1: Basic tool logging (Easy)

Write a @wrap_tool_call that prints the name, args, and result of every tool. Try it with two tools.

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} — Python was created in 1991."

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

def wrap_tool_call(tool_call, config, call_next):
    print(f"[TOOL] {tool_call['name']}({tool_call['args']})")
    result = call_next(tool_call, config)
    print(f"[TOOL] → {str(result)[:100]}")
    return result

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[search, calculator],
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({"messages": [("user", "When was Python created, and what is 3**5?")]})
print(f"\nAnswer: {result['messages'][-1].content}")

# Expected output:
# [TOOL] search({'query': 'when was Python created'})
# [TOOL] → Result: when was Python created — Python was created in 1991.
# [TOOL] calculator({'expression': '3**5'})
# [TOOL] → 243
#
# Answer: Python was created in 1991. 3^5 = 243.

Explanation: The middleware intercepts every tool call, prints its details, runs the tool, and shows the result.

Exercise 2: Timing per tool (Easy)

Write middleware that measures how long each tool takes and prints the time in milliseconds.

See solution
from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import create_agent
from langchain_core.tools import tool

@tool
def instant_lookup(key: str) -> str:
    """Look up a local value."""
    return f"Value: {key} = 42"

@tool
def slow_search(query: str) -> str:
    """Search a slow external service."""
    time.sleep(0.5)
    return f"Result for '{query}': detailed information."

def wrap_tool_call(tool_call, config, call_next):
    start = time.time()
    result = call_next(tool_call, config)
    elapsed_ms = (time.time() - start) * 1000
    print(f"[TIMING] {tool_call['name']}: {elapsed_ms:.1f}ms")
    return result

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[instant_lookup, slow_search],
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({
    "messages": [("user", "Look up 'pi' locally and search for info about calculus")]
})
print(f"\nAnswer: {result['messages'][-1].content}")

# Expected output:
# [TIMING] instant_lookup: 0.2ms
# [TIMING] slow_search: 502.3ms
#
# Answer: The value for pi is 42. On calculus: detailed information.

Explanation: The gap between the two timings identifies slow_search as the bottleneck.

Exercise 3: Retry with exponential backoff (Medium)

Write a tool that fails the first 2 times and middleware with retry (3 attempts, backoff). Check that the third attempt succeeds.

See solution
from dotenv import load_dotenv
load_dotenv()

import time
from langchain.agents import create_agent
from langchain_core.tools import tool

call_counter = 0

@tool
def unreliable_api(query: str) -> str:
    """An API that fails intermittently."""
    global call_counter
    call_counter += 1
    if call_counter < 3:
        raise ConnectionError(f"Timeout on attempt {call_counter}")
    return f"Data for '{query}': successful response."

def wrap_tool_call(tool_call, config, call_next):
    max_retries = 3
    for attempt in range(max_retries):
        try:
            result = call_next(tool_call, config)
            if attempt > 0:
                print(f"[RETRY] OK on attempt {attempt + 1}")
            return result
        except Exception as e:
            if attempt < max_retries - 1:
                wait = 2 ** attempt
                print(f"[RETRY] Attempt {attempt + 1}/{max_retries} failed: {e}. Waiting {wait}s...")
                time.sleep(wait)
            else:
                return f"Error: {tool_call['name']} unavailable after {max_retries} attempts."

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[unreliable_api],
    wrap_tool_call=wrap_tool_call,
)

call_counter = 0
result = agent.invoke({"messages": [("user", "Query the API about LangChain")]})
print(f"\nAnswer: {result['messages'][-1].content}")

# Expected output:
# [RETRY] Attempt 1/3 failed: Timeout on attempt 1. Waiting 1s...
# [RETRY] Attempt 2/3 failed: Timeout on attempt 2. Waiting 2s...
# [RETRY] OK on attempt 3
#
# Answer: On LangChain: successful response.

Explanation: The backoff (1s, 2s) gives the service time to recover. The third attempt goes through.

Exercise 4: Blocking dangerous tools (Medium)

Create an agent with read, write, and delete tools. The middleware blocks delete by returning a friendly error. Check that the others still work.

See solution
from dotenv import load_dotenv
load_dotenv()

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

@tool
def read_data(key: str) -> str:
    """Read data from the system."""
    return f"Data for '{key}': value_123"

@tool
def write_data(key: str, value: str) -> str:
    """Write data to the system."""
    return f"Written: {key} = {value}"

@tool
def delete_data(key: str) -> str:
    """Delete data from the system."""
    return f"Deleted: {key}"

BLOCKED = {"delete_data"}

def wrap_tool_call(tool_call, config, call_next):
    if tool_call["name"] in BLOCKED:
        return f"'{tool_call['name']}' is blocked. Only reads and writes are allowed."
    return call_next(tool_call, config)

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[read_data, write_data, delete_data],
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({
    "messages": [("user", "Read 'config', write 'status'='active', and delete 'temp'")]
})
print(result["messages"][-1].content)

# Expected output:
# I read 'config' (value_123) and wrote 'status' = 'active'.
# Deleting 'temp' is blocked — only reads and writes are allowed.

Explanation: read_data and write_data go through. delete_data gets intercepted and returns a message the model folds into its answer.

Exercise 5: Full observability — model + tools (Challenge)

Combine @wrap_model_call and @wrap_tool_call. The model middleware counts calls and tokens. The tool middleware records timing and status. Print a dashboard at the end.

See solution
from dotenv import load_dotenv
load_dotenv()

import time
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} — relevant data."

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

dashboard = {
    "model_calls": 0, "model_time": 0.0,
    "input_tokens": 0, "output_tokens": 0,
    "tool_calls": 0, "tool_time": 0.0, "tool_errors": 0,
}

def wrap_model_call(messages, config, call_next):
    start = time.time()
    response = call_next(messages, config)
    dashboard["model_calls"] += 1
    dashboard["model_time"] += time.time() - start
    usage = response.response_metadata.get("token_usage", {})
    dashboard["input_tokens"] += usage.get("prompt_tokens", 0)
    dashboard["output_tokens"] += usage.get("completion_tokens", 0)
    return response

def wrap_tool_call(tool_call, config, call_next):
    start = time.time()
    try:
        result = call_next(tool_call, config)
        dashboard["tool_calls"] += 1
        dashboard["tool_time"] += time.time() - start
        return result
    except Exception as e:
        dashboard["tool_errors"] += 1
        dashboard["tool_calls"] += 1
        return f"Error: {e}"

agent = create_agent(
    "openai:gpt-4.1-mini",
    tools=[search, calculator],
    wrap_model_call=wrap_model_call,
    wrap_tool_call=wrap_tool_call,
)

result = agent.invoke({"messages": [("user", "What is AI, and what is 256 * 4?")]})
print(f"Answer: {result['messages'][-1].content}\n")

total_cost = (
    (dashboard["input_tokens"] / 1000) * 0.00015
    + (dashboard["output_tokens"] / 1000) * 0.0006
)
print("=== DASHBOARD ===")
print(f"Model: {dashboard['model_calls']} calls | {dashboard['model_time']:.2f}s")
print(f"Tokens: {dashboard['input_tokens']} in / {dashboard['output_tokens']} out")
print(f"Cost: ${total_cost:.6f} USD")
print(f"Tools: {dashboard['tool_calls']} calls | {dashboard['tool_time']:.3f}s | {dashboard['tool_errors']} errors")

# Expected output:
# Answer: AI (Artificial Intelligence) is... 256 × 4 = 1,024.
#
# === DASHBOARD ===
# Model: 2 calls | 1.23s
# Tokens: ~200 in / ~60 out
# Cost: $0.000066 USD
# Tools: 2 calls | 0.002s | 0 errors

Explanation: Both middleware feed a shared dashboard. Together they give you the full performance picture: the LLM (tokens, cost, latency) and the tools (executions, errors, timing).


Summary

In this capsule you learned:

  • @wrap_tool_call receives (tool_call, config, call_next) and wraps every tool execution
  • The tool_call is a dict with name, args, and id
  • Retry logic: retrying tools with exponential backoff to survive unstable APIs
  • Custom error handling: catching exceptions by type and returning messages the model can use
  • Tool-level logging: recording every execution with a timestamp, args, result, and status
  • Tool-level timing: measuring duration to find bottlenecks
  • Blocking tools: preventing dangerous tools from running by returning error messages
  • Full observability: combining @wrap_model_call + @wrap_tool_call for total visibility
  • The fallback pattern (retry → error message) beats letting exceptions propagate

Next capsule: Dynamic models — you'll learn to use middleware to automatically pick which model to use based on complexity, optimizing for both cost and latency.


Further reading

  1. create_agent API Reference — Full documentation of the middleware parameters
  2. LangGraph Agents — Tool Middleware — The official guide to tool middleware
  3. Custom Error Handling in Agents — How-to for custom error handling
  4. Tool Calling — LangChain — Tool calling concepts
  5. Exponential Backoff — Google Cloud — Reference on exponential backoff

Module 4 — LangChain & LangGraph: From Chains to Agents