server.py

import brotli
import fastapi
import requests

from bucket import Bucket
from gate import AIBudgetGate

def create_app(gate: AIBudgetGate) -> fastapi.FastAPI:
    app = fastapi.FastAPI()
    session = requests.Session()
    base_uri_to_bucket = {
        bucket.base_uri: bucket 
        for bucket in gate.buckets
    }

    @app.api_route("/{base_uri}/{path:path}", methods = ["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD", "OPTIONS", "TRACE"])
    async def api_proxy(base_uri: str, path: str, request: fastapi.Request) -> fastapi.Response:
        bucket = base_uri_to_bucket.get("/" + base_uri)
        if bucket is None:
            return fastapi.Response(status_code = 404)
        if bucket.exceeds_threshold(gate.budget_threshold):
            return fastapi.Response(status_code = 429)

        method = request.method.upper()
        url = f"{bucket.upstream_url}/{path}"
        params = request.url.query
        data = await request.body()
        headers = {k: v for k, v in request.headers.items() if k.lower() != "host"}

        response = session.request(
            method = method,
            url = url,
            params = params,
            data = data,
            headers = headers
        )

        if response.headers.get('content-encoding') == 'br':
            try:
                response._content = brotli.decompress(response.content)
            except Exception:
                pass
            response.headers['content-encoding'] = 'identity'

        if response.headers.get('content-type') == 'application/json':
            update_bucket_usage(bucket, response)

        return fastapi.Response(
            content = response.content,
            status_code = response.status_code,
            headers = response.headers
        )

    def update_bucket_usage(bucket: Bucket, response: requests.Response) -> None:
        data = response.json()
        if isinstance(data, dict):
            usage = data.get("usage")
            if isinstance(usage, dict):
                total_tokens = usage.get("total_tokens")
                if isinstance(total_tokens, int):
                    bucket.increment_usage(total_tokens)

    return app