228 lines
8.1 KiB
Python
228 lines
8.1 KiB
Python
"""Real OpenAI-compatible chat client for the phase-2 LLM layer (P5.1).
|
|
|
|
``OpenAILLMClient`` speaks the OpenAI chat-completions wire format over HTTPS
|
|
using only the standard library (``urllib.request``), so the repository ships
|
|
no HTTP dependency. The base URL, model and API key are resolved from the
|
|
environment (:mod:`pf1e_simulator.config`, section 15); the key is never
|
|
embedded in code.
|
|
|
|
Error taxonomy (shared with :mod:`pf1e_simulator.client`):
|
|
|
|
- ``ProviderError`` for transient failures (transport errors, HTTP 429, any
|
|
5xx) -- retried by ``LLMClient``'s backoff loop and counted by the breaker;
|
|
- ``ProviderRejectedError`` for permanent 4xx rejections (bad key, unknown
|
|
model, quota) -- raised immediately, never retried;
|
|
- ``ProviderResponseError`` for a 2xx body that is not a valid
|
|
chat-completions payload -- raised immediately, never retried.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
import urllib.request
|
|
from typing import TYPE_CHECKING, Protocol, Self, TypeGuard
|
|
from urllib.error import HTTPError
|
|
from urllib.request import Request
|
|
|
|
from pf1e_simulator.client import (
|
|
CircuitBreaker,
|
|
ClientError,
|
|
ClientParams,
|
|
Completion,
|
|
CostBudget,
|
|
LLMClient,
|
|
Pricing,
|
|
ProviderError,
|
|
ResponseCache,
|
|
TokenUsage,
|
|
estimate_tokens,
|
|
)
|
|
from pf1e_simulator.combat import OBSERVATION_SCHEMA_VERSION
|
|
from pf1e_simulator.config import (
|
|
ENV_LLM_API_KEY,
|
|
llm_api_key,
|
|
llm_base_url,
|
|
llm_model,
|
|
)
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Callable
|
|
|
|
# Status-code buckets of the chat-completions wire (RFC 9110): 2xx is a
|
|
# success, 429 and every 5xx mean "try again later" (the only statuses safe
|
|
# to retry), and any other 4xx is a permanent rejection.
|
|
_OK_MIN = 200
|
|
_OK_EXCLUSIVE = 300
|
|
_RATE_LIMITED = 429
|
|
_SERVER_ERROR_MIN = 500
|
|
|
|
|
|
class ProviderRejectedError(ClientError):
|
|
"""A permanent 4xx provider rejection - never retried."""
|
|
|
|
|
|
class ProviderResponseError(ClientError):
|
|
"""A 2xx body that is not a valid chat-completions payload."""
|
|
|
|
|
|
class HttpResponse(Protocol):
|
|
"""The minimal ``urllib.response`` surface the wire call uses."""
|
|
|
|
status: int
|
|
|
|
def read(self) -> bytes: ...
|
|
|
|
def __enter__(self) -> Self: ...
|
|
|
|
def __exit__(self, _exc_type: object, _exc_value: object, _traceback: object) -> None: ...
|
|
|
|
|
|
def _is_object_dict(value: object) -> TypeGuard[dict[str, object]]:
|
|
return isinstance(value, dict)
|
|
|
|
|
|
def _is_object_list(value: object) -> TypeGuard[list[object]]:
|
|
return isinstance(value, list)
|
|
|
|
|
|
def _parse_completion(payload: dict[str, object]) -> str:
|
|
"""Extract the completion text from a chat-completions payload."""
|
|
choices_raw = payload.get("choices")
|
|
if not _is_object_list(choices_raw) or not choices_raw:
|
|
msg = "provider response has no choices"
|
|
raise ProviderResponseError(msg)
|
|
choice = choices_raw[0]
|
|
if not _is_object_dict(choice):
|
|
msg = "provider response choice is not a JSON object"
|
|
raise ProviderResponseError(msg)
|
|
message_raw = choice.get("message")
|
|
if not _is_object_dict(message_raw):
|
|
msg = "provider response message is not a JSON object"
|
|
raise ProviderResponseError(msg)
|
|
content = message_raw.get("content")
|
|
if not isinstance(content, str):
|
|
msg = "provider response content is not a string"
|
|
raise ProviderResponseError(msg)
|
|
return content
|
|
|
|
|
|
def _parse_usage(payload: dict[str, object], prompt: str, text: str) -> TokenUsage:
|
|
"""Read the usage block, or estimate tokens when the provider omits it."""
|
|
usage_raw = payload.get("usage")
|
|
if not _is_object_dict(usage_raw):
|
|
return TokenUsage(
|
|
prompt_tokens=estimate_tokens(prompt),
|
|
completion_tokens=estimate_tokens(text),
|
|
)
|
|
prompt_tokens = usage_raw.get("prompt_tokens")
|
|
completion_tokens = usage_raw.get("completion_tokens")
|
|
if not isinstance(prompt_tokens, int) or not isinstance(completion_tokens, int):
|
|
msg = "provider usage block is malformed"
|
|
raise ProviderResponseError(msg)
|
|
return TokenUsage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens)
|
|
|
|
|
|
class OpenAILLMClient(LLMClient):
|
|
"""Real chat-completions client (OpenAI or any compatible gateway).
|
|
|
|
Endpoint, model and key default to the environment (section 15); all may
|
|
be overridden explicitly. ``urlopen`` and ``sleep`` are injectable so
|
|
tests exercise the full retry/breaker machinery offline.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
prompt_version: str,
|
|
model: str | None = None,
|
|
api_key: str | None = None,
|
|
base_url: str | None = None,
|
|
timeout_s: float = 30.0,
|
|
urlopen: Callable[..., HttpResponse] = urllib.request.urlopen,
|
|
schema_version: int = OBSERVATION_SCHEMA_VERSION,
|
|
cache: ResponseCache | None = None,
|
|
budget: CostBudget | None = None,
|
|
pricing: Pricing | None = None,
|
|
retries: int = 2,
|
|
backoff_base_s: float = 0.25,
|
|
breaker: CircuitBreaker | None = None,
|
|
sleep: Callable[[float], None] | None = None,
|
|
) -> None:
|
|
key = api_key if api_key is not None else llm_api_key()
|
|
if key is None or key == "":
|
|
msg = (
|
|
f"no LLM API key: set the {ENV_LLM_API_KEY} environment "
|
|
"variable or pass an explicit api_key"
|
|
)
|
|
raise ValueError(msg)
|
|
if timeout_s <= 0:
|
|
msg = f"timeout_s must be > 0, got {timeout_s}"
|
|
raise ValueError(msg)
|
|
super().__init__(
|
|
model=model if model is not None else llm_model(),
|
|
prompt_version=prompt_version,
|
|
schema_version=schema_version,
|
|
cache=cache,
|
|
budget=budget,
|
|
pricing=pricing,
|
|
retries=retries,
|
|
backoff_base_s=backoff_base_s,
|
|
sleep=sleep if sleep is not None else time.sleep,
|
|
breaker=breaker,
|
|
)
|
|
resolved_base_url = base_url if base_url is not None else llm_base_url()
|
|
self._api_key = key
|
|
self._base_url = resolved_base_url.rstrip("/")
|
|
self._timeout_s = timeout_s
|
|
self._urlopen = urlopen
|
|
|
|
def _call_once(self, prompt: str, *, params: ClientParams) -> Completion:
|
|
"""One chat-completions attempt; raise ``ProviderError`` when retryable."""
|
|
body = {
|
|
"model": self.model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"temperature": params.temperature,
|
|
"max_tokens": params.max_tokens,
|
|
"top_p": params.top_p,
|
|
"response_format": {"type": "json_object"},
|
|
}
|
|
request = Request(
|
|
f"{self._base_url}/chat/completions",
|
|
data=json.dumps(body).encode("utf-8"),
|
|
headers={
|
|
"Authorization": f"Bearer {self._api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
method="POST",
|
|
)
|
|
try:
|
|
with self._urlopen(request, timeout=self._timeout_s) as response:
|
|
status = response.status
|
|
data = response.read()
|
|
except HTTPError as exc:
|
|
status = exc.code
|
|
data = exc.read()
|
|
except OSError as exc:
|
|
msg = f"provider unreachable: {exc}"
|
|
raise ProviderError(msg) from exc
|
|
|
|
if status == _RATE_LIMITED or status >= _SERVER_ERROR_MIN:
|
|
msg = f"provider transient failure: HTTP {status}"
|
|
raise ProviderError(msg)
|
|
if not _OK_MIN <= status < _OK_EXCLUSIVE:
|
|
msg = f"provider rejected request: HTTP {status}"
|
|
raise ProviderRejectedError(msg)
|
|
|
|
try:
|
|
payload = json.loads(data.decode("utf-8"))
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as exc:
|
|
msg = f"provider response is not valid JSON: {exc}"
|
|
raise ProviderResponseError(msg) from exc
|
|
if not _is_object_dict(payload):
|
|
msg = "provider payload is not a JSON object"
|
|
raise ProviderResponseError(msg)
|
|
text = _parse_completion(payload)
|
|
usage = _parse_usage(payload, prompt, text)
|
|
return Completion(text=text, usage=usage)
|