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

1#! /usr/bin/env python3 

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

3 

4from typing import List, Dict, Any, Optional 

5from dataclasses import dataclass, field 

6import tomllib 

7import os 

8 

9from loguru import logger 

10 

11@dataclass 

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

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

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

167if __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}")