233 lines
10 KiB
Python
233 lines
10 KiB
Python
"""
|
|
Rate limiting, TTL caching, and retry with backoff for Super Search MCP server.
|
|
|
|
Provides:
|
|
- TokenBucket: async-aware token bucket rate limiter (per provider)
|
|
- TTLCache: dict-based cache with per-key expiration
|
|
- rate_limited_search: decorator that applies rate limiting + caching + retry
|
|
"""
|
|
|
|
import asyncio
|
|
import time
|
|
import logging
|
|
from functools import wraps
|
|
from typing import Any, Callable, Optional
|
|
|
|
logger = logging.getLogger("super-search.ratelimit")
|
|
|
|
# ── Rate limits (tokens per interval in seconds) ────────────────────────────────
|
|
RATE_LIMITS: dict[str, dict[str, float]] = {
|
|
"searxng": {"tokens": 30, "interval": 1.0}, # local, high limit
|
|
"exa": {"tokens": 10, "interval": 1.0}, # paid tier
|
|
"firecrawl": {"tokens": 5, "interval": 60.0}, # 1k/mo, be careful
|
|
"opencorporates": {"tokens": 1, "interval": 1.0}, # ~1/sec free tier
|
|
"courtlistener": {"tokens": 10, "interval": 6.0}, # ~100/min free tier
|
|
"duckduckgo": {"tokens": 5, "interval": 1.0}, # free, be respectful
|
|
"wikipedia": {"tokens": 30, "interval": 1.0}, # public API, generous
|
|
}
|
|
|
|
# ── Cache TTLs (seconds) ────────────────────────────────────────────────────────
|
|
CACHE_TTL: dict[str, int] = {
|
|
"searxng": 120, # 2 min — web results change
|
|
"exa": 300, # 5 min
|
|
"opencorporates": 3600, # 1 hour — company records don't change often
|
|
"courtlistener": 3600, # 1 hour — court cases are infrequent
|
|
"duckduckgo": 300, # 5 min — instant answers can change
|
|
"wikipedia": 3600, # 1 hour — encyclopedia articles are stable
|
|
}
|
|
|
|
# ── Exceptions considered "retryable" ───────────────────────────────────────────
|
|
RETRYABLE_STATUSES: set[int] = {429, 503, 502, 504}
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
# Token Bucket Rate Limiter
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
|
|
class TokenBucket:
|
|
"""Async-aware token bucket rate limiter.
|
|
|
|
Tokens refill at a constant rate (tokens per interval). Each acquire()
|
|
consumes one token; if none are available the caller must wait.
|
|
"""
|
|
|
|
def __init__(self, tokens: float, interval: float) -> None:
|
|
self._capacity = float(tokens)
|
|
self._tokens = float(tokens)
|
|
self._interval = float(interval)
|
|
self._refill_rate = self._capacity / self._interval
|
|
self._last_refill = time.monotonic()
|
|
self._lock = asyncio.Lock()
|
|
|
|
def _refill(self) -> None:
|
|
now = time.monotonic()
|
|
elapsed = now - self._last_refill
|
|
self._tokens = min(self._capacity, self._tokens + elapsed * self._refill_rate)
|
|
self._last_refill = now
|
|
|
|
async def acquire(self) -> None:
|
|
"""Acquire one token, waiting if necessary until one is available."""
|
|
async with self._lock:
|
|
self._refill()
|
|
if self._tokens >= 1.0:
|
|
self._tokens -= 1.0
|
|
return
|
|
# Calculate wait time until next token
|
|
wait = (1.0 - self._tokens) / self._refill_rate
|
|
logger.debug("TokenBucket: waiting %.2fs for token", wait)
|
|
# Sleep outside the lock so other consumers can also wait
|
|
await asyncio.sleep(wait)
|
|
# Retry after sleep
|
|
async with self._lock:
|
|
self._refill()
|
|
self._tokens -= 1.0
|
|
|
|
|
|
# ── Global bucket registry ──────────────────────────────────────────────────────
|
|
|
|
_buckets: dict[str, TokenBucket] = {}
|
|
|
|
def get_bucket(provider: str) -> TokenBucket:
|
|
"""Get or create a token bucket for a provider."""
|
|
if provider not in _buckets:
|
|
limits = RATE_LIMITS.get(provider, {"tokens": 5, "interval": 1.0})
|
|
_buckets[provider] = TokenBucket(
|
|
tokens=limits["tokens"],
|
|
interval=limits["interval"],
|
|
)
|
|
return _buckets[provider]
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
# TTL Cache
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
|
|
class TTLCache:
|
|
"""Simple dict-based TTL cache with per-key expiration."""
|
|
|
|
def __init__(self) -> None:
|
|
self._cache: dict[str, tuple[Any, float]] = {} # key → (value, expires_at)
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def get(self, key: str) -> Optional[Any]:
|
|
"""Return cached value if not expired, else None."""
|
|
async with self._lock:
|
|
entry = self._cache.get(key)
|
|
if entry is None:
|
|
return None
|
|
value, expires_at = entry
|
|
if time.monotonic() >= expires_at:
|
|
del self._cache[key]
|
|
return None
|
|
return value
|
|
|
|
async def set(self, key: str, value: Any, ttl: int) -> None:
|
|
"""Store a value with the given TTL in seconds."""
|
|
async with self._lock:
|
|
self._cache[key] = (value, time.monotonic() + ttl)
|
|
|
|
async def invalidate(self, key: str) -> None:
|
|
"""Remove an entry from the cache."""
|
|
async with self._lock:
|
|
self._cache.pop(key, None)
|
|
|
|
|
|
# ── Global cache instance ───────────────────────────────────────────────────────
|
|
|
|
_cache = TTLCache()
|
|
|
|
|
|
def cache_key(provider: str, query: str, limit: int) -> str:
|
|
"""Build a deterministic cache key for a search."""
|
|
return f"{provider}:{query.strip().lower()}:{limit}"
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
# Retry with exponential backoff
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
|
|
import httpx # for HTTPStatusError type hint
|
|
|
|
async def _retry_with_backoff(
|
|
coro_factory: Callable[[], Any],
|
|
provider: str,
|
|
max_retries: int = 3,
|
|
base_delay: float = 1.0,
|
|
) -> Any:
|
|
"""Execute a coroutine with retry on 429/5xx, using exponential backoff.
|
|
|
|
Respects Retry-After header when present.
|
|
"""
|
|
for attempt in range(max_retries + 1):
|
|
try:
|
|
result = await coro_factory()
|
|
return result
|
|
except httpx.HTTPStatusError as e:
|
|
status = e.response.status_code
|
|
if status not in RETRYABLE_STATUSES or attempt == max_retries:
|
|
raise
|
|
|
|
# Check for Retry-After header
|
|
retry_after = e.response.headers.get("Retry-After", "")
|
|
if retry_after and retry_after.isdigit():
|
|
delay = float(retry_after)
|
|
else:
|
|
delay = base_delay * (2 ** attempt)
|
|
|
|
logger.warning(
|
|
"Provider %s returned %d (attempt %d/%d). Retrying in %.1fs…",
|
|
provider, status, attempt + 1, max_retries, delay,
|
|
)
|
|
await asyncio.sleep(delay)
|
|
|
|
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
# Rate-limited + cached search wrapper
|
|
# ══════════════════════════════════════════════════════════════════════════════════
|
|
|
|
def rate_limited_search(provider: str):
|
|
"""Decorator: wraps a search coroutine with rate limiting + TTL caching.
|
|
|
|
Usage:
|
|
@rate_limited_search("opencorporates")
|
|
async def _opencorporates_search(query, limit=5) -> list[dict]:
|
|
...
|
|
"""
|
|
|
|
def decorator(func):
|
|
@wraps(func)
|
|
async def wrapper(query: str, limit: int = 5, **kwargs):
|
|
# 1. Check cache
|
|
key = cache_key(provider, query, limit)
|
|
cached = await _cache.get(key)
|
|
if cached is not None:
|
|
logger.debug("Cache HIT for %s:%s", provider, key)
|
|
return cached
|
|
|
|
# 2. Acquire rate-limit token
|
|
bucket = get_bucket(provider)
|
|
await bucket.acquire()
|
|
|
|
# 3. Execute with retry
|
|
ttl = CACHE_TTL.get(provider, 300)
|
|
|
|
async def _call():
|
|
return await func(query, limit=limit, **kwargs)
|
|
|
|
try:
|
|
result = await _retry_with_backoff(_call, provider)
|
|
except Exception:
|
|
# Don't cache failures
|
|
raise
|
|
|
|
# 4. Cache result
|
|
if result:
|
|
await _cache.set(key, result, ttl)
|
|
logger.debug("Cache SET for %s:%s (TTL=%ds)", provider, key, ttl)
|
|
|
|
return result
|
|
|
|
return wrapper
|
|
|
|
return decorator
|