109 lines
5.2 KiB
Python
109 lines
5.2 KiB
Python
"""Dependency-free ModelForge capability client. Python 3.11+, no automatic inference retries."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
import re
|
|
from http.client import HTTPException
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from urllib.error import HTTPError, URLError
|
|
from urllib.parse import quote, urlsplit
|
|
from urllib.request import HTTPRedirectHandler, Request, build_opener
|
|
|
|
MAX_RESPONSE_BYTES = 8 * 1024 * 1024
|
|
|
|
|
|
class ModelForgeError(Exception):
|
|
def __init__(self, code: str, message: str, status: int = 0, correlation_id: str | None = None):
|
|
super().__init__(message)
|
|
self.code = code
|
|
self.status = status
|
|
self.correlation_id = correlation_id
|
|
|
|
|
|
class _NoRedirect(HTTPRedirectHandler):
|
|
def redirect_request(self, req: Any, fp: Any, code: int, msg: str, headers: Any, newurl: str) -> None:
|
|
return None
|
|
|
|
|
|
class ModelForgeClient:
|
|
def __init__(self, base_url: str, token: str, capability: str, *, timeout: float = 60, allow_http: bool = False):
|
|
url = urlsplit(base_url)
|
|
if (url.scheme not in ({"http", "https"} if allow_http else {"https"})
|
|
or not url.hostname or url.username or url.password or url.query or url.fragment
|
|
or url.path not in {"", "/"}):
|
|
raise ValueError("Use a gateway origin without userinfo, path, query or fragment; HTTP requires allow_http=True")
|
|
if not re.fullmatch(r"[a-z][a-z0-9_.]*@[1-9][0-9]*", capability):
|
|
raise ValueError("Use an exact versioned capability key")
|
|
if not re.fullmatch(r"mf_[A-Za-z0-9_-]{48}", token):
|
|
raise ValueError("Use a scoped ModelForge service credential")
|
|
if not math.isfinite(timeout) or not 0 < timeout <= 300:
|
|
raise ValueError("timeout must be between 0 and 300 seconds")
|
|
self.base_url = base_url.rstrip("/")
|
|
self.capability = capability
|
|
self.timeout = timeout
|
|
self._token = token
|
|
self._opener = build_opener(_NoRedirect())
|
|
|
|
@classmethod
|
|
def from_env(cls, *, timeout: float = 60) -> ModelForgeClient:
|
|
with Path(os.environ["MODELFORGE_TOKEN_FILE"]).open("r", encoding="ascii") as handle:
|
|
token = handle.read(513).strip()
|
|
return cls(os.environ["MODELFORGE_URL"], token, os.environ["MODELFORGE_CAPABILITY"],
|
|
timeout=timeout, allow_http=os.getenv("MODELFORGE_ALLOW_HTTP") == "true")
|
|
|
|
def _request(self, suffix: str, body: dict[str, Any] | None = None) -> Any:
|
|
path = f"/api/v1/capabilities/{quote(self.capability, safe='')}/{suffix}"
|
|
request = Request(self.base_url + path, data=json.dumps(body).encode() if body is not None else None,
|
|
headers={"Authorization": f"Bearer {self._token}", "Content-Type": "application/json"},
|
|
method="POST" if body is not None else "GET")
|
|
try:
|
|
response = self._opener.open(request, timeout=self.timeout)
|
|
except HTTPError as error:
|
|
correlation = error.headers.get("X-Correlation-ID")
|
|
try:
|
|
with error:
|
|
raw = error.read(MAX_RESPONSE_BYTES + 1)
|
|
except (TimeoutError, OSError, HTTPException):
|
|
raise ModelForgeError("TRANSPORT_ERROR", "Gateway response failed or timed out",
|
|
error.code, correlation) from None
|
|
try:
|
|
payload = json.loads(raw) if len(raw) <= MAX_RESPONSE_BYTES else {}
|
|
except (ValueError, UnicodeError):
|
|
payload = {}
|
|
if not isinstance(payload, dict):
|
|
payload = {}
|
|
nested_error = payload.get("error")
|
|
error_payload = nested_error if isinstance(nested_error, dict) else payload
|
|
raise ModelForgeError(
|
|
str(error_payload.get("code", "HTTP_ERROR")),
|
|
str(error_payload.get("message", "Gateway request refused")),
|
|
error.code,
|
|
error_payload.get("correlation_id") or correlation,
|
|
) from None
|
|
except (URLError, TimeoutError, OSError, HTTPException):
|
|
raise ModelForgeError("TRANSPORT_ERROR", "Gateway connection failed or timed out") from None
|
|
try:
|
|
with response:
|
|
raw = response.read(MAX_RESPONSE_BYTES + 1)
|
|
correlation = response.headers.get("X-Correlation-ID")
|
|
except (TimeoutError, OSError, HTTPException):
|
|
raise ModelForgeError("TRANSPORT_ERROR", "Gateway response failed or timed out") from None
|
|
if len(raw) > MAX_RESPONSE_BYTES:
|
|
raise ModelForgeError("RESPONSE_TOO_LARGE", "Gateway response exceeds the client limit", correlation_id=correlation)
|
|
try:
|
|
return json.loads(raw)
|
|
except (ValueError, UnicodeError):
|
|
raise ModelForgeError("INVALID_RESPONSE", "Gateway response is not JSON", correlation_id=correlation) from None
|
|
|
|
def health(self) -> dict[str, Any]:
|
|
"""Authenticate the exact scope without executing inference."""
|
|
return self._request("check")
|
|
|
|
def invoke(self, payload: dict[str, Any]) -> Any:
|
|
"""Execute once. The client never changes capability, deployment or scheduling priority."""
|
|
return self._request("invoke", payload)
|