iov_ana/core/config.py

"""
配置文件 —— 支持环境变量 + .env 文件
"""
import os


def _load_dotenv(path: str = ".env"):
    """简易 .env 加载器，不依赖 python-dotenv"""
    if not os.path.isfile(path):
        return
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" not in line:
                continue
            key, _, val = line.partition("=")
            key, val = key.strip(), val.strip().strip('"').strip("'")
            if key and key not in os.environ:  # 环境变量优先
                os.environ[key] = val


# 项目根目录（先定义，.env 加载需要用到）
PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))

# 加载 .env（项目根目录优先，其次当前目录）
_load_dotenv(os.path.join(PROJECT_ROOT, ".env"))
_load_dotenv(".env")

# LLM 配置（兼容 OpenAI API 格式，包括 Ollama / vLLM / DeepSeek 等）
LLM_CONFIG = {
    "api_key": os.getenv("LLM_API_KEY", ""),
    "base_url": os.getenv("LLM_BASE_URL", "https://api.openai.com/v1"),
    "model": os.getenv("LLM_MODEL", "gpt-4o-mini"),
}

# 沙箱安全规则
SANDBOX_RULES = {
    "max_result_rows": 1000,
    "round_floats": 2,
    "suppress_small_n": 5,
    "banned_keywords": [
        "SELECT *", "INSERT", "UPDATE", "DELETE",
        "DROP", "ALTER", "CREATE", "ATTACH", "PRAGMA",
    ],
    "require_aggregation": True,
}

# 项目根目录
PROJECT_ROOT = os.path.dirname(os.path.dirname(__file__))

# 数据库路径
DB_PATH = os.getenv("DB_PATH", os.path.join(PROJECT_ROOT, "demo.db"))

# Playbook 目录
PLAYBOOK_DIR = os.getenv("PLAYBOOK_DIR", os.path.join(PROJECT_ROOT, "playbooks"))

# 图表输出目录
CHARTS_DIR = os.getenv("CHARTS_DIR", os.path.join(PROJECT_ROOT, "charts"))

# 分析控制
MAX_EXPLORATION_ROUNDS = int(os.getenv("MAX_ROUNDS", "6"))