Coverage for aipyapp/aipy/config.py: 20%

1 import sys

2 import re

3 import io

4 import datetime

5 from pathlib import Path

6 import traceback

7

8 from dynaconf import Dynaconf

9 from rich import print

10 import tomli_w

11

12 from .. import __respath__, T

13 from .trustoken import TrustToken

14

15 __PACKAGE_NAME__ = "aipyapp"

16

17 OLD_SETTINGS_FILES = [

18 Path.home() / '.aipy.toml',

19 Path('aipython.toml').resolve(),

20 Path('.aipy.toml').resolve(),

21 Path('aipy.toml').resolve(),

22]

23

24 CONFIG_FILE_NAME = f"{__PACKAGE_NAME__}.toml"

25 USER_CONFIG_FILE_NAME = "user_config.toml"

26 MCP_CONFIG_FILE_NAME = "mcp.json"

27

28 def init_config_dir():

29 """

30 获取平台相关的配置目录，并确保目录存在

31 """

32 config_dir = Path.home() / f".{__PACKAGE_NAME__}"

33 # 确保目录存在

34 try:

35 config_dir.mkdir(parents=True, exist_ok=True)

36 except PermissionError:

37 print(T("Permission denied to create directory: {}").format(config_dir))

38 raise

39 except Exception as e:

40 print(T("Error creating configuration directory: {}").format(config_dir, str(e)))

41 raise

42

43 return config_dir

44

45

46 CONFIG_DIR = init_config_dir()

47 PLUGINS_DIR = CONFIG_DIR / "plugins"

48 ROLES_DIR = CONFIG_DIR / "roles"

49

50 def get_config_file_path(config_dir=None, file_name=CONFIG_FILE_NAME, create=True):

51 """

52 获取配置文件的完整路径

53 :return: 配置文件的完整路径

54 """

55 if config_dir:

56 config_dir = Path(config_dir)

57 else:

58 config_dir = CONFIG_DIR

59

60 config_file_path = config_dir / file_name

61

62 # 如果配置文件不存在，则创建一个空文件

63 if not config_file_path.exists() and create:

64 try:

65 config_file_path.touch()

66 except Exception as e:

67 print(T("Error creating configuration directory: {}").format(config_file_path))

68 raise

69

70 return config_file_path

71

72

73 def lowercase_keys(d):

74 """递归地将字典中的所有键转换为小写"""

75 if not isinstance(d, dict):

76 return d

77 return {k.lower(): lowercase_keys(v) for k, v in d.items()}

78

79

80 def is_valid_api_key(api_key):

81 """

82 校验是否为有效的 API Key 格式。

83 API Key 格式为字母、数字、减号、下划线的组合，长度在 8 到 128 之间

84 :param api_key: 待校验的 API Key 字符串

85 :return: 如果格式有效返回 True，否则返回 False

86 """

87 pattern = r"^[A-Za-z0-9_-]{8,128}$"

88 return bool(re.match(pattern, api_key))

89

90

91 def get_mcp_config_file(config_dir=None):

92 mcp_config_file = get_config_file_path(

93 config_dir, MCP_CONFIG_FILE_NAME, create=False

94 )

95 # exists and not empty

96 if not mcp_config_file.exists() or mcp_config_file.stat().st_size == 0:

97 return None

98 return mcp_config_file

99

100

101 def get_tt_api_key(settings=None) -> str:

102 """获取 TrustToken API Key

103 :param settings: 配置对象

104 :return: API Key 字符串

105 """

106 if not settings or not isinstance(settings, Dynaconf):

107 return ""

108

109 key = settings.get('llm', {}).get('Trustoken', {}).get('api_key')

110 if not key:

111 key = settings.get('llm', {}).get('trustoken', {}).get('api_key')

112 if not key:

113 return ""

114 return key

115

116

117 class ConfigManager:

118 def __init__(self, config_dir=None):

119 self.config_file = get_config_file_path(config_dir)

120 self.user_config_file = get_config_file_path(config_dir, USER_CONFIG_FILE_NAME)

121 self.default_config = __respath__ / "default.toml"

122 self.config = self._load_config()

123

124 self.config.update({'_config_dir': config_dir})

125

126 self.trust_token = TrustToken()

127 # print(self.config.to_dict())

128

129 def get_work_dir(self):

130 if self.config.workdir:

131 return Path.cwd() / self.config.workdir

