server.py

import os
import time
import logging
from contextlib import asynccontextmanager
from pathlib import Path

import uvicorn
from fastapi import Depends, FastAPI, Request
from fastapi.exceptions import RequestValidationError
import json
import typing

from fastapi.responses import JSONResponse
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from slowapi import Limiter
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.middleware.cors import CORSMiddleware

from redactor import PIIRedactor
from schemas import ErrorResponse, RedactRequest, RedactResponse

env_file = Path(__file__).parent / ".env"
if env_file.exists():
    for line in env_file.read_text().splitlines():
        line = line.strip()
        if line and not line.startswith("#") and "=" in line:
            key, _, value = line.partition("=")
            os.environ.setdefault(key.strip(), value.strip())

DEV_MODE = os.environ.get("DEV_MODE", "false").lower() == "true"
AUTH_KEYS = [k.strip() for k in os.environ.get("AUTH_KEYS", "").split(",") if k.strip()]
RATE_LIMIT = os.environ.get("RATE_LIMIT", "60/minute")
CORS_ORIGINS = [o.strip() for o in os.environ.get("CORS_ORIGINS", "").split(",") if o.strip()]
PRETTY_JSON = os.environ.get("PRETTY_JSON", "false").lower() == "true"

logging.basicConfig(
    level=logging.DEBUG if DEV_MODE else logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger("privacy-filter")

logging.getLogger("httpx").setLevel(logging.WARNING)
logging.getLogger("httpcore").setLevel(logging.WARNING)

redactor: PIIRedactor
security = HTTPBearer(auto_error=False)
limiter = Limiter(key_func=get_remote_address, default_limits=[RATE_LIMIT])


class RequestLoggerMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        start = time.perf_counter()
        response = await call_next(request)
        elapsed_ms = (time.perf_counter() - start) * 1000
        logger.info(
            "%s %s -> %d (%.1fms)",
            request.method,
            request.url.path,
            response.status_code,
            elapsed_ms,
        )
        return response


class AuthError(Exception):
    def __init__(self, status_code: int, detail: str):
        self.status_code = status_code
        self.detail = detail


async def verify_auth(credentials: HTTPAuthorizationCredentials | None = Depends(security)):
    if not AUTH_KEYS:
        return
    if credentials is None or credentials.credentials not in AUTH_KEYS:
        raise AuthError(status_code=401, detail="Invalid or missing API key")


async def auth_error_handler(_request: Request, exc: AuthError):
    return JSONResponse(status_code=exc.status_code, content=ErrorResponse(error="auth_error", detail=exc.detail))


async def rate_limit_handler(_request: Request, _exc: RateLimitExceeded):
    return JSONResponse(status_code=429, content=ErrorResponse(error="rate_limit", detail="Too many requests"))


async def validation_error_handler(_request: Request, exc: RequestValidationError):
    return JSONResponse(status_code=422, content=ErrorResponse(error="validation_error", detail=str(exc)))


async def unhandled_error_handler(request: Request, exc: Exception):
    logger.error("Unhandled error on %s %s: %s", request.method, request.url.path, exc)
    return JSONResponse(status_code=500, content=ErrorResponse(error="internal_error"))


@asynccontextmanager
async def lifespan(_app: FastAPI):
    global redactor
    device = os.environ.get("DEVICE", "cpu")
    redactor = PIIRedactor(device=device)
    logger.info("Model loaded on %s", device)
    yield


class PrettyJSONResponse(JSONResponse):
    def render(self, content: typing.Any) -> bytes:
        return json.dumps(content, ensure_ascii=False, allow_nan=False, indent=2).encode("utf-8")


app = FastAPI(title="Privacy Filter Proxy", lifespan=lifespan, default_response_class=PrettyJSONResponse if PRETTY_JSON else JSONResponse)
app.state.limiter = limiter
app.add_exception_handler(AuthError, auth_error_handler)
app.add_exception_handler(RateLimitExceeded, rate_limit_handler)
app.add_exception_handler(RequestValidationError, validation_error_handler)
app.add_exception_handler(Exception, unhandled_error_handler)

app.add_middleware(RequestLoggerMiddleware)
if CORS_ORIGINS:
    app.add_middleware(
        CORSMiddleware,
        allow_origins=CORS_ORIGINS,
        allow_credentials=True,
        allow_methods=["GET", "POST"],
        allow_headers=["Authorization", "Content-Type"],
    )


@app.get("/health")
async def health():
    return {"status": "ok"}


@app.post("/redact", response_model=RedactResponse)
@limiter.limit(RATE_LIMIT)
async def redact(request: Request, req: RedactRequest, _=Depends(verify_auth)):
    logger.debug("Input text length: %d chars", len(req.text))
    result = redactor.redact(req.text)
    logger.info(
        "Detected %d spans: %s",
        result.summary.total_spans,
        result.summary.by_label,
    )
    return result


if __name__ == "__main__":
    host = os.environ.get("HOST", "0.0.0.0")
    port = int(os.environ.get("PORT", "8000"))
    uvicorn.run(app, host=host, port=port, log_level="debug" if DEV_MODE else "info")