89 lines
3.3 KiB
Python
89 lines
3.3 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""配置文件加载与校验。"""
|
||
from __future__ import annotations
|
||
|
||
import sys
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import yaml
|
||
|
||
_REQUIRED_KEYS = {
|
||
"database": [
|
||
"driver", "server", "port", "database", "schema", "username", "password",
|
||
"table", "id_column", "param_column", "connect_timeout", "query_timeout",
|
||
],
|
||
"llm": [
|
||
"base_url", "api_key", "model", "temperature", "max_tokens",
|
||
"disable_thinking", "timeout", "max_retry", "retry_backoff_seconds",
|
||
],
|
||
"business": ["max_workers", "log_level", "default_mode", "log_dir", "format_retry"],
|
||
}
|
||
_VALID_MODES = {"coarse", "fine"}
|
||
|
||
|
||
class ConfigError(Exception):
|
||
"""配置文件缺失或格式错误。"""
|
||
|
||
|
||
def load_config(path: str | Path = "config.yaml") -> dict[str, Any]:
|
||
"""读取 YAML 配置文件并做基础字段校验。
|
||
|
||
校验失败会抛出 ConfigError 并说明缺哪个字段，避免程序运行到一半才因为
|
||
配置缺失而报出难以理解的异常。
|
||
"""
|
||
p = Path(path)
|
||
if not p.exists():
|
||
raise ConfigError(f"配置文件不存在: {p.resolve()}")
|
||
|
||
with p.open("r", encoding="utf-8") as f:
|
||
try:
|
||
cfg = yaml.safe_load(f)
|
||
except yaml.YAMLError as e:
|
||
raise ConfigError(f"配置文件 YAML 解析失败: {e}") from e
|
||
|
||
if not isinstance(cfg, dict):
|
||
raise ConfigError("配置文件顶层必须是一个字典 (database/llm/business)")
|
||
|
||
missing_sections = [s for s in _REQUIRED_KEYS if s not in cfg]
|
||
if missing_sections:
|
||
raise ConfigError(f"配置文件缺少顶层配置块: {missing_sections}")
|
||
|
||
for section, keys in _REQUIRED_KEYS.items():
|
||
section_cfg = cfg[section]
|
||
if not isinstance(section_cfg, dict):
|
||
raise ConfigError(f"配置块 '{section}' 必须是字典")
|
||
missing = [k for k in keys if k not in section_cfg]
|
||
if missing:
|
||
raise ConfigError(f"配置块 '{section}' 缺少字段: {missing}")
|
||
|
||
default_mode = cfg["business"]["default_mode"]
|
||
if default_mode not in _VALID_MODES:
|
||
raise ConfigError(
|
||
f"business.default_mode 必须是 {_VALID_MODES} 之一，实际为: {default_mode!r}"
|
||
)
|
||
|
||
# enable_other_category 是后加的可选字段，不放进 _REQUIRED_KEYS——避免
|
||
# 旧配置文件因为没有这个字段就直接报"缺字段"。不写就按 False（关闭）处理；
|
||
# 一旦写了就必须是布尔值，防止 "false"(字符串) 这类误写被当成真值静默生效。
|
||
if "enable_other_category" in cfg["business"]:
|
||
other_flag = cfg["business"]["enable_other_category"]
|
||
if not isinstance(other_flag, bool):
|
||
raise ConfigError(
|
||
f"business.enable_other_category 必须是布尔值 true/false，实际为: {other_flag!r}"
|
||
)
|
||
|
||
# 简单的占位符提醒，防止用户忘记改配置直接运行
|
||
placeholders = []
|
||
if cfg["database"]["password"] == "CHANGE_ME":
|
||
placeholders.append("database.password")
|
||
if cfg["llm"]["api_key"] == "CHANGE_ME":
|
||
placeholders.append("llm.api_key")
|
||
if placeholders:
|
||
print(
|
||
f"[警告] 以下配置项仍是占位符，请修改 config.yaml: {placeholders}",
|
||
file=sys.stderr,
|
||
)
|
||
|
||
return cfg
|