Fix file_id propagation through task chain
- Pass file_id as parameter through all task chains - Update process_document, extract_metadata_with_gpt, embed_metadata_into_pdf, finalize_document_storage, send_to_all_destinations, rotate_pdf_pages, and process_with_azure_document_intelligence - Remove unreliable LIKE queries, use explicit file_id passing instead Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
@@ -46,7 +46,7 @@ def persist_metadata(metadata, final_pdf_path):
|
|||||||
return json_path
|
return json_path
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||||
def embed_metadata_into_pdf(self, local_file_path: str, extracted_text: str, metadata: dict):
|
def embed_metadata_into_pdf(self, local_file_path: str, extracted_text: str, metadata: dict, file_id: int = None):
|
||||||
"""
|
"""
|
||||||
Embeds extracted metadata into the PDF's standard metadata fields.
|
Embeds extracted metadata into the PDF's standard metadata fields.
|
||||||
The mapping is as follows:
|
The mapping is as follows:
|
||||||
@@ -62,11 +62,14 @@ def embed_metadata_into_pdf(self, local_file_path: str, extracted_text: str, met
|
|||||||
"""
|
"""
|
||||||
task_id = self.request.id
|
task_id = self.request.id
|
||||||
logger.info(f"[{task_id}] Starting metadata embedding for: {local_file_path}")
|
logger.info(f"[{task_id}] Starting metadata embedding for: {local_file_path}")
|
||||||
log_task_progress(task_id, "embed_metadata_into_pdf", "in_progress", f"Embedding metadata into {os.path.basename(local_file_path)}")
|
log_task_progress(task_id, "embed_metadata_into_pdf", "in_progress", f"Embedding metadata into {os.path.basename(local_file_path)}", file_id=file_id)
|
||||||
|
|
||||||
# Get file_id from database
|
# Get file_id from database if not provided
|
||||||
file_id = None
|
if file_id is None:
|
||||||
# Check for file existence; if not found, try the known shared tmp directory.
|
with SessionLocal() as db:
|
||||||
|
file_record = db.query(FileRecord).filter_by(local_filename=local_file_path).first()
|
||||||
|
if file_record:
|
||||||
|
file_id = file_record.id
|
||||||
if not os.path.exists(local_file_path):
|
if not os.path.exists(local_file_path):
|
||||||
alt_path = os.path.join(settings.workdir, "tmp", os.path.basename(local_file_path))
|
alt_path = os.path.join(settings.workdir, "tmp", os.path.basename(local_file_path))
|
||||||
if os.path.exists(alt_path):
|
if os.path.exists(alt_path):
|
||||||
@@ -76,10 +79,8 @@ def embed_metadata_into_pdf(self, local_file_path: str, extracted_text: str, met
|
|||||||
log_task_progress(task_id, "embed_metadata_into_pdf", "failure", "File not found")
|
log_task_progress(task_id, "embed_metadata_into_pdf", "failure", "File not found")
|
||||||
return {"error": "File not found"}
|
return {"error": "File not found"}
|
||||||
|
|
||||||
with SessionLocal() as db:
|
|
||||||
file_record = db.query(FileRecord).filter_by(local_filename=local_file_path).first()
|
# Check for file existence; if not found, try the known shared tmp directory.
|
||||||
if file_record:
|
|
||||||
file_id = file_record.id
|
|
||||||
|
|
||||||
# Work on a safe copy in a secure temporary directory
|
# Work on a safe copy in a secure temporary directory
|
||||||
original_file = local_file_path
|
original_file = local_file_path
|
||||||
@@ -149,7 +150,7 @@ def embed_metadata_into_pdf(self, local_file_path: str, extracted_text: str, met
|
|||||||
# Trigger the next step: final storage.
|
# Trigger the next step: final storage.
|
||||||
logger.info(f"[{task_id}] Queueing final storage task")
|
logger.info(f"[{task_id}] Queueing final storage task")
|
||||||
log_task_progress(task_id, "embed_metadata_into_pdf", "success", "Metadata embedded, queuing finalization", file_id=file_id)
|
log_task_progress(task_id, "embed_metadata_into_pdf", "success", "Metadata embedded, queuing finalization", file_id=file_id)
|
||||||
finalize_document_storage.delay(original_file, final_file_path, metadata)
|
finalize_document_storage.delay(original_file, final_file_path, metadata, file_id)
|
||||||
|
|
||||||
# After triggering final storage, delete the original file if it is in workdir/tmp.
|
# After triggering final storage, delete the original file if it is in workdir/tmp.
|
||||||
workdir_tmp = os.path.join(settings.workdir, "tmp")
|
workdir_tmp = os.path.join(settings.workdir, "tmp")
|
||||||
|
|||||||
@@ -46,21 +46,21 @@ def extract_json_from_text(text):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||||
def extract_metadata_with_gpt(self, filename: str, cleaned_text: str):
|
def extract_metadata_with_gpt(self, filename: str, cleaned_text: str, file_id: int = None):
|
||||||
"""Uses OpenAI to classify document metadata."""
|
"""Uses OpenAI to classify document metadata."""
|
||||||
task_id = self.request.id
|
task_id = self.request.id
|
||||||
logger.info(f"[{task_id}] Starting metadata extraction for: {filename}")
|
logger.info(f"[{task_id}] Starting metadata extraction for: {filename}")
|
||||||
log_task_progress(task_id, "extract_metadata_with_gpt", "in_progress", f"Extracting metadata for {filename}")
|
log_task_progress(task_id, "extract_metadata_with_gpt", "in_progress", f"Extracting metadata for {filename}", file_id=file_id)
|
||||||
|
|
||||||
# Get file_id from database
|
# Get file_id from database if not provided
|
||||||
file_id = None
|
if file_id is None:
|
||||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||||
file_path = os.path.join(tmp_dir, filename)
|
file_path = os.path.join(tmp_dir, filename)
|
||||||
if os.path.exists(file_path):
|
if os.path.exists(file_path):
|
||||||
with SessionLocal() as db:
|
with SessionLocal() as db:
|
||||||
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
||||||
if file_record:
|
if file_record:
|
||||||
file_id = file_record.id
|
file_id = file_record.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.
|
||||||
@@ -123,7 +123,7 @@ Return only valid JSON with no additional commentary.
|
|||||||
# Trigger the next step: embedding metadata into the PDF
|
# Trigger the next step: embedding metadata into the PDF
|
||||||
logger.info(f"[{task_id}] Queueing metadata embedding task")
|
logger.info(f"[{task_id}] Queueing metadata embedding task")
|
||||||
log_task_progress(task_id, "extract_metadata_with_gpt", "success", "Metadata extracted, queuing embed task", file_id=file_id)
|
log_task_progress(task_id, "extract_metadata_with_gpt", "success", "Metadata extracted, queuing embed task", file_id=file_id)
|
||||||
embed_metadata_into_pdf.delay(filename, cleaned_text, metadata)
|
embed_metadata_into_pdf.delay(filename, cleaned_text, metadata, file_id)
|
||||||
|
|
||||||
return {"s3_file": filename, "metadata": metadata}
|
return {"s3_file": filename, "metadata": metadata}
|
||||||
|
|
||||||
|
|||||||
@@ -17,29 +17,29 @@ logger = logging.getLogger(__name__)
|
|||||||
|
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||||
def finalize_document_storage(self, original_file: str, processed_file: str, metadata: dict):
|
def finalize_document_storage(self, original_file: str, processed_file: str, metadata: dict, file_id: int = None):
|
||||||
"""
|
"""
|
||||||
Final storage step after embedding metadata.
|
Final storage step after embedding metadata.
|
||||||
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
|
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
|
||||||
"""
|
"""
|
||||||
task_id = self.request.id
|
task_id = self.request.id
|
||||||
logger.info(f"[{task_id}] Finalizing document storage for {processed_file}")
|
logger.info(f"[{task_id}] Finalizing document storage for {processed_file}")
|
||||||
log_task_progress(task_id, "finalize_document_storage", "in_progress", f"Finalizing: {os.path.basename(processed_file)}")
|
log_task_progress(task_id, "finalize_document_storage", "in_progress", f"Finalizing: {os.path.basename(processed_file)}", file_id=file_id)
|
||||||
|
|
||||||
# Get file_id from database
|
# Get file_id from database if not provided
|
||||||
file_id = None
|
if file_id is None:
|
||||||
with SessionLocal() as db:
|
with SessionLocal() as db:
|
||||||
# Try to find by the processed file path first
|
# Try to find by the original file path
|
||||||
file_record = db.query(FileRecord).filter(
|
file_record = db.query(FileRecord).filter(
|
||||||
FileRecord.local_filename.like(f"%{os.path.basename(original_file)}%")
|
FileRecord.local_filename.like(f"%{os.path.basename(original_file)}%")
|
||||||
).first()
|
).first()
|
||||||
if file_record:
|
if file_record:
|
||||||
file_id = file_record.id
|
file_id = file_record.id
|
||||||
|
|
||||||
# 2) Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
# 2) Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
||||||
logger.info(f"[{task_id}] Queueing uploads to all destinations")
|
logger.info(f"[{task_id}] Queueing uploads to all destinations")
|
||||||
log_task_progress(task_id, "finalize_document_storage", "success", "Queuing uploads to destinations", file_id=file_id)
|
log_task_progress(task_id, "finalize_document_storage", "success", "Queuing uploads to destinations", file_id=file_id)
|
||||||
send_to_all_destinations.delay(processed_file)
|
send_to_all_destinations.delay(processed_file, True, file_id)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"status": "Completed",
|
"status": "Completed",
|
||||||
|
|||||||
@@ -130,12 +130,12 @@ def process_document(self, original_local_file: str):
|
|||||||
# Call metadata extraction directly
|
# Call metadata extraction directly
|
||||||
logger.info(f"[{task_id}] Queueing metadata extraction")
|
logger.info(f"[{task_id}] Queueing metadata extraction")
|
||||||
log_task_progress(task_id, "process_document", "success", "Queued for metadata extraction", file_id=new_record.id)
|
log_task_progress(task_id, "process_document", "success", "Queued for metadata extraction", file_id=new_record.id)
|
||||||
extract_metadata_with_gpt.delay(new_filename, extracted_text)
|
extract_metadata_with_gpt.delay(new_filename, extracted_text, new_record.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 Azure Document Intelligence processing
|
# 3. If no embedded text, queue Azure Document Intelligence processing
|
||||||
logger.info(f"[{task_id}] No embedded text found. Queueing Azure Document Intelligence processing")
|
logger.info(f"[{task_id}] No embedded text found. Queueing Azure Document Intelligence processing")
|
||||||
log_task_progress(task_id, "check_text", "success", "No embedded text, queuing OCR", file_id=new_record.id)
|
log_task_progress(task_id, "check_text", "success", "No embedded text, queuing OCR", file_id=new_record.id)
|
||||||
log_task_progress(task_id, "process_document", "success", "Queued for OCR processing", file_id=new_record.id)
|
log_task_progress(task_id, "process_document", "success", "Queued for OCR processing", file_id=new_record.id)
|
||||||
process_with_azure_document_intelligence.delay(new_filename)
|
process_with_azure_document_intelligence.delay(new_filename, new_record.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}
|
||||||
|
|||||||
@@ -76,7 +76,7 @@ def check_page_rotation(result, filename):
|
|||||||
return rotation_data
|
return rotation_data
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
def process_with_azure_document_intelligence(filename: str):
|
def process_with_azure_document_intelligence(filename: str, file_id: int = None):
|
||||||
"""
|
"""
|
||||||
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
|
||||||
the local temporary file (stored under <workdir>/tmp).
|
the local temporary file (stored under <workdir>/tmp).
|
||||||
@@ -88,6 +88,10 @@ def process_with_azure_document_intelligence(filename: str):
|
|||||||
3. Saves the OCR-processed PDF locally in the same location as before.
|
3. Saves the OCR-processed PDF locally in the same location as before.
|
||||||
4. Checks for page rotation and triggers page rotation if needed.
|
4. Checks for page rotation and triggers page rotation if needed.
|
||||||
5. Triggers downstream metadata extraction.
|
5. Triggers downstream metadata extraction.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
filename: Name of the file to process
|
||||||
|
file_id: Optional file ID to pass through to subsequent tasks
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
tmp_file_path = os.path.join(settings.workdir, "tmp", filename)
|
tmp_file_path = os.path.join(settings.workdir, "tmp", filename)
|
||||||
@@ -139,7 +143,7 @@ def process_with_azure_document_intelligence(filename: str):
|
|||||||
logger.info(f"Extracted text for {filename}: {len(extracted_text)} characters")
|
logger.info(f"Extracted text for {filename}: {len(extracted_text)} characters")
|
||||||
|
|
||||||
# Trigger page rotation task if rotation is detected, otherwise proceed to metadata extraction
|
# Trigger page rotation task if rotation is detected, otherwise proceed to metadata extraction
|
||||||
rotate_pdf_pages.delay(filename, extracted_text, rotation_data)
|
rotate_pdf_pages.delay(filename, extracted_text, rotation_data, file_id)
|
||||||
|
|
||||||
return {"file": filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
|
return {"file": filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": extracted_text}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -47,7 +47,7 @@ def determine_rotation_angle(detected_angle):
|
|||||||
return rotation_value
|
return rotation_value
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None, file_id: int = None):
|
||||||
"""
|
"""
|
||||||
Rotates pages in a PDF document based on detected rotation angles.
|
Rotates pages in a PDF document based on detected rotation angles.
|
||||||
|
|
||||||
@@ -55,6 +55,7 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
|||||||
filename: The name of the file to rotate
|
filename: The name of the file to rotate
|
||||||
extracted_text: The extracted text from the document
|
extracted_text: The extracted text from the document
|
||||||
rotation_data: Optional rotation data dictionary {page_index: angle}
|
rotation_data: Optional rotation data dictionary {page_index: angle}
|
||||||
|
file_id: Optional file ID to pass through to subsequent tasks
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
pdf_path = os.path.join(settings.workdir, "tmp", filename)
|
pdf_path = os.path.join(settings.workdir, "tmp", filename)
|
||||||
@@ -64,7 +65,7 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
|||||||
# Skip rotation if no rotation data provided
|
# Skip rotation if no rotation data provided
|
||||||
if not rotation_data:
|
if not rotation_data:
|
||||||
logger.info(f"No rotation data provided for {filename}, proceeding with metadata extraction")
|
logger.info(f"No rotation data provided for {filename}, proceeding with metadata extraction")
|
||||||
extract_metadata_with_gpt.delay(filename, extracted_text)
|
extract_metadata_with_gpt.delay(filename, extracted_text, file_id)
|
||||||
return {"file": filename, "status": "no_rotation_needed"}
|
return {"file": filename, "status": "no_rotation_needed"}
|
||||||
|
|
||||||
# Standardize rotation_data keys to integers
|
# Standardize rotation_data keys to integers
|
||||||
@@ -77,7 +78,7 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
|||||||
|
|
||||||
if not any(abs(angle) > 0 for angle in normalized_rotation_data.values()):
|
if not any(abs(angle) > 0 for angle in normalized_rotation_data.values()):
|
||||||
logger.info(f"No significant rotations detected in {filename}, proceeding with metadata extraction")
|
logger.info(f"No significant rotations detected in {filename}, proceeding with metadata extraction")
|
||||||
extract_metadata_with_gpt.delay(filename, extracted_text)
|
extract_metadata_with_gpt.delay(filename, extracted_text, file_id)
|
||||||
return {"file": filename, "status": "no_rotation_needed"}
|
return {"file": filename, "status": "no_rotation_needed"}
|
||||||
|
|
||||||
logger.info(f"Rotating {len(normalized_rotation_data)} pages in {filename}")
|
logger.info(f"Rotating {len(normalized_rotation_data)} pages in {filename}")
|
||||||
@@ -117,7 +118,7 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
|||||||
logger.info(f"Detected rotations in {filename} but no rotations were actually applied (angles too small or not multiples of 90°)")
|
logger.info(f"Detected rotations in {filename} but no rotations were actually applied (angles too small or not multiples of 90°)")
|
||||||
|
|
||||||
# Continue with metadata extraction
|
# Continue with metadata extraction
|
||||||
extract_metadata_with_gpt.delay(filename, extracted_text)
|
extract_metadata_with_gpt.delay(filename, extracted_text, file_id)
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"file": filename,
|
"file": filename,
|
||||||
@@ -129,5 +130,5 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error rotating PDF {filename}: {e}")
|
logger.error(f"Error rotating PDF {filename}: {e}")
|
||||||
# Continue with metadata extraction despite rotation failure
|
# Continue with metadata extraction despite rotation failure
|
||||||
extract_metadata_with_gpt.delay(filename, extracted_text)
|
extract_metadata_with_gpt.delay(filename, extracted_text, file_id)
|
||||||
return {"file": filename, "status": "rotation_failed", "error": str(e)}
|
return {"file": filename, "status": "rotation_failed", "error": str(e)}
|
||||||
|
|||||||
+12
-11
@@ -108,7 +108,7 @@ def get_configured_services_from_validator():
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||||
def send_to_all_destinations(self, file_path: str, use_validator=True):
|
def send_to_all_destinations(self, file_path: str, use_validator=True, file_id: int = None):
|
||||||
"""
|
"""
|
||||||
Distribute a file to all configured storage destinations.
|
Distribute a file to all configured storage destinations.
|
||||||
|
|
||||||
@@ -116,25 +116,26 @@ def send_to_all_destinations(self, file_path: str, use_validator=True):
|
|||||||
file_path: Path to the file to distribute
|
file_path: Path to the file to distribute
|
||||||
use_validator: Whether to use the config validator to determine enabled services
|
use_validator: Whether to use the config validator to determine enabled services
|
||||||
(if False, falls back to individual checks)
|
(if False, falls back to individual checks)
|
||||||
|
file_id: Optional file ID to associate with logs
|
||||||
"""
|
"""
|
||||||
task_id = self.request.id
|
task_id = self.request.id
|
||||||
|
|
||||||
if not os.path.exists(file_path):
|
if not os.path.exists(file_path):
|
||||||
logger.error(f"[{task_id}] File not found: {file_path}")
|
logger.error(f"[{task_id}] File not found: {file_path}")
|
||||||
log_task_progress(task_id, "send_to_all_destinations", "failure", "File not found")
|
log_task_progress(task_id, "send_to_all_destinations", "failure", "File not found", file_id=file_id)
|
||||||
raise FileNotFoundError(f"File not found: {file_path}")
|
raise FileNotFoundError(f"File not found: {file_path}")
|
||||||
|
|
||||||
logger.info(f"[{task_id}] Sending {file_path} to all configured destinations")
|
logger.info(f"[{task_id}] Sending {file_path} to all configured destinations")
|
||||||
log_task_progress(task_id, "send_to_all_destinations", "in_progress", f"Distributing: {os.path.basename(file_path)}")
|
log_task_progress(task_id, "send_to_all_destinations", "in_progress", f"Distributing: {os.path.basename(file_path)}", file_id=file_id)
|
||||||
|
|
||||||
# Get file_id from database
|
# Get file_id from database if not provided
|
||||||
file_id = None
|
if file_id is None:
|
||||||
with SessionLocal() as db:
|
with SessionLocal() as db:
|
||||||
file_record = db.query(FileRecord).filter(
|
file_record = db.query(FileRecord).filter(
|
||||||
FileRecord.local_filename.like(f"%{os.path.basename(file_path)}%")
|
FileRecord.local_filename.like(f"%{os.path.basename(file_path)}%")
|
||||||
).first()
|
).first()
|
||||||
if file_record:
|
if file_record:
|
||||||
file_id = file_record.id
|
file_id = file_record.id
|
||||||
|
|
||||||
results = {}
|
results = {}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user