AI Budget Tracker and Rate Limiter

AI Budget Tracker and Rate Limiter

Problem Statement

Teams deploying LLM-powered features quickly discover that token costs are easy to underestimate. A single runaway loop, a misconfigured retry policy, or a chatty internal tool can burn through hundreds of dollars in API credits before anyone notices. Worse, because most OpenAI SDK calls succeed silently, there's no natural "circuit breaker" — the application happily keeps spending until the billing dashboard screams.

This script solves two related problems in one place. First, it enforces a daily spend ceiling by estimating the cost of each request before it is sent, and refusing to proceed if the projected total would exceed the budget. Second, it acts as a rate limiter using Redis so that multiple workers or containers share a single token bucket, preventing bursts of parallel requests from overwhelming the API or your wallet. It also emits threshold alerts (e.g. at 80% of budget) so operators get a warning before the hard stop kicks in.

How It Works

The script is structured as a small library with a CLI wrapper. At the core is a BudgetTracker class that maintains a Redis key per UTC day (ai_budget:YYYY-MM-DD) holding the accumulated cost in micro-dollars (integers avoid floating-point drift). Before every OpenAI call, we estimate the cost from the model's per-1K-token pricing and the prompt length, then run a Lua script on Redis that atomically checks-and-increments the counter. If the increment would exceed the daily limit, the script aborts with a BudgetExceededError. A second Redis key implements a fixed-window rate limiter (ai_budget:rate:<minute>) using INCR + EXPIRE.

Design decisions worth noting: (1) costs are stored in micro-dollars as integers to avoid rounding errors; (2) the check-and-increment is atomic via Lua, so parallel workers can't both pass the budget check simultaneously; (3) the CLI is a thin wrapper that demonstrates a full request lifecycle — estimate, check, call, reconcile actual cost from the API response — so it can be dropped into an existing codebase with minimal changes.

Full Script

#!/usr/bin/env python3
"""
AI Budget Tracker and Rate Limiter

Enforces a daily USD spend cap and a per-minute request rate limit for
OpenAI API calls using Redis as shared state. Estimates cost before each
call, refuses requests that would exceed the budget, and reconciles the
actual cost from the API response afterward.

Usage:
    python3 ai_budget.py --daily-limit 5.00 --alert-threshold 80 --model gpt-4o --prompt "Hello"
"""

from __future__ import annotations

import argparse
import os
import sys
import time
import logging
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Final

import redis
from openai import OpenAI, OpenAIError

# ---------------------------------------------------------------------------
# Configuration (env vars with sensible defaults)
# ---------------------------------------------------------------------------
REDIS_URL: Final[str] = os.environ.get("REDIS_URL", "redis://localhost:6379/0")
OPENAI_API_KEY: Final[str | None] = os.environ.get("OPENAI_API_KEY")
DEFAULT_DAILY_LIMIT: Final[float] = float(os.environ.get("AI_DAILY_LIMIT", "5.00"))
DEFAULT_ALERT_THRESHOLD: Final[int] = int(os.environ.get("AI_ALERT_THRESHOLD", "80"))
DEFAULT_RATE_LIMIT: Final[int] = int(os.environ.get("AI_RATE_LIMIT_PER_MIN", "60"))
MICROS: Final[int] = 1_000_000  # micro-dollars per dollar

# Per-1K-token pricing in USD (input, output). Update as OpenAI changes prices.
MODEL_PRICING: Final[dict[str, tuple[float, float]]] = {
    "gpt-4o": (0.005, 0.015),
    "gpt-4o-mini": (0.00015, 0.0006),
    "gpt-4-turbo": (0.01, 0.03),
    "gpt-3.5-turbo": (0.0005, 0.0015),
}

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
log = logging.getLogger("ai_budget")


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------
class BudgetExceededError(RuntimeError):
    """Raised when a request would push spend past the daily limit."""


class RateLimitError(RuntimeError):
    """Raised when the per-minute request quota is exhausted."""


