lean-bench/leanbench/config.py

"""Runtime configuration, read from the environment (and an optional .env file).

Everything the pipeline needs to talk to an LLM or to Lean is settable without
touching code, so the same image runs against any OpenAI-compatible endpoint and
any Lean project.
"""
from __future__ import annotations

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

DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"
DEFAULT_MODEL = "qwen-max"
DEFAULT_API_KEY_ENV = "DASHSCOPE_API_KEY"


def load_dotenv(path: Path | str = ".env") -> None:
    """Load ``KEY=VALUE`` lines into os.environ without overwriting what is set."""
    path = Path(path)
    if not path.is_file():
        return
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, value = line.partition("=")
        os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))


def _env(name: str, default: str) -> str:
    value = os.environ.get(name)
    return default if value is None or value == "" else value


def _env_int(name: str, default: int) -> int:
    try:
        return int(_env(name, str(default)))
    except ValueError:
        return default


def _env_float(name: str, default: float) -> float:
    try:
        return float(_env(name, str(default)))
    except ValueError:
        return default


@dataclass
class LLMSettings:
    model: str = DEFAULT_MODEL
    base_url: str = DEFAULT_BASE_URL
    api_key_env: str = DEFAULT_API_KEY_ENV
    timeout: float = 300.0
    max_retries: int = 4
    temperature: float = 0.0

    @classmethod
    def from_env(cls) -> LLMSettings:
        return cls(
            model=_env("LEANBENCH_LLM_MODEL", DEFAULT_MODEL),
            base_url=_env("LEANBENCH_LLM_BASE_URL", DEFAULT_BASE_URL),
            api_key_env=_env("LEANBENCH_LLM_API_KEY_ENV", DEFAULT_API_KEY_ENV),
            timeout=_env_float("LEANBENCH_LLM_TIMEOUT", 300.0),
            max_retries=_env_int("LEANBENCH_LLM_MAX_RETRIES", 4),
            temperature=_env_float("LEANBENCH_LLM_TEMPERATURE", 0.0),
        )


@dataclass
class LeanSettings:
    """How to invoke Lean. ``command`` is an argv prefix that is given a .lean path."""

    command: list[str] = field(default_factory=lambda: ["lake", "env", "lean"])
    project: str | None = None
    timeout: float = 120.0
    header: str = "import Mathlib\n"

    @classmethod
    def from_env(cls) -> LeanSettings:
        raw = _env("LEANBENCH_LEAN_COMMAND", '["lake","env","lean"]')
        try:
            command = json.loads(raw)
            if not isinstance(command, list) or not all(isinstance(x, str) for x in command):
                raise ValueError
        except ValueError:
            command = raw.split()
        return cls(
            command=command,
            project=os.environ.get("LEANBENCH_LEAN_PROJECT") or None,
            timeout=_env_float("LEANBENCH_LEAN_TIMEOUT", 120.0),
            header=_env("LEANBENCH_LEAN_HEADER", "import Mathlib\n"),
        )


@dataclass
class OcrSettings:
    """Reading pages as pictures. Off by default; ``auto`` is the useful setting."""

    mode: str = "text"          # text | auto | ocr
    backend: str = "vision"     # vision | tesseract
    model: str = "qwen-vl-max"
    dpi: int = 200
    threshold: float = 0.75     # below this quality, `auto` reaches for the camera
    cache_dir: str = ".leanbench-cache/ocr"
    language: str = "eng"       # tesseract only

    @classmethod
    def from_env(cls) -> OcrSettings:
        return cls(
            mode=_env("LEANBENCH_EXTRACTOR", "text").lower(),
            backend=_env("LEANBENCH_OCR_BACKEND", "vision").lower(),
            model=_env("LEANBENCH_OCR_MODEL", "qwen-vl-max"),
            dpi=_env_int("LEANBENCH_OCR_DPI", 200),
            threshold=_env_float("LEANBENCH_OCR_THRESHOLD", 0.75),
            cache_dir=_env("LEANBENCH_OCR_CACHE", ".leanbench-cache/ocr"),
            language=_env("LEANBENCH_OCR_LANGUAGE", "eng"),
        )


@dataclass
class Settings:
    llm: LLMSettings = field(default_factory=LLMSettings)
    lean: LeanSettings = field(default_factory=LeanSettings)
    ocr: OcrSettings = field(default_factory=OcrSettings)
    database: str = "database/leanbench.sqlite"
    max_window_pages: int = 3
    repair_attempts: int = 1

    @classmethod
    def from_env(cls, dotenv: Path | str | None = ".env") -> Settings:
        if dotenv is not None:
            load_dotenv(dotenv)
        return cls(
            llm=LLMSettings.from_env(),
            lean=LeanSettings.from_env(),
            ocr=OcrSettings.from_env(),
            database=_env("LEANBENCH_DATABASE", "database/leanbench.sqlite"),
            max_window_pages=_env_int("LEANBENCH_MAX_WINDOW_PAGES", 3),
            repair_attempts=_env_int("LEANBENCH_REPAIR_ATTEMPTS", 1),
        )