Coverage for config_manager.py: 62%

1""" 

2配置管理模块 

3 

4统一管理用户配置，所有配置文件存放在 ~/.qrclaw/ 目录下： 

5- ~/.qrclaw/config.yaml 用户配置（API Key 等） 

6- ~/.qrclaw/permissions.yaml 权限配置 

7- ~/.qrclaw/MEMORY.md 中期记忆 

8- ~/.qrclaw/sessions/ 会话历史 

9- ~/.qrclaw/logs/ 日志文件 

10""" 

11 

12import os 

13import yaml 

14from pathlib import Path 

15from typing import Any 

16from qrclaw.logger import get_logger 

17 

18logger = get_logger("qrclaw.config_manager") 

19 

20# 配置目录 

21CONFIG_DIR = Path.home() / ".qrclaw" 

22CONFIG_FILE = CONFIG_DIR / "config.yaml" 

23 

24# 默认配置 

25DEFAULT_CONFIG = { 

26 "agent": { 

27 "name": "QRClaw", 

28 "max_iterations": 50, 

29 }, 

30 "llm": { 

31 "provider": "openai", # openai | vertex 

32 "openai": { 

33 "api_key": "", 

34 "model": "gpt-4o", 

35 "base_url": "", 

36 "max_tokens": 128000, 

37 }, 

38 "vertex": { 

39 "api_key": "", 

40 }, 

41 }, 

42 "search": { 

43 "tavily_api_key": "", 

44 }, 

45 "log": { 

46 "level": "INFO", 

47 "max_days": 30, 

48 "to_file": True, 

49 "to_console": True, 

50 "console_level": "WARNING", 

51 }, 

52 "heartbeat": { 

53 "enabled": True, 

54 "interval": 3600, 

55 }, 

56 "compress": { 

57 "threshold_ratio": 0.6, 

58 "target_min_ratio": 0.20, 

59 "target_max_ratio": 0.25, 

60 }, 

61} 

62 

63# 配置注释 

64CONFIG_HEADER = """# ═══════════════════════════════════════════════════════════════ 

65# QRClaw 配置文件 

66# ═══════════════════════════════════════════════════════════════ 

67# 文档：https://github.com/fu-qingrong/qrclaw 

68# 修改配置后无需重启，下次启动自动生效。 

69# ═══════════════════════════════════════════════════════════════ 

70 

71""" 

72 

73 

74def ensure_config_dir(): 

75 """确保配置目录存在""" 

76 CONFIG_DIR.mkdir(parents=True, exist_ok=True) 

77 

78 

79def _write_config(config: dict): 

80 """写入配置文件""" 

81 ensure_config_dir() 

82 with open(CONFIG_FILE, "w", encoding="utf-8") as f: 

83 f.write(CONFIG_HEADER) 

84 yaml.dump(config, f, default_flow_style=False, allow_unicode=True, sort_keys=False) 

85 

86 

87def _deep_merge(base: dict, override: dict) -> dict: 

88 """深度合并两个字典""" 

89 result = base.copy() 

90 for key, value in override.items(): 

91 if key in result and isinstance(result[key], dict) and isinstance(value, dict): 

92 result[key] = _deep_merge(result[key], value) 

93 else: 

94 result[key] = value 

95 return result 

96 

97 

98def init_config(): 

99 """初始化配置文件（不存在则创建默认配置）""" 

100 ensure_config_dir() 

101 

102 if CONFIG_FILE.exists(): 

103 logger.debug(f"使用配置文件: {CONFIG_FILE}") 

104 return 

105 

106 _write_config(DEFAULT_CONFIG) 

107 logger.info(f"创建默认配置文件: {CONFIG_FILE}") 

108 logger.info("请编辑配置文件，填入你的 API Key") 

109 

110 

111def load_config(): 

112 """加载配置并注入到环境变量""" 

113 if not CONFIG_FILE.exists(): 

114 init_config() 

115 

116 try: 

117 with open(CONFIG_FILE, "r", encoding="utf-8") as f: 

118 config = yaml.safe_load(f) or {} 

119 

120 config = _deep_merge(DEFAULT_CONFIG, config) 

121 _inject_to_env(config) 

122 logger.debug(f"已加载配置: {CONFIG_FILE}") 

123 

124 except Exception as e: 

125 logger.error(f"加载配置失败: {e}", exc_info=True) 

126 

127 

128def _inject_to_env(config: dict): 

129 """将配置注入到环境变量""" 

130 os.environ.setdefault("AGENT_NAME", config["agent"]["name"]) 

131 os.environ.setdefault("MAX_ITERATIONS", str(config["agent"]["max_iterations"])) 

132 

133 os.environ.setdefault("LLM_PROVIDER", config["llm"]["provider"]) 

134 os.environ.setdefault("OPENAI_API_KEY", config["llm"]["openai"]["api_key"]) 

135 os.environ.setdefault("OPENAI_MODEL", config["llm"]["openai"]["model"]) 

136 os.environ.setdefault("OPENAI_BASE_URL", config["llm"]["openai"]["base_url"]) 

137 os.environ.setdefault("MODEL_MAX_TOKENS", str(config["llm"]["openai"]["max_tokens"])) 

138 os.environ.setdefault("VERTEX_API_KEY", config["llm"]["vertex"]["api_key"]) 

139 

140 os.environ.setdefault("TAVILY_API_KEY", config["search"]["tavily_api_key"]) 

141 

142 os.environ.setdefault("LOG_LEVEL", config["log"]["level"]) 

143 os.environ.setdefault("LOG_MAX_DAYS", str(config["log"]["max_days"])) 

144 os.environ.setdefault("LOG_TO_FILE", str(config["log"]["to_file"]).lower()) 

145 os.environ.setdefault("LOG_TO_CONSOLE", str(config["log"]["to_console"]).lower()) 

146 os.environ.setdefault("LOG_CONSOLE_LEVEL", config["log"]["console_level"]) 

147 

148 os.environ.setdefault("HEARTBEAT_ENABLED", str(config["heartbeat"]["enabled"]).lower()) 

149 os.environ.setdefault("HEARTBEAT_INTERVAL", str(config["heartbeat"]["interval"])) 

150 

151 

152def get_config() -> dict: 

153 """获取完整配置字典""" 

154 if not CONFIG_FILE.exists(): 

155 return DEFAULT_CONFIG.copy() 

156 

157 with open(CONFIG_FILE, "r", encoding="utf-8") as f: 

158 config = yaml.safe_load(f) or {} 

159 

160 return _deep_merge(DEFAULT_CONFIG, config) 

161 

162 

163def get(key: str, default: Any = None) -> Any: 

164 """获取配置项（支持点号路径，如 llm.openai.api_key）""" 

165 config = get_config() 

166 keys = key.split(".") 

167 value = config 

168 

169 for k in keys: 

170 if isinstance(value, dict) and k in value: 

171 value = value[k] 

172 else: 

173 return default 

174 

175 return value 

176 

177 

178def set_config(key: str, value: Any): 

179 """设置配置项（支持点号路径）""" 

180 config = get_config() 

181 keys = key.split(".") 

182 target = config 

183 

184 for k in keys[:-1]: 

185 if k not in target: 

186 target[k] = {} 

187 target = target[k] 

188 

189 target[keys[-1]] = value 

190 _write_config(config) 

191 

192 

193def get_config_path() -> Path: 

194 """获取配置文件路径""" 

195 return CONFIG_FILE 

196 

197 

198def get_config_dir() -> Path: 

199 """获取配置目录路径""" 

200 return CONFIG_DIR