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

VERSION

@@ -1 +1 @@
-0.4.0-dev
+0.4.1-dev

app/celery_worker.py

@@ -9,8 +9,8 @@ from app.celery_app import celery
 from app import tasks  # <— This imports app/tasks.py so Celery can register tasks
 
 # **Ensure all tasks are imported before Celery starts**
-from app.tasks.process_document import process_document  # Updated import
+from app.tasks.process_document import process_document
-from app.tasks.process_with_textract import process_with_textract
+from app.tasks.process_with_azure_document_intelligence import process_with_azure_document_intelligence
 from app.tasks.refine_text_with_gpt import refine_text_with_gpt
 from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
 from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf
@@ -47,7 +47,8 @@ celery.conf.beat_schedule = {
     "poll-inboxes-every-minute": {
         "task": "app.tasks.imap_tasks.pull_all_inboxes",
         "schedule": crontab(minute="*/1"),  # every 1 minute
-    },
+        "options": {"expires": 55},  # Ensure tasks don't pile up
+    } if (settings.imap1_host or settings.imap2_host) else None,
     # Add Uptime Kuma ping task if configured
     "ping-uptime-kuma": {
         "task": "app.tasks.uptime_kuma_tasks.ping_uptime_kuma",

app/tasks/__init__.py

@@ -0,0 +1,3 @@
+# Import tasks so they can be discovered by Celery
+from app.tasks.process_document import process_document
+from app.tasks.process_with_azure_document_intelligence import process_with_azure_document_intelligence

app/tasks/extract_metadata_with_gpt.py

@@ -34,7 +34,7 @@ def extract_json_from_text(text):
     return None
 
 @celery.task(base=BaseTaskWithRetry)
-def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str):
+def extract_metadata_with_gpt(filename: str, cleaned_text: str):
     """Uses OpenAI to classify document metadata."""
     prompt = f"""
 You are a specialized document analyzer trained to extract structured metadata from documents.
@@ -69,7 +69,7 @@ Return only valid JSON with no additional commentary.
 """
 
     try:
-        print(f"[DEBUG] Sending classification request for {s3_filename}...")
+        print(f"[DEBUG] Sending classification request for {filename}...")
         completion = client.chat.completions.create(
             model=settings.openai_model,
             messages=[
@@ -80,21 +80,21 @@ Return only valid JSON with no additional commentary.
         )
 
         content = completion.choices[0].message.content
-        print(f"[DEBUG] Raw classification response for {s3_filename}: {content}")
+        print(f"[DEBUG] Raw classification response for {filename}: {content}")
 
         json_text = extract_json_from_text(content)
         if not json_text:
-            print(f"[ERROR] Could not find valid JSON in GPT response for {s3_filename}.")
+            print(f"[ERROR] Could not find valid JSON in GPT response for {filename}.")
             return {}
 
         metadata = json.loads(json_text)
         print(f"[DEBUG] Extracted metadata: {metadata}")
 
         # Trigger the next step: embedding metadata into the PDF
-        embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata)
+        embed_metadata_into_pdf.delay(filename, cleaned_text, metadata)
 
-        return {"s3_file": s3_filename, "metadata": metadata}
+        return {"s3_file": filename, "metadata": metadata}
 
     except Exception as e:
-        print(f"[ERROR] OpenAI classification failed for {s3_filename}: {e}")
+        print(f"[ERROR] OpenAI classification failed for {filename}: {e}")
         return {}

app/tasks/process_document.py

@@ -8,7 +8,7 @@ import fitz  # PyMuPDF for checking embedded text
 
 from app.config import settings
 from app.tasks.retry_config import BaseTaskWithRetry
-from app.tasks.process_with_textract import process_with_textract
+from app.tasks.process_with_azure_document_intelligence import process_with_azure_document_intelligence
 from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
 from app.celery_app import celery
 from app.database import SessionLocal
@@ -99,6 +99,6 @@ def process_document(original_local_file: str):
         extract_metadata_with_gpt.delay(new_filename, extracted_text)
         return {"file": new_local_path, "status": "Text extracted locally"}
 
-    # 3. If no embedded text, queue Textract processing
+    # 3. If no embedded text, queue Azure Document Intelligence processing
-    process_with_textract.delay(new_filename)
+    process_with_azure_document_intelligence.delay(new_filename)
     return {"file": new_local_path, "status": "Queued for OCR"}

app/tasks/process_with_azure_document_intelligence.py

