#!/usr/bin/env python3 """Unified OCR processing task for DocuElevate. This task replaces the single-provider ``process_with_azure_document_intelligence`` task with a multi-engine OCR pipeline that: 1. Runs every OCR provider listed in ``OCR_PROVIDERS`` (default: ``azure``). 2. Merges/cross-checks the results using the configured AI model when more than one provider is active (see ``OCR_MERGE_STRATEGY``). 3. Writes the best searchable PDF back to the working directory. 4. Hands off to the page-rotation and metadata-extraction pipeline exactly as the legacy Azure task did. """ import logging import os from typing import Optional from app.celery_app import celery from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.rotate_pdf_pages import rotate_pdf_pages from app.utils import log_task_progress from app.utils.ocr_provider import OCRResult, get_ocr_providers, merge_ocr_results logger = logging.getLogger(__name__) @celery.task(base=BaseTaskWithRetry, bind=True) def process_with_ocr(self, filename: str, file_id: Optional[int] = None): """Run the configured OCR providers on *filename* and continue the pipeline. When multiple OCR providers are configured the results are merged using the AI model (or a simpler strategy controlled by ``OCR_MERGE_STRATEGY``). Args: filename: Base name of the file inside ``/tmp/``. file_id: Optional database record ID passed through to downstream tasks. """ task_id = self.request.id log_task_progress( task_id, "process_with_ocr", "in_progress", f"Starting OCR for {filename}", file_id=file_id, ) 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}") providers = get_ocr_providers() provider_names = [p.name for p in providers] logger.info(f"[{task_id}] Running {len(providers)} OCR provider(s): {provider_names}") log_task_progress( task_id, "run_ocr_providers", "in_progress", f"Running OCR providers: {', '.join(provider_names)}", file_id=file_id, ) results = [] errors = [] for provider in providers: pname = provider.__class__.__name__ try: logger.info(f"[{task_id}] Running {pname} on {filename}") result: OCRResult = provider.process(tmp_file_path) results.append(result) logger.info(f"[{task_id}] {pname} extracted {len(result.text)} chars") except Exception as exc: logger.error(f"[{task_id}] {pname} failed for {filename}: {exc}") errors.append(f"{pname}: {exc}") if not results: error_summary = "; ".join(errors) log_task_progress( task_id, "run_ocr_providers", "failure", "All OCR providers failed", file_id=file_id, detail=error_summary, ) raise RuntimeError(f"All OCR providers failed for {filename}: {error_summary}") if errors: logger.warning(f"[{task_id}] Some OCR providers failed: {'; '.join(errors)}") log_task_progress( task_id, "run_ocr_providers", "success", f"{len(results)} of {len(providers)} OCR provider(s) succeeded", file_id=file_id, ) # Merge results (no-op when only one provider succeeded) extracted_text, searchable_pdf_path, rotation_data = merge_ocr_results(results, filename) logger.info( f"[{task_id}] Merged OCR text: {len(extracted_text)} chars, " f"pdf={'yes' if searchable_pdf_path else 'no'}, " f"rotations={len(rotation_data)}" ) log_task_progress( task_id, "process_with_ocr", "success", f"OCR complete for {filename}", file_id=file_id, detail=f"Extracted {len(extracted_text)} chars using {len(results)} provider(s)", ) # Continue pipeline: rotate pages (if needed), then extract metadata rotate_pdf_pages.delay(filename, extracted_text, rotation_data, file_id) return { "file": filename, "searchable_pdf": searchable_pdf_path or tmp_file_path, "cleaned_text": extracted_text, "providers_used": [r.provider for r in results], } except Exception as exc: logger.error(f"[{task_id}] OCR failed for {filename}: {exc}") log_task_progress( task_id, "process_with_ocr", "failure", f"OCR failed for {filename}", file_id=file_id, detail=str(exc), ) raise