Update main.py

main.py

@@ -6,11 +6,12 @@ from pathlib import Path
 from typing import List, Union, Optional, Dict, Any
 
 from fastapi import FastAPI, File, UploadFile, HTTPException, status
-from pydantic import BaseModel, Field, field_validator
-from PIL import Image # For type hinting, though not directly used in API models
+from pydantic import BaseModel, Field # Removed field_validator as it wasn't used
+from PIL import Image # For type hinting
 
-# Import necessary components from the monolithic script
-from mdr_pdf_parse import (
+# --- CORRECTED IMPORT ---
+# Import directly from the monolithic script file name
+from magic_pdf_processor import (
     MagicPDFProcessor,
     MDRStructuredBlock,
     MDRTextBlock,
@@ -163,7 +164,7 @@ class MDRFigureBlockModel(MDRBasicBlockModel):
         )
 
 # Union type for the response model
-MDRStructuredBlockModel = Union[MDRTextBlockModel, MDRTableBlockModel, MDRFormulaBlockModel, MDRFigureBlockModel]
+MDRStructuredBlockModelAPI = Union[MDRTextBlockModel, MDRTableBlockModel, MDRFormulaBlockModel, MDRFigureBlockModel] # Renamed API Union type
 
 # --- FastAPI App ---
 app = FastAPI(
@@ -187,7 +188,7 @@ async def health_check():
     return {"status": "ok", "message": "MagicPDFProcessor is running."}
 
 @app.post("/process-pdf/",
-          response_model=List[MDRStructuredBlockModel], # Use the Union type
+          response_model=List[MDRStructuredBlockModelAPI], # Use the Union type
           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.")):
@@ -201,6 +202,7 @@ async def process_pdf_endpoint(file: UploadFile = File(..., description="The PDF
         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
     try:
         # Create a temporary directory if it doesn't exist
         temp_dir = Path("./temp_uploads")
@@ -215,7 +217,7 @@ async def process_pdf_endpoint(file: UploadFile = File(..., description="The PDF
     except Exception as e:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to save uploaded file: {e}")
 
-    extracted_blocks_api: List[MDRStructuredBlockModel] = []
+    extracted_blocks_api: List[MDRStructuredBlockModelAPI] = []
     start_process_time = time.time()
 
     try:
@@ -246,11 +248,12 @@ async def process_pdf_endpoint(file: UploadFile = File(..., description="The PDF
         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
-        try:
-            os.remove(temp_pdf_path)
-            print(f"Cleaned up temporary file: {temp_pdf_path}")
-        except OSError as e:
-            print(f"Warning: Could not remove temporary file {temp_pdf_path}: {e}")
+        if temp_pdf_path and os.path.exists(temp_pdf_path):
+            try:
+                os.remove(temp_pdf_path)
+                print(f"Cleaned up temporary file: {temp_pdf_path}")
+            except OSError as e:
+                print(f"Warning: Could not remove temporary file {temp_pdf_path}: {e}")
 
     return extracted_blocks_api
 
@@ -261,12 +264,4 @@ async def read_root():
         "message": "Welcome to the MagicDataReadiness PDF Processor API!",
         "docs_url": "/docs",
         "health_url": "/health"
-    }
-
-# --- Run with Uvicorn (for local testing) ---
-# This part is usually not included when deploying with Docker,
-# as Docker CMD handles running uvicorn.
-# if __name__ == "__main__":
-#     import uvicorn
-#     print("Starting Uvicorn server locally...")
-#     uvicorn.run(app, host="0.0.0.0", port=8000)
+    }