wehub-resource-sync
348 lines
12 KiB
Python
348 lines
12 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""OpenAI-compatible external endpoint client for admin benchmarks.
|
|
|
|
Shared by the throughput and accuracy benchmarks to run against a remote
|
|
/chat/completions endpoint instead of a local engine. Token counts always
|
|
come from the endpoint's usage payload — SSE chunks are never counted as
|
|
tokens because providers batch multiple tokens per chunk.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import Any, Optional
|
|
|
|
import httpx
|
|
from pydantic import BaseModel, SecretStr, field_validator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# read=3600 covers both the largest between-chunk gap on streams (TTFT of a
|
|
# very long prefill on a slow remote) and the full-response wait for
|
|
# non-streaming accuracy calls. Benchmarks are supervised and cancellable,
|
|
# so a generous ceiling beats spurious failures; connect=15 still fails
|
|
# dead endpoints fast.
|
|
DEFAULT_TIMEOUT = httpx.Timeout(connect=15.0, read=3600.0, write=120.0, pool=30.0)
|
|
|
|
_ERROR_DETAIL_MAX_CHARS = 300
|
|
|
|
|
|
class ExternalEndpointConfig(BaseModel):
|
|
"""Connection settings for an external OpenAI-compatible endpoint.
|
|
|
|
api_key is a SecretStr so the key never leaks through repr() or logs.
|
|
"""
|
|
|
|
base_url: str
|
|
api_key: SecretStr = SecretStr("")
|
|
model: str
|
|
|
|
@field_validator("base_url")
|
|
@classmethod
|
|
def validate_base_url(cls, v: str) -> str:
|
|
v = v.strip().rstrip("/")
|
|
if not v.startswith(("http://", "https://")):
|
|
raise ValueError("base_url must start with http:// or https://")
|
|
return v
|
|
|
|
@field_validator("model")
|
|
@classmethod
|
|
def validate_model(cls, v: str) -> str:
|
|
v = v.strip()
|
|
if not v:
|
|
raise ValueError("model must not be empty")
|
|
return v
|
|
|
|
|
|
class ExternalEndpointError(Exception):
|
|
"""User-presentable failure talking to an external endpoint."""
|
|
|
|
|
|
@dataclass
|
|
class StreamStats:
|
|
"""Timing and token stats from one streamed chat completion."""
|
|
|
|
prompt_tokens: int
|
|
completion_tokens: int
|
|
cached_tokens: int
|
|
start_time: float
|
|
first_content_time: float
|
|
last_content_time: float
|
|
end_time: float
|
|
text: str
|
|
|
|
|
|
@dataclass
|
|
class ChatResult:
|
|
"""Non-streaming chat completion result."""
|
|
|
|
text: str
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
|
|
|
|
def _extract_error_detail(body: str) -> str:
|
|
"""Pull a short human-readable message out of an error response body."""
|
|
try:
|
|
data = json.loads(body)
|
|
if isinstance(data, dict):
|
|
err = data.get("error")
|
|
if isinstance(err, dict) and err.get("message"):
|
|
return str(err["message"])[:_ERROR_DETAIL_MAX_CHARS]
|
|
for key in ("message", "detail"):
|
|
if data.get(key):
|
|
return str(data[key])[:_ERROR_DETAIL_MAX_CHARS]
|
|
except ValueError:
|
|
pass
|
|
text = body.strip()
|
|
if "<" in text and ">" in text:
|
|
return f"unexpected non-JSON response ({len(body)} bytes)"
|
|
return text[:_ERROR_DETAIL_MAX_CHARS] or "no response body"
|
|
|
|
|
|
class ExternalAPIClient:
|
|
"""Async client for an external OpenAI-compatible /chat/completions API.
|
|
|
|
Only a whitelist of request fields is ever sent (model, messages,
|
|
max_tokens, stream, stream_options, temperature) because providers
|
|
commonly reject unknown parameters with HTTP 400.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
config: ExternalEndpointConfig,
|
|
timeout: httpx.Timeout = DEFAULT_TIMEOUT,
|
|
transport: Optional[httpx.AsyncBaseTransport] = None,
|
|
):
|
|
self._config = config
|
|
self._chat_url = f"{config.base_url}/chat/completions"
|
|
headers = {}
|
|
key = config.api_key.get_secret_value()
|
|
if key:
|
|
headers["Authorization"] = f"Bearer {key}"
|
|
# transport is injectable for tests (httpx.MockTransport).
|
|
self._client = httpx.AsyncClient(
|
|
headers=headers,
|
|
timeout=timeout,
|
|
limits=httpx.Limits(max_connections=64),
|
|
transport=transport,
|
|
)
|
|
|
|
async def aclose(self) -> None:
|
|
await self._client.aclose()
|
|
|
|
def _build_body(
|
|
self,
|
|
messages: list[dict],
|
|
max_tokens: int,
|
|
temperature: Optional[float],
|
|
stream: bool,
|
|
) -> dict:
|
|
body: dict[str, Any] = {
|
|
"model": self._config.model,
|
|
"messages": messages,
|
|
"max_tokens": max_tokens,
|
|
}
|
|
if temperature is not None:
|
|
body["temperature"] = temperature
|
|
if stream:
|
|
body["stream"] = True
|
|
body["stream_options"] = {"include_usage": True}
|
|
return body
|
|
|
|
def _map_transport_error(self, exc: httpx.HTTPError) -> ExternalEndpointError:
|
|
base_url = self._config.base_url
|
|
if isinstance(exc, httpx.ConnectTimeout):
|
|
return ExternalEndpointError(
|
|
f"Timed out connecting to external endpoint {base_url}"
|
|
)
|
|
if isinstance(exc, httpx.ConnectError):
|
|
return ExternalEndpointError(
|
|
f"Cannot connect to external endpoint {base_url}: {exc}"
|
|
)
|
|
if isinstance(exc, httpx.ReadTimeout):
|
|
return ExternalEndpointError(
|
|
"External endpoint timed out while waiting for a response"
|
|
)
|
|
return ExternalEndpointError(
|
|
f"External endpoint request failed: {type(exc).__name__}: {exc}"
|
|
)
|
|
|
|
def _status_error(self, status: int, body_text: str) -> ExternalEndpointError:
|
|
if status in (401, 403):
|
|
return ExternalEndpointError(
|
|
f"External endpoint rejected the API key (HTTP {status})"
|
|
)
|
|
detail = _extract_error_detail(body_text)
|
|
return ExternalEndpointError(
|
|
f"External endpoint returned HTTP {status}: {detail}"
|
|
)
|
|
|
|
async def chat_completion(
|
|
self,
|
|
messages: list[dict],
|
|
max_tokens: int,
|
|
temperature: Optional[float],
|
|
) -> ChatResult:
|
|
"""Send a non-streaming chat completion request."""
|
|
body = self._build_body(messages, max_tokens, temperature, stream=False)
|
|
try:
|
|
response = await self._client.post(self._chat_url, json=body)
|
|
except httpx.HTTPError as e:
|
|
raise self._map_transport_error(e) from e
|
|
if response.status_code != 200:
|
|
raise self._status_error(response.status_code, response.text)
|
|
try:
|
|
data = response.json()
|
|
choice = data["choices"][0]
|
|
except (KeyError, IndexError, TypeError, ValueError) as e:
|
|
raise ExternalEndpointError(
|
|
f"External endpoint returned an unexpected response shape: {e}"
|
|
) from e
|
|
text = (choice.get("message") or {}).get("content") or ""
|
|
usage = data.get("usage") or {}
|
|
return ChatResult(
|
|
text=text,
|
|
prompt_tokens=int(usage.get("prompt_tokens") or 0),
|
|
completion_tokens=int(usage.get("completion_tokens") or 0),
|
|
)
|
|
|
|
async def stream_chat_completion(
|
|
self,
|
|
messages: list[dict],
|
|
max_tokens: int,
|
|
temperature: Optional[float],
|
|
) -> StreamStats:
|
|
"""Send a streaming chat completion request and collect stats.
|
|
|
|
Requires the endpoint to return usage via stream_options
|
|
(include_usage); raises ExternalEndpointError otherwise because
|
|
token counts cannot be measured accurately without it.
|
|
"""
|
|
body = self._build_body(messages, max_tokens, temperature, stream=True)
|
|
start_time = time.perf_counter()
|
|
first_content_time: Optional[float] = None
|
|
last_content_time: Optional[float] = None
|
|
usage: Optional[dict] = None
|
|
text_parts: list[str] = []
|
|
|
|
try:
|
|
async with self._client.stream(
|
|
"POST", self._chat_url, json=body
|
|
) as response:
|
|
if response.status_code != 200:
|
|
error_body = await response.aread()
|
|
raise self._status_error(
|
|
response.status_code,
|
|
error_body.decode("utf-8", errors="replace"),
|
|
)
|
|
async for line in response.aiter_lines():
|
|
if not line.startswith("data:"):
|
|
continue
|
|
payload = line[len("data:") :].strip()
|
|
if payload == "[DONE]":
|
|
break
|
|
try:
|
|
chunk = json.loads(payload)
|
|
except ValueError:
|
|
continue
|
|
chunk_usage = chunk.get("usage")
|
|
if chunk_usage:
|
|
usage = chunk_usage
|
|
choices = chunk.get("choices") or []
|
|
if not choices:
|
|
continue
|
|
delta = choices[0].get("delta") or {}
|
|
content = delta.get("content")
|
|
if content:
|
|
text_parts.append(content)
|
|
if content or delta.get("reasoning_content"):
|
|
now = time.perf_counter()
|
|
if first_content_time is None:
|
|
first_content_time = now
|
|
last_content_time = now
|
|
except httpx.HTTPError as e:
|
|
raise self._map_transport_error(e) from e
|
|
|
|
end_time = time.perf_counter()
|
|
|
|
if (
|
|
usage is None
|
|
or usage.get("prompt_tokens") is None
|
|
or usage.get("completion_tokens") is None
|
|
):
|
|
raise ExternalEndpointError(
|
|
"External endpoint does not support stream usage "
|
|
"(stream_options.include_usage); cannot measure token counts"
|
|
)
|
|
if first_content_time is None:
|
|
first_content_time = end_time
|
|
if last_content_time is None:
|
|
last_content_time = end_time
|
|
details = usage.get("prompt_tokens_details") or {}
|
|
return StreamStats(
|
|
prompt_tokens=int(usage["prompt_tokens"]),
|
|
completion_tokens=int(usage["completion_tokens"]),
|
|
cached_tokens=int(details.get("cached_tokens") or 0),
|
|
start_time=start_time,
|
|
first_content_time=first_content_time,
|
|
last_content_time=last_content_time,
|
|
end_time=end_time,
|
|
text="".join(text_parts),
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class _AdapterOutput:
|
|
"""Minimal GenerationOutput stand-in; eval code only reads .text."""
|
|
|
|
text: str
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
|
|
|
|
class ExternalChatAdapter:
|
|
"""Duck-typed engine for accuracy benchmarks against an external API.
|
|
|
|
eval/base.py only touches engine.model_type and engine.chat(), so this
|
|
adapter maps both onto ExternalAPIClient. Sampling comes from the
|
|
constructor-injected profile: the temperature/penalty defaults that
|
|
_eval_single injects via setdefault cannot be told apart from
|
|
profile-supplied values, so all sampling kwargs are accepted and
|
|
dropped here — "deterministic" sends temperature 0 and
|
|
"model_settings" sends no sampling params (remote server defaults).
|
|
"""
|
|
|
|
model_type = None
|
|
|
|
def __init__(self, client: ExternalAPIClient, sampling_profile: str):
|
|
self._client = client
|
|
self._sampling_profile = sampling_profile
|
|
|
|
async def preflight(self) -> None:
|
|
"""Fail fast on auth/URL/model errors before a long evaluation."""
|
|
await self._client.chat_completion(
|
|
messages=[{"role": "user", "content": "Say OK"}],
|
|
max_tokens=4,
|
|
temperature=None,
|
|
)
|
|
|
|
async def chat(
|
|
self,
|
|
messages: list[dict],
|
|
max_tokens: int = 256,
|
|
**kwargs: Any,
|
|
) -> _AdapterOutput:
|
|
temperature = 0.0 if self._sampling_profile == "deterministic" else None
|
|
result = await self._client.chat_completion(
|
|
messages=messages,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
)
|
|
return _AdapterOutput(
|
|
text=result.text,
|
|
prompt_tokens=result.prompt_tokens,
|
|
completion_tokens=result.completion_tokens,
|
|
)
|