Tiered budget enforcement with graceful downgrade for the LiteLLM gateway.
Checks each team's accumulated spend against thresholds (75% warn → 90% warn → 100% enforce). At 100%, it either downgrades the request to a cheaper model or refuses it, depending on which team is asking.
Please refer to insights.md or exploration/ directory for Data analysis.
Please refer to setup.py for initial setup as shown below.
Please refer to demo.py for a simplistic demo on how the whole pipeline should work.
Requirements: Python 3.10+, Docker.
# 1. Install dependencies
python3 -m venv venv
source venv/bin/activate
pip install -r requirements.txt
# 2. Start Postgres
docker compose up -d
# 3. See it work
python demo.pydemo.py loads the supplied 30 days of spend, computes every team's budget
position, and shows what the gateway would do with each team's next request.
It's the fastest way to understand what this does.
Run the tests:
pytest tests/ -v15 tests. They use SQLite in-memory, so no database is needed to run them.
Three things, in order:
1. Some of the bill cannot be governed at all.
5 rows in spend_30d.csv have no team attached including key-personal-mhuber,
a personal API key that spent $2,003 on the most expensive model. That's 5.5%
of the entire bill, with no owner and no budget. These rows are quarantined, not
dropped, because dropna() would make the numbers tidy and the platform wrong.
2. Where each team stands. DevAgent is at 336% of budget, there's a cost runaway in this data starting 16 November.
3. What happens to the next request from each team. Same thresholds, different outcomes:
| Team | At 100% of budget | Why |
|---|---|---|
| DigestBot, Research, Marketing | silently downgraded | Internal, batch, not a record. Nobody is harmed by a cheaper model. |
| DevAgent | downgraded, and told | An agent that doesn't know it's weaker will thrash mid-chain. |
| AdvisorChat | blocked | Customer-facing financial Q&A. Quietly answering a money question with a weaker model is a conduct risk. |
| KYC | blocked | Regulated. Substituting the model changes the evidentiary record of a compliance decision. |
The rule: silent downgrade is only permitted where the output is not customer-facing, not regulated, and not a record.
Reading order: policy.py → guard.py → worker.py → db.py → demo.py
| File | What it does |
|---|---|
policy.py |
The downgrade matrix and thresholds. The judgment lives here. |
db.py |
Tables: spend ledger, reservations, policy state, decisions, budgets, prices, quarantine |
seed.py |
Loads spend_30d.csv, quarantining rows it can't govern |
seed_config.py |
Loads budgets and model prices into the database |
worker.py |
Slow half. Background job: sums spend, computes each team's tier |
guard.py |
Fast half. Runs per request: decides allow / downgrade / block |
litellm_hook.py |
Plugs the guard into the LiteLLM gateway |
demo.py |
End-to-end run against the real data |
tests/ |
15 pytest tests (SQLite in-memory, no infra needed) |
exploration/ |
Data analysis scratch scripts used during development |
worker.py (background, ~60s) guard.py (every request)
───────────────────────────── ────────────────────────────
SUM(spend) GROUP BY team ──▶ read the precomputed tier
work out each team's tier + add in-flight reservations
write it to policy_state + add this request's estimate
│
▼
allow / downgrade / block
The SUM() over the ledger is too slow to run on every request, so the worker
does it once a minute in the background. The request path just reads one row.
Slow work in the background. Fast decisions on the request.
Spend lands in the ledger after a call finishes. So if 20 requests arrive at once, they all read the same stale number, all see room, and all get allowed — a classic check-then-act race.
So each request reserves its worst-case cost before the call: effective_spend = settled_ledger + open_reservations + this_request
Now concurrent requests can see what each other are about to spend, and the burst gets cut off partway instead of all sailing through.
test_concurrent_burst_cannot_blow_the_budget proves this: it fires 10 requests
without letting any settle, so the ledger never moves meaning the reservation
is the only thing that can stop them. Comment out the reservation write and
all 10 blow through.
Reservations bound the overshoot. They don't eliminate it see below.
Budgets and model prices live in database tables, not in code:
# Finance raises DevAgent's budget. No deploy, no code change.
docker compose exec postgres psql -U platform -d aiplatform \
-c "UPDATE team_budgets SET monthly_budget_usd = 25000 WHERE team = 'DevAgent';"
python worker.py # DevAgent flips from 'exceeded' to 'warn'But the downgrade mode (silent / signalled / hard_stop) stays in policy.py,
behind a pull request. Changing KYC from hard_stop to silent isn't a config
tweak, it's a compliance decision.
Money moves in the UI. Safety moves through review.
- Budgets are invented. The brief supplies spend but no budgets, so these are derived from each team's November baseline plus headroom. Real budgets need an owner in Finance.
- Model prices are fitted from the supplied CSV. In production these would be synced from a price book, a stale price silently corrupts every cost estimate the control depends on.
- Input tokens are estimated at ~4 characters per token. This only needs to be a safe upper bound for the reservation. Production would use the provider's tokenizer.
- The LiteLLM hook is written to the real
CustomLoggersignature, but has not been run against a live gateway, the tests exercise the guard directly.
This control is a backstop, not the primary defence against a runaway.
I tested it against the actual data. DevAgent's runaway starts on 16 November. The 100% threshold trips on 20 November four days and $5,810 later. The month ends at 358% of budget.
That isn't a bug. A budget threshold measures a total, and a total can only tell you something is wrong after the money is committed. It's a fuel gauge, not a leak detector.
What would have caught it on day one:
- A
max_tokenscap. 85% of DevAgent's cost is output tokens, and the runaway was output length growing 4.8×. A per request cap stops that at the request, not at the invoice. The hook already sets one. - Rate of change detection. "Today's spend is 2.8× normal" fires on 16 November, after ~$400 instead of $5,810.
And no budget control can ever catch key-personal-mhuber you cannot enforce
a budget against a key that has no team. That fix is upstream: enforce attribution
at key issuance, so no key exists without a team, an owner, and a budget.
Things I chose not to build for this submission, and why.
No warn/critical notification.
The tiers exist (75% warn, 90% critical) and the guard returns them, but nothing
fires, no Slack message, no webhook, no structured log event. The decision is
logged to the decisions table and the architecture puts Datadog downstream of
it, but the emit isn't wired. In production this is a day-one requirement.
No monthly budget period.
The worker sums spend across all time, there is no date filter. This works for
the 30-day demo dataset, but in production a team that spent $5k last month and
$100 this month would show $5,100 against a monthly budget. Adding a
WHERE date >= start_of_month to the worker query and a budget-period column to
team_budgets is straightforward; I scoped it out to keep the demo focused on
the enforcement logic.
settle() does not reconcile actual cost.
The parameter actual_cost exists on settle() but is unused. Reservations are
marked settled, but the real cost is never compared to the estimate. Overshoot
from pessimistic worst case estimates accumulates until the worker refreshes
from the ledger. In production the settle path would write the delta back so the
guard's effective spend calculation stays tight between worker cycles.
Worker has no scheduler.
The README says "~60s" but there is no loop, no cron job, no Celery beat. The
worker runs once when invoked. In production this would be a periodic ECS task
or a simple while True: refresh(); sleep(60) loop with health checks.
1. Ship the per request circuit breaker.
A per request cost ceiling, alongside the max_tokens cap the hook already sets.
85% of DevAgent's cost is output tokens and the runaway was output length growing
4.8×, so a cap on what a single call can generate is higher-leverage than anything
in this component, and it is a hard ceiling, which a budget can never be.
2. Ship rate of change detection. The signal that catches the regression on day one instead of day four. It must require both a relative jump and absolute materiality: Research spikes 13× relative to its own median, but its peak day is $67, a naive percentage rule would page on call for a rounding error.
3. Enforce team attribution at key issuance.
No key exists without a team, an owner, a budget, and a data class. This closes the
key-personal-mhuber hole $2,003 of spend that no amount of budget checking can
ever see, because the key has no team to check against.
4. Add a row lock on the budget check.
SELECT ... FOR UPDATE on the team's row, to close the remaining microsecond race
between reading the reservations and writing one. Reservations shrink the race from
seconds to microseconds; a lock closes it.
5. Fix the unpriced spend path. One row in the supplied data has 9,100 requests and 24M tokens but a NULL cost. Spend that happened and was never priced means the ledger silently understates a team's position — which quietly weakens every threshold in this component.
6. Make the downgrade target dynamic, not hardcoded. Each team currently falls back to a fixed model name. Providers change prices and cheaper models keep arriving, so the fallback should be "the cheapest model meeting this quality bar," resolved against the live price table, rather than a name baked into config.
7. Route by task difficulty, not just by budget. DevAgent runs 100% of its traffic on gpt 5.4, the most expensive model in the fleet. Most agent turns are simple tool calls that do not need frontier reasoning. Tiered routing, cheap model for simple turns, frontier for planning, cuts cost without waiting for a threshold to fire. That is a saving available today, not a penalty applied at 100%.
8. Flag upgrades, not just downgrades. Model prices drop and better models arrive constantly. A team that chose gpt-5.4 six months ago may now be paying more for a worse model than what is available today, and nobody rechecks. The platform already holds a price table, so it knows when the price/capability frontier moves, it should surface that: "you can switch to X for the same money and get better results."
The framing matters: this is not "spend the headroom because we have it," which just ratchets budgets upward. It is "stop overpaying for the model you are on." The same price table that powers the downgrade powers this, and it makes the platform something that serves teams rather than only policing them.
9. Move the policy state read to a cache.
If the hotpath read proves expensive under load, policy_state is a natural Redis
candidate, it is written once a minute and read on every request.