c762564d40
- Fix 71 invalid-syntax files (class-body newline-broken assignments) - Add from/None chain to 307 B904 raise-without-from sites - Add B008 ignore to ruff.toml (already in pyproject.toml) - Noqa F401 on __init__.py re-exports (137 sites) - Noqa E402 on deferred imports (63 sites) - Bulk-add stdlib/FastAPI/project imports for F821 (127 sites) - Replace ×→x, –→-, …→... in docstrings (4093 chars) - Manual refactor of 5 SIM103/SIM116 patterns Tests: 791 passed (66 deselected due to pre-existing Redis issues in test_rag.py) Co-authored-by: opencode <opencode@rugmunch.io>
82 lines
2.7 KiB
Python
82 lines
2.7 KiB
Python
"""Cost tracking middleware - per-tenant/per-route cost enforcement.
|
|
|
|
Per RMIV5 v4.0 §T31. Tracks per-request cost in USD and enforces
|
|
budget caps per tenant. Stub implementation: records costs but
|
|
doesn't enforce caps yet.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from collections import defaultdict
|
|
from typing import Any
|
|
|
|
from fastapi import Request, Response
|
|
from starlette.middleware.base import BaseHTTPMiddleware
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class CostBuffer:
|
|
"""In-memory buffer of recent costs.
|
|
|
|
Production version writes to Redis for durability across pods.
|
|
"""
|
|
|
|
def __init__(self, max_size: int = 10_000) -> None:
|
|
self._buffer: asyncio.Queue[tuple[str, float, float, float]] = asyncio.Queue(maxsize=max_size)
|
|
self._totals: dict[str, float] = defaultdict(float)
|
|
|
|
def record(self, tenant_id: str, route: str, cost_usd: float, latency_s: float) -> None:
|
|
"""Record a single cost event."""
|
|
try: # noqa: SIM105
|
|
self._buffer.put_nowait((tenant_id, cost_usd, latency_s, time.monotonic()))
|
|
except asyncio.QueueFull:
|
|
# Drop oldest if full (acceptable - we keep running totals)
|
|
pass
|
|
self._totals[tenant_id] += cost_usd
|
|
|
|
def get_total(self, tenant_id: str) -> float:
|
|
"""Get cumulative cost for a tenant."""
|
|
return self._totals.get(tenant_id, 0.0)
|
|
|
|
def get_all_totals(self) -> dict[str, float]:
|
|
"""Snapshot of all tenant totals."""
|
|
return dict(self._totals)
|
|
|
|
|
|
class CostTrackingMiddleware(BaseHTTPMiddleware):
|
|
"""Records per-request cost in the CostBuffer.
|
|
|
|
Currently uses a simple heuristic: x402 endpoints cost $0.01-$1.00,
|
|
free endpoints cost $0. Stub does NOT enforce caps yet.
|
|
"""
|
|
|
|
def __init__(self, app, buffer: CostBuffer) -> None:
|
|
super().__init__(app)
|
|
self.buffer = buffer
|
|
|
|
async def dispatch(self, request: Request, call_next: Any) -> Response:
|
|
start = time.monotonic()
|
|
response = await call_next(request)
|
|
elapsed = time.monotonic() - start
|
|
|
|
# Simple heuristic for cost estimation
|
|
path = request.url.path
|
|
if path.startswith("/api/v1/x402/"):
|
|
cost_usd = 0.05 # typical x402 tool call
|
|
elif path.startswith("/api/v1/reports/"):
|
|
cost_usd = 5.0 # $5 report
|
|
else:
|
|
cost_usd = 0.0
|
|
|
|
# Tenant ID from header (if present)
|
|
tenant_id = request.headers.get("X-Tenant-ID", "anonymous")
|
|
|
|
if cost_usd > 0:
|
|
self.buffer.record(tenant_id, path, cost_usd, elapsed)
|
|
|
|
# Surface total cost in response header for ops visibility
|
|
response.headers["X-Cost-Total-USD"] = f"{self.buffer.get_total(tenant_id):.4f}"
|
|
return response
|