113 lines
3.7 KiB
Python
113 lines
3.7 KiB
Python
import asyncio
|
||
import json
|
||
from typing import Annotated, Optional
|
||
from fastapi import APIRouter, HTTPException, Depends, Query
|
||
from fastapi.responses import StreamingResponse
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
from sqlalchemy import select
|
||
|
||
from app.database import get_db
|
||
from app.models.transcript import Transcript
|
||
from app.state import task_manager
|
||
from app.services.extractor import run_extraction
|
||
|
||
router = APIRouter(prefix="/extraction", tags=["故事提取"])
|
||
|
||
|
||
@router.post("/start")
|
||
async def start_extraction(
|
||
task_id: Optional[str] = Query(None),
|
||
db: Annotated[AsyncSession, Depends(get_db)] = None,
|
||
):
|
||
"""启动提取任务，支持多用户并发"""
|
||
# 如果传了 task_id，复用已有任务
|
||
if task_id:
|
||
state = task_manager.get_task(task_id)
|
||
if state and state.is_running:
|
||
return {"task_id": task_id, "message": "任务正在进行中"}
|
||
else:
|
||
state = None
|
||
|
||
# 检查是否有待处理的文字稿
|
||
result = await db.execute(select(Transcript).where(Transcript.status == "pending"))
|
||
pending = result.scalars().all()
|
||
|
||
if not pending:
|
||
# 没有 pending 的，把所有非 pending 的文件重置为 pending（包括 completed/processing/error）
|
||
result2 = await db.execute(
|
||
select(Transcript).where(Transcript.status != "pending")
|
||
)
|
||
others = result2.scalars().all()
|
||
if others:
|
||
for t in others:
|
||
t.status = "pending"
|
||
t.error_message = ""
|
||
await db.commit()
|
||
pending = others
|
||
else:
|
||
raise HTTPException(400, "没有可提取的文字稿，请先上传文件")
|
||
|
||
total = len(pending)
|
||
|
||
# 创建新任务
|
||
state = task_manager.create_task()
|
||
state.reset()
|
||
state.is_running = True
|
||
|
||
asyncio.create_task(run_extraction(state))
|
||
|
||
return {"task_id": state.task_id, "total_files": total}
|
||
|
||
|
||
@router.get("/status")
|
||
async def get_status(task_id: Optional[str] = Query(None)):
|
||
"""获取提取进度"""
|
||
if task_id:
|
||
state = task_manager.get_task(task_id)
|
||
if state:
|
||
return {**state.progress, "task_id": state.task_id}
|
||
return {"status": "idle", "percent": 0}
|
||
|
||
|
||
@router.get("/status/stream")
|
||
async def stream_status(task_id: Optional[str] = Query(None)):
|
||
"""SSE 实时进度推送（支持多用户）"""
|
||
if not task_id:
|
||
raise HTTPException(400, "缺少 task_id 参数")
|
||
|
||
state = task_manager.get_task(task_id)
|
||
if not state:
|
||
raise HTTPException(404, "任务不存在")
|
||
|
||
async def event_generator():
|
||
# 立即推送当前状态
|
||
yield f"event: progress\ndata: {json.dumps({**state.progress, 'task_id': state.task_id}, ensure_ascii=False)}\n\n"
|
||
if state.progress.get("status") in ("completed", "error", "stopped"):
|
||
return
|
||
while True:
|
||
data = await state.get_update()
|
||
yield f"event: progress\ndata: {json.dumps({**data, 'task_id': state.task_id}, ensure_ascii=False)}\n\n"
|
||
if data.get("status") in ("completed", "error", "stopped"):
|
||
break
|
||
|
||
return StreamingResponse(
|
||
event_generator(),
|
||
media_type="text/event-stream",
|
||
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
|
||
)
|
||
|
||
|
||
@router.post("/stop")
|
||
async def stop_extraction(task_id: Optional[str] = Query(None)):
|
||
"""停止提取任务"""
|
||
if task_id:
|
||
state = task_manager.get_task(task_id)
|
||
else:
|
||
state = None
|
||
|
||
if not state or not state.is_running:
|
||
raise HTTPException(400, "没有正在进行的提取任务")
|
||
|
||
state.cancel_event.set()
|
||
return {"message": "正在停止...", "task_id": state.task_id}
|