feat: add rate limiting, CORS, queue loop in lifespan
This commit is contained in:
@@ -1,8 +1,17 @@
|
|||||||
from fastapi import Header, HTTPException
|
from fastapi import Header, HTTPException, Request
|
||||||
|
from slowapi import Limiter
|
||||||
|
from slowapi.util import get_remote_address
|
||||||
|
|
||||||
from src.backend.config import settings
|
from src.backend.config import settings
|
||||||
|
|
||||||
|
|
||||||
|
def _get_api_key(request: Request) -> str:
|
||||||
|
return request.headers.get("X-Api-Key", get_remote_address(request))
|
||||||
|
|
||||||
|
|
||||||
|
limiter = Limiter(key_func=_get_api_key)
|
||||||
|
|
||||||
|
|
||||||
async def verify_api_key(x_api_key: str | None = Header(default=None)):
|
async def verify_api_key(x_api_key: str | None = Header(default=None)):
|
||||||
if x_api_key is None or x_api_key != settings.api_key:
|
if x_api_key is None or x_api_key != settings.api_key:
|
||||||
raise HTTPException(status_code=403, detail="API key non valida")
|
raise HTTPException(status_code=403, detail="API key non valida")
|
||||||
|
|||||||
@@ -1,11 +1,28 @@
|
|||||||
from fastapi import APIRouter, Depends
|
from fastapi import APIRouter, Depends
|
||||||
|
|
||||||
from src.backend.api.dependencies import verify_api_key
|
from src.backend.api.dependencies import verify_api_key
|
||||||
|
from src.backend.db.database import get_db
|
||||||
from src.backend.models.api import ok
|
from src.backend.models.api import ok
|
||||||
|
from src.backend.services.worker_client import WorkerClient
|
||||||
|
|
||||||
router = APIRouter(dependencies=[Depends(verify_api_key)])
|
router = APIRouter(dependencies=[Depends(verify_api_key)])
|
||||||
|
|
||||||
|
|
||||||
@router.get("/health", summary="Health check", tags=["system"])
|
@router.get("/health", summary="Health check", tags=["system"])
|
||||||
async def health():
|
async def health():
|
||||||
return ok({"status": "running"})
|
db_ok = True
|
||||||
|
try:
|
||||||
|
db = await get_db()
|
||||||
|
await db.execute("SELECT 1")
|
||||||
|
await db.close()
|
||||||
|
except Exception:
|
||||||
|
db_ok = False
|
||||||
|
|
||||||
|
worker = WorkerClient()
|
||||||
|
worker_ok = await worker.health()
|
||||||
|
|
||||||
|
return ok({
|
||||||
|
"status": "running",
|
||||||
|
"database": db_ok,
|
||||||
|
"worker": worker_ok,
|
||||||
|
})
|
||||||
|
|||||||
@@ -1,7 +1,8 @@
|
|||||||
from fastapi import APIRouter, Depends, Query
|
from fastapi import APIRouter, Depends, Query, Request as FastAPIRequest
|
||||||
from starlette.responses import JSONResponse
|
from starlette.responses import JSONResponse
|
||||||
|
|
||||||
from src.backend.api.dependencies import verify_api_key
|
from src.backend.api.dependencies import verify_api_key, limiter
|
||||||
|
from src.backend.config import settings
|
||||||
from src.backend.db.database import (
|
from src.backend.db.database import (
|
||||||
get_db,
|
get_db,
|
||||||
create_request,
|
create_request,
|
||||||
@@ -27,7 +28,8 @@ async def _get_db():
|
|||||||
|
|
||||||
|
|
||||||
@router.post("", summary="Create request")
|
@router.post("", summary="Create request")
|
||||||
async def create(body: RequestCreate, db=Depends(_get_db)):
|
@limiter.limit(lambda: f"{settings.rate_limit_per_minute}/minute;{settings.rate_limit_per_hour}/hour")
|
||||||
|
async def create(request: FastAPIRequest, body: RequestCreate, db=Depends(_get_db)):
|
||||||
req = await create_request(
|
req = await create_request(
|
||||||
db,
|
db,
|
||||||
prompt=body.prompt,
|
prompt=body.prompt,
|
||||||
|
|||||||
+36
-2
@@ -1,21 +1,55 @@
|
|||||||
|
import asyncio
|
||||||
|
import logging
|
||||||
from contextlib import asynccontextmanager
|
from contextlib import asynccontextmanager
|
||||||
|
|
||||||
from fastapi import FastAPI
|
from fastapi import FastAPI
|
||||||
|
from fastapi.middleware.cors import CORSMiddleware
|
||||||
|
from slowapi import _rate_limit_exceeded_handler
|
||||||
|
from slowapi.errors import RateLimitExceeded
|
||||||
|
|
||||||
|
from src.backend.api.dependencies import limiter
|
||||||
from src.backend.api.routes.health import router as health_router
|
from src.backend.api.routes.health import router as health_router
|
||||||
from src.backend.api.routes.topics import router as topics_router
|
|
||||||
from src.backend.api.routes.requests import router as requests_router
|
from src.backend.api.routes.requests import router as requests_router
|
||||||
from src.backend.api.routes.sessions import router as sessions_router
|
from src.backend.api.routes.sessions import router as sessions_router
|
||||||
|
from src.backend.api.routes.topics import router as topics_router
|
||||||
from src.backend.db.database import init_db
|
from src.backend.db.database import init_db
|
||||||
|
from src.backend.services.queue import queue_loop
|
||||||
|
from src.backend.services.worker_client import WorkerClient
|
||||||
|
|
||||||
|
logging.basicConfig(level=logging.INFO)
|
||||||
|
|
||||||
|
|
||||||
@asynccontextmanager
|
@asynccontextmanager
|
||||||
async def lifespan(app: FastAPI):
|
async def lifespan(app: FastAPI):
|
||||||
await init_db()
|
await init_db()
|
||||||
|
worker = WorkerClient()
|
||||||
|
task = asyncio.create_task(queue_loop(worker))
|
||||||
yield
|
yield
|
||||||
|
task.cancel()
|
||||||
|
try:
|
||||||
|
await task
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
app = FastAPI(title="OpusAgent", version="0.1.0", lifespan=lifespan)
|
app = FastAPI(
|
||||||
|
title="OpusAgent",
|
||||||
|
version="0.1.0",
|
||||||
|
docs_url="/doc",
|
||||||
|
redoc_url="/redoc",
|
||||||
|
openapi_url="/openapi.json",
|
||||||
|
lifespan=lifespan,
|
||||||
|
)
|
||||||
|
|
||||||
|
app.state.limiter = limiter
|
||||||
|
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
|
||||||
|
|
||||||
|
app.add_middleware(
|
||||||
|
CORSMiddleware,
|
||||||
|
allow_origins=["*"],
|
||||||
|
allow_methods=["*"],
|
||||||
|
allow_headers=["*"],
|
||||||
|
)
|
||||||
|
|
||||||
app.include_router(health_router, prefix="/api")
|
app.include_router(health_router, prefix="/api")
|
||||||
app.include_router(topics_router, prefix="/api")
|
app.include_router(topics_router, prefix="/api")
|
||||||
|
|||||||
Reference in New Issue
Block a user