AgentCoreService/pqai_agent_server/task_server.py

import threading
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Dict

from pyarrow.dataset import dataset
from sqlalchemy import func

from pqai_agent import logging_service
from pqai_agent.data_models.agent_configuration import AgentConfiguration
from pqai_agent.data_models.agent_test_task import AgentTestTask
from pqai_agent.data_models.agent_test_task_conversations import AgentTestTaskConversations
from pqai_agent_server.const.status_enum import TestTaskConversationsStatus, TestTaskStatus, get_test_task_status_desc

logger = logging_service.logger


class TaskManager:
    """任务管理器"""

    def __init__(self, session_maker, dataset_server, max_workers: int = 10):
        self.session_maker = session_maker
        self.dataset_server = dataset_server
        self.task_events = {}  # 任务ID -> Event (用于取消任务)
        self.task_locks = {}  # 任务ID -> Lock (用于任务状态同步)
        self.running_tasks = set()
        self.executor = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix='TaskWorker')
        self.task_futures = {}  # 任务ID -> Future

    def get_test_task_list(self, page_num: int, page_size: int) -> Dict:
        with self.session_maker() as session:
            # 计算偏移量
            offset = (page_num - 1) * page_size
            # 查询分页数据
            result = (session.query(AgentTestTask, AgentConfiguration)
                      .outerjoin(AgentConfiguration, AgentTestTask.agent_id == AgentConfiguration.id)
                      .limit(page_size).offset(offset).all())
            # 查询总记录数
            total = session.query(func.count(AgentTestTask.id)).scalar()

            total_page = total // page_size + 1 if total % page_size > 0 else total // page_size
            total_page = 1 if total_page <= 0 else total_page
            response_data = [
                {
                    "id": agent_test_task.id,
                    "agentName": agent_configuration.name,
                    "createUser": agent_test_task.create_user,
                    "updateUser": agent_test_task.update_user,
                    "statusName": get_test_task_status_desc(agent_test_task.status),
                    "createTime": agent_test_task.create_time.strftime("%Y-%m-%d %H:%M:%S"),
                    "updateTime": agent_test_task.update_time.strftime("%Y-%m-%d %H:%M:%S")
                }
                for agent_test_task, agent_configuration in result
            ]
            return {
                "currentPage": page_num,
                "pageSize": page_size,
                "totalSize": total_page,
                "total": total,
                "list": response_data,
            }

    def get_test_task_conversations(self, task_id: int, page_num: int, page_size: int) -> Dict:
        with self.session_maker() as session:
            # 计算偏移量
            offset = (page_num - 1) * page_size
            # 查询分页数据
            result = (session.query(AgentTestTaskConversations, AgentConfiguration)
                      .outerjoin(AgentConfiguration, AgentTestTaskConversations.agent_id == AgentConfiguration.id)
                      .filter(AgentTestTaskConversations.task_id == task_id)
                      .limit(page_size).offset(offset).all())
            # 查询总记录数
            total = session.query(func.count(AgentTestTaskConversations.id)).scalar()

            total_page = total // page_size + 1 if total % page_size > 0 else total // page_size
            total_page = 1 if total_page <= 0 else total_page
            response_data = [
                {
                    "id": agent_test_task_conversation.id,
                    "agentName": agent_configuration.name,
                    "input": agent_test_task_conversation.input,
                    "output": agent_test_task_conversation.output,
                    "score": agent_test_task_conversation.score,
                    "statusName": get_test_task_status_desc(agent_test_task_conversation.status),
                    "createTime": agent_test_task_conversation.create_time.strftime("%Y-%m-%d %H:%M:%S"),
                    "updateTime": agent_test_task_conversation.update_time.strftime("%Y-%m-%d %H:%M:%S")
                }
                for agent_test_task_conversation, agent_configuration in result
            ]
            return {
                "currentPage": page_num,
                "pageSize": page_size,
                "totalSize": total_page,
                "total": total,
                "list": response_data,
            }

    def create_task(self, agent_id: int, module_id: int) -> Dict:
        """创建新任务"""
        with (self.session_maker() as session):
            with session.begin():
                agent_test_task = AgentTestTask(agent_id=agent_id, module_id=module_id)
                session.add(agent_test_task)
                session.flush()  # 强制SQL执行，但不提交事务
                task_id = agent_test_task.id
                agent_test_task_conversations = []
                datasets_list = self.dataset_server.get_dataset_list_by_module(module_id)
                for datasets in datasets_list:
                    conversation_datas = self.dataset_server.get_conversation_data_list_by_dataset(datasets.id)
                    for conversation_data in conversation_datas:
                        agent_test_task_conversation = AgentTestTaskConversations(task_id=task_id, agent_id=agent_id,
                                                                                  dataset_id=datasets.id,
                                                                                  conversation_id=conversation_data.id)
                        agent_test_task_conversations.append(agent_test_task_conversation)
                session.add_all(agent_test_task_conversations)
        # 异步执行任务
        self._execute_task(task_id)
        return self.get_task(task_id)

    def get_task(self, task_id: int):
        """获取任务信息"""
        with self.session_maker() as session:
            return session.query(AgentTestTask).filter(AgentTestTask.id == task_id).one()

    def get_task_conversations(self, task_id: int):
        """获取任务的所有子任务"""
        with self.session_maker() as session:
            return session.query(AgentTestTaskConversations).filter(AgentTestTaskConversations.task_id == task_id).all()

    def get_pending_task_conversations(self, task_id: int):
        """获取待处理的子任务"""
        with self.session_maker() as session:
            return session.query(AgentTestTaskConversations).filter(
                AgentTestTaskConversations.task_id == task_id).filter(
                AgentTestTaskConversations.status == TestTaskConversationsStatus.PENDING.value).all()

    def update_task_status(self, task_id: int, status: int):
        """更新任务状态"""
        with self.session_maker() as session:
            session.query(AgentTestTask).filter(AgentTestTask.id == task_id).update({"status": status})
            session.commit()

    def update_task_conversations_status(self, task_conversations_id: int, status: int):
        """更新子任务状态"""
        with self.session_maker() as session:
            session.query(AgentTestTaskConversations).filter(
                AgentTestTaskConversations.id == task_conversations_id).update({"status": status})
            session.commit()

    def update_task_conversations_res(self, task_conversations_id: int, status: int, score: str):
        """更新子任务结果"""
        with self.session_maker() as session:
            session.query(AgentTestTaskConversations).filter(
                AgentTestTaskConversations.id == task_conversations_id).update({"status": status, "score": score})
            session.commit()

    def cancel_task(self, task_id: int):
        """取消任务（带事务支持）"""
        # 设置取消事件
        if task_id in self.task_events:
            self.task_events[task_id].set()
        # 如果任务正在执行，尝试取消Future
        if task_id in self.task_futures:
            self.task_futures[task_id].cancel()

        with self.session_maker() as session:
            with session.begin():
                session.query(AgentTestTask).filter(AgentTestTask.id == task_id).update(
                    {"status": TestTaskStatus.CANCELLED.value})
                session.query(AgentTestTaskConversations).filter(AgentTestTaskConversations.task_id == task_id).filter(
                    AgentTestTaskConversations.status == TestTaskConversationsStatus.PENDING.value).update(
                    {"status": TestTaskConversationsStatus.CANCELLED.value})
                session.commit()

    def resume_task(self, task_id: int) -> bool:
        """恢复已取消的任务"""
        task = self.get_task(task_id)
        if not task or task.status != TestTaskStatus.CANCELLED.value:
            return False

        with self.session_maker() as session:
            with session.begin():
                session.query(AgentTestTask).filter(AgentTestTask.id == task_id).update(
                    {"status": TestTaskStatus.NOT_STARTED.value})
                session.query(AgentTestTaskConversations).filter(AgentTestTaskConversations.task_id == task_id).filter(
                    AgentTestTaskConversations.status == TestTaskConversationsStatus.CANCELLED.value).update(
                    {"status": TestTaskConversationsStatus.PENDING.value})
                session.commit()

        # 重新执行任务
        self._execute_task(task_id)
        logger.info(f"Resumed task {task_id}")
        return True

    def _execute_task(self, task_id: int):
        """提交任务到线程池执行"""
        # 确保任务状态一致性
        if task_id in self.running_tasks:
            return

        # 创建任务事件和锁
        if task_id not in self.task_events:
            self.task_events[task_id] = threading.Event()
        if task_id not in self.task_locks:
            self.task_locks[task_id] = threading.Lock()

        # 提交到线程池
        future = self.executor.submit(self._process_task, task_id)
        self.task_futures[task_id] = future

        # 标记任务为运行中
        with self.task_locks[task_id]:
            self.running_tasks.add(task_id)

    def _process_task(self, task_id: int):
        """处理任务的所有子任务"""
        try:
            # 更新任务状态为运行中
            self.update_task_status(task_id, TestTaskStatus.IN_PROGRESS.value)

            # 获取所有待处理的子任务
            task_conversations = self.get_pending_task_conversations(task_id)

            # 执行每个子任务
            for task_conversation in task_conversations:
                # 检查任务是否被取消
                if self.task_events[task_id].is_set():
                    break

                # 更新子任务状态为运行中
                self.update_task_conversations_status(task_conversation.id,
                                                      TestTaskConversationsStatus.RUNNING.value)
                try:
                    conversation_data = self.dataset_server.get_conversation_data_by_id(
                        task_conversation.conversation_id)
                    user_profile_data = self.dataset_server.get_user_profile_data(conversation_data.user_id)
                    staff_profile_data = self.dataset_server.get_staff_profile_data(conversation_data.staff_id)


                    # 模拟任务执行 - 在实际应用中替换为实际业务逻辑
                    # TODO 后续改成实际任务执行
                    time.sleep(1)
                    score = '{"score":0.05}'
                    # 更新子任务状态为已完成
                    self.update_task_conversations_res(task_conversation.id,
                                                       TestTaskConversationsStatus.SUCCESS.value, score)
                except Exception as e:
                    logger.error(f"Error executing task {task_id}: {str(e)}")
                    self.update_task_conversations_status(task_conversation.id,
                                                          TestTaskConversationsStatus.FAILED.value)

            # 检查任务是否完成
            task_conversations = self.get_task_conversations(task_id)
            all_completed = all(task_conversation.status in
                                (TestTaskConversationsStatus.SUCCESS.value, TestTaskConversationsStatus.FAILED.value)
                                for task_conversation in task_conversations)
            any_pending = any(task_conversation.status in
                              (TestTaskConversationsStatus.PENDING.value, TestTaskConversationsStatus.RUNNING.value)
                              for task_conversation in task_conversations)

            if all_completed:
                self.update_task_status(task_id, TestTaskStatus.COMPLETED.value)
                logger.info(f"Task {task_id} completed")
            elif not any_pending:
                # 没有待处理子任务但未全部完成（可能是取消了）
                current_status = self.get_task(task_id).status
                if current_status != TestTaskStatus.CANCELLED.value:
                    self.update_task_status(task_id, TestTaskStatus.COMPLETED.value
                    if all_completed else TestTaskStatus.CANCELLED.value)
        except Exception as e:
            logger.error(f"Error executing task {task_id}: {str(e)}")
            self.update_task_status(task_id, TestTaskStatus.COMPLETED.value)
        finally:
            # 清理资源
            with self.task_locks[task_id]:
                if task_id in self.running_tasks:
                    self.running_tasks.remove(task_id)
                if task_id in self.task_events:
                    del self.task_events[task_id]
                if task_id in self.task_futures:
                    del self.task_futures[task_id]

    def shutdown(self):
        """关闭执行器"""
        self.executor.shutdown(wait=False)
        logger.info("Task executor shutdown")