# Budget-aware RAG (/docs/sample-apps/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](https://app.tokenpolice.ai). The firewall runs in `enforce` throughout.
## What TokenPolice adds [#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.
```python
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:
```python
except tp.TokenPoliceBlockedError:
# the wall: this session's budget is exhausted — degrade gracefully
...
```
## Prerequisite — create the wall (one-time, in the dashboard) [#prerequisite--create-the-wall-one-time-in-the-dashboard]
On the Firewall page, click **New rule from template** and pick **Limit spending per session**:
it groups by `session_id`, so every login gets its own allowance. Set a USD limit and the idle
window, keep the action on **Block**, and switch the rule to **Enforce**. **Without this rule
the agent always reads "healthy" and is never blocked.** Sizing tips are in the app's README.
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](/docs/rules/templates).
## Run it [#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.
```bash
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 [#what-youll-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 [#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:
| Rule mode | Action | A wall? | Agent behaviour |
| --------- | ------- | ------- | -------------------------------------------------------- |
| Enforce | Block | **Yes** | winds down — the only enforced case |
| Enforce | Reroute | No | advisory — the session keeps running, on a cheaper model |
| Enforce | Notify | No | advisory — an alert is raised, the call still runs |
| Dry-run | any | No | advisory — the block is projected, not applied |
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](/docs/concepts/dry-run-vs-enforce) made concrete.
## Related [#related]