feat: add api key authentication to /predict endpoint

api/main.py

@@ -3,8 +3,9 @@
 from contextlib import asynccontextmanager
 from time import perf_counter
 
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Depends, status
 from fastapi.responses import JSONResponse, PlainTextResponse
+from fastapi.security import APIKeyHeader
 from prometheus_client import CONTENT_TYPE_LATEST, Counter, Histogram, generate_latest
 
 from mlops_nlp.config import load_config
@@ -20,9 +21,21 @@
 REQUEST_LATENCY = Histogram("request_latency_seconds", "HTTP request latency", ["path", "method"])
 PROMETHEUS_ENABLED = APP_CONFIG.monitoring.enable_prometheus
 
+API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)
+
 pipeline: InferencePipeline | None = None
 
 
+def get_api_key(api_key_header: str = Depends(API_KEY_HEADER)):
+    if APP_CONFIG.api.api_key:
+        if api_key_header != APP_CONFIG.api.api_key:
+            raise HTTPException(
+                status_code=status.HTTP_403_FORBIDDEN,
+                detail="Could not validate credentials",
+            )
+    return api_key_header
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     global pipeline
@@ -81,7 +94,10 @@ def health() -> JSONResponse:
 
 
 @app.post("/predict", response_model=PredictionResponse)
-def predict(payload: PredictionRequest) -> PredictionResponse:
+def predict(
+    payload: PredictionRequest, 
+    _api_key: str = Depends(get_api_key)
+) -> PredictionResponse:
     if pipeline is None:
         raise HTTPException(status_code=503, detail="Model is not loaded. Train model first.")
     prediction, confidence = pipeline.predict(payload.text)

src/mlops_nlp/config.py

@@ -49,6 +49,7 @@ class ApiConfig(BaseModel):
     title: str
     version: str
     description: str
+    api_key: str | None = None
 
 
 class MonitoringConfig(BaseModel):
@@ -88,4 +89,8 @@ def load_config(config_path: str | Path = "configs/config.yaml") -> AppConfig:
     if enable_prometheus is not None:
         config.monitoring.enable_prometheus = _parse_bool(enable_prometheus)
 
+    api_key = os.getenv("MLOPS_API_KEY")
+    if api_key:
+        config.api.api_key = api_key
+
     return config

tests/test_api_auth.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import pytest
+from fastapi.testclient import TestClient
+from api.main import app
+import os
+
+client = TestClient(app)
+
+def test_predict_no_auth():
+    # If MLOPS_API_KEY is not set, it should allow access (default behavior for local dev if not configured)
+    # But for this test, we want to ensure that IF it is set, it blocks.
+    # However, the current implementation in main.py checks APP_CONFIG.api.api_key which is loaded at start.
+    pass
+
+@pytest.fixture
+def mock_api_key(monkeypatch):
+    monkeypatch.setenv("MLOPS_API_KEY", "test-secret-key")
+    # We need to reload the config or the app to pick up the env var if it's already loaded.
+    # In api/main.py, APP_CONFIG = load_config() happens at module level.
+    # For a robust test, we might need to override the dependency.
+    return "test-secret-key"
+
+def test_predict_with_auth_dependency_override():
+    from api.main import get_api_key, app
+
+    # Mocking the actual config value for the test
+    import api.main
+    original_key = api.main.APP_CONFIG.api.api_key
+    api.main.APP_CONFIG.api.api_key = "test-secret-key"
+
+    try:
+        # 1. No key provided
+        response = client.post("/predict", json={"text": "hello"})
+        assert response.status_code == 403
+
+        # 2. Wrong key provided
+        response = client.post("/predict", json={"text": "hello"}, headers={"X-API-Key": "wrong-key"})
+        assert response.status_code == 403
+
+        # 3. Correct key provided
+        # We need a model loaded for a full 200, otherwise it might be 503 if no model exists.
+        response = client.post("/predict", json={"text": "hello"}, headers={"X-API-Key": "test-secret-key"})
+        # We expect 503 (Model not loaded) or 200 (Success), but NOT 403.
+        assert response.status_code != 403
+
+    finally:
+        api.main.APP_CONFIG.api.api_key = original_key