S SmartDocs
Série: RAG AI python 74 lignes · Mis à jour 2026-05-08

config.py

RAG_AI/src/config.py

"""集中管理所有設定。所有可調參數由 .env 控制,避免散落在各處。"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path

from dotenv import load_dotenv

load_dotenv()

ROOT_DIR = Path(__file__).resolve().parent.parent


def _env_bool(key: str, default: bool = False) -> bool:
    return os.getenv(key, str(default)).lower() in {"1", "true", "yes", "on"}


def _env_int(key: str, default: int) -> int:
    try:
        return int(os.getenv(key, str(default)))
    except ValueError:
        return default


@dataclass(frozen=True)
class Config:
    """所有專案設定的單一真實來源(Single Source of Truth)。"""

    # --- 路徑 ---
    data_dir: Path = field(default_factory=lambda: ROOT_DIR / "data")
    db_dir: Path = field(default_factory=lambda: ROOT_DIR / "chroma_db")
    cache_dir: Path = field(default_factory=lambda: ROOT_DIR / "embedding_cache")
    record_db: Path = field(default_factory=lambda: ROOT_DIR / "record_manager.db")
    collection_name: str = "documents"

    # --- 模型 ---
    embedding_model: str = field(
        default_factory=lambda: os.getenv("EMBEDDING_MODEL", "text-embedding-3-small")
    )
    chat_model: str = field(
        default_factory=lambda: os.getenv("CHAT_MODEL", "gpt-4o-mini")
    )
    reranker_model: str = field(
        default_factory=lambda: os.getenv("RERANKER_MODEL", "BAAI/bge-reranker-v2-m3")
    )

    # --- 行為旗標 ---
    use_reranker: bool = field(default_factory=lambda: _env_bool("USE_RERANKER", False))

    # --- 檢索參數 ---
    top_k_retrieve: int = field(default_factory=lambda: _env_int("TOP_K_RETRIEVE", 10))
    top_n_rerank: int = field(default_factory=lambda: _env_int("TOP_N_RERANK", 4))
    chunk_size: int = field(default_factory=lambda: _env_int("CHUNK_SIZE", 500))
    chunk_overlap: int = field(default_factory=lambda: _env_int("CHUNK_OVERLAP", 50))

    # --- BM25 與 Vector 的混合權重 ---
    bm25_weight: float = 0.3
    vector_weight: float = 0.7

    @property
    def has_openai_key(self) -> bool:
        return bool(os.getenv("OPENAI_API_KEY"))


config = Config()


def ensure_dirs() -> None:
    """確保所有需要的目錄存在。"""
    config.data_dir.mkdir(parents=True, exist_ok=True)
    config.db_dir.mkdir(parents=True, exist_ok=True)
    config.cache_dir.mkdir(parents=True, exist_ok=True)

Articles liés