refactor: rename upload_to_s3 to process_document

- Rename upload_to_s3.py to process_document.py to better reflect its purpose
- Update all import statements across the codebase to use new module name
- Remove S3-specific code and references
- Keep the core document processing logic intact
- Update docstrings and comments to reflect new functionality

This change is part of removing AWS S3 dependencies and simplifying the
document processing pipeline.
This commit is contained in:
Christian Krakau-Louis
2025-03-27 14:50:51 +01:00
parent e1caa144f2
commit de06fd1286
7 changed files with 42 additions and 52 deletions
+1 -1
View File
@@ -9,7 +9,7 @@ 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.upload_to_s3 import upload_to_s3
from app.tasks.process_document import process_document # Updated import
from app.tasks.process_with_textract import process_with_textract
from app.tasks.refine_text_with_gpt import refine_text_with_gpt
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
+7 -7
View File
@@ -13,7 +13,7 @@ from pathlib import Path
from app.database import init_db
from app.config import settings
from app.tasks.upload_to_s3 import upload_to_s3
from app.tasks.process_document import process_document # Updated import
from app.tasks.upload_to_dropbox import upload_to_dropbox
from app.tasks.upload_to_paperless import upload_to_paperless
from app.tasks.upload_to_nextcloud import upload_to_nextcloud
@@ -57,7 +57,7 @@ def on_startup():
def process(file_path: str):
"""
API Endpoint to start document processing.
This enqueues the first task (upload_to_s3), which handles the full pipeline.
This enqueues document processing which handles the full pipeline.
"""
if not os.path.isabs(file_path):
file_path = os.path.join(settings.workdir, file_path)
@@ -67,7 +67,7 @@ def process(file_path: str):
status_code=400, detail=f"File {file_path} not found."
)
task = upload_to_s3.delay(file_path)
task = process_document.delay(file_path) # Updated function call
return {"task_id": task.id, "status": "queued"}
@app.post("/send_to_dropbox/")
@@ -122,7 +122,7 @@ def send_to_all_destinations_endpoint(file_path: str):
@app.post("/processall")
def process_all_pdfs_in_workdir():
"""
Finds all .pdf files in <workdir> and enqueues them for upload_to_s3.
Finds all .pdf files in <workdir> and enqueues them for processing.
"""
target_dir = settings.workdir
if not os.path.exists(target_dir):
@@ -141,7 +141,7 @@ def process_all_pdfs_in_workdir():
task_ids = []
for pdf in pdf_files:
file_path = os.path.join(target_dir, pdf)
task = upload_to_s3.delay(file_path)
task = process_document.delay(file_path) # Updated function call
task_ids.append(task.id)
return {
@@ -152,7 +152,7 @@ def process_all_pdfs_in_workdir():
@app.post("/ui-upload")
async def ui_upload(file: UploadFile = File(...)):
"""Endpoint to accept a user-uploaded file and enqueue it to S3."""
"""Endpoint to accept a user-uploaded file and enqueue it for processing."""
workdir = "/workdir"
target_path = os.path.join(workdir, file.filename)
try:
@@ -165,7 +165,7 @@ async def ui_upload(file: UploadFile = File(...)):
detail=f"Failed to save file: {e}"
)
task = upload_to_s3.delay(target_path)
task = process_document.delay(target_path) # Updated function call
return {"task_id": task.id, "status": "queued"}
# Custom 404 - we can still return the Jinja2 template, or the old static file:
+5 -4
View File
@@ -41,8 +41,9 @@ class FileRecord(Base):
class ProcessingLog(Base):
__tablename__ = "processing_logs"
id = Column(Integer, primary_key=True, index=True)
file_id = Column(Integer, ForeignKey("files.id"))
step_name = Column(String) # e.g. "OCR", "convert_to_pdf", "upload_s3"
status = Column(String) # "success" / "failure"
message = Column(String) # error text or success note
file_id = Column(Integer, ForeignKey("files.id"), nullable=True) # Optional file association
task_id = Column(String, index=True) # Celery task ID
step_name = Column(String) # e.g., "OCR", "convert_to_pdf", "upload_s3"
status = Column(String) # "pending", "in_progress", "success", "failure"
message = Column(String, nullable=True) # Error text or success note
timestamp = Column(DateTime(timezone=True), server_default=func.now())
+2 -2
View File
@@ -5,7 +5,7 @@ import logging
import mimetypes
from celery import shared_task
from app.config import settings
from app.tasks.upload_to_s3 import upload_to_s3
from app.tasks.process_document import process_document # Updated import
logger = logging.getLogger(__name__)
@@ -64,7 +64,7 @@ def convert_to_pdf(file_path):
with open(converted_file_path, "wb") as out_file:
out_file.write(response.content)
logger.info(f"Converted file saved as PDF: {converted_file_path}")
upload_to_s3.delay(converted_file_path)
process_document.delay(converted_file_path) # Updated function call
return converted_file_path
else:
logger.error(f"Conversion failed for {file_path}. Status code: {response.status_code}")
+2 -2
View File
@@ -9,7 +9,7 @@ import re
from datetime import datetime, timedelta, timezone
from celery import shared_task
from app.config import settings
from app.tasks.upload_to_s3 import upload_to_s3
from app.tasks.process_document import process_document # Updated import
from app.tasks.convert_to_pdf import convert_to_pdf # new conversion task
logger = logging.getLogger(__name__)
@@ -306,7 +306,7 @@ def fetch_attachments_and_enqueue(email_message):
f.write(part.get_payload(decode=True))
if mime_type == "application/pdf":
upload_to_s3.delay(file_path)
process_document.delay(file_path) # Updated function call
logger.info("Enqueued PDF for upload: %s", filename)
elif mime_type in ALLOWED_MIME_TYPES:
# Enqueue conversion to PDF using the Gotenberg service.
@@ -2,7 +2,6 @@
import os
import uuid
import boto3
import shutil
import mimetypes
import fitz # PyMuPDF for checking embedded text
@@ -12,39 +11,24 @@ from app.tasks.retry_config import BaseTaskWithRetry
from app.tasks.process_with_textract import process_with_textract
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
from app.celery_app import celery
# NEW imports for the DB
from app.database import SessionLocal
from app.models import FileRecord
from app.utils import hash_file
# Initialize S3 client
s3_client = boto3.client(
"s3",
aws_access_key_id=settings.aws_access_key_id,
aws_secret_access_key=settings.aws_secret_access_key,
region_name=settings.aws_region,
)
@celery.task(base=BaseTaskWithRetry)
def upload_to_s3(original_local_file: str):
def process_document(original_local_file: str):
"""
Uploads a file to S3 with a UUID-based filename and triggers processing.
Process a document file and trigger appropriate text extraction.
Steps:
1. Check if we have a FileRecord entry (via SHA-256 hash). If found, skip re-processing.
2. If not found, insert a new DB row and continue with the pipeline:
- Copy file to /workdir/tmp
- Check for embedded text. If present, skip S3 and run local GPT extraction
- Otherwise, upload to S3 and queue Textract-based OCR
- Check for embedded text. If present, run local GPT extraction
- Otherwise, queue Textract-based OCR
"""
bucket_name = settings.s3_bucket_name
if not bucket_name:
print("[ERROR] S3 bucket name not set.")
return {"error": "Missing S3 bucket name"}
if not os.path.exists(original_local_file):
print(f"[ERROR] File {original_local_file} not found.")
return {"error": "File not found"}
@@ -102,7 +86,7 @@ def upload_to_s3(original_local_file: str):
pdf_doc.close()
if has_text:
print(f"[INFO] PDF {original_local_file} contains embedded text. Skipping Textract.")
print(f"[INFO] PDF {original_local_file} contains embedded text. Processing locally.")
# Extract text locally
extracted_text = ""
@@ -113,20 +97,8 @@ def upload_to_s3(original_local_file: str):
# Call metadata extraction directly
extract_metadata_with_gpt.delay(new_filename, extracted_text)
return {"file": new_local_path, "status": "Text extracted locally"}
# 3. If no embedded text, upload to S3 and queue Textract processing
try:
print(f"[INFO] Uploading {new_local_path} to s3://{bucket_name}/{new_filename}...")
s3_client.upload_file(new_local_path, bucket_name, new_filename)
print(f"[INFO] File uploaded successfully: {new_filename}")
# Trigger Textract processing if no embedded text was found
process_with_textract.delay(new_filename)
return {"file": new_local_path, "s3_key": new_filename, "status": "Uploaded to S3 for OCR"}
except Exception as e:
print(f"[ERROR] Failed to upload {new_local_path} to S3: {e}")
return {"error": str(e)}
# 3. If no embedded text, queue Textract processing
process_with_textract.delay(new_filename)
return {"file": new_local_path, "status": "Queued for OCR"}
+17
View File
@@ -1,5 +1,7 @@
# app/utils.py
import hashlib
from app.database import SessionLocal
from app.models import ProcessingLog
def hash_file(filepath, chunk_size=65536):
"""
@@ -14,3 +16,18 @@ def hash_file(filepath, chunk_size=65536):
break
sha256.update(data)
return sha256.hexdigest()
def log_task_progress(task_id, step_name, status, message=None, file_id=None):
"""
Logs the progress of a Celery task to the database.
"""
with SessionLocal() as db:
log_entry = ProcessingLog(
task_id=task_id,
step_name=step_name,
status=status,
message=message,
file_id=file_id,
)
db.add(log_entry)
db.commit()