11. Token-Bucket Rate Limiter for the Completions API

Problem

A lab's chat completions endpoint is shared across thousands of API keys. Without a limiter, one misbehaving or overly-aggressive integration could monopolize GPU capacity that other customers' requests are also queued on. The limiter needs to allow each key a steady sustained rate of requests while still tolerating short bursts (e.g. an agent that fires off several tool-call completions back-to-back), rather than rejecting anything above a rigid per-second cap.

Implement a per-key token-bucket limiter: each key's bucket holds up to capacity tokens, refills continuously at rate tokens per second, and each admitted call spends cost tokens (default 1). A call is admitted only if enough tokens are available.

Source: src/11_token_bucket_rate_limiter.py

class TokenBucketLimiter:
    def __init__(self, rate: float, capacity: int): ...
    def allow(self, key: str, timestamp: float, cost: int = 1) -> bool: ...

>>> limiter = TokenBucketLimiter(rate=1.0, capacity=3)
>>> limiter.allow("key1", 0.0)
True
>>> limiter.allow("key1", 0.0)
True
>>> limiter.allow("key1", 0.0)
True
>>> limiter.allow("key1", 0.0)
False
>>> limiter.allow("key1", 2.0)
True

Step-by-Step Approach

1. Store, per key, a pair (tokens_available, last_refill_timestamp) lazily — a brand-new key starts with a full bucket (capacity tokens) "refilled" as of the current call's timestamp.
2. On each allow(key, timestamp, cost) call, compute elapsed = timestamp - last_refill_timestamp and refill: tokens = min(capacity, tokens + elapsed * rate). The min caps the bucket so idle time doesn't let a key bank unlimited tokens for one giant future burst.
3. If the refilled tokens >= cost, admit the call: subtract cost, store the new (tokens, timestamp), and return True.
4. Otherwise deny the call, but still store the refilled (tokens, timestamp) — the refill that happened this call is real and must not be recomputed (double-counted) on the next call.
5. No explicit background timer or thread is needed: because refill is computed lazily from elapsed time at each call, the bucket is always correct as of "now" without any polling.

The key insight is that a token bucket never needs a background refill loop — each call lazily "catches up" the bucket by computing tokens earned since the last touch from elapsed time alone, which keeps each call O(1) regardless of how long the key has been idle.

Reference solution

class TokenBucketLimiter:
    def __init__(self, rate: float, capacity: int):
        self.rate = rate
        self.capacity = capacity
        # key -> (tokens_available, last_refill_timestamp); created lazily per key
        self._buckets: dict[str, tuple[float, float]] = {}

    def allow(self, key: str, timestamp: float, cost: int = 1) -> bool:
        # an unseen key starts with a full bucket "refilled" as of now
        tokens, last_ts = self._buckets.get(key, (float(self.capacity), timestamp))
        # lazy refill: earn tokens proportional to elapsed time, capped at capacity
        # so an idle key can't bank an unbounded future burst
        elapsed = timestamp - last_ts
        tokens = min(self.capacity, tokens + elapsed * self.rate)
        if tokens >= cost:
            tokens -= cost
            self._buckets[key] = (tokens, timestamp)
            return True
        # denied, but the refill that happened this call is still real --
        # store it so the next call doesn't double-count this elapsed window
        self._buckets[key] = (tokens, timestamp)
        return False

Key Functions & Tricks

• dict.get(key, default) — supplies a lazily-computed full bucket for a never-seen key without a separate init pass.
• elapsed * rate — converts idle time directly into earned tokens — the core of lazy (as opposed to timer-driven) refill.
• min(capacity, tokens + earned) — caps the bucket so a long-idle key can't accumulate an unbounded burst allowance.
• cost parameter — lets a single call spend more than one token, useful for weighting expensive operations differently from cheap ones.
• storing tokens on denial too — prevents re-crediting the same elapsed window twice on the next call.

How to Recognize This Pattern