neuron-tai/packages/gateway/meshnet_gateway/server.py

"""Gateway HTTP server — accepts OpenAI-format requests and routes to inference nodes."""

import http.server
import json
import threading
import time
import urllib.error
import urllib.parse
import urllib.request


class _GatewayHTTPServer(http.server.HTTPServer):
    def __init__(
        self,
        addr,
        handler,
        inference_route: list[str] | None,
        tracker_url: str | None,
        model_presets: list[str] | None,
    ):
        super().__init__(addr, handler)
        self.inference_route = inference_route
        self.tracker_url = tracker_url
        self.model_presets = model_presets


class _GatewayHandler(http.server.BaseHTTPRequestHandler):
    def log_message(self, fmt, *args):  # noqa: suppress request logs in tests
        pass

    def do_GET(self):
        if self.path == "/v1/health":
            self._handle_health()
        elif self.path == "/v1/models":
            self._handle_models()
        else:
            self.send_response(404)
            self.end_headers()

    def do_POST(self):
        if self.path == "/v1/chat/completions":
            self._handle_chat_completions()
        else:
            self.send_response(404)
            self.end_headers()

    def _send_json(self, status: int, data: dict) -> None:
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_json_error(self, status: int, message: str) -> None:
        self._send_json(status, {"error": message})

    def _send_openai_error(self, status: int, message: str, code: str) -> None:
        self._send_json(status, {
            "error": {
                "message": message,
                "type": "invalid_request_error",
                "param": "model",
                "code": code,
            }
        })

    def _handle_health(self) -> None:
        self._send_json(200, {"status": "ok"})

    def _handle_models(self) -> None:
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        created = int(time.time())
        if server.tracker_url is not None:
            try:
                models_resp = _get_json(f"{server.tracker_url}/v1/models")
                self._send_json(200, models_resp)
                return
            except Exception:
                self._send_json_error(503, "tracker unavailable")
                return
        presets = server.model_presets or ["stub-model"]
        data = [
            {"id": p, "object": "model", "created": created, "owned_by": "meshnet"}
            for p in presets
        ]
        self._send_json(200, {"object": "list", "data": data})

    def _handle_chat_completions(self):
        server: _GatewayHTTPServer = self.server  # type: ignore[assignment]
        length = int(self.headers.get("Content-Length", 0))
        try:
            body = json.loads(self.rfile.read(length))
        except (json.JSONDecodeError, ValueError):
            self._send_openai_error(400, "invalid JSON body", "invalid_request_error")
            return
        streaming = bool(body.get("stream", False))

        if server.tracker_url is not None:
            model = body.get("model", "stub-model")
            tracker_path = f"{server.tracker_url}/v1/route?model={urllib.parse.quote(model)}"
            try:
                tracker_resp = _get_json(tracker_path)
                route = tracker_resp["route"]
                if (
                    not isinstance(route, list)
                    or not route
                    or not all(isinstance(node_url, str) and node_url for node_url in route)
                ):
                    raise ValueError("invalid tracker route")
            except urllib.error.HTTPError as exc:
                error_body = _safe_error_body(exc)
                self._send_openai_error(
                    503,
                    error_body.get("error", "model not available"),
                    "model_not_available",
                )
                return
            except (urllib.error.URLError, TimeoutError, json.JSONDecodeError, KeyError, TypeError, ValueError):
                self._send_json_error(503, "tracker unavailable")
                return
        else:
            route = server.inference_route  # type: ignore[assignment]

        messages = body.get("messages", [])

        try:
            # Send messages to the first node in the inference route.
            node_resp = _post_json(f"{route[0]}/v1/infer", {"messages": messages})

            # Forward activations through remaining nodes in the route.
            for node_url in route[1:]:
                node_resp = _post_json(f"{node_url}/v1/infer", {"activations": node_resp["activations"]})
        except (urllib.error.URLError, TimeoutError, KeyError, json.JSONDecodeError) as exc:
            self._send_openai_error(503, f"inference pipeline error: {exc}", "pipeline_error")
            return

        text = node_resp["text"]
        model = body.get("model", "stub-model")

        if streaming:
            self._send_sse_response(text, model)
        else:
            self._send_json(200, {
                "id": "chatcmpl-stub",
                "object": "chat.completion",
                "created": int(time.time()),
                "model": model,
                "choices": [{
                    "index": 0,
                    "message": {
                        "role": "assistant",
                        "content": text,
                    },
                    "finish_reason": "stop",
                }],
                "usage": {
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0,
                },
            })

    def _send_sse_response(self, text: str, model: str) -> None:
        chunk_id = "chatcmpl-stub"
        created = int(time.time())

        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream; charset=utf-8")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("X-Accel-Buffering", "no")
        self.end_headers()

        def _emit(data: str) -> None:
            self.wfile.write(f"data: {data}\n\n".encode())
            self.wfile.flush()

        # Role delta
        _emit(json.dumps({
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}, "finish_reason": None}],
        }))

        # Content delta
        _emit(json.dumps({
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}],
        }))

        # Stop delta
        _emit(json.dumps({
            "id": chunk_id,
            "object": "chat.completion.chunk",
            "created": created,
            "model": model,
            "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
        }))

        # SSE sentinel
        self.wfile.write(b"data: [DONE]\n\n")
        self.wfile.flush()