# ---------------------------------------------------------------------------
# Redis-backed budget + rate limiter
# ---------------------------------------------------------------------------
class BudgetTracker:
    # Lua script: atomically check-and-increment the daily counter.
    # KEYS[1] = daily key, ARGV[1] = cost in micro-dollars, ARGV[2] = limit in micro-dollars
    _LUA_CHECK_AND_ADD = """
    local current = tonumber(redis.call('GET', KEYS[1]) or '0')
    local add = tonumber(ARGV[1])
    local limit = tonumber(ARGV[2])
    if current + add > limit then
        return {0, current}
    end
    local new = redis.call('INCRBY', KEYS[1], add)
    redis.call('EXPIRE', KEYS[1], 172800)  -- 48h TTL, safely past UTC midnight
    return {1, new}
    """

    def __init__(
        self,
        redis_client: redis.Redis,
        daily_limit_usd: float,
        alert_threshold_pct: int = 80,
        rate_limit_per_min: int = 60,
    ) -> None:
        self.r = redis_client
        self.daily_limit_micros = int(daily_limit_usd * MICROS)
        self.alert_threshold_micros = int(
            self.daily_limit_micros * alert_threshold_pct / 100
        )
        self.rate_limit_per_min = rate_limit_per_min
        self._script = self.r.register_script(self._LUA_CHECK_AND_ADD)

    # ---- keys -------------------------------------------------------------
    @staticmethod
    def _daily_key() -> str:
        day = datetime.now(timezone.utc).strftime("%Y-%m-%d")
        return f"ai_budget:daily:{day}"

    @staticmethod
    def _rate_key() -> str:
        minute = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M")
        return f"ai_budget:rate:{minute}"

    # ---- public API -------------------------------------------------------
    def check_rate(self) -> None:
        """Fixed-window per-minute limiter. Raises RateLimitError if exceeded."""
        key = self._rate_key()
        pipe = self.r.pipeline()
        pipe.incr(key)
        pipe.expire(key, 120)
        count, _ = pipe.execute()
        if count > self.rate_limit_per_min:
            raise RateLimitError(
                f"Rate limit hit: {count}/{self.rate_limit_per_min} requests this minute"
            )

    def reserve(self, estimated_cost_usd: float) -> int:
        """Atomically reserve estimated cost. Returns new total in micro-dollars."""
        cost_micros = max(1, int(estimated_cost_usd * MICROS))
        ok, total = self._script(
            keys=[self._daily_key()],
            args=[cost_micros, self.daily_limit_micros],
        )
        if not ok:
            raise BudgetExceededError(
                f"Daily budget would be exceeded. "
                f"Current: ${total / MICROS:.4f}, "
                f"Request: ${cost_micros / MICROS:.4f}, "
                f"Limit: ${self.daily_limit_micros / MICROS:.2f}"
            )
        if total >= self.alert_threshold_micros:
            log.warning(
                "Budget alert: $%.4f of $%.2f used (%.1f%%)",
                total / MICROS,
                self.daily_limit_micros / MICROS,
                100 * total / self.daily_limit_micros,
            )
        return total

    def reconcile(self, estimated_usd: float, actual_usd: float) -> None:
        """Adjust the counter by the delta between estimate and actual."""
        delta_micros = int((actual_usd - estimated_usd) * MICROS)
        if delta_micros != 0:
            self.r.incrby(self._daily_key(), delta_micros)

    def total_spent_usd(self) -> float:
        raw = self.r.get(self._daily_key()) or 0
        return int(raw) / MICROS


# ---------------------------------------------------------------------------
# Cost estimation helpers
# ---------------------------------------------------------------------------
def estimate_prompt_tokens(text: str) -> int:
    """Rough heuristic: ~4 characters per token for English text."""
    return max(1, len(text) // 4)


def estimate_cost_usd(model: str, prompt: str, max_output_tokens: int) -> float:
    if model not in MODEL_PRICING:
        raise ValueError(
            f"Unknown model {model!r}. Known: {', '.join(MODEL_PRICING)}"
        )
    in_price, out_price = MODEL_PRICING[model]
    in_tokens = estimate_prompt_tokens(prompt)
    # Assume worst case: full max_output_tokens consumed
    return (in_tokens / 1000) * in_price + (max_output_tokens / 1000) * out_price


def actual_cost_usd(model: str, usage) -> float:
    in_price, out_price = MODEL_PRICING[model]
    return (
        usage.prompt_tokens / 1000 * in_price
        + usage.completion_tokens / 1000 * out_price
    )


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="AI budget tracker and rate limiter")
    p.add_argument("--daily-limit", type=float, default=DEFAULT_DAILY_LIMIT,
                   help="Daily USD spend cap (default: %(default)s)")
    p.add_argument("--alert-threshold", type=int, default=DEFAULT_ALERT_THRESHOLD,
                   help="Percent of budget at which to warn (default: %(default)s)")
    p.add_argument("--rate-limit", type=int, default=DEFAULT_RATE_LIMIT,
                   help="Max requests per minute (default: %(default)s)")
    p.add_argument("--model", default="gpt-4o", choices=sorted(MODEL_PRICING),
                   help="OpenAI model to use")
    p.add_argument("--prompt", required=True, help="Prompt text")
    p.add_argument("--max-output-tokens", type=int, default=256,
                   help="Max output tokens (used for pre-flight estimate)")
    p.add_argument("--dry-run", action="store_true",
                   help="Estimate cost and check budget without calling OpenAI")
    return p.parse_args()


