OpenAI-compatible API server 深度解析

42MODEL_ID = "deltafin-kimi-k3"
43MAX_TOKENS_CAP = int(os.environ.get("K3_SERVER_MAX_TOKENS", "0"))   # 0 = no cap
44RESPONSE_MEMO_ENTRIES = int(
45    os.environ.get("K3_RESPONSE_MEMO_ENTRIES", "32"))
46RESPONSE_MARKER = "<|open|>response<|sep|>"
47THINK_CLOSE = "<|close|>think<|sep|>"
48
49_lock = threading.Lock()
50_tok = None
51_layers = None
52_embed = None
53_memo = DeterministicResponseMemo(RESPONSE_MEMO_ENTRIES)
54
55
56def _boot():
57    global _tok, _layers, _embed
58    from transformers import AutoTokenizer
59    print("[serve] loading tokenizer + layer skeletons...", flush=True)
60    _tok = AutoTokenizer.from_pretrained(
61        os.path.join(kr.ROOT, "k3-meta"), trust_remote_code=True)
62    kr.check_expert_pool()
63    _layers = kr.build_layers()
64    _embed = kr.LazyEmbed()
65    print("[serve] ready", flush=True)

128    def do_POST(self):
129        try:
130            body = json.loads(self.rfile.read(int(self.headers.get("Content-Length", 0)) or 0) or b"{}")
131        except json.JSONDecodeError:
132            return self._err(400, "invalid JSON body")
133        chat = self.path in ("/v1/chat/completions", "/chat/completions")
134        comp = self.path in ("/v1/completions", "/completions")
135        if not (chat or comp):
136            return self._err(404, f"no route {self.path}")
137
138        if chat:
139            messages = body.get("messages")
140            if not messages:
141                return self._err(400, "messages required")
142            ids = _tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True)
143        else:
144            prompt = body.get("prompt")
145            if not isinstance(prompt, str):
146                return self._err(400, "prompt (string) required")
147            ids = _tok.encode(prompt)
148        # OpenAI semantics: omitted max_tokens means the model decides. Chat ends
149        # naturally at EOS; raw completions have no terminator, so only THEY get
150        # a default cap (256) — an explicit max_tokens is always honored.
151        req_max = body.get("max_tokens")
152        if req_max:
153            max_new = int(req_max)
154        else:
155            max_new = 1_000_000 if chat else 256
156        if MAX_TOKENS_CAP:
157            max_new = min(max_new, MAX_TOKENS_CAP)
158        stream = bool(body.get("stream"))
159        rid = ("chatcmpl-" if chat else "cmpl-") + uuid.uuid4().hex[:20]
160        created = int(time.time())
161
162        if not _lock.acquire(timeout=5):
163            return self._err(429, "a generation is already running (Deltafin serves one at a time)")
164        try:
165            mode = "chat" if chat else "completion"
166            cached = _memo.get(mode, ids, max_new)
167            if cached is not None:
168                print(f"[serve] deterministic response memo hit "
169                      f"({len(cached.token_ids)} tokens)", flush=True)
170            if stream:

170            if stream:
171                self.send_response(200)
172                self.send_header("Content-Type", "text/event-stream")
173                self.send_header("Cache-Control", "no-cache")
174                self.send_header("Connection", "close")
175                self.end_headers()
176
177                def sse(obj):
178                    self.wfile.write(b"data: " + json.dumps(obj).encode() + b"\n\n")
179                    self.wfile.flush()
180
181                def on_delta(delta):
182                    if chat:
183                        sse({"id": rid, "object": "chat.completion.chunk", "created": created,
184                             "model": MODEL_ID,
185                             "choices": [{"index": 0, "delta": {"content": delta},
186                                          "finish_reason": None}]})
187                    else:
188                        sse({"id": rid, "object": "text_completion", "created": created,
189                             "model": MODEL_ID,
190                             "choices": [{"index": 0, "text": delta, "finish_reason": None}]})
191
192                if cached is None:
193                    out, text, finish = _gen(
194                        ids, max_new, on_delta=on_delta)
195                    _memo.put(mode, ids, max_new, out, text, finish)
196                else:
197                    out = list(cached.token_ids)
198                    text = cached.text
199                    finish = cached.finish_reason
200                    if text:
201                        on_delta(text)
202                key = "delta" if chat else "text"
203                sse({"id": rid, "object": "chat.completion.chunk" if chat else "text_completion",
204                     "created": created, "model": MODEL_ID,
205                     "choices": [{"index": 0, key: {} if chat else "", "finish_reason": finish}]})
206                self.wfile.write(b"data: [DONE]\n\n")
207                self.wfile.flush()
208                return
209
210            if cached is None:
211                out, text, finish = _gen(ids, max_new)
212                _memo.put(mode, ids, max_new, out, text, finish)
213            else:
214                out = list(cached.token_ids)
215                text = cached.text
216                finish = cached.finish_reason
217            usage = {"prompt_tokens": len(ids), "completion_tokens": len(out),
218                     "total_tokens": len(ids) + len(out)}
219            if chat:
220                reasoning, content = _split_reasoning(text)
221                msg = {"role": "assistant", "content": content}
222                if reasoning:
223                    msg["reasoning_content"] = reasoning
224                self._json(200, {"id": rid, "object": "chat.completion", "created": created,
225                                 "model": MODEL_ID, "usage": usage,
226                                 "choices": [{"index": 0, "message": msg,
227                                              "finish_reason": finish}]})
228            else:
229                self._json(200, {"id": rid, "object": "text_completion", "created": created,
230                                 "model": MODEL_ID, "usage": usage,
231                                 "choices": [{"index": 0, "text": text,
232                                              "finish_reason": finish}]})