Stop One User From Eating Your AI Budget: Build a Per-User Spending Guard in Python
Stop One User From Eating Your AI Budget: Build a Per-User Spending Guard in Python
Logging tells you what happened. It does not stop anything. Our last tutorial built a logging wrapper that records every call's price — and that wrapper is the right first step, because you cannot control spending you cannot see. But sooner or later every developer with a public-facing AI feature hits the next problem: one user, one script, or one runaway loop spends in a weekend what the rest of your users spend in a month, and you only find out from the invoice.
This tutorial builds the second half: a spending guard that sits in front of your API calls and refuses a call before it is made, when the user behind it is already over budget. It is about two hundred lines of Python, it uses only the standard library, and it stores its data in a small SQLite file you can inspect with any database tool. There is no new service to run, no dashboard to learn, and no pricing data to trust: you supply the rates from your provider's pricing page, and the guard does arithmetic on your own token counts.
One honesty note before the code. Every dollar figure in the demo below comes from made-up toy rates, labelled as such, chosen so the arithmetic is easy to follow. They are not any provider's prices and not measurements from a real app. In production you replace them with the rates on your provider's pricing page on the day you deploy — our OpenAI pricing guide and DeepSeek pricing guide show where those numbers live, but the guard itself never guesses a rate.
The problem with after-the-fact controls
Most AI cost controls work after the money is spent. A dashboard total, a daily email summary, or a log file all answer the same question: what did we spend? None of them can answer the question that matters when a bill is running away: should this next call be allowed? By the time a human reads an alert, the calls that caused it have already been billed. Our bill-diagnosis checklist walks through the usual causes of a surprise invoice — a retry loop, a growing system prompt, an agent loop with no step cap — and nearly all of them share one property: they are patterns that repeat a cheap call many times. A guard aimed at the individual user is the cheapest way to bound that repetition, because it turns an unbounded loop into a bounded one: at most one budget's worth of calls, and then a refusal your code can handle gracefully.
There are three design decisions that make this guard practical instead of annoying. First, it checks before the call, using an estimate, and records the real token counts after the call. Estimates are always wrong in detail, so the guard never treats an estimate as a bill — it treats it as a gate, and the recorded actuals are what the next check trusts. Second, budgets are monthly windows that reset on the first of the month, which matches how most API billing works and means you never have to write a reset job. Third, a user with no budget set is never blocked. That last rule matters more than it looks: it means you can roll the guard out for one heavy user, one expensive feature, or one trial tier, without changing behaviour for everyone else on day one.
What the guard does, in plain terms
A budget is a pair: a user id and a scope. The scope is any label you already use for a feature — "chat", "summariser", "agent", or anything else. You set a monthly dollar limit for a user on a scope, and the guard tracks what that user has spent on that scope this month. There is one special scope, "*", which means a user-wide budget: it is checked on every call for that user, whatever the feature is. So a user can have a tight budget on an expensive agent feature and a looser budget overall, and the guard enforces whichever one runs out first.
Every call goes through three steps. Before the call, check() adds the estimated cost to what the user has already spent this month and refuses the call if the total would pass the budget, raising a BudgetExceeded exception your code catches and turns into a friendly message, a degraded answer, or a queue-and-retry-later — your choice, not the guard's. If spending has passed eighty percent of a budget, check() returns a warning string instead of blocking, so you can warn the user or alert yourself while the call still goes ahead. After the call, record() stores the real input and output token counts and the real cost, computed from your rates. The helper run_guarded() ties the three steps together: give it a function that makes the call and returns the real token counts, and it checks, runs, and records in one place.
Two things the guard deliberately does not do. It does not decide what a call should cost — that is your rates table, and if a rate is wrong, only the dollar column is wrong while the token counts stay right. And it does not store prompt or answer text, user content, or anything but user ids, scope labels, token counts, costs, and timestamps. A budget database should be boring enough to keep forever and safe enough to show a teammate.
The code: a guard you can drop in
Save this as budget_guard.py next to your app. It needs Python 3.8 or newer and nothing to install — sqlite3 and datetime are in the standard library. By default it stores budgets and usage in a file called budgets.db in the working directory; pass ":memory:" instead if you want a throwaway database for tests, as the demo does.
"""budget_guard.py -- Per-user / per-feature AI spending guard.
Standard library only (sqlite3). No API calls are made by this module:
it only tracks budgets and costs. You supply the rates from your
provider's pricing page, and you pass in the token counts your API
client reports after each call.
Core idea, in three steps for every call:
1. guard.check(...) -- before the call, refuse it if the user is
already over budget (or the estimate alone would push them over).
2. Make the API call.
3. guard.record(...) -- after the call, store the real token counts
and cost, so the next check uses actuals, not guesses.
"""
import sqlite3
from datetime import datetime, timezone
class BudgetExceeded(Exception):
"""Raised when a call would push a user past their budget."""
def _month_start_utc():
"""Timestamp (seconds) of the first day of the current month, UTC."""
now = datetime.now(timezone.utc)
return datetime(now.year, now.month, 1, tzinfo=timezone.utc).timestamp()
class BudgetGuard:
"""Monthly spending budgets, per user and per feature scope.
A budget is identified by (user_id, scope). Scope is any label you
already use, e.g. "chat", "summariser", "agent". Two scopes are
special:
"*" -- a user-wide budget, checked on every call for that user,
whatever the feature scope is.
A user with no budget set is never blocked: check() simply passes.
"""
def __init__(self, db_path="budgets.db"):
self.db = sqlite3.connect(db_path)
self.db.execute(
"CREATE TABLE IF NOT EXISTS budgets ("
" user_id TEXT, scope TEXT, limit_usd REAL,"
" PRIMARY KEY (user_id, scope))"
)
self.db.execute(
"CREATE TABLE IF NOT EXISTS usage ("
" id INTEGER PRIMARY KEY AUTOINCREMENT,"
" user_id TEXT, scope TEXT,"
" input_tokens INTEGER, output_tokens INTEGER,"
" cost_usd REAL, ts REAL)"
)
self.db.commit()
# -- budgets ----------------------------------------------------
def set_budget(self, user_id, scope, limit_usd):
"""Set (or replace) a monthly budget in US dollars."""
self.db.execute(
"INSERT OR REPLACE INTO budgets VALUES (?, ?, ?)",
(user_id, scope, float(limit_usd)),
)
self.db.commit()
def get_budget(self, user_id, scope):
row = self.db.execute(
"SELECT limit_usd FROM budgets WHERE user_id=? AND scope=?",
(user_id, scope),
).fetchone()
return row[0] if row else None
# -- spending ---------------------------------------------------
def spent(self, user_id, scope=None):
"""Dollars spent this month. scope=None means all scopes."""
if scope is None:
row = self.db.execute(
"SELECT COALESCE(SUM(cost_usd), 0) FROM usage"
" WHERE user_id=? AND ts>=?",
(user_id, _month_start_utc()),
).fetchone()
else:
row = self.db.execute(
"SELECT COALESCE(SUM(cost_usd), 0) FROM usage"
" WHERE user_id=? AND scope=? AND ts>=?",
(user_id, scope, _month_start_utc()),
).fetchone()
return row[0]
def remaining(self, user_id, scope):
"""Dollars left on a budget, or None if no budget is set."""
limit = self.get_budget(user_id, scope)
if limit is None:
return None
base = self.spent(user_id) if scope == "*" else self.spent(user_id, scope)
return limit - base
# -- pricing ----------------------------------------------------
@staticmethod
def estimate_cost(input_tokens, output_tokens,
price_in_per_million, price_out_per_million):
"""Cost of a call from token counts and per-million rates."""
return (
input_tokens * price_in_per_million
+ output_tokens * price_out_per_million
) / 1_000_000
# -- the guard ----------------------------------------------------
def check(self, user_id, scope, estimated_cost, warn_at=0.8):
"""Pre-call check. Returns a warning string, or None.
Raises BudgetExceeded if the estimated cost would push the user
past any budget that applies: the feature scope budget and/or
the user-wide ("*") budget. Users with no budget are never
blocked. warn_at=0.8 means: warn once spending passes 80% of
a budget (the call still goes ahead).
"""
for budget_scope in (scope, "*"):
limit = self.get_budget(user_id, budget_scope)
if limit is None:
continue
base = (self.spent(user_id) if budget_scope == "*"
else self.spent(user_id, scope))
if base + estimated_cost > limit:
raise BudgetExceeded(
f"user {user_id!r} scope {budget_scope!r}: "
f"spent ${base:.4f} + estimated ${estimated_cost:.4f} "
f"exceeds monthly budget ${limit:.4f}"
)
if base >= warn_at * limit:
return (
f"warning: user {user_id!r} has used "
f"${base:.4f} of ${limit:.4f} "
f"({100 * base / limit:.0f}%) on {budget_scope!r} "
f"this month"
)
return None
def record(self, user_id, scope, input_tokens, output_tokens,
price_in_per_million, price_out_per_million):
"""Post-call true-up. Stores real token counts and cost.
Returns the cost of this call in dollars.
"""
cost = self.estimate_cost(input_tokens, output_tokens,
price_in_per_million,
price_out_per_million)
self.db.execute(
"INSERT INTO usage (user_id, scope, input_tokens,"
" output_tokens, cost_usd, ts) VALUES (?, ?, ?, ?, ?, ?)",
(user_id, scope, int(input_tokens), int(output_tokens),
cost, datetime.now(timezone.utc).timestamp()),
)
self.db.commit()
return cost
def run_guarded(self, user_id, scope, call_fn,
est_input_tokens, est_output_tokens,
price_in_per_million, price_out_per_million,
warn_at=0.8):
"""Run call_fn() under the guard; returns (result, cost, warn
).
call_fn must return (result, actual_input_tokens,
actual_output_tokens). The estimate only gates the call; the
recorded cost always comes from the actual token counts.
"""
est = self.estimate_cost(est_input_tokens, est_output_tokens,
price_in_per_million,
price_out_per_million)
warn = self.check(user_id, scope, est, warn_at=warn_at)
result, actual_in, actual_out = call_fn()
cost = self.record(user_id, scope, actual_in, actual_out,
price_in_per_million,
price_out_per_million)
return result, cost, warn
# -- demo (fake client, no network, no API key)
# ------------------------------------------------------------
class FakeClient:
"""Stands in for a real API client in the demo."""
def __init__(self, input_tokens=2000, output_tokens=500):
self.input_tokens = input_tokens
self.output_tokens = output_tokens
def chat(self, _prompt):
return ("<fake answer>", self.input_tokens, self.output_tokens)
def demo():
# Toy rates for the demo ONLY -- replace these with the rates on
# your provider's pricing page before using this for real traffic.
PRICE_IN, PRICE_OUT = 1.00, 3.00 # dollars per million tokens
guard = BudgetGuard(":memory:") # in-memory DB, nothing is saved
client = FakeClient()
print("== alice, feature budget $0.012 ==")
guard.set_budget("alice", "chat", 0.012)
for i in range(1, 5):
try:
_res, cost, warn = guard.run_guarded(
"alice", "chat", lambda: client.chat("hello"),
2000, 500, PRICE_IN, PRICE_OUT)
print(f"call {i}: ok, cost ${cost:.4f}, "
f"spent ${guard.spent('alice', 'chat'):.4f}"
+ (f" -- {warn}" if warn else ""))
except BudgetExceeded as exc:
print(f"call {i}: BLOCKED -- {exc}")
print("== bob, feature budget $0.004 (warning demo) ==")
guard.set_budget("bob", "chat", 0.004)
_res, cost, _warn = guard.run_guarded(
"bob", "chat", lambda: client.chat("hello"),
2000, 500, PRICE_IN, PRICE_OUT)
print(f"call 1: ok, cost ${cost:.4f}, "
f"spent ${guard.spent('bob', 'chat'):.4f}")
# A check with a zero estimate passes the budget but still reports
# the warning, because bob is already past 80% of his budget.
warn = guard.check("bob", "chat", 0.0)
print(f"pre-call check for call 2: {warn}")
if __name__ == "__main__":
demo()
How to run it
There is nothing to install and no API key to configure. Save the file, then run it:
python3 budget_guard.pyBecause the demo uses a fake client and an in-memory database, it runs anywhere and touches nothing: no network calls, no files written, no money spent. Its output is the guard doing exactly the two things it exists for — letting calls through while there is budget left, and refusing the call that would cross the line:
== alice, feature budget $0.012 ==
call 1: ok, cost $0.0035, spent $0.0035
call 2: ok, cost $0.0035, spent $0.0070
call 3: ok, cost $0.0035, spent $0.0105
call 4: BLOCKED -- user 'alice' scope 'chat': spent $0.0105 + estimated $0.0035 exceeds monthly budget $0.0120
== bob, feature budget $0.004 (warning demo) ==
call 1: ok, cost $0.0035, spent $0.0035
pre-call check for call 2: warning: user 'bob' has used $0.0035 of $0.0040 (88%) on 'chat' this monthRead the alice sequence carefully, because it shows the estimate rule working as designed. The fourth call is refused before it runs, not after it bills: at that point alice has spent $0.0105, and the next call's estimate of $0.0035 would take her to $0.0140, past her $0.0120 budget. The guard does not need to know what the call would actually have cost — the estimate alone is enough to refuse it, and refusing early is the whole point. The bob sequence shows the warning side: at 88% of his budget, a check still allows the call but hands you a warning string you can show the user or send to your own alert channel.
One detail worth noticing: the guard's monthly window is computed from the current date in UTC, so spending resets automatically on the first of the month and old rows simply fall out of the window. Nothing deletes them, which is deliberate — last month's rows are still there if you ever want to query them, but they no longer count against this month's budgets.
Wiring it to a real client
The demo's fake client stands in for whatever SDK you actually use. The only thing the guard needs from a real call is the token counts the provider already reports. With OpenAI's Python SDK, a chat completion response carries a usage object with prompt_tokens and completion_tokens, exactly the two numbers run_guarded wants back from your function. So the production version of the demo's one-liner looks like this:
def make_call():
response = client.chat.completions.create(
model="your-model-id", messages=messages)
return (response,
response.usage.prompt_tokens,
response.usage.completion_tokens)
try:
response, cost, warn = guard.run_guarded(
user_id, "chat", make_call,
est_input_tokens, est_output_tokens,
PRICE_IN, PRICE_OUT)
except BudgetExceeded:
return "You have used this month's AI budget. It resets on the 1st."Two practical notes on that wiring. First, the estimate before the call: for a chat feature you can estimate input tokens from the prompt you are about to send — most providers publish a tokenizer, or you can start with a rough character-based guess and let the post-call true-up correct it, since the recorded cost always comes from the real counts. What you must not do is skip the pre-call check for calls you cannot estimate: estimate generously instead, because an overestimate only refuses a call slightly early, while no estimate at all refuses nothing. Second, streaming: if you stream a response, the usage object arrives with the final chunk, so record after the stream finishes rather than per chunk — the same rule our logging wrapper follows, for the same reason. A guard that records per chunk would count one call many times over and block users who did nothing wrong.
When a call is blocked, catch BudgetExceeded at the edge of your feature and answer the user plainly. "You have used this month's AI budget" beats a bare error, and it beats silently degrading to a worse answer without explanation. If you sell tiers, the block is also the natural place to offer the upgrade: the user who just hit a limit is the user most likely to understand what a higher limit is worth.
How to adapt it
The file is intentionally small so that adapting it is a matter of changing one function, not redesigning it. Four adaptations cover nearly every real deployment. First, the database: SQLite is one file and one writer at a time, which is fine for a single-process app but not for several web workers. To move to Postgres or MySQL, replace the five SQL statements with your driver's equivalents — the table shapes stay the same, and nothing else in the file knows or cares which database it is. Second, the period: if your billing or your product works on daily budgets instead of monthly ones, change only _month_start_utc() to return the start of the current day; every query reads through that one function. Third, the warning: check() returns a string, and what you do with it is up to you — log it, email it, or send it to a webhook. A one-line change at the call site turns a returned warning into a Slack message the first time a user crosses eighty percent. Fourth, cached input: this version prices two token kinds, input and output. If your provider bills cached input at a different rate, add a third rate and a third token count to record() and estimate_cost() — the same split our logging wrapper records, for the same reason. Do not add a rate you cannot verify on a pricing page: a guard with a guessed rate enforces a fictional budget, which is worse than no guard at all.
One more adaptation deserves a warning of its own. It is tempting to give the guard a back door — an admin override, a "just this once" flag, a second code path that skips the check for internal users. Every back door becomes the path all traffic eventually takes, usually via a forgotten debug flag. If some callers genuinely need no limit, give them no budget instead: the guard already treats "no budget set" as "never blocked", so an unlimited internal job is just a caller with no budget row, visible in the same table as everything else, instead of a secret exception in code.
FAQ
Do I need exact token counts before the call?
No. The pre-call check runs on an estimate, and the estimate only needs to be in the right neighbourhood. If it is too high, a call gets refused slightly early; if it is too low, the post-call true-up still records the real cost, so the next check is computed from real spending. The failure mode to avoid is skipping the check entirely for calls you cannot estimate — estimate generously and let the true-up correct you.
What happens when the month changes?
Nothing you need to do. The spending window starts on the first of the current month in UTC, so on the first, last month's rows fall outside the window and every user starts at zero again. Old rows are kept, not deleted, so you can still query last month's totals — they just no longer count against any budget.
Can I use this for something that is not a user?
Yes. The guard does not know what a user id is — it is just a label. Use a job name for a nightly pipeline, a team name for a shared key, or a customer id for a paying account. A background job with a monthly budget behaves exactly like a user with one: it runs until the budget is gone, and then it stops instead of running up a bill while nobody watches.
Logging showed you where the money goes; this guard decides where it stops. Build it this week, set a budget for your heaviest user or your most expensive feature first, and let the rest of your traffic run unbudgeted until the numbers tell you where the next budget belongs. The goal is not to spend less in general — it is to make sure no single user, job, or loop can ever spend without a limit again.