Build a usage logging system with token + cost capture

Introduction

When your compound AI router dispatches inference across OpenAI, Anthropic, and Gemini, every provider returns token counts in a slightly different shape and charges a different per-token rate — and if you don't capture that cost at the moment of each call, the monthly bill arrives as a number you can't decompose by model, by user, or by feature. Without per-request tracking you can't enforce budgets, can't diagnose runaway prompts, and can't tell whether your fallback tier is being abused. By the end of this lesson you will be able to attach a LiteLLM async success callback that extracts normalized usage metadata, computes per-request cost from LiteLLM's pricing tables, and persists a structured record to PostgreSQL for every successful inference.

Key Terminology

• ModelResponse: LiteLLM's normalized response object returned from acompletion, exposing a provider-agnostic usage attribute with prompt_tokens, completion_tokens, and total_tokens.
• async_success_callback: LiteLLM's asynchronous hook that fires after every successful inference call, receiving kwargs and the ModelResponse so a tracker can extract usage and persist a record without blocking the request path.
• completion_cost: LiteLLM helper (litellm.completion_cost) that computes per-request USD cost from a model identifier plus prompt and completion token counts, using LiteLLM's internal pricing registry.

Concepts

How LiteLLM exposes token usage metadata

When you call litellm.acompletion, the returned ModelResponse object carries a usage attribute that LiteLLM populates by normalizing each provider's native response format. For OpenAI, this maps directly from the usage field in the API response. For Anthropic, LiteLLM translates input_tokens and output_tokens into the OpenAI-compatible prompt_tokens and completion_tokens keys. For Gemini, the usageMetadata block undergoes the same normalization. The result is that your tracking code never needs provider-specific branching—a single access pattern works across all four providers in your gateway.

Beyond raw token counts, LiteLLM provides a _hidden_params dictionary on the response containing model_id (the internal identifier used by litellm.Router), api_base (the endpoint URL), and critically, the actual model string that was sent to the provider. This distinction matters because your router might map the alias "fast-tier" to "gemini/gemini-2.5-flash", and you need to log the resolved model name for accurate cost calculation. LiteLLM also exposes response_cost directly in the response metadata when you enable cost tracking, computed from its internal pricing table that covers over 300 model variants.

Key considerations for production deployments

• Batch inserts for high throughput: If your gateway handles thousands of requests per second, individual INSERT statements per callback will saturate your database connection pool. Buffer records in an in-memory queue and flush in batches of 100-500 using session.add_all(). A background asyncio.Task can drain the queue on a timer or when the buffer reaches a size threshold.

• Streaming response handling: When using stream=True with acompletion, token usage is not available on intermediate chunks. LiteLLM accumulates usage and attaches it to the final chunk's usage field only when stream_options={"include_usage": True} is set. Your callback must detect streaming responses and handle the case where usage is populated only on the aggregated response.

• Cost accuracy for fine-tuned models: LiteLLM's built-in pricing covers standard model variants but may not reflect custom fine-tuned model pricing. Pass a custom_pricing dictionary to override per-token rates: litellm.register_model({"custom-ft-model": {"input_cost_per_token": 0.000003, "output_cost_per_token": 0.000015}}).

• Table partitioning: For gateways logging millions of requests daily, partition the llm_usage_log table by created_at using PostgreSQL's native range partitioning. Monthly partitions enable efficient pruning of old data and keep index sizes manageable for real-time queries against recent records.

• Correlation with Router metrics: When using litellm.Router with fallback chains (covered in another goal), the kwargs dictionary includes model_group (the deployment group name) and litellm_model_name (the specific deployment that handled the request). Log both to distinguish between primary and fallback model usage—if 30% of your traffic hits the fallback provider, that signals a reliability problem with your primary.

Code Walkthrough

Architecture of the callback-based tracking pipeline

The UsageLogger integrates with LiteLLM through its callback mechanism rather than wrapping every call site. LiteLLM supports both synchronous success_callback and asynchronous async_success_callback hooks. For a production gateway handling concurrent requests, the async variant is essential—it avoids blocking the event loop while database writes complete. The callback receives two arguments: kwargs (the original request parameters including model, messages, and any metadata you attached) and response (the full ModelResponse object). Your logger extracts token counts from response.usage, computes cost via litellm.completion_cost(), and writes the record to PostgreSQL.

