Middleware and Callbacks in LangChain Agents

What are Callbacks?

LangChain fires events at the start and end of every chain, LLM call, tool invocation, and retriever call. A BaseCallbackHandler subscribes to these events — giving you a clean way to add cross-cutting concerns without touching business logic.

LLM events

on_llm_start
on_llm_new_token
on_llm_end
on_llm_error

Chain & tool events

on_chain_start / end
on_tool_start / end
on_retriever_start / end
on_agent_action / finish

Writing a Custom Handler

Subclass BaseCallbackHandler and override only the events you care about. Methods you skip are no-ops by default.

python

from langchain_core.callbacks import BaseCallbackHandler

class LoggingHandler(BaseCallbackHandler):
    def on_llm_start(self, serialized, prompts, **kwargs):
        print(f"[LLM] Starting with {len(prompts)} prompt(s)")

    def on_llm_new_token(self, token: str, **kwargs):
        print(token, end="", flush=True)

    def on_llm_end(self, response, **kwargs):
        usage = response.llm_output.get("token_usage", {})
        print(f"\n[LLM] Tokens — prompt: {usage.get('prompt_tokens')}, "
              f"completion: {usage.get('completion_tokens')}")

    def on_chain_error(self, error, **kwargs):
        print(f"[ERROR] {error}")

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", callbacks=[LoggingHandler()])
llm.invoke("What is LangChain?")

Streaming Tokens

LangChain has two streaming patterns: callback-based (push) and iterator-based (pull). Use iterator-based streaming in web servers for the simplest code.

python

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", streaming=True)

# --- Iterator streaming (pull) ---
for chunk in llm.stream("Explain LangChain in 3 sentences."):
    print(chunk.content, end="", flush=True)

# --- Async streaming (FastAPI / async servers) ---
import asyncio

async def stream_response():
    async for chunk in llm.astream("Explain LangChain."):
        print(chunk.content, end="", flush=True)

asyncio.run(stream_response())

# --- Full chain streaming ---
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

chain = ChatPromptTemplate.from_template("Tell me about {topic}") | llm | StrOutputParser()

for chunk in chain.stream({"topic": "LangGraph"}):
    print(chunk, end="", flush=True)

Cost & Token Tracking

Use the built-in get_openai_callback context manager to measure token usage and estimated cost for a block of code.

python

from langchain_community.callbacks import get_openai_callback
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini")

with get_openai_callback() as cb:
    llm.invoke("What is LangChain?")
    llm.invoke("What is LangGraph?")

print(f"Prompt tokens:     {cb.prompt_tokens}")
print(f"Completion tokens: {cb.completion_tokens}")
print(f"Total tokens:      {cb.total_tokens}")
print(f"Total cost (USD):  ${cb.total_cost:.6f}")

Note: get_openai_callback works for any OpenAI-compatible model. For other providers, implement a custom handler that reads response.llm_output["token_usage"] in on_llm_end.

Runtime Callback Configuration

Attach callbacks at invoke time with RunnableConfig so you can pass request-specific handlers without changing the chain definition.

python

from langchain_core.runnables import RunnableConfig

class RequestLogger(BaseCallbackHandler):
    def __init__(self, request_id: str):
        self.request_id = request_id

    def on_llm_start(self, serialized, prompts, **kwargs):
        print(f"[{self.request_id}] LLM start")

chain = prompt | llm | parser

# Per-request callback — safe in concurrent environments
result = chain.invoke(
    {"topic": "LangChain"},
    config=RunnableConfig(callbacks=[RequestLogger("req-abc123")])
)

Retry & Fallback Middleware

LangChain LCEL has built-in retry and fallback support — these act as middleware without any callback boilerplate.

python

from langchain_openai import ChatOpenAI

fast_llm = ChatOpenAI(model="gpt-4o-mini")
smart_llm = ChatOpenAI(model="gpt-4o")

# Retry on transient errors (rate limits, timeouts)
resilient_llm = fast_llm.with_retry(
    stop_after_attempt=3,
    wait_exponential_jitter=True,
)

# Fallback: try gpt-4o-mini, fall back to gpt-4o on error
fallback_chain = fast_llm.with_fallbacks([smart_llm])

# Combine retry + fallback
chain = (fast_llm
    .with_retry(stop_after_attempt=2)
    .with_fallbacks([smart_llm.with_retry(stop_after_attempt=2)])
)

result = chain.invoke("Tell me a joke.")

Async Callbacks

For async applications use AsyncCallbackHandler so callback methods don't block the event loop.

python

from langchain_core.callbacks import AsyncCallbackHandler

class AsyncStreamingHandler(AsyncCallbackHandler):
    def __init__(self, queue: asyncio.Queue):
        self.queue = queue

    async def on_llm_new_token(self, token: str, **kwargs):
        await self.queue.put(token)

    async def on_llm_end(self, response, **kwargs):
        await self.queue.put(None)  # signal end

# FastAPI SSE example
async def generate(prompt: str):
    queue = asyncio.Queue()
    handler = AsyncStreamingHandler(queue)
    asyncio.create_task(llm.ainvoke(prompt, config={"callbacks": [handler]}))
    while True:
        token = await queue.get()
        if token is None:
            break
        yield token

Callbacks pair naturally with the Tools and Agents lesson, where tracing every tool call helps you debug agent loops.

Middleware and Callbacks FAQ

What is agent middleware in LangChain?

Agent middleware is cross-cutting logic that wraps a LangChain run without changing the chain or agent definition. In LCEL this includes built-in behavior such as with_retry and with_fallbacks, plus callback handlers that observe and react to events. Middleware lets you add logging, tracing, cost tracking, retries, and request shaping in one place so your core agent code stays clean.

What are callbacks in LangChain?

Callbacks are event hooks that LangChain fires at the start and end of every chain, LLM call, tool invocation, and retriever call. You subclass BaseCallbackHandler (or AsyncCallbackHandler for async apps) and override only the events you care about, such as on_llm_start, on_llm_new_token, and on_llm_end. Callbacks are how you observe an agent run without editing its logic.

How do I log and trace a LangChain agent with callbacks?

Write a BaseCallbackHandler that prints or records data inside on_llm_start, on_chain_start, on_tool_start, and on_llm_end, then pass it via callbacks=[handler] on the model or per request with RunnableConfig. Per-request handlers are safe in concurrent servers. For full tracing across a run, attach the handler at the top-level invoke so every nested step inherits it.

How do I modify a request before the model runs in LangChain?

Shape the input before it reaches the model by composing a step ahead of the LLM in your LCEL chain, for example a RunnableLambda or prompt template that rewrites or enriches the messages. Callbacks observe events but should not mutate them; for request shaping put the transformation in the chain itself so the modified request flows into the model deterministically.

How do I track token usage and cost with callbacks?

Use the get_openai_callback context manager to sum prompt tokens, completion tokens, total tokens, and estimated cost for any block of OpenAI calls. For other providers, write a custom handler that reads response.llm_output['token_usage'] inside on_llm_end and accumulates the totals. Attach it at invoke time so it captures every call in the run.

What are common middleware use cases for LangChain agents?

Common middleware use cases include structured logging and tracing, streaming tokens to a UI, per-request cost and token tracking, retrying transient errors with with_retry, falling back to a stronger model with with_fallbacks, and attaching request-scoped handlers with RunnableConfig. These behaviors are added around the agent so the core chain code stays focused on the task.