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)
|
||||||
@@ -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"
|
||||||
Reference in New Issue
Block a user