cleveragents-core/robot/e2e/check_openai_key.py

"""Probe the OpenAI API to verify that the key is functional.

This script is used by the E2E test suite setup (via ``Resolve LLM Actor``
in ``common_e2e.resource``) to decide which LLM actor to use before any test
runs.  It replaces the naive key-presence check that would select the OpenAI
actor even when the key exists but is quota-exhausted.

Usage::

    python check_openai_key.py

Exit codes:
    0 — OpenAI API returned HTTP 200; key is functional.
    1 — Key is missing, quota-exhausted (HTTP 429), unauthorised (HTTP 401),
        or any other error occurred; caller should fall back to Anthropic.

The script uses only Python standard-library modules (``urllib.request``,
``json``, ``os``) — no third-party dependencies are required.

The probe sends the cheapest possible request:
    model: gpt-4o-mini
    messages: [{"role": "user", "content": "Hi"}]
    max_tokens: 1
This costs a fraction of a cent and adds < 5 s to suite setup time.
"""

from __future__ import annotations

import contextlib
import json
import os
import urllib.error
import urllib.request

_OPENAI_URL = "https://api.openai.com/v1/chat/completions"
_PROBE_MODEL = "gpt-4o-mini"
_TIMEOUT_SECONDS = 15


def _probe(api_key: str) -> tuple[bool, str]:
    """Send a minimal chat-completion request to the OpenAI API.

    Returns ``(True, "ok")`` when the API responds with HTTP 200.
    Returns ``(False, reason)`` for any other outcome.
    """
    payload = json.dumps(
        {
            "model": _PROBE_MODEL,
            "messages": [{"role": "user", "content": "Hi"}],
            "max_tokens": 1,
        }
    ).encode()

    req = urllib.request.Request(
        _OPENAI_URL,
        data=payload,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=_TIMEOUT_SECONDS) as resp:
            if resp.status == 200:
                return True, "ok"
            # Unexpected non-200 success-range status
            return False, f"unexpected HTTP {resp.status}"
    except urllib.error.HTTPError as exc:
        body = ""
        with contextlib.suppress(Exception):
            body = exc.read().decode(errors="replace")
        return False, f"HTTP {exc.code}: {body[:200]}"
    except urllib.error.URLError as exc:
        return False, f"network error: {exc.reason}"
    except TimeoutError:
        return False, f"timed out after {_TIMEOUT_SECONDS}s"
    except Exception as exc:
        return False, f"unexpected error: {exc}"


def main() -> int:
    """Entrypoint.  Returns the process exit code."""
    api_key = os.environ.get("OPENAI_API_KEY", "")
    if not api_key:
        print("OPENAI_API_KEY is not set")
        return 1

    ok, reason = _probe(api_key)
    print(reason)
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())