logic-skeleton-bench/logicbench/client.py

from __future__ import annotations

import json
import os
import socket
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Any


@dataclass(frozen=True)
class LLMClientConfig:
    model: str
    api_key: str
    base_url: str = "https://api.openai.com/v1"
    api_style: str = "responses"
    timeout_seconds: int = 240
    enable_thinking: bool | None = None

    @classmethod
    def from_env(cls, role: str) -> "LLMClientConfig":
        prefix = role.upper()
        model = os.getenv(f"{prefix}_MODEL", "").strip()
        api_key = os.getenv(f"{prefix}_API_KEY", "").strip()
        base_url = os.getenv(f"{prefix}_BASE_URL", "").strip()
        api_style = os.getenv(f"{prefix}_API_STYLE", "").strip() or "responses"
        timeout_text = os.getenv(f"{prefix}_TIMEOUT_SECONDS", "").strip()
        thinking_text = os.getenv(f"{prefix}_ENABLE_THINKING", "").strip().lower()
        if not model:
            raise ValueError(f"{prefix}_MODEL is required")
        if not api_key:
            raise ValueError(f"{prefix}_API_KEY is required")
        if not base_url:
            raise ValueError(f"{prefix}_BASE_URL is required")
        if api_style not in {"responses", "chat_completions"}:
            raise ValueError(f"{prefix}_API_STYLE must be responses or chat_completions")
        timeout_seconds = int(timeout_text) if timeout_text else 240
        if timeout_seconds < 1:
            raise ValueError(f"{prefix}_TIMEOUT_SECONDS must be positive")
        if thinking_text not in {"", "true", "false"}:
            raise ValueError(f"{prefix}_ENABLE_THINKING must be true or false")
        enable_thinking = None if not thinking_text else thinking_text == "true"
        return cls(
            model=model,
            api_key=api_key,
            base_url=base_url.rstrip("/"),
            api_style=api_style,
            timeout_seconds=timeout_seconds,
            enable_thinking=enable_thinking,
        )


class LLMClient:
    """Small dependency-free client for OpenAI Responses and compatible chat APIs."""

    def __init__(self, config: LLMClientConfig):
        self.config = config

    def request_json(
        self,
        *,
        system: str,
        user: str,
        schema_name: str,
        schema: dict[str, Any],
    ) -> dict[str, Any]:
        if self.config.api_style == "responses":
            path = "/responses"
            body = {
                "model": self.config.model,
                "input": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "text": {
                    "format": {
                        "type": "json_schema",
                        "name": schema_name,
                        "strict": True,
                        "schema": schema,
                    }
                },
                "max_output_tokens": 12000,
            }
        else:
            path = "/chat/completions"
            body = {
                "model": self.config.model,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "response_format": {
                    "type": "json_schema",
                    "json_schema": {"name": schema_name, "strict": True, "schema": schema},
                },
                "max_tokens": 12000,
            }

        if self.config.enable_thinking is not None:
            body["enable_thinking"] = self.config.enable_thinking

        text = self._post(path, body)
        try:
            result = json.loads(text)
        except json.JSONDecodeError as exc:
            raise RuntimeError("LLM API returned non-JSON structured output") from exc
        if not isinstance(result, dict):
            raise RuntimeError("LLM API structured output must be an object")
        return result

    def request_text(self, *, system: str, user: str) -> str:
        if self.config.api_style == "responses":
            path = "/responses"
            body = {
                "model": self.config.model,
                "input": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "max_output_tokens": 1024,
            }
        else:
            path = "/chat/completions"
            body = {
                "model": self.config.model,
                "messages": [
                    {"role": "system", "content": system},
                    {"role": "user", "content": user},
                ],
                "max_tokens": 1024,
            }
        if self.config.enable_thinking is not None:
            body["enable_thinking"] = self.config.enable_thinking
        return self._post(path, body)

    def _post(self, path: str, body: dict[str, Any]) -> str:
        request = urllib.request.Request(
            self.config.base_url + path,
            data=json.dumps(body).encode("utf-8"),
            headers={
                "Authorization": f"Bearer {self.config.api_key}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        try:
            with urllib.request.urlopen(request, timeout=self.config.timeout_seconds) as response:
                payload = json.loads(response.read())
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", errors="replace")[:1000]
            raise RuntimeError(f"LLM API returned HTTP {exc.code}: {detail}") from exc
        except urllib.error.URLError as exc:
            raise RuntimeError(f"LLM API request failed: {exc.reason}") from exc
        except (TimeoutError, socket.timeout) as exc:
            raise RuntimeError(
                f"LLM API request timed out after {self.config.timeout_seconds} seconds"
            ) from exc

        return self._extract_text(payload)

    def _extract_text(self, payload: dict[str, Any]) -> str:
        if self.config.api_style == "chat_completions":
            return payload["choices"][0]["message"]["content"]
        if isinstance(payload.get("output_text"), str):
            return payload["output_text"]
        for item in payload.get("output", []):
            if item.get("type") == "message":
                for content in item.get("content", []):
                    if content.get("type") == "output_text":
                        return content["text"]
        raise RuntimeError("LLM API response did not contain output text")