Add PDF check before processing it

Author: gabriel-piles

docker-compose.yml

@@ -31,17 +31,25 @@ services:
       - ENVIRONMENT=${ENVIRONMENT:-development}
       - SENTRY_DSN=${SENTRY_DSN:-}
       - PYTHONUNBUFFERED=1
+      - USE_FAST=false
     network_mode: host
     depends_on:
       - mongo-pdf-layout
 
   worker-pdf-layout-gpu:
-    container_name: "worker-pdf-layout-no-gpu"
+    container_name: "worker-pdf-layout-gpu"
     entrypoint: [ "gunicorn", "-k", "uvicorn.workers.UvicornWorker", "--chdir", "./src", "app:app", "--bind", "0.0.0.0:5060", "--timeout", "10000"]
     image: ghcr.io/huridocs/pdf-document-layout-analysis:0.0.23
     init: true
     restart: unless-stopped
     network_mode: host
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: 1
+              capabilities: [ gpu ]
     volumes:
       - data:/app/xmls
 

src/drivers/queues_processor/start_queue_processor.py

@@ -18,6 +18,7 @@
     service_logger,
     QUEUES_NAMES,
 )
+from domain.PdfFile import PdfFile
 from domain.ResultMessage import ResultMessage
 from domain.Task import Task
 from use_cases.extract_segments_use_case import ocr_pdf, get_xml_name, extract_segments
@@ -33,6 +34,23 @@ def get_failed_results_message(task: Task, message: str) -> ResultMessage:
     )
 
 
+def is_valid_pdf(filepath):
+    try:
+        with open(filepath, "rb") as f:
+            header = f.read(5)
+            if header != b"%PDF-":
+                return False
+            f.seek(-1024, 2)
+            end_bytes = f.read(1024)
+            if b"%%EOF" not in end_bytes:
+                return False
+        return True
+    except FileNotFoundError:
+        return False
+    except Exception:
+        return False
+
+
 def ocr_pdf_task(task):
     ocr_pdf(task)
 
@@ -53,24 +71,26 @@ def process(message):
     try:
         task = Task(**message)
     except ValidationError:
-        service_logger.error(f"validation error: {message}", exc_info=True)
+        service_logger.error(f"The message was incorrectly formatted: {message}")
         return None
 
     try:
         service_logger.info(f"Processing Redis message: {message}")
 
+        if not is_valid_pdf(PdfFile(task.tenant).get_path(task.params.filename)):
+            extraction_message = get_failed_results_message(task, f"The file does not appear to be a valid PDF")
+            service_logger.info(extraction_message.model_dump_json())
+            return extraction_message.model_dump_json()
+
         if task.task == "ocr":
             return ocr_pdf_task(task)
 
         return process_task(task).model_dump_json()
-    except RuntimeError:
-        extraction_message = get_failed_results_message(task, "Error processing PDF")
-        service_logger.error(extraction_message.model_dump_json(), exc_info=True)
     except FileNotFoundError:
-        extraction_message = get_failed_results_message(task, "Error FileNotFoundError")
+        extraction_message = get_failed_results_message(task, "The PDF could not be found")
         service_logger.error(extraction_message.model_dump_json(), exc_info=True)
-    except Exception:
-        extraction_message = get_failed_results_message(task, "Error")
+    except (RuntimeError, Exception):
+        extraction_message = get_failed_results_message(task, "An unexpected error occurred")
         service_logger.error(extraction_message.model_dump_json(), exc_info=True)
 
     return extraction_message.model_dump_json()

src/tests/test_end_to_end.py

@@ -140,7 +140,7 @@ def test_error_ocr(self):
         self.assertEqual(False, extraction_message.success)
 
     @staticmethod
-    def get_redis_message() -> ResultMessage:
+    def get_redis_message() -> ResultMessage | None:
         queues_names = ["segmentation", "ocr"]
 
         for i in range(160):
@@ -151,3 +151,5 @@ def get_redis_message() -> ResultMessage:
                 if message:
                     queue.deleteMessage(id=message["id"]).execute()
                     return ResultMessage(**json.loads(message["message"]))
+
+        return None