162 lines
5.4 KiB
Python
162 lines
5.4 KiB
Python
"""
|
|
任务队列管理器
|
|
管理并行任务（生图、动态表格维护等）的状态和生命周期
|
|
"""
|
|
import asyncio
|
|
from typing import Dict, List, Optional
|
|
from enum import Enum
|
|
from datetime import datetime
|
|
|
|
|
|
class TaskStatus(Enum):
|
|
"""任务状态枚举"""
|
|
PENDING = "pending" # 等待中
|
|
RUNNING = "running" # 进行中
|
|
COMPLETED = "completed" # 已完成
|
|
FAILED = "failed" # 失败
|
|
CANCELLED = "cancelled" # 已取消
|
|
|
|
|
|
class TaskType(Enum):
|
|
"""任务类型枚举"""
|
|
IMAGE_WORKFLOW = "image_workflow"
|
|
DYNAMIC_TABLE = "dynamic_table"
|
|
|
|
|
|
class TaskItem:
|
|
"""任务项"""
|
|
|
|
def __init__(self, task_id: str, task_type: TaskType, chat_id: str):
|
|
self.task_id = task_id
|
|
self.task_type = task_type
|
|
self.chat_id = chat_id
|
|
self.status = TaskStatus.PENDING
|
|
self.created_at = datetime.now()
|
|
self.started_at = None
|
|
self.completed_at = None
|
|
self.error = None
|
|
self.metadata = {} # 用于存储提示词、修改内容等
|
|
|
|
def to_dict(self):
|
|
"""转换为字典格式（前端友好）"""
|
|
return {
|
|
"taskId": self.task_id,
|
|
"taskType": self.task_type.value,
|
|
"chatId": self.chat_id,
|
|
"status": self.status.value,
|
|
"createdAt": self.created_at.isoformat(),
|
|
"startedAt": self.started_at.isoformat() if self.started_at else None,
|
|
"completedAt": self.completed_at.isoformat() if self.completed_at else None,
|
|
"error": self.error,
|
|
"metadata": self.metadata
|
|
}
|
|
|
|
|
|
class TaskQueueManager:
|
|
"""
|
|
全局任务队列管理器
|
|
|
|
功能：
|
|
- 管理所有并行任务的生命周期
|
|
- 支持按聊天ID查询任务
|
|
- 支持取消任务
|
|
- 自动清理已完成的任务
|
|
"""
|
|
|
|
def __init__(self):
|
|
self.tasks: Dict[str, TaskItem] = {}
|
|
self.chat_tasks: Dict[str, List[str]] = {} # chat_id -> [task_ids]
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def add_task(self, task_id: str, task_type: TaskType, chat_id: str) -> TaskItem:
|
|
"""添加任务到队列"""
|
|
async with self._lock:
|
|
task = TaskItem(task_id, task_type, chat_id)
|
|
self.tasks[task_id] = task
|
|
|
|
if chat_id not in self.chat_tasks:
|
|
self.chat_tasks[chat_id] = []
|
|
self.chat_tasks[chat_id].append(task_id)
|
|
|
|
return task
|
|
|
|
async def start_task(self, task_id: str):
|
|
"""标记任务开始执行"""
|
|
async with self._lock:
|
|
if task_id in self.tasks:
|
|
self.tasks[task_id].status = TaskStatus.RUNNING
|
|
self.tasks[task_id].started_at = datetime.now()
|
|
|
|
async def complete_task(self, task_id: str, metadata: dict = None):
|
|
"""标记任务完成"""
|
|
async with self._lock:
|
|
if task_id in self.tasks:
|
|
self.tasks[task_id].status = TaskStatus.COMPLETED
|
|
self.tasks[task_id].completed_at = datetime.now()
|
|
if metadata:
|
|
self.tasks[task_id].metadata.update(metadata)
|
|
|
|
async def fail_task(self, task_id: str, error: str):
|
|
"""标记任务失败"""
|
|
async with self._lock:
|
|
if task_id in self.tasks:
|
|
self.tasks[task_id].status = TaskStatus.FAILED
|
|
self.tasks[task_id].completed_at = datetime.now()
|
|
self.tasks[task_id].error = error
|
|
|
|
async def cancel_task(self, task_id: str) -> bool:
|
|
"""
|
|
取消任务
|
|
|
|
Returns:
|
|
bool: 是否成功取消
|
|
"""
|
|
async with self._lock:
|
|
if task_id in self.tasks:
|
|
task = self.tasks[task_id]
|
|
if task.status in [TaskStatus.PENDING, TaskStatus.RUNNING]:
|
|
task.status = TaskStatus.CANCELLED
|
|
task.completed_at = datetime.now()
|
|
return True
|
|
return False
|
|
|
|
async def get_chat_tasks(self, chat_id: str, include_completed: bool = False) -> List[dict]:
|
|
"""
|
|
获取某个聊天的所有任务
|
|
|
|
Args:
|
|
chat_id: 聊天ID
|
|
include_completed: 是否包含已完成的任务
|
|
|
|
Returns:
|
|
List[dict]: 任务列表
|
|
"""
|
|
async with self._lock:
|
|
task_ids = self.chat_tasks.get(chat_id, [])
|
|
tasks = []
|
|
for task_id in task_ids:
|
|
if task_id in self.tasks:
|
|
task = self.tasks[task_id]
|
|
# 根据参数决定是否包含已完成的任务
|
|
if include_completed or task.status in [TaskStatus.PENDING, TaskStatus.RUNNING]:
|
|
tasks.append(task.to_dict())
|
|
return tasks
|
|
|
|
async def cleanup_completed_tasks(self, chat_id: str):
|
|
"""清理已完成的任务"""
|
|
async with self._lock:
|
|
if chat_id in self.chat_tasks:
|
|
task_ids = self.chat_tasks[chat_id]
|
|
completed_ids = [
|
|
tid for tid in task_ids
|
|
if tid in self.tasks and
|
|
self.tasks[tid].status in [TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED]
|
|
]
|
|
for tid in completed_ids:
|
|
del self.tasks[tid]
|
|
self.chat_tasks[chat_id].remove(tid)
|
|
|
|
|
|
# 全局实例
|
|
task_queue_manager = TaskQueueManager()
|