完成请求推送,但组装mes还有问题
This commit is contained in:
161
backend/services/task_queue_manager.py
Normal file
161
backend/services/task_queue_manager.py
Normal file
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
任务队列管理器
|
||||
管理并行任务(生图、动态表格维护等)的状态和生命周期
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user