Update main.py

main.py

@@ -2,67 +2,61 @@ import os
 import time
 import tempfile
 import traceback
+import logging # Added
 from pathlib import Path
 from typing import List, Union, Optional, Dict, Any, Literal
+import uuid # For request IDs
 
-from fastapi import FastAPI, File, UploadFile, HTTPException, status
+from fastapi import FastAPI, File, UploadFile, HTTPException, status, Request # Added Request
 from pydantic import BaseModel, Field
-from PIL import Image
+# PIL.Image is not directly used in this file after refactoring,
+# but mdr_pdf_parser might use it, so keep if necessary for that.
 
 # --- IMPORTS ---
 from mdr_pdf_parser import (
     MagicPDFProcessor,
-    MDRStructuredBlock,
+    MDRStructuredBlock, # Assuming this is the base type for the others
     MDRTextBlock,
     MDRTableBlock,
     MDRFormulaBlock,
     MDRFigureBlock,
-    MDRTextKind,
-    MDRTableFormat,
-    MDRRectangle,
-    MDRTextSpan,
+    MDRTextKind,       # Used by MDRTextBlockModel
+    MDRTableFormat,    # Used by MDRTableBlockModel
+    MDRRectangle,      # Used by MDRRectangleModel
+    MDRTextSpan,       # Used by MDRTextSpanModel
     MDRExtractedTableFormat # For configuration
 )
 
+# --- Logging Configuration ---
+LOG_LEVEL_STR = os.environ.get("LOG_LEVEL", "INFO").upper()
+LOG_LEVEL = getattr(logging, LOG_LEVEL_STR, logging.INFO)
+
+logging.basicConfig(
+    level=LOG_LEVEL,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger("mdr_fastapi_service")
+
 # --- Configuration ---
 # Read from environment variables, falling back to defaults
-MODEL_DIR = os.environ.get("MDR_MODEL_DIR", "/models") # Default path inside container
-DEVICE = os.environ.get("MDR_DEVICE", "cuda") # Default to cuda, processor will fallback if needed
-TABLE_FORMAT_STR = os.environ.get("MDR_TABLE_FORMAT", "MARKDOWN") # Default table format
+MODEL_DIR = os.environ.get("MDR_MODEL_DIR", "/models")
+DEVICE = os.environ.get("MDR_DEVICE", "cuda")
+TABLE_FORMAT_STR = os.environ.get("MDR_TABLE_FORMAT", "MARKDOWN")
+DEBUG_DIR_PATH_STR = os.environ.get("MDR_DEBUG_DIR_PATH", None) # For processor's debug output
 
 # Convert table format string to Enum
 try:
     TABLE_FORMAT = MDRExtractedTableFormat[TABLE_FORMAT_STR.upper()]
 except KeyError:
-    print(f"Warning: Invalid MDR_TABLE_FORMAT '{TABLE_FORMAT_STR}'. Defaulting to DISABLE.")
+    logger.warning(f"Invalid MDR_TABLE_FORMAT '{TABLE_FORMAT_STR}'. Defaulting to DISABLE.")
     TABLE_FORMAT = MDRExtractedTableFormat.DISABLE
 
-# --- Initialize Processor ---
-# This happens once when the service starts
-print("Initializing MagicPDFProcessor...")
-print(f"Model Directory: {MODEL_DIR}")
-print(f"Target Device: {DEVICE}")
-print(f"Table Format: {TABLE_FORMAT.name}")
-start_time = time.time()
-try:
-    mdr_processor = MagicPDFProcessor(
-        device=DEVICE,
-        model_dir_path=MODEL_DIR,
-        extract_table_format=TABLE_FORMAT,
-        # Set debug_dir_path=None for production service
-        debug_dir_path=None
-    )
-    print(f"MagicPDFProcessor initialized successfully ({time.time() - start_time:.2f}s)")
-except Exception as e:
-    print(f"FATAL ERROR: Failed to initialize MagicPDFProcessor during startup: {e}")
-    print("Service cannot start.")
-    traceback.print_exc()
-    # Optionally exit or raise to prevent FastAPI from starting incorrectly
-    mdr_processor = None # Ensure processor is None if init fails
+# --- Global Processor Variable ---
+mdr_processor: Optional[MagicPDFProcessor] = None
 
 # --- API Models (Pydantic) ---
-# Define models for API input/output for validation and documentation
-
+# (Pydantic models remain largely the same, ensure they are correct and complete)
 class MDRPointModel(BaseModel):
     x: float
     y: float
@@ -96,18 +90,17 @@ class MDRTextSpanModel(BaseModel):
         )
 
 class MDRBasicBlockModel(BaseModel):