def main() -> int:
    args = parse_args()

    try:
        r = redis.Redis.from_url(REDIS_URL, decode_responses=False)
        r.ping()
    except redis.RedisError as exc:
        log.error("Cannot reach Redis at %s: %s", REDIS_URL, exc)
        return 2

    tracker = BudgetTracker(
        redis_client=r,
        daily_limit_usd=args.daily_limit,
        alert_threshold_pct=args.alert_threshold,
        rate_limit_per_min=args.rate_limit,
    )

    # 1. Rate limit first — cheapest possible rejection.
    try:
        tracker.check_rate()
    except RateLimitError as exc:
        log.error("%s", exc)
        return 3

    # 2. Pre-flight cost estimate + atomic reservation.
    estimated = estimate_cost_usd(args.model, args.prompt, args.max_output_tokens)
    log.info("Estimated cost for this call: $%.6f", estimated)
    try:
        total = tracker.reserve(estimated)
    except BudgetExceededError as exc:
        log.error("%s", exc)
        return 4
    log.info("Reserved. Daily total now: $%.6f", total / MICROS)

    if args.dry_run:
        log.info("Dry-run complete; no API call made.")
        return 0

    # 3. Call OpenAI and reconcile the actual cost.
    if not OPENAI_API_KEY:
        log.error("OPENAI_API_KEY is not set")
        return 5

    client = OpenAI(api_key=OPENAI_API_KEY)
    try:
        resp = client.chat.completions.create(
            model=args.model,
            messages=[{"role": "user", "content": args.prompt}],
            max_tokens=args.max_output_tokens,
        )
    except OpenAIError as exc:
        # Refund the reservation if the call failed — we didn't consume tokens.
        tracker.reconcile(estimated, 0.0)
        log.error("OpenAI request failed: %s", exc)
        return 6

    actual = actual_cost_usd(args.model, resp.usage)
    tracker.reconcile(estimated, actual)
    log.info("Actual cost: $%.6f | Daily total: $%.6f",
             actual, tracker.total_spent_usd())
    print(resp.choices[0].message.content)
    return 0


if __name__ == "__main__":
    sys.exit(main())

Configuration

Keep MODEL_PRICING in sync with OpenAI's published rates — the script trusts these constants for estimating cost. If a model isn't listed, the script fails fast with a ValueError rather than guessing.

Usage Examples

1. A normal call:

$ export OPENAI_API_KEY=sk-...
$ python3 ai_budget.py --daily-limit 5.00 --alert-threshold 80 --model gpt-4o \
    --prompt "Summarize the CAP theorem in one sentence."
2024-06-01 12:00:01 [INFO] Estimated cost for this call: $0.005125
2024-06-01 12:00:01 [INFO] Reserved. Daily total now: $0.005125
2024-06-01 12:00:02 [INFO] Actual cost: $0.004310 | Daily total: $0.004310
CAP theorem states that a distributed system can simultaneously provide at most two of consistency, availability, and partition tolerance.

2. Hitting the alert threshold:

$ python3 ai_budget.py --daily-limit 1.00 --alert-threshold 50 --model gpt-4-turbo \
    --prompt "Write a haiku about Redis."
2024-06-01 12:05:11 [INFO] Estimated cost for this call: $0.017750
2024-06-01 12:05:11 [WARNING] Budget alert: $0.5120 of $1.00 used (51.2%)
2024-06-01 12:05:13 [INFO] Actual cost: $0.015200 | Daily total: $0.5272

3. Budget rejection (exit code 4):

$ python3 ai_budget.py --daily-limit 0.01 --model gpt-4o --prompt "Hello"
2024-06-01 12:10:00 [INFO] Estimated cost for this call: $0.003840
2024-06-01 12:10:00 [ERROR] Daily budget would be exceeded. Current: $0.008000, Request: $0.003840, Limit: $0.01
$ echo $?
4

Error Handling

• Redis unreachable → redis.RedisError caught at startup, logs the failure, exits with code 2 before burning any API tokens.
• Rate limit exceeded → RateLimitError, exit code 3. The counter still increments so abusers don't get free retries.
• Budget exceeded → the Lua script returns {0, current} without incrementing, so a rejected request doesn't corrupt the counter. Exit code 4.
• Missing API key → detected before the OpenAI client is constructed; exit code 5. The reservation is not refunded automatically here, so you may want to add a tracker.reconcile(estimated, 0.0) if you care about that edge case.
• OpenAI API failure → caught as OpenAIError; the reservation is refunded via reconcile(estimated, 0.0) so transient failures don't eat your budget. Exit code 6.
• Unknown model → estimate_cost_usd raises ValueError early, before any Redis writes.
• Counter drift → because we reserve the worst case and then reconcile to the actual, the counter can briefly over-report. That's intentional: under-reporting would let you silently blow past the cap.

Extending Ideas

• Persistent audit log: push each (timestamp, model, estimated, actual, total) tuple to a Redis stream or Postgres table so you can build a spend dashboard and reconcile against OpenAI's billing exports.
• Per-user budgets: swap the daily key for ai_budget:daily:{date}:{user_id} and pass the user ID through the CLI or as an environment variable. The Lua script is unchanged.
• Token-bucket rate limiter: replace the fixed-window INCR+EXPIRE with a sliding-window or leaky-bucket Lua script for smoother throughput and fewer boundary bursts.
• Async/streaming support: wrap BudgetTracker in an asyncio-friendly adapter using redis.asyncio, and reconcile cost when the stream finishes rather than after the full response — useful for chat UIs where responses stream token by token.