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 a real value.

    A variable that is set but *empty* counts as unset. This is not pedantry: the
    Docker image declares ``ENV DASHSCOPE_API_KEY=`` so the variable is documented
    and overridable at ``docker run``, which left it present-and-empty in every
    container. With plain ``setdefault`` that empty string shadowed the .env file,
    and the only symptom was every OCR page failing with "No API key" -- a whole
    25-page run spent on text the pipeline had already judged unreadable.
    """
    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("=")
        key = key.strip()
        if os.environ.get(key):  # a non-empty value wins; empty or absent does not
            continue
        os.environ[key] = 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),
        )