gh-christianlouis-docuelevate/app/tasks/process_with_azure_document_intelligence.py

import os
import logging
import PyPDF2
from azure.core.credentials import AzureKeyCredential
from azure.ai.documentintelligence import DocumentIntelligenceClient
from azure.ai.documentintelligence.models import AnalyzeOutputOption, AnalyzeResult
import azure.core.exceptions

from app.config import settings
from app.tasks.retry_config import BaseTaskWithRetry
from app.tasks.rotate_pdf_pages import rotate_pdf_pages
from app.celery_app import celery

logger = logging.getLogger(__name__)

# Initialize Azure Document Intelligence client with error handling
try:
    document_intelligence_client = DocumentIntelligenceClient(
        endpoint=settings.azure_endpoint,
        credential=AzureKeyCredential(settings.azure_ai_key)
    )
    logger.info("Azure Document Intelligence client initialized successfully")
except (ValueError, azure.core.exceptions.ClientAuthenticationError) as e:
    logger.error(f"Failed to initialize Azure Document Intelligence client: {e}")
    document_intelligence_client = None
except Exception as e:
    logger.error(f"Unexpected error initializing Azure Document Intelligence client: {e}")
    document_intelligence_client = None

# Azure Document Intelligence service limits for Standard S0 tier
AZURE_DOC_INTELLIGENCE_LIMITS = {
    "max_file_size_bytes": 500 * 1024 * 1024,  # 500 MB
    "max_pages": 2000,
}

def get_pdf_page_count(file_path):
    """Get the number of pages in a PDF file."""
    try:
        with open(file_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            return len(pdf_reader.pages)
    except Exception as e:
        logger.error(f"Error getting PDF page count: {e}")
        return None

def check_page_rotation(result, filename):
    """
    Checks if pages in the document are rotated and logs the rotation information.
    
    Args:
        result: The AnalyzeResult from Azure Document Intelligence API
        filename: The name of the file being processed
        
    Returns:
        dict: Dictionary mapping page indices (integers) to rotation angles
    """
    logger.error(f"Checking rotation for document: {filename}")
    rotation_data = {}
    
    if not hasattr(result, 'pages') or not result.pages:
        logger.error(f"No page information available for rotation check: {filename}")
        return rotation_data
        
    for i, page in enumerate(result.pages):
        if hasattr(page, 'angle'):
            rotation_angle = page.angle
            if rotation_angle != 0:
                logger.error(f"Page {i+1} is rotated by {rotation_angle} degrees")
                # Store page index as integer, not string
                rotation_data[i] = rotation_angle
            else:
                logger.error(f"Page {i+1} has no rotation (0 degrees)")
        else:
            logger.error(f"Page {i+1} rotation information not available")
            
    return rotation_data

@celery.task(base=BaseTaskWithRetry)
def process_with_azure_document_intelligence(filename: str, file_id: int = None):
    """
    Processes a PDF document using Azure Document Intelligence and overlays OCR text onto
    the local temporary file (stored under <workdir>/tmp).
    
    Steps:
      0. Verify the file meets Azure Document Intelligence service limits
      1. Uploads the document for OCR using Azure Document Intelligence.
      2. Retrieves the processed PDF with embedded text.
      3. Saves the OCR-processed PDF locally in the same location as before.
      4. Checks for page rotation and triggers page rotation if needed.
      5. Triggers downstream metadata extraction.
      
    Args:
        filename: Name of the file to process
        file_id: Optional file ID to pass through to subsequent tasks
    """
    try:
        tmp_file_path = os.path.join(settings.workdir, "tmp", filename)
        if not os.path.exists(tmp_file_path):
            raise FileNotFoundError(f"Local file not found: {tmp_file_path}")

        # Check file size against service limits
        file_size = os.path.getsize(tmp_file_path)
        if file_size > AZURE_DOC_INTELLIGENCE_LIMITS["max_file_size_bytes"]:
            error_msg = f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds Azure Document Intelligence limit of 500 MB"
            logger.error(error_msg)
            return {"error": error_msg, "file": filename, "status": "Failed - Size limit exceeded"}

        # For PDF files, check page count against service limits
        # "Fail open" approach: only reject if we're sure it exceeds the limit
        if filename.lower().endswith('.pdf'):
            page_count = get_pdf_page_count(tmp_file_path)
            if page_count is not None and page_count > AZURE_DOC_INTELLIGENCE_LIMITS["max_pages"]:
                error_msg = f"PDF page count ({page_count}) exceeds Azure Document Intelligence limit of 2000 pages"
                logger.error(error_msg)
                return {"error": error_msg, "file": filename, "status": "Failed - Page limit exceeded"}
            if page_count is None:
                logger.warning(f"Could not determine page count for {filename}, proceeding with processing anyway")

        logger.info(f"Processing {filename} with Azure Document Intelligence OCR.")

        # Open and send the document for processing
        with open(tmp_file_path, "rb") as f:
            poller = document_intelligence_client.begin_analyze_document(
                "prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF]
            )
        result: AnalyzeResult = poller.result()
        operation_id = poller.details["operation_id"]

        # Check and log page rotation information
        rotation_data = check_page_rotation(result, filename)

        # Retrieve the processed searchable PDF
        response = document_intelligence_client.get_analyze_result_pdf(
            model_id=result.model_id, result_id=operation_id
        )
        searchable_pdf_path = tmp_file_path  # Overwrite the original PDF location
        with open(searchable_pdf_path, "wb") as writer:
            writer.writelines(response)
        logger.info(f"Searchable PDF saved at: {searchable_pdf_path}")

        # Extract raw text content from the result
        extracted_text = result.content if result.content else ""
        logger.info(f"Extracted text for {filename}: {len(extracted_text)} characters")

        # Trigger page rotation task if rotation is detected, otherwise proceed to metadata extraction
        rotate_pdf_pages.delay(filename, extracted_text, rotation_data, file_id)

        return {"file": filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
    except Exception as e:
        logger.error(f"Error processing {filename} with Azure Document Intelligence: {e}")
        raise