update

AzA march 2026/aza_rate_limit.py

@@ -0,0 +1,52 @@
+import time
+from dataclasses import dataclass
+from typing import Dict
+
+from fastapi import HTTPException, status
+
+
+@dataclass
+class Bucket:
+    tokens: float
+    last_ts: float
+
+
+class TokenBucketRateLimiter:
+    """
+    Simple in-memory token bucket limiter.
+    capacity: max burst
+    refill_rate: tokens per second
+    """
+    def __init__(self, capacity: int, refill_rate: float):
+        self.capacity = float(capacity)
+        self.refill_rate = float(refill_rate)
+        self._buckets: Dict[str, Bucket] = {}
+
+    def _get_bucket(self, key: str) -> Bucket:
+        now = time.time()
+        b = self._buckets.get(key)
+        if not b:
+            b = Bucket(tokens=self.capacity, last_ts=now)
+            self._buckets[key] = b
+            return b
+
+        elapsed = now - b.last_ts
+        b.tokens = min(self.capacity, b.tokens + elapsed * self.refill_rate)
+        b.last_ts = now
+        return b
+
+    def consume(self, key: str, cost: float = 1.0) -> None:
+        b = self._get_bucket(key)
+        if b.tokens < cost:
+            raise HTTPException(
+                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
+                detail="Rate limit exceeded",
+            )
+        b.tokens -= cost
+
+
+# Defaults:
+# - Per token: ~30 req/min (0.5 tokens/sec), burst 10
+# - Per IP:    ~60 req/min (1.0 tokens/sec), burst 20
+default_token_limiter = TokenBucketRateLimiter(capacity=10, refill_rate=0.5)
+default_ip_limiter = TokenBucketRateLimiter(capacity=20, refill_rate=1.0)