MCP server source
This commit is contained in:
root
@@ -0,0 +1,8 @@
|
||||
FROM python:3.13-slim
|
||||
WORKDIR /app
|
||||
COPY requirements.txt .
|
||||
RUN pip install --no-cache-dir -r requirements.txt
|
||||
COPY server.py ratelimit.py ./
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
EXPOSE 8899
|
||||
CMD ["python3", "server.py"]
|
||||
+232
@@ -0,0 +1,232 @@
|
||||
"""
|
||||
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
|
||||
@@ -0,0 +1,5 @@
|
||||
fastmcp>=2.0.0
|
||||
httpx>=0.27.0
|
||||
trafilatura>=1.12.0
|
||||
python-dotenv>=1.0.0
|
||||
firecrawl-py>=1.0.0
|
||||
Reference in New Issue
Block a user