Skip to content
Draft
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion server.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,12 @@
from pathlib import Path
from typing import Any, Dict, List

from fastapi import FastAPI
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from starlette.middleware.base import BaseHTTPMiddleware

# ── Project root setup ────────────────────────────────────────────
import sys
Expand All @@ -46,12 +47,49 @@

from src.pipelines.ingest import IngestPipeline
from src.pipelines.retrieval import RetrievalPipeline
from src.config import settings


# ═══════════════════════════════════════════════════════════════════
# Log capture — collects pipeline log messages during a run
# ═══════════════════════════════════════════════════════════════════


# ═══════════════════════════════════════════════════════════════════
# Rate Limiting Middleware
# ═══════════════════════════════════════════════════════════════════

class RateLimitMiddleware(BaseHTTPMiddleware):
def __init__(self, app):
super().__init__(app)
self.rate_limit_records: Dict[str, List[float]] = {}

async def dispatch(self, request: Request, call_next):
client_ip = request.client.host if request.client else "unknown"
current_time = time.time()

if client_ip not in self.rate_limit_records:
self.rate_limit_records[client_ip] = []

# Filter timestamps older than 60 seconds
self.rate_limit_records[client_ip] = [
t for t in self.rate_limit_records[client_ip]
if current_time - t < 60
]

if len(self.rate_limit_records[client_ip]) >= settings.rate_limit:
return JSONResponse(
status_code=429,
content={
"error": "Too many requests",
"detail": f"Rate limit exceeded: {settings.rate_limit} per minute"
},
headers={"Retry-After": "60"}
)

self.rate_limit_records[client_ip].append(current_time)
return await call_next(request)

class StepCapture(logging.Handler):
"""Captures log records into a list of step dicts."""

Expand Down Expand Up @@ -123,6 +161,8 @@ async def lifespan(app: FastAPI):

app = FastAPI(title="Xmem Test Frontend", lifespan=lifespan)

app.add_middleware(RateLimitMiddleware)

app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
Expand Down