vllm_agent_strategy/app/settings.py

from dataclasses import dataclass
import json
import os
from pathlib import Path
from typing import Any


def _load_config_file() -> dict[str, Any]:
    explicit = os.getenv("STRATEGY_CONFIG_FILE")
    candidates = [explicit] if explicit else ["config.local.json", "config.json"]
    for item in candidates:
        if not item:
            continue
        path = Path(item)
        if path.exists():
            with path.open("r", encoding="utf-8") as fh:
                return json.load(fh)
    return {}


def _value(config: dict[str, Any], name: str, default: Any = None) -> Any:
    if name in os.environ:
        return os.environ[name]
    if name in config:
        return config[name]
    lower = name.lower()
    if lower in config:
        return config[lower]
    return default


def _bool_value(config: dict[str, Any], name: str, default: bool) -> bool:
    raw = _value(config, name, default)
    if isinstance(raw, bool):
        return raw
    if raw is None:
        return default
    return str(raw).strip().lower() in {"1", "true", "yes", "y", "on"}


def _int_value(config: dict[str, Any], name: str, default: int) -> int:
    raw = _value(config, name, default)
    if raw is None or raw == "":
        return default
    return int(raw)


@dataclass(frozen=True)
class Settings:
    auth_token: str
    hf_token: str
    contest_api_token: str
    email: str
    user_id: str
    contributors: str
    strategy_id: str

    modelhub_url: str = "https://modelhub.org.cn"
    modelhub_api_base: str = "https://modelhub.org.cn/api"
    modelhub_adminapi_base: str = "https://modelhub.org.cn/adminApi"

    db_path: str = "./state.db"
    host: str = "0.0.0.0"
    port: int = 8080
    poll_interval_seconds: int = 60
    download_poll_interval_seconds: int = 300
    request_timeout_seconds: int = 60
    max_self_downloads: int = 8
    max_user_tasks: int = 2000
    default_gpu_aliases: str = "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8"
    default_max_model_len: int = 1024

    submit_dry_run: bool = False
    inject_strategy_id: bool = True
    contest_task_strategy_field: str = "strategyId"
    bounty_allow_direct_submit: bool = True
    auto_load_seed_path: str = ""

    enable_not_adapted_crawler: bool = True
    target_machine_names: str = ""
    target_task_level: str = "文本生成"
    modelhub_page_size: int = 100
    crawler_refresh_seconds: int = 3600
    crawler_max_pages: int = 400

    enable_download_success_crawler: bool = True
    download_success_page_size: int = 50
    download_success_max_pages: int = 100

    config_file_loaded: str = ""

    @property
    def default_gpu_alias_list(self) -> list[str]:
        return [item.strip() for item in self.default_gpu_aliases.split(",") if item.strip()]


def load_settings(require_secrets: bool = True) -> Settings:
    config = _load_config_file()
    explicit = os.getenv("STRATEGY_CONFIG_FILE")
    loaded_path = ""
    if explicit and Path(explicit).exists():
        loaded_path = explicit
    elif Path("config.local.json").exists():
        loaded_path = "config.local.json"
    elif Path("config.json").exists():
        loaded_path = "config.json"

    hf_token = _value(config, "HF_TOKEN", None)
    if hf_token is None:
        hf_token = _value(config, "HFTOKEN", "")

    settings = Settings(
        auth_token=str(_value(config, "AUTH_TOKEN", "") or ""),
        hf_token=str(hf_token or ""),
        contest_api_token=str(_value(config, "CONFTEST_API_TOKEN", "") or ""),
        email=str(_value(config, "EMAIL", "") or ""),
        user_id=str(_value(config, "USER_ID", "") or ""),
        contributors=str(_value(config, "CONTRIBUTORS", "") or ""),
        strategy_id=str(_value(config, "STRATEGY_ID", "") or ""),
        modelhub_url=str(_value(config, "MODELHUB_URL", "https://modelhub.org.cn")),
        modelhub_api_base=str(_value(config, "MODELHUB_API_BASE", "https://modelhub.org.cn/api")),
        modelhub_adminapi_base=str(_value(config, "MODELHUB_ADMINAPI_BASE", "https://modelhub.org.cn/adminApi")),
        db_path=str(_value(config, "DB_PATH", "./state.db")),
        host=str(_value(config, "HOST", "0.0.0.0")),
        port=_int_value(config, "PORT", 8080),
        poll_interval_seconds=_int_value(config, "POLL_INTERVAL_SECONDS", 60),
        download_poll_interval_seconds=_int_value(config, "DOWNLOAD_POLL_INTERVAL_SECONDS", 300),
        request_timeout_seconds=_int_value(config, "REQUEST_TIMEOUT_SECONDS", 60),
        max_self_downloads=_int_value(config, "MAX_SELF_DOWNLOADS", 8),
        max_user_tasks=_int_value(config, "MAX_USER_TASKS", 2000),
        default_gpu_aliases=str(_value(config, "DEFAULT_GPU_ALIASES", "910b,k100,p800,166m,bi100,bi150,c500,s4000,mrv100,mlu370-x4,mlu370-x8")),
        default_max_model_len=_int_value(config, "DEFAULT_MAX_MODEL_LEN", 1024),
        submit_dry_run=_bool_value(config, "SUBMIT_DRY_RUN", False),
        inject_strategy_id=_bool_value(config, "INJECT_STRATEGY_ID", True),
        contest_task_strategy_field=str(_value(config, "CONTEST_TASK_STRATEGY_FIELD", "strategyId")),
        bounty_allow_direct_submit=_bool_value(config, "BOUNTY_ALLOW_DIRECT_SUBMIT", True),
        auto_load_seed_path=str(_value(config, "AUTO_LOAD_SEED_PATH", "") or ""),
        enable_not_adapted_crawler=_bool_value(config, "ENABLE_NOT_ADAPTED_CRAWLER", True),
        target_machine_names=str(_value(config, "TARGET_MACHINE_NAMES", "") or ""),
        target_task_level=str(_value(config, "TARGET_TASK_LEVEL", "文本生成")),
        modelhub_page_size=_int_value(config, "MODELHUB_PAGE_SIZE", 100),
        crawler_refresh_seconds=_int_value(config, "CRAWLER_REFRESH_SECONDS", 3600),
        crawler_max_pages=_int_value(config, "CRAWLER_MAX_PAGES", 400),
        enable_download_success_crawler=_bool_value(config, "ENABLE_DOWNLOAD_SUCCESS_CRAWLER", True),
        download_success_page_size=_int_value(config, "DOWNLOAD_SUCCESS_PAGE_SIZE", 50),
        download_success_max_pages=_int_value(config, "DOWNLOAD_SUCCESS_MAX_PAGES", 100),
        config_file_loaded=loaded_path,
    )
    if require_secrets:
        missing = []
        required_names = [
            ("AUTH_TOKEN", settings.auth_token),
            ("CONFTEST_API_TOKEN", settings.contest_api_token),
            ("EMAIL", settings.email),
            ("USER_ID", settings.user_id),
            ("CONTRIBUTORS", settings.contributors),
        ]
        for name, value in required_names:
            if not value:
                missing.append(name)
        if settings.inject_strategy_id and not settings.strategy_id:
            missing.append("STRATEGY_ID")
        if missing:
            raise RuntimeError(f"Missing required settings: {', '.join(missing)}")
    return settings