From 9823935927a9a326ecb80e23db8fcb0b75b81333 Mon Sep 17 00:00:00 2001
From: Christian Krakau-Louis <christianlouis@gmail.com>
Date: Wed, 9 Apr 2025 07:29:42 +0200
Subject: [PATCH] feat: Replace PyMuPDF with PyPDF2 for PDF metadata editing
 and text extraction

---
 app/tasks/embed_metadata_into_pdf.py | 35 +++++++++++++++++-----------
 app/tasks/process_document.py        | 20 +++++++++-------
 requirements.txt                     |  3 +--
 3 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/app/tasks/embed_metadata_into_pdf.py b/app/tasks/embed_metadata_into_pdf.py
index d5d59254..6828e9a0 100644
--- a/app/tasks/embed_metadata_into_pdf.py
+++ b/app/tasks/embed_metadata_into_pdf.py
@@ -2,7 +2,7 @@
 
 import os
 import shutil
-import fitz  # PyMuPDF for PDF metadata editing
+import PyPDF2  # Replace fitz with PyPDF2
 import json
 from app.config import settings
 from app.tasks.retry_config import BaseTaskWithRetry
@@ -51,7 +51,6 @@ def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata:
     After processing, the file is moved to
       <workdir>/processed/<suggested_filename.pdf>
     where <suggested_filename.pdf> is derived from metadata["filename"].
-    The output PDF is saved incrementally while preserving its original encryption.
     Additionally, the metadata is persisted to a JSON file with the same base name.
     """
     # Check for file existence; if not found, try the known shared tmp directory.
@@ -74,18 +73,26 @@ def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata:
     try:
         print(f"[DEBUG] Embedding metadata into {processed_file}...")
 
-        # Open the PDF
-        doc = fitz.open(processed_file)
-        # Set PDF metadata using only the standard keys.
-        doc.set_metadata({
-            "title": metadata.get("filename", "Unknown Document"),
-            "author": metadata.get("absender", "Unknown"),
-            "subject": metadata.get("document_type", "Unknown"),
-            "keywords": ", ".join(metadata.get("tags", []))
-        })
-        # Save incrementally and preserve encryption
-        doc.save(processed_file, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
-        doc.close()
+        # Open the PDF and modify metadata
+        with open(processed_file, 'rb') as file:
+            pdf_reader = PyPDF2.PdfReader(file)
+            pdf_writer = PyPDF2.PdfWriter()
+            
+            # Copy all pages from the reader to the writer
+            for page in pdf_reader.pages:
+                pdf_writer.add_page(page)
+            
+            # Set PDF metadata
+            pdf_writer.add_metadata({
+                "/Title": metadata.get("filename", "Unknown Document"),
+                "/Author": metadata.get("absender", "Unknown"),
+                "/Subject": metadata.get("document_type", "Unknown"),
+                "/Keywords": ", ".join(metadata.get("tags", []))
+            })
+            
+            # Write the modified PDF
+            with open(processed_file, 'wb') as output_file:
+                pdf_writer.write(output_file)
 
         print(f"[INFO] Metadata embedded successfully in {processed_file}")
 
diff --git a/app/tasks/process_document.py b/app/tasks/process_document.py
index 880264d2..fa795600 100644
--- a/app/tasks/process_document.py
+++ b/app/tasks/process_document.py
@@ -4,7 +4,7 @@ import os
 import uuid
 import shutil
 import mimetypes
-import fitz  # PyMuPDF for checking embedded text
+import PyPDF2  # Replace fitz with PyPDF2
 
 from app.config import settings
 from app.tasks.retry_config import BaseTaskWithRetry
@@ -81,19 +81,23 @@ def process_document(original_local_file: str):
         db.commit()
 
     # 2. Check for embedded text (outside the DB session to avoid long open transactions)
-    pdf_doc = fitz.open(new_local_path)
-    has_text = any(page.get_text() for page in pdf_doc)
-    pdf_doc.close()
+    with open(new_local_path, 'rb') as file:
+        pdf_reader = PyPDF2.PdfReader(file)
+        has_text = False
+        for page in pdf_reader.pages:
+            if page.extract_text().strip():
+                has_text = True
+                break
 
     if has_text:
         print(f"[INFO] PDF {original_local_file} contains embedded text. Processing locally.")
 
         # Extract text locally
         extracted_text = ""
-        pdf_doc = fitz.open(new_local_path)
-        for page in pdf_doc:
-            extracted_text += page.get_text("text") + "\n"
-        pdf_doc.close()
+        with open(new_local_path, 'rb') as file:
+            pdf_reader = PyPDF2.PdfReader(file)
+            for page in pdf_reader.pages:
+                extracted_text += page.extract_text() + "\n"
 
         # Call metadata extraction directly
         extract_metadata_with_gpt.delay(new_filename, extracted_text)
diff --git a/requirements.txt b/requirements.txt
index eebca058..7457c3f6 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -5,8 +5,7 @@ redis  # Message broker for Celery
 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 and now also for rotation
+PyPDF2>=3.0.0  # PDF processing for text extraction, metadata editing and rotation (replaces PyMuPDF)
 requests  # HTTP client
 dropbox>=11.36.0  # Dropbox integration
 azure-ai-documentintelligence  # Azure OCR service