Back to Blog

Building a token-metered billing layer for LLM features

· 7 min read
engineering ai llm langchain billing architecture bunnydesk

If you ship more than one AI feature, sooner or later you need to know what your users are spending — and you need to bill for it. Here’s how I built that for BunnyDesk.

The constraints were familiar: multiple LLM providers and models, multiple features each making their own calls, per-workspace credits with plan-based grants and free-trial credits, an audit trail so customers (and us) could see exactly where their credits went, and zero tolerance for double-counting or losing a deduction.

This post walks through the design.

The shape of the problem

A simplified call site looks like this:

result = llm.invoke(prompt)

That call could be:

  • The agent orchestrator picking a tool.
  • The researcher sub-agent generating embeddings.
  • The page editor’s “Ask AI” button.
  • A background job rewriting a docs page.
  • A user-facing chat turn.

Each uses a different model. Each has a different cost-per-token, sometimes a different rate for input vs output. Each happens in a different request context. Each needs to: check the workspace has credits, deduct the credits after the call returns, and record what happened.

If you naively wrap every call site, you end up with duplicated logic, missed call sites, and inconsistent reporting.

Three pieces

The system has three pieces:

  1. A model registry that knows what every supported model costs.
  2. A credit ledger that records every deduction as an immutable row.
  3. A LangChain callback that hooks into every LLM call and writes to the ledger.

That’s the whole architecture.

The model registry

A ModelConfig table maps a model identifier (e.g. openai/gpt-5-mini, gemini-2.5-flash) to its prices — two prices because input and output tokens cost different amounts:

# Pseudocode
class ModelConfig(models.Model):
    identifier = models.CharField(unique=True)   # "openai/gpt-5-mini"
    purpose = models.CharField()                 # "agent", "embedding", "image", ...
    input_price_per_million = models.DecimalField()
    output_price_per_million = models.DecimalField()
    is_active = models.BooleanField(default=True)

Why a database table and not a config dict? Because model prices change, new models get added, and we want to do that without a code deploy. Also: per-purpose routing. The agent uses one model for now; a different purpose (say, “classify this topic”) routes to a cheaper-but-good-enough model. Routing lives in data, not code.

A get_llm(purpose) factory reads from this table and returns the right ChatOpenAI / ChatGoogleGenerativeAI instance. Call sites only ever name a purpose:

# Pseudocode
llm = get_llm("agent", workspace=workspace, metadata={"task_id": task.id})
result = llm.invoke(prompt)

That’s the only thing call sites need to know.

The credit ledger

The ledger is write-once. Every deduction (or grant, or refund) is a row that’s never updated:

# Pseudocode
class CreditTransaction(models.Model):
    workspace = models.ForeignKey(Workspace)
    amount = models.DecimalField()              # negative = deduction, positive = grant
    kind = models.CharField()                   # "llm_usage", "plan_grant", "trial_grant", ...
    model = models.ForeignKey(ModelConfig, null=True)
    purpose = models.CharField(null=True)
    tokens_in = models.IntegerField(null=True)
    tokens_out = models.IntegerField(null=True)
    metadata = models.JSONField()               # request_id, task_id, ...
    created_at = models.DateTimeField(auto_now_add=True)

The workspace’s current balance is SUM(amount) over all transactions. We cache it on the Workspace row for fast reads, but the ledger is the source of truth.

Why write-once?

  • Auditability. You can show a customer every credit they spent, what model, what purpose, when. There’s no “edited 5 minutes ago” footnote on any row.
  • Atomicity. Inserting a row is one INSERT. There’s no read-modify-write window where two concurrent requests can race on a balance update. The balance cache may lag by a second; the ledger never lies.

A negative balance is allowed in the data model — we don’t lose any deductions to a “your balance was already zero” race. Enforcement happens at the call site: if the cached balance is too low to afford the expected cost of an operation, we refuse to start the operation. Once started, the real deduction happens after the call.

The LangChain callback

