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

1#!/usr/bin/env python 

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

3 

4import uuid 

5import time 

6import asyncio 

7from typing import Dict, Any, Optional 

8from concurrent.futures import ThreadPoolExecutor 

9from datetime import datetime 

10 

11from loguru import logger 

12 

13from .taskmgr import TaskManager 

14from .task import Task 

15 

16class 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 

43class 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)