-    block_type: str # To distinguish block types in the union
+    block_type: str
     rect: MDRRectangleModel
-    texts: List[MDRTextSpanModel] = Field(default_factory=list) # Captions/footnotes
+    texts: List[MDRTextSpanModel] = Field(default_factory=list)
     font_size: float
 
 class MDRTextBlockModel(MDRBasicBlockModel):
     block_type: Literal["TextBlock"] = "TextBlock"
-    kind: str # Use string representation of enum
+    kind: str
     has_paragraph_indentation: bool
     last_line_touch_end: bool
-    # Override texts field specifically for TextBlock
-    texts: List[MDRTextSpanModel] # Text content itself
+    texts: List[MDRTextSpanModel]
 
     @classmethod
     def from_mdr_text_block(cls, block: MDRTextBlock):
@@ -115,7 +108,7 @@ class MDRTextBlockModel(MDRBasicBlockModel):
             rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
             texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts],
             font_size=block.font_size,
-            kind=block.kind.name, # Convert enum to string name
+            kind=block.kind.name,
             has_paragraph_indentation=block.has_paragraph_indentation,
             last_line_touch_end=block.last_line_touch_end
         )
@@ -123,47 +116,43 @@ class MDRTextBlockModel(MDRBasicBlockModel):
 class MDRTableBlockModel(MDRBasicBlockModel):
     block_type: Literal["TableBlock"] = "TableBlock"
     content: str
-    format: str # Use string representation of enum
-    # Omit 'image' field from API response
+    format: str
 
     @classmethod
     def from_mdr_table_block(cls, block: MDRTableBlock):
         return cls(
             rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
-            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts], # Captions
+            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts],
             font_size=block.font_size,
             content=block.content,
-            format=block.format.name # Convert enum to string name
+            format=block.format.name
         )
 
 class MDRFormulaBlockModel(MDRBasicBlockModel):
     block_type: Literal["FormulaBlock"] = "FormulaBlock"
     content: Optional[str] = None
-    # Omit 'image' field from API response
 
     @classmethod
     def from_mdr_formula_block(cls, block: MDRFormulaBlock):
         return cls(
             rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
-            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts], # Captions
+            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts],
             font_size=block.font_size,
             content=block.content
         )
 
 class MDRFigureBlockModel(MDRBasicBlockModel):
     block_type: Literal["FigureBlock"] = "FigureBlock"
-    # Omit 'image' field from API response
 
     @classmethod
     def from_mdr_figure_block(cls, block: MDRFigureBlock):
         return cls(
             rect=MDRRectangleModel.from_mdr_rectangle(block.rect),
-            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts], # Captions
+            texts=[MDRTextSpanModel.from_mdr_text_span(span) for span in block.texts],
             font_size=block.font_size
         )
 
