init

backend/app/api/extraction.py

@@ -0,0 +1,112 @@
+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}