diff --git a/app/celery_worker.py b/app/celery_worker.py index 926ee7e1..7b31b934 100644 --- a/app/celery_worker.py +++ b/app/celery_worker.py @@ -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 diff --git a/app/main.py b/app/main.py index 9dc08856..cff56a6e 100644 --- a/app/main.py +++ b/app/main.py @@ -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 and enqueues them for upload_to_s3. + Finds all .pdf files in 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: diff --git a/app/models.py b/app/models.py index ea07a074..05cd4a7a 100644 --- a/app/models.py +++ b/app/models.py @@ -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()) diff --git a/app/tasks/convert_to_pdf.py b/app/tasks/convert_to_pdf.py index f03d039e..2842c984 100644 --- a/app/tasks/convert_to_pdf.py +++ b/app/tasks/convert_to_pdf.py @@ -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}") diff --git a/app/tasks/imap_tasks.py b/app/tasks/imap_tasks.py index b41572c2..ba8dff63 100644 --- a/app/tasks/imap_tasks.py +++ b/app/tasks/imap_tasks.py @@ -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. diff --git a/app/tasks/upload_to_s3.py b/app/tasks/process_document.py similarity index 69% rename from app/tasks/upload_to_s3.py rename to app/tasks/process_document.py index 45e09ecb..d31e4199 100644 --- a/app/tasks/upload_to_s3.py +++ b/app/tasks/process_document.py @@ -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"} diff --git a/app/utils.py b/app/utils.py index 05d0026e..cfd6284f 100644 --- a/app/utils.py +++ b/app/utils.py @@ -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()