""" TB-Guard-XAI Production Backend Complete end-to-end implementation with: - Proper dependency injection (no global state) - Structured logging and monitoring - Error handling and rate limiting - CORS security - Request/response validation - Async/await for non-blocking operations - Audit trail for HIPAA compliance """ import os import sys from pathlib import Path # CRITICAL: Add local directory to sys.path FIRST, before any other imports # This prevents torchxrayvision.config from being imported instead of local config _local_dir = str(Path(__file__).parent) if _local_dir not in sys.path: sys.path.insert(0, _local_dir) import logging import json import traceback import base64 import shutil import time import asyncio import io import hashlib from datetime import datetime from typing import Optional from contextlib import asynccontextmanager try: import torch TORCH_AVAILABLE = True except ImportError: TORCH_AVAILABLE = False torch = None import uvicorn from PIL import Image # FastAPI imports from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Depends, Request from fastapi.responses import HTMLResponse, FileResponse, StreamingResponse, JSONResponse from fastapi.templating import Jinja2Templates from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from fastapi.exception_handlers import RequestValidationError from pydantic import ValidationError # Now import local modules AFTER sys.path is configured from config import settings from schemas import ( AnalysisRequest, AnalysisResponse, BatchAnalysisResponse, HealthCheckResponse, PredictionType, AnalysisMode, UncertaintyLevel, TranscribeResponse, ConsultResponse, ConsultRequest, ErrorResponse ) from errors import ( TBGuardException, ModelNotLoadedError, InvalidImageError, FileTooLargeError, PreprocessingError, ModelError, InvalidInputError, handle_exception ) from rate_limiter import rate_limiter try: from mistral_explainer import MistralExplainer logger_init = logging.getLogger("tb_guard") logger_init.info("✓ MistralExplainer imported successfully") except ImportError as e: logger_init = logging.getLogger("tb_guard") logger_init.warning(f"⚠ MistralExplainer import deferred: {e} (will try again on startup)") MistralExplainer = None except Exception as e: logger_init = logging.getLogger("tb_guard") logger_init.error(f"✗ Failed to import MistralExplainer: {e}", exc_info=True) MistralExplainer = None # ============ Logging Configuration ============ logging.basicConfig( level=logging.INFO, format="[%(levelname)s] %(asctime)s %(name)s: %(message)s", stream=sys.stdout ) logger = logging.getLogger("tb_guard") # ============ Application State ============ class AppState: """Application state holder (replaces global variables)""" explainer: Optional[MistralExplainer] = None device: str = "cuda" if (TORCH_AVAILABLE and torch.cuda.is_available()) else "cpu" app_state = AppState() # ============ FastAPI Lifespan (Startup/Shutdown) ============ @asynccontextmanager async def lifespan(app: FastAPI): """ Lifespan context manager handles startup and shutdown Issues #7, #8: Replaces global state mutation """ # ===== STARTUP ===== logger.info("=" * 70) logger.info("TB-Guard-XAI Backend Starting") logger.info("=" * 70) logger.info(f"Device: {app_state.device}") logger.info(f"Model path: {settings.model_path}") logger.info(f"Environment: {settings.env}") logger.info(f"Host: {settings.host}:{settings.port}") # Clean up old batch reports reports_dir = Path("batch_reports") if reports_dir.exists(): now = time.time() cleanup_count = 0 for f in reports_dir.iterdir(): if f.is_file() and f.suffix == '.json': try: if now - f.stat().st_mtime > 86400: # 24 hours f.unlink() cleanup_count += 1 except Exception: pass if cleanup_count > 0: logger.info(f"Cleaned up {cleanup_count} old batch reports") # Load models try: if not TORCH_AVAILABLE: logger.warning("⚠ torch not installed - model loading unavailable") logger.warning(" (Install with: pip install torch torchvision)") logger.info(" API will run in demo mode (no ML predictions)") elif MistralExplainer is None: logger.error("✗ MistralExplainer not available (import failed)") raise ModelNotLoadedError() else: model_path = Path(settings.model_path) if not model_path.exists(): logger.warning(f"⚠ Model not found: {model_path}") logger.info(" API will run in demo mode (no ML predictions)") else: logger.info(f"Loading model from: {model_path} ({model_path.stat().st_size / 1024 / 1024:.1f} MB)") app_state.explainer = MistralExplainer(model_path=str(model_path)) logger.info("✓ Models loaded successfully") logger.info(f" - CNN Ensemble: Ready") logger.info(f" - Mistral LLM: {'Ready' if app_state.explainer.mistral else 'Not configured'}") logger.info(f" - Qdrant RAG: {'Ready' if app_state.explainer.rag else 'Not available'}") except Exception as e: logger.error(f"✗ Model loading failed: {e}", exc_info=True) raise logger.info("✓ Startup complete. Server ready for requests.") logger.info("=" * 70) yield # Server running # ===== SHUTDOWN ===== logger.info("=" * 70) logger.info("TB-Guard-XAI Backend Shutting Down") logger.info("=" * 70) if app_state.explainer: del app_state.explainer logger.info("✓ Models unloaded") logger.info("✓ Shutdown complete") logger.info("=" * 70) # ============ FastAPI Application ============ app = FastAPI( title="TB-Guard-XAI Clinical AI", description="Enterprise-grade AI screening for Tuberculosis", version="3.0.0", lifespan=lifespan ) # ============ CORS Middleware ============ # Issue #9: CORS limited to configured origins app.add_middleware( CORSMiddleware, allow_origins=settings.cors_origins, allow_credentials=False, allow_methods=settings.cors_methods, allow_headers=["Content-Type", "Authorization"], ) # ============ Setup Templates & Static Files ============ BASE_DIR = Path(__file__).resolve().parent templates = Jinja2Templates(directory=BASE_DIR / "templates") try: app.mount("/static", StaticFiles(directory=BASE_DIR / "static"), name="static") except Exception as e: logger.warning(f"Could not mount static files: {e}") # ============ Dependency Injection ============ # Issue #7: Proper dependency injection instead of globals def get_explainer() -> MistralExplainer: """Get loaded explainer model""" if app_state.explainer is None: raise ModelNotLoadedError() return app_state.explainer def get_request_id() -> str: """Generate request ID for logging""" import uuid return str(uuid.uuid4())[:12] async def check_rate_limit(request: Request) -> str: """ Check rate limit for request IP Issue #16: Rate limiting per IP and API key """ client_ip = request.client.host if request.client else "unknown" api_key = request.headers.get("X-API-Key", "default") allowed, reason = rate_limiter.is_allowed(client_ip, api_key) if not allowed: logger.warning(f"Rate limit exceeded for {client_ip}: {reason}") raise HTTPException(status_code=429, detail=reason) return client_ip # ============ Exception Handlers ============ @app.exception_handler(ValidationError) async def validation_exception_handler(request: Request, exc: ValidationError): """Handle Pydantic validation errors""" logger.warning(f"Validation error: {exc}") return JSONResponse( status_code=400, content={ "error": "Validation error", "code": "VALIDATION_ERROR", "details": exc.errors() } ) @app.exception_handler(TBGuardException) async def tb_guard_exception_handler(request: Request, exc: TBGuardException): """Handle TB-Guard custom exceptions""" exc.log(exc_info=False) return JSONResponse( status_code=exc.status_code, content={ "error": exc.message, "code": exc.code, "error_id": exc.error_id, "timestamp": exc.timestamp } ) @app.exception_handler(Exception) async def generic_exception_handler(request: Request, exc: Exception): """Catch-all for unexpected exceptions""" error_response = handle_exception(exc) return JSONResponse( status_code=error_response.status_code, content=error_response.detail ) # ============ Routes: Health & Status ============ @app.get("/", response_class=HTMLResponse) async def home(request: Request): """Render home page""" return templates.TemplateResponse(request, "index.html") @app.get("/health") async def health_check() -> HealthCheckResponse: """ Health check endpoint (for load balancers, monitoring) Issue #21: Structured logging """ rag_connected = False if app_state.explainer and hasattr(app_state.explainer, 'rag'): rag_connected = app_state.explainer.rag is not None return HealthCheckResponse( status="ok" if app_state.explainer else "error", model_device=app_state.device, rag_ready=rag_connected, timestamp=datetime.utcnow().isoformat() ) @app.get("/status") async def status(explainer: MistralExplainer = Depends(get_explainer)) -> dict: """Get system status""" rag_connected = explainer.rag is not None return { "status": "online", "model_device": app_state.device, "rag_ready": rag_connected } # ============ Routes: X-Ray Analysis ============ @app.post("/analyze", response_model=AnalysisResponse) async def analyze_xray( file: UploadFile = File(...), symptoms: str = Form(""), threshold: float = Form(0.5), age_group: str = Form("Adult (18-64)"), force_offline: bool = Form(False), report_type: str = Form("clinical"), mrn: str = Form(""), patient_name: str = Form(""), age: str = Form(""), sex: str = Form(""), institution: str = Form(""), explainer: MistralExplainer = Depends(get_explainer), request_id: str = Depends(get_request_id), client_ip: str = Depends(check_rate_limit) ) -> AnalysisResponse: """ Analyze chest X-ray for tuberculosis Issue #6: Using asyncio.to_thread() for blocking operations Issue #13: Response model with proper typing Issue #15: Input validation Issue #21: Performance logging Issue #25: Drift monitoring integration """ logger.info(f"[REQ {request_id}] Analyze request: symptoms={len(symptoms)} chars, " f"threshold={threshold}, age_group={age_group}, force_offline={force_offline}, " f"report_type={report_type}") start_time = time.time() try: # ===== File Validation ===== if not file.filename: raise InvalidImageError("File has no name") file_ext = Path(file.filename).suffix.lower() allowed_ext = ['.jpg', '.jpeg', '.png', '.webp'] if file_ext not in allowed_ext: raise InvalidImageError(f"Invalid file type: {file_ext}. Allowed: {', '.join(allowed_ext)}") contents = await file.read() if len(contents) == 0: raise InvalidImageError("Empty file") if len(contents) > settings.max_file_size_bytes: raise FileTooLargeError( file_size_mb=len(contents) / 1024 / 1024, max_size_mb=settings.max_file_size_bytes ) # ===== Image Validation ===== try: img = Image.open(io.BytesIO(contents)) img.verify() img = Image.open(io.BytesIO(contents)) w, h = img.size if w < settings.min_image_width or h < settings.min_image_height: raise InvalidImageError( f"Image too small: {w}x{h}. Min: {settings.min_image_width}x{settings.min_image_height}" ) if w > settings.max_image_width or h > settings.max_image_height: raise InvalidImageError( f"Image too large: {w}x{h}. Max: {settings.max_image_width}x{settings.max_image_height}" ) except InvalidImageError: raise except Exception as e: raise InvalidImageError(f"Image validation failed: {str(e)}") # ===== Input Validation ===== symptoms = symptoms.strip()[:500] # Sanitize if len(symptoms) > 0: forbidden = ["