Okay, so I've been deep in the weeds building a few different chatbots with LangChain lately, and while it's fantastic for prototyping, my cloud bill started giving me a little side-eye 😅. I realized I had absolutely no visibility into which chains, which models, or even which users were driving the cost. LangChain obscures the actual LLM calls a bit, so you can't just slap a simple wrapper on one client.
I decided I needed to track *every* LLM call—OpenAI, Anthropic, maybe even local models if I swap them in—and log the cost, tokens, and context for later analysis. After a bunch of tinkering and comparing a few methods, here's the step-by-step I landed on that's been working beautifully. It's all about using LangChain's **callbacks**.
The core idea is to create a custom callback handler that hooks into the `on_llm_start` and `on_llm_end` events. At the start, you note the time and prompt. At the end, you get the response and—critically—the `llm_output` which contains token usage for supported models.
Here's the blueprint for the handler:
```python
from langchain.callbacks.base import BaseCallbackHandler
import time
from typing import Any, Dict, List
from decimal import Decimal
class CostTrackingCallbackHandler(BaseCallbackHandler):
def __init__(self):
self.cost_records = []
self.prompt_tokens = 0
self.completion_tokens = 0
self.start_time = None
def on_llm_start(self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any) -> None:
self.start_time = time.time()
# Store the first prompt for this call; you might want to hash or truncate it.
self.current_prompt = prompts[0]
def on_llm_end(self, response: Any, **kwargs: Any) -> None:
latency = time.time() - self.start_time
llm_output = response.llm_output
token_usage = llm_output.get('token_usage', {}) if llm_output else {}
self.prompt_tokens = token_usage.get('prompt_tokens', 0)
self.completion_tokens = token_usage.get('completion_tokens', 0)
# Calculate cost based on model name (from serialized info or response)
# You'll need a pricing lookup function here.
model_name = response.llm_output.get('model_name', 'unknown')
cost = self._calculate_cost(model_name, self.prompt_tokens, self.completion_tokens)
record = {
'model': model_name,
'prompt': self.current_prompt[:200], # Truncated for logging
'prompt_tokens': self.prompt_tokens,
'completion_tokens': self.completion_tokens,
'latency': latency,
'cost': cost
}
self.cost_records.append(record)
def _calculate_cost(self, model_name: str, prompt_tokens: int, completion_tokens: int) -> Decimal:
# Your pricing logic here. Example for GPT-4:
pricing = {
'gpt-4': {'input': 0.03, 'output': 0.06},
'gpt-4-turbo-preview': {'input': 0.01, 'output': 0.03},
'gpt-3.5-turbo': {'input': 0.0005, 'output': 0.0015},
'claude-3-opus': {'input': 0.015, 'output': 0.075},
}
rates = pricing.get(model_name, {'input': 0, 'output': 0})
cost = (prompt_tokens / 1000) * rates['input'] + (completion_tokens / 1000) * rates['output']
return Decimal(cost).quantize(Decimal('0.000001'))
```
Then, you just attach this handler when you run your chain:
```python
cost_handler = CostTrackingCallbackHandler()
# For a single LLM call
result = llm.invoke("Tell me a joke", config={'callbacks': [cost_handler]})
# For a chain
chain.invoke({"input": "What's the weather?"}, config={'callbacks': [cost_handler]})
# Later, analyze your records
print(f"Total calls: {len(cost_handler.cost_records)}")
print(f"Total estimated cost: ${sum(r['cost'] for r in cost_handler.cost_records):.4f}")
```
Some pitfalls and extra steps I had to figure out:
* **Model Name Consistency:** The model name in `llm_output` can be tricky. Some providers pack it in metadata, some don't. You might need to parse it from the serialized dictionary in `on_llm_start` as a fallback.
* **Streaming:** If you use streaming responses, the token usage might not be populated in `llm_output`. You'd need a different strategy, maybe estimating from character count or using a separate tokenizer.
* **Aggregation & Export:** I hooked mine up to a simple SQLite database to log by session/user. You could also fire events to a monitoring service like Datadog or Prometheus.
* **Non-OpenAI Models:** For Anthropic, Bedrock, etc., the token usage field might be different. You'll need to adjust the key and pricing lookup.
The real power move? Wrap this in a context manager or a decorator so you can easily apply it to any chain without modifying your core logic. Now I can break down costs per feature and identify which prompts are token-hungry. It’s been a total game-changer for moving from "oh this is fine" to actual, accountable budgeting.
Has anyone else built something similar? I'm curious how you handled pricing for obscure/local models or if you found a way to get perfect token counts without relying on the provider's output.