lean-bench/leanbench/llm.py

"""The LLM connector: step 1 of the pipeline, and the thing everything else sits on.

Deliberately dependency-free (stdlib ``urllib`` only) and provider-neutral: any
OpenAI-compatible ``/chat/completions`` endpoint works -- DashScope/Qwen by default,
but also vLLM, Ollama, OpenRouter, DeepSeek, or OpenAI itself. Swap it out entirely
by passing any object with a ``complete(messages, **kwargs) -> str`` method.
"""
from __future__ import annotations

import json
import os
import random
import re
import time
import urllib.error
import urllib.request
from collections.abc import Iterable
from dataclasses import dataclass, field
from typing import Any, Protocol

from .config import LLMSettings

RETRY_STATUS = {408, 409, 425, 429, 500, 502, 503, 504}
_FENCE = re.compile(r"```(?:json|lean4|lean)?\s*(.*?)```", re.DOTALL)


class LLMError(RuntimeError):
    """Any failure to obtain a usable answer from the provider."""


class ChatClient(Protocol):
    """The whole contract the rest of the pipeline depends on."""

    def complete(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        ...


@dataclass
class Usage:
    calls: int = 0
    prompt_tokens: int = 0
    completion_tokens: int = 0

    def add(self, data: dict[str, Any]) -> None:
        self.calls += 1
        usage = data.get("usage") or {}
        self.prompt_tokens += int(usage.get("prompt_tokens") or 0)
        self.completion_tokens += int(usage.get("completion_tokens") or 0)

    def as_dict(self) -> dict[str, int]:
        return {
            "calls": self.calls,
            "prompt_tokens": self.prompt_tokens,
            "completion_tokens": self.completion_tokens,
        }


@dataclass
class OpenAICompatibleClient:
    """POST to ``{base_url}/chat/completions``, with retries and JSON helpers."""

    model: str
    base_url: str
    api_key_env: str = "DASHSCOPE_API_KEY"
    timeout: float = 300.0
    max_retries: int = 4
    temperature: float = 0.0
    usage: Usage = field(default_factory=Usage)

    @classmethod
    def from_settings(cls, settings: LLMSettings) -> OpenAICompatibleClient:
        return cls(
            model=settings.model,
            base_url=settings.base_url,
            api_key_env=settings.api_key_env,
            timeout=settings.timeout,
            max_retries=settings.max_retries,
            temperature=settings.temperature,
        )

    # -- transport ---------------------------------------------------------

    def _api_key(self) -> str:
        key = os.environ.get(self.api_key_env, "").strip()
        if not key:
            raise LLMError(
                f"No API key: set {self.api_key_env} in the environment or in .env"
            )
        return key

    def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
        url = self.base_url.rstrip("/") + "/chat/completions"
        body = json.dumps(payload).encode("utf-8")
        headers = {
            "Content-Type": "application/json",
            "Authorization": "Bearer " + self._api_key(),
        }
        last: Exception | None = None
        for attempt in range(self.max_retries + 1):
            request = urllib.request.Request(url, data=body, headers=headers, method="POST")
            try:
                with urllib.request.urlopen(request, timeout=self.timeout) as response:
                    return json.loads(response.read().decode("utf-8"))
            except urllib.error.HTTPError as error:
                detail = error.read().decode("utf-8", "replace")[:600]
                last = LLMError(f"HTTP {error.code} from {url}: {detail}")
                if error.code not in RETRY_STATUS:
                    raise last from error
            except (urllib.error.URLError, TimeoutError, OSError) as error:
                last = LLMError(f"Network failure talking to {url}: {error}")
            if attempt < self.max_retries:
                time.sleep(min(2 ** attempt, 30) + random.random())
        raise last or LLMError("Request failed for an unknown reason")

    # -- api ---------------------------------------------------------------

    def complete(
        self,
        messages: list[dict[str, str]],
        *,
        temperature: float | None = None,
        json_object: bool = False,
        max_tokens: int | None = None,
        **extra: Any,
    ) -> str:
        payload: dict[str, Any] = {
            "model": self.model,
            "messages": messages,
            "temperature": self.temperature if temperature is None else temperature,
        }
        if max_tokens:
            payload["max_tokens"] = max_tokens
        if json_object:
            payload["response_format"] = {"type": "json_object"}
        payload.update(extra)

        try:
            data = self._post(payload)
        except LLMError as error:
            # Not every OpenAI-compatible server implements response_format.
            if json_object and "response_format" in str(error):
                payload.pop("response_format", None)
                data = self._post(payload)
            else:
                raise
        self.usage.add(data)
        try:
            content = data["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError) as error:
            raise LLMError(f"Unexpected response shape: {json.dumps(data)[:600]}") from error
        if isinstance(content, list):  # some servers return content parts
            content = "".join(part.get("text", "") for part in content if isinstance(part, dict))
        if not isinstance(content, str) or not content.strip():
            raise LLMError("Provider returned empty content")
        return content

    def complete_json(
        self,
        messages: list[dict[str, str]],
        *,
        attempts: int = 3,
        temperature: float | None = None,
        **extra: Any,
    ) -> Any:
        """Ask for JSON and keep asking until it parses (or give up loudly)."""
        conversation = list(messages)
        error_text = ""
        for attempt in range(attempts):
            raw = self.complete(
                conversation,
                temperature=temperature,
                json_object=attempt == 0,
                **extra,
            )
            try:
                return parse_json(raw)
            except ValueError as error:
                error_text = str(error)
                conversation = list(messages) + [
                    {"role": "assistant", "content": raw[:4000]},
                    {
                        "role": "user",
                        "content": (
                            f"That was not valid JSON ({error_text}). Reply with the JSON "
                            "object only: no prose, no markdown fences."
                        ),
                    },
                ]
        raise LLMError(f"Model never produced valid JSON: {error_text}")


def parse_json(raw: str) -> Any:
    """Parse JSON out of a model reply, tolerating fences and surrounding prose."""
    text = raw.strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    fenced = _FENCE.search(text)
    if fenced:
        try:
            return json.loads(fenced.group(1).strip())
        except json.JSONDecodeError:
            text = fenced.group(1).strip()
    candidates = [i for i in (text.find("{"), text.find("[")) if i != -1]
    if candidates:
        start = min(candidates)
        end = max(text.rfind("}"), text.rfind("]"))
        if end > start:
            try:
                return json.loads(text[start : end + 1])
            except json.JSONDecodeError as error:
                raise ValueError(str(error)) from error
    raise ValueError("no JSON found in reply")


def extract_code(raw: str, language: str = "lean") -> str:
    """Pull a fenced code block out of a reply, or return the reply unchanged."""
    fenced = re.search(r"```(?:" + language + r"4?)?[ \t]*\n(.*?)```", raw, re.DOTALL)
    return (fenced.group(1) if fenced else raw).strip()


class RecordingClient:
    """Wraps a client and keeps every exchange -- handy for debugging a bad run."""

    def __init__(self, inner: ChatClient) -> None:
        self.inner = inner
        self.log: list[dict[str, Any]] = []

    def complete(self, messages: list[dict[str, str]], **kwargs: Any) -> str:
        reply = self.inner.complete(messages, **kwargs)
        self.log.append({"messages": messages, "reply": reply})
        return reply

    def complete_json(self, messages: list[dict[str, str]], **kwargs: Any) -> Any:
        result = self.inner.complete_json(messages, **kwargs)  # type: ignore[attr-defined]
        self.log.append({"messages": messages, "reply": result})
        return result


class ScriptedClient:
    """A canned-reply client for tests; replies are consumed in order."""

    def __init__(self, replies: Iterable[str]) -> None:
        self.replies = list(replies)
        self.calls: list[list[dict[str, str]]] = []

    def complete(self, messages: list[dict[str, str]], **_: Any) -> str:
        self.calls.append(messages)
        if not self.replies:
            raise LLMError("ScriptedClient ran out of replies")
        return self.replies.pop(0)

    def complete_json(self, messages: list[dict[str, str]], **_: Any) -> Any:
        return parse_json(self.complete(messages))


def build_client(settings: LLMSettings | None = None) -> OpenAICompatibleClient:
    return OpenAICompatibleClient.from_settings(settings or LLMSettings.from_env())