Add environment variable support for API key and max workers in OCR API; implement API key verification for secure access

Dockers/ocr-api/main.py

@@ -7,18 +7,31 @@ from pathlib import Path
 from concurrent.futures import ThreadPoolExecutor
 import time
 
-from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks
+from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks, Depends, Header
 from fastapi.responses import JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
 from pydantic import BaseModel
 import pytesseract
 from PIL import Image
 from pdf2image import convert_from_path
 import uvicorn
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# Environment variables
+DEFAULT_MAX_WORKERS = int(os.getenv("DEFAULT_MAX_WORKERS", "4"))
+API_KEY = os.getenv("API_KEY")
+REQUIRE_API_KEY = os.getenv("REQUIRE_API_KEY", "true").lower() == "true"
+
+# Security
+security = HTTPBearer(auto_error=False)
+
 app = FastAPI(
     title="OCR API",
     description="API for extracting text from PDF files using OCR",
@@ -35,6 +48,33 @@ class ErrorResponse(BaseModel):
     error: str
     detail: str
 
+async def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
+    """
+    Verify API key authentication
+    """
+    if not REQUIRE_API_KEY:
+        return True
+
+    if not API_KEY:
+        logger.warning("API_KEY not set in environment variables")
+        return True
+
+    if not credentials:
+        raise HTTPException(
+            status_code=401,
+            detail="API key required",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    if credentials.credentials != API_KEY:
+        raise HTTPException(
+            status_code=401,
+            detail="Invalid API key",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    return True
+
 @app.get("/")
 async def root():
     return {"message": "OCR API is running. Use /docs for API documentation."}
@@ -100,7 +140,7 @@ async def process_pdf_with_ocr(pdf_path: str, max_workers: int = None) -> OCRRes
 
         # Determine optimal number of workers
         if max_workers is None:
-            max_workers = min(len(images), os.cpu_count() or 4)
+            max_workers = min(len(images), DEFAULT_MAX_WORKERS)
 
         logger.info(f"Processing {len(images)} pages with {max_workers} workers")
 
@@ -152,13 +192,14 @@ async def process_pdf_with_ocr(pdf_path: str, max_workers: int = None) -> OCRRes
 async def extract_text_from_pdf(
     file: UploadFile = File(...),
     max_workers: int = None,
-    background_tasks: BackgroundTasks = None
+    background_tasks: BackgroundTasks = None,
+    _: bool = Depends(verify_api_key)
 ):
     """
     Extract text from a PDF file using OCR with concurrent processing
 
     - **file**: PDF file to process
-    - **max_workers**: Maximum number of concurrent workers (default: min(pages, CPU cores))
+    - **max_workers**: Maximum number of concurrent workers (default: from DEFAULT_MAX_WORKERS env var)
     - Returns extracted text with metadata
     """
 
@@ -206,7 +247,8 @@ async def extract_text_from_pdf(
 @app.post("/ocr/image", response_model=OCRResponse)
 async def extract_text_from_image(
     file: UploadFile = File(...),
-    background_tasks: BackgroundTasks = None
+    background_tasks: BackgroundTasks = None,
+    _: bool = Depends(verify_api_key)
 ):
     """
     Extract text from an image file using OCR