When you need to trace token usage langchain pipeline across multi-step chains, the default debug output hides per-call breaks behind aggregated logs. Building a custom callback handler gives you exact prompt, completion, and total counts for every LLM and embedding call. This walkthrough instruments a real summarization-plus-QA pipeline from scratch and shows how to verify the numbers.
Step 1: Install and import the core packages
Use LangChain’s modular packages rather than the legacy langchain umbrella import. The examples below target langchain-core>=0.3 and langchain-openai>=0.1.
pip install langchain-core langchain-openaifrom langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParserIf you still run langchain<0.2, swap langchain_core.callbacks for langchain.callbacks.base. The handler interface is identical.
Step 2: Write a token-tracking callback handler
Subclass BaseCallbackHandler and override on_llm_end. OpenAI-compatible chat models attach a token_usage dict to response.llm_output. Aggregate from there.
class TokenUsageCallback(BaseCallbackHandler):
def __init__(self):
self.prompt_tokens = 0
self.completion_tokens = 0
self.total_tokens = 0
self.calls = []
def on_llm_end(self, response: LLMResult, **kwargs):
usage = (response.llm_output or {}).get("token_usage", {})
p = usage.get("prompt_tokens", 0)
c = usage.get("completion_tokens", 0)
t = usage.get("total_tokens", p + c)
self.prompt_tokens += p
self.completion_tokens += c
self.total_tokens += t
self.calls.append({"prompt": p, "completion": c, "total": t})The handler stores a per-call breakdown in calls so you can later attribute tokens to specific chain steps. If your pipeline requests n>1 completions per call, sum them inside the loop—token_usage from OpenAI already reflects the total for that request.
Step 3: Configure the LLM to expose token metadata
Instantiate ChatOpenAI normally. Token usage is reported by default; you do not need extra kwargs. If you route through an OpenAI-compatible gateway such as n4n.ai, set base_url and api_key—the gateway returns per-token usage metering in the same llm_output field and handles provider fallback automatically when a backend is degraded.
llm = ChatOpenAI(
model="gpt-4o-mini",
temperature=0,
# base_url="https://api.n4n.ai/v1",
# api_key="your-gateway-key",
)Avoid passing streaming=True without a plan: some versions omit llm_output in on_llm_end and instead emit usage only via on_llm_stream. For tracing, either disable streaming or accumulate deltas in on_llm_stream.
Step 4: Build a multi-step pipeline
Construct two chained sequences: a summarizer and a QA step that consumes the summary. This mirrors a realistic trace token usage langchain pipeline scenario where one document triggers multiple LLM calls.
summary_prompt = ChatPromptTemplate.from_messages([
("system", "Write a concise summary of the following text."),
("human", "{text}"),
])
summary_chain = summary_prompt | llm | StrOutputParser()
qa_prompt = ChatPromptTemplate.from_messages([
("system", "Answer the question using only the summary: {summary}"),
("human", "{question}"),
])
qa_chain = qa_prompt | llm | StrOutputParser()Each | composes a runnable. The parser strips the message to a string, but the callback still fires on the upstream LLM.
Step 5: Attach the handler and run
Pass the handler through the config dict. LangChain threads callbacks to every child runnable unless you explicitly silence them.
handler = TokenUsageCallback()
long_text = "..." * 200 # simulate a long document
summary = summary_chain.invoke(
{"text": long_text},
config={"callbacks": [handler]},
)
answer = qa_chain.invoke(
{"summary": summary, "question": "What is the primary claim?"},
config={"callbacks": [handler]},
)Both invoke calls share the same handler instance, so totals accumulate across the pipeline.
Step 6: Aggregate and report per step
After execution, inspect the handler. The calls list maps one entry per LLM invocation in order.
print(f"LLM calls: {len(handler.calls)}")
for i, call in enumerate(handler.calls, 1):
print(f" step {i}: prompt={call['prompt']} completion={call['completion']} total={call['total']}")
print(f"Aggregate prompt={handler.prompt_tokens} completion={handler.completion_tokens} total={handler.total_tokens}")You now have a complete trace token usage langchain pipeline output: step 1 is the summarizer, step 2 is the QA call. If you later insert a third refinement chain, it appears as step 3 without code changes.
Step 7: Verify success
Run the script and check three things:
len(handler.calls) == 2(one per chain).- Each
totalis greater than zero; a zero indicates the model didn’t returnllm_output(often a streaming or base-URL misconfig). handler.total_tokens == sum(c['total'] for c in handler.calls).
For an independent cross-check, set langchain.debug = True before the run. The verbose log prints raw provider responses containing token_usage. Match those numbers to your handler. If they diverge, your handler is parsing the wrong field—inspect response.llm_output directly.
Step 8: Extend to embeddings and retrievers
Summarization and QA are only half the story. If your pipeline uses a vector store, embedding calls consume tokens too. Capture them by overriding on_embedding_end:
def on_embedding_end(self, response, **kwargs):
usage = (getattr(response, "llm_output", None) or {}).get("token_usage", {})
p = usage.get("prompt_tokens", 0)
t = usage.get("total_tokens", p)
self.prompt_tokens += p
self.total_tokens += t
self.calls.append({"embedding_prompt": p, "total": t})Attach the same handler when calling embeddings.embed_documents(...). Retrieval itself (the similarity search) is free of token cost, but the embedding step is not, and teams often miss it when they trace token usage langchain pipeline spend.
Pitfalls and gotchas
Nested chains: If you compose summary_chain inside another chain via RunnableLambda, the callback still propagates. Do not instantiate a fresh handler inside a child runnable or you’ll lose the aggregate.
Multiple models: When mixing a cheap summarizer with a premium QA model, extend the handler to record kwargs.get("model", "unknown") from the run’s metadata. on_llm_end receives run_id and parent_run_id; use kwargs["invocation_params"].get("model_name") to label each step.
Streaming: With streaming=True, implement on_llm_stream and sum usage from the final chunk. Some providers send usage only in the last delta.
Cache hits: LangChain’s RunnableWithMessageHistory or semantic cache can skip LLM calls entirely. Your handler will show fewer calls than expected—that’s correct, not a bug.
Provider variance: Non-OpenAI models may nest usage under different keys. Print response.llm_output once during development to confirm the path before relying on it.
A solid token trace is the difference between guessing at LLM cost and enforcing budgets in CI. The callback pattern above drops into any LangChain deployment with zero changes to your business logic.