AI budget control for Python + FastAPI

AI providers bill per token. Without per-user limits, a single user can exhaust your entire monthly budget – through prompt attacks, runaway loops, or just heavy legitimate use.

A token bucket rate limit maps directly onto how AI billing works. You estimate the cost of each request in tokens, deduct it from the user’s bucket, and deny requests when the bucket is empty. The bucket refills over time, giving each user a sustained allowance without sharp rate-limit cliffs.

Alternatively, you can use a fixed window or sliding window limit to enforce a hard cap on spend per user per day, week, or month. For details on different approaches, see the rate limiting algorithms reference.

Arcjet handles bucket state across all instances of your application – no Redis or external state store required.

Get started

In this example we use LangChain to create a simple AI chat server with FastAPI, and Arcjet to enforce per-user token budgets to prevent cost overruns. The same principles can be applied to any AI application, including those built with other frameworks.

Set up the environment and install dependencies (uses uv, but you can also use pip to install the Arcjet Python SDK):

# Export your Arcjet API key from https://console.arcjet.com
export ARCJET_KEY="ajkey_..."
export ARCJET_ENV=development

# Export your OpenAI API key (used by LangChain)
export OPENAI_API_KEY="sk-..."

# Install dependencies
uv add arcjet fastapi uvicorn langchain langchain-openai

Create the chat server:

import logging
import math
import os

from arcjet import Mode, arcjet, token_bucket
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from pydantic import BaseModel

app = FastAPI()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

arcjet_key = os.getenv("ARCJET_KEY")
if not arcjet_key:
    raise RuntimeError("ARCJET_KEY is required. Get one at https://console.arcjet.com")

openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
    raise RuntimeError(
        "OPENAI_API_KEY is required. Get one at https://platform.openai.com"
    )

llm = ChatOpenAI(model="gpt-4o-mini", api_key=openai_api_key)

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are a helpful assistant."),
        ("human", "{message}"),
    ]
)

chain = prompt | llm | StrOutputParser()


class ChatRequest(BaseModel):
    message: str


# Create a single Arcjet client at startup and reuse it across requests
aj = arcjet(
    key=arcjet_key,  # Get your key from https://console.arcjet.com
    rules=[
        # Token bucket rate limiting is best for AI budget control
        token_bucket(
            mode=Mode.LIVE,  # Blocks requests. Use Mode.DRY_RUN to log only
            # Track budgets per user — replace "userId" with any stable
            # identifier. Removing this falls back to IP-based rate limiting.
            characteristics=["userId"],
            refill_rate=2_000,  # Refill 2,000 tokens per interval
            interval=3_600,  # Refill every hour (in seconds)
            capacity=5_000,  # Maximum 5,000 tokens in the bucket
        ),
    ],
)


@app.post("/chat")
async def chat(request: Request, body: ChatRequest):
    # Replace with your session/auth lookup to get a stable user ID
    user_id = "user-123"

    # Estimate token cost: ~1 token per 4 characters of text (rough heuristic).
    # For accurate counts use https://github.com/openai/tiktoken
    estimate = math.ceil(len(body.message) / 4)

    # Deduct the estimated tokens from the user's budget
    decision = await aj.protect(
        request,
        requested=estimate,
        characteristics={"userId": user_id},
    )

    if decision.is_denied():
        # The token_bucket rule is the only rule configured, so the only
        # possible denial reason is RATE_LIMIT (429).
        return JSONResponse({"error": "AI usage limit exceeded"}, status_code=429)

    reply = await chain.ainvoke({"message": body.message})

    return {"reply": reply}

Then run the server:

uv run uvicorn main:app --reload

And send a message to the API endpoint:

curl -X POST http://localhost:8000/chat \
  -H "Content-Type: application/json" \
  -d '{"message": "What is the capital of France?"}'

Requests appear in your Arcjet dashboard in real time.

Configure the rate limit

characteristics: ["userId"] - Tracks the bucket per user. Replace "userId" with the characteristic that identifies a unique user in your application, such as a session token, API key, or authenticated user ID. Pass the value to aj.protect() as a named argument.

refillRate and interval - Set the sustained allowance. refillRate: 2_000, interval: "1h" gives each user 2,000 tokens per hour. Adjust to match your AI provider’s pricing and your cost targets. These are hard coded in this example, but you can also calculate them dynamically based on user subscription level or other factors. Pass the calculated values to the rule.

capacity - The maximum tokens a user can accumulate. Setting capacity: 5_000 with refillRate: 2_000 lets users burst up to 5,000 tokens if they haven’t used their allowance recently.

Token estimation

The example uses a characters / 4 heuristic (~1 token per 4 characters for common English text). This is a reasonable starting point – it avoids introducing extra dependencies and works well enough for budget enforcement where a small margin of error is acceptable.

For accurate counts, use a tokenizer:

JavaScript / TypeScript: tiktoken
Python: tiktoken
Anthropic: provides a token counting API

AI budget control for Python + FastAPI

Get started

Configure the rate limit

Token estimation

Learn more