Added first working version of the code. Processes

PDF files, no upload yet.
2025-02-11 19:42:23 +01:00
parent 22f8f60f12
commit 1ad9425102
17 changed files with 563 additions and 66 deletions
@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+
+import os
+import shutil
+import fitz  # PyMuPDF for PDF metadata editing
+import json
+from app.config import settings
+from app.tasks.retry_config import BaseTaskWithRetry
+from app.tasks.finalize_document_storage import finalize_document_storage
+
+# Import the shared Celery instance
+from app.celery_app import celery
+
+def unique_filepath(directory, base_filename, extension=".pdf"):
+    """
+    Returns a unique filepath in the specified directory.
+    If 'base_filename.pdf' exists, it will append an underscore and counter.
+    """
+    candidate = os.path.join(directory, base_filename + extension)
+    if not os.path.exists(candidate):
+        return candidate
+    counter = 1
+    while True:
+        candidate = os.path.join(directory, f"{base_filename}_{counter}{extension}")
+        if not os.path.exists(candidate):
+            return candidate
+        counter += 1
+
+def persist_metadata(metadata, final_pdf_path):
+    """
+    Saves the metadata dictionary to a JSON file with the same base name as the final PDF.
+    For example, if final_pdf_path is "/var/docparse/working/processed/MyFile.pdf",
+    the metadata will be saved as "/var/docparse/working/processed/MyFile.json".
+    """
+    base, _ = os.path.splitext(final_pdf_path)
+    json_path = base + ".json"
+    with open(json_path, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, ensure_ascii=False, indent=2)
+    return json_path
+
+@celery.task(base=BaseTaskWithRetry)
+def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata: dict):
+    """
+    Embeds extracted metadata into the PDF's standard metadata fields.
+    The mapping is as follows:
+      - title: uses the extracted metadata "filename"
+      - author: uses "absender" (or "Unknown" if missing)
+      - subject: uses "document_type" (or "Unknown")
+      - keywords: a comma‐separated list from the "tags" field
+
+    After processing, the file is moved to
+      /var/docparse/working/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 directory.
+    if not os.path.exists(local_file_path):
+        alt_path = os.path.join("/var/docparse/working/tmp", os.path.basename(local_file_path))
+        if os.path.exists(alt_path):
+            local_file_path = alt_path
+        else:
+            print(f"[ERROR] Local file {local_file_path} not found, cannot embed metadata.")
+            return {"error": "File not found"}
+
+    # Work on a safe copy in /tmp
+    tmp_dir = "/tmp"
+    original_file = local_file_path
+    processed_file = os.path.join(tmp_dir, f"processed_{os.path.basename(local_file_path)}")
+
+    # Create a safe copy to work on
+    shutil.copy(original_file, processed_file)
+
+    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()
+
+        print(f"[INFO] Metadata embedded successfully in {processed_file}")
+
+        # Use the suggested filename from metadata; if not provided, use the original basename.
+        suggested_filename = metadata.get("filename", os.path.splitext(os.path.basename(local_file_path))[0])
+        # Remove any extension and then add .pdf
+        suggested_filename = os.path.splitext(suggested_filename)[0]
+        # Define the final directory and ensure it exists.
+        final_dir = "/var/docparse/working/processed"
+        os.makedirs(final_dir, exist_ok=True)
+        # Get a unique filepath in case of collisions.
+        final_file_path = unique_filepath(final_dir, suggested_filename, extension=".pdf")
+
+        # Move the processed file using shutil.move to handle cross-device moves.
+        shutil.move(processed_file, final_file_path)
+
+        # Persist the metadata into a JSON file with the same base name.
+        json_path = persist_metadata(metadata, final_file_path)
+        print(f"[INFO] Metadata persisted to {json_path}")
+
+        # Trigger the next step: final storage.
+        finalize_document_storage.delay(original_file, final_file_path, metadata)
+
+        return {"file": final_file_path, "metadata_file": json_path, "status": "Metadata embedded"}
+
+    except Exception as e:
+        print(f"[ERROR] Failed to embed metadata into {processed_file}: {e}")
+        return {"error": str(e)}
+