58 lines
2.3 KiB
Python
58 lines
2.3 KiB
Python
from fastapi import APIRouter, HTTPException, status
|
|
|
|
from app.deps import AuthDep
|
|
from app.schemas.tasks import MediaOcrTaskRequest, OcrTaskRequest, TaskStatusResponse
|
|
from app.services.supabase_rest import supabase_rest
|
|
from app.services.task_tracker import task_tracker
|
|
from app.workers.tasks import media_ocr_pipeline, ocr_pipeline
|
|
from app.workers.utils import dispatch_task
|
|
|
|
router = APIRouter(prefix="/tasks")
|
|
|
|
|
|
@router.post("/ocr", response_model=TaskStatusResponse)
|
|
async def enqueue_ocr_task(payload: OcrTaskRequest, auth: AuthDep) -> TaskStatusResponse:
|
|
"""记录任务并投递 Celery，阶段 0 返回占位任务。"""
|
|
task = task_tracker.create_task(user_id=auth.user_id, document_id=payload.document_id, task_type="ocr")
|
|
dispatch_task(
|
|
ocr_pipeline,
|
|
task_id=task.task_id,
|
|
document_id=payload.document_id,
|
|
file_url=str(payload.file_url),
|
|
user_id=auth.user_id,
|
|
)
|
|
return task
|
|
|
|
|
|
@router.post("/media-ocr")
|
|
async def enqueue_media_ocr(payload: MediaOcrTaskRequest, auth: AuthDep) -> dict:
|
|
"""触发媒体资产 OCR，占位实现先写入模拟文本。"""
|
|
asset = supabase_rest.select_one("media_assets", {"id": payload.asset_id})
|
|
if not asset:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Media asset not found")
|
|
|
|
document_id = asset.get("document_id")
|
|
if not document_id:
|
|
raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Media asset missing document")
|
|
|
|
document = supabase_rest.select_one("documents", {"id": document_id})
|
|
if not document or document.get("user_id") != auth.user_id:
|
|
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Forbidden")
|
|
|
|
supabase_rest.update("media_assets", {"id": payload.asset_id}, {"ocr_status": "processing"})
|
|
dispatch_task(
|
|
media_ocr_pipeline,
|
|
asset_id=payload.asset_id,
|
|
user_id=auth.user_id,
|
|
document_id=document_id,
|
|
)
|
|
return {"asset_id": payload.asset_id, "status": "queued"}
|
|
|
|
|
|
@router.get("/{task_id}", response_model=TaskStatusResponse)
|
|
async def get_task_status(task_id: str, auth: AuthDep) -> TaskStatusResponse:
|
|
task = task_tracker.get_task(task_id=task_id, user_id=auth.user_id)
|
|
if not task:
|
|
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Task not found")
|
|
return task
|