refactor: enhance task logging in text refinement and metadata extraction processes
This commit is contained in:
Christian Krakau-Louis
@@ -13,10 +13,11 @@ from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
|
||||
from app.celery_app import celery
|
||||
from app.database import SessionLocal
|
||||
from app.models import FileRecord
|
||||
from app.utils import hash_file, log_task_progress, task_step_logging
|
||||
from app.utils import hash_file, task_logger, log_task
|
||||
|
||||
|
||||
@celery.task(base=BaseTaskWithRetry)
|
||||
@log_task("process_document")
|
||||
def process_document(original_local_file: str):
|
||||
"""
|
||||
Process a document file and trigger appropriate text extraction.
|
||||
@@ -29,95 +30,96 @@ def process_document(original_local_file: str):
|
||||
- Otherwise, queue Textract-based OCR
|
||||
"""
|
||||
task_id = process_document.request.id
|
||||
log_task_progress(task_id, "process_document", "pending", f"Processing {original_local_file}", file_path=original_local_file)
|
||||
task_logger(f"Processing {original_local_file}", step_name="process_document", task_id=task_id, file_path=original_local_file)
|
||||
|
||||
if not os.path.exists(original_local_file):
|
||||
log_task_progress(task_id, "process_document", "failure", f"File {original_local_file} not found.", file_path=original_local_file)
|
||||
task_logger(f"File {original_local_file} not found.", level="error", step_name="process_document", task_id=task_id)
|
||||
return {"error": "File not found"}
|
||||
|
||||
# 0. Compute the file hash and check for duplicates
|
||||
with task_step_logging(task_id, "compute_hash", file_path=original_local_file):
|
||||
filehash = hash_file(original_local_file)
|
||||
original_filename = os.path.basename(original_local_file)
|
||||
file_size = os.path.getsize(original_local_file)
|
||||
mime_type, _ = mimetypes.guess_type(original_local_file)
|
||||
if not mime_type:
|
||||
mime_type = "application/octet-stream"
|
||||
task_logger(f"Computing hash for {original_local_file}", step_name="compute_hash", task_id=task_id)
|
||||
filehash = hash_file(original_local_file)
|
||||
original_filename = os.path.basename(original_local_file)
|
||||
file_size = os.path.getsize(original_local_file)
|
||||
mime_type, _ = mimetypes.guess_type(original_local_file)
|
||||
if not mime_type:
|
||||
mime_type = "application/octet-stream"
|
||||
|
||||
# Acquire DB session in the task
|
||||
new_record = None
|
||||
with SessionLocal() as db:
|
||||
with task_step_logging(task_id, "check_duplicates", file_path=original_local_file):
|
||||
existing = db.query(FileRecord).filter_by(filehash=filehash).one_or_none()
|
||||
if existing:
|
||||
log_task_progress(task_id, "process_document", "success",
|
||||
f"Duplicate file detected (hash={filehash[:10]}...). Skipping processing.",
|
||||
file_id=existing.id)
|
||||
return {
|
||||
"status": "duplicate_file",
|
||||
"file_id": existing.id,
|
||||
"detail": "File already processed."
|
||||
}
|
||||
task_logger(f"Checking for duplicate files", step_name="check_duplicates", task_id=task_id)
|
||||
existing = db.query(FileRecord).filter_by(filehash=filehash).one_or_none()
|
||||
if existing:
|
||||
task_logger(f"Duplicate file detected (hash={filehash[:10]}...). Skipping processing.",
|
||||
step_name="process_document", task_id=task_id, file_id=existing.id, status="success")
|
||||
return {
|
||||
"status": "duplicate_file",
|
||||
"file_id": existing.id,
|
||||
"detail": "File already processed."
|
||||
}
|
||||
|
||||
# Not a duplicate -> insert a new record
|
||||
with task_step_logging(task_id, "create_file_record", file_path=original_local_file):
|
||||
new_record = FileRecord(
|
||||
filehash=filehash,
|
||||
original_filename=original_filename,
|
||||
local_filename="", # Will fill in after we move it
|
||||
file_size=file_size,
|
||||
mime_type=mime_type,
|
||||
)
|
||||
db.add(new_record)
|
||||
db.commit()
|
||||
db.refresh(new_record)
|
||||
task_logger(f"Creating file record for {original_local_file}", step_name="create_file_record", task_id=task_id)
|
||||
new_record = FileRecord(
|
||||
filehash=filehash,
|
||||
original_filename=original_filename,
|
||||
local_filename="", # Will fill in after we move it
|
||||
file_size=file_size,
|
||||
mime_type=mime_type,
|
||||
)
|
||||
db.add(new_record)
|
||||
db.commit()
|
||||
db.refresh(new_record)
|
||||
|
||||
# 1. Generate a UUID-based filename and place it in /workdir/tmp
|
||||
with task_step_logging(task_id, "copy_to_workdir", file_id=new_record.id, file_path=original_local_file):
|
||||
file_ext = os.path.splitext(original_local_file)[1]
|
||||
file_uuid = str(uuid.uuid4())
|
||||
new_filename = f"{file_uuid}{file_ext}"
|
||||
task_logger(f"Copying to workdir", step_name="copy_to_workdir", task_id=task_id, file_id=new_record.id)
|
||||
file_ext = os.path.splitext(original_local_file)[1]
|
||||
file_uuid = str(uuid.uuid4())
|
||||
new_filename = f"{file_uuid}{file_ext}"
|
||||
|
||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||
os.makedirs(tmp_dir, exist_ok=True)
|
||||
new_local_path = os.path.join(tmp_dir, new_filename)
|
||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||
os.makedirs(tmp_dir, exist_ok=True)
|
||||
new_local_path = os.path.join(tmp_dir, new_filename)
|
||||
|
||||
# Copy the file instead of moving it
|
||||
shutil.copy(original_local_file, new_local_path)
|
||||
# Copy the file instead of moving it
|
||||
shutil.copy(original_local_file, new_local_path)
|
||||
|
||||
# Update the DB with final local filename
|
||||
new_record.local_filename = new_local_path
|
||||
db.commit()
|
||||
# Update the DB with final local filename
|
||||
new_record.local_filename = new_local_path
|
||||
db.commit()
|
||||
|
||||
# 2. Check for embedded text (outside the DB session to avoid long open transactions)
|
||||
with task_step_logging(task_id, "check_embedded_text", file_id=new_record.id, file_path=new_local_path):
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
has_text = any(page.get_text() for page in pdf_doc)
|
||||
pdf_doc.close()
|
||||
task_logger(f"Checking for embedded text", step_name="check_embedded_text", task_id=task_id, file_id=new_record.id)
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
has_text = any(page.get_text() for page in pdf_doc)
|
||||
pdf_doc.close()
|
||||
|
||||
if has_text:
|
||||
log_task_progress(task_id, "process_document", "in_progress",
|
||||
f"PDF {original_local_file} contains embedded text. Processing locally.",
|
||||
file_id=new_record.id)
|
||||
task_logger(f"PDF {original_local_file} contains embedded text. Processing locally.",
|
||||
step_name="process_document", task_id=task_id, file_id=new_record.id)
|
||||
|
||||
# Extract text locally
|
||||
extracted_text = ""
|
||||
with task_step_logging(task_id, "extract_text_locally", file_id=new_record.id, file_path=new_local_path):
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
for page in pdf_doc:
|
||||
extracted_text += page.get_text("text") + "\n"
|
||||
pdf_doc.close()
|
||||
task_logger(f"Extracting text locally", step_name="extract_text_locally", task_id=task_id, file_id=new_record.id)
|
||||
pdf_doc = fitz.open(new_local_path)
|
||||
for page in pdf_doc:
|
||||
extracted_text += page.get_text("text") + "\n"
|
||||
pdf_doc.close()
|
||||
|
||||
# Call metadata extraction directly
|
||||
log_task_progress(task_id, "process_document", "success",
|
||||
"Text extracted locally. Queuing for metadata extraction.",
|
||||
file_id=new_record.id)
|
||||
extract_metadata_with_gpt.delay(new_filename, extracted_text)
|
||||
task_logger(f"Text extracted locally. Queuing for metadata extraction.",
|
||||
step_name="process_document", task_id=task_id, file_id=new_record.id, status="success")
|
||||
metadata_task = extract_metadata_with_gpt.delay(new_filename, extracted_text)
|
||||
task_logger(f"Triggered metadata extraction task: {metadata_task.id}",
|
||||
step_name="process_document", task_id=task_id)
|
||||
|
||||
return {"file": new_local_path, "status": "Text extracted locally", "file_id": new_record.id}
|
||||
|
||||
# 3. If no embedded text, queue Textract processing
|
||||
log_task_progress(task_id, "process_document", "success",
|
||||
"No embedded text found. Queuing for OCR.",
|
||||
file_id=new_record.id)
|
||||
process_with_textract.delay(new_filename)
|
||||
task_logger(f"No embedded text found. Queuing for OCR.",
|
||||
step_name="process_document", task_id=task_id, file_id=new_record.id, status="success")
|
||||
ocr_task = process_with_textract.delay(new_filename)
|
||||
task_logger(f"Triggered OCR task: {ocr_task.id}", step_name="process_document", task_id=task_id)
|
||||
|
||||
return {"file": new_local_path, "status": "Queued for OCR", "file_id": new_record.id}
|
||||
|
||||
Reference in New Issue
Block a user