132 return Path.cwd()

133

134 def _load_config(self, settings_files=[]):

135 """加载配置文件

136 :param settings_files: 配置文件列表

137 :return: 配置对象

138 """

139 if not settings_files:

140 # 新版本配置文件

141 settings_files = [

142 self.default_config,

143 self.user_config_file,

144 self.config_file,

145 ]

146 # 读取配置文件

147 try:

148 config = Dynaconf(

149 settings_files=settings_files, envvar_prefix="AIPY", merge_enabled=True

150 )

151

152 # check if it's a valid config

153 assert config.to_dict()

154 except Exception as e:

155 # 配置加载异常处理

156 # print(T('error_loading_config'), str(e))

157 # 回退到一个空配置实例，避免后续代码因 config 未定义而出错

158 config = Dynaconf(

159 settings_files=[], envvar_prefix="AIPY", merge_enabled=True

160 )

161 return config

162

163 def reload_config(self):

164 self.config = self._load_config()

165 return self.config

166

167 def get_config(self):

168 return self.config

169

170 def update_sys_config(self, new_config, overwrite=False):

171 """更新aipyapp.toml配置文件

172 :param new_config: 新配置字典, 如 {"workdir": "/path/to/workdir"}

173 """

174 # 加载系统配置文件

175 assert isinstance(new_config, dict)

176

177 if overwrite:

178 # 如果需要覆盖，则直接使用新的配置

179 config = Dynaconf(

180 settings_files=[], envvar_prefix="AIPY", merge_enabled=True

181 )

182 else:

183 config = self._load_config(settings_files=[self.config_file])

184

185 config.update(new_config)

186

187 # 保存到配置文件

188 header_comments = [

189 f"# Configuration file for {__PACKAGE_NAME__}",

190 "# Auto-generated on "

191 + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),

192 f"# 请勿直接修改此文件，除非您了解具体配置格式，如果自定义配置，请放到{self.user_config_file}",

193 f"# Please do not edit this file directly unless you understand the format. If you want to customize the configuration, please edit {self.user_config_file}",

194 "",

195 ]

196 footer_comments = ["", "# End of configuration file"]

197 # print(config.to_dict())

198 cfg_dict = lowercase_keys(config.to_dict())

199

200 with open(self.config_file, "w", encoding="utf-8") as f:

201 # 1. 写入头部注释

202 f.write("\n".join(header_comments) + "\n")

203

204 # 2. 写入 TOML 内容到临时内存文件

205 temp_buffer = io.BytesIO()

206 tomli_w.dump(cfg_dict, temp_buffer, multiline_strings=True)

207 toml_content = temp_buffer.getvalue().decode('utf-8')

208

209 # 3. 写入 TOML 内容

210 f.write(toml_content)

211

212 # 4. 写入尾部注释

213 f.write("\n".join(footer_comments))

214 return config

215

216 def update_api_config(self, new_api_config):

217 """更新配置文件中的API定义"""

218 assert isinstance(new_api_config, dict)

219 assert 'api' in new_api_config

220 config = self._load_config(settings_files=[self.config_file])

221 cfg_dict = lowercase_keys(config.to_dict())

222 cfg_dict.pop('api', None)

223 cfg_dict.update(new_api_config)

224

225 # 配置overwrite

226 self.update_sys_config(cfg_dict, overwrite=True)

227

228 def save_tt_config(self, api_key):

229 config = {

230 'llm': {

231 'trustoken': {

232 'api_key': api_key,

233 'type': 'trust',

234 'base_url': T("https://sapi.trustoken.ai/v1"),

235 'model': 'auto',

236 'default': True,

237 'enable': True,

238 }

239 }

240 }

241 self.update_sys_config(config)

242 return config

243

244 def check_llm(self):

245 """检查是否有可用的LLM配置。

246 只要有可用的配置，就不强制要求trustoken配置。

247 """

248 llm = self.config.get("llm")

249 if not llm:

250 print(T("Missing 'llm' configuration."))

251

252 llms = {}

253 for name, config in self.config.get('llm', {}).items():

254 if config.get("enable", True):

255 llms[name] = config

256

257 return llms

258

259 def fetch_config(self):

260 """从tt获取配置并保存到配置文件中。"""

261 self.trust_token.fetch_token(self.save_tt_config)

262

263 def check_config(self, gui=False):

