Update professional engineer-grade serve_openai_api.py (CLI chat + OpenAI server + capability tags)

serve_openai_api.py

@@ -0,0 +1,224 @@
+#!/usr/bin/env python3
+"""
+MarsNext Phobos-1 — OpenAI-Compatible REST API Server (`/v1/chat/completions` & `/v1/models`)
+(`phobos-1/models/phobos_scratch_v1/serve_openai_api.py`)
+==============================================================================================
+Zero-dependency HTTP server (`ThreadingHTTPServer`) exposing OpenAI-compatible endpoints:
+- `GET  /health`                -> Server & model status
+- `GET  /v1/models`             -> Lists `marsnext/Phobos-1`
+- `POST /v1/chat/completions`   -> Supports standard JSON (`stream: false`) and SSE (`stream: true`)
+- `POST /v1/completions`        -> Legacy prompt completion endpoint
+
+Compatible with the official `openai` Python/TS SDK, `curl`, OpenWebUI, LangChain, and Cursor.
+Strictly <= 220 lines (`RULE[user_global]`).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+import uuid
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+from typing import Any, Dict, List
+
+import torch
+
+from chat_cli import (
+    load_phobos_engine,
+    resolve_default_checkpoint,
+    resolve_default_tokenizer,
+    sample_next_token,
+)
+
+MODEL_ID = "marsnext/Phobos-1"
+_ENGINE: Dict[str, Any] = {}
+
+
+def _format_messages(messages: List[Dict[str, str]], force_think: bool = False) -> str:
+    """Formats OpenAI chat messages into Phobos-1 `<|user|>` / `<|assistant|>` sequence."""
+    parts: List[str] = ["<|bos|>"]
+    sys_prompt = ""
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = str(msg.get("content", "")).strip()
+        if role == "system":
+            sys_prompt = content
+        elif role == "user":
+            if sys_prompt:
+                content = f"{sys_prompt}\n\n{content}"
+                sys_prompt = ""
+            parts.append(f"<|user|>\n{content}\n")
+        elif role == "assistant":
+            parts.append(f"<|assistant|>\n{content}<|eos|>\n")
+    parts.append("<|assistant|>\n" + ("<think>\n" if force_think else ""))
+    return "".join(parts)
+
+
+class PhobosOpenAIHandler(BaseHTTPRequestHandler):
+    """Handles OpenAI REST requests for `marsnext/Phobos-1`."""
+
+    def _send_json(self, code: int, payload: Dict[str, Any]) -> None:
+        raw = json.dumps(payload, ensure_ascii=False).encode("utf-8")
+        self.send_response(code)
+        self.send_header("Content-Type", "application/json; charset=utf-8")
+        self.send_header("Access-Control-Allow-Origin", "*")
+        self.send_header("Content-Length", str(len(raw)))
+        self.end_headers()
+        self.wfile.write(raw)
+
+    def do_OPTIONS(self) -> None:  # noqa: N802
+        self.send_response(204)
+        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")
+        self.end_headers()
+
+    def do_GET(self) -> None:  # noqa: N802
+        if self.path.startswith("/health"):
+            self._send_json(200, {"status": "ok", "model": MODEL_ID, "device": str(_ENGINE["device"])})
+        elif self.path.startswith("/v1/models"):
+            self._send_json(
+                200,
+                {
+                    "object": "list",
+                    "data": [
+                        {
+                            "id": MODEL_ID,
+                            "object": "model",
+                            "created": int(time.time()),
+                            "owned_by": "marsnext",
+                        }
+                    ],
+                },
+            )
+        else:
+            self._send_json(404, {"error": {"message": f"Unknown endpoint {self.path}"}})
+
+    def do_POST(self) -> None:  # noqa: N802
+        if not (self.path.startswith("/v1/chat/completions") or self.path.startswith("/v1/completions")):
+            self._send_json(404, {"error": {"message": f"Unknown endpoint {self.path}"}})
+            return
+
+        length = int(self.headers.get("Content-Length", "0"))
+        body = json.loads(self.rfile.read(length).decode("utf-8") or "{}")
+
+        max_tokens = int(body.get("max_tokens", 256))
+        temperature = float(body.get("temperature", 0.0))
+        top_p = float(body.get("top_p", 0.9))
+        rep_penalty = float(body.get("repetition_penalty", 1.15))
+        stream = bool(body.get("stream", False))
+        force_think = bool(body.get("think", False))
+
+        if "messages" in body:
+            formatted = _format_messages(body["messages"], force_think=force_think)
+        else:
+            prompt_str = str(body.get("prompt", "")).strip()
+            formatted = f"<|bos|><|user|>\n{prompt_str}\n<|assistant|>\n" + ("<think>\n" if force_think else "")
+
+        model = _ENGINE["model"]
+        tok = _ENGINE["tok"]
+        device = _ENGINE["device"]
+
+        input_ids = tok.encode(formatted, max_len=1536)
+        ids = torch.tensor([input_ids], dtype=torch.long, device=device)
+        gen_ids: List[int] = tok.encode("<think>\n", max_len=16) if force_think else []
+        req_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+        created = int(time.time())
+
+        if stream:
+            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("Access-Control-Allow-Origin", "*")
+            self.end_headers()
+            prev_text = ""
+            with torch.no_grad():
+                for _ in range(max_tokens):
+                    logits = model(ids[:, -2048:])["logits"][0, -1, :].float()
+                    next_id = sample_next_token(logits, ids[0].tolist(), temperature, top_p, rep_penalty)
+                    if next_id == tok.EOS_ID:
+                        break
+                    gen_ids.append(next_id)
+                    ids = torch.cat([ids, torch.tensor([[next_id]], dtype=torch.long, device=device)], dim=1)
+                    cur_text = tok.decode(gen_ids, skip_special_tokens=False)
+                    delta = cur_text[len(prev_text):]
+                    if delta:
+                        prev_text = cur_text
+                        chunk = {
+                            "id": req_id,
+                            "object": "chat.completion.chunk",
+                            "created": created,
+                            "model": MODEL_ID,
+                            "choices": [{"index": 0, "delta": {"content": delta}, "finish_reason": None}],
+                        }
+                        self.wfile.write(f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n".encode("utf-8"))
+                        self.wfile.flush()
+            done_chunk = {
+                "id": req_id,
+                "object": "chat.completion.chunk",
+                "created": created,
+                "model": MODEL_ID,
+                "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
+            }
+            self.wfile.write(f"data: {json.dumps(done_chunk)}\n\ndata: [DONE]\n\n".encode("utf-8"))
+            self.wfile.flush()
+            return
+
+        with torch.no_grad():
+            for _ in range(max_tokens):
+                logits = model(ids[:, -2048:])["logits"][0, -1, :].float()
+                next_id = sample_next_token(logits, ids[0].tolist(), temperature, top_p, rep_penalty)
+                if next_id == tok.EOS_ID:
+                    break
+                gen_ids.append(next_id)
+                ids = torch.cat([ids, torch.tensor([[next_id]], dtype=torch.long, device=device)], dim=1)
+
+        reply_text = tok.decode(gen_ids, skip_special_tokens=False)
+        self._send_json(
+            200,
+            {
+                "id": req_id,
+                "object": "chat.completion",
+                "created": created,
+                "model": MODEL_ID,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": reply_text},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": len(input_ids),
+                    "completion_tokens": len(gen_ids),
+                    "total_tokens": len(input_ids) + len(gen_ids),
+                },
+            },
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="MarsNext Phobos-1 OpenAI-Compatible API Server")
+    parser.add_argument("--ckpt", "--weights", dest="ckpt", type=Path, default=None, help="Path to Phobos-1 .pt checkpoint")
+    parser.add_argument("--tokenizer", type=Path, default=None, help="Path to phobos_bpe_32k.json")
+    parser.add_argument("--device", type=str, default=None, help="Device: cuda, mps, or cpu")
+    parser.add_argument("--host", type=str, default="0.0.0.0", help="Bind host")
+    parser.add_argument("--port", type=int, default=8000, help="Bind port")
+    args = parser.parse_args()
+
+    ckpt_path = args.ckpt or resolve_default_checkpoint()
+    tok_path = args.tokenizer or resolve_default_tokenizer()
+    model, tok, device = load_phobos_engine(ckpt_path, tok_path, args.device)
+    _ENGINE["model"] = model
+    _ENGINE["tok"] = tok
+    _ENGINE["device"] = device
+
+    server = ThreadingHTTPServer((args.host, args.port), PhobosOpenAIHandler)
+    print(f"[MarsNext Phobos-1 OpenAI Server] Listening on http://{args.host}:{args.port}/v1 (device={device})", flush=True)
+    server.serve_forever()
+
+
+if __name__ == "__main__":
+    main()