Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b3a85fe9c2 | |||
| 7481581e4e | |||
| aae9a89d6b | |||
| 59db28d27b | |||
| a92cace662 | |||
| ffc049196b |
@@ -0,0 +1,78 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
import sqlite3
|
||||||
|
import os
|
||||||
|
import logging
|
||||||
|
from app.config import settings
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
def run_migrations():
|
||||||
|
"""
|
||||||
|
Run database migrations to add missing columns or make other schema changes.
|
||||||
|
"""
|
||||||
|
logger.info("Running database migrations...")
|
||||||
|
|
||||||
|
# Parse the DATABASE_URL to get the SQLite database path
|
||||||
|
db_url = settings.database_url
|
||||||
|
if not db_url.startswith("sqlite:///"):
|
||||||
|
logger.warning(f"Non-SQLite database detected: {db_url}. Migrations may need to be adapted.")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Extract the database path from the URL
|
||||||
|
db_path = db_url.replace("sqlite:///", "")
|
||||||
|
if not os.path.exists(db_path):
|
||||||
|
logger.error(f"Database file not found at {db_path}")
|
||||||
|
return
|
||||||
|
|
||||||
|
# Connect to the database
|
||||||
|
conn = None
|
||||||
|
try:
|
||||||
|
conn = sqlite3.connect(db_path)
|
||||||
|
cursor = conn.cursor()
|
||||||
|
|
||||||
|
# Check if the processing_logs table exists
|
||||||
|
cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='processing_logs';")
|
||||||
|
if not cursor.fetchone():
|
||||||
|
logger.info("Creating processing_logs table...")
|
||||||
|
cursor.execute("""
|
||||||
|
CREATE TABLE processing_logs (
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
file_id INTEGER,
|
||||||
|
step_name VARCHAR,
|
||||||
|
status VARCHAR,
|
||||||
|
message TEXT,
|
||||||
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
||||||
|
FOREIGN KEY (file_id) REFERENCES files (id)
|
||||||
|
);
|
||||||
|
""")
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
# Check if the task_id column exists in processing_logs
|
||||||
|
cursor.execute("PRAGMA table_info(processing_logs);")
|
||||||
|
columns = [row[1] for row in cursor.fetchall()]
|
||||||
|
|
||||||
|
if 'task_id' not in columns:
|
||||||
|
logger.info("Adding task_id column to processing_logs table...")
|
||||||
|
cursor.execute("ALTER TABLE processing_logs ADD COLUMN task_id VARCHAR;")
|
||||||
|
conn.commit()
|
||||||
|
logger.info("Created task_id column in processing_logs")
|
||||||
|
|
||||||
|
# Create an index on task_id for faster lookups
|
||||||
|
cursor.execute("CREATE INDEX idx_processing_logs_task_id ON processing_logs (task_id);")
|
||||||
|
conn.commit()
|
||||||
|
logger.info("Created index on task_id column")
|
||||||
|
|
||||||
|
logger.info("Database migrations completed successfully.")
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Error during database migration: {e}")
|
||||||
|
if conn:
|
||||||
|
conn.rollback()
|
||||||
|
finally:
|
||||||
|
if conn:
|
||||||
|
conn.close()
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# Configure logging
|
||||||
|
logging.basicConfig(level=logging.INFO)
|
||||||
|
run_migrations()
|
||||||
@@ -12,6 +12,7 @@ from uvicorn.middleware.proxy_headers import ProxyHeadersMiddleware
|
|||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
from app.database import init_db
|
from app.database import init_db
|
||||||
|
from app.db_migration import run_migrations
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
from app.tasks.process_document import process_document # Updated import
|
from app.tasks.process_document import process_document # Updated import
|
||||||
from app.tasks.upload_to_dropbox import upload_to_dropbox
|
from app.tasks.upload_to_dropbox import upload_to_dropbox
|
||||||
@@ -52,6 +53,7 @@ app.mount("/static", StaticFiles(directory=frontend_static_dir), name="static")
|
|||||||
@app.on_event("startup")
|
@app.on_event("startup")
|
||||||
def on_startup():
|
def on_startup():
|
||||||
init_db() # Create tables if they don't exist
|
init_db() # Create tables if they don't exist
|
||||||
|
run_migrations() # Run migrations to add any missing columns
|
||||||
|
|
||||||
@app.post("/process/")
|
@app.post("/process/")
|
||||||
def process(file_path: str):
|
def process(file_path: str):
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import json
|
|||||||
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.finalize_document_storage import finalize_document_storage
|
from app.tasks.finalize_document_storage import finalize_document_storage
|
||||||
|
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
|
||||||
@@ -39,6 +40,7 @@ def persist_metadata(metadata, final_pdf_path):
|
|||||||
return json_path
|
return json_path
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
|
@log_task("embed_metadata")
|
||||||
def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata: dict):
|
def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata: dict):
|
||||||
"""
|
"""
|
||||||
Embeds extracted metadata into the PDF's standard metadata fields.
|
Embeds extracted metadata into the PDF's standard metadata fields.
|
||||||
@@ -59,8 +61,10 @@ def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata:
|
|||||||
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):
|
||||||
local_file_path = alt_path
|
local_file_path = alt_path
|
||||||
|
task_logger(f"Using alternative path: {local_file_path}", step_name="embed_metadata")
|
||||||
else:
|
else:
|
||||||
print(f"[ERROR] Local file {local_file_path} not found, cannot embed metadata.")
|
task_logger(f"Local file {local_file_path} not found, cannot embed metadata.",
|
||||||
|
level="error", step_name="embed_metadata")
|
||||||
return {"error": "File not found"}
|
return {"error": "File not found"}
|
||||||
|
|
||||||
# Work on a safe copy in /tmp
|
# Work on a safe copy in /tmp
|
||||||
@@ -70,9 +74,10 @@ def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata:
|
|||||||
|
|
||||||
# Create a safe copy to work on
|
# Create a safe copy to work on
|
||||||
shutil.copy(original_file, processed_file)
|
shutil.copy(original_file, processed_file)
|
||||||
|
task_logger(f"Created working copy at {processed_file}", step_name="embed_metadata")
|
||||||
|
|
||||||
try:
|
try:
|
||||||
print(f"[DEBUG] Embedding metadata into {processed_file}...")
|
task_logger(f"Embedding metadata into {processed_file}", step_name="embed_metadata")
|
||||||
|
|
||||||
# Open the PDF
|
# Open the PDF
|
||||||
doc = fitz.open(processed_file)
|
doc = fitz.open(processed_file)
|
||||||
@@ -87,7 +92,7 @@ def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata:
|
|||||||
doc.save(processed_file, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
|
doc.save(processed_file, incremental=True, encryption=fitz.PDF_ENCRYPT_KEEP)
|
||||||
doc.close()
|
doc.close()
|
||||||
|
|
||||||
print(f"[INFO] Metadata embedded successfully in {processed_file}")
|
task_logger("Metadata embedded successfully", step_name="embed_metadata")
|
||||||
|
|
||||||
# Use the suggested filename from metadata; if not provided, use the original basename.
|
# Use the suggested filename from metadata; if not provided, use the original basename.
|
||||||
suggested_filename = metadata.get("filename", os.path.splitext(os.path.basename(local_file_path))[0])
|
suggested_filename = metadata.get("filename", os.path.splitext(os.path.basename(local_file_path))[0])
|
||||||
@@ -101,28 +106,34 @@ def embed_metadata_into_pdf(local_file_path: str, extracted_text: str, metadata:
|
|||||||
|
|
||||||
# Move the processed file using shutil.move to handle cross-device moves.
|
# Move the processed file using shutil.move to handle cross-device moves.
|
||||||
shutil.move(processed_file, final_file_path)
|
shutil.move(processed_file, final_file_path)
|
||||||
|
task_logger(f"Moved processed file to {final_file_path}", step_name="embed_metadata")
|
||||||
|
|
||||||
# Ensure the temporary file is deleted if it still exists.
|
# Ensure the temporary file is deleted if it still exists.
|
||||||
if os.path.exists(processed_file):
|
if os.path.exists(processed_file):
|
||||||
os.remove(processed_file)
|
os.remove(processed_file)
|
||||||
|
|
||||||
# Persist the metadata into a JSON file with the same base name.
|
# Persist the metadata into a JSON file with the same base name.
|
||||||
json_path = persist_metadata(metadata, final_file_path)
|
json_path = persist_metadata(metadata, final_file_path)
|
||||||
print(f"[INFO] Metadata persisted to {json_path}")
|
task_logger(f"Metadata persisted to {json_path}", step_name="embed_metadata")
|
||||||
|
|
||||||
# Trigger the next step: final storage.
|
# Trigger the next step: final storage.
|
||||||
finalize_document_storage.delay(original_file, final_file_path, metadata)
|
finalize_doc_task = finalize_document_storage.delay(original_file, final_file_path, metadata)
|
||||||
|
task_logger(f"Triggered final document storage with task ID: {finalize_doc_task.id}",
|
||||||
|
step_name="embed_metadata")
|
||||||
|
|
||||||
# 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")
|
||||||
if original_file.startswith(workdir_tmp) and os.path.exists(original_file):
|
if original_file.startswith(workdir_tmp) and os.path.exists(original_file):
|
||||||
try:
|
try:
|
||||||
os.remove(original_file)
|
os.remove(original_file)
|
||||||
print(f"[INFO] Deleted original file from {original_file}")
|
task_logger(f"Deleted original file from {original_file}", step_name="embed_metadata")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"[ERROR] Could not delete original file {original_file}: {e}")
|
task_logger(f"Could not delete original file {original_file}: {e}",
|
||||||
|
level="warning", step_name="embed_metadata")
|
||||||
|
|
||||||
return {"file": final_file_path, "metadata_file": json_path, "status": "Metadata embedded"}
|
return {"file": final_file_path, "metadata_file": json_path, "status": "Metadata embedded"}
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
print(f"[ERROR] Failed to embed metadata into {processed_file}: {e}")
|
task_logger(f"Failed to embed metadata into {processed_file}: {e}",
|
||||||
|
level="error", step_name="embed_metadata")
|
||||||
return {"error": str(e)}
|
return {"error": str(e)}
|
||||||
|
|||||||
@@ -2,9 +2,13 @@
|
|||||||
|
|
||||||
import json
|
import json
|
||||||
import re
|
import re
|
||||||
|
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 task_logger, log_task
|
||||||
|
from app.database import SessionLocal
|
||||||
|
from app.models import FileRecord
|
||||||
|
|
||||||
# Import the shared Celery instance
|
# Import the shared Celery instance
|
||||||
from app.celery_app import celery
|
from app.celery_app import celery
|
||||||
@@ -34,8 +38,15 @@ def extract_json_from_text(text):
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
|
@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()
|
||||||
|
try:
|
||||||
|
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.
|
||||||
@@ -68,8 +79,7 @@ Extracted text:
|
|||||||
Return only valid JSON with no additional commentary.
|
Return only valid JSON with no additional commentary.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
try:
|
task_logger(f"Sending classification request for {s3_filename}", step_name="extract_metadata")
|
||||||
print(f"[DEBUG] Sending classification request for {s3_filename}...")
|
|
||||||
completion = client.chat.completions.create(
|
completion = client.chat.completions.create(
|
||||||
model=settings.openai_model,
|
model=settings.openai_model,
|
||||||
messages=[
|
messages=[
|
||||||
@@ -80,21 +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}",
|
||||||
|
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)
|
||||||
|
task_logger(f"Triggered embed_metadata task with ID: {embed_task.id}", step_name="extract_metadata")
|
||||||
|
|
||||||
return {"s3_file": s3_filename, "metadata": metadata}
|
# Update database record
|
||||||
|
file_record = session.query(FileRecord).filter(FileRecord.local_filename.like(f'%{s3_filename}')).first()
|
||||||
|
if file_record:
|
||||||
|
# Since we can't store dict directly, you might want to store it as JSON string
|
||||||
|
# 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 {"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}",
|
||||||
|
level="error", step_name="extract_metadata")
|
||||||
return {}
|
return {}
|
||||||
|
finally:
|
||||||
|
session.close()
|
||||||
|
|||||||
@@ -2,25 +2,29 @@
|
|||||||
|
|
||||||
from app.config import settings
|
from app.config import settings
|
||||||
from app.tasks.retry_config import BaseTaskWithRetry
|
from app.tasks.retry_config import BaseTaskWithRetry
|
||||||
# Import the shared Celery instance
|
|
||||||
from app.celery_app import celery
|
from app.celery_app import celery
|
||||||
|
from app.utils import task_logger, log_task
|
||||||
|
|
||||||
# 1) Import the aggregator task
|
# Import the aggregator task
|
||||||
from app.tasks.send_to_all import send_to_all_destinations
|
from app.tasks.send_to_all import send_to_all_destinations
|
||||||
|
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
|
@log_task("finalize_storage")
|
||||||
def finalize_document_storage(original_file: str, processed_file: str, metadata: dict):
|
def finalize_document_storage(original_file: str, processed_file: str, metadata: dict):
|
||||||
"""
|
"""
|
||||||
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.
|
||||||
"""
|
"""
|
||||||
print(f"[INFO] Finalizing document storage for {processed_file}")
|
task_logger(f"Finalizing document storage for {processed_file}", step_name="finalize_storage")
|
||||||
|
|
||||||
# 2) Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
# Enqueue uploads to all destinations (Dropbox, Nextcloud, Paperless)
|
||||||
send_to_all_destinations.delay(processed_file)
|
send_task = send_to_all_destinations.delay(processed_file)
|
||||||
|
|
||||||
|
task_logger(f"Triggered send to all destinations with task ID: {send_task.id}",
|
||||||
|
step_name="finalize_storage", status="success")
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"status": "Completed",
|
"status": "Completed",
|
||||||
"file": processed_file
|
"file": processed_file,
|
||||||
|
"send_task_id": send_task.id
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
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.
|
||||||
@@ -28,12 +29,15 @@ def process_document(original_local_file: str):
|
|||||||
- Check for embedded text. If present, run local GPT extraction
|
- Check for embedded text. If present, run local GPT extraction
|
||||||
- Otherwise, queue Textract-based OCR
|
- Otherwise, queue Textract-based OCR
|
||||||
"""
|
"""
|
||||||
|
task_id = process_document.request.id
|
||||||
|
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):
|
||||||
print(f"[ERROR] File {original_local_file} not found.")
|
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
|
||||||
|
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)
|
||||||
@@ -43,16 +47,19 @@ def process_document(original_local_file: str):
|
|||||||
|
|
||||||
# Acquire DB session in the task
|
# Acquire DB session in the task
|
||||||
with SessionLocal() as db:
|
with SessionLocal() as db:
|
||||||
existing = db.query(FileRecord).filter_by(filehash=filehash).one_or_none()
|
# Keep all FileRecord operations within this session scope
|
||||||
if existing:
|
task_logger(f"Checking for duplicate files", step_name="check_duplicates", task_id=task_id)
|
||||||
print(f"[INFO] Duplicate file detected (hash={filehash[:10]}...) Skipping processing.")
|
existing_record = db.query(FileRecord).filter(FileRecord.filehash == filehash).one_or_none()
|
||||||
|
if existing_record:
|
||||||
|
task_logger(f"Duplicate file detected (hash={filehash[:10]}...). Skipping processing.",
|
||||||
|
step_name="process_document", task_id=task_id, file_id=existing_record.id, status="success")
|
||||||
return {
|
return {
|
||||||
"status": "duplicate_file",
|
"status": "duplicate_file",
|
||||||
"file_id": existing.id,
|
"file_id": existing_record.id,
|
||||||
"detail": "File already processed."
|
"detail": "File already processed."
|
||||||
}
|
}
|
||||||
|
else:
|
||||||
# Not a duplicate -> insert a new record
|
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,
|
||||||
@@ -65,6 +72,7 @@ def process_document(original_local_file: str):
|
|||||||
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
|
||||||
|
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}"
|
||||||
@@ -80,25 +88,39 @@ def process_document(original_local_file: str):
|
|||||||
new_record.local_filename = new_local_path
|
new_record.local_filename = new_local_path
|
||||||
db.commit()
|
db.commit()
|
||||||
|
|
||||||
|
# Perform all further interactions with existing_record/new_record here
|
||||||
|
|
||||||
# 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)
|
||||||
|
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:
|
||||||
print(f"[INFO] PDF {original_local_file} contains embedded text. Processing locally.")
|
task_logger(f"PDF {original_local_file} contains embedded text. Processing locally.",
|
||||||
|
step_name="process_document", task_id=task_id, file_id=new_record.id)
|
||||||
|
|
||||||
# Extract text locally
|
# Extract text locally
|
||||||
extracted_text = ""
|
extracted_text = ""
|
||||||
|
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
|
||||||
extract_metadata_with_gpt.delay(new_filename, extracted_text)
|
task_logger(f"Text extracted locally. Queuing for metadata extraction.",
|
||||||
return {"file": new_local_path, "status": "Text extracted locally"}
|
step_name="process_document", task_id=task_id, file_id=new_record.id, status="success")
|
||||||
|
metadata_task = 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}
|
||||||
|
|
||||||
# 3. If no embedded text, queue Textract processing
|
# 3. If no embedded text, queue Textract processing
|
||||||
process_with_textract.delay(new_filename)
|
task_logger(f"No embedded text found. Queuing for OCR.",
|
||||||
return {"file": new_local_path, "status": "Queued for OCR"}
|
step_name="process_document", task_id=task_id, file_id=new_record.id, status="success")
|
||||||
|
ocr_task = 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}
|
||||||
|
|||||||
@@ -8,6 +8,9 @@ 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 task_logger, log_task
|
||||||
|
from app.database import SessionLocal
|
||||||
|
from app.models import FileRecord
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -18,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
|
||||||
@@ -30,12 +34,31 @@ def process_with_textract(s3_filename: str):
|
|||||||
4. Extracts the text content for metadata processing.
|
4. Extracts the text content for metadata processing.
|
||||||
5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
|
5. Triggers downstream metadata extraction by calling extract_metadata_with_gpt.
|
||||||
"""
|
"""
|
||||||
try:
|
task_id = process_with_textract.request.id
|
||||||
tmp_file_path = os.path.join(settings.workdir, "tmp", s3_filename)
|
tmp_file_path = os.path.join(settings.workdir, "tmp", s3_filename)
|
||||||
|
|
||||||
|
# Get the file_id from the database
|
||||||
|
file_id = None
|
||||||
|
with SessionLocal() as db:
|
||||||
|
file_record = db.query(FileRecord).filter(
|
||||||
|
FileRecord.local_filename == tmp_file_path
|
||||||
|
).first()
|
||||||
|
if file_record:
|
||||||
|
file_id = file_record.id
|
||||||
|
|
||||||
|
task_logger(f"Starting OCR for {s3_filename}", step_name="process_with_textract",
|
||||||
|
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):
|
||||||
|
task_logger(f"Local file not found: {tmp_file_path}", level="error",
|
||||||
|
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}")
|
||||||
|
|
||||||
logger.info(f"Processing {s3_filename} with Azure Document Intelligence OCR.")
|
try:
|
||||||
|
task_logger(f"Sending document to Azure Document Intelligence",
|
||||||
|
step_name="azure_document_intelligence", task_id=task_id,
|
||||||
|
file_id=file_id, file_path=tmp_file_path)
|
||||||
|
|
||||||
# Open and send the document for processing
|
# Open and send the document for processing
|
||||||
with open(tmp_file_path, "rb") as f:
|
with open(tmp_file_path, "rb") as f:
|
||||||
@@ -45,23 +68,34 @@ def process_with_textract(s3_filename: str):
|
|||||||
result: AnalyzeResult = poller.result()
|
result: AnalyzeResult = poller.result()
|
||||||
operation_id = poller.details["operation_id"]
|
operation_id = poller.details["operation_id"]
|
||||||
|
|
||||||
|
task_logger(f"Azure Document Intelligence processing complete, operation ID: {operation_id}",
|
||||||
|
step_name="azure_document_intelligence", task_id=task_id)
|
||||||
|
|
||||||
# Retrieve the processed searchable PDF
|
# Retrieve the processed searchable PDF
|
||||||
|
task_logger(f"Retrieving searchable PDF", step_name="retrieve_pdf", task_id=task_id)
|
||||||
response = document_intelligence_client.get_analyze_result_pdf(
|
response = document_intelligence_client.get_analyze_result_pdf(
|
||||||
model_id=result.model_id, result_id=operation_id
|
model_id=result.model_id, result_id=operation_id
|
||||||
)
|
)
|
||||||
searchable_pdf_path = tmp_file_path # Overwrite the original PDF location
|
searchable_pdf_path = tmp_file_path # Overwrite the original PDF location
|
||||||
with open(searchable_pdf_path, "wb") as writer:
|
with open(searchable_pdf_path, "wb") as writer:
|
||||||
writer.writelines(response)
|
writer.writelines(response)
|
||||||
logger.info(f"Searchable PDF saved at: {searchable_pdf_path}")
|
|
||||||
|
|
||||||
# Extract raw text content from the result
|
# Extract raw text content from the result
|
||||||
extracted_text = result.content if result.content else ""
|
extracted_text = result.content if result.content else ""
|
||||||
logger.info(f"Extracted text for {s3_filename}: {len(extracted_text)} characters")
|
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
|
||||||
extract_metadata_with_gpt.delay(s3_filename, extracted_text)
|
task_logger(f"OCR completed. Queueing metadata extraction for {s3_filename}",
|
||||||
|
step_name="process_with_textract", task_id=task_id, status="success")
|
||||||
|
|
||||||
return {"s3_file": s3_filename, "searchable_pdf": searchable_pdf_path, "cleaned_text": 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 {"file": s3_filename, "searchable_pdf": searchable_pdf_path, "text_length": text_length}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error processing {s3_filename} with Azure Document Intelligence: {e}")
|
task_logger(f"Error processing with Azure Document Intelligence: {e}",
|
||||||
|
level="error", step_name="process_with_textract", task_id=task_id)
|
||||||
raise
|
raise
|
||||||
|
|||||||
@@ -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}
|
||||||
|
|
||||||
|
|||||||
@@ -4,18 +4,31 @@ from app.celery_app import celery
|
|||||||
from app.tasks.upload_to_dropbox import upload_to_dropbox
|
from app.tasks.upload_to_dropbox import upload_to_dropbox
|
||||||
from app.tasks.upload_to_nextcloud import upload_to_nextcloud
|
from app.tasks.upload_to_nextcloud import upload_to_nextcloud
|
||||||
from app.tasks.upload_to_paperless import upload_to_paperless
|
from app.tasks.upload_to_paperless import upload_to_paperless
|
||||||
|
from app.utils import task_logger, log_task
|
||||||
|
|
||||||
@celery.task
|
@celery.task
|
||||||
|
@log_task("send_to_all_destinations")
|
||||||
def send_to_all_destinations(file_path: str):
|
def send_to_all_destinations(file_path: str):
|
||||||
"""
|
"""
|
||||||
Fires off tasks to upload a single file to Dropbox, Nextcloud, and Paperless.
|
Fires off tasks to upload a single file to Dropbox, Nextcloud, and Paperless.
|
||||||
These tasks run in parallel (Celery returns immediately from each .delay()).
|
These tasks run in parallel (Celery returns immediately from each .delay()).
|
||||||
"""
|
"""
|
||||||
upload_to_dropbox.delay(file_path)
|
task_logger(f"Sending {file_path} to all destinations", step_name="send_to_all")
|
||||||
upload_to_nextcloud.delay(file_path)
|
|
||||||
upload_to_paperless.delay(file_path)
|
dropbox_task = upload_to_dropbox.delay(file_path)
|
||||||
|
nextcloud_task = upload_to_nextcloud.delay(file_path)
|
||||||
|
paperless_task = upload_to_paperless.delay(file_path)
|
||||||
|
|
||||||
|
task_logger(f"Enqueued file for all destinations: Dropbox (task: {dropbox_task.id}), "
|
||||||
|
f"Nextcloud (task: {nextcloud_task.id}), Paperless (task: {paperless_task.id})",
|
||||||
|
step_name="send_to_all", status="success")
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"status": "All upload tasks enqueued",
|
"status": "All upload tasks enqueued",
|
||||||
"file_path": file_path
|
"file_path": file_path,
|
||||||
|
"task_ids": {
|
||||||
|
"dropbox": dropbox_task.id,
|
||||||
|
"nextcloud": nextcloud_task.id,
|
||||||
|
"paperless": paperless_task.id
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -6,6 +6,7 @@ import dropbox
|
|||||||
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.celery_app import celery
|
from app.celery_app import celery
|
||||||
|
from app.utils import task_logger, log_task
|
||||||
|
|
||||||
def get_dropbox_access_token():
|
def get_dropbox_access_token():
|
||||||
"""Refresh the Dropbox access token using the stored refresh token from ENV."""
|
"""Refresh the Dropbox access token using the stored refresh token from ENV."""
|
||||||
@@ -25,14 +26,16 @@ def get_dropbox_access_token():
|
|||||||
return response.json()["access_token"]
|
return response.json()["access_token"]
|
||||||
else:
|
else:
|
||||||
error_msg = f"Failed to refresh Dropbox token: {response.status_code} - {response.text}"
|
error_msg = f"Failed to refresh Dropbox token: {response.status_code} - {response.text}"
|
||||||
print(f"[ERROR] {error_msg}")
|
task_logger(error_msg, level="error", step_name="dropbox_auth")
|
||||||
raise Exception(error_msg)
|
raise Exception(error_msg)
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
|
@log_task("upload_to_dropbox")
|
||||||
def upload_to_dropbox(file_path: str):
|
def upload_to_dropbox(file_path: str):
|
||||||
"""Uploads a file to Dropbox using the API."""
|
"""Uploads a file to Dropbox using the API."""
|
||||||
|
|
||||||
if not os.path.exists(file_path):
|
if not os.path.exists(file_path):
|
||||||
|
task_logger(f"File not found: {file_path}", level="error", step_name="dropbox_upload")
|
||||||
raise FileNotFoundError(f"File not found: {file_path}")
|
raise FileNotFoundError(f"File not found: {file_path}")
|
||||||
|
|
||||||
# Extract filename and set target path
|
# Extract filename and set target path
|
||||||
@@ -41,16 +44,20 @@ def upload_to_dropbox(file_path: str):
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
# Get fresh access token
|
# Get fresh access token
|
||||||
|
task_logger(f"Getting Dropbox access token", step_name="dropbox_auth")
|
||||||
access_token = get_dropbox_access_token()
|
access_token = get_dropbox_access_token()
|
||||||
dbx = dropbox.Dropbox(access_token)
|
dbx = dropbox.Dropbox(access_token)
|
||||||
|
|
||||||
file_size = os.path.getsize(file_path)
|
file_size = os.path.getsize(file_path)
|
||||||
chunk_size = 4 * 1024 * 1024 # 4MB chunk size
|
chunk_size = 4 * 1024 * 1024 # 4MB chunk size
|
||||||
|
|
||||||
|
task_logger(f"Starting upload of {filename} ({file_size} bytes) to Dropbox", step_name="dropbox_upload")
|
||||||
|
|
||||||
with open(file_path, "rb") as file_data:
|
with open(file_path, "rb") as file_data:
|
||||||
if file_size <= chunk_size:
|
if file_size <= chunk_size:
|
||||||
dbx.files_upload(file_data.read(), dropbox_path)
|
dbx.files_upload(file_data.read(), dropbox_path)
|
||||||
else:
|
else:
|
||||||
|
task_logger(f"Using chunked upload for {filename}", step_name="dropbox_upload")
|
||||||
upload_session_start_result = dbx.files_upload_session_start(file_data.read(chunk_size))
|
upload_session_start_result = dbx.files_upload_session_start(file_data.read(chunk_size))
|
||||||
cursor = dropbox.files.UploadSessionCursor(
|
cursor = dropbox.files.UploadSessionCursor(
|
||||||
session_id=upload_session_start_result.session_id,
|
session_id=upload_session_start_result.session_id,
|
||||||
@@ -65,10 +72,10 @@ def upload_to_dropbox(file_path: str):
|
|||||||
dbx.files_upload_session_append_v2(file_data.read(chunk_size), cursor)
|
dbx.files_upload_session_append_v2(file_data.read(chunk_size), cursor)
|
||||||
cursor.offset = file_data.tell()
|
cursor.offset = file_data.tell()
|
||||||
|
|
||||||
print(f"[INFO] Successfully uploaded {filename} to Dropbox at {dropbox_path}.")
|
task_logger(f"Successfully uploaded {filename} to Dropbox at {dropbox_path}", step_name="dropbox_upload", status="success")
|
||||||
return {"status": "Completed", "file": file_path}
|
return {"status": "Completed", "file": file_path}
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
error_msg = f"[ERROR] Failed to upload {filename} to Dropbox: {str(e)}"
|
error_msg = f"Failed to upload {filename} to Dropbox: {str(e)}"
|
||||||
print(error_msg)
|
task_logger(error_msg, level="error", step_name="dropbox_upload", status="failure")
|
||||||
raise Exception(error_msg)
|
raise Exception(error_msg)
|
||||||
|
|||||||
@@ -5,12 +5,15 @@ import requests
|
|||||||
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.celery_app import celery
|
from app.celery_app import celery
|
||||||
|
from app.utils import task_logger, log_task
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
|
@log_task("upload_to_nextcloud")
|
||||||
def upload_to_nextcloud(file_path: str):
|
def upload_to_nextcloud(file_path: str):
|
||||||
"""Uploads a file to Nextcloud in the configured folder."""
|
"""Uploads a file to Nextcloud in the configured folder."""
|
||||||
|
|
||||||
if not os.path.exists(file_path):
|
if not os.path.exists(file_path):
|
||||||
|
task_logger(f"File not found: {file_path}", level="error", step_name="nextcloud_upload")
|
||||||
raise FileNotFoundError(f"File not found: {file_path}")
|
raise FileNotFoundError(f"File not found: {file_path}")
|
||||||
|
|
||||||
# Extract filename
|
# Extract filename
|
||||||
@@ -19,6 +22,8 @@ def upload_to_nextcloud(file_path: str):
|
|||||||
# Construct the full upload URL
|
# Construct the full upload URL
|
||||||
nextcloud_url = f"{settings.nextcloud_upload_url}/{settings.nextcloud_folder}/{filename}"
|
nextcloud_url = f"{settings.nextcloud_upload_url}/{settings.nextcloud_folder}/{filename}"
|
||||||
|
|
||||||
|
task_logger(f"Starting upload of {filename} to Nextcloud", step_name="nextcloud_upload")
|
||||||
|
|
||||||
# Read file content
|
# Read file content
|
||||||
with open(file_path, "rb") as file_data:
|
with open(file_path, "rb") as file_data:
|
||||||
response = requests.put(
|
response = requests.put(
|
||||||
@@ -29,9 +34,10 @@ def upload_to_nextcloud(file_path: str):
|
|||||||
|
|
||||||
# Check if upload was successful
|
# Check if upload was successful
|
||||||
if response.status_code in (200, 201):
|
if response.status_code in (200, 201):
|
||||||
print(f"[INFO] Successfully uploaded {filename} to Nextcloud at {nextcloud_url}.")
|
task_logger(f"Successfully uploaded {filename} to Nextcloud at {nextcloud_url}",
|
||||||
|
step_name="nextcloud_upload", status="success")
|
||||||
return {"status": "Completed", "file": file_path}
|
return {"status": "Completed", "file": file_path}
|
||||||
else:
|
else:
|
||||||
error_msg = f"[ERROR] Failed to upload {filename} to Nextcloud: {response.status_code} - {response.text}"
|
error_msg = f"Failed to upload {filename} to Nextcloud: {response.status_code} - {response.text}"
|
||||||
print(error_msg)
|
task_logger(error_msg, level="error", step_name="nextcloud_upload", status="failure")
|
||||||
raise Exception(error_msg)
|
raise Exception(error_msg)
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ from typing import Dict, Any
|
|||||||
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.celery_app import celery
|
from app.celery_app import celery
|
||||||
|
from app.utils import task_logger, log_task
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -46,13 +47,15 @@ def poll_task_for_document_id(task_id: str) -> int:
|
|||||||
|
|
||||||
while attempts < POLL_MAX_ATTEMPTS:
|
while attempts < POLL_MAX_ATTEMPTS:
|
||||||
try:
|
try:
|
||||||
|
task_logger(f"Polling Paperless for task {task_id}, attempt {attempts+1}/{POLL_MAX_ATTEMPTS}",
|
||||||
|
step_name="paperless_poll")
|
||||||
resp = requests.get(url, headers=_get_headers(), params={"task_id": task_id})
|
resp = requests.get(url, headers=_get_headers(), params={"task_id": task_id})
|
||||||
resp.raise_for_status()
|
resp.raise_for_status()
|
||||||
tasks_data = resp.json()
|
tasks_data = resp.json()
|
||||||
except requests.exceptions.RequestException as exc:
|
except requests.exceptions.RequestException as exc:
|
||||||
logger.warning(
|
task_logger(
|
||||||
"Failed to poll for task_id='%s'. Attempt=%d Error=%s",
|
f"Failed to poll for task_id='{task_id}'. Attempt={attempts + 1}/{POLL_MAX_ATTEMPTS} Error={exc}",
|
||||||
task_id, attempts + 1, exc
|
level="warning", step_name="paperless_poll"
|
||||||
)
|
)
|
||||||
time.sleep(POLL_INTERVAL_SEC)
|
time.sleep(POLL_INTERVAL_SEC)
|
||||||
attempts += 1
|
attempts += 1
|
||||||
@@ -67,21 +70,29 @@ def poll_task_for_document_id(task_id: str) -> int:
|
|||||||
if status == "SUCCESS":
|
if status == "SUCCESS":
|
||||||
doc_str = task_info.get("related_document")
|
doc_str = task_info.get("related_document")
|
||||||
if doc_str:
|
if doc_str:
|
||||||
|
task_logger(f"Task {task_id} completed successfully with document ID: {doc_str}",
|
||||||
|
step_name="paperless_poll", status="success")
|
||||||
return int(doc_str)
|
return int(doc_str)
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
f"Task {task_id} completed but no doc ID found. Task info: {task_info}"
|
f"Task {task_id} completed but no doc ID found. Task info: {task_info}"
|
||||||
)
|
)
|
||||||
elif status == "FAILURE":
|
elif status == "FAILURE":
|
||||||
raise RuntimeError(f"Task {task_id} failed: {task_info.get('result')}")
|
error_msg = f"Task {task_id} failed: {task_info.get('result')}"
|
||||||
|
task_logger(error_msg, level="error", step_name="paperless_poll", status="failure")
|
||||||
|
raise RuntimeError(error_msg)
|
||||||
|
else:
|
||||||
|
task_logger(f"Task {task_id} status: {status}, waiting {POLL_INTERVAL_SEC}s",
|
||||||
|
step_name="paperless_poll")
|
||||||
|
|
||||||
attempts += 1
|
attempts += 1
|
||||||
time.sleep(POLL_INTERVAL_SEC)
|
time.sleep(POLL_INTERVAL_SEC)
|
||||||
|
|
||||||
raise TimeoutError(
|
timeout_msg = f"Task {task_id} didn't reach SUCCESS within {POLL_MAX_ATTEMPTS} attempts."
|
||||||
f"Task {task_id} didn't reach SUCCESS within {POLL_MAX_ATTEMPTS} attempts."
|
task_logger(timeout_msg, level="error", step_name="paperless_poll", status="failure")
|
||||||
)
|
raise TimeoutError(timeout_msg)
|
||||||
|
|
||||||
@celery.task(base=BaseTaskWithRetry)
|
@celery.task(base=BaseTaskWithRetry)
|
||||||
|
@log_task("upload_to_paperless")
|
||||||
def upload_to_paperless(file_path: str) -> Dict[str, Any]:
|
def upload_to_paperless(file_path: str) -> Dict[str, Any]:
|
||||||
"""
|
"""
|
||||||
Uploads a PDF to Paperless with minimal metadata (filename and date only).
|
Uploads a PDF to Paperless with minimal metadata (filename and date only).
|
||||||
@@ -93,9 +104,11 @@ def upload_to_paperless(file_path: str) -> Dict[str, Any]:
|
|||||||
Returns a dict with status, the paperless_task_id, paperless_document_id, and file_path.
|
Returns a dict with status, the paperless_task_id, paperless_document_id, and file_path.
|
||||||
"""
|
"""
|
||||||
if not os.path.exists(file_path):
|
if not os.path.exists(file_path):
|
||||||
|
task_logger(f"File not found: {file_path}", level="error", step_name="paperless_upload")
|
||||||
raise FileNotFoundError(f"File not found: {file_path}")
|
raise FileNotFoundError(f"File not found: {file_path}")
|
||||||
|
|
||||||
base_name = os.path.basename(file_path)
|
base_name = os.path.basename(file_path)
|
||||||
|
task_logger(f"Starting upload of {base_name} to Paperless", step_name="paperless_upload")
|
||||||
|
|
||||||
# Upload the PDF
|
# Upload the PDF
|
||||||
post_url = _paperless_api_url("/api/documents/post_document/")
|
post_url = _paperless_api_url("/api/documents/post_document/")
|
||||||
@@ -106,22 +119,21 @@ def upload_to_paperless(file_path: str) -> Dict[str, Any]:
|
|||||||
data = {"title": base_name} # Title = Filename (no additional metadata)
|
data = {"title": base_name} # Title = Filename (no additional metadata)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
logger.debug("Posting document to Paperless: file=%s", base_name)
|
task_logger(f"Posting document to Paperless: file={base_name}", step_name="paperless_upload")
|
||||||
resp = requests.post(post_url, headers=_get_headers(), files=files, data=data)
|
resp = requests.post(post_url, headers=_get_headers(), files=files, data=data)
|
||||||
resp.raise_for_status()
|
resp.raise_for_status()
|
||||||
except requests.exceptions.RequestException as exc:
|
except requests.exceptions.RequestException as exc:
|
||||||
logger.error(
|
error_msg = f"Failed to upload document '{file_path}' to Paperless. Error: {exc}. Response={getattr(exc.response, 'text', '<no response>')}"
|
||||||
"Failed to upload document '%s' to Paperless. Error: %s. Response=%s",
|
task_logger(error_msg, level="error", step_name="paperless_upload", status="failure")
|
||||||
file_path, exc, getattr(exc.response, "text", "<no response>")
|
|
||||||
)
|
|
||||||
raise
|
raise
|
||||||
|
|
||||||
raw_task_id = resp.text.strip().strip('"').strip("'")
|
raw_task_id = resp.text.strip().strip('"').strip("'")
|
||||||
logger.info(f"Received Paperless task ID: {raw_task_id}")
|
task_logger(f"Received Paperless task ID: {raw_task_id}", step_name="paperless_upload")
|
||||||
|
|
||||||
# Poll tasks until success/fail => get doc_id
|
# Poll tasks until success/fail => get doc_id
|
||||||
doc_id = poll_task_for_document_id(raw_task_id)
|
doc_id = poll_task_for_document_id(raw_task_id)
|
||||||
logger.info(f"Document {file_path} successfully ingested => ID={doc_id}")
|
task_logger(f"Document {file_path} successfully ingested => ID={doc_id}",
|
||||||
|
step_name="paperless_upload", status="success")
|
||||||
|
|
||||||
return {
|
return {
|
||||||
"status": "Completed",
|
"status": "Completed",
|
||||||
|
|||||||
+145
-3
@@ -1,7 +1,14 @@
|
|||||||
# app/utils.py
|
# app/utils.py
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import logging
|
||||||
|
import contextlib
|
||||||
|
from functools import wraps
|
||||||
|
from typing import Optional, Callable
|
||||||
|
from celery import Task
|
||||||
from app.database import SessionLocal
|
from app.database import SessionLocal
|
||||||
from app.models import ProcessingLog
|
from app.models import ProcessingLog, FileRecord
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
def hash_file(filepath, chunk_size=65536):
|
def hash_file(filepath, chunk_size=65536):
|
||||||
"""
|
"""
|
||||||
@@ -18,12 +25,29 @@ def hash_file(filepath, chunk_size=65536):
|
|||||||
return sha256.hexdigest()
|
return sha256.hexdigest()
|
||||||
|
|
||||||
|
|
||||||
|
def log_task_progress(task_id: str, step_name: str, status: str, message: Optional[str] = None,
|
||||||
def log_task_progress(task_id, step_name, status, message=None, file_id=None):
|
file_id: Optional[int] = None, file_path: Optional[str] = None):
|
||||||
"""
|
"""
|
||||||
Logs the progress of a Celery task to the database.
|
Logs the progress of a Celery task to the database.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
task_id (str): The Celery task ID
|
||||||
|
step_name (str): Name of the processing step
|
||||||
|
status (str): Status of the step ("pending", "in_progress", "success", "failure")
|
||||||
|
message (str, optional): Additional message or error details
|
||||||
|
file_id (int, optional): ID of associated FileRecord
|
||||||
|
file_path (str, optional): Path to file - will attempt to find file_id from path
|
||||||
"""
|
"""
|
||||||
|
try:
|
||||||
with SessionLocal() as db:
|
with SessionLocal() as db:
|
||||||
|
# If file_path is provided but not file_id, try to look up the file_id
|
||||||
|
if not file_id and file_path:
|
||||||
|
file_record = db.query(FileRecord).filter(
|
||||||
|
FileRecord.local_filename == file_path
|
||||||
|
).first()
|
||||||
|
if file_record:
|
||||||
|
file_id = file_record.id
|
||||||
|
|
||||||
log_entry = ProcessingLog(
|
log_entry = ProcessingLog(
|
||||||
task_id=task_id,
|
task_id=task_id,
|
||||||
step_name=step_name,
|
step_name=step_name,
|
||||||
@@ -33,3 +57,121 @@ def log_task_progress(task_id, step_name, status, message=None, file_id=None):
|
|||||||
)
|
)
|
||||||
db.add(log_entry)
|
db.add(log_entry)
|
||||||
db.commit()
|
db.commit()
|
||||||
|
logger.info(f"Task {task_id} - {step_name}: {status} {message or ''}")
|
||||||
|
return log_entry.id
|
||||||
|
except Exception as e:
|
||||||
|
logger.error(f"Failed to log task progress: {e}")
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
@contextlib.contextmanager
|
||||||
|
def task_step_logging(task_id: str, step_name: str, file_id: Optional[int] = None, file_path: Optional[str] = None):
|
||||||
|
"""
|
||||||
|
Context manager for logging the beginning and end of a task step.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
with task_step_logging(task.request.id, "extract_text", file_path=pdf_path):
|
||||||
|
# Do the actual work
|
||||||
|
text = extract_text_from_pdf(pdf_path)
|
||||||
|
"""
|
||||||
|
log_id = log_task_progress(task_id, step_name, "in_progress",
|
||||||
|
"Starting processing step", file_id, file_path)
|
||||||
|
try:
|
||||||
|
yield
|
||||||
|
log_task_progress(task_id, step_name, "success",
|
||||||
|
"Successfully completed", file_id, file_path)
|
||||||
|
except Exception as e:
|
||||||
|
log_task_progress(task_id, step_name, "failure",
|
||||||
|
f"Error: {str(e)}", file_id, file_path)
|
||||||
|
raise # Re-raise the exception after logging
|
||||||
|
|
||||||
|
|
||||||
|
def log_task(step_name: str):
|
||||||
|
"""
|
||||||
|
Decorator for Celery tasks to automatically log progress.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
@celery.task
|
||||||
|
@log_task("process_pdf")
|
||||||
|
def process_pdf(file_path):
|
||||||
|
# Task implementation
|
||||||
|
"""
|
||||||
|
def decorator(func: Callable):
|
||||||
|
@wraps(func)
|
||||||
|
def wrapper(*args, **kwargs):
|
||||||
|
# Get task_id from Celery's current task
|
||||||
|
task = wrapper.request if hasattr(wrapper, 'request') else None
|
||||||
|
task_id = task.id if task else "unknown_task"
|
||||||
|
|
||||||
|
# Try to determine file_id or file_path from arguments
|
||||||
|
file_path = None
|
||||||
|
if args and isinstance(args[0], str):
|
||||||
|
file_path = args[0] # Assume first arg is file path
|
||||||
|
|
||||||
|
# Log start
|
||||||
|
log_task_progress(task_id, step_name, "pending", "Task queued", file_path=file_path)
|
||||||
|
|
||||||
|
try:
|
||||||
|
# Log in_progress
|
||||||
|
log_task_progress(task_id, step_name, "in_progress", "Task started", file_path=file_path)
|
||||||
|
|
||||||
|
# Execute the task
|
||||||
|
result = func(*args, **kwargs)
|
||||||
|
|
||||||
|
# Log success
|
||||||
|
log_task_progress(task_id, step_name, "success", "Task completed", file_path=file_path)
|
||||||
|
|
||||||
|
return result
|
||||||
|
except Exception as e:
|
||||||
|
# Log failure
|
||||||
|
log_task_progress(task_id, step_name, "failure", f"Error: {str(e)}", file_path=file_path)
|
||||||
|
raise # Re-raise the exception
|
||||||
|
|
||||||
|
return wrapper
|
||||||
|
return decorator
|
||||||
|
|
||||||
|
|
||||||
|
def task_logger(message: str, level: str = "info", task_id: str = None, step_name: str = None,
|
||||||
|
status: str = None, file_path: Optional[str] = None, file_id: Optional[int] = None):
|
||||||
|
"""
|
||||||
|
Unified logging function that logs to both console and database.
|
||||||
|
This replaces print() statements in tasks with proper logging.
|
||||||
|
|
||||||
|
Parameters:
|
||||||
|
message: The log message
|
||||||
|
level: Log level (info, error, warning, debug)
|
||||||
|
task_id: Celery task ID (tries to get from current task if None)
|
||||||
|
step_name: Name of the processing step
|
||||||
|
status: Status for database logging (pending, in_progress, success, failure)
|
||||||
|
file_path: Path to the file being processed
|
||||||
|
file_id: ID of the FileRecord
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
task_logger("Processing file", task_id=task.request.id, step_name="process_pdf")
|
||||||
|
task_logger("Error processing file", level="error")
|
||||||
|
"""
|
||||||
|
# Get task_id from current task if not provided
|
||||||
|
if task_id is None:
|
||||||
|
from celery._state import get_current_task
|
||||||
|
current_task = get_current_task()
|
||||||
|
task_id = current_task.request.id if current_task else "unknown_task"
|
||||||
|
|
||||||
|
# Default step name if not provided
|
||||||
|
if step_name is None:
|
||||||
|
step_name = "general"
|
||||||
|
|
||||||
|
# Default status if not provided
|
||||||
|
if status is None:
|
||||||
|
if level == "error":
|
||||||
|
status = "failure"
|
||||||
|
elif level == "warning":
|
||||||
|
status = "warning"
|
||||||
|
else:
|
||||||
|
status = "in_progress"
|
||||||
|
|
||||||
|
# Log to console
|
||||||
|
log_method = getattr(logger, level.lower(), logger.info)
|
||||||
|
log_method(f"[{step_name}] {message}")
|
||||||
|
|
||||||
|
# Log to database
|
||||||
|
return log_task_progress(task_id, step_name, status, message, file_id, file_path)
|
||||||
|
|||||||
@@ -0,0 +1,23 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
|
||||||
|
"""
|
||||||
|
Script to run database migrations manually.
|
||||||
|
Run this script to update the database schema before starting the application.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python migrate_db.py
|
||||||
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
from app.db_migration import run_migrations
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# Configure logging
|
||||||
|
logging.basicConfig(
|
||||||
|
level=logging.INFO,
|
||||||
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||||
|
)
|
||||||
|
|
||||||
|
print("Running database migrations...")
|
||||||
|
run_migrations()
|
||||||
|
print("Database migrations completed.")
|
||||||
Reference in New Issue
Block a user