copilot-swe-agent[bot]
ad74bc852d
Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
267 lines
8.0 KiB
Python
267 lines
8.0 KiB
Python
"""
|
|
Utility functions for managing file processing step status.
|
|
|
|
This module provides functions to initialize, update, and query the status
|
|
of file processing steps using the FileProcessingStep model.
|
|
"""
|
|
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.models import FileProcessingStep
|
|
|
|
# Define the expected processing steps for a standard file workflow
|
|
MAIN_PROCESSING_STEPS = [
|
|
"hash_file",
|
|
"create_file_record",
|
|
"check_text",
|
|
"extract_text", # Or "process_with_azure_document_intelligence"
|
|
"extract_metadata_with_gpt",
|
|
"embed_metadata_into_pdf",
|
|
"finalize_document_storage",
|
|
"send_to_all_destinations",
|
|
]
|
|
|
|
|
|
def initialize_file_steps(db: Session, file_id: int, include_uploads: bool = False) -> None:
|
|
"""
|
|
Initialize processing steps for a file with 'pending' status.
|
|
|
|
Args:
|
|
db: Database session
|
|
file_id: ID of the file
|
|
include_uploads: Whether to include upload destination steps (set after destinations are known)
|
|
"""
|
|
# Initialize main processing steps
|
|
for step_name in MAIN_PROCESSING_STEPS:
|
|
step = FileProcessingStep(file_id=file_id, step_name=step_name, status="pending")
|
|
db.add(step)
|
|
|
|
db.commit()
|
|
|
|
|
|
def add_upload_steps(db: Session, file_id: int, upload_destinations: List[str]) -> None:
|
|
"""
|
|
Add upload destination steps for a file.
|
|
|
|
Args:
|
|
db: Database session
|
|
file_id: ID of the file
|
|
upload_destinations: List of upload destination names (e.g., ["dropbox", "s3", "nextcloud"])
|
|
"""
|
|
for destination in upload_destinations:
|
|
# Add both queue and upload steps for each destination
|
|
for prefix in ["queue_", "upload_to_"]:
|
|
step_name = f"{prefix}{destination}"
|
|
# Check if step already exists
|
|
existing = (
|
|
db.query(FileProcessingStep)
|
|
.filter(FileProcessingStep.file_id == file_id, FileProcessingStep.step_name == step_name)
|
|
.first()
|
|
)
|
|
if not existing:
|
|
step = FileProcessingStep(file_id=file_id, step_name=step_name, status="pending")
|
|
db.add(step)
|
|
|
|
db.commit()
|
|
|
|
|
|
def update_step_status(
|
|
db: Session,
|
|
file_id: int,
|
|
step_name: str,
|
|
status: str,
|
|
error_message: Optional[str] = None,
|
|
started_at: Optional[datetime] = None,
|
|
completed_at: Optional[datetime] = None,
|
|
) -> None:
|
|
"""
|
|
Update the status of a processing step.
|
|
|
|
Args:
|
|
db: Database session
|
|
file_id: ID of the file
|
|
step_name: Name of the processing step
|
|
status: New status ("pending", "in_progress", "success", "failure", "skipped")
|
|
error_message: Error message if status is "failure"
|
|
started_at: When the step started (for "in_progress")
|
|
completed_at: When the step completed (for "success" or "failure")
|
|
"""
|
|
step = (
|
|
db.query(FileProcessingStep)
|
|
.filter(FileProcessingStep.file_id == file_id, FileProcessingStep.step_name == step_name)
|
|
.first()
|
|
)
|
|
|
|
if not step:
|
|
# Create the step if it doesn't exist
|
|
step = FileProcessingStep(
|
|
file_id=file_id,
|
|
step_name=step_name,
|
|
status=status,
|
|
started_at=started_at,
|
|
completed_at=completed_at,
|
|
error_message=error_message,
|
|
)
|
|
db.add(step)
|
|
else:
|
|
# Update existing step
|
|
step.status = status
|
|
if error_message is not None:
|
|
step.error_message = error_message
|
|
if started_at is not None:
|
|
step.started_at = started_at
|
|
if completed_at is not None:
|
|
step.completed_at = completed_at
|
|
|
|
db.commit()
|
|
|
|
|
|
def get_file_step_status(db: Session, file_id: int) -> Dict[str, Dict]:
|
|
"""
|
|
Get the current status of all processing steps for a file.
|
|
|
|
Args:
|
|
db: Database session
|
|
file_id: ID of the file
|
|
|
|
Returns:
|
|
Dictionary mapping step_name to status info:
|
|
{
|
|
"step_name": {
|
|
"status": "success",
|
|
"started_at": datetime,
|
|
"completed_at": datetime,
|
|
"error_message": None
|
|
},
|
|
...
|
|
}
|
|
"""
|
|
steps = db.query(FileProcessingStep).filter(FileProcessingStep.file_id == file_id).all()
|
|
|
|
result = {}
|
|
for step in steps:
|
|
result[step.step_name] = {
|
|
"status": step.status,
|
|
"started_at": step.started_at,
|
|
"completed_at": step.completed_at,
|
|
"error_message": step.error_message,
|
|
}
|
|
|
|
return result
|
|
|
|
|
|
def get_file_overall_status(db: Session, file_id: int) -> Dict:
|
|
"""
|
|
Get the overall processing status for a file based on its steps.
|
|
|
|
Args:
|
|
db: Database session
|
|
file_id: ID of the file
|
|
|
|
Returns:
|
|
Dictionary with overall status info:
|
|
{
|
|
"status": "completed", # "pending", "processing", "completed", "failed"
|
|
"has_errors": False,
|
|
"total_steps": 10,
|
|
"completed_steps": 8,
|
|
"failed_steps": 0,
|
|
"in_progress_steps": 2
|
|
}
|
|
"""
|
|
steps = db.query(FileProcessingStep).filter(FileProcessingStep.file_id == file_id).all()
|
|
|
|
if not steps:
|
|
return {
|
|
"status": "pending",
|
|
"has_errors": False,
|
|
"total_steps": 0,
|
|
"completed_steps": 0,
|
|
"failed_steps": 0,
|
|
"in_progress_steps": 0,
|
|
}
|
|
|
|
total_steps = len(steps)
|
|
completed_steps = sum(1 for s in steps if s.status == "success")
|
|
failed_steps = sum(1 for s in steps if s.status == "failure")
|
|
in_progress_steps = sum(1 for s in steps if s.status == "in_progress")
|
|
skipped_steps = sum(1 for s in steps if s.status == "skipped")
|
|
|
|
has_errors = failed_steps > 0
|
|
|
|
# Determine overall status
|
|
if has_errors:
|
|
status = "failed"
|
|
elif in_progress_steps > 0:
|
|
status = "processing"
|
|
elif completed_steps + skipped_steps == total_steps:
|
|
status = "completed"
|
|
else:
|
|
status = "pending"
|
|
|
|
return {
|
|
"status": status,
|
|
"has_errors": has_errors,
|
|
"total_steps": total_steps,
|
|
"completed_steps": completed_steps,
|
|
"failed_steps": failed_steps,
|
|
"in_progress_steps": in_progress_steps,
|
|
"skipped_steps": skipped_steps,
|
|
}
|
|
|
|
|
|
def get_step_summary(db: Session, file_id: int) -> Dict:
|
|
"""
|
|
Get a summary of main steps vs upload steps with status counts.
|
|
|
|
Args:
|
|
db: Database session
|
|
file_id: ID of the file
|
|
|
|
Returns:
|
|
Dictionary with step counts:
|
|
{
|
|
"main": {"queued": 0, "in_progress": 1, "success": 7, "failure": 0, "skipped": 0},
|
|
"uploads": {"queued": 2, "in_progress": 0, "success": 4, "failure": 0, "skipped": 0},
|
|
"total_main_steps": 8,
|
|
"total_upload_tasks": 6
|
|
}
|
|
"""
|
|
steps = db.query(FileProcessingStep).filter(FileProcessingStep.file_id == file_id).all()
|
|
|
|
main_counts = {"queued": 0, "in_progress": 0, "success": 0, "failure": 0, "skipped": 0}
|
|
upload_counts = {"queued": 0, "in_progress": 0, "success": 0, "failure": 0, "skipped": 0}
|
|
|
|
upload_prefixes = ["upload_to_", "queue_"]
|
|
|
|
main_steps_count = 0
|
|
upload_steps_count = 0
|
|
|
|
for step in steps:
|
|
# Normalize status
|
|
status = step.status.lower()
|
|
if status == "pending":
|
|
status = "queued"
|
|
|
|
# Check if it's an upload task
|
|
is_upload = any(step.step_name.startswith(prefix) for prefix in upload_prefixes)
|
|
|
|
if is_upload:
|
|
if status in upload_counts:
|
|
upload_counts[status] += 1
|
|
upload_steps_count += 1
|
|
elif step.step_name in MAIN_PROCESSING_STEPS:
|
|
if status in main_counts:
|
|
main_counts[status] += 1
|
|
main_steps_count += 1
|
|
|
|
return {
|
|
"main": main_counts,
|
|
"uploads": upload_counts,
|
|
"total_main_steps": main_steps_count,
|
|
"total_upload_tasks": upload_steps_count,
|
|
}
|