minimind/vllm_native/chat_server.py

94 lines
3.4 KiB
Python

"""Loopback-only web chat and same-origin proxy for the local vLLM API."""
from __future__ import annotations
import argparse
import json
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
CHAT_PAGE = Path(__file__).with_name("chat.html")
class ChatHandler(BaseHTTPRequestHandler):
backend = "http://127.0.0.1:8000"
def log_message(self, fmt: str, *args: object) -> None:
print(f"[{self.log_date_time_string()}] {fmt % args}")
def _respond(self, status: int, body: bytes, content_type: str) -> None:
self.send_response(status)
self.send_header("Content-Type", content_type)
self.send_header("Content-Length", str(len(body)))
self.send_header("Cache-Control", "no-store")
self.end_headers()
self.wfile.write(body)
def _proxy(self, path: str, body: bytes | None = None) -> None:
headers = {"Accept": "application/json"}
if body is not None:
headers["Content-Type"] = "application/json"
request = Request(
f"{self.backend}{path}",
data=body,
headers=headers,
method="POST" if body is not None else "GET",
)
try:
with urlopen(request, timeout=600) as response:
self._respond(
response.status,
response.read(),
response.headers.get("Content-Type", "application/json"),
)
except HTTPError as error:
self._respond(error.code, error.read(), "application/json; charset=utf-8")
except (TimeoutError, URLError, OSError) as error:
message = json.dumps(
{"error": f"无法连接到 vLLM 服务:{error}"}, ensure_ascii=False
).encode("utf-8")
self._respond(502, message, "application/json; charset=utf-8")
def do_GET(self) -> None:
if self.path == "/":
try:
self._respond(200, CHAT_PAGE.read_bytes(), "text/html; charset=utf-8")
except OSError as error:
self._respond(500, str(error).encode(), "text/plain; charset=utf-8")
elif self.path == "/health":
self._proxy("/health")
else:
self._respond(404, b"Not found", "text/plain; charset=utf-8")
def do_POST(self) -> None:
if self.path != "/api/chat":
self._respond(404, b"Not found", "text/plain; charset=utf-8")
return
try:
length = int(self.headers.get("Content-Length", "0"))
except ValueError:
length = 0
if not 0 < length <= 1_048_576:
self._respond(400, b'{"error":"Invalid request size"}', "application/json")
return
self._proxy("/v1/chat/completions", self.rfile.read(length))
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, default=8080)
parser.add_argument("--backend", default="http://127.0.0.1:8000")
args = parser.parse_args()
ChatHandler.backend = args.backend.rstrip("/")
server = ThreadingHTTPServer((args.host, args.port), ChatHandler)
print(f"MiniMind web chat: http://{args.host}:{args.port}", flush=True)
server.serve_forever()
if __name__ == "__main__":
main()