| """vLLM backend — OpenAI-compatible HTTP client to a local vLLM-ROCm server. |
| |
| Reads `MINDXTRAIN_VLLM_BASE_URL` (default `http://localhost:8000/v1`). |
| Streams via the OpenAI `text/event-stream` protocol vLLM exposes. |
| """ |
|
|
| 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("vllm") |
| class VllmBackend(ModelBackend): |
| name = "vllm" |
|
|
| def __init__( |
| self, |
| base_url: str | None = None, |
| timeout_s: float = 120.0, |
| ) -> None: |
| self.base_url = (base_url or os.environ.get("MINDXTRAIN_VLLM_BASE_URL", "http://localhost:8000/v1")).rstrip("/") |
| self.timeout_s = timeout_s |
|
|
| 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), |
| ) |
| 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), |
| ) 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() |
|
|
|
|
| |
| VLLMBackend = VllmBackend |
|
|
|
|
| __all__ = ["VLLMBackend", "VllmBackend"] |
|
|
