Building a custom LangChain callback for cost tracking

When you need fine-grained visibility into LLM spend inside your Python applications, building a langchain custom callback cost tracking handler is the most direct way to attribute token usage to specific chains, users, or product features. LangChain’s callback system exposes lifecycle hooks that receive token counts straight from provider responses, so you can compute cost without wrapping every call by hand.

Prerequisites

• Python 3.10 or newer
• langchain core plus a model integration (langchain-openai used here)
• An API key for an OpenAI-compatible provider

pip install "langchain>=0.2.0" langchain-openai python-dotenv

Set your key in the environment:

import os
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY", "sk-...")

Why not just use the provider dashboard

Provider dashboards aggregate by account, not by feature or user. If a retrieval chain silently doubles its context size, you will see the bill rise but won’t know which code path caused it. In-process attribution closes that gap. A langchain custom callback cost tracking handler runs in the same process as your chain, so it can tag spend with arbitrary metadata.

The pricing table

Hardcode public list prices (USD per 1M tokens) as of mid-2024. Negotiated enterprise rates will differ, but the structure stays the same.

# pricing as of mid-2024, USD per 1M tokens
MODEL_PRICING = {
    "gpt-4o": {"prompt": 5.0, "completion": 15.0},
    "gpt-3.5-turbo": {"prompt": 0.5, "completion": 1.5},
    "claude-3-sonnet": {"prompt": 3.0, "completion": 15.0},
}

Implementing the callback

Subclass BaseCallbackHandler and override on_llm_end. The LLMResult object carries llm_output with token_usage for OpenAI-compatible backends.

from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import LLMResult

class CostTrackingCallback(BaseCallbackHandler):
    def __init__(self, pricing: dict, model_name: str = None):
        self.pricing = pricing
        self.model_name = model_name
        self.total_prompt_tokens = 0
        self.total_completion_tokens = 0
        self.total_cost = 0.0
        self.calls = 0

    def on_llm_end(self, response: LLMResult, **kwargs):
        llm_output = response.llm_output or {}
        usage = llm_output.get("token_usage", {})
        prompt_tokens = usage.get("prompt_tokens", 0)
        completion_tokens = usage.get("completion_tokens", 0)
        model = self.model_name or kwargs.get("model", "gpt-3.5-turbo")
        price = self.pricing.get(model)
        if not price:
            return
        cost = (prompt_tokens / 1_000_000) * price["prompt"] + \
               (completion_tokens / 1_000_000) * price["completion"]
        self.total_prompt_tokens += prompt_tokens
        self.total_completion_tokens += completion_tokens
        self.total_cost += cost
        self.calls += 1

    def report(self) -> str:
        return (
            f"Calls: {self.calls} | "
            f"Prompt tokens: {self.total_prompt_tokens} | "
            f"Completion tokens: {self.total_completion_tokens} | "
            f"Est. cost: ${self.total_cost:.4f}"
        )

In recent LangChain versions, on_llm_end fires for both chat and text models, so one handler covers ChatOpenAI and OpenAI.

Wiring it into a chain

Create a chat model, attach the callback, and run a trivial prompt.

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

callback = CostTrackingCallback(MODEL_PRICING, model_name="gpt-4o")
llm = ChatOpenAI(model="gpt-4o", temperature=0, callbacks=[callback])

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a terse assistant."),
    ("user", "{input}")
])

chain = prompt | llm
result = chain.invoke({"input": "What is the capital of France?"})
print(result.content)
print(callback.report())

Expected output at this checkpoint:

Paris
Calls: 1 | Prompt tokens: 15 | Completion tokens: 1 | Est. cost: $0.0001

Token counts depend on the exact tokenizer; treat the numbers as illustrative.

Tracking across nested chains

Pass the same callback instance via config to guarantee propagation into sub-chains and tools.

chain.invoke(
    {"input": "Explain quantum computing in one sentence."},
    config={"callbacks": [callback]}
)
print(callback.report())

Because the handler holds mutable accumulators, every LLM call inside the run updates the same totals.

Handling streaming and async

Streaming does not break usage accounting: LangChain still calls on_llm_end with the final token_usage payload. For async apps, subclass AsyncCallbackHandler if you need to write to a database without blocking, but the sync handler works inside ainvoke via the default executor.

from langchain.callbacks.base import AsyncCallbackHandler

class AsyncCostTrackingCallback(AsyncCallbackHandler):
    async def on_llm_end(self, response, **kwargs):
        # identical accounting logic, awaits omitted for brevity
        ...

Integrating with a gateway

If you route through an OpenAI-compatible endpoint that already returns usage, the langchain custom callback cost tracking code above works unchanged. A gateway like n4n.ai provides per-token usage metering at the edge and automatic fallback across 240+ models, but attaching this callback still lets you slice spend by LangChain chain or end-user ID without polling a separate dashboard.

Gotchas

• Some local wrappers (e.g., certain llama.cpp bindings) omit token_usage. Estimate with a local tokenizer like tiktoken and patch llm_output in a custom LLM class.
• Retries fire on_llm_end per attempt. Dedup on run_id if you only want billed attempts.
• Cost is an estimate. Reconcile with provider invoices at month end.

Final thoughts

Building a langchain custom callback cost tracking handler takes about 30 lines and pays off the first time a chain unexpectedly burns tokens. Drop it into your middleware, tag runs with user IDs, and log callback.report() to your metrics stack.