-# Union type for the response model
-MDRStructuredBlockModelAPI = Union[MDRTextBlockModel, MDRTableBlockModel, MDRFormulaBlockModel, MDRFigureBlockModel] # Renamed API Union type
+MDRStructuredBlockModelAPI = Union[MDRTextBlockModel, MDRTableBlockModel, MDRFormulaBlockModel, MDRFigureBlockModel]
 
 # --- FastAPI App ---
 app = FastAPI(
@@ -172,95 +161,184 @@ app = FastAPI(
     version="1.0.0"
 )
 
+# --- Helper Functions ---
+def _convert_block_to_api_model(block: MDRStructuredBlock) -> Optional[MDRStructuredBlockModelAPI]:
+    """Converts internal MDR block to an API model."""
+    if isinstance(block, MDRTextBlock):
+        return MDRTextBlockModel.from_mdr_text_block(block)
+    elif isinstance(block, MDRTableBlock):
+        return MDRTableBlockModel.from_mdr_table_block(block)
+    elif isinstance(block, MDRFormulaBlock):
+        return MDRFormulaBlockModel.from_mdr_formula_block(block)
+    elif isinstance(block, MDRFigureBlock):
+        return MDRFigureBlockModel.from_mdr_figure_block(block)
+    logger.warning(f"Unknown block type encountered: {type(block)}. Skipping conversion.")
+    return None
+
+# --- Application Lifecycle Events ---
 @app.on_event("startup")
 async def startup_event():
+    global mdr_processor
+    logger.info("Application startup sequence initiated.")
+    logger.info("--- Configuration ---")
+    logger.info(f"  MDR_MODEL_DIR: {MODEL_DIR}")
+    logger.info(f"  MDR_DEVICE: {DEVICE}")
+    logger.info(f"  MDR_TABLE_FORMAT: {TABLE_FORMAT.name}")
+    logger.info(f"  MDR_DEBUG_DIR_PATH: {DEBUG_DIR_PATH_STR if DEBUG_DIR_PATH_STR else 'Not set'}")
+    logger.info(f"  LOG_LEVEL: {LOG_LEVEL_STR}")
+    logger.info("---------------------")
+
+    logger.info("Initializing MagicPDFProcessor...")
+    init_start_time = time.time()
+    try:
+        mdr_processor = MagicPDFProcessor(
+            device=DEVICE,
+            model_dir_path=MODEL_DIR,
+            extract_table_format=TABLE_FORMAT,
+            debug_dir_path=DEBUG_DIR_PATH_STR # Pass the actual path or None
+        )
+        init_duration = time.time() - init_start_time
+        logger.info(f"MagicPDFProcessor initialized successfully ({init_duration:.2f}s)")
+    except Exception as e:
+        logger.critical(f"Failed to initialize MagicPDFProcessor: {e}", exc_info=True)
+        # mdr_processor will remain None, startup_event_check will handle this
+        # No need to print traceback here, logger.critical with exc_info=True does it.
+
+@app.on_event("startup") # Separate event to check after initialization attempt
+async def startup_event_check():
     if mdr_processor is None:
-        # This prevents the app from starting if initialization failed
-        raise RuntimeError("MagicPDFProcessor failed to initialize. Service cannot start.")
-    print("MagicDataReadiness Service is ready.")
+        logger.error("MagicPDFProcessor is not initialized. Service cannot function correctly.")
+        # Depending on deployment, you might want to exit or let it run in a degraded state.
+        # For now, it will allow FastAPI to start but /health and /process-pdf will fail.
+        # raise RuntimeError("MagicPDFProcessor failed to initialize. Service cannot start.") # This would stop FastAPI
+    else:
+        logger.info("MagicDataReadiness Service is ready and processor is available.")
+
 
-@app.get("/health")
+# --- API Endpoints ---
+@app.get("/health", summary="Health Check")
 async def health_check():
     """Simple health check endpoint."""
     if mdr_processor is None:
+        logger.warning("/health endpoint called but processor is not initialized.")
         raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Processor not initialized")
     return {"status": "ok", "message": "MagicPDFProcessor is running."}
 
 @app.post("/process-pdf/",
-          response_model=List[MDRStructuredBlockModelAPI], # Use the Union type
+          response_model=List[MDRStructuredBlockModelAPI],
           summary="Process a PDF file",
           description="Upload a PDF file to extract structured blocks (text, tables, figures, formulas).")
-async def process_pdf_endpoint(file: UploadFile = File(..., description="The PDF file to process.")):
+async def process_pdf_endpoint(request: Request, file: UploadFile = File(..., description="The PDF file to process.")):
     """
     Handles PDF file upload, processing, and returns extracted blocks.
     """
+    request_id = str(uuid.uuid4())
+    client_host = request.client.host if request.client else "unknown"
+    logger.info(f"RID-{request_id}: Received /process-pdf request from {client_host} for file: '{file.filename}' (type: {file.content_type}, size: {file.size})")
+
     if mdr_processor is None:
+        logger.error(f"RID-{request_id}: Processor not initialized. Cannot process request.")
         raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Processor not initialized")
 
-    if not file.filename.lower().endswith(".pdf"):
+    if not file.filename or not file.filename.lower().endswith(".pdf"):
+        logger.warning(f"RID-{request_id}: Invalid file type uploaded: '{file.filename}'")
         raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Invalid file type. Please upload a PDF.")
 
-    # Save uploaded file temporarily
-    temp_pdf_path = "" # Initialize path
+    temp_pdf_path_obj = None # To ensure it's always defined for finally block
     try:
+        # Save uploaded file temporarily
+        save_start_time = time.time()
         # Create a temporary directory if it doesn't exist
-        temp_dir = Path("./temp_uploads")
-        temp_dir.mkdir(exist_ok=True)
+        temp_dir = Path("./temp_uploads") # Consider making this configurable
+        temp_dir.mkdir(parents=True, exist_ok=True)
 
-        # Use a temporary file within the directory
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir=temp_dir) as temp_file:
+        # Use a temporary file with a unique name to avoid collisions
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf", dir=temp_dir, prefix=f"req_{request_id}_") as temp_file:
             content = await file.read()
             temp_file.write(content)
