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

1#!/usr/bin/env python

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

3

4 import os

5 from pathlib import Path

6 from collections import deque, namedtuple

7 from dataclasses import dataclass

8 from typing import Optional, Any

9

10 from loguru import logger

11

12 from .task import Task

13 from .plugins import PluginManager

14 from .prompts import Prompts

15 from .diagnose import Diagnose

16 from .llm import ClientManager

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

18 from .role import RoleManager

19 from .mcp_tool import MCPToolManager

20

21@dataclass

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

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