@@ -1,67 +1,103 @@
-import os
+import os
-import logging
+import logging
-from azure.core.credentials import AzureKeyCredential
+import PyPDF2
-from azure.ai.documentintelligence import DocumentIntelligenceClient
+from azure.core.credentials import AzureKeyCredential
-from azure.ai.documentintelligence.models import AnalyzeOutputOption, AnalyzeResult
+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.config import settings
-from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
+from app.tasks.retry_config import BaseTaskWithRetry
-from app.celery_app import celery
+from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
-
+from app.celery_app import celery
-logger = logging.getLogger(__name__)
+
-
+logger = logging.getLogger(__name__)
-# Initialize Azure Document Intelligence client
+
-document_intelligence_client = DocumentIntelligenceClient(
+# Initialize Azure Document Intelligence client
-    endpoint=settings.azure_endpoint,
+document_intelligence_client = DocumentIntelligenceClient(
-    credential=AzureKeyCredential(settings.azure_ai_key)
+    endpoint=settings.azure_endpoint,
-)
+    credential=AzureKeyCredential(settings.azure_ai_key)
-
+)
-@celery.task(base=BaseTaskWithRetry)
+
-def process_with_textract(s3_filename: str):
+# Azure Document Intelligence service limits for Standard S0 tier
-    """
+AZURE_DOC_INTELLIGENCE_LIMITS = {
-    Processes a PDF document using Azure Document Intelligence and overlays OCR text onto
+    "max_file_size_bytes": 500 * 1024 * 1024,  # 500 MB
-    the local temporary file (stored under <workdir>/tmp).
+    "max_pages": 2000,
-    
+}
-    Steps:
+
-      1. Uploads the document for OCR using Azure Document Intelligence.
+def get_pdf_page_count(file_path):
-      2. Retrieves the processed PDF with embedded text.
+    """Get the number of pages in a PDF file."""
-      3. Saves the OCR-processed PDF locally in the same location as before.
+    try:
-      4. Extracts the text content for metadata processing.
+        with open(file_path, 'rb') as file:
-      5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
+            pdf_reader = PyPDF2.PdfReader(file)
-    """
+            return len(pdf_reader.pages)
-    try:
+    except Exception as e:
-        tmp_file_path = os.path.join(settings.workdir, "tmp", s3_filename)
+        logger.error(f"Error getting PDF page count: {e}")
-        if not os.path.exists(tmp_file_path):
+        return None
-            raise FileNotFoundError(f"Local file not found: {tmp_file_path}")
+
-
+@celery.task(base=BaseTaskWithRetry)
-        logger.info(f"Processing {s3_filename} with Azure Document Intelligence OCR.")
+def process_with_azure_document_intelligence(filename: str):
-
+    """
-        # Open and send the document for processing
+    Processes a PDF document using Azure Document Intelligence and overlays OCR text onto
-        with open(tmp_file_path, "rb") as f:
+    the local temporary file (stored under <workdir>/tmp).
-            poller = document_intelligence_client.begin_analyze_document(
+    
-                "prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF]
+    Steps:
-            )
+      0. Verify the file meets Azure Document Intelligence service limits
-        result: AnalyzeResult = poller.result()
+      1. Uploads the document for OCR using Azure Document Intelligence.
-        operation_id = poller.details["operation_id"]
+      2. Retrieves the processed PDF with embedded text.
-
+      3. Saves the OCR-processed PDF locally in the same location as before.
-        # Retrieve the processed searchable PDF
+      4. Extracts the text content for metadata processing.
-        response = document_intelligence_client.get_analyze_result_pdf(
+      5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
-            model_id=result.model_id, result_id=operation_id
+    """
-        )
+    try:
-        searchable_pdf_path = tmp_file_path  # Overwrite the original PDF location
+        tmp_file_path = os.path.join(settings.workdir, "tmp", filename)
-        with open(searchable_pdf_path, "wb") as writer:
+        if not os.path.exists(tmp_file_path):
-            writer.writelines(response)
+            raise FileNotFoundError(f"Local file not found: {tmp_file_path}")
-        logger.info(f"Searchable PDF saved at: {searchable_pdf_path}")
+
-
+        # Check file size against service limits
-        # Extract raw text content from the result
+        file_size = os.path.getsize(tmp_file_path)
-        extracted_text = result.content if result.content else ""
+        if file_size > AZURE_DOC_INTELLIGENCE_LIMITS["max_file_size_bytes"]:
-        logger.info(f"Extracted text for {s3_filename}: {len(extracted_text)} characters")
+            error_msg = f"File size ({file_size / (1024 * 1024):.2f} MB) exceeds Azure Document Intelligence limit of 500 MB"
-
+            logger.error(error_msg)
-        # Trigger downstream metadata extraction
+            return {"error": error_msg, "file": filename, "status": "Failed - Size limit exceeded"}
-        extract_metadata_with_gpt.delay(s3_filename, extracted_text)
+
-
+        # For PDF files, check page count against service limits
-        return {"s3_file": s3_filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
+        # "Fail open" approach: only reject if we're sure it exceeds the limit
-    except Exception as e:
+        if filename.lower().endswith('.pdf'):
-        logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {e}")
+            page_count = get_pdf_page_count(tmp_file_path)
-        raise
+            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

app/tasks/refine_text_with_gpt.py

@@ -14,7 +14,7 @@ client = openai.OpenAI(
 )
 
 @celery.task(base=BaseTaskWithRetry)
-def refine_text_with_gpt(s3_filename: str, raw_text: str):
+def refine_text_with_gpt(filename: str, raw_text: str):
     """Uses OpenAI to clean and refine OCR text."""
     response = client.chat.completions.create(
         model=settings.openai_model,
@@ -28,7 +28,7 @@ def refine_text_with_gpt(s3_filename: str, raw_text: str):
 
     # Trigger next task (import locally if needed to avoid circular imports)
     from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
-    extract_metadata_with_gpt.delay(s3_filename, cleaned_text)
+    extract_metadata_with_gpt.delay(filename, cleaned_text)
 
-    return {"s3_file": s3_filename, "cleaned_text": cleaned_text}
+    return {"s3_file": filename, "cleaned_text": cleaned_text}
 

requirements.txt

@@ -6,6 +6,7 @@ sqlalchemy  # Database ORM
 pydantic  # Data validation
 openai  # GPT integration for metadata extraction
 pymupdf  # PDF processing, text extraction, and detection (imported as 'fitz')
+PyPDF2  # PDF processing for page counting
 requests  # HTTP client
 dropbox  # Dropbox integration
 azure-ai-documentintelligence  # Azure OCR service