The following diagram illustrates the data flow from an acompletion call through the callback pipeline into PostgreSQL:

Every LLM call costs money, and without per-request tracking, budgets blow up fast. This sequence diagram traces the full lifecycle of a litellm.acompletion call—from the Application Code dispatching a request with model, messages, and metadata, through provider-normalized ModelResponse delivery, to the UsageLogger.async_success_callback extracting token counts, computing cost, and persisting each record to the llm_usage_log table in PostgreSQL. This asynchronous callback pattern ensures usage logging never blocks inference latency.

Loading diagram...

Notice that the callback fires after the response has already been returned to the application. This non-blocking design means tracking latency never adds to user-perceived latency. If the database write fails, the application request is unaffected—you handle logging failures with a retry queue or dead-letter log, never by raising into the caller.

Schema and UsageLogger callback

Before building the logger, you need a PostgreSQL schema that supports both real-time cost queries ("how much has user X spent today?") and analytical aggregations ("what's the average cost per model tier this week?"). The UsageRecord model stores the resolved model name, provider, token counts, computed cost in USD, request latency, and an optional request_fingerprint column that links back to your application's request tracing system. The metadata_json column uses PostgreSQL's native JSONB type to store arbitrary key-value pairs — such as the Instructor retry count or the compound router's complexity classification — without schema migrations.

With the schema defined, the UsageLogger class implements LiteLLM's async_success_callback interface. The constructor accepts an async SQLAlchemy AsyncSession factory (typically created via async_sessionmaker), and the callback method extracts all necessary fields from kwargs and response. The cost calculation delegates to litellm.completion_cost(), which looks up the model's per-token pricing from LiteLLM's internal registry covering OpenAI, Anthropic, Google, and open-source model pricing. For custom or self-hosted models where LiteLLM has no pricing data, the function returns 0.0, and you can override with a custom pricing dictionary passed to the logger. The _extract_provider helper parses the provider prefix from LiteLLM's model string convention (e.g., extracting "gemini" from "gemini/gemini-2.5-flash"), falling back to "openai" when the model string lacks a prefix.

Code snippetpython
1import litellm
2import time
3import uuid
4from datetime import datetime, timezone
5
6from sqlalchemy import Column, Integer, String, Float, DateTime
7from sqlalchemy.dialects.postgresql import JSONB
8from sqlalchemy.orm import declarative_base
9from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
10
11Base = declarative_base()
12
13class UsageRecord(Base):
14    __tablename__ = "llm_usage_log"
15
16    id = Column(Integer, primary_key=True, autoincrement=True)
17    request_id = Column(String(64), nullable=False, index=True)
18    model = Column(String(128), nullable=False, index=True)
19    provider = Column(String(64), nullable=False)
20    prompt_tokens = Column(Integer, nullable=False, default=0)
21    completion_tokens = Column(Integer, nullable=False, default=0)
22    total_tokens = Column(Integer, nullable=False, default=0)
23    cost_usd = Column(Float, nullable=False, default=0.0)
24    latency_ms = Column(Float, nullable=True)
25    request_fingerprint = Column(String(128), nullable=True, index=True)
26    metadata_json = Column(JSONB, nullable=True)
27    created_at = Column(
28        DateTime(timezone=True),
29        nullable=False,
30        default=lambda: datetime.now(timezone.utc),
31    )
32
33class UsageLogger:
34    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
35        self.session_factory = session_factory
36
37    async def async_success_callback(self, kwargs: dict, response, *args, **cb_kwargs):
38        usage = getattr(response, "usage", None)
39        if usage is None:
40            return
41
42        model = kwargs.get("model", "unknown")
43        request_id = kwargs.get("litellm_call_id", str(uuid.uuid4()))
44        metadata = kwargs.get("metadata", {}) or {}
45        start_time = kwargs.get("start_time")
46
47        try:
48            cost = litellm.completion_cost(
49                model=model,
50                prompt_tokens=usage.prompt_tokens,
51                completion_tokens=usage.completion_tokens,
52            )
53        except Exception:
54            cost = 0.0
55
56        latency_ms = None
57        if start_time is not None:
58            latency_ms = (time.time() - start_time) * 1000
59
60        record = UsageRecord(
61            request_id=request_id,
62            model=model,
63            provider=self._extract_provider(model),
64            prompt_tokens=usage.prompt_tokens,
65            completion_tokens=usage.completion_tokens,
66            total_tokens=usage.total_tokens,
67            cost_usd=cost,
68            latency_ms=latency_ms,
69            request_fingerprint=metadata.get("request_fingerprint"),
70            metadata_json=metadata,
71        )
72
73        async with self.session_factory() as session:
74            async with session.begin():
75                session.add(record)
76
77    @staticmethod
78    def _extract_provider(model: str) -> str:
79        if "/" in model:
80            return model.split("/", 1)[0]
81        return "openai"