This is the actual integration point. LangChain’s callback system fires on every LLM call:

# Pseudocode
class CreditTrackingCallback(BaseCallbackHandler):
    def __init__(self, workspace_id, model_config, purpose, metadata):
        self.workspace_id = workspace_id
        self.model_config = model_config
        self.purpose = purpose
        self.metadata = metadata

    def on_llm_end(self, response, **kwargs):
        usage = response.llm_output.get("token_usage") or {}
        tokens_in = usage.get("prompt_tokens", 0)
        tokens_out = usage.get("completion_tokens", 0)
        cost = (
            tokens_in * self.model_config.input_price_per_million / 1_000_000
            + tokens_out * self.model_config.output_price_per_million / 1_000_000
        )
        CreditTransaction.objects.create(
            workspace_id=self.workspace_id,
            amount=-cost,
            kind="llm_usage",
            model=self.model_config,
            purpose=self.purpose,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            metadata=self.metadata,
        )

The get_llm(purpose, ...) factory attaches one of these callbacks to every LLM it constructs. Call sites never see it.

The nice thing about this pattern: it doesn’t matter how many LLM calls a feature makes. If the agent calls the LLM 12 times in one task, you get 12 ledger rows — no per-feature accounting code. The deductions are itemized by tool/purpose because the callback carries that context. Customer support can answer “why did this task burn so many credits?” by pointing at the ledger.

Plan grants and free trials

The same ledger handles grants. A new workspace gets a kind="trial_grant" row with a positive amount. A paying workspace on a plan gets a recurring kind="plan_grant" row each billing cycle. Plan changes? Another row.

# Pseudocode — runs on Stripe `invoice.payment_succeeded`
CreditTransaction.objects.create(
    workspace=workspace,
    amount=plan.monthly_credit_grant,
    kind="plan_grant",
    metadata={"stripe_invoice_id": invoice.id, "period_start": ..., "period_end": ...},
)

Refunds (an agent crashed before producing useful output, say) are positive rows with kind="refund". Nothing in the system ever subtracts from history; everything is an additional row.

Plan gating vs credit gating

Plan-level feature gating is separate from credit balance. Some features are off entirely on the free plan even if you have credits; some are throttled. That’s a Plan.features check, not a credit check. The two systems compose: a feature must be both allowed by your plan and affordable from your workspace balance.

Keeping these two concerns separate has been one of the better calls. They change for different reasons (plans change on business decisions; credits change on pricing decisions) and gluing them together would have made both harder to reason about.

Gotchas

Streaming responses. LangChain’s on_llm_end fires once per call even for streaming — but you have to make sure stream_usage=True is set so the final chunk includes the token totals. We learned this the boring way (we’d deducted nothing for streamed calls until I noticed a workspace with an “always free” usage pattern).

Embeddings have a different callback path. Embeddings don’t hit on_llm_end; they hit on_embedding_end. Same callback, different method. Subclass cleanly and override both.

Pre-charge estimates. For long-running agent runs, we estimate the upper bound cost before starting, and check the balance against that. If the actual run uses less, the actual deduction is less. This avoids starting a 3-minute task only to fail at the last LLM call when credits run out mid-run.

Multi-LLM tasks. A single task can fan out to several LLM calls across different models. Each gets its own callback instance with its own model config — but they all carry the same task_id in metadata, so the ledger can group them later.

What I’d do differently

Two things, honestly:

  • I’d put the cost calculation in a single pure function that takes a ModelConfig and a usage dict, and unit-test it. We’ve had two off-by-1,000,000 bugs in the price math (per-million vs per-thousand). A pure function with tests would have caught both.
  • I’d add a “dry run” mode to the callback for development, so local LLM calls don’t write to the ledger but still log what they would have done. We just disabled the callback locally for a while, which let one real-money path slip through testing.

The whole system is maybe 250 lines of code if you exclude the data model. The hard part wasn’t writing the code — it was deciding to put it in one place, with one callback, behind one factory function. Everything else fell out from that.