96 lines
2.5 KiB
Python
96 lines
2.5 KiB
Python
import asyncio
|
|
import json
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
|
|
@dataclass
|
|
class TaskState:
|
|
"""单个提取任务的状态"""
|
|
|
|
task_id: str = ""
|
|
is_running: bool = False
|
|
cancel_event: asyncio.Event = field(default_factory=asyncio.Event)
|
|
progress: dict = field(default_factory=lambda: {
|
|
"percent": 0,
|
|
"status": "idle",
|
|
"current_file": "",
|
|
"action": "",
|
|
"files_total": 0,
|
|
"files_completed": 0,
|
|
"stories_found": 0,
|
|
"phase": "",
|
|
"phase_percent": 0,
|
|
"phase_detail": "",
|
|
})
|
|
_queue: asyncio.Queue = field(default_factory=asyncio.Queue)
|
|
|
|
def update(self, **kwargs):
|
|
"""更新进度，同时推送到 SSE 队列"""
|
|
self.progress.update(kwargs)
|
|
try:
|
|
self._queue.put_nowait(self.progress.copy())
|
|
except asyncio.QueueFull:
|
|
pass
|
|
|
|
async def get_update(self) -> dict:
|
|
"""SSE 消费端：等待新的进度更新"""
|
|
return await self._queue.get()
|
|
|
|
def reset(self):
|
|
"""重置状态"""
|
|
self.is_running = False
|
|
self.cancel_event = asyncio.Event()
|
|
self.progress = {
|
|
"percent": 0,
|
|
"status": "idle",
|
|
"current_file": "",
|
|
"action": "",
|
|
"files_total": 0,
|
|
"files_completed": 0,
|
|
"stories_found": 0,
|
|
"phase": "",
|
|
"phase_percent": 0,
|
|
"phase_detail": "",
|
|
}
|
|
self._queue = asyncio.Queue()
|
|
|
|
|
|
class TaskManager:
|
|
"""多任务管理器：支持多个用户同时提取"""
|
|
|
|
def __init__(self):
|
|
self._tasks: dict[str, TaskState] = {}
|
|
|
|
def create_task(self) -> TaskState:
|
|
"""创建新任务"""
|
|
task_id = uuid.uuid4().hex[:12]
|
|
state = TaskState(task_id=task_id)
|
|
self._tasks[task_id] = state
|
|
return state
|
|
|
|
def get_task(self, task_id: str) -> Optional[TaskState]:
|
|
"""获取任务状态"""
|
|
return self._tasks.get(task_id)
|
|
|
|
def remove_task(self, task_id: str):
|
|
"""清理已完成的任务"""
|
|
self._tasks.pop(task_id, None)
|
|
|
|
def cleanup_old_tasks(self):
|
|
"""清理所有非运行中的任务"""
|
|
to_remove = [
|
|
tid for tid, state in self._tasks.items()
|
|
if not state.is_running
|
|
]
|
|
for tid in to_remove:
|
|
self._tasks.pop(tid, None)
|
|
|
|
|
|
# 全局任务管理器
|
|
task_manager = TaskManager()
|
|
|
|
# 向后兼容：默认任务（单用户场景）
|
|
extraction_state = TaskState()
|