Counting tokens per customer when your API calls models through LiteLLM

[supertj]

This is about a setup where you sell an API to your own customers, and behind it you call models through the LiteLLM Python SDK. LiteLLM tells you how many tokens each call used. What it can't know is which of your customers the call was for, or what their plan allows. That's the part we handle in our own API layer, and the usage object LiteLLM returns gives us the number to count.

Here's a FastAPI endpoint. The UsageRail middleware checks the caller's API key, rate limit and quota before the endpoint runs. The app is set to count units (Settings > What counts as usage in the console), so track() records the tokens against that customer; in an app that counts requests, track() does nothing.

import os

import litellm
from fastapi import FastAPI, Request
from usagerail.fastapi import UsageRailMiddleware

app = FastAPI()
app.add_middleware(
    UsageRailMiddleware,
    app_key=os.environ["USAGERAIL_APP_KEY"],
    auth="api-key",  # reads Authorization: Bearer <key>, or x-api-key
)


@app.post("/v1/chat")
async def chat(request: Request):
    body = await request.json()
    response = await litellm.acompletion(
        model=body.get("model", "gpt-4o-mini"),
        messages=body["messages"],
    )
    # Prompt plus completion tokens, as reported by the provider
    request.state.usagerail.track(response.usage.total_tokens)
    return {"content": response.choices[0].message.content}

To try it without a provider key, pass mock_response="Hello" to acompletion. LiteLLM then returns a canned response with a token count, and the rest of the flow works the same.

Notes from wiring this up:

• Count after the call, reserve before it. The middleware has to decide whether to let a request in before the model runs, when the token count isn't known yet. In apps that count units, each route has a "Reserve with track()" setting: set it to the most tokens one request can use (your max_tokens plus the largest prompt you accept). If a request uses more than it reserved, the extra still counts, but it can take the customer past their quota.
• track() keeps the last value. It replaces, it doesn't add. If one request makes several model calls (a tool loop, say), add up usage.total_tokens yourself and call track() with the total.
• Failed calls don't count. If acompletion raises and your endpoint returns a 4xx or 5xx, the request is recorded but doesn't count toward the customer's usage or quota, whatever you passed to track().
• Weight tokens if models cost different amounts. One track() number per request means a token on a large model counts the same as one on a small model. If that's wrong for your pricing, multiply by a per-model factor before you pass the number in. (Route weights don't help here: a request that calls track() counts the tracked number, not the weight.)
• If you also run the LiteLLM proxy. Its virtual keys can have budgets and rate limits too. In this setup the customer's plan and quota are enforced in your API layer, by the middleware. You can use both.

Disclosure: I work on UsageRail. FastAPI setup and options: https://usagerail.com/docs/fastapi