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:
2026-07-05 00:52:48 +08:00
parent cfc41f8af5
commit 138c5c1881
2 changed files with 150 additions and 0 deletions
+68
View File
@@ -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)
+82
View File
@@ -0,0 +1,82 @@
"""配置加载模块测试."""
import tempfile
from pathlib import Path
import pytest
from src.core.config import AppConfig, EmbedConfig, ChunkConfig, load_config
class TestEmbedConfig:
"""嵌入配置解析."""
def test_local_mode_defaults(self):
"""默认 local 模式,带默认模型名."""
data = {"embed": {"mode": "local"}}
cfg = AppConfig(**data)
assert cfg.embed.mode == "local"
assert cfg.embed.local_model == "BAAI/bge-small-zh-v1.5"
assert cfg.embed.api_base == ""
def test_api_mode_fields(self):
"""api 模式下 api_base 和 api_key 可设置."""
data = {
"embed": {
"mode": "api",
"api_base": "https://api.openai.com/v1",
"api_key": "sk-test",
}
}
cfg = AppConfig(**data)
assert cfg.embed.mode == "api"
assert cfg.embed.api_base == "https://api.openai.com/v1"
assert cfg.embed.api_key == "sk-test"
class TestChunkConfig:
"""分块配置解析."""
def test_default_values(self):
"""分块默认值正确."""
data = {"chunk": {}}
cfg = AppConfig(**data)
assert cfg.chunk.max_size == 1000
assert cfg.chunk.overlap == 100
class TestLoadConfig:
"""load_config 函数测试."""
def test_load_from_yaml_file(self):
"""从 YAML 文件加载配置."""
yaml_content = """
chroma:
persist_dir: /tmp/test_data
collection_name: test_collection
embed:
mode: local
chunk:
max_size: 500
server:
port: 9000
"""
with tempfile.NamedTemporaryFile(
mode="w", suffix=".yaml", delete=False
) as f:
f.write(yaml_content)
tmp_path = f.name
try:
cfg = load_config(tmp_path)
assert cfg.chroma.persist_dir == "/tmp/test_data"
assert cfg.chroma.collection_name == "test_collection"
assert cfg.chunk.max_size == 500
assert cfg.server.port == 9000
finally:
Path(tmp_path).unlink()
def test_load_missing_file_uses_defaults(self):
"""配置文件不存在时使用默认值."""
cfg = load_config("/nonexistent/config.yaml")
assert cfg.chroma.persist_dir == "./data"
assert cfg.embed.mode == "local"