-            temp_pdf_path = temp_file.name
-        print(f"Received file '{file.filename}', saved temporarily to '{temp_pdf_path}'")
+            temp_pdf_path_obj = Path(temp_file.name)
+        save_duration = time.time() - save_start_time
+        logger.info(f"RID-{request_id}: File '{file.filename}' saved temporarily to '{temp_pdf_path_obj}' ({save_duration:.2f}s)")
+
     except Exception as e:
+        logger.error(f"RID-{request_id}: Failed to save uploaded file '{file.filename}': {e}", exc_info=True)
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to save uploaded file: {e}")
 
     extracted_blocks_api: List[MDRStructuredBlockModelAPI] = []
-    start_process_time = time.time()
+    processing_start_time = time.time()
 
     try:
-        print(f"Processing '{temp_pdf_path}'...")
+        logger.info(f"RID-{request_id}: Starting PDF processing for '{temp_pdf_path_obj}'...")
+
+        # Define a progress callback for the processor
+        def log_progress(completed_pages: int, total_pages: int):
+            if total_pages > 0:
+                percent_done = (completed_pages / total_pages) * 100
+                logger.info(f"RID-{request_id}: Processor progress - Page {completed_pages}/{total_pages} ({percent_done:.1f}%) scanned.")
+            else:
+                logger.info(f"RID-{request_id}: Processor progress - Page {completed_pages} scanned (total pages unknown or zero).")
+
+
         # Process the document using the temporary file path
