From a97392f421881e1fd0b2da32c9627416d54ca5d2 Mon Sep 17 00:00:00 2001 From: Bram Date: Thu, 17 Jul 2025 19:26:11 +0200 Subject: [PATCH] Add environment variable support for API key and max workers in OCR API; implement API key verification for secure access --- Dockers/ocr-api/docker-compose.yml | 5 +++ Dockers/ocr-api/env.example | 13 ++++++++ Dockers/ocr-api/main.py | 52 +++++++++++++++++++++++++++--- Dockers/ocr-api/test_client.py | 6 ++++ 4 files changed, 71 insertions(+), 5 deletions(-) create mode 100644 Dockers/ocr-api/env.example diff --git a/Dockers/ocr-api/docker-compose.yml b/Dockers/ocr-api/docker-compose.yml index 214a0cb..1e2b28d 100644 --- a/Dockers/ocr-api/docker-compose.yml +++ b/Dockers/ocr-api/docker-compose.yml @@ -7,6 +7,11 @@ services: - "8000:8000" environment: - PYTHONUNBUFFERED=1 + - DEFAULT_MAX_WORKERS=${DEFAULT_MAX_WORKERS:-4} + - API_KEY=${API_KEY:-} + - REQUIRE_API_KEY=${REQUIRE_API_KEY:-true} + env_file: + - .env volumes: # Optional: mount a directory for testing with sample files - ./test_files:/app/test_files:ro diff --git a/Dockers/ocr-api/env.example b/Dockers/ocr-api/env.example new file mode 100644 index 0000000..b5bb7aa --- /dev/null +++ b/Dockers/ocr-api/env.example @@ -0,0 +1,13 @@ +# OCR API Configuration + +# API Authentication +API_KEY=your-secret-api-key-here +REQUIRE_API_KEY=true + +# Performance Configuration +DEFAULT_MAX_WORKERS=4 + +# Optional: Override default settings +# LOG_LEVEL=INFO +# HOST=0.0.0.0 +# PORT=8000 \ No newline at end of file diff --git a/Dockers/ocr-api/main.py b/Dockers/ocr-api/main.py index 021c70c..bb0e0a8 100644 --- a/Dockers/ocr-api/main.py +++ b/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 diff --git a/Dockers/ocr-api/test_client.py b/Dockers/ocr-api/test_client.py index b45103a..4cc09b7 100644 --- a/Dockers/ocr-api/test_client.py +++ b/Dockers/ocr-api/test_client.py @@ -7,10 +7,16 @@ This script demonstrates how to use the OCR API endpoints import requests import json import sys +import os from pathlib import Path +from dotenv import load_dotenv + +# Load environment variables +load_dotenv() # API base URL BASE_URL = "http://localhost:8000" +API_KEY = os.getenv("API_KEY") def test_health(): """Test the health endpoint"""