feat: add config loading module with YAML support
实现应用配置加载模块,支持从 YAML 文件或默认值加载配置。 覆盖 ChromaDB、嵌入模型、文档分块和 HTTP 服务配置。 - src/core/config.py: 配置数据类 (AppConfig/ChromaConfig/EmbedConfig/ChunkConfig/ServerConfig) + load_config 函数 - tests/test_config.py: 5 个单元测试覆盖默认值、API 模式和文件加载 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,68 @@
|
||||
"""应用配置加载模块."""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChromaConfig:
|
||||
"""ChromaDB 配置."""
|
||||
|
||||
persist_dir: str = "./data"
|
||||
collection_name: str = "markdown_docs"
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbedConfig:
|
||||
"""嵌入模型配置."""
|
||||
|
||||
mode: str = "local" # "local" | "api"
|
||||
local_model: str = "BAAI/bge-small-zh-v1.5"
|
||||
api_base: str = ""
|
||||
api_key: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ChunkConfig:
|
||||
"""文档分块配置."""
|
||||
|
||||
max_size: int = 1000
|
||||
overlap: int = 100
|
||||
|
||||
|
||||
@dataclass
|
||||
class ServerConfig:
|
||||
"""HTTP 服务配置."""
|
||||
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8000
|
||||
|
||||
|
||||
@dataclass
|
||||
class AppConfig:
|
||||
"""应用总配置."""
|
||||
|
||||
chroma: ChromaConfig = field(default_factory=ChromaConfig)
|
||||
embed: EmbedConfig = field(default_factory=EmbedConfig)
|
||||
chunk: ChunkConfig = field(default_factory=ChunkConfig)
|
||||
server: ServerConfig = field(default_factory=ServerConfig)
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.chroma = ChromaConfig(**kwargs.get("chroma", {}))
|
||||
self.embed = EmbedConfig(**kwargs.get("embed", {}))
|
||||
self.chunk = ChunkConfig(**kwargs.get("chunk", {}))
|
||||
self.server = ServerConfig(**kwargs.get("server", {}))
|
||||
|
||||
|
||||
def load_config(path: str | None = None) -> AppConfig:
|
||||
"""从 YAML 文件加载配置, 若文件不存在则返回默认配置."""
|
||||
config_path = path or "config.yaml"
|
||||
if not Path(config_path).exists():
|
||||
return AppConfig()
|
||||
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f) or {}
|
||||
return AppConfig(**data)
|
||||
Reference in New Issue
Block a user