Coverage for aipyapp/aipy/taskmgr.py: 37%

1#!/usr/bin/env python 

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

3 

4import os 

5from pathlib import Path 

6from collections import deque, namedtuple 

7from dataclasses import dataclass 

8from typing import Optional, Any 

9 

10from loguru import logger 

11 

12from .task import Task 

13from .plugins import PluginManager 

14from .prompts import Prompts 

15from .diagnose import Diagnose 

16from .llm import ClientManager 

17from .config import PLUGINS_DIR, ROLES_DIR, get_mcp_config_file, get_tt_api_key 

18from .role import RoleManager 

19from .mcp_tool import MCPToolManager 

20 

21@dataclass 

22class TaskContext: 

23 """任务上下文，包含创建任务所需的所有信息""" 

24 settings: Any 

25 cwd: Path 

26 plugin_manager: PluginManager 

27 display_manager: Any 

28 client_manager: ClientManager 

29 role_manager: RoleManager 

30 diagnose: Diagnose 

31 mcp: Optional[MCPToolManager] 

32 prompts: Prompts 

33 

34class TaskManager: 

35 MAX_TASKS = 16 

36 

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

38 # 核心配置 

39 self.settings = settings 

40 self.display_manager = display_manager 

41 self.log = logger.bind(src='taskmgr') 

42 

43 # 任务管理 

44 self.tasks = deque(maxlen=self.MAX_TASKS) 

45 

46 # 工作环境 

47 self._init_workenv() 

48 

49 # 初始化各种管理器 

50 self._init_managers() 

51 

52 # 创建任务上下文 

53 self.task_context = self._create_task_context() 

54 

55 def _init_workenv(self): 

56 """初始化工作环境""" 

57 # 环境变量 

58 envs = self.settings.get('environ', {}) 

59 for name, value in envs.items(): 

60 os.environ[name] = value 

61 

62 if self.settings.workdir: 

63 workdir = Path.cwd() / self.settings.workdir 

64 workdir.mkdir(parents=True, exist_ok=True) 

65 os.chdir(workdir) 

66 self.cwd = workdir 

67 else: 

68 self.cwd = Path.cwd() 

69 

70 def _init_managers(self): 

71 """初始化各种管理器""" 

72 # 插件管理器 

73 plugin_manager = PluginManager() 

74 plugin_manager.add_plugin_directory(PLUGINS_DIR) 

75 plugin_manager.load_all_plugins() 

76 self.plugin_manager = plugin_manager 

77 

78 if self.display_manager: 

79 for plugin in plugin_manager.get_display_plugins(): 

80 self.display_manager.register_plugin(plugin) 

81 

82 # 诊断器 

83 self.diagnose = Diagnose.create(self.settings) 

84 

85 # 客户端管理器 

86 self.client_manager = ClientManager(self.settings) 

87 

88 # 角色管理器 

89 api_conf = self.settings.get('api', {}) 

90 self.role_manager = RoleManager(ROLES_DIR, api_conf) 

91 self.role_manager.load_roles() 

92 self.role_manager.use(self.settings.get('role', 'aipy')) 

93 

94 # MCP 工具管理器 

95 mcp_config_file = get_mcp_config_file(self.settings.get('_config_dir')) 

96 self.mcp = MCPToolManager(mcp_config_file, get_tt_api_key(self.settings)) 

97 

98 # 提示管理器 

99 self.prompts = Prompts() 

100 

101 def get_status(self): 

102 status = { 

103 'tasks': len(self.tasks), 

104 'workdir': str(self.cwd), 

105 'role': self.role_manager.current_role.name, 

106 'client': repr(self.client_manager.current), 

107 'llm': self.client_manager.current.name, 

108 'display': self.display_manager.style, 

109 'mcp_enabled': self.mcp.is_mcp_enabled, 

110 } 

111 return status 

112 

113 def _create_task_context(self) -> TaskContext: 

114 """创建任务上下文""" 

115 return TaskContext( 

116 settings=self.settings, 

117 cwd=self.cwd, 

118 plugin_manager=self.plugin_manager, 

119 display_manager=self.display_manager, 

120 client_manager=self.client_manager, 

121 role_manager=self.role_manager, 

122 diagnose=self.diagnose, 

123 mcp=self.mcp, 

124 prompts=self.prompts 

125 ) 

126 

127 @property 

128 def workdir(self): 

129 return str(self.cwd) 

130 

131 def get_tasks(self): 

132 return list(self.tasks) 

133 

134 def list_llms(self): 

135 return self.client_manager.to_records() 

136 

137 def list_roles(self): 

138 RoleRecord = namedtuple('RoleRecord', ['Name', 'Description', 'Tips', 'Current']) 

139 rows = [] 

140 for name, role in self.role_manager.roles.items(): 

141 current = '*' if role == self.role_manager.current_role else '' 

142 rows.append(RoleRecord(name, role.short, len(role.tips), current)) 

143 return rows 

144 

145 def list_envs(self): 

146 EnvRecord = namedtuple('EnvRecord', ['Name', 'Description', 'Value']) 

147 rows = [] 

148 for name, (value, desc) in self.role_manager.current_role.envs.items(): 

149 rows.append(EnvRecord(name, desc, value[:32])) 

150 return rows 

151 

152 def list_tasks(self): 

153 rows = [] 

154 for task in self.tasks: 

155 rows.append(task.to_record()) 

156 return rows 

157 

158 def get_task_by_id(self, task_id): 

159 for task in self.tasks: 

160 if task.task_id == task_id: 

161 return task 

162 return None 

163 

164 def get_update(self, force=False): 

165 return self.diagnose.check_update(force) 

166 

167 def use(self, llm=None, role=None): 

168 rets = {} 

169 if llm: 

170 ret = self.client_manager.use(llm) 

171 rets['llm'] = ret 

172 if role: 

173 ret = self.role_manager.use(role) 

174 rets['role'] = ret 

175 return rets 

176 

177 def new_task(self): 

178 """创建新任务""" 

179 # 创建新任务 

180 task = Task(self.task_context) 

181 self.tasks.append(task) 

182 self.log.info('New task created', task_id=task.task_id) 

183 return task 

184 

185 def load_task(self, task_state): 

186 """从任务状态加载任务""" 

187 task = Task(self.task_context) 

188 task.restore_state(task_state) 

189 self.tasks.append(task) 

190 self.log.info('Task loaded', task_id=task.task_id) 

191 return task