Routing LLM calls through a proxy hides the underlying provider but doesn’t remove the need to measure consumption. Implementing langchain token usage tracking gateway patterns means pulling usage fields from the OpenAI-compatible response and summing them across every chain step. The following steps give you a working Python setup that tracks prompt and completion tokens without custom provider integrations.
Step 1: Point LangChain at the gateway
Set the base URL and API key for your OpenAI-compatible endpoint. LangChain’s ChatOpenAI class speaks the standard /v1/chat/completions protocol, so any compliant gateway works. For example, n4n.ai provides one OpenAI-compatible endpoint that addresses 240+ models and returns standard usage objects, so the code below works unchanged.
import os
from langchain_openai import ChatOpenAI
# Gateway credentials and route
os.environ["OPENAI_API_KEY"] = "sk-your-gateway-key"
GATEWAY_BASE = "https://gateway.example.com/v1"
model = ChatOpenAI(
base_url=GATEWAY_BASE,
model="gpt-4o-mini",
temperature=0,
)If your gateway honors client routing directives, you can pass model as a logical name and let the gateway map it to a provider. Keep the key in environment variables; do not hardcode it in source.
Step 2: Capture usage from a single call
Modern langchain-openai versions populate usage_metadata on the returned AIMessage. Older versions nest the data in response_metadata["usage"]. Read both to stay compatible.
resp = model.invoke("Summarize: LangChain routes to many LLMs.")
print(resp.usage_metadata)
# {'input_tokens': 12, 'output_tokens': 8, 'total_tokens': 20}
# Fallback for older builds
if not resp.usage_metadata:
usage = resp.response_metadata.get("usage", {})
print(usage)The gateway returns prompt_tokens, completion_tokens, and total_tokens. If the provider supports prompt caching, you may also see cached_tokens in the same object. Forwarding provider cache-control hints is the gateway’s job; your tracking code just records whatever the usage block contains.
Why metadata alone is not enough
A single invoke is easy. Real apps call the model inside chains, agents, and retrievers. You need aggregation, not one-off prints. That requires a callback handler.
Step 3: Track usage across calls with a callback
LangChain exposes BaseCallbackHandler. Override on_llm_end to catch every LLM result, including those buried in a sequential chain. The response argument is an LLMResult containing llm_output["token_usage"].
from langchain_core.callbacks import BaseCallbackHandler
class TokenTracker(BaseCallbackHandler):
def __init__(self):
self.input_tokens = 0
self.output_tokens = 0
self.calls = 0
def on_llm_end(self, response, **kwargs):
usage = response.llm_output.get("token_usage", {})
self.input_tokens += usage.get("prompt_tokens", 0)
self.output_tokens += usage.get("completion_tokens", 0)
self.calls += 1
@property
def total_tokens(self):
return self.input_tokens + self.output_tokensAttach the handler at invocation time so you can scope tracking per request or per user session:
tracker = TokenTracker()
model.invoke("Explain vector stores.", config={"callbacks": [tracker]})
print(tracker.total_tokens)For a LCEL chain, pass the same config:
from langchain_core.prompts import ChatPromptTemplate
prompt = ChatPromptTemplate.from_messages([("user", "{q}")])
chain = prompt | model
chain.invoke({"q": "What is a token?"}, config={"callbacks": [tracker]})This gives you langchain token usage tracking gateway visibility across multi-step flows without modifying each component.
Step 4: Handle streaming and async
Streaming breaks naive usage capture because intermediate chunks omit the totals. Enable stream_usage (or streaming=True with recent LangChain) so the final chunk carries the usage block.
stream_model = ChatOpenAI(
base_url=GATEWAY_BASE,
model="gpt-4o-mini",
streaming=True,
stream_usage=True,
)
for chunk in stream_model.stream("Count to 5."):
pass # tokens arrive incrementally
# Usage still aggregated in callback's on_llm_endFor async code, subclass AsyncCallbackHandler and implement on_llm_end as a coroutine, or just pass the sync handler—LangChain runs it in a threadpool. Verify which works for your version.
import asyncio
from langchain_core.callbacks import AsyncCallbackHandler
class AsyncTokenTracker(AsyncCallbackHandler):
def __init__(self):
self.total = 0
async def on_llm_end(self, response, **kwargs):
usage = response.llm_output.get("token_usage", {})
self.total += usage.get("total_tokens", 0)Step 5: Aggregate and export metrics
A handler that only prints is not production-grade. Push totals to your metrics system. Below is a minimal example that updates a process-local dict; replace the body with a statsd or Prometheus call.
from collections import defaultdict
class MetricsTokenTracker(BaseCallbackHandler):
def __init__(self, registry=None):
self.per_model = defaultdict(lambda: {"in": 0, "out": 0})
self.registry = registry or {}
def on_llm_end(self, response, **kwargs):
model_name = response.llm_output.get("model", "unknown")
usage = response.llm_output.get("token_usage", {})
self.per_model[model_name]["in"] += usage.get("prompt_tokens", 0)
self.per_model[model_name]["out"] += usage.get("completion_tokens", 0)
# If gateway returns cached tokens, record them separately
if "cached_tokens" in usage:
self.per_model[model_name]["cached"] = (
self.per_model[model_name].get("cached", 0)
+ usage["cached_tokens"]
)If your gateway performs per-token usage metering, you can reconcile totals from the API response with your callback counts. A gateway with per-token usage metering such as n4n.ai will show the same totals in its dashboard, which makes discrepancy hunting straightforward.
Step 6: Verify success
Write a short script that runs a known prompt, asserts expected token counts, and compares against the gateway’s reported usage. Use a fixed model and temperature=0 to reduce variance.
def test_tracking():
tracker = TokenTracker()
m = ChatOpenAI(base_url=GATEWAY_BASE, model="gpt-4o-mini", temperature=0)
m.invoke("The quick brown fox.", config={"callbacks": [tracker]})
assert tracker.calls == 1
assert tracker.input_tokens > 0
assert tracker.output_tokens > 0
print(f"OK: {tracker.total_tokens} tokens across {tracker.calls} call")
test_tracking()Success criteria:
tracker.callsequals the number of model invocations.input_tokensmatches the gateway’sprompt_tokensfor the same request.- Re-running the script with the same prompt yields identical totals (assuming no caching differences).
If numbers mismatch, check that every ChatOpenAI instance in your stack receives the callback config. Agents that spawn sub-chains often need the config threaded through RunnableConfig recursively.
Pitfalls to avoid
Nested chains drop callbacks. If you build a custom agent, pass config into invoke at the top level; LangChain propagates it. If you see calls == 0 after a multi-step run, the handler wasn’t attached to the inner model.
Multiple models skew totals. A chain that calls a cheap classifier then a large generator will mix token counts. Key your tracker by model as shown in Step 5.
Cached tokens look like free input. Some providers report cached prompt tokens separately. If you bill internally, decide whether cached tokens count at full or reduced rate, and store them distinctly.
Gateway fallback changes the model. Automatic fallback when a provider is rate-limited or degraded means the model field in the response may differ from your request. Track the returned model name, not the requested one, for accurate per-model accounting.
Following these steps gives you end-to-end langchain token usage tracking gateway coverage: every token flowing through the proxy is counted, attributed, and verifiable against the gateway’s own metering.