Coverage for aipyapp/aipy/role.py: 26%

1#! /usr/bin/env python3

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

3

4 from typing import List, Dict, Any, Optional

5 from dataclasses import dataclass, field

6 import tomllib

7 import os

8

9 from loguru import logger

10

11@dataclass

12 class Tip:

13 """提示信息对象"""

14 name: str

15 short: str

16 detail: str

17

18 @classmethod

19 def from_dict(cls, name: str, data: Dict[str, str]) -> 'Tip':

20 """从字典创建提示信息对象"""

21 return cls(

22 name=name,

23 short=data.get('short', ''),

24 detail=data.get('detail', '')

25 )

26

27 def __str__(self):

28 return f"<tip name=\"{self.name}\">\n{self.detail.strip()}\n</tip>"

29

30 class Role:

31 """提示信息管理器"""

32 def __init__(self):

33 self.name: str = ''

34 self.short: str = ''

35 self.detail: str = ''

36 self.envs: Dict[str, tuple[str, str]] = {}

37 self.packages: Dict[str, set[str]] = {}

38 self.tips: Dict[str, Tip] = {}

39 self.plugins: Dict[str, Dict[str, Any]] = {}

40

41 def get_tip(self, name: str) -> Optional[Tip]:

42 """获取指定名称的提示信息"""

43 return self.tips.get(name, None)

44

45 def add_env(self, name: str, value: str, desc: str):

46 self.envs[name] = (value, desc)

47

48 def add_package(self, name: str, packages: List[str]):

49 self.packages[name] = set(packages)

50

51 def add_tip(self, name: str, short: str, detail: str):

52 self.tips[name] = Tip(name, short, detail)

53

54 def add_plugin(self, name: str, data: Dict[str, Any]):

55 self.plugins[name] = data

56

57 def __iter__(self):

58 return iter(self.tips.items())

59

60 def __len__(self):

61 return len(self.tips)

62

63 def __getitem__(self, name: str) -> Tip:

64 return self.tips[name]

65

66 @classmethod

67 def from_dict(cls, data: Dict[str, Any]) -> 'Role':

68 """从字典创建角色对象"""

69 role = cls()

70

71 # 设置角色基本信息

72 role.name = data.get('name', '')

73 role.short = data.get('short', '')

74 role.detail = data.get('detail', '')

75

76 # 加载环境变量

77 env_data = data.get('envs', {})

78 for env_name, env_info in env_data.items():

79 if isinstance(env_info, list):

80 value = env_info[0]

81 desc = env_info[1]

82 else:

83 value = str(env_info)

84 desc = ''

85 role.add_env(env_name, value, desc)

86

87 # 加载包信息

88 packages_data = data.get('packages', {})

89 for lang, packages in packages_data.items():

90 role.add_package(lang, packages)

91

92 # 加载提示信息

93 tips_data = data.get('tips', {})

94 for tip_name, tip_data in tips_data.items():

95 short = tip_data.get('short', '')

96 detail = tip_data.get('detail', '')

97 role.add_tip(tip_name, short, detail)

98

99 # 加载插件信息

100 plugins_data = data.get('plugins', {})

101 for plugin_name, plugin_data in plugins_data.items():

102 role.add_plugin(plugin_name, plugin_data)

103

104 return role

105

106 @classmethod

107 def load(cls, toml_path: str) -> 'Role':

108 """从 TOML 文件加载角色信息

109

110 Args:

111 toml_path: TOML 文件路径

112

113 Returns:

114 Role: 角色对象

115 """

116 with open(toml_path, 'rb') as f:

117 data = tomllib.load(f)

118

119 return cls.from_dict(data)

120

121 class RoleManager:

122 def __init__(self, roles_dir: str = None, api_conf: Dict[str, Dict[str, Any]] = None):

123 self.roles_dir = roles_dir

124 self.roles: Dict[str, Role] = {}

125 self.default_role: Role = None

126 self.current_role: Role = None

127 self.log = logger.bind(src='roles')

128 self.api_conf = api_conf

129

130 def _add_api(self, role: Role):

131 for api_name, api_conf in self.api_conf.items():

132 desc = api_conf.get('desc')

133 if not desc:

134 self.log.warning(f"API {api_name} has no description")

135 continue

136 role.add_tip(api_name, '', desc)

137 envs = api_conf.get('env')

138 if not envs:

139 continue

140 for name, (value, desc) in envs.items():

141 role.add_env(name, value, desc)

142

143 def load_roles(self):

144 sys_roles_dir = os.path.join(os.path.dirname(__file__), '..', 'res', 'roles')

145 for roles_dir in [sys_roles_dir, self.roles_dir]:

146 if not roles_dir or not os.path.exists(roles_dir):

147 continue

148 for fname in os.listdir(roles_dir):

149 if fname.endswith(".toml") and not fname.startswith("_"):

150 role = Role.load(os.path.join(roles_dir, fname))

151 self.log.info(f"Loaded role: {role.name}/{len(role)}")

152 self.roles[role.name.lower()] = role

153 self._add_api(role)

154

155 if self.roles:

156 self.default_role = list(self.roles.values())[0]

157 self.current_role = self.default_role

158

159 def use(self, name: str):

160 name = name.lower()

161 if name in self.roles:

162 self.log.info(f"Using role: {name}")

163 self.current_role = self.roles[name]

164 return True

165 return False

166

167 if __name__ == '__main__':

168 # 创建角色管理器实例

169 role_manager = RoleManager()

170 role_manager.load_roles()

171

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

173 # 打印角色信息

174 print(f"角色名称: {role.name}")

175 print(f"简短描述: {role.short}")

176 print(f"详细描述: {role.detail}")

177 print("-" * 100)

178

179 print(role.envs)

180 print(role.packages)

181

182 # 打印所有提示信息

183 print("\n提示信息:")

184 for name, tip in role:

185 print(f"\n{tip.name}:")

186 print(f"简短描述: {tip.short}")

187 print(f"详细描述: {tip.detail}")