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:
copilot-swe-agent[bot]
@@ -46,7 +46,7 @@ def persist_metadata(metadata, final_pdf_path):
|
||||
return json_path
|
||||
|
||||
@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.
|
||||
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
|
||||
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
|
||||
file_id = None
|
||||
# Check for file existence; if not found, try the known shared tmp directory.
|
||||
# Get file_id from database if not provided
|
||||
if file_id is None:
|
||||
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):
|
||||
alt_path = os.path.join(settings.workdir, "tmp", os.path.basename(local_file_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")
|
||||
return {"error": "File not found"}
|
||||
|
||||
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
|
||||
|
||||
# Check for file existence; if not found, try the known shared tmp directory.
|
||||
|
||||
# Work on a safe copy in a secure temporary directory
|
||||
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.
|
||||
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)
|
||||
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.
|
||||
workdir_tmp = os.path.join(settings.workdir, "tmp")
|
||||
|
||||
@@ -46,21 +46,21 @@ def extract_json_from_text(text):
|
||||
return None
|
||||
|
||||
@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."""
|
||||
task_id = self.request.id
|
||||
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
|
||||
file_id = None
|
||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||
file_path = os.path.join(tmp_dir, filename)
|
||||
if os.path.exists(file_path):
|
||||
with SessionLocal() as db:
|
||||
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
||||
if file_record:
|
||||
file_id = file_record.id
|
||||
# Get file_id from database if not provided
|
||||
if file_id is None:
|
||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||
file_path = os.path.join(tmp_dir, filename)
|
||||
if os.path.exists(file_path):
|
||||
with SessionLocal() as db:
|
||||
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
||||
if file_record:
|
||||
file_id = file_record.id
|
||||
|
||||
prompt = f"""
|
||||
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
|
||||
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)
|
||||
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}
|
||||
|
||||
|
||||
@@ -17,29 +17,29 @@ logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@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.
|
||||
We will now call 'send_to_all_destinations' to push the final PDF to Dropbox/Nextcloud/Paperless.
|
||||
"""
|
||||
task_id = self.request.id
|
||||
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
|
||||
file_id = None
|
||||
with SessionLocal() as db:
|
||||
# Try to find by the processed file path first
|
||||
file_record = db.query(FileRecord).filter(
|
||||
FileRecord.local_filename.like(f"%{os.path.basename(original_file)}%")
|
||||
).first()
|
||||
if file_record:
|
||||
file_id = file_record.id
|
||||
# Get file_id from database if not provided
|
||||
if file_id is None:
|
||||
with SessionLocal() as db:
|
||||
# Try to find by the original file path
|
||||
file_record = db.query(FileRecord).filter(
|
||||
FileRecord.local_filename.like(f"%{os.path.basename(original_file)}%")
|
||||
).first()
|
||||
if file_record:
|
||||
file_id = file_record.id
|
||||
|
||||
# 2) Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
||||
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)
|
||||
send_to_all_destinations.delay(processed_file)
|
||||
send_to_all_destinations.delay(processed_file, True, file_id)
|
||||
|
||||
return {
|
||||
"status": "Completed",
|
||||
|
||||
@@ -130,12 +130,12 @@ def process_document(self, original_local_file: str):
|
||||
# Call metadata extraction directly
|
||||
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)
|
||||
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}
|
||||
|
||||
# 3. If no embedded text, queue 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, "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}
|
||||
|
||||
@@ -76,7 +76,7 @@ def check_page_rotation(result, filename):
|
||||
return rotation_data
|
||||
|
||||
@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
|
||||
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.
|
||||
4. Checks for page rotation and triggers page rotation if needed.
|
||||
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:
|
||||
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")
|
||||
|
||||
# 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}
|
||||
except Exception as e:
|
||||
|
||||
@@ -47,7 +47,7 @@ def determine_rotation_angle(detected_angle):
|
||||
return rotation_value
|
||||
|
||||
@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.
|
||||
|
||||
@@ -55,6 +55,7 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
||||
filename: The name of the file to rotate
|
||||
extracted_text: The extracted text from the document
|
||||
rotation_data: Optional rotation data dictionary {page_index: angle}
|
||||
file_id: Optional file ID to pass through to subsequent tasks
|
||||
"""
|
||||
try:
|
||||
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
|
||||
if not rotation_data:
|
||||
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"}
|
||||
|
||||
# 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()):
|
||||
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"}
|
||||
|
||||
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°)")
|
||||
|
||||
# Continue with metadata extraction
|
||||
extract_metadata_with_gpt.delay(filename, extracted_text)
|
||||
extract_metadata_with_gpt.delay(filename, extracted_text, file_id)
|
||||
|
||||
return {
|
||||
"file": filename,
|
||||
@@ -129,5 +130,5 @@ def rotate_pdf_pages(filename: str, extracted_text: str, rotation_data=None):
|
||||
except Exception as e:
|
||||
logger.error(f"Error rotating PDF {filename}: {e}")
|
||||
# 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)}
|
||||
|
||||
+12
-11
@@ -108,7 +108,7 @@ def get_configured_services_from_validator():
|
||||
return result
|
||||
|
||||
@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.
|
||||
|
||||
@@ -116,25 +116,26 @@ def send_to_all_destinations(self, file_path: str, use_validator=True):
|
||||
file_path: Path to the file to distribute
|
||||
use_validator: Whether to use the config validator to determine enabled services
|
||||
(if False, falls back to individual checks)
|
||||
file_id: Optional file ID to associate with logs
|
||||
"""
|
||||
task_id = self.request.id
|
||||
|
||||
if not os.path.exists(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}")
|
||||
|
||||
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
|
||||
file_id = None
|
||||
with SessionLocal() as db:
|
||||
file_record = db.query(FileRecord).filter(
|
||||
FileRecord.local_filename.like(f"%{os.path.basename(file_path)}%")
|
||||
).first()
|
||||
if file_record:
|
||||
file_id = file_record.id
|
||||
# Get file_id from database if not provided
|
||||
if file_id is None:
|
||||
with SessionLocal() as db:
|
||||
file_record = db.query(FileRecord).filter(
|
||||
FileRecord.local_filename.like(f"%{os.path.basename(file_path)}%")
|
||||
).first()
|
||||
if file_record:
|
||||
file_id = file_record.id
|
||||
|
||||
results = {}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user