• The metadata_json JSONB column lets you record arbitrary per-request context (router tier, retry count, feature flag) without schema migrations as your tracking dimensions evolve.
• The created_at timestamp with timezone defaults to UTC, enabling time-range partitioning for tables that grow to millions of rows per day.
• The constructor accepts an async_sessionmaker factory rather than a raw session, ensuring each callback invocation gets its own session scope — critical for concurrent request handling where multiple callbacks fire simultaneously.
• The guard clause if usage is None: return handles streaming responses where usage may arrive only on the final aggregated chunk; the callback safely no-ops on chunks without usage.
• _extract_provider splits the model string on the first / character, following LiteLLM's "provider/model-name" convention, defaulting to "openai" when bare model strings like "gpt-4o" are passed.

Registering the logger with LiteLLM

Wiring the UsageLogger into your gateway requires two steps: creating the async session factory and registering the callback. The following code shows how to initialize the logger at application startup—typically inside your FastAPI lifespan handler or an equivalent async context manager. The litellm.success_callback list accepts both function references and string names of built-in integrations (like "langfuse" or "lunary"). By appending your logger's method to this list, every successful acompletion and completion call automatically triggers the tracking pipeline. Failed requests (provider errors, timeouts) hit the separate failure_callback hook, which you can wire to a similar logger that records error metadata for reliability analysis.

1from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
2import litellm
3
4DATABASE_URL = "postgresql+asyncpg://user:pass@localhost:5432/llm_gateway"
5
6engine = create_async_engine(DATABASE_URL, pool_size=10, max_overflow=20)
7session_factory = async_sessionmaker(engine, expire_on_commit=False)
8
9logger = UsageLogger(session_factory=session_factory)
10litellm.success_callback = [logger.async_success_callback]
11
12# Every acompletion call now triggers tracking automatically
13response = await litellm.acompletion(
14    model="anthropic/claude-sonnet-4-20250514",
15    messages=[{"role": "user", "content": "Explain circuit breakers"}],
16    metadata={
17        "request_fingerprint": "doc-summary-abc123",
18        "complexity_tier": "medium",
19    },
20)

Do's and Don'ts

Do's

1. ✓Do register your logger via async_success_callback — using the asynchronous variant instead of success_callback ensures PostgreSQL writes from UsageLogger never block the event loop, so tracking latency is invisible to callers of litellm.acompletion.
2. ✓Do delegate cost calculation to litellm.completion_cost() — it looks up per-token pricing for OpenAI, Anthropic, and Gemini from LiteLLM's internal registry, giving you a single consistent cost signal in USD across all providers without maintaining your own rate tables.
3. ✓Do use PostgreSQL's JSONB type for metadata_json — storing variable-shape callback context (Instructor retry counts, compound router complexity classifications, request fingerprints) in the JSONB column lets you persist arbitrary per-request telemetry without issuing schema migrations each time a new metadata key is introduced.

Don'ts

1. ✗Don't raise exceptions inside async_success_callback — if the INSERT INTO llm_usage_log write fails, absorb the error with a retry queue or dead-letter log; letting the exception propagate into the callback will surface as an inference failure to the caller even though the LLM response was already returned successfully.
2. ✗Don't assume litellm.completion_cost() returns a nonzero value for custom or self-hosted models — it returns 0.0 when the model string isn't in LiteLLM's pricing registry, so you must supply a custom pricing dictionary to UsageLogger for any model whose cost would otherwise silently record as free.
3. ✗Don't parse the provider out of LiteLLM's model string with ad-hoc splitting — the _extract_provider helper handles the "provider/model-name" convention (e.g., extracting "gemini" from "gemini/gemini-2.5-flash") and falls back to "openai" for bare model names; reimplementing this inline breaks on any model string that omits the prefix.