feat: update version to 0.4.1-dev, refactor task imports, and replace Textract processing with Azure Document Intelligence

app/tasks/process_with_azure_document_intelligence.py

@@ -0,0 +1,103 @@
+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
+
+from app.config import settings
+from app.tasks.retry_config import BaseTaskWithRetry
+from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
+from app.celery_app import celery
+
+logger = logging.getLogger(__name__)
+
+# Initialize Azure Document Intelligence client
+document_intelligence_client = DocumentIntelligenceClient(
+    endpoint=settings.azure_endpoint,
+    credential=AzureKeyCredential(settings.azure_ai_key)
+)
+
+# 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
+
+@celery.task(base=BaseTaskWithRetry)
+def process_with_azure_document_intelligence(filename: str):
+    """
+    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. Extracts the text content for metadata processing.
+      5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
+    """
+    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"]
+
+        # 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 downstream metadata extraction
+        extract_metadata_with_gpt.delay(filename, extracted_text)
+
+        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