Budget-aware RAG

This is the app that shows the firewall bite end to end. A RAG chat (FastAPI + LangGraph + chromadb, folder budget_aware_rag_python/, port 6003) where the agent watches its own per-session budget: it warns you in plain language as spend runs low, winds down, and when you keep going past the limit TokenPolice blocks the call — and the app answers with a friendly "start a fresh session" message instead of crashing.

Log in with a name and ask questions about a public-domain book (The Adventures of Sherlock Holmes, downloaded and indexed locally on first run). Every login gets a stable session_id, so the whole conversation — query embeddings and chat calls alike — is billed to one session in your dashboard. The firewall runs in enforce throughout.

What TokenPolice adds

The base integration is the same as any app: one init() before the LLM imports, plus workflow wrappers that attach the user and session identity.

import token_police as tp

tp.init(
    api_key=os.environ.get("TOKEN_POLICE_API_KEY", "tp_sk_test"),
    base_url=os.environ.get("TOKEN_POLICE_URL", "http://localhost:3001"),
    firewall="enforce",
    log_errors=True,
)

# agent.py — one turn = one trace; the query embedding and the chat call
# are both billed to this session.
@tp.workflow(name="budget_aware_rag_python", kind="agent")
async def run_turn(user_id, session_id, provider, message, history, metadata):
    ...

# rag.py — the one-time index build has no session_id, so it never
# spends a user's session budget.
@tp.workflow(name="budget_aware_rag_index_build")
def _build(collection):
    ...

On top of that, the app layers a self-governing pattern: each turn the agent calls a check_budget tool, which asks TokenPolice how much of the session budget is used and hands the model back an abstract status word — never a dollar amount or a percentage. The system prompt tells the model when and how to wind down. That policy is app code, not an SDK feature; read it in the repo.

The block itself is caught explicitly:

except tp.TokenPoliceBlockedError:
    # the wall: this session's budget is exhausted — degrade gracefully
    ...

Prerequisite — create the wall (one-time, in the dashboard)

That template uses a sliding window: a session's budget resets only after the configured number of minutes with no spend, so it is an inactivity window, not a rolling "last N minutes" one — a session that keeps chatting keeps accumulating. See Rule templates.

Run it

You'll need a TokenPolice API key (tp_sk_…), an OpenAI key (always required — embeddings are always OpenAI, whatever PROVIDER you pick), and, for a non-OpenAI PROVIDER, a key for that provider too.

cd budget_aware_rag_python
cp env.example .env        # then fill in your keys
bash setup.sh              # creates .venv and installs deps incl. the TokenPolice SDK
source .venv/bin/activate
python main.py             # first run downloads the book and builds the index
# open http://localhost:6003

env.example ships pointing TOKEN_POLICE_URL at a local stack (http://localhost:3001) — set it to https://collect.tokenpolice.ai for the hosted service. Set PROVIDER to openai, anthropic, or gemini (default openai); it's read at startup and shown read-only in the UI header. The README covers the app's own tuning knobs (warning thresholds, index settings, model choice).

What you'll see

A single-page chat over the book. Ask a few questions ("Who is Irene Adler?", "What happens in The Red-Headed League?"). As spend crosses the app's warning threshold the assistant tells you it's running low — an amber notice — and offers to summarize or wrap up, with no numbers. Nearer the limit it gives a closing summary, still amber. Keep going anyway and TokenPolice blocks the call: the reply turns red and reads

We've reached this reading session's resource limit, so I can't look anything else up right now. Please start a fresh session to keep exploring the book — thanks for reading!

The app never crashes, and a fresh login starts a new session budget.

Enforced vs advisory

The agent only winds down for a budget that will actually halt the session — a rule in Enforce whose action is Block. Every other combination is advisory: spend is still tracked and the agent still checks its budget every turn, but it answers normally and sends no "running low" notice.

This app hard-codes firewall="enforce" in init(), so the SDK dial is always on and the table below turns only on the rule:

In your own app both dials matter: a rule in Enforce does nothing while the SDK is in dry-run. That's Dry-run vs enforce made concrete.

Related