diff --git a/app/api/files.py b/app/api/files.py index ad5ad663..0e81f43b 100644 --- a/app/api/files.py +++ b/app/api/files.py @@ -456,35 +456,18 @@ def _retry_pipeline_step(file_record: FileRecord, step_name: str, db: Session) - filename = os.path.basename(file_record.local_filename) task = extract_metadata_with_gpt.delay(filename, extracted_text, file_id) elif step_name == "embed_metadata_into_pdf": - from app.tasks.embed_metadata_into_pdf import embed_metadata_into_pdf - - # Retrieve the last successful metadata extraction result from processing logs - last_metadata_log = ( - db.query(ProcessingLog) - .filter( - ProcessingLog.file_id == file_id, - ProcessingLog.step_name == "extract_metadata_with_gpt", - ProcessingLog.status == "success", - ) - .order_by(ProcessingLog.timestamp.desc()) - .first() - ) - if not last_metadata_log: - raise HTTPException( - status_code=400, - detail="No successful metadata extraction found. Retry extract_metadata_with_gpt first.", - ) + from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt as extract_metadata_task + # Retrying embed requires re-running metadata extraction first, because + # embed_metadata_into_pdf needs the actual metadata dict (not empty). + # Re-trigger extract_metadata_with_gpt which will chain into embed_metadata_into_pdf. if not file_record.local_filename or not os.path.exists(file_record.local_filename): raise HTTPException( status_code=400, detail="Local file not found on disk. Cannot retry metadata embedding." ) extracted_text = _extract_text_from_pdf(file_record.local_filename) filename = os.path.basename(file_record.local_filename) - # Empty metadata dict: embed_metadata_into_pdf will re-run with the provided text. - # The GPT extraction step must succeed first (validated above) to ensure - # the pipeline can produce new metadata during the subsequent re-extraction. - task = embed_metadata_into_pdf.delay(filename, extracted_text, {}, file_id) + task = extract_metadata_task.delay(filename, extracted_text, file_id) else: raise HTTPException(status_code=400, detail=f"Unsupported pipeline step: {step_name}") diff --git a/app/tasks/extract_metadata_with_gpt.py b/app/tasks/extract_metadata_with_gpt.py index e673e3bf..03b3231c 100644 --- a/app/tasks/extract_metadata_with_gpt.py +++ b/app/tasks/extract_metadata_with_gpt.py @@ -126,17 +126,18 @@ def extract_metadata_with_gpt(self, filename: str, cleaned_text: str, file_id: i # SECURITY: Validate filename format from GPT to prevent path traversal # The prompt requests filenames with only letters, numbers, periods, and underscores # Enforce this constraint to prevent malicious filenames - import re - filename = metadata.get("filename", "") - if filename: + suggested_filename = metadata.get("filename", "") + if suggested_filename: # Check if filename contains only safe characters AND explicitly check for ".." # Defense in depth: While the regex [\w\-\. ]+ already excludes / and \, # we explicitly reject ".." to guard against: # 1. Potential locale-specific \w behavior # 2. Files literally named ".." which are valid but problematic # 3. Future code changes that might relax the regex - if not re.match(r'^[\w\-\. ]+$', filename) or ".." in filename: - logger.warning(f"[{task_id}] Invalid filename format from GPT: '{filename}', using fallback") + if not re.match(r'^[\w\-\. ]+$', suggested_filename) or ".." in suggested_filename: + logger.warning( + f"[{task_id}] Invalid filename format from GPT: '{suggested_filename}', using fallback" + ) # Reset to empty to trigger fallback to original filename metadata["filename"] = "" @@ -146,6 +147,7 @@ def extract_metadata_with_gpt(self, filename: str, cleaned_text: str, file_id: i ) # Trigger the next step: embedding metadata into the PDF + # Pass the original filename (UUID-based) so embed_metadata_into_pdf can find the file on disk 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 diff --git a/tests/test_file_detail_endpoints.py b/tests/test_file_detail_endpoints.py index 40cb894a..6f897ee8 100644 --- a/tests/test_file_detail_endpoints.py +++ b/tests/test_file_detail_endpoints.py @@ -213,8 +213,11 @@ class TestSubtaskRetry: assert data["status"] == "success" assert data["subtask_name"] == "extract_metadata_with_gpt" - def test_retry_pipeline_step_embed_no_metadata(self, client: TestClient, db_session, sample_pdf_path): - """Test retrying embed_metadata_into_pdf without prior metadata extraction.""" + def test_retry_pipeline_step_embed_metadata(self, client: TestClient, db_session, sample_pdf_path): + """Test retrying embed_metadata_into_pdf re-triggers metadata extraction.""" + mock_task = MagicMock() + mock_task.id = "embed-retry-task" + file_record = FileRecord( filehash="pipeline_retry4", original_filename="embed_retry.pdf", @@ -226,10 +229,17 @@ class TestSubtaskRetry: db_session.commit() db_session.refresh(file_record) - # Try to retry embed without a successful metadata extraction log - response = client.post(f"/api/files/{file_record.id}/retry-subtask?subtask_name=embed_metadata_into_pdf") - assert response.status_code == 400 - assert "retry extract_metadata_with_gpt first" in response.json()["detail"].lower() + with patch("app.tasks.extract_metadata_with_gpt.extract_metadata_with_gpt") as mock_extract: + mock_extract.delay.return_value = mock_task + response = client.post( + f"/api/files/{file_record.id}/retry-subtask?subtask_name=embed_metadata_into_pdf" + ) + assert response.status_code == 200 + data = response.json() + assert data["status"] == "success" + assert data["subtask_name"] == "embed_metadata_into_pdf" + # Verify extract_metadata_with_gpt.delay was called (which chains into embed) + mock_extract.delay.assert_called_once() def test_retry_pipeline_step_missing_local_file(self, client: TestClient, db_session): """Test retrying a pipeline step when local file is missing."""