|
|
@@ -1,5 +1,6 @@
|
|
|
import json
|
|
|
import threading
|
|
|
+import concurrent.futures
|
|
|
import time
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
from datetime import datetime
|
|
|
@@ -15,7 +16,7 @@ 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_server.const.status_enum import TestTaskConversationsStatus, TestTaskStatus, get_test_task_status_desc
|
|
|
-
|
|
|
+from concurrent.futures import ThreadPoolExecutor
|
|
|
logger = logging_service.logger
|
|
|
|
|
|
|
|
|
@@ -324,107 +325,270 @@ class TaskManager:
|
|
|
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
|
|
|
- agent = MultiModalChatAgent(model=agent_configuration.execution_model,
|
|
|
- system_prompt=agent_configuration.system_prompt,
|
|
|
- tools=json.loads(agent_configuration.tools))
|
|
|
- # 执行每个子任务
|
|
|
- for task_conversation in task_conversations:
|
|
|
- # 检查任务是否被取消
|
|
|
- if self.task_events[task_id].is_set():
|
|
|
- break
|
|
|
-
|
|
|
- # 更新子任务状态为运行中
|
|
|
- if task_conversation.status == TestTaskConversationsStatus.PENDING.value:
|
|
|
- self.update_task_conversations_status(task_conversation.id,
|
|
|
- TestTaskConversationsStatus.RUNNING.value)
|
|
|
- try:
|
|
|
- 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_data = 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"])
|
|
|
- push_time = int(last_timestamp / 1000) + 24 * 3600
|
|
|
- push_dt = datetime.fromtimestamp(push_time).strftime('%Y-%m-%d %H:%M:%S')
|
|
|
- push_message = agent._generate_message(
|
|
|
- context={
|
|
|
- "formatted_staff_profile": staff_profile_data,
|
|
|
- "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": push_dt
|
|
|
- },
|
|
|
- dialogue_history=conversations,
|
|
|
- query_prompt_template=query_prompt_template
|
|
|
+
|
|
|
+ # 使用线程池执行子任务
|
|
|
+ with ThreadPoolExecutor(max_workers=20) 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_conversation,
|
|
|
+ query_prompt_template,
|
|
|
+ agent_configuration
|
|
|
)
|
|
|
- # TODO 获取打分
|
|
|
- score = '{"score":0.05}'
|
|
|
- # 更新子任务状态为已完成
|
|
|
- self.update_task_conversations_res(task_conversation.id,
|
|
|
- TestTaskConversationsStatus.SUCCESS.value,
|
|
|
- json.dumps(conversations, ensure_ascii=False),
|
|
|
- json.dumps(push_message, ensure_ascii=False),
|
|
|
- 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)
|
|
|
+ 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 executing task {task_id}: {str(e)}")
|
|
|
+ logger.error(f"Error processing task {task_id}: {str(e)}")
|
|
|
self.update_task_status(task_id, TestTaskStatus.FAILED.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]
|
|
|
+ self._cleanup_task_resources(task_id)
|
|
|
+
|
|
|
+ def _process_single_conversation(self, task_id, 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_data = 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"])
|
|
|
+ push_time = int(last_timestamp / 1000) + 24 * 3600
|
|
|
+ push_dt = datetime.fromtimestamp(push_time).strftime('%Y-%m-%d %H:%M:%S')
|
|
|
+ push_message = agent._generate_message(
|
|
|
+ context={
|
|
|
+ "formatted_staff_profile": staff_profile_data,
|
|
|
+ "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": push_dt
|
|
|
+ },
|
|
|
+ dialogue_history=conversations,
|
|
|
+ query_prompt_template=query_prompt_template
|
|
|
+ )
|
|
|
+
|
|
|
+ # 获取打分（TODO: 实际实现）
|
|
|
+ score = '{"score":0.05}'
|
|
|
+
|
|
|
+ # 更新子任务结果
|
|
|
+ self.update_task_conversations_res(
|
|
|
+ task_conversation.id,
|
|
|
+ TestTaskConversationsStatus.SUCCESS.value,
|
|
|
+ json.dumps(conversations, ensure_ascii=False),
|
|
|
+ json.dumps(push_message, ensure_ascii=False),
|
|
|
+ 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 _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)
|
|
|
+ #
|
|
|
+ # agent_configuration = self.get_agent_configuration_by_task_id(task_id)
|
|
|
+ # query_prompt_template = agent_configuration.task_prompt
|
|
|
+ # agent = MultiModalChatAgent(model=agent_configuration.execution_model,
|
|
|
+ # system_prompt=agent_configuration.system_prompt,
|
|
|
+ # tools=json.loads(agent_configuration.tools))
|
|
|
+ # # 执行每个子任务
|
|
|
+ # for task_conversation in task_conversations:
|
|
|
+ # # 检查任务是否被取消
|
|
|
+ # if self.task_events[task_id].is_set():
|
|
|
+ # break
|
|
|
+ #
|
|
|
+ # # 更新子任务状态为运行中
|
|
|
+ # if task_conversation.status == TestTaskConversationsStatus.PENDING.value:
|
|
|
+ # self.update_task_conversations_status(task_conversation.id,
|
|
|
+ # TestTaskConversationsStatus.RUNNING.value)
|
|
|
+ # try:
|
|
|
+ # 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_data = 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"])
|
|
|
+ # push_time = int(last_timestamp / 1000) + 24 * 3600
|
|
|
+ # push_dt = datetime.fromtimestamp(push_time).strftime('%Y-%m-%d %H:%M:%S')
|
|
|
+ # push_message = agent._generate_message(
|
|
|
+ # context={
|
|
|
+ # "formatted_staff_profile": staff_profile_data,
|
|
|
+ # "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": push_dt
|
|
|
+ # },
|
|
|
+ # dialogue_history=conversations,
|
|
|
+ # query_prompt_template=query_prompt_template
|
|
|
+ # )
|
|
|
+ # # TODO 获取打分
|
|
|
+ # score = '{"score":0.05}'
|
|
|
+ # # 更新子任务状态为已完成
|
|
|
+ # self.update_task_conversations_res(task_conversation.id,
|
|
|
+ # TestTaskConversationsStatus.SUCCESS.value,
|
|
|
+ # json.dumps(conversations, ensure_ascii=False),
|
|
|
+ # json.dumps(push_message, ensure_ascii=False),
|
|
|
+ # 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.FAILED.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):
|
|
|
"""关闭执行器"""
|
xueyiming