flexinference 0.1.0__py3-none-any.whl

flexinference/__init__.py

@@ -0,0 +1,6 @@
+from __future__ import annotations
+
+from . import models
+from ._client import DEFAULT_BASE_URL, FlexInference, FlexInferenceError
+
+__all__ = ["DEFAULT_BASE_URL", "FlexInference", "FlexInferenceError", "models"]

flexinference/_client.py

@@ -0,0 +1,156 @@
+from __future__ import annotations
+
+import json
+from collections.abc import Iterator, Mapping
+from typing import Any, Literal, cast, overload
+
+import httpx
+
+from .models import (
+    CreateChatCompletionRequest,
+    CreateChatCompletionResponse,
+    CreateChatCompletionStreamResponse,
+    CreateResponse,
+    Response,
+    ResponseStreamEvent,
+)
+
+DEFAULT_BASE_URL = "https://api.flexinference.com/v1"
+
+
+class FlexInferenceError(Exception):
+    """An error returned by FlexInference or by the upstream provider, passed through."""
+
+    def __init__(self, status: int, body: Mapping[str, Any] | None, fallback: str) -> None:
+        err = body.get("error") if isinstance(body, Mapping) else None
+        err = err if isinstance(err, Mapping) else {}
+        super().__init__(err.get("message") or fallback)
+        self.status = status
+        self.type: str | None = err.get("type")
+        self.code: str | None = err.get("code")
+        self.param: str | None = err.get("param")
+
+
+class FlexInference:
+    def __init__(
+        self,
+        *,
+        api_key: str,
+        base_url: str = DEFAULT_BASE_URL,
+        client: httpx.Client | None = None,
+    ) -> None:
+        if not api_key:
+            raise ValueError("FlexInference: `api_key` is required.")
+        self._base_url = base_url.rstrip("/")
+        # Inference calls routinely run far longer than httpx's 5s default; match the
+        # OpenAI SDK's generous read budget while keeping a short connect timeout.
+        self._client = client or httpx.Client(timeout=httpx.Timeout(600.0, connect=10.0))
+        self._headers = {
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream, application/json",
+        }
+        self.responses = _Responses(self)
+        self.chat = _Chat(self)
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> FlexInference:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def _post_json(self, path: str, payload: Mapping[str, Any]) -> Any:
+        r = self._client.post(
+            f"{self._base_url}{path}", headers=self._headers, content=json.dumps(payload)
+        )
+        if r.status_code >= 400:
+            raise FlexInferenceError(r.status_code, _safe_json(r), f"HTTP {r.status_code}")
+        return r.json()
+
+    def _post_stream(self, path: str, payload: Mapping[str, Any]) -> Iterator[Any]:
+        with self._client.stream(
+            "POST", f"{self._base_url}{path}", headers=self._headers, content=json.dumps(payload)
+        ) as r:
+            if r.status_code >= 400:
+                r.read()
+                raise FlexInferenceError(r.status_code, _safe_json(r), f"HTTP {r.status_code}")
+            yield from _parse_sse(r)
+
+
+def _safe_json(r: httpx.Response) -> Mapping[str, Any] | None:
+    try:
+        data = r.json()
+    except (json.JSONDecodeError, ValueError):
+        return None
+    return data if isinstance(data, Mapping) else None
+
+
+def _parse_sse(r: httpx.Response) -> Iterator[Any]:
+    data_lines: list[str] = []
+    for line in r.iter_lines():
+        if line == "":
+            if data_lines:
+                data = "\n".join(data_lines)
+                data_lines = []
+                if data == "[DONE]":
+                    return
+                yield json.loads(data)
+            continue
+        if line.startswith("data:"):
+            data_lines.append(line[5:].lstrip(" "))
+    if data_lines:
+        data = "\n".join(data_lines)
+        if data != "[DONE]":
+            yield json.loads(data)
+
+
+class _Responses:
+    def __init__(self, parent: FlexInference) -> None:
+        self._parent = parent
+
+    @overload
+    def create(self, body: CreateResponse, *, stream: Literal[False] = False) -> Response: ...
+    @overload
+    def create(
+        self, body: CreateResponse, *, stream: Literal[True]
+    ) -> Iterator[ResponseStreamEvent]: ...
+    def create(
+        self, body: CreateResponse, *, stream: bool = False
+    ) -> Response | Iterator[ResponseStreamEvent]:
+        payload = dict(body)
+        payload["stream"] = stream
+        if stream:
+            return self._parent._post_stream("/responses", payload)
+        return cast(Response, self._parent._post_json("/responses", payload))
+
+
+class _Chat:
+    def __init__(self, parent: FlexInference) -> None:
+        self.completions = _ChatCompletions(parent)
+
+
+class _ChatCompletions:
+    def __init__(self, parent: FlexInference) -> None:
+        self._parent = parent
+
+    @overload
+    def create(
+        self, body: CreateChatCompletionRequest, *, stream: Literal[False] = False
+    ) -> CreateChatCompletionResponse: ...
+    @overload
+    def create(
+        self, body: CreateChatCompletionRequest, *, stream: Literal[True]
+    ) -> Iterator[CreateChatCompletionStreamResponse]: ...
+    def create(
+        self, body: CreateChatCompletionRequest, *, stream: bool = False
+    ) -> CreateChatCompletionResponse | Iterator[CreateChatCompletionStreamResponse]:
+        payload = dict(body)
+        payload["stream"] = stream
+        if stream:
+            return self._parent._post_stream("/chat/completions", payload)
+        return cast(
+            CreateChatCompletionResponse, self._parent._post_json("/chat/completions", payload)
+        )