gitea-admin/gh-christianlouis-docuelevate
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: enhance task logging in text refinement and metadata extraction processes

Browse Source
This commit is contained in:
Christian Krakau-Louis
2025-03-28 16:29:41 +01:00
parent a92cace662
commit 59db28d27b
4 changed files with 140 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
app/tasks/extract_metadata_with_gpt.py
+21 -15
View File
@@ -6,7 +6,7 @@ import os
from app.config import settings from app.config import settings
from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.retry_config import BaseTaskWithRetry
from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf 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.database import SessionLocal
from app.models import FileRecord from app.models import FileRecord
@@ -41,9 +41,12 @@ def extract_json_from_text(text):
@log_task("extract_metadata") @log_task("extract_metadata")
def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str): def extract_metadata_with_gpt(s3_filename: str, cleaned_text: str):
"""Uses OpenAI to classify document metadata.""" """Uses OpenAI to classify document metadata."""
task_id = extract_metadata_with_gpt.request.id
session = SessionLocal() session = SessionLocal()
try: 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""" prompt = f"""
You are a specialized document analyzer trained to extract structured metadata from documents. 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. 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. 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( completion = client.chat.completions.create(
model=settings.openai_model, model=settings.openai_model,
messages=[ messages=[
@@ -87,32 +90,35 @@ Return only valid JSON with no additional commentary.
) )
content = completion.choices[0].message.content 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) json_text = extract_json_from_text(content)
if not json_text: if not json_text:
print(f"[ERROR] Could not find valid JSON in GPT response for {s3_filename}.") task_logger(f"Could not find valid JSON in GPT response for {s3_filename}",
log_task_progress(session, s3_filename, "Failed to extract valid JSON") level="error", step_name="extract_metadata")
return {} return {}
metadata = json.loads(json_text) 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 # Trigger the next step: embedding metadata into the PDF
embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata) embed_task = embed_metadata_into_pdf.delay(s3_filename, cleaned_text, metadata)
log_task_progress(session, s3_filename, "Metadata extraction completed") task_logger(f"Triggered embed_metadata task with ID: {embed_task.id}", step_name="extract_metadata")
# Update database record # 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: if file_record:
file_record.metadata = metadata # Since we can't store dict directly, you might want to store it as JSON string
session.commit() # 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: except Exception as e:
print(f"[ERROR] OpenAI classification failed for {s3_filename}: {e}") task_logger(f"OpenAI classification failed for {s3_filename}: {e}",
log_task_progress(session, s3_filename, f"Error: {e}") level="error", step_name="extract_metadata")
return {} return {}
finally: finally:
session.close() session.close()
app/tasks/process_document.py
+66 -64
View File
@@ -13,10 +13,11 @@ from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
from app.celery_app import celery from app.celery_app import celery
from app.database import SessionLocal from app.database import SessionLocal
from app.models import FileRecord 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) @celery.task(base=BaseTaskWithRetry)
@log_task("process_document")
def process_document(original_local_file: str): def process_document(original_local_file: str):
""" """
Process a document file and trigger appropriate text extraction. 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 - Otherwise, queue Textract-based OCR
""" """
task_id = process_document.request.id 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): 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"} return {"error": "File not found"}
# 0. Compute the file hash and check for duplicates # 0. Compute the file hash and check for duplicates
with task_step_logging(task_id, "compute_hash", file_path=original_local_file): task_logger(f"Computing hash for {original_local_file}", step_name="compute_hash", task_id=task_id)
filehash = hash_file(original_local_file) filehash = hash_file(original_local_file)
original_filename = os.path.basename(original_local_file) original_filename = os.path.basename(original_local_file)
file_size = os.path.getsize(original_local_file) file_size = os.path.getsize(original_local_file)
mime_type, _ = mimetypes.guess_type(original_local_file) mime_type, _ = mimetypes.guess_type(original_local_file)
if not mime_type: if not mime_type:
mime_type = "application/octet-stream" mime_type = "application/octet-stream"
# Acquire DB session in the task # Acquire DB session in the task
new_record = None new_record = None
with SessionLocal() as db: with SessionLocal() as db:
with task_step_logging(task_id, "check_duplicates", file_path=original_local_file): 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() existing = db.query(FileRecord).filter_by(filehash=filehash).one_or_none()
if existing: if existing:
log_task_progress(task_id, "process_document", "success", task_logger(f"Duplicate file detected (hash={filehash[:10]}...). Skipping processing.",
f"Duplicate file detected (hash={filehash[:10]}...). Skipping processing.", step_name="process_document", task_id=task_id, file_id=existing.id, status="success")
file_id=existing.id) return {
return { "status": "duplicate_file",
"status": "duplicate_file", "file_id": existing.id,
"file_id": existing.id, "detail": "File already processed."
"detail": "File already processed." }
}
# Not a duplicate -> insert a new record # Not a duplicate -> insert a new record
with task_step_logging(task_id, "create_file_record", file_path=original_local_file): task_logger(f"Creating file record for {original_local_file}", step_name="create_file_record", task_id=task_id)
new_record = FileRecord( new_record = FileRecord(
filehash=filehash, filehash=filehash,
original_filename=original_filename, original_filename=original_filename,
local_filename="", # Will fill in after we move it local_filename="", # Will fill in after we move it
file_size=file_size, file_size=file_size,
mime_type=mime_type, mime_type=mime_type,
) )
db.add(new_record) db.add(new_record)
db.commit() db.commit()
db.refresh(new_record) db.refresh(new_record)
# 1. Generate a UUID-based filename and place it in /workdir/tmp # 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): 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_ext = os.path.splitext(original_local_file)[1]
file_uuid = str(uuid.uuid4()) file_uuid = str(uuid.uuid4())
new_filename = f"{file_uuid}{file_ext}" new_filename = f"{file_uuid}{file_ext}"
tmp_dir = os.path.join(settings.workdir, "tmp") tmp_dir = os.path.join(settings.workdir, "tmp")
os.makedirs(tmp_dir, exist_ok=True) os.makedirs(tmp_dir, exist_ok=True)
new_local_path = os.path.join(tmp_dir, new_filename) new_local_path = os.path.join(tmp_dir, new_filename)
# Copy the file instead of moving it # Copy the file instead of moving it
shutil.copy(original_local_file, new_local_path) shutil.copy(original_local_file, new_local_path)
# Update the DB with final local filename # Update the DB with final local filename
new_record.local_filename = new_local_path new_record.local_filename = new_local_path
db.commit() db.commit()
# 2. Check for embedded text (outside the DB session to avoid long open transactions) # 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): 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) pdf_doc = fitz.open(new_local_path)
has_text = any(page.get_text() for page in pdf_doc) has_text = any(page.get_text() for page in pdf_doc)
pdf_doc.close() pdf_doc.close()
if has_text: if has_text:
log_task_progress(task_id, "process_document", "in_progress", task_logger(f"PDF {original_local_file} contains embedded text. Processing locally.",
f"PDF {original_local_file} contains embedded text. Processing locally.", step_name="process_document", task_id=task_id, file_id=new_record.id)
file_id=new_record.id)
# Extract text locally # Extract text locally
extracted_text = "" extracted_text = ""
with task_step_logging(task_id, "extract_text_locally", file_id=new_record.id, file_path=new_local_path): 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) pdf_doc = fitz.open(new_local_path)
for page in pdf_doc: for page in pdf_doc:
extracted_text += page.get_text("text") + "\n" extracted_text += page.get_text("text") + "\n"
pdf_doc.close() pdf_doc.close()
# Call metadata extraction directly # Call metadata extraction directly
log_task_progress(task_id, "process_document", "success", task_logger(f"Text extracted locally. Queuing for metadata extraction.",
"Text extracted locally. Queuing for metadata extraction.", step_name="process_document", task_id=task_id, file_id=new_record.id, status="success")
file_id=new_record.id) metadata_task = extract_metadata_with_gpt.delay(new_filename, extracted_text)
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} return {"file": new_local_path, "status": "Text extracted locally", "file_id": new_record.id}
# 3. If no embedded text, queue Textract processing # 3. If no embedded text, queue Textract processing
log_task_progress(task_id, "process_document", "success", task_logger(f"No embedded text found. Queuing for OCR.",
"No embedded text found. Queuing for OCR.", step_name="process_document", task_id=task_id, file_id=new_record.id, status="success")
file_id=new_record.id) ocr_task = process_with_textract.delay(new_filename)
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} return {"file": new_local_path, "status": "Queued for OCR", "file_id": new_record.id}
app/tasks/process_with_textract.py
+44 -33
View File
@@ -8,7 +8,7 @@ from app.config import settings
from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.retry_config import BaseTaskWithRetry
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt
from app.celery_app import celery 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.database import SessionLocal
from app.models import FileRecord from app.models import FileRecord
@@ -21,6 +21,7 @@ document_intelligence_client = DocumentIntelligenceClient(
) )
@celery.task(base=BaseTaskWithRetry) @celery.task(base=BaseTaskWithRetry)
@log_task("process_with_textract")
def process_with_textract(s3_filename: str): def process_with_textract(s3_filename: str):
""" """
Processes a PDF document using Azure Document Intelligence and overlays OCR text onto 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: if file_record:
file_id = file_record.id file_id = file_record.id
log_task_progress(task_id, "process_with_textract", "pending", task_logger(f"Starting OCR for {s3_filename}", step_name="process_with_textract",
f"Starting OCR for {s3_filename}", file_id, tmp_file_path) task_id=task_id, file_id=file_id, file_path=tmp_file_path)
if not os.path.exists(tmp_file_path): if not os.path.exists(tmp_file_path):
log_task_progress(task_id, "process_with_textract", "failure", task_logger(f"Local file not found: {tmp_file_path}", level="error",
f"Local file not found: {tmp_file_path}", file_id, tmp_file_path) 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}") raise FileNotFoundError(f"Local file not found: {tmp_file_path}")
try: try:
with task_step_logging(task_id, "azure_document_intelligence", file_id, tmp_file_path): task_logger(f"Sending document to Azure Document Intelligence",
# Open and send the document for processing step_name="azure_document_intelligence", task_id=task_id,
with open(tmp_file_path, "rb") as f: file_id=file_id, file_path=tmp_file_path)
poller = document_intelligence_client.begin_analyze_document(
"prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF] # Open and send the document for processing
) with open(tmp_file_path, "rb") as f:
result: AnalyzeResult = poller.result() poller = document_intelligence_client.begin_analyze_document(
operation_id = poller.details["operation_id"] "prebuilt-read", body=f, output=[AnalyzeOutputOption.PDF]
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
) )
searchable_pdf_path = tmp_file_path # Overwrite the original PDF location result: AnalyzeResult = poller.result()
with open(searchable_pdf_path, "wb") as writer: operation_id = poller.details["operation_id"]
writer.writelines(response)
task_logger(f"Azure Document Intelligence processing complete, operation ID: {operation_id}",
# Extract raw text content from the result step_name="azure_document_intelligence", task_id=task_id)
extracted_text = result.content if result.content else ""
log_task_progress(task_id, "process_with_textract", "in_progress", # Retrieve the processed searchable PDF
f"Extracted {len(extracted_text)} characters of text", file_id, tmp_file_path) 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 # Trigger downstream metadata extraction
log_task_progress(task_id, "process_with_textract", "success", task_logger(f"OCR completed. Queueing metadata extraction for {s3_filename}",
"OCR completed. Queueing metadata extraction.", file_id, tmp_file_path) step_name="process_with_textract", task_id=task_id, status="success")
extract_metadata_with_gpt.delay(s3_filename, extracted_text)
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: except Exception as e:
log_task_progress(task_id, "process_with_textract", "failure", task_logger(f"Error processing with Azure Document Intelligence: {e}",
f"Error processing with Azure Document Intelligence: {e}", file_id, tmp_file_path) level="error", step_name="process_with_textract", task_id=task_id)
logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {e}")
raise raise
app/tasks/refine_text_with_gpt.py
+9 -2
View File
@@ -3,6 +3,7 @@
from app.config import settings from app.config import settings
import openai import openai
from app.tasks.retry_config import BaseTaskWithRetry from app.tasks.retry_config import BaseTaskWithRetry
from app.utils import task_logger, log_task
# Import the shared Celery instance # Import the shared Celery instance
from app.celery_app import celery from app.celery_app import celery
@@ -14,8 +15,12 @@ client = openai.OpenAI(
) )
@celery.task(base=BaseTaskWithRetry) @celery.task(base=BaseTaskWithRetry)
@log_task("refine_text")
def refine_text_with_gpt(s3_filename: str, raw_text: str): def refine_text_with_gpt(s3_filename: str, raw_text: str):
"""Uses OpenAI to clean and refine OCR text.""" """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( response = client.chat.completions.create(
model=settings.openai_model, model=settings.openai_model,
messages=[ messages=[
@@ -25,10 +30,12 @@ def refine_text_with_gpt(s3_filename: str, raw_text: str):
) )
cleaned_text = response.choices[0].message.content 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) # Trigger next task (import locally if needed to avoid circular imports)
from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt 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}