Coverage for aipyapp/aipy/agent_taskmgr.py: 0%

1#!/usr/bin/env python

2# -*- coding: utf-8 -*-

3

4 import uuid

5 import time

6 import asyncio

7 from typing import Dict, Any, Optional

8 from concurrent.futures import ThreadPoolExecutor

9 from datetime import datetime

10

11 from loguru import logger

12

13 from .taskmgr import TaskManager

14 from .task import Task

15

16 class AgentTask:

17 """Agent任务封装"""

18

19 def __init__(self, task_id: str, instruction: str, task: Task, display: Any):

20 self.task_id = task_id

21 self.instruction = instruction

22 self.task = task

23 self.display = display

24 self.status = 'pending'

25 self.created_at = datetime.now()

26 self.started_at = None

27 self.completed_at = None

28 self.error = None

29

30 def to_dict(self) -> Dict[str, Any]:

31 """转换为字典格式"""

32 return {

33 'task_id': self.task_id,

34 'instruction': self.instruction,

35 'status': self.status,

36 'created_at': self.created_at.isoformat(),

37 'started_at': self.started_at.isoformat() if self.started_at else None,

38 'completed_at': self.completed_at.isoformat() if self.completed_at else None,

39 'error': self.error,

40 'captured_data': self.display.get_captured_data() if self.display else None

41 }

42

43 class AgentTaskManager(TaskManager):

44 """Agent模式任务管理器"""

45

46 def __init__(self, settings, /, display_manager=None):

47 # 强制使用agent显示模式和headless设置

48 super().__init__(settings, display_manager=display_manager)

49

50 # Agent特有属性

51 self.agent_tasks: Dict[str, AgentTask] = {}

52 self.executor = ThreadPoolExecutor(max_workers=4) # 支持并发

53 self.log = logger.bind(src='agent_taskmgr')

54

55 async def submit_task(self, instruction: str, metadata: Dict[str, Any] = None) -> str:

56 """提交新任务"""

57 task_id = str(uuid.uuid4())

58

59 try:

60 # 创建任务

61 task = super().new_task()

62

63 # 获取Task内部的display对象（这个已经注册到事件系统）

64 display = task.display

65

66 # 确保display是DisplayAgent类型

67 if display and hasattr(display, 'captured_data'):

68 # 添加元数据

69 if metadata:

70 display.captured_data['metadata'].update(metadata)

71 # 清空之前可能的数据

72 display.clear_captured_data()

73 if metadata:

74 display.captured_data['metadata'].update(metadata)

75

76 # 创建Agent任务封装

77 agent_task = AgentTask(task_id, instruction, task, display)

78 agent_task.status = 'pending'

79

80 self.agent_tasks[task_id] = agent_task

81 self.log.info(f"Task submitted: {task_id}")

82

83 return task_id

84

85 except Exception as e:

86 self.log.error(f"Failed to submit task: {e}")

87 raise

88

89 async def execute_task(self, task_id: str) -> Dict[str, Any]:

90 """执行任务"""

91 if task_id not in self.agent_tasks:

92 raise ValueError(f"Task {task_id} not found")

93

94 agent_task = self.agent_tasks[task_id]

95 if agent_task.status != 'pending':

96 raise ValueError(f"Task {task_id} is not in pending status")

97

98 agent_task.status = 'running'

99 agent_task.started_at = datetime.now()

100

101 try:

102 # 在线程池中执行任务

103 loop = asyncio.get_event_loop()

104 await loop.run_in_executor(

105 self.executor,

106 self._run_task_sync,

107 agent_task

108 )

109

110 agent_task.status = 'completed'

111 agent_task.completed_at = datetime.now()

112

113 except Exception as e:

114 agent_task.status = 'error'

115 agent_task.error = str(e)

116 agent_task.completed_at = datetime.now()

117 self.log.error(f"Task {task_id} failed: {e}")

118

119 return agent_task.to_dict()

120

121 def _run_task_sync(self, agent_task: AgentTask):

122 """同步执行任务（在线程池中运行）"""

123 try:

124 # 执行任务

125 agent_task.task.run(agent_task.instruction)

126

127 # 确保任务完成

128 agent_task.task.done()

129

130 except Exception as e:

131 # 捕获异常并记录到display中

132 if hasattr(agent_task.display, 'on_exception'):

133 from ..interface import Event

134 event = Event('exception', msg=str(e), exception=e)

135 agent_task.display.on_exception(event)

136 raise

137

138 async def get_task_status(self, task_id: str) -> Dict[str, Any]:

139 """获取任务状态"""

140 if task_id not in self.agent_tasks:

141 raise ValueError(f"Task {task_id} not found")

142

143 return self.agent_tasks[task_id].to_dict()

144

145 async def get_task_result(self, task_id: str) -> Dict[str, Any]:

146 """获取任务结果"""

147 if task_id not in self.agent_tasks:

148 raise ValueError(f"Task {task_id} not found")

149

150 agent_task = self.agent_tasks[task_id]

151 result = agent_task.to_dict()

152

153 # 添加详细的执行结果

154 if agent_task.status == 'completed':

155 captured_data = agent_task.display.get_captured_data()

156 result['output'] = {

157 'messages': captured_data['messages'],

158 'results': captured_data['results'],

159 'errors': captured_data['errors'],

160 'metadata': captured_data['metadata']

161 }

162

163 return result

164

165 async def list_tasks(self) -> Dict[str, Any]:

166 """列出所有任务"""

167 tasks = {}

168 for task_id, agent_task in self.agent_tasks.items():

169 tasks[task_id] = {

170 'task_id': task_id,

171 'instruction': agent_task.instruction,

172 'status': agent_task.status,

173 'created_at': agent_task.created_at.isoformat(),

174 'started_at': agent_task.started_at.isoformat() if agent_task.started_at else None,

175 'completed_at': agent_task.completed_at.isoformat() if agent_task.completed_at else None,

176 }

177 return tasks

178

179 async def cancel_task(self, task_id: str) -> bool:

180 """取消任务"""

181 if task_id not in self.agent_tasks:

182 return False

183

184 agent_task = self.agent_tasks[task_id]

185 if agent_task.status == 'running':

186 # 尝试停止任务

187 if hasattr(agent_task.task, 'stop'):

188 agent_task.task.stop()

189 agent_task.status = 'cancelled'

190 agent_task.completed_at = datetime.now()

191 return True

192

193 return False

194

195 def cleanup_completed_tasks(self, max_age_hours: int = 24):

196 """清理完成的任务"""

197 current_time = datetime.now()

198 to_remove = []

199

200 for task_id, agent_task in self.agent_tasks.items():

201 if agent_task.status in ['completed', 'error', 'cancelled']:

202 if agent_task.completed_at:

203 age = (current_time - agent_task.completed_at).total_seconds() / 3600

204 if age > max_age_hours:

205 to_remove.append(task_id)

206

207 for task_id in to_remove:

208 del self.agent_tasks[task_id]

209 self.log.info(f"Cleaned up task: {task_id}")

210

211 return len(to_remove)