openai

azure_bootstrap.openai

SDK-agnostic Azure OpenAI / Anthropic usage tracker.

Records tokens + cost per deployment in three sliding windows (60s, 60m, 24h). Optional soft TPM cap via acquire(). Threshold-based cost alerts via check_thresholds_and_alert(). Pricing table includes both OpenAI and Anthropic defaults; apps override per-deployment via env vars or register_pricing(...).

Classes:

Name	Description
AiUsageTracker	Class facade around the module-level singleton.

Functions:

Name	Description
record_usage	Record one call. Best-effort, never raises.
record_rate_limit_event	Mark a rate-limit-hit event + fire an ERROR alert (deduped).
acquire	Soft TPM cap. No-op when AI_TPM_LIMIT is unset/0.
usage_snapshot	Returns by_deployment + totals with windowed metrics.
check_thresholds_and_alert	Compare windowed metrics against env thresholds; fire alerts on breach.
reset_state	Test-only. Refuses unless AZURE_BOOTSTRAP_ALLOW_RESET=1.

AiUsageTracker

Class facade around the module-level singleton.

Mostly here for callers that prefer to inject a tracker instance; the module-level functions are the canonical entry points.

record_usage

record_usage(deployment: str, prompt_tokens: int, completion_tokens: int) -> None

Record one call. Best-effort, never raises.

Source code in azure_bootstrap/openai/__init__.py

def record_usage(
    deployment: str,
    prompt_tokens: int,
    completion_tokens: int,
) -> None:
    """Record one call. Best-effort, never raises."""
    try:
        prompt_tokens = max(0, int(prompt_tokens))
        completion_tokens = max(0, int(completion_tokens))
        cost = _compute_cost(deployment, prompt_tokens, completion_tokens)
        now = time.monotonic()
        entry = _UsageEntry(
            ts=now,
            deployment=deployment,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            cost_usd=cost,
        )
        with _state.lock:
            _state.recent.append(entry)
            cum = _state.cumulative.setdefault(deployment, _DeploymentCumulative())
            cum.prompt_tokens += prompt_tokens
            cum.completion_tokens += completion_tokens
            cum.total_tokens += prompt_tokens + completion_tokens
            cum.cost_usd += cost
            cum.calls += 1
        bump_counter("ai.tokens.total", prompt_tokens + completion_tokens)
        bump_counter("ai.cost_usd_micros", int(cost * 1_000_000))
        bump_counter("ai.calls")
    except Exception:
        pass

record_rate_limit_event

record_rate_limit_event(deployment: str, source: str) -> None

Mark a rate-limit-hit event + fire an ERROR alert (deduped).

Source code in azure_bootstrap/openai/__init__.py

def record_rate_limit_event(deployment: str, source: str) -> None:
    """Mark a rate-limit-hit event + fire an ERROR alert (deduped)."""
    try:
        with _state.lock:
            cum = _state.cumulative.setdefault(deployment, _DeploymentCumulative())
            cum.rate_limit_events += 1
        bump_counter("ai.rate_limit_events")
        try:
            from azure_bootstrap.alerts import AlertSeverity, alert_dev_team

            alert_dev_team(
                AlertSeverity.ERROR,
                subject=f"AI rate limit triggered ({source}) for {deployment}",
                context={"deployment": deployment, "source": source},
                dedup_key=f"ai_usage.rate_limit:{deployment}:{source}",
            )
        except Exception:
            pass
    except Exception:
        pass

acquire

acquire(deployment: str, estimated_tokens: int, timeout: float | None = None) -> None

Soft TPM cap. No-op when AI_TPM_LIMIT is unset/0.

On max-wait reached: fires CRITICAL alert and lets the call through — a slow, billed call beats a stranded customer.

Source code in azure_bootstrap/openai/__init__.py

def acquire(
    deployment: str,
    estimated_tokens: int,
    timeout: float | None = None,
) -> None:
    """Soft TPM cap. No-op when ``AI_TPM_LIMIT`` is unset/0.

    On max-wait reached: fires CRITICAL alert and lets the call through —
    a slow, billed call beats a stranded customer.
    """
    try:
        limit = _tpm_limit_for(deployment)
        if limit <= 0:
            return
        max_wait = (
            timeout
            if timeout is not None
            else float(os.environ.get("AI_RATE_LIMIT_MAX_WAIT_SECONDS", "60"))
        )
        start = time.monotonic()
        recorded = False
        while True:
            now = time.monotonic()
            window_tokens, oldest = _tokens_in_window(deployment, now, _RATE_WINDOW_SECONDS)
            if window_tokens + max(0, int(estimated_tokens)) <= limit:
                return
            elapsed = now - start
            if elapsed >= max_wait:
                record_rate_limit_event(deployment, source="proactive_timeout")
                try:
                    from azure_bootstrap.alerts import AlertSeverity, alert_dev_team

                    alert_dev_team(
                        AlertSeverity.CRITICAL,
                        subject=(
                            f"AI rate limit max-wait exceeded for {deployment}; "
                            "allowing call through"
                        ),
                        context={
                            "deployment": deployment,
                            "limit_tpm": limit,
                            "elapsed_seconds": round(elapsed, 2),
                        },
                        dedup_key=f"ai_usage.max_wait_exceeded:{deployment}",
                    )
                except Exception:
                    pass
                return
            if not recorded:
                record_rate_limit_event(deployment, source="proactive")
                recorded = True
            window_clear = _RATE_WINDOW_SECONDS - (now - oldest) + 0.05
            sleep_for = max(0.1, min(window_clear, max_wait - elapsed))
            time.sleep(sleep_for)
    except Exception:
        # Governance must never break the calling path.
        return

