File size: 3,623 Bytes
dfb775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """OpenAI-compatible backend — works against OpenAI, Anthropic-via-proxy,
Groq, ZAI, or any HTTP endpoint that follows the OpenAI Chat Completions
protocol.
Reads `MINDXTRAIN_OPENAI_BASE_URL` (default `https://api.openai.com/v1`) and
`MINDXTRAIN_OPENAI_API_KEY`.
"""
from __future__ import annotations
import json
import os
from collections.abc import AsyncIterator
import httpx
from mindxtrain.models.registry import ChatRequest, ChatResponse, ModelBackend, register_backend
@register_backend("openai_compat")
class OpenAICompatBackend(ModelBackend):
name = "openai_compat"
def __init__(
self,
base_url: str | None = None,
api_key: str | None = None,
timeout_s: float = 120.0,
) -> None:
self.base_url = (base_url or os.environ.get("MINDXTRAIN_OPENAI_BASE_URL", "https://api.openai.com/v1")).rstrip("/")
self.api_key = api_key or os.environ.get("MINDXTRAIN_OPENAI_API_KEY", "")
self.timeout_s = timeout_s
def _headers(self) -> dict[str, str]:
h = {"Content-Type": "application/json"}
if self.api_key:
h["Authorization"] = f"Bearer {self.api_key}"
return h
def _payload(self, request: ChatRequest, *, stream: bool) -> dict[str, object]:
return {
"model": request.model,
"messages": [m.model_dump() for m in request.messages],
"temperature": request.temperature,
"max_tokens": request.max_tokens,
"stream": stream,
}
async def chat(self, request: ChatRequest) -> ChatResponse:
async with httpx.AsyncClient(timeout=self.timeout_s) as client:
resp = await client.post(
f"{self.base_url}/chat/completions",
json=self._payload(request, stream=False),
headers=self._headers(),
)
resp.raise_for_status()
data = resp.json()
choice = (data.get("choices") or [{}])[0]
usage = data.get("usage") or {}
return ChatResponse(
model=data.get("model", request.model),
content=(choice.get("message") or {}).get("content", "") or "",
finish_reason=choice.get("finish_reason", "stop") or "stop",
prompt_tokens=int(usage.get("prompt_tokens", 0)),
completion_tokens=int(usage.get("completion_tokens", 0)),
)
async def stream_chat(self, request: ChatRequest) -> AsyncIterator[str]:
async def _gen() -> AsyncIterator[str]:
async with httpx.AsyncClient(timeout=self.timeout_s) as client:
async with client.stream(
"POST",
f"{self.base_url}/chat/completions",
json=self._payload(request, stream=True),
headers=self._headers(),
) as resp:
resp.raise_for_status()
async for raw in resp.aiter_lines():
if not raw or not raw.startswith("data:"):
continue
data = raw[5:].strip()
if data == "[DONE]":
return
try:
chunk = json.loads(data)
except json.JSONDecodeError:
continue
delta = (chunk.get("choices") or [{}])[0].get("delta", {})
token = delta.get("content")
if token:
yield token
return _gen()
__all__ = ["OpenAICompatBackend"]
|