diff --git a/VERSION b/VERSION index 25a5afd8..1351681f 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -0.4.0-dev +0.4.1-dev diff --git a/app/celery_worker.py b/app/celery_worker.py index 769a13c4..c2ea5f0c 100644 --- a/app/celery_worker.py +++ b/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_with_textract import process_with_textract +from app.tasks.process_document import process_document +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", diff --git a/app/tasks/__init__.py b/app/tasks/__init__.py index e69de29b..0717aa1f 100644 --- a/app/tasks/__init__.py +++ b/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 diff --git a/app/tasks/extract_metadata_with_gpt.py b/app/tasks/extract_metadata_with_gpt.py index 18c44aee..ea20440d 100644 --- a/app/tasks/extract_metadata_with_gpt.py +++ b/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 {} diff --git a/app/tasks/process_document.py b/app/tasks/process_document.py index d31e4199..95658aa5 100644 --- a/app/tasks/process_document.py +++ b/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 - process_with_textract.delay(new_filename) + # 3. If no embedded text, queue Azure Document Intelligence processing + process_with_azure_document_intelligence.delay(new_filename) return {"file": new_local_path, "status": "Queued for OCR"} diff --git a/app/tasks/process_with_textract.py b/app/tasks/process_with_azure_document_intelligence.py similarity index 50% rename from app/tasks/process_with_textract.py rename to app/tasks/process_with_azure_document_intelligence.py index d3319545..f0dbf562 100644 --- a/app/tasks/process_with_textract.py +++ b/app/tasks/process_with_azure_document_intelligence.py @@ -1,67 +1,103 @@ -import os -import logging -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) -) - -@celery.task(base=BaseTaskWithRetry) -def process_with_textract(s3_filename: str): - """ - Processes a PDF document using Azure Document Intelligence and overlays OCR text onto - the local temporary file (stored under /tmp). - - Steps: - 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", s3_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.") - - # 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 {s3_filename}: {len(extracted_text)} characters") - - # Trigger downstream metadata extraction - extract_metadata_with_gpt.delay(s3_filename, extracted_text) - - return {"s3_file": s3_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}") - raise +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 /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 diff --git a/app/tasks/refine_text_with_gpt.py b/app/tasks/refine_text_with_gpt.py index 950117c8..b01cce3a 100644 --- a/app/tasks/refine_text_with_gpt.py +++ b/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} diff --git a/requirements.txt b/requirements.txt index 4b75c542..9eb9a356 100644 --- a/requirements.txt +++ b/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