usage_snapshot

usage_snapshot() -> dict[str, Any]

Returns by_deployment + totals with windowed metrics.

Source code in azure_bootstrap/openai/__init__.py

def usage_snapshot() -> dict[str, Any]:
    """Returns by_deployment + totals with windowed metrics."""
    _prune_old()
    now = time.monotonic()
    by_dep: dict[str, Any] = {}
    total_calls = 0
    total_tokens = 0
    total_cost = 0.0
    total_rate_limits = 0
    with _state.lock:
        deployments = list(_state.cumulative.keys())
    for dep in deployments:
        with _state.lock:
            cum = _state.cumulative[dep]
            cum_dict = {
                "prompt_tokens": cum.prompt_tokens,
                "completion_tokens": cum.completion_tokens,
                "total_tokens": cum.total_tokens,
                "cost_usd": round(cum.cost_usd, 6),
                "calls": cum.calls,
                "rate_limit_events": cum.rate_limit_events,
            }
        by_dep[dep] = {
            "cumulative": cum_dict,
            "window_60s": _window_summary(dep, now, _RATE_WINDOW_SECONDS),
            "window_60m": _window_summary(dep, now, _HOURLY_WINDOW_SECONDS),
            "window_24h": _window_summary(dep, now, _DAILY_WINDOW_SECONDS),
            "tpm_limit": _tpm_limit_for(dep),
        }
        total_calls += cum_dict["calls"]  # type: ignore[assignment]
        total_tokens += cum_dict["total_tokens"]  # type: ignore[assignment]
        total_cost += cum.cost_usd
        total_rate_limits += cum_dict["rate_limit_events"]  # type: ignore[assignment]
    return {
        "by_deployment": by_dep,
        "totals": {
            "calls": total_calls,
            "total_tokens": total_tokens,
            "cost_usd": round(total_cost, 6),
            "rate_limit_events": total_rate_limits,
        },
    }

check_thresholds_and_alert

check_thresholds_and_alert() -> dict[str, Any]

Compare windowed metrics against env thresholds; fire alerts on breach.

Designed to run every 10 minutes via APScheduler.

Source code in azure_bootstrap/openai/__init__.py

def check_thresholds_and_alert() -> dict[str, Any]:
    """Compare windowed metrics against env thresholds; fire alerts on breach.

    Designed to run every 10 minutes via APScheduler.
    """
    snap = usage_snapshot()
    fired: list[dict[str, Any]] = []
    try:
        hourly_cost_limit = float(os.environ.get("AI_COST_ALERT_HOURLY_DOLLARS", "0"))
    except ValueError:
        hourly_cost_limit = 0.0
    try:
        daily_cost_limit = float(os.environ.get("AI_COST_ALERT_DAILY_DOLLARS", "0"))
    except ValueError:
        daily_cost_limit = 0.0
    try:
        hourly_tokens_limit = int(os.environ.get("AI_HIGH_USAGE_TOKENS_HOURLY", "0"))
    except ValueError:
        hourly_tokens_limit = 0
    for dep, payload in snap["by_deployment"].items():
        hourly = payload["window_60m"]
        daily = payload["window_24h"]
        if hourly_cost_limit > 0 and hourly["cost_usd"] > hourly_cost_limit:
            key = f"ai_usage.cost_hourly:{dep}"
            if _fire_threshold_alert(
                key,
                subject=f"AI hourly cost over threshold for {dep}: ${hourly['cost_usd']}",
                context={
                    "deployment": dep,
                    "window": "60m",
                    "cost_usd": hourly["cost_usd"],
                    "threshold": hourly_cost_limit,
                },
            ):
                fired.append({"key": key, "subject": "hourly_cost", "cost_usd": hourly["cost_usd"]})
        if daily_cost_limit > 0 and daily["cost_usd"] > daily_cost_limit:
            key = f"ai_usage.cost_daily:{dep}"
            if _fire_threshold_alert(
                key,
                subject=f"AI daily cost over threshold for {dep}: ${daily['cost_usd']}",
                context={
                    "deployment": dep,
                    "window": "24h",
                    "cost_usd": daily["cost_usd"],
                    "threshold": daily_cost_limit,
                },
            ):
                fired.append({"key": key, "subject": "daily_cost", "cost_usd": daily["cost_usd"]})
        if hourly_tokens_limit > 0 and hourly["total_tokens"] > hourly_tokens_limit:
            key = f"ai_usage.tokens_hourly:{dep}"
            if _fire_threshold_alert(
                key,
                subject=f"AI hourly token use over threshold for {dep}: {hourly['total_tokens']}",
                context={
                    "deployment": dep,
                    "window": "60m",
                    "total_tokens": hourly["total_tokens"],
                    "threshold": hourly_tokens_limit,
                },
            ):
                fired.append(
                    {
                        "key": key,
                        "subject": "hourly_tokens",
                        "total_tokens": hourly["total_tokens"],
                    }
                )
    return {"snapshot": snap, "fired": fired}

reset_state

reset_state() -> None

Test-only. Refuses unless AZURE_BOOTSTRAP_ALLOW_RESET=1.

Source code in azure_bootstrap/openai/__init__.py

def reset_state() -> None:
    """Test-only. Refuses unless AZURE_BOOTSTRAP_ALLOW_RESET=1."""
    if os.environ.get("AZURE_BOOTSTRAP_ALLOW_RESET") != "1":
        raise RuntimeError("reset_state is test-only — set AZURE_BOOTSTRAP_ALLOW_RESET=1")
    with _state.lock:
        _state.recent.clear()
        _state.cumulative.clear()
        _state.last_alert_fired.clear()