diff --git a/src/core/config.py b/src/core/config.py new file mode 100644 index 0000000..de0f201 --- /dev/null +++ b/src/core/config.py @@ -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) diff --git a/tests/test_config.py b/tests/test_config.py new file mode 100644 index 0000000..8739ac9 --- /dev/null +++ b/tests/test_config.py @@ -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"