264 """检查配置文件是否存在，并加载配置。

265 如果配置文件不存在，则创建一个新的配置文件。

266 """

267 try:

268 if not self.config:

269 print(T("Configuration not loaded."))

270 return

271

272 if self.check_llm():

273 # 有状态为 enable 的配置文件，则不需要强制要求 trustoken 配置。

274 return

275

276 # 尝试从旧版本配置迁移

277 if self._migrate_config():

278 # 迁移完成后重新加载配置

279 self.reload_config()

280

281 # 如果仍然没有可用的 LLM 配置，则从网络拉取

282 if not self.check_llm():

283 if gui:

284 return 'TrustToken'

285 self.fetch_config()

286 self.reload_config()

287

288 if not self.check_llm():

289 print(T("Missing 'llm' configuration."))

290 sys.exit(1)

291

292 except Exception as e:

293 traceback.print_exc()

294 sys.exit(1)

295

296 def _migrate_config(self):

297 """

298 Migrates configuration from old settings files (OLD_SETTINGS_FILES)

299 to the new user_config.toml and potentially creates the main aipyapp.toml

300 if a TrustToken configuration is found.

301 """

302 combined_toml_content = ""

303 migrated_files = []

304 backup_files = []

305

306 for path in OLD_SETTINGS_FILES:

307 if not path.exists():

308 continue

309

310 try:

311 config = Dynaconf(settings_files=[path])

312 config_dict = config.to_dict()

313 assert config_dict

314

315 try:

316 content = path.read_text(encoding='utf-8')

317 except UnicodeDecodeError:

318 try:

319 content = path.read_text(encoding='gbk')

320 except Exception as e:

321 print(f"Error reading file {path}: {e}")

322 continue

323

324 combined_toml_content += content + "\n\n" # Add separator

325

326 # 文件内容、格式都正常，则准备迁移

327 migrated_files.append(str(path))

328 # Backup the old file

329 backup_path = path.with_name(f"{path.stem}-backup{path.suffix}")

330 try:

331 path.rename(backup_path)

332 backup_files.append(str(backup_path))

333 except Exception as e:

334 pass

335

336 except Exception as e:

337 pass

338 if not combined_toml_content:

339 return

340

341 print(T("""Found old configuration files: {}

342 Attempting to migrate configuration from these files...

343 After migration, these files will be backed up to {}, please check them.""").format(

344 ', '.join(migrated_files), ', '.join(backup_files)

345 )

346 )

347

348 # Write combined content to user_config.toml

349 try:

350 with open(self.user_config_file, "w", encoding="utf-8") as f:

351 f.write(f"# Migrated from: {', '.join(migrated_files)}\n")

352 f.write(f"# Original files backed up to: {', '.join(backup_files)}\n\n")

353 f.write(combined_toml_content)

354 print(T("Successfully migrated old version user configuration to {}").format(self.user_config_file))

355 except Exception as e:

356 return

357

358 # Now, load the newly created user config to find TT key

359 try:

360 temp_config = self._load_config(settings_files=[self.user_config_file])

361 llm_config = temp_config.get('llm', {})

362

363 for section_name, section_data in llm_config.items():

364 if isinstance(section_data, dict) and self._is_tt_config(

365 section_name, section_data

366 ):

367 api_key = section_data.get('api_key', section_data.get('api-key'))

368 if api_key:

369 # print("Token found:", api_key)

370 self.save_tt_config(api_key)

371 print(T("Successfully migrated old version trustoken configuration to {}").format(self.config_file))

372 break

373

374 except Exception as e:

375 pass

376

377 return True

378

379 def _is_tt_config(self, name, config):

380 """

381 判断配置是否符合特定条件

382

383 参数:

384 name: 配置名称

385 config: 配置内容字典

386

387 返回: 如果符合条件返回True

388 """

389 # 条件1: 配置名称包含目标关键字

390 if any(keyword in name.lower() for keyword in ['trustoken', 'trust']):

391 return True

392

393 base_url = config.get('base_url', config.get('base-url', '')).lower()

394 # 条件2: base_url包含目标域名

395 if isinstance(config, dict) and base_url:

396 if 'trustoken.' in base_url:

397 return True

398

399 # 条件3: 其他特定标记

400 # type == trust, 且没有base_url.

401 if isinstance(config, dict) and config.get('type') == 'trust' and not base_url:

402 return True

403

404 return False