diff --git a/backend/app/routers/tasks.py b/backend/app/routers/tasks.py index fc304ab..3a50347 100644 --- a/backend/app/routers/tasks.py +++ b/backend/app/routers/tasks.py @@ -9,9 +9,11 @@ from pathlib import Path import shutil import hashlib -from fastapi import APIRouter, Depends, HTTPException, status, Query, UploadFile, File +from fastapi import APIRouter, Depends, HTTPException, status, Query, UploadFile, File, BackgroundTasks from fastapi.responses import FileResponse from sqlalchemy.orm import Session +import json +from datetime import datetime from app.core.deps import get_db, get_current_user from app.core.config import settings @@ -29,12 +31,96 @@ from app.schemas.task import ( ) from app.services.task_service import task_service from app.services.file_access_service import file_access_service +from app.services.ocr_service import OCRService logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/v2/tasks", tags=["Tasks"]) +def process_task_ocr(task_id: str, task_db_id: int, file_path: str, filename: str): + """ + Background task to process OCR for a task + + Args: + task_id: Task UUID string + task_db_id: Task database ID + file_path: Path to uploaded file + filename: Original filename + """ + from app.core.database import SessionLocal + from app.models.task import Task + + db = SessionLocal() + start_time = datetime.now() + + try: + logger.info(f"Starting OCR processing for task {task_id}, file: {filename}") + + # Get task directly by database ID (bypass user isolation for background task) + task = db.query(Task).filter(Task.id == task_db_id).first() + if not task: + logger.error(f"Task {task_id} not found in database") + return + + # Initialize OCR service + ocr_service = OCRService() + + # Process the file with OCR + ocr_result = ocr_service.process_image( + image_path=Path(file_path), + lang='ch', + detect_layout=True + ) + + # Calculate processing time + processing_time_ms = int((datetime.now() - start_time).total_seconds() * 1000) + + # Create result directory + result_dir = Path(settings.result_dir) / task_id + result_dir.mkdir(parents=True, exist_ok=True) + + # Save JSON result + json_path = result_dir / f"{Path(filename).stem}_result.json" + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(ocr_result, f, ensure_ascii=False, indent=2) + + # Save Markdown result + markdown_path = result_dir / f"{Path(filename).stem}_result.md" + markdown_content = ocr_result.get('markdown', '') + with open(markdown_path, 'w', encoding='utf-8') as f: + f.write(markdown_content) + + # Update task with results (direct database update) + task.result_json_path = str(json_path) + task.result_markdown_path = str(markdown_path) + task.processing_time_ms = processing_time_ms + task.status = TaskStatus.COMPLETED + task.completed_at = datetime.utcnow() + task.updated_at = datetime.utcnow() + + db.commit() + + logger.info(f"OCR processing completed for task {task_id} in {processing_time_ms}ms") + + except Exception as e: + logger.exception(f"OCR processing failed for task {task_id}") + + # Update task status to failed (direct database update) + try: + task = db.query(Task).filter(Task.id == task_db_id).first() + if task: + task.status = TaskStatus.FAILED + task.error_message = str(e) + task.updated_at = datetime.utcnow() + db.commit() + except Exception as update_error: + logger.error(f"Failed to update task status: {update_error}") + + finally: + db.close() + + @router.post("/", response_model=TaskResponse, status_code=status.HTTP_201_CREATED) async def create_task( task_data: TaskCreate, @@ -425,6 +511,7 @@ async def download_pdf( @router.post("/{task_id}/start", response_model=TaskResponse, summary="Start task processing") async def start_task( task_id: str, + background_tasks: BackgroundTasks, db: Session = Depends(get_db), current_user: User = Depends(get_current_user) ): @@ -434,11 +521,11 @@ async def start_task( - **task_id**: Task UUID """ try: - task = task_service.update_task_status( + # Get task details + task = task_service.get_task_by_id( db=db, task_id=task_id, - user_id=current_user.id, - status=TaskStatus.PROCESSING + user_id=current_user.id ) if not task: @@ -447,7 +534,39 @@ async def start_task( detail="Task not found" ) - logger.info(f"Started task {task_id} for user {current_user.email}") + # Check if task is in pending status + if task.status != TaskStatus.PENDING: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Cannot start task in '{task.status.value}' status" + ) + + # Get task file + task_file = db.query(TaskFile).filter(TaskFile.task_id == task.id).first() + if not task_file: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail="Task file not found" + ) + + # Update task status to processing + task = task_service.update_task_status( + db=db, + task_id=task_id, + user_id=current_user.id, + status=TaskStatus.PROCESSING + ) + + # Start OCR processing in background + background_tasks.add_task( + process_task_ocr, + task_id=task_id, + task_db_id=task.id, + file_path=task_file.stored_path, + filename=task_file.original_name + ) + + logger.info(f"Started OCR processing task {task_id} for user {current_user.email}") return task except HTTPException: