openmask.py

#!/usr/bin/env python3
"""OpenMask — 纯 Python 本地大模型网关 (OpenAI 兼容)。

把多个 LLM 后端 (本地 llama-server / ollama / 在线 OpenAI 兼容 API) 收拢到
同一个 OpenAI 兼容的 /v1 入口后面。零第三方依赖，只用标准库。

与仓库里的 new-api 封装是两套并列的实现：new-api 功能全(多用户/后台UI)，
本文件是轻量、自包含、纯 Python 的替代方案。

运行:
    python openmask.py --config openmask.json
    python openmask.py --host 127.0.0.1 --port 3000 --token sk-demo

配置 (openmask.json) 见 openmask.example.json。
"""
from __future__ import annotations

import argparse
import http.client
import json
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlparse

# ---------------------------------------------------------------------------
# 配置
# ---------------------------------------------------------------------------

DEFAULT_CONFIG = {
    "host": "127.0.0.1",
    "port": 3000,
    "tokens": [],          # 接受的客户端令牌；空列表 = 接受任意(含无令牌)
    "log_file": "openmask.log",
    "models": {},          # 统一模型名 -> [后端列表]
}


def load_config(path):
    cfg = dict(DEFAULT_CONFIG)
    if path:
        with open(path, "r", encoding="utf-8") as f:
            cfg.update(json.load(f))
    cfg.setdefault("models", {})
    cfg.setdefault("tokens", [])
    return cfg


def ordered_backends(backends):
    """按 priority 升序(数字越小越优先)；同优先级保持稳定顺序。"""
    return sorted(backends, key=lambda b: b.get("priority", 100))


# ---------------------------------------------------------------------------
# 日志
# ---------------------------------------------------------------------------

_log_lock = threading.Lock()
_log_fp = None


def log_line(msg):
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line, flush=True)
    if _log_fp is not None:
        with _log_lock:
            try:
                _log_fp.write(line + "\n")
                _log_fp.flush()
            except Exception:
                pass


# ---------------------------------------------------------------------------
# 路由 + 代理核心
# ---------------------------------------------------------------------------

def select_backend(model_name, cfg):
    backends = cfg["models"].get(model_name)
    if not backends:
        return None
    return ordered_backends(backends)


def build_upstream_payload(payload, backend):
    """把网关请求体重写为指向某个后端：替换 model 为后端真实名。"""
    p = dict(payload)
    p["model"] = backend.get("model", p.get("model"))
    return p


def upstream_headers(backend, content_type="application/json"):
    hdrs = {"Content-Type": content_type}
    key = backend.get("api_key") or ""
    if key:
        hdrs["Authorization"] = f"Bearer {key}"
    return hdrs


def _open_conn(parsed, timeout):
    if parsed.scheme == "https":
        return http.client.HTTPSConnection(parsed.hostname, parsed.port or 443, timeout=timeout)
    return http.client.HTTPConnection(parsed.hostname, parsed.port or 80, timeout=timeout)


def proxy_once(backend, payload):
    """对单个后端发起一次请求，返回 (resp, conn) 或抛异常。"""
    parsed = urlparse(backend["url"])
    path = parsed.path or "/"
    if parsed.query:
        path += "?" + parsed.query
    body = json.dumps(build_upstream_payload(payload, backend)).encode("utf-8")
    conn = _open_conn(parsed, timeout=600)
    conn.request("POST", path, body=body, headers=upstream_headers(backend))
    resp = conn.getresponse()
    return resp, conn


def try_proxy(backends, payload):
    """依次尝试后端，返回第一个成功(200)的 (resp, conn, backend)。全失败抛最后一个异常。"""
    last_exc = None
    for backend in backends:
        try:
            resp, conn = proxy_once(backend, payload)
            if resp.status == 200:
                return resp, conn, backend
            # 非 200：读掉 body 再尝试下一个
            try:
                resp.read()
            except Exception:
                pass
            conn.close()
            log_line(f"backend {backend['url']} -> HTTP {resp.status}; failover")
        except Exception as e:  # noqa: BLE001
            last_exc = e
            log_line(f"backend {backend['url']} error: {e}; failover")
    if last_exc:
        raise last_exc
    raise RuntimeError("all backends failed")


# ---------------------------------------------------------------------------
# HTTP 处理器
# ---------------------------------------------------------------------------

class Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"
    cfg = None

    def _cors(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Authorization")

    def _auth_ok(self):
        tokens = self.cfg.get("tokens") or []
        if not tokens:
            return True
        auth = self.headers.get("Authorization", "")
        if auth.lower().startswith("bearer "):
            tok = auth[7:].strip()
            return tok in tokens
        return False

    def _read_body(self):
        length = int(self.headers.get("Content-Length", 0) or 0)
        if length <= 0:
            return b""
        return self.rfile.read(length)

    def do_OPTIONS(self):
        self.send_response(204)
        self._cors()
        self.end_headers()

    def do_GET(self):
        if self.path.rstrip("/") in ("", "/health", "/healthz"):
            self._json(200, {"status": "ok", "mode": "pure-python"})
            return
        if self.path.rstrip("/") == "/v1/models":
            self._handle_models()
            return
        self._json(404, {"error": {"message": f"not found: {self.path}", "type": "not_found"}})

    def _handle_models(self):
        data = {
            "object": "list",
            "data": [
                {
                    "id": name,
                    "object": "model",
                    "created": 0,
                    "owned_by": "openmask",
                    "backends": len(bs),
                }
                for name, bs in self.cfg["models"].items()
            ],
        }
        self._json(200, data)

    def do_POST(self):
        if self.path.rstrip("/") == "/v1/chat/completions":
            self._handle_chat()
            return
        self._json(404, {"error": {"message": f"not found: {self.path}", "type": "not_found"}})

    def _handle_chat(self):
        if not self._auth_ok():
            self._json(401, {"error": {"message": "invalid or missing API key", "type": "auth_error"}})
            return
        raw = self._read_body()
        try:
            payload = json.loads(raw.decode("utf-8"))
        except Exception:
            self._json(400, {"error": {"message": "invalid JSON body", "type": "invalid_request"}})
            return

        model_name = payload.get("model")
        backends = select_backend(model_name, self.cfg)
        if not backends:
            self._json(404, {"error": {"message": f"model not found: {model_name}", "type": "model_not_found"}})
            return

        stream = bool(payload.get("stream"))
        try:
            resp, conn, backend = try_proxy(backends, payload)
        except Exception as e:  # noqa: BLE001
            self._json(502, {"error": {"message": f"all backends failed: {e}", "type": "upstream_error"}})
            return

        log_line(f"serving model={model_name} via backend={backend['url']} stream={stream}")
        if stream:
            self._stream_response(resp, conn)
        else:
            self._buffer_response(resp, conn)

    def _stream_response(self, resp, conn):
        self.send_response(200)
        self.send_header("Content-Type", "text/event-stream")
        self.send_header("Cache-Control", "no-cache")
        self.send_header("Connection", "keep-alive")
        self._cors()
        self.end_headers()
        try:
            for line in resp:
                if not line:
                    continue
                self.wfile.write(line)
                if not line.endswith(b"\n"):
                    self.wfile.write(b"\n")
                self.wfile.flush()
        except Exception as e:  # noqa: BLE001
            log_line(f"stream interrupted: {e}")
        finally:
            conn.close()

    def _buffer_response(self, resp, conn):
        body = resp.read()
        conn.close()
        self.send_response(resp.status)
        self.send_header("Content-Type", "application/json")
        self._cors()
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _json(self, code, obj):
        body = json.dumps(obj).encode("utf-8")
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self._cors()
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def log_message(self, fmt, *args):  # 静音默认访问日志(我们用自己的 log_line)
        return


# ---------------------------------------------------------------------------
# 入口
# ---------------------------------------------------------------------------

def main(argv=None):
    global _log_fp
    ap = argparse.ArgumentParser(description="OpenMask pure-Python LLM gateway")
    ap.add_argument("--config", help="path to openmask.json")
    ap.add_argument("--host", help="listen host (default 127.0.0.1)")
    ap.add_argument("--port", type=int, help="listen port (default 3000)")
    ap.add_argument("--token", action="append", default=[], help="allowed client token (repeatable)")
    args = ap.parse_args(argv)

    cfg = load_config(args.config)
    if args.host:
        cfg["host"] = args.host
    if args.port:
        cfg["port"] = args.port
    if args.token:
        cfg["tokens"] = args.token

    log_path = cfg.get("log_file")
    if log_path:
        try:
            _log_fp = open(log_path, "a", encoding="utf-8")
        except Exception:
            _log_fp = None

    Handler.cfg = cfg
    server = ThreadingHTTPServer((cfg["host"], cfg["port"]), Handler)
    n_models = len(cfg["models"])
    log_line(f"OpenMask (pure-python) listening on http://{cfg['host']}:{cfg['port']}/v1  models={n_models}")
    log_line(f"client tokens: {'ANY' if not cfg['tokens'] else ', '.join(cfg['tokens'])}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        log_line("shutting down")
    finally:
        server.server_close()
        if _log_fp:
            _log_fp.close()


if __name__ == "__main__":
    main()