169 lines
6.4 KiB
Python
169 lines
6.4 KiB
Python
# config.py
|
||
|
||
import configparser
|
||
import os
|
||
import logging
|
||
|
||
# 获取根日志器，确保日志配置已在utils.py中完成
|
||
logger = logging.getLogger(__name__)
|
||
|
||
CONFIG_FILE = "config.ini"
|
||
|
||
# 定义默认配置，如果config.ini不存在或缺失，将使用这些值
|
||
DEFAULT_CONFIG = {
|
||
"API": {
|
||
"GEMINI_API_BASE_URL": "https://generativelanguage.googleapis.com/v1beta",
|
||
"GEMINI_API_KEY": "YOUR_GEMINI_API_KEY_HERE", # 默认值，强烈建议用户手动修改
|
||
"DEFAULT_GEMINI_MODEL": "gemini-pro",
|
||
"MEMORY_UPDATE_MODEL": "gemini-pro",
|
||
},
|
||
"Application": {
|
||
"CONTEXT_WINDOW_SIZE": "6",
|
||
"MEMORY_RETENTION_TURNS": "18",
|
||
"MAX_SHORT_TERM_EVENTS": "20",
|
||
"FEATURES_DIR": "data/features",
|
||
"MEMORIES_DIR": "data/memories",
|
||
"CHAT_LOGS_DIR": "data/chat_logs",
|
||
"APP_LOG_FILE_PATH": "logs/app.log",
|
||
},
|
||
"Session": {
|
||
"CURRENT_ROLE_ID": "default_role",
|
||
"CURRENT_MEMORY_ID": "default_memory",
|
||
},
|
||
}
|
||
|
||
|
||
def load_config():
|
||
"""
|
||
从 config.ini 文件加载配置。
|
||
如果文件不存在或某些配置项缺失，则使用默认值并尝试创建/更新文件。
|
||
:return: 包含所有配置的字典。
|
||
"""
|
||
config = configparser.ConfigParser()
|
||
|
||
# 检查config.ini是否存在，如果不存在则创建并写入默认配置
|
||
if not os.path.exists(CONFIG_FILE):
|
||
logger.info(f"config.ini 文件不存在，正在创建默认配置文件：{CONFIG_FILE}")
|
||
for section, options in DEFAULT_CONFIG.items():
|
||
config.add_section(section)
|
||
for key, value in options.items():
|
||
config.set(section, key, value)
|
||
try:
|
||
with open(CONFIG_FILE, "w", encoding="utf-8") as f:
|
||
config.write(f)
|
||
logger.info("默认 config.ini 文件创建成功。")
|
||
except IOError as e:
|
||
logger.error(f"无法写入 config.ini 文件: {e}")
|
||
# 即使写入失败，也继续使用内存中的默认配置
|
||
else:
|
||
logger.info(f"从 {CONFIG_FILE} 加载配置...")
|
||
|
||
# 读取配置
|
||
try:
|
||
config.read(CONFIG_FILE, encoding="utf-8")
|
||
except Exception as e:
|
||
logger.error(f"读取 config.ini 文件时发生错误: {e}，将使用默认配置。")
|
||
# 如果读取失败，清空 configparser 对象，让后续逻辑使用默认值
|
||
config = configparser.ConfigParser()
|
||
|
||
# 将配置转换为字典格式，并确保所有默认值都存在
|
||
app_config = {}
|
||
for section, options in DEFAULT_CONFIG.items():
|
||
app_config[section] = {}
|
||
for key, default_value in options.items():
|
||
# 使用 get() 方法获取值，如果不存在则使用默认值
|
||
# 对于整型或布尔型配置，需要手动转换
|
||
value = config.get(section, key, fallback=default_value)
|
||
|
||
# 尝试将特定配置项转换为数字类型
|
||
if section == "Application" and key in [
|
||
"CONTEXT_WINDOW_SIZE",
|
||
"MEMORY_RETENTION_TURNS",
|
||
"MAX_SHORT_TERM_EVENTS",
|
||
]:
|
||
try:
|
||
app_config[section][key] = int(value)
|
||
except ValueError:
|
||
logger.warning(
|
||
f"配置项 {section}.{key} 的值 '{value}' 不是有效数字，使用默认值 {default_value}。"
|
||
)
|
||
app_config[section][key] = int(default_value)
|
||
elif section == "API" and key == "GEMINI_API_KEY":
|
||
# 优先从环境变量读取 GEMINI_API_KEY
|
||
env_key = os.getenv("GEMINI_API_KEY")
|
||
if env_key:
|
||
app_config[section][key] = env_key
|
||
logger.info("已从环境变量加载 GEMINI_API_KEY。")
|
||
else:
|
||
app_config[section][key] = value
|
||
else:
|
||
app_config[section][key] = value
|
||
|
||
# 如果存在config.ini但缺少某些默认值，更新它
|
||
try:
|
||
updated = False
|
||
for section, options in DEFAULT_CONFIG.items():
|
||
if not config.has_section(section):
|
||
config.add_section(section)
|
||
updated = True
|
||
for key, default_value in options.items():
|
||
if not config.has_option(section, key):
|
||
config.set(section, key, str(default_value)) # 确保写入的是字符串
|
||
updated = True
|
||
if updated:
|
||
with open(CONFIG_FILE, "w", encoding="utf-8") as f:
|
||
config.write(f)
|
||
logger.info("config.ini 文件已更新，补齐缺失配置项。")
|
||
except IOError as e:
|
||
logger.error(f"无法更新 config.ini 文件以补齐缺失配置项: {e}")
|
||
|
||
logger.info("配置加载完成。")
|
||
return app_config
|
||
|
||
|
||
def save_config(new_config):
|
||
"""
|
||
将给定的配置字典保存到 config.ini 文件。
|
||
它会读取现有配置，然后只更新 new_config 中提供的部分。
|
||
:param new_config: 包含要保存的所有配置的字典。
|
||
"""
|
||
config = configparser.ConfigParser()
|
||
|
||
# 首先读取现有配置，以保留未在 new_config 中提供的项
|
||
if os.path.exists(CONFIG_FILE):
|
||
try:
|
||
config.read(CONFIG_FILE, encoding="utf-8")
|
||
except Exception as e:
|
||
logger.error(f"读取现有 config.ini 文件时发生错误: {e}，将只保存新配置。")
|
||
|
||
for section, options in new_config.items():
|
||
if not config.has_section(section):
|
||
config.add_section(section)
|
||
for key, value in options.items():
|
||
# 确保将值转换为字符串，因为 configparser 存储的是字符串
|
||
config.set(section, key, str(value))
|
||
|
||
try:
|
||
with open(CONFIG_FILE, "w", encoding="utf-8") as f:
|
||
config.write(f)
|
||
logger.info(f"配置已成功保存到 {CONFIG_FILE}")
|
||
except IOError as e:
|
||
logger.error(f"保存配置到 {CONFIG_FILE} 失败: {e}")
|
||
|
||
|
||
# 全局变量，存储当前加载的配置
|
||
app_config = load_config()
|
||
|
||
|
||
def get_config():
|
||
"""获取当前加载的配置"""
|
||
return app_config
|
||
|
||
|
||
def set_config(new_config):
|
||
"""更新内存中的配置并保存到文件"""
|
||
global app_config
|
||
app_config = new_config
|
||
save_config(app_config)
|
||
logger.info("内存配置已更新并保存。")
|