Skip to content
UsageRail
Frameworks / Flask

Flask extension

The Flask extension checks each request's API key, rate limit and quota in a before_request hook, before your view runs, and counts the request when the server has finished sending the response.

bash
pip install usagerail
app.py
import os

from flask import Flask, g, request
from usagerail.flask import UsageRail

app = Flask(__name__)
UsageRail(app, app_key=os.environ["USAGERAIL_APP_KEY"], auth="api-key")


@app.post("/v1/generate")
def generate():
    result = run_model(request.json)
    g.usagerail.track(result.tokens)  # optional: count tokens
    return result.body

With an application factory, create the extension once and call init_app in the factory:

python
rail = UsageRail(app_key=os.environ["USAGERAIL_APP_KEY"], auth="api-key")


def create_app():
    app = Flask(__name__)
    rail.init_app(app)
    return app

The SDK needs Python 3.10 or newer and has no dependencies of its own. The key comes from the Authorization: Bearer header, or from x-api-key when the request has no bearer token.

Options

UsageRail() takes the same options as the FastAPI middleware: app_key, auth, consumer, on_reject, on_key_store_unavailable, stale_key_snapshot, on_quota_unavailable, snapshot, flush_on_exit and logger. See the table on the FastAPI page. Here consumer and on_reject get Flask's request, and on_reject returns a Flask response.

Count units instead of requests

In apps that count units, call g.usagerail.track(n) in the view with a whole number from 0 to 10,000,000. Route weights and reservations work as described under quotas.

Use your own authentication

Register your own before_request hook first, so the caller is known when UsageRail's hook runs:

python
app.before_request(load_user)  # your own hook, sets g.user

UsageRail(
    app,
    app_key=os.environ["USAGERAIL_APP_KEY"],
    auth="custom",
    consumer=lambda request: g.user.account_id if g.get("user") else None,
)

Returning None makes the request anonymous: it's recorded, but not rate limited, not counted against a quota and never billed.

Rejected requests

Refused requests get the statuses and JSON body listed on the Express page. Pass on_reject to send your own response:

python
from flask import jsonify


def reject(rejection, request):
    response = jsonify(code=rejection.error)
    response.status_code = rejection.status
    response.headers.update(rejection.headers)
    return response

What gets counted

  • Each request is counted under the URL rule Flask matched, such as /v1/users/<int:user_id>, with its method and outcome.
  • The status is the one Flask finally sent, after all your after_request hooks have run.
  • A streamed response counts once its last chunk has been sent.
  • 2xx and 3xx responses are billable; 4xx and 5xx responses and refused requests aren't. The redirect Flask sends to add a trailing slash is recorded under _unmatched and billed as 1 unit.
  • A view that raises an exception all the way out is counted as a 500.
  • Flask can't tell when a client disconnects early, so requests are always counted by their status.

Running in production

Counts are sent every 60 seconds and once more when the process exits normally, which is what gunicorn does on SIGTERM. A process killed outright loses the counts it hadn't sent. With gunicorn --preload, each worker becomes its own SDK instance after the fork.