Coverage for aipyapp/aipy/config.py: 20%
212 statements
« prev ^ index » next coverage.py v7.10.3, created at 2025-08-11 13:03 +0200
1import sys
2import re
3import io
4import datetime
5from pathlib import Path
6import traceback
8from dynaconf import Dynaconf
9from rich import print
10import tomli_w
12from .. import __respath__, T
13from .trustoken import TrustToken
15__PACKAGE_NAME__ = "aipyapp"
17OLD_SETTINGS_FILES = [
18 Path.home() / '.aipy.toml',
19 Path('aipython.toml').resolve(),
20 Path('.aipy.toml').resolve(),
21 Path('aipy.toml').resolve(),
22]
24CONFIG_FILE_NAME = f"{__PACKAGE_NAME__}.toml"
25USER_CONFIG_FILE_NAME = "user_config.toml"
26MCP_CONFIG_FILE_NAME = "mcp.json"
28def 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
43 return config_dir
46CONFIG_DIR = init_config_dir()
47PLUGINS_DIR = CONFIG_DIR / "plugins"
48ROLES_DIR = CONFIG_DIR / "roles"
50def 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
60 config_file_path = config_dir / file_name
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
70 return config_file_path
73def 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()}
80def 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))
91def 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
101def 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 ""
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
117class 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()
124 self.config.update({'_config_dir': config_dir})
126 self.trust_token = TrustToken()
127 # print(self.config.to_dict())
129 def get_work_dir(self):
130 if self.config.workdir:
131 return Path.cwd() / self.config.workdir
132 return Path.cwd()
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 )
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
163 def reload_config(self):
164 self.config = self._load_config()
165 return self.config
167 def get_config(self):
168 return self.config
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)
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])
185 config.update(new_config)
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())
200 with open(self.config_file, "w", encoding="utf-8") as f:
201 # 1. 写入头部注释
202 f.write("\n".join(header_comments) + "\n")
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')
209 # 3. 写入 TOML 内容
210 f.write(toml_content)
212 # 4. 写入尾部注释
213 f.write("\n".join(footer_comments))
214 return config
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)
225 # 配置overwrite
226 self.update_sys_config(cfg_dict, overwrite=True)
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
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."))
252 llms = {}
253 for name, config in self.config.get('llm', {}).items():
254 if config.get("enable", True):
255 llms[name] = config
257 return llms
259 def fetch_config(self):
260 """从tt获取配置并保存到配置文件中。"""
261 self.trust_token.fetch_token(self.save_tt_config)
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
272 if self.check_llm():
273 # 有状态为 enable 的配置文件，则不需要强制要求 trustoken 配置。
274 return
276 # 尝试从旧版本配置迁移
277 if self._migrate_config():
278 # 迁移完成后重新加载配置
279 self.reload_config()
281 # 如果仍然没有可用的 LLM 配置，则从网络拉取
282 if not self.check_llm():
283 if gui:
284 return 'TrustToken'
285 self.fetch_config()
286 self.reload_config()
288 if not self.check_llm():
289 print(T("Missing 'llm' configuration."))
290 sys.exit(1)
292 except Exception as e:
293 traceback.print_exc()
294 sys.exit(1)
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 = []
306 for path in OLD_SETTINGS_FILES:
307 if not path.exists():
308 continue
310 try:
311 config = Dynaconf(settings_files=[path])
312 config_dict = config.to_dict()
313 assert config_dict
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
324 combined_toml_content += content + "\n\n" # Add separator
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
336 except Exception as e:
337 pass
338 if not combined_toml_content:
339 return
341 print(T("""Found old configuration files: {}
342Attempting to migrate configuration from these files...
343After migration, these files will be backed up to {}, please check them.""").format(
344 ', '.join(migrated_files), ', '.join(backup_files)
345 )
346 )
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
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', {})
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
374 except Exception as e:
375 pass
377 return True
379 def _is_tt_config(self, name, config):
380 """
381 判断配置是否符合特定条件
383 参数:
384 name: 配置名称
385 config: 配置内容字典
387 返回: 如果符合条件返回True
388 """
389 # 条件1: 配置名称包含目标关键字
390 if any(keyword in name.lower() for keyword in ['trustoken', 'trust']):
391 return True
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
399 # 条件3: 其他特定标记
400 # type == trust, 且没有base_url.
401 if isinstance(config, dict) and config.get('type') == 'trust' and not base_url:
402 return True
404 return False