LlamaIndex token counting and cost tracking across providers

Most teams adopt LlamaIndex and treat token consumption as a mystery until the invoice shows up. Reliable llamaindex token counting cost tracking across providers requires capturing the exact usage objects returned by each model, normalizing them into a single schema, and mapping them to a pricing table you control. This guide walks through a callback-based pipeline that works whether you call OpenAI directly, route through an OpenAI-compatible gateway, or mix Anthropic and Mistral models in the same index.

Step 1: Capture real provider usage instead of local estimates

LlamaIndex ships a TokenCountingHandler that tallies tokens with a local tokenizer. That number is an estimate and diverges from what providers bill, especially with cached prompts or provider-specific tokenization. For accurate llamaindex token counting cost tracking you need the usage block from the provider response.

Write a callback handler that listens to LLM events and pulls the raw response:

from llama_index.core.callbacks import BaseCallbackHandler, CBEventType, EventPayload

class UsageTrackingHandler(BaseCallbackHandler):
    def __init__(self):
        # Only subscribe to LLM completion events
        super().__init__(events=[CBEventType.LLM])
        self.records = []

    def on_event_start(self, event_type, payload=None, event_id=None, parent_id=None, **kwargs):
        return

    def on_event_end(self, event_type, payload=None, event_id=None, parent_id=None, **kwargs):
        if event_type != CBEventType.LLM:
            return
        response = payload.get(EventPayload.RESPONSE)
        if response is None:
            return
        raw = getattr(response, "raw", None)
        if raw is None:
            return
        usage = getattr(raw, "usage", None)
        if usage is None:
            # Some providers nest usage under a different key
            usage = getattr(raw, "token_usage", None)
        if usage is None:
            return
        self.records.append({
            "model": getattr(response, "model", "unknown"),
            "prompt_tokens": getattr(usage, "prompt_tokens", 0),
            "completion_tokens": getattr(usage, "completion_tokens", 0),
            "total_tokens": getattr(usage, "total_tokens", 0),
        })

The raw attribute on a LlamaIndex CompletionResponse is the SDK object from the underlying provider. For OpenAI-compatible endpoints it mirrors the OpenAI Python SDK, so usage.prompt_tokens is reliable. For Anthropic via LlamaIndex the field names are similar but you should confirm against the installed provider package.

Step 2: Normalize model identities and provider tags

A single query engine may call gpt-4o-mini, claude-3-5-sonnet, or a hosted Mistral model. Your cost table needs a stable key. Strip quantization suffixes and prefix with the provider if you route through a gateway that renames models.

def normalize_record(rec):
    model = rec["model"]
    # Example: "n4n/llama-3.1-70b" -> provider "n4n", family "llama-3.1-70b"
    if "/" in model:
        provider, family = model.split("/", 1)
    else:
        # Heuristic: openai models start with gpt, anthropic with claude
        if model.startswith("gpt"):
            provider = "openai"
        elif model.startswith("claude"):
            provider = "anthropic"
        else:
            provider = "unknown"
        family = model
    rec["provider"] = provider
    rec["family"] = family
    return rec

records = [normalize_record(r) for r in handler.records]

Doing this normalization inside the callback or immediately after collection keeps your later aggregation code free of provider branches.

Step 3: Wire the handler into LlamaIndex and run a query

Create a CallbackManager, attach the handler, and bind it to the LLM and global Settings. If you use a gateway, point the api_base at its OpenAI-compatible URL.

from llama_index.core.callbacks import CallbackManager
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI

handler = UsageTrackingHandler()
callback_manager = CallbackManager([handler])

llm = OpenAI(
    model="gpt-4o-mini",
    api_key="YOUR_KEY",
    api_base="https://api.n4n.ai/v1",  # single endpoint, 240+ models, per-token metering
    callback_manager=callback_manager,
)
Settings.llm = llm

# Assume `index` is a prebuilt VectorStoreIndex
query_engine = index.as_query_engine()
response = query_engine.query("What is the refund policy?")

Routing through n4n.ai here means one OpenAI-compatible endpoint addresses 240+ models and returns standardized usage metadata, so the handler in Step 1 needs no per-provider code. The gateway also forwards cache-control hints, which can show up as reduced prompt_tokens when a provider honors prompt caching.

Step 4: Load a pricing table and compute spend

Never hardcode prices in the tracking logic. Keep them in a JSON file you update when providers change rates. Use placeholder values below; replace with current numbers from provider docs.

{
  "openai/gpt-4o-mini": {"prompt": 0.00000015, "completion": 0.0000006},
  "anthropic/claude-3-5-sonnet": {"prompt": 0.000003, "completion": 0.000015},
  "n4n/llama-3.1-70b": {"prompt": 0.0000009, "completion": 0.0000009}
}

Load and compute:

import json

with open("pricing.json") as f:
    pricing = json.load(f)

def cost_for(rec):
    key = f"{rec['provider']}/{rec['family']}"
    rate = pricing.get(key)
    if not rate:
        return 0.0
    return rec["prompt_tokens"] * rate["prompt"] + rec["completion_tokens"] * rate["completion"]

total = 0.0
for rec in records:
    c = cost_for(rec)
    rec["cost_usd"] = c
    total += c

print(f"Total tracked spend: ${total:.6f}")

This separation lets you audit llamaindex token counting cost tracking without touching application code when GPT-4o-mini pricing shifts.

Step 5: Emit metrics for observability

A list of records is fine for a script, but production needs durable signals. Append each record to a JSONL file and optionally increment Prometheus counters.

import json

with open("usage.jsonl", "a") as f:
    for rec in records:
        f.write(json.dumps(rec) + "\n")

If you run a Prometheus client:

from prometheus_client import Counter

PROMPT_TOKENS = Counter("llm_prompt_tokens", "Prompt tokens", ["provider", "family"])
COMPLETION_TOKENS = Counter("llm_completion_tokens", "Completion tokens", ["provider", "family"])

for rec in records:
    PROMPT_TOKENS.labels(rec["provider"], rec["family"]).inc(rec["prompt_tokens"])
    COMPLETION_TOKENS.labels(rec["provider"], rec["family"]).inc(rec["completion_tokens"])

Grafana dashboards can then show spend per model family, which is the whole point of llamaindex token counting cost tracking across providers.

Step 6: Verify the pipeline end to end

Run a minimal script that executes one query and asserts the handler captured usage.

# verify.py
from your_module import handler, query_engine

response = query_engine.query("Test query")
assert len(handler.records) > 0, "No LLM usage captured"
rec = handler.records[0]
assert rec["total_tokens"] > 0, "Token count is zero"
print("Captured:", rec)

Execute it:

python verify.py

Success looks like a printed record with non-zero token counts and a cost_usd field after Step 4 runs. If records is empty, check that the LLM object actually received the callback_manager and that the provider returns a usage field. Some mock LLMs used in tests omit it; point at a real endpoint.

For continuous validation, add a unit test that feeds a fake response object with a raw usage stub into on_event_end and checks the record shape. That catches schema drift when you upgrade LlamaIndex or swap providers.

The pattern above gives you exact, provider-spanning visibility into token flow. You can extend the handler to capture embedding calls by subscribing to EMBEDDING events, or add a request_id from kwargs to join with application logs. The core requirement—accurate llamaindex token counting cost tracking—stays intact because you trust the provider’s usage block, not a local guess.