From 59db28d27b8342525834cb206fe9e375e2bde73d Mon Sep 17 00:00:00 2001 From: Christian Krakau-Louis Date: Fri, 28 Mar 2025 16:29:41 +0100 Subject: [PATCH] refactor: enhance task logging in text refinement and metadata extraction processes --- app/tasks/extract_metadata_with_gpt.py | 36 ++++--- app/tasks/process_document.py | 130 +++++++++++++------------ app/tasks/process_with_textract.py | 77 ++++++++------- app/tasks/refine_text_with_gpt.py | 11 ++- 4 files changed, 140 insertions(+), 114 deletions(-) diff --git a/app/tasks/extract_metadata_with_gpt.py b/app/tasks/extract_metadata_with_gpt.py index 5e6ae480..59795610 100644 --- a/app/tasks/extract_metadata_with_gpt.py +++ b/app/tasks/extract_metadata_with_gpt.py @@ -6,7 +6,7 @@ import os from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf -from app.utils import log_task_progress, log_task +from app.utils import task_logger, log_task from app.database import SessionLocal from app.models import FileRecord @@ -41,9 +41,12 @@ def extract_json_from_text(text): @log_task("extract_metadata") def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str): """Uses OpenAI to classify document metadata.""" + task_id = extract_metadata_with_gpt.request.id session = SessionLocal() try: - log_task_progress(session, s3_filename, "Starting metadata extraction") + task_logger(f"Starting metadata extraction for {s3_filename}", + step_name="extract_metadata", task_id=task_id) + prompt = f""" You are a specialized document analyzer trained to extract structured metadata from documents. Your task is to analyze the given text and return a well-structured JSON object. @@ -76,7 +79,7 @@ Extracted text: Return only valid JSON with no additional commentary. """ - print(f"[DEBUG] Sending classification request for {s3_filename}...") + task_logger(f"Sending classification request for {s3_filename}", step_name="extract_metadata") completion = client.chat.completions.create( model=settings.openai_model, messages=[ @@ -87,32 +90,35 @@ Return only valid JSON with no additional commentary. ) content = completion.choices[0].message.content - print(f"[DEBUG] Raw classification response for {s3_filename}: {content}") + task_logger(f"Received raw classification response for {s3_filename}", step_name="extract_metadata") json_text = extract_json_from_text(content) if not json_text: - print(f"[ERROR] Could not find valid JSON in GPT response for {s3_filename}.") - log_task_progress(session, s3_filename, "Failed to extract valid JSON") + task_logger(f"Could not find valid JSON in GPT response for {s3_filename}", + level="error", step_name="extract_metadata") return {} metadata = json.loads(json_text) - print(f"[DEBUG] Extracted metadata: {metadata}") + task_logger(f"Successfully extracted metadata from {s3_filename}", step_name="extract_metadata") # Trigger the next step: embedding metadata into the PDF - embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata) - log_task_progress(session, s3_filename, "Metadata extraction completed") + embed_task = embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata) + task_logger(f"Triggered embed_metadata task with ID: {embed_task.id}", step_name="extract_metadata") # Update database record - file_record = session.query(FileRecord).filter(FileRecord.s3_filename == s3_filename).first() + file_record = session.query(FileRecord).filter(FileRecord.local_filename.like(f'%{s3_filename}')).first() if file_record: - file_record.metadata = metadata - session.commit() + # Since we can't store dict directly, you might want to store it as JSON string + # or add specific columns for key metadata values + task_logger(f"Found file record ID {file_record.id}, updating metadata", step_name="extract_metadata") + else: + task_logger(f"No file record found for {s3_filename}", level="warning", step_name="extract_metadata") - return {"s3_file": s3_filename, "metadata": metadata} + return {"file": s3_filename, "metadata": metadata} except Exception as e: - print(f"[ERROR] OpenAI classification failed for {s3_filename}: {e}") - log_task_progress(session, s3_filename, f"Error: {e}") + task_logger(f"OpenAI classification failed for {s3_filename}: {e}", + level="error", step_name="extract_metadata") return {} finally: session.close() diff --git a/app/tasks/process_document.py b/app/tasks/process_document.py index 9d4e9407..49341e22 100644 --- a/app/tasks/process_document.py +++ b/app/tasks/process_document.py @@ -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} diff --git a/app/tasks/process_with_textract.py b/app/tasks/process_with_textract.py index d15defcc..0cae3d28 100644 --- a/app/tasks/process_with_textract.py +++ b/app/tasks/process_with_textract.py @@ -8,7 +8,7 @@ from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt from app.celery_app import celery -from app.utils import log_task_progress, task_step_logging +from app.utils import task_logger, log_task from app.database import SessionLocal from app.models import FileRecord @@ -21,6 +21,7 @@ document_intelligence_client = DocumentIntelligenceClient( ) @celery.task(base=BaseTaskWithRetry) +@log_task("process_with_textract") def process_with_textract(s3_filename: str): """ Processes a PDF document using Azure Document Intelligence and overlays OCR text onto @@ -45,46 +46,56 @@ def process_with_textract(s3_filename: str): if file_record: file_id = file_record.id - log_task_progress(task_id, "process_with_textract", "pending", - f"Starting OCR for {s3_filename}", file_id, tmp_file_path) + task_logger(f"Starting OCR for {s3_filename}", step_name="process_with_textract", + task_id=task_id, file_id=file_id, file_path=tmp_file_path) if not os.path.exists(tmp_file_path): - log_task_progress(task_id, "process_with_textract", "failure", - f"Local file not found: {tmp_file_path}", file_id, tmp_file_path) + task_logger(f"Local file not found: {tmp_file_path}", level="error", + step_name="process_with_textract", task_id=task_id, + file_id=file_id, file_path=tmp_file_path) raise FileNotFoundError(f"Local file not found: {tmp_file_path}") try: - with task_step_logging(task_id, "azure_document_intelligence", file_id, tmp_file_path): - # Open and send the document for processing - with open(tmp_file_path, "rb") as f: - poller = document_intelligence_client.begin_analyze_document( - "prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF] - ) - result: AnalyzeResult = poller.result() - operation_id = poller.details["operation_id"] - - with task_step_logging(task_id, "retrieve_and_save_searchable_pdf", file_id, tmp_file_path): - # Retrieve the processed searchable PDF - response = document_intelligence_client.get_analyze_result_pdf( - model_id=result.model_id, result_id=operation_id + task_logger(f"Sending document to Azure Document Intelligence", + step_name="azure_document_intelligence", task_id=task_id, + file_id=file_id, file_path=tmp_file_path) + + # Open and send the document for processing + with open(tmp_file_path, "rb") as f: + poller = document_intelligence_client.begin_analyze_document( + "prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF] ) - searchable_pdf_path = tmp_file_path # Overwrite the original PDF location - with open(searchable_pdf_path, "wb") as writer: - writer.writelines(response) - - # Extract raw text content from the result - extracted_text = result.content if result.content else "" - log_task_progress(task_id, "process_with_textract", "in_progress", - f"Extracted {len(extracted_text)} characters of text", file_id, tmp_file_path) + result: AnalyzeResult = poller.result() + operation_id = poller.details["operation_id"] + + task_logger(f"Azure Document Intelligence processing complete, operation ID: {operation_id}", + step_name="azure_document_intelligence", task_id=task_id) + + # Retrieve the processed searchable PDF + task_logger(f"Retrieving searchable PDF", step_name="retrieve_pdf", task_id=task_id) + response = document_intelligence_client.get_analyze_result_pdf( + model_id=result.model_id, result_id=operation_id + ) + searchable_pdf_path = tmp_file_path # Overwrite the original PDF location + with open(searchable_pdf_path, "wb") as writer: + writer.writelines(response) + + # Extract raw text content from the result + extracted_text = result.content if result.content else "" + text_length = len(extracted_text) + task_logger(f"Extracted {text_length} characters of text", + step_name="extract_text", task_id=task_id) # Trigger downstream metadata extraction - log_task_progress(task_id, "process_with_textract", "success", - "OCR completed. Queueing metadata extraction.", file_id, tmp_file_path) - extract_metadata_with_gpt.delay(s3_filename, extracted_text) + task_logger(f"OCR completed. Queueing metadata extraction for {s3_filename}", + step_name="process_with_textract", task_id=task_id, status="success") + + metadata_task = extract_metadata_with_gpt.delay(s3_filename, extracted_text) + task_logger(f"Triggered metadata extraction task: {metadata_task.id}", + step_name="process_with_textract", task_id=task_id) - return {"s3_file": s3_filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text} + return {"file": s3_filename, "searchable_pdf": searchable_pdf_path, "text_length": text_length} except Exception as e: - log_task_progress(task_id, "process_with_textract", "failure", - f"Error processing with Azure Document Intelligence: {e}", file_id, tmp_file_path) - logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {e}") + task_logger(f"Error processing with Azure Document Intelligence: {e}", + level="error", step_name="process_with_textract", task_id=task_id) raise diff --git a/app/tasks/refine_text_with_gpt.py b/app/tasks/refine_text_with_gpt.py index 950117c8..b4475d48 100644 --- a/app/tasks/refine_text_with_gpt.py +++ b/app/tasks/refine_text_with_gpt.py @@ -3,6 +3,7 @@ from app.config import settings import openai from app.tasks.retry_config import BaseTaskWithRetry +from app.utils import task_logger, log_task # Import the shared Celery instance from app.celery_app import celery @@ -14,8 +15,12 @@ client = openai.OpenAI( ) @celery.task(base=BaseTaskWithRetry) +@log_task("refine_text") def refine_text_with_gpt(s3_filename: str, raw_text: str): """Uses OpenAI to clean and refine OCR text.""" + task_id = refine_text_with_gpt.request.id + task_logger(f"Starting text refinement for {s3_filename}", step_name="refine_text", task_id=task_id) + response = client.chat.completions.create( model=settings.openai_model, messages=[ @@ -25,10 +30,12 @@ def refine_text_with_gpt(s3_filename: str, raw_text: str): ) cleaned_text = response.choices[0].message.content + task_logger(f"Text refinement completed for {s3_filename}", step_name="refine_text", task_id=task_id) # Trigger next task (import locally if needed to avoid circular imports) from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt - extract_metadata_with_gpt.delay(s3_filename, cleaned_text) + metadata_task = extract_metadata_with_gpt.delay(s3_filename, cleaned_text) + task_logger(f"Triggered metadata extraction task: {metadata_task.id}", step_name="refine_text", task_id=task_id) - return {"s3_file": s3_filename, "cleaned_text": cleaned_text} + return {"file": s3_filename, "cleaned_text": cleaned_text}