-        # Note: process_document returns a generator, collect all blocks
-        all_blocks = list(mdr_processor.process_document(pdf_input=temp_pdf_path))
-        print(f"Extracted {len(all_blocks)} raw blocks.")
-
-        # Convert internal block types to API response models
-        for block in all_blocks:
-            if isinstance(block, MDRTextBlock):
-                extracted_blocks_api.append(MDRTextBlockModel.from_mdr_text_block(block))
-            elif isinstance(block, MDRTableBlock):
-                extracted_blocks_api.append(MDRTableBlockModel.from_mdr_table_block(block))
-            elif isinstance(block, MDRFormulaBlock):
-                extracted_blocks_api.append(MDRFormulaBlockModel.from_mdr_formula_block(block))
-            elif isinstance(block, MDRFigureBlock):
-                extracted_blocks_api.append(MDRFigureBlockModel.from_mdr_figure_block(block))
-            # Add more elif clauses if there are other block types
-
-        process_time = time.time() - start_process_time
-        print(f"Processing finished in {process_time:.2f}s. Returning {len(extracted_blocks_api)} blocks.")
+        # MagicPDFProcessor.process_document returns a generator
+        all_blocks_internal = list(mdr_processor.process_document(
+            pdf_input=str(temp_pdf_path_obj),
+            report_progress=log_progress # Pass the progress logger
+        ))
+        collection_duration = time.time() - processing_start_time
+        logger.info(f"RID-{request_id}: Extracted {len(all_blocks_internal)} raw blocks from processor ({collection_duration:.2f}s).")
+
+        conversion_start_time = time.time()
+        for i, block_internal in enumerate(all_blocks_internal):
+            logger.debug(f"RID-{request_id}: Converting internal block {i+1}/{len(all_blocks_internal)} of type {type(block_internal)} to API model.")
+            api_block = _convert_block_to_api_model(block_internal)
+            if api_block:
+                extracted_blocks_api.append(api_block)
+        conversion_duration = time.time() - conversion_start_time
+        logger.info(f"RID-{request_id}: Converted {len(extracted_blocks_api)} blocks to API models ({conversion_duration:.2f}s).")
+
+        total_processing_duration = time.time() - processing_start_time
+        logger.info(f"RID-{request_id}: PDF processing finished in {total_processing_duration:.2f}s. Returning {len(extracted_blocks_api)} blocks.")
 
     except Exception as e:
-        print(f"ERROR during PDF processing: {e}")
-        traceback.print_exc()
+        logger.error(f"RID-{request_id}: Error during PDF processing for '{temp_pdf_path_obj}': {e}", exc_info=True)
+        # Ensure traceback is logged by logger.error with exc_info=True
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An error occurred during PDF processing: {e}")
     finally:
         # Clean up the temporary file
-        if temp_pdf_path and os.path.exists(temp_pdf_path):
+        if temp_pdf_path_obj and temp_pdf_path_obj.exists():
             try:
-                os.remove(temp_pdf_path)
-                print(f"Cleaned up temporary file: {temp_pdf_path}")
+                os.remove(temp_pdf_path_obj)
+                logger.info(f"RID-{request_id}: Cleaned up temporary file: {temp_pdf_path_obj}")
             except OSError as e:
-                print(f"Warning: Could not remove temporary file {temp_pdf_path}: {e}")
+                logger.warning(f"RID-{request_id}: Could not remove temporary file {temp_pdf_path_obj}: {e}")
+        elif temp_pdf_path_obj: # Path was set but file doesn't exist (e.g. save failed)
+             logger.info(f"RID-{request_id}: Temporary file {temp_pdf_path_obj} not found for cleanup (may have failed to save).")
+
 
     return extracted_blocks_api
 
-# Optional: Add root endpoint for basic info/docs link
-@app.get("/")
+@app.get("/", summary="Root Endpoint")
 async def read_root():
+    """Provides basic information about the API."""
     return {
         "message": "Welcome to the MagicDataReadiness PDF Processor API!",
-        "docs_url": "/docs",
-        "health_url": "/health"
-    }
+        "docs_url": "/docs", # FastAPI default
+        "redoc_url": "/redoc", # FastAPI default
+        "health_url": "/health",
+        "active_configuration": {
+            "model_directory": MODEL_DIR,
+            "target_device": DEVICE, # This is the configured device, actual might differ if fallback
+            "table_format": TABLE_FORMAT.name,
+            "log_level": LOG_LEVEL_STR,
+            "processor_debug_output": DEBUG_DIR_PATH_STR if DEBUG_DIR_PATH_STR else "Disabled"
+        }
+    }
+
+# --- Main execution for local testing (optional) ---
+if __name__ == "__main__":
+    # This block is for running the app directly with uvicorn for local development.
+    # It's not strictly necessary if you always run with `uvicorn main:app`.
+    import uvicorn
+    logger.info("Starting Uvicorn server for local development...")
+    uvicorn.run(app, host="0.0.0.0", port=8000, log_level=LOG_LEVEL_STR.lower())