def _safe_error_body(exc: urllib.error.HTTPError) -> dict:
    try:
        body = json.loads(exc.read())
    except json.JSONDecodeError:
        return {}
    return body if isinstance(body, dict) else {}


def _post_json(url: str, payload: dict, timeout: float = 5.0) -> dict:
    data = json.dumps(payload).encode()
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return json.loads(r.read())


def _get_json(url: str, timeout: float = 5.0) -> dict:
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return json.loads(r.read())


class GatewayServer:
    """HTTP gateway that routes /v1/chat/completions through an ordered inference route.

    Routing modes (mutually exclusive):
    - ``tracker_url``: gateway queries the tracker per request using the model name.
    - ``inference_route``: static ordered list of node base URLs.
    - ``node_url``: single-node shortcut (US-001 compat), equivalent to inference_route=[node_url].

    Optional ``model_presets`` lists the model names returned by GET /v1/models when using
    a static inference route. Ignored when using tracker_url (tracker provides the list).
    """

    def __init__(
        self,
        node_url: str | None = None,
        *,
        inference_route: list[str] | None = None,
        tracker_url: str | None = None,
        model_presets: list[str] | None = None,
        host: str = "127.0.0.1",
        port: int = 0,
    ):
        routing_modes = sum(x is not None for x in (node_url, inference_route, tracker_url))
        if routing_modes > 1:
            raise ValueError("Provide exactly one of: node_url, inference_route, or tracker_url")
        if routing_modes == 0:
            raise ValueError("Provide one of: node_url, inference_route, or tracker_url")

        if node_url is not None:
            inference_route = [node_url]
        if inference_route is not None and not inference_route:
            raise ValueError("inference_route must contain at least one node URL")

        self._inference_route: list[str] | None = inference_route
        self._tracker_url: str | None = tracker_url.rstrip("/") if tracker_url is not None else None
        self._model_presets: list[str] | None = model_presets
        self._host = host
        self._requested_port = port
        self._server: _GatewayHTTPServer | None = None
        self._thread: threading.Thread | None = None
        self.port: int | None = None

    def start(self) -> int:
        if self._server is not None:
            raise RuntimeError("GatewayServer is already running")

        self._server = _GatewayHTTPServer(
            (self._host, self._requested_port),
            _GatewayHandler,
            self._inference_route,
            self._tracker_url,
            self._model_presets,
        )
        self.port = self._server.server_address[1]
        self._thread = threading.Thread(target=self._server.serve_forever, daemon=True)
        self._thread.start()
        return self.port

    def stop(self) -> None:
        if self._server is None:
            return

        self._server.shutdown()
        self._server.server_close()
        if self._thread is not None:
            self._thread.join(timeout=1)
        self._server = None
        self._thread = None
        self.port = None