refactor(tasks): extract filename regex to shared constant
Move the valid filename regex pattern to a shared constant in `app/utils/filename_utils.py` and update both the task logic and security tests to use it. This eliminates duplication and ensures consistency across the codebase. Also normalized line endings to LF in affected files to ensure CI compatibility. Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
google-labs-jules[bot]
@@ -1,191 +1,192 @@
|
|||||||
#!/usr/bin/env python3
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
|
|
||||||
# Import the shared Celery instance
|
# Import the shared Celery instance
|
||||||
from app.celery_app import celery
|
from app.celery_app import celery
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
from app.database import SessionLocal
|
from app.database import SessionLocal
|
||||||
from app.models import FileRecord
|
from app.models import FileRecord
|
||||||
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.tasks.retry_config import BaseTaskWithRetry
|
from app.tasks.retry_config import BaseTaskWithRetry
|
||||||
from app.utils import log_task_progress
|
from app.utils import log_task_progress
|
||||||
from app.utils.ai_provider import get_ai_provider
|
from app.utils.ai_provider import get_ai_provider
|
||||||
|
from app.utils.filename_utils import VALID_FILENAME_RE
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
def extract_json_from_text(text):
|
|
||||||
"""
|
def extract_json_from_text(text):
|
||||||
Try to extract a JSON object from the text.
|
"""
|
||||||
- First, check for a JSON block inside triple backticks.
|
Try to extract a JSON object from the text.
|
||||||
- If not found, try to extract text from the first '{' to the last '}'.
|
- First, check for a JSON block inside triple backticks.
|
||||||
"""
|
- If not found, try to extract text from the first '{' to the last '}'.
|
||||||
pattern = r"```(?:json)?\s*(\{.*?\})\s*```"
|
"""
|
||||||
match = re.search(pattern, text, re.DOTALL)
|
pattern = r"```(?:json)?\s*(\{.*?\})\s*```"
|
||||||
if match:
|
match = re.search(pattern, text, re.DOTALL)
|
||||||
return match.group(1)
|
if match:
|
||||||
else:
|
return match.group(1)
|
||||||
start = text.find("{")
|
else:
|
||||||
end = text.rfind("}")
|
start = text.find("{")
|
||||||
if start != -1 and end != -1 and end > start:
|
end = text.rfind("}")
|
||||||
return text[start : end + 1]
|
if start != -1 and end != -1 and end > start:
|
||||||
return None
|
return text[start : end + 1]
|
||||||
|
return None
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry, bind=True)
|
|
||||||
def extract_metadata_with_gpt(self, filename: str, cleaned_text: str, file_id: int = None):
|
@celery.task(base=BaseTaskWithRetry, bind=True)
|
||||||
"""
|
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.
|
||||||
Args:
|
|
||||||
filename: Can be either a basename (e.g., "file.pdf") or a full path (e.g., "/workdir/processed/file.pdf")
|
Args:
|
||||||
cleaned_text: The extracted text from the document
|
filename: Can be either a basename (e.g., "file.pdf") or a full path (e.g., "/workdir/processed/file.pdf")
|
||||||
file_id: Optional file ID for tracking
|
cleaned_text: The extracted text from the document
|
||||||
"""
|
file_id: Optional file ID for tracking
|
||||||
task_id = self.request.id
|
"""
|
||||||
logger.info(f"[{task_id}] Starting metadata extraction for: {filename}")
|
task_id = self.request.id
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Starting metadata extraction for: {filename}")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"extract_metadata_with_gpt",
|
task_id,
|
||||||
"in_progress",
|
"extract_metadata_with_gpt",
|
||||||
f"Extracting metadata for {os.path.basename(filename)}",
|
"in_progress",
|
||||||
file_id=file_id,
|
f"Extracting metadata for {os.path.basename(filename)}",
|
||||||
)
|
file_id=file_id,
|
||||||
|
)
|
||||||
# Get file_id from database if not provided
|
|
||||||
if file_id is None:
|
# Get file_id from database if not provided
|
||||||
tmp_dir = os.path.join(settings.workdir, "tmp")
|
if file_id is None:
|
||||||
# Handle both basename and full path
|
tmp_dir = os.path.join(settings.workdir, "tmp")
|
||||||
if os.path.isabs(filename):
|
# Handle both basename and full path
|
||||||
file_path = filename
|
if os.path.isabs(filename):
|
||||||
else:
|
file_path = filename
|
||||||
file_path = os.path.join(tmp_dir, filename)
|
else:
|
||||||
if os.path.exists(file_path):
|
file_path = os.path.join(tmp_dir, filename)
|
||||||
with SessionLocal() as db:
|
if os.path.exists(file_path):
|
||||||
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
with SessionLocal() as db:
|
||||||
if file_record:
|
file_record = db.query(FileRecord).filter_by(local_filename=file_path).first()
|
||||||
file_id = file_record.id
|
if file_record:
|
||||||
|
file_id = file_record.id
|
||||||
prompt = (
|
|
||||||
"You are a specialized document analyzer trained to extract structured metadata from documents.\n"
|
prompt = (
|
||||||
"Your task is to analyze the given text and return a well-structured JSON object.\n\n"
|
"You are a specialized document analyzer trained to extract structured metadata from documents.\n"
|
||||||
"Extract and return the following fields:\n"
|
"Your task is to analyze the given text and return a well-structured JSON object.\n\n"
|
||||||
"1. **filename**: Machine-readable filename "
|
"Extract and return the following fields:\n"
|
||||||
"(YYYY-MM-DD_DescriptiveTitle, use only letters, numbers, periods, and underscores).\n"
|
"1. **filename**: Machine-readable filename "
|
||||||
'2. **empfaenger**: The recipient, or "Unknown" if not found.\n'
|
"(YYYY-MM-DD_DescriptiveTitle, use only letters, numbers, periods, and underscores).\n"
|
||||||
'3. **absender**: The sender, or "Unknown" if not found.\n'
|
'2. **empfaenger**: The recipient, or "Unknown" if not found.\n'
|
||||||
"4. **correspondent**: The entity or company that issued the document "
|
'3. **absender**: The sender, or "Unknown" if not found.\n'
|
||||||
'(shortest possible name, e.g., "Amazon" instead of "Amazon EU SARL, German branch").\n'
|
"4. **correspondent**: The entity or company that issued the document "
|
||||||
"5. **kommunikationsart**: One of [Behoerdlicher_Brief, Rechnung, Kontoauszug, Vertrag, "
|
'(shortest possible name, e.g., "Amazon" instead of "Amazon EU SARL, German branch").\n'
|
||||||
"Quittung, Privater_Brief, Einladung, Gewerbliche_Korrespondenz, Newsletter, Werbung, Sonstiges].\n"
|
"5. **kommunikationsart**: One of [Behoerdlicher_Brief, Rechnung, Kontoauszug, Vertrag, "
|
||||||
"6. **kommunikationskategorie**: One of [Amtliche_Postbehoerdliche_Dokumente, "
|
"Quittung, Privater_Brief, Einladung, Gewerbliche_Korrespondenz, Newsletter, Werbung, Sonstiges].\n"
|
||||||
"Finanz_und_Vertragsdokumente, Geschaeftliche_Kommunikation, "
|
"6. **kommunikationskategorie**: One of [Amtliche_Postbehoerdliche_Dokumente, "
|
||||||
"Private_Korrespondenz, Sonstige_Informationen].\n"
|
"Finanz_und_Vertragsdokumente, Geschaeftliche_Kommunikation, "
|
||||||
"7. **document_type**: Precise classification (e.g., Invoice, Contract, Information, Unknown).\n"
|
"Private_Korrespondenz, Sonstige_Informationen].\n"
|
||||||
"8. **tags**: A list of up to 4 relevant thematic keywords.\n"
|
"7. **document_type**: Precise classification (e.g., Invoice, Contract, Information, Unknown).\n"
|
||||||
'9. **language**: Detected document language (ISO 639-1 code, e.g., "de" or "en").\n'
|
"8. **tags**: A list of up to 4 relevant thematic keywords.\n"
|
||||||
"10. **title**: A human-readable title summarizing the document content.\n"
|
'9. **language**: Detected document language (ISO 639-1 code, e.g., "de" or "en").\n'
|
||||||
"11. **confidence_score**: A numeric value (0-100) indicating the confidence level "
|
"10. **title**: A human-readable title summarizing the document content.\n"
|
||||||
"of the extracted metadata.\n"
|
"11. **confidence_score**: A numeric value (0-100) indicating the confidence level "
|
||||||
"12. **reference_number**: Extracted invoice/order/reference number if available.\n"
|
"of the extracted metadata.\n"
|
||||||
"13. **monetary_amounts**: A list of key monetary values detected in the document.\n\n"
|
"12. **reference_number**: Extracted invoice/order/reference number if available.\n"
|
||||||
"### Important Rules:\n"
|
"13. **monetary_amounts**: A list of key monetary values detected in the document.\n\n"
|
||||||
"- **OCR Correction**: Assume the text has been corrected for OCR errors.\n"
|
"### Important Rules:\n"
|
||||||
"- **Tagging**: Max 4 tags, avoiding generic or overly specific terms.\n"
|
"- **OCR Correction**: Assume the text has been corrected for OCR errors.\n"
|
||||||
"- **Title**: Concise, no addresses, and contains key identifying features.\n"
|
"- **Tagging**: Max 4 tags, avoiding generic or overly specific terms.\n"
|
||||||
"- **Date Selection**: Use the most relevant date if multiple are found.\n"
|
"- **Title**: Concise, no addresses, and contains key identifying features.\n"
|
||||||
"- **Output Language**: Maintain the document's original language.\n\n"
|
"- **Date Selection**: Use the most relevant date if multiple are found.\n"
|
||||||
f"Extracted text:\n{cleaned_text}\n\n"
|
"- **Output Language**: Maintain the document's original language.\n\n"
|
||||||
"Return only valid JSON with no additional commentary.\n"
|
f"Extracted text:\n{cleaned_text}\n\n"
|
||||||
)
|
"Return only valid JSON with no additional commentary.\n"
|
||||||
|
)
|
||||||
try:
|
|
||||||
logger.info(f"[{task_id}] Sending classification request for {filename}...")
|
try:
|
||||||
log_task_progress(task_id, "call_ai_provider", "in_progress", "Calling AI provider API", file_id=file_id)
|
logger.info(f"[{task_id}] Sending classification request for {filename}...")
|
||||||
provider = get_ai_provider()
|
log_task_progress(task_id, "call_ai_provider", "in_progress", "Calling AI provider API", file_id=file_id)
|
||||||
model = settings.ai_model or settings.openai_model
|
provider = get_ai_provider()
|
||||||
content = provider.chat_completion(
|
model = settings.ai_model or settings.openai_model
|
||||||
messages=[
|
content = provider.chat_completion(
|
||||||
{"role": "system", "content": "You are an intelligent document classifier."},
|
messages=[
|
||||||
{"role": "user", "content": prompt},
|
{"role": "system", "content": "You are an intelligent document classifier."},
|
||||||
],
|
{"role": "user", "content": prompt},
|
||||||
model=model,
|
],
|
||||||
temperature=0,
|
model=model,
|
||||||
)
|
temperature=0,
|
||||||
|
)
|
||||||
logger.info(f"[{task_id}] Raw classification response for {filename}: {content[:200]}...")
|
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Raw classification response for {filename}: {content[:200]}...")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"call_ai_provider",
|
task_id,
|
||||||
"success",
|
"call_ai_provider",
|
||||||
"Received AI provider response",
|
"success",
|
||||||
file_id=file_id,
|
"Received AI provider response",
|
||||||
detail=f"Raw classification response:\n{content}",
|
file_id=file_id,
|
||||||
)
|
detail=f"Raw classification response:\n{content}",
|
||||||
|
)
|
||||||
json_text = extract_json_from_text(content)
|
|
||||||
if not json_text:
|
json_text = extract_json_from_text(content)
|
||||||
logger.error(f"[{task_id}] Could not find valid JSON in GPT response for {filename}.")
|
if not json_text:
|
||||||
log_task_progress(
|
logger.error(f"[{task_id}] Could not find valid JSON in GPT response for {filename}.")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"extract_metadata_with_gpt",
|
task_id,
|
||||||
"failure",
|
"extract_metadata_with_gpt",
|
||||||
"Invalid JSON in response",
|
"failure",
|
||||||
file_id=file_id,
|
"Invalid JSON in response",
|
||||||
detail=f"Could not parse valid JSON from GPT response.\nRaw response:\n{content}",
|
file_id=file_id,
|
||||||
)
|
detail=f"Could not parse valid JSON from GPT response.\nRaw response:\n{content}",
|
||||||
return {}
|
)
|
||||||
|
return {}
|
||||||
metadata = json.loads(json_text)
|
|
||||||
|
metadata = json.loads(json_text)
|
||||||
# SECURITY: Validate filename format from GPT to prevent path traversal
|
|
||||||
# The prompt requests filenames with only letters, numbers, periods, and underscores
|
# SECURITY: Validate filename format from GPT to prevent path traversal
|
||||||
# Enforce this constraint to prevent malicious filenames
|
# The prompt requests filenames with only letters, numbers, periods, and underscores
|
||||||
suggested_filename = metadata.get("filename", "")
|
# Enforce this constraint to prevent malicious filenames
|
||||||
if suggested_filename:
|
suggested_filename = metadata.get("filename", "")
|
||||||
# Check if filename contains only safe characters AND explicitly check for ".."
|
if suggested_filename:
|
||||||
# Defense in depth: While the regex [\w\-\. ]+ already excludes / and \,
|
# Check if filename contains only safe characters AND explicitly check for ".."
|
||||||
# we explicitly reject ".." to guard against:
|
# Defense in depth: While the regex VALID_FILENAME_PATTERN already excludes / and \,
|
||||||
# 1. Potential locale-specific \w behavior
|
# we explicitly reject ".." to guard against:
|
||||||
# 2. Files literally named ".." which are valid but problematic
|
# 1. Potential locale-specific \w behavior
|
||||||
# 3. Future code changes that might relax the regex
|
# 2. Files literally named ".." which are valid but problematic
|
||||||
if not re.match(r"^[\w\-\. ]+$", suggested_filename) or ".." in suggested_filename:
|
# 3. Future code changes that might relax the regex
|
||||||
logger.warning(f"[{task_id}] Invalid filename format from GPT: '{suggested_filename}', using fallback")
|
if not VALID_FILENAME_RE.match(suggested_filename) or ".." in suggested_filename:
|
||||||
# Reset to empty to trigger fallback to original filename
|
logger.warning(f"[{task_id}] Invalid filename format from GPT: '{suggested_filename}', using fallback")
|
||||||
metadata["filename"] = ""
|
# Reset to empty to trigger fallback to original filename
|
||||||
|
metadata["filename"] = ""
|
||||||
logger.info(f"[{task_id}] Extracted metadata: {metadata}")
|
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Extracted metadata: {metadata}")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"parse_metadata",
|
task_id,
|
||||||
"success",
|
"parse_metadata",
|
||||||
f"Parsed metadata: {list(metadata.keys())}",
|
"success",
|
||||||
file_id=file_id,
|
f"Parsed metadata: {list(metadata.keys())}",
|
||||||
detail=f"Extracted metadata:\n{json.dumps(metadata, ensure_ascii=False, indent=2)}",
|
file_id=file_id,
|
||||||
)
|
detail=f"Extracted metadata:\n{json.dumps(metadata, ensure_ascii=False, indent=2)}",
|
||||||
|
)
|
||||||
# Trigger the next step: embedding metadata into the PDF
|
|
||||||
# Pass the filename (can be basename or full path) so embed_metadata_into_pdf can find the file on disk
|
# Trigger the next step: embedding metadata into the PDF
|
||||||
logger.info(f"[{task_id}] Queueing metadata embedding task")
|
# Pass the filename (can be basename or full path) so embed_metadata_into_pdf can find the file on disk
|
||||||
log_task_progress(
|
logger.info(f"[{task_id}] Queueing metadata embedding task")
|
||||||
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, file_id)
|
)
|
||||||
|
embed_metadata_into_pdf.delay(filename, cleaned_text, metadata, file_id)
|
||||||
return {"s3_file": os.path.basename(filename), "metadata": metadata}
|
|
||||||
|
return {"s3_file": os.path.basename(filename), "metadata": metadata}
|
||||||
except Exception as e:
|
|
||||||
logger.exception(f"[{task_id}] AI provider classification failed for {filename}: {e}")
|
except Exception as e:
|
||||||
log_task_progress(
|
logger.exception(f"[{task_id}] AI provider classification failed for {filename}: {e}")
|
||||||
task_id,
|
log_task_progress(
|
||||||
"extract_metadata_with_gpt",
|
task_id,
|
||||||
"failure",
|
"extract_metadata_with_gpt",
|
||||||
f"Exception: {str(e)}",
|
"failure",
|
||||||
file_id=file_id,
|
f"Exception: {str(e)}",
|
||||||
detail=f"AI provider classification failed for {filename}.\nException: {str(e)}",
|
file_id=file_id,
|
||||||
)
|
detail=f"AI provider classification failed for {filename}.\nException: {str(e)}",
|
||||||
return {}
|
)
|
||||||
|
return {}
|
||||||
|
|||||||
@@ -8,6 +8,11 @@ from pathlib import Path
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Pattern for valid filenames (alphanumeric, dash, underscore, period, and space)
|
||||||
|
# Used for validating GPT-provided filenames and other inputs
|
||||||
|
VALID_FILENAME_PATTERN = r"^[\w\-\. ]+$"
|
||||||
|
VALID_FILENAME_RE = re.compile(VALID_FILENAME_PATTERN)
|
||||||
|
|
||||||
|
|
||||||
def get_unique_filename(original_path: str, check_exists_func: Callable[[str], bool] | None = None) -> str:
|
def get_unique_filename(original_path: str, check_exists_func: Callable[[str], bool] | None = None) -> str:
|
||||||
"""
|
"""
|
||||||
|
|||||||
@@ -212,9 +212,10 @@ class TestExtractMetadataFilenameValidation:
|
|||||||
"""Test that invalid filename formats are rejected."""
|
"""Test that invalid filename formats are rejected."""
|
||||||
import re
|
import re
|
||||||
|
|
||||||
# Valid pattern from extract_metadata_with_gpt.py
|
from app.utils.filename_utils import VALID_FILENAME_PATTERN
|
||||||
# TODO: Consider extracting this to a shared constant to avoid duplication
|
|
||||||
valid_pattern = r"^[\w\-\. ]+$"
|
# Valid pattern from app.utils.filename_utils
|
||||||
|
valid_pattern = VALID_FILENAME_PATTERN
|
||||||
|
|
||||||
# Test valid filenames
|
# Test valid filenames
|
||||||
valid_filenames = [
|
valid_filenames = [
|
||||||
|
|||||||
Reference in New Issue
Block a user