llm_client.py

#!/usr/bin/env python3
"""Minimal, stdlib-only chat client used by the ComposeBench harness.

It talks to either an OpenAI-compatible ``/v1/chat/completions`` endpoint or an
Anthropic-compatible ``/v1/messages`` endpoint, and reads its configuration from
environment variables or an optional YAML file (``LLM_YAML``). No third-party
dependencies (no OpenAI/Anthropic SDK, no PyYAML) so the benchmark runs in a
bare Python environment.

Configuration (first match wins):
  1. A YAML file at $LLM_YAML (see llm.example.yaml), keyed by model name.
  2. Environment variables:
       OPENAI_API_BASE / OPENAI_API_KEY                (OpenAI-compatible)
       ANTHROPIC_BASE_URL / ANTHROPIC_API_KEY          (Anthropic-compatible)
     Optional: LLM_API_STYLE=openai|anthropic, LLM_HTTP_TIMEOUT, ANTHROPIC_MAX_TOKENS.

Token accounting is exposed via the module-level ``USAGE`` dict.
"""
import json
import os
import pathlib
import socket
import time
import urllib.error
import urllib.request

USAGE = {"in": 0, "out": 0, "calls": 0}
HTTP_TIMEOUT = float(os.environ.get("LLM_HTTP_TIMEOUT", "90"))
socket.setdefaulttimeout(HTTP_TIMEOUT)


def _strip_comment(line):
    in_quote = None
    esc = False
    out = []
    for ch in line:
        if esc:
            out.append(ch)
            esc = False
            continue
        if ch == "\\" and in_quote:
            out.append(ch)
            esc = True
            continue
        if ch in ("'", '"'):
            if in_quote == ch:
                in_quote = None
            elif not in_quote:
                in_quote = ch
            out.append(ch)
            continue
        if ch == "#" and not in_quote:
            break
        out.append(ch)
    return "".join(out).rstrip()


def _split_kv(line):
    if ":" not in line:
        return None, None
    key, value = line.split(":", 1)
    key = key.strip().strip('"').strip("'")
    return key, value.strip()


def _scalar(value):
    value = value.strip()
    if not value:
        return ""
    if (value[0], value[-1:]) in (('"', '"'), ("'", "'")):
        value = value[1:-1]
    if value.startswith("${") and value.endswith("}"):
        return os.environ.get(value[2:-1], "")
    return value


def load_config(path=None):
    """Parse the small YAML subset used for model config.

    Supports top-level scalar keys plus ``models: {name: {scalar keys...}}``.
    Deliberately avoids a PyYAML dependency.
    """
    path = path or os.environ.get("LLM_YAML")
    cfg = {"models": {}}
    if not path or not pathlib.Path(path).exists():
        return cfg
    section = None
    current_model = None
    for raw in pathlib.Path(path).read_text().splitlines():
        line = _strip_comment(raw)
        if not line.strip():
            continue
        indent = len(line) - len(line.lstrip(" "))
        key, value = _split_kv(line)
        if key is None:
            continue
        if indent == 0:
            current_model = None
            if key == "models":
                section = "models"
            else:
                section = None
                if value:
                    cfg[key] = _scalar(value)
            continue
        if section == "models" and indent == 2:
            current_model = key
            cfg["models"].setdefault(current_model, {})
            continue
        if section == "models" and indent >= 4 and current_model and value:
            cfg["models"][current_model][key] = _scalar(value)
    return cfg


def resolve_model(model, config_path=None):
    cfg = load_config(config_path)
    entry = dict(cfg.get("models", {}).get(model, {}))
    env_base = os.environ.get("ANTHROPIC_BASE_URL") or os.environ.get("OPENAI_API_BASE")
    env_key = (os.environ.get("ANTHROPIC_API_KEY")
               or os.environ.get("OPENAI_API_KEY"))
    base_url = entry.get("baseURL") or env_base or cfg.get("baseURL")
    api_key = entry.get("apiKey") or env_key or cfg.get("apiKey")
    api_model = entry.get("modelId") or model
    style = entry.get("apiStyle") or os.environ.get("LLM_API_STYLE")
    if not style:
        b = (base_url or "").rstrip("/")
        style = "anthropic" if "anthropic.com" in b else "openai"
    return {"model": api_model, "baseURL": base_url, "apiKey": api_key, "style": style.lower()}


def has_credentials(model, config_path=None):
    r = resolve_model(model, config_path)
    return bool(r.get("baseURL") and r.get("apiKey"))


def _endpoint(base_url, style):
    base = (base_url or "").rstrip("/")
    if style == "anthropic":
        return base + "/messages" if base.endswith("/v1") else base + "/v1/messages"
    return base + "/chat/completions" if base.endswith("/v1") else base + "/v1/chat/completions"


def _chat_anthropic(resolved, messages, max_tokens, temperature):
    # Anthropic Messages requires max_tokens; keep it out of the benchmark logic.
    if max_tokens is None:
        max_tokens = int(os.environ.get("ANTHROPIC_MAX_TOKENS", "8192"))
    payload = json.dumps({
        "model": resolved["model"],
        "max_tokens": max_tokens,
        "temperature": temperature,
        "messages": messages,
    }).encode()
    req = urllib.request.Request(
        _endpoint(resolved["baseURL"], "anthropic"),
        data=payload,
        headers={
            "content-type": "application/json",
            "anthropic-version": "2023-06-01",
            "x-api-key": resolved["apiKey"],
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as r:
        data = json.loads(r.read())
    usage = data.get("usage", {}) or {}
    USAGE["in"] += usage.get("input_tokens", 0)
    USAGE["out"] += usage.get("output_tokens", 0)
    USAGE["calls"] += 1
    return "".join(b.get("text", "") for b in data.get("content", []) if b.get("type") == "text")


def _chat_openai(resolved, messages, max_tokens, temperature):
    payload = json.dumps({
        "model": resolved["model"],
        "temperature": temperature,
        "messages": messages,
    }).encode()
    req = urllib.request.Request(
        _endpoint(resolved["baseURL"], "openai"),
        data=payload,
        headers={
            "content-type": "application/json",
            "authorization": "Bearer " + resolved["apiKey"],
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=HTTP_TIMEOUT) as r:
        data = json.loads(r.read())
    usage = data.get("usage", {}) or {}
    USAGE["in"] += usage.get("prompt_tokens", usage.get("input_tokens", 0))
    USAGE["out"] += usage.get("completion_tokens", usage.get("output_tokens", 0))
    USAGE["calls"] += 1
    choice = (data.get("choices") or [{}])[0]
    msg = choice.get("message") or {}
    return msg.get("content") or choice.get("text", "")


def chat(model, messages, max_tokens=None, retries=8, temperature=0):
    resolved = resolve_model(model)
    if not resolved.get("baseURL") or not resolved.get("apiKey"):
        raise RuntimeError(f"No API base/key configured for model {model!r}")
    last = None
    for attempt in range(retries):
        try:
            if resolved["style"] == "anthropic":
                return _chat_anthropic(resolved, messages, max_tokens, temperature)
            return _chat_openai(resolved, messages, max_tokens, temperature)
        except urllib.error.HTTPError as e:
            body = e.read().decode(errors="replace")[:300]
            last = f"HTTP {e.code}: {body}"
        except Exception as e:  # noqa: BLE001 - retry any transient failure
            last = str(e)
        time.sleep(min(60, 2 ** attempt) + (attempt % 5) * 0.13)
    raise RuntimeError(f"API failed after {retries} retries: {last}")