AgentCoreService/pqai_agent_server/task_server.py

import json
import threading
import concurrent.futures
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from typing import Dict

from sqlalchemy import func

from pqai_agent import logging_service
from pqai_agent.agents.message_push_agent import MessagePushAgent
from pqai_agent.agents.multimodal_chat_agent import MultiModalChatAgent
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.data_models.service_module import ServiceModule
from pqai_agent.utils.prompt_utils import format_agent_profile
from pqai_agent_server.const.status_enum import TestTaskConversationsStatus, TestTaskStatus, get_test_task_status_desc
from concurrent.futures import ThreadPoolExecutor

from scripts.evaluate_agent import evaluate_agent

logger = logging_service.logger


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

    def __init__(self, session_maker, dataset_service):
        self.session_maker = session_maker
        self.dataset_service = dataset_service
        self.task_events = {}  # 任务ID -> Event (用于取消任务)
        self.task_locks = {}  # 任务ID -> Lock (用于任务状态同步)
        self.running_tasks = set()
        self.executor = ThreadPoolExecutor(max_workers=20, thread_name_prefix='TaskWorker')
        self.create_task_executor = ThreadPoolExecutor(max_workers=10, thread_name_prefix='CreateTaskWorker')
        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": MultiModalChatAgent.compose_dialogue(json.loads(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, evaluate_type: int) -> Dict:
        """创建新任务"""
        with self.session_maker() as session:
            agent_test_task = AgentTestTask(agent_id=agent_id, module_id=module_id, evaluate_type=evaluate_type,
                                            status=TestTaskStatus.CREATING.value)
            session.add(agent_test_task)
            session.commit()  # 显式提交
            task_id = agent_test_task.id
        # 异步执行创建任务
        self.create_task_executor.submit(self._generate_agent_test_task_conversation_batch, task_id, agent_id,
                                         module_id)
        return self.get_task(task_id)

    def _generate_agent_test_task_conversation_batch(self, task_id: int, agent_id: int, module_id: int):
        """异步生成子任务"""
        try:
            # 获取数据集列表
            dataset_module_list = self.dataset_service.get_dataset_module_list_by_module(module_id)

            # 批量处理数据集 - 减少数据库交互
            batch_size = 100  # 每批处理100个子任务
            agent_test_task_conversation_batch = []

            for dataset_module in dataset_module_list:
                # 获取对话数据列表
                conversation_datas = self.dataset_service.get_conversation_data_list_by_dataset(
                    dataset_module.dataset_id)

                for conversation_data in conversation_datas:
                    # 创建子任务对象
                    agent_test_task_conversation = AgentTestTaskConversations(
                        task_id=task_id,
                        agent_id=agent_id,
                        dataset_id=dataset_module.dataset_id,
                        conversation_id=conversation_data.id,
                        status=TestTaskConversationsStatus.PENDING.value
                    )
                    agent_test_task_conversation_batch.append(agent_test_task_conversation)

                    # 批量提交
                    if len(agent_test_task_conversation_batch) >= batch_size:
                        self.save_agent_test_task_conversation_batch(agent_test_task_conversation_batch)
                        agent_test_task_conversation_batch = []

            # 提交剩余的子任务
            if agent_test_task_conversation_batch:
                self.save_agent_test_task_conversation_batch(agent_test_task_conversation_batch)

            # 更新主任务状态为未开始
            self.update_task_status(task_id, TestTaskStatus.NOT_STARTED.value)

            # 自动提交任务执行
            self._execute_task(task_id)

        except Exception as e:
            logger.error(f"生成子任务失败: {str(e)}")
            # 更新任务状态为失败
            self.update_task_status(task_id, TestTaskStatus.CREATED_FAIL.value)

    def save_agent_test_task_conversation_batch(self, agent_test_task_conversation_batch: list):
        """批量保存子任务到数据库"""
        try:
            with self.session_maker() as session:
                with session.begin():
                    session.add_all(agent_test_task_conversation_batch)
        except Exception as e:
            logger.error(e)

    def get_agent_configuration_by_task_id(self, task_id: int):
        """获取指定任务ID对应的Agent配置信息"""
        with self.session_maker() as session:
            return session.query(AgentConfiguration) \
                .join(AgentTestTask, AgentTestTask.agent_id == AgentConfiguration.id) \
                .filter(AgentTestTask.id == task_id) \
                .one_or_none()  # 返回单个对象或None（如果未找到）

    def get_service_module_by_task_id(self, task_id: int):
        """获取指定任务ID对应的Agent配置信息"""
        with self.session_maker() as session:
            return session.query(ServiceModule) \
                .join(AgentTestTask, AgentTestTask.module_id == ServiceModule.id) \
                .filter(AgentTestTask.id == task_id) \
                .one_or_none()  # 返回单个对象或None（如果未找到）

    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_in_progress_task(self):
        """获取执行中任务"""
        with self.session_maker() as session:
            return session.query(AgentTestTask).filter(AgentTestTask.status == TestTaskStatus.IN_PROGRESS.value).all()

    def get_creating_task(self):
        """获取执行中任务"""
        with self.session_maker() as session:
            return session.query(AgentTestTask).filter(AgentTestTask.status == TestTaskStatus.CREATING.value).all()

    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 del_task_conversations(self, task_id: int):
        with self.session_maker() as session:
            session.query(AgentTestTaskConversations).filter(AgentTestTaskConversations.task_id == task_id).delete()
            # 提交事务生效
            session.commit()

    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.in_([
                    TestTaskConversationsStatus.PENDING.value,
                    TestTaskConversationsStatus.RUNNING.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, "update_time": datetime.now()})
            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, "update_time": datetime.now()})
            session.commit()

    def update_task_conversations_res(self, task_conversations_id: int, status: int, input: str, output: str,
                                      score: str):
        """更新子任务结果"""
        with self.session_maker() as session:
            session.query(AgentTestTaskConversations).filter(
                AgentTestTaskConversations.id == task_conversations_id).update(
                {"status": status, "input": input, "output": output, "score": score, "update_time": datetime.now()})
            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 recover_tasks(self):
        """服务启动时恢复未完成的任务"""

        creating = self.get_creating_task()
        for task in creating:
            task_id = task.id
            agent_id = task.agent_id
            module_id = task.module_id
            self.del_task_conversations(task_id)
            # 重新提交任务
            # 异步执行创建任务
            self.create_task_executor.submit(self._generate_agent_test_task_conversation_batch, task_id, agent_id,
                                             module_id)

        # 获取所有进行中的任务ID（根据实际状态定义查询）
        in_progress_tasks = self.get_in_progress_task()

        for task in in_progress_tasks:
            task_id = task.id
            # 重新提交任务
            self._execute_task(task_id)

    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)

            if not task_conversations:
                self.update_task_status(task_id, TestTaskStatus.COMPLETED.value)
                return

            agent_configuration = self.get_agent_configuration_by_task_id(task_id)
            query_prompt_template = agent_configuration.task_prompt

            task = self.get_task(task_id)

            # 使用线程池执行子任务
            with ThreadPoolExecutor(max_workers=8) as executor:  # 可根据需要调整并发数
                futures = {}
                for task_conversation in task_conversations:
                    if self.task_events[task_id].is_set():
                        break  # 检查任务取消事件

                    # 提交子任务到线程池
                    future = executor.submit(
                        self._process_single_conversation,
                        task_id,
                        task,
                        task_conversation,
                        query_prompt_template,
                        agent_configuration
                    )
                    futures[future] = task_conversation.id

                # 等待所有子任务完成或取消
                for future in concurrent.futures.as_completed(futures):
                    conv_id = futures[future]
                    try:
                        future.result()  # 获取结果（如有异常会在此抛出）
                    except Exception as e:
                        logger.error(f"Subtask {conv_id} failed: {str(e)}")
                        self.update_task_conversations_status(
                            conv_id,
                            TestTaskConversationsStatus.FAILED.value
                        )

            # 检查最终任务状态
            self._update_final_task_status(task_id)

        except Exception as e:
            logger.error(f"Error processing task {task_id}: {str(e)}")
            self.update_task_status(task_id, TestTaskStatus.FAILED.value)
        finally:
            self._cleanup_task_resources(task_id)

    def _process_single_conversation(self, task_id, task, task_conversation, query_prompt_template,
                                     agent_configuration):
        """处理单个对话子任务（线程安全）"""
        # 检查任务是否被取消
        if self.task_events[task_id].is_set():
            return

        # 更新子任务状态
        if task_conversation.status == TestTaskConversationsStatus.PENDING.value:
            self.update_task_conversations_status(
                task_conversation.id,
                TestTaskConversationsStatus.RUNNING.value
            )

        try:
            # 创建独立的agent实例（确保线程安全）
            agent = MultiModalChatAgent(
                model=agent_configuration.execution_model,
                system_prompt=agent_configuration.system_prompt,
                tools=json.loads(agent_configuration.tools)
            )

            # 获取对话数据
            conversation_data = self.dataset_service.get_conversation_data_by_id(
                task_conversation.conversation_id)
            user_profile_data = self.dataset_service.get_user_profile_data(
                conversation_data.user_id,
                conversation_data.version_date.replace("-", ""))
            user_profile = json.loads(user_profile_data['profile_data_v1'])
            avatar = user_profile_data['iconurl']
            staff_profile = self.dataset_service.get_staff_profile_data(
                conversation_data.staff_id).agent_profile
            conversations = self.dataset_service.get_chat_conversation_list_by_ids(
                json.loads(conversation_data.conversation),
                conversation_data.staff_id
            )
            conversations = sorted(conversations, key=lambda i: i['timestamp'], reverse=False)

            # 生成推送消息
            last_timestamp = int(conversations[-1]["timestamp"])
            match task.evaluate_type:
                case 0:
                    send_timestamp = int(last_timestamp / 1000) + 10
                case 1:
                    send_timestamp = int(last_timestamp / 1000) + 24 * 3600
                case _:
                    raise ValueError("evaluate_type must be 0 or 1")
            send_time = datetime.fromtimestamp(send_timestamp).strftime('%Y-%m-%d %H:%M:%S')
            message = agent._generate_message(
                context={
                    "formatted_staff_profile": staff_profile,
                    "nickname": user_profile['nickname'],
                    "name": user_profile['name'],
                    "avatar": avatar,
                    "preferred_nickname": user_profile['preferred_nickname'],
                    "gender": user_profile['gender'],
                    "age": user_profile['age'],
                    "region": user_profile['region'],
                    "health_conditions": user_profile['health_conditions'],
                    "medications": user_profile['medications'],
                    "interests": user_profile['interests'],
                    "current_datetime": send_time
                },
                dialogue_history=conversations,
                query_prompt_template=query_prompt_template
            )

            if not message:
                self.update_task_conversations_status(
                    task_conversation.id,
                    TestTaskConversationsStatus.MESSAGE_FAILED.value
                )
                return

            param = {}
            param["dialogue_history"] = conversations
            param["message"] = message
            param["send_time"] = send_time
            param["agent_profile"] = json.loads(staff_profile)
            param["user_profile"] = user_profile
            score = evaluate_agent(param, task.evaluate_type)

            if not score:
                self.update_task_conversations_status(
                    task_conversation.id,
                    TestTaskConversationsStatus.SCORE_FAILED.value
                )
                return

            # 更新子任务结果
            self.update_task_conversations_res(
                task_conversation.id,
                TestTaskConversationsStatus.SUCCESS.value,
                json.dumps(conversations, ensure_ascii=False),
                json.dumps(message, ensure_ascii=False),
                json.dumps(score)
            )

        except Exception as e:
            logger.error(f"Subtask {task_conversation.id} failed: {str(e)}")
            self.update_task_conversations_status(
                task_conversation.id,
                TestTaskConversationsStatus.FAILED.value
            )
            raise  # 重新抛出异常以便主线程捕获

    def _update_final_task_status(self, task_id):
        """更新任务的最终状态"""
        task_conversations = self.get_task_conversations(task_id)
        all_completed = all(
            conv.status in (TestTaskConversationsStatus.SUCCESS.value,
                            TestTaskConversationsStatus.FAILED.value)
            for conv 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(
                conv.status in (TestTaskConversationsStatus.PENDING.value,
                                TestTaskConversationsStatus.RUNNING.value)
                for conv in task_conversations
        ):
            current_status = self.get_task(task_id).status
            if current_status != TestTaskStatus.CANCELLED.value:
                new_status = TestTaskStatus.COMPLETED.value if all_completed else TestTaskStatus.CANCELLED.value
                self.update_task_status(task_id, new_status)

    def _cleanup_task_resources(self, task_id):
        """清理任务资源（线程安全）"""
        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")