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,5 +1,6 @@
 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
@@ -17,13 +18,30 @@ document_intelligence_client = DocumentIntelligenceClient(
     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_textract(s3_filename: str):
+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.
@@ -31,11 +49,29 @@ def process_with_textract(s3_filename: str):
       5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
     """
     try:
-        tmp_file_path = os.path.join(settings.workdir, "tmp", s3_filename)
+        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}")
 
-        logger.info(f"Processing {s3_filename} with Azure Document Intelligence OCR.")
+        # 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:
@@ -56,12 +92,12 @@ def process_with_textract(s3_filename: str):
 
         # Extract raw text content from the result
         extracted_text = result.content if result.content else ""
-        logger.info(f"Extracted text for {s3_filename}: {len(extracted_text)} characters")
+        logger.info(f"Extracted text for {filename}: {len(extracted_text)} characters")
 
         # Trigger downstream metadata extraction
-        extract_metadata_with_gpt.delay(s3_filename, extracted_text)
+        extract_metadata_with_gpt.delay(filename, extracted_text)
 
-        return {"s3_file": s3_filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
+        return {"file": filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
     except Exception as e:
-        logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {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