diff --git a/.env.example b/.env.example index b2ec8a9..faecbdf 100644 --- a/.env.example +++ b/.env.example @@ -4,3 +4,4 @@ COMPASS_DOCS_PATH=./data/docs COMPASS_WORKSPACE=./data/index COMPASS_ALLOWED_ORIGINS=http://localhost:3000,http://localhost:5173 COMPASS_API_KEY= # leave empty to disable auth (local dev) +COMPASS_RATE_LIMIT=60/minute # requests per IP per minute — raise for trusted networks, lower for public exposure diff --git a/backend/main.py b/backend/main.py index 4b5d6a2..f041c52 100644 --- a/backend/main.py +++ b/backend/main.py @@ -9,6 +9,9 @@ from fastapi.security.api_key import APIKeyHeader from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel +from slowapi import Limiter, _rate_limit_exceeded_handler +from slowapi.util import get_remote_address +from slowapi.errors import RateLimitExceeded load_dotenv() @@ -20,12 +23,29 @@ logging.basicConfig(level=logging.INFO, format="%(levelname)s — %(name)s — %(message)s") logger = logging.getLogger(__name__) -# Config +# --------------------------------------------------------------------------- +# Config — all values come from environment variables (see .env.example) +# --------------------------------------------------------------------------- + OLLAMA_API_BASE = os.getenv("OLLAMA_API_BASE", "http://localhost:11434") COMPASS_MODEL = os.getenv("COMPASS_MODEL", "ollama/gemma4:e4b") DOCS_PATH = Path(os.getenv("COMPASS_DOCS_PATH", "./data/docs")) WORKSPACE = Path(os.getenv("COMPASS_WORKSPACE", "./data/index")) -API_KEY = os.getenv("COMPASS_API_KEY") # None = auth disabled + +# Optional API key — if unset, authentication is disabled (safe for local dev). +# Set COMPASS_API_KEY in .env to require the X-API-Key header on all requests. +API_KEY = os.getenv("COMPASS_API_KEY") + +# Rate limit — all endpoints share the same limit (requests per minute per IP). +# All requests over the limit receive HTTP 429. Adjust COMPASS_RATE_LIMIT in +# .env to tune (e.g. "30/minute", "200/minute"). Default: 60/minute. +RATE_LIMIT = os.getenv("COMPASS_RATE_LIMIT", "60/minute") + +os.environ["OLLAMA_API_BASE"] = OLLAMA_API_BASE + +# --------------------------------------------------------------------------- +# Auth +# --------------------------------------------------------------------------- _api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False) @@ -34,8 +54,17 @@ async def verify_api_key(key: str = Security(_api_key_header)): if API_KEY and key != API_KEY: raise HTTPException(status_code=401, detail="Invalid or missing API key.") -os.environ["OLLAMA_API_BASE"] = OLLAMA_API_BASE +# --------------------------------------------------------------------------- +# Rate limiter — keyed by client IP +# --------------------------------------------------------------------------- + +limiter = Limiter(key_func=get_remote_address) + + +# --------------------------------------------------------------------------- +# App lifecycle +# --------------------------------------------------------------------------- def get_indexer(request: Request) -> CompassIndexer: return request.app.state.indexer @@ -61,6 +90,10 @@ async def lifespan(app: FastAPI): app.state.watcher.join() +# --------------------------------------------------------------------------- +# App +# --------------------------------------------------------------------------- + app = FastAPI( title="Compass", description="El cerebro operativo de tu empresa.", @@ -69,6 +102,9 @@ async def lifespan(app: FastAPI): dependencies=[Security(verify_api_key)], ) +app.state.limiter = limiter +app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) + ALLOWED_ORIGINS = [ o.strip() for o in os.getenv("COMPASS_ALLOWED_ORIGINS", "http://localhost:3000,http://localhost:5173").split(",") @@ -79,16 +115,25 @@ async def lifespan(app: FastAPI): CORSMiddleware, allow_origins=ALLOWED_ORIGINS, allow_methods=["GET", "POST"], - allow_headers=["Content-Type"], + allow_headers=["Content-Type", "X-API-Key"], ) +# --------------------------------------------------------------------------- +# Models +# --------------------------------------------------------------------------- + class ChatRequest(BaseModel): question: str session_id: str = "default" +# --------------------------------------------------------------------------- +# Routes +# --------------------------------------------------------------------------- + @app.get("/health") +@limiter.limit(RATE_LIMIT) def health(request: Request): indexer = get_indexer(request) return { @@ -99,6 +144,7 @@ def health(request: Request): @app.post("/chat") +@limiter.limit(RATE_LIMIT) async def chat(body: ChatRequest, request: Request): indexer = get_indexer(request) if not indexer.documents: @@ -118,6 +164,7 @@ async def chat(body: ChatRequest, request: Request): @app.post("/upload", status_code=201) +@limiter.limit(RATE_LIMIT) async def upload(file: UploadFile = File(...), request: Request = None): indexer = get_indexer(request) @@ -139,6 +186,7 @@ async def upload(file: UploadFile = File(...), request: Request = None): @app.get("/documents") +@limiter.limit(RATE_LIMIT) def list_documents(request: Request): indexer = get_indexer(request) return {"documents": indexer.list_documents(), "total": len(indexer.documents)} diff --git a/requirements.txt b/requirements.txt index e17ab4c..beefb59 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,7 @@ pymupdf==1.26.4 watchdog==6.0.0 PyPDF2==3.0.1 pyyaml==6.0.2 +slowapi==0.1.9 # testing pytest==8.3.5 diff --git a/tests/test_api.py b/tests/test_api.py index 4db0edf..e4be0f3 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -50,6 +50,23 @@ def test_auth_applies_to_documents(self, client_empty): assert r.status_code == 401 +# --------------------------------------------------------------------------- +# Rate limiting +# --------------------------------------------------------------------------- + +class TestRateLimit: + def test_rate_limit_handler_is_registered(self, client_empty): + """Verify slowapi's 429 handler is registered on the app.""" + from slowapi.errors import RateLimitExceeded + from backend.main import app + assert RateLimitExceeded in app.exception_handlers + + def test_limiter_attached_to_app(self, client_empty): + """Verify the limiter is attached to app.state so slowapi can find it.""" + from backend.main import app, limiter + assert app.state.limiter is limiter + + # --------------------------------------------------------------------------- # GET /health # ---------------------------------------------------------------------------