diff --git a/docs/DETAILS.md b/docs/DETAILS.md index e2565af6f98..9a4d1215649 100644 --- a/docs/DETAILS.md +++ b/docs/DETAILS.md @@ -406,6 +406,7 @@ The server listens on `http://127.0.0.1:8080` (`--port` in setup, or edit the ru | API | Endpoint | | --- | --- | | OpenAI Chat Completions (stream and non-stream, tools) | `POST /v1/chat/completions` | +| OpenAI text completions (stream and non-stream, raw prompt) | `POST /v1/completions` | | Anthropic Messages (stream and non-stream, tools) | `POST /v1/messages` | | OpenAI Responses (stream and non-stream, tools; stateless, [below](#the-responses-api-and-codex-cli)) | `POST /v1/responses` | | Model list and health | `GET /v1/models`, `GET /models`, `GET /health` | @@ -549,6 +550,15 @@ applies at the next start. Network, key, MCP and program keys are not editable t `--parallel`, only a `max_tokens` 1 request gets them: it is the one a batch slot decodes through the single-token path, which writes an LP line. A longer reply there is a 400, since the slot's `BT` lines carry none -- entries for a prefix only, which a client cannot tell from the model's own answer. +- **logprobs on `/v1/completions` (vLLM's shape).** `"logprobs": N` (0..20, or `true` with `top_logprobs`) adds + `choices[0].logprobs` with `tokens`, `token_logprobs`, `top_logprobs` (one `{token: logprob}` per token) and + `text_offset`. Every generated token is listed, the stop token too, so a one-token read whose most likely token is + the stop still has its distribution. `"allowed_token_ids"` lists those ids with their full-vocabulary logprobs, not + renormalised (unlike chat), so a caller can tell how much of the mass its labels hold; + `"return_tokens_as_token_ids": true` keys every token as `token_id:`. The prompt may also be a list of token + ids. This is what a decision client reads from `vllm serve`: a typed choice from one next-token distribution + (AnyJev's `StrataBackend`). The rules are the chat ones (no cap of its own on `max_tokens`; with `--parallel`, only + `max_tokens` 1), and one of its own: logprobs are not streamed on this route, `stream: true` with them is a 400. - `STRATA_SPEC_PROB=1` with the default `--spec-min-p` gate: the gate tests the token the draft head drew from its distribution q. A draft that fails the gate still counts as a proposal. The verifier keeps q's mass on the tokens that pass the gate (not renormalised; call it mu), accepts an offered token D with probability min(1, p(D)/mu(D)), and draws the replacement diff --git a/serve/frontend.py b/serve/frontend.py index 63e319b82c2..f48a4275371 100644 --- a/serve/frontend.py +++ b/serve/frontend.py @@ -65,6 +65,18 @@ def render(self, messages: list[dict], tools: list[dict] | None = None, add_gene return self.template.render(messages=messages, tools=tools, add_generation_prompt=add_generation_prompt, **kwargs) + def gen_prompt_tail(self, enable_thinking: bool) -> str: + """What add_generation_prompt adds to this template: the generation prompt's tail. + + The raw /v1/completions prompt is passed through untouched; this is what a chat render would have + appended when the client asks thinking off. The probe message keeps the render the template needs; + both renders get the same kwargs, so the difference is the generation prompt alone. + """ + probe = [{"role": "user", "content": "x"}] + base = self.render(probe, add_generation_prompt=False, enable_thinking=enable_thinking) + full = self.render(probe, add_generation_prompt=True, enable_thinking=enable_thinking) + return full[len(base):] + def _detect_caps(self) -> dict[str, bool]: """llama.cpp's capability names, checked at load time against this template and Strata's tool-call format. These are rendering hints, not a guarantee that the model will follow a request.""" diff --git a/serve/server.py b/serve/server.py index 8c5213b1f80..3d4b6064ed6 100644 --- a/serve/server.py +++ b/serve/server.py @@ -3840,6 +3840,28 @@ def batch_slots_on(engine) -> bool: return any(int(getattr(e, "batch", 0) or 0) for e in (getattr(engine, "engines", None) or [engine])) +def check_logprobs(req: dict, svc, max_req: int) -> None: + """The checks a logprobs request passes before anything is sent, on /v1/chat/completions and /v1/completions + alike. No cap of our own on the reply's length: max_tokens (or the rest of the context) already bounds it.""" + if top_logprobs(req) is None: + return + if batch_slots_on(svc.engine) and max_req != 1: + # --parallel/--batch: a slot's "BT " lines carry no LP line, and the batched loop never + # calls emit_lp_line (src/program/generate.cpp); an admission -- a max_tokens 1 request -- is read + # as GEN 1, which does write them. Past that one token a reply would hold entries for a prefix + # only, which a client cannot tell from the model's own answer, so the longer one is refused. + raise ValueError('logprobs is available with "parallel" (batch slots) only for max_tokens 1; ' + 'run the server without "--parallel" for longer replies') + n_vocab = len(getattr(svc.tok, "tokens", ())) or None + if n_vocab: + # emit_lp_line indexes the logits row with each named id as given, so an id past the vocabulary + # would read out of bounds: refused here, before anything is sent + bad = [i for i in allowed_ids(req) if i >= n_vocab] + if bad: + raise ValueError(f"allowed_token_ids: {bad[0]} is outside the vocabulary " + f"(0..{n_vocab - 1})") + + NO_LINE = object() # StrataEngine.generate: no line held after a token @@ -4846,7 +4868,7 @@ def scaled(v, unit, digits=0): # one request at a time (more wait their turn), or "parallel": N batch slots (#465) "concurrency": {"serving": max(1, int(getattr(self.engine, "batch", 0) or 0)), "requested": getattr(self, "requested_slots", max(1, int(getattr(self.engine, "batch", 0) or 0)))}, - "dialects": ["/v1/chat/completions", "/v1/messages", "/v1/responses"], + "dialects": ["/v1/chat/completions", "/v1/completions", "/v1/messages", "/v1/responses"], "vision": {"enabled": images, "available": images, "error": None}, "image_generation": self.images.status() if self.images is not None else {"enabled": False}, "activity": {"requests": totals["requests"] + running, @@ -6275,6 +6297,57 @@ def openai_collect(chunks) -> dict: return out +class CompletionLogprobs: + """The engine's LP lines as the legacy /v1/completions `logprobs` object (vLLM's shape): + `tokens`, `token_logprobs`, `top_logprobs` (one {token: logprob} per token) and `text_offset`. + + Unlike LogprobEntries every generated token is listed, the stop token too, so a one-token decision whose most + likely token is a stop token still has its distribution. `allowed_token_ids` reports the named ids' + full-vocabulary logprobs, not renormalised (vLLM also reports the unmasked distribution), so a caller can tell + how much of the mass its labels hold. `return_tokens_as_token_ids` keys every token as "token_id:", which + keeps two ids that decode to the same text apart.""" + + def __init__(self, svc, req, owner): + self.tok, self.stop_ids, self.owner = svc.tok, svc.stop_ids, owner + self.allowed = allowed_ids(req) + self.as_ids = req.get("return_tokens_as_token_ids") is True + self.seen = [] # this request's LPLines, in order (one per generate pass) + + def watch(self, eng): + # a batch admission keeps the request's lines on its cancel Event, as LogprobEntries.take reads them + own = getattr(self.owner, "_strata_lp_lines", None) + cur = own if own is not None else getattr(eng, "lp_lines", None) + if getattr(cur, "owner", None) is self.owner and not any(cur is s for s in self.seen): + self.seen.append(cur) + + def text(self, t) -> str: + raw = self.tok.token_bytes(t) if hasattr(self.tok, "token_bytes") else self.tok.decode([t]).encode("utf-8") + return raw.decode("utf-8", "replace") + + def name(self, t) -> str: + return f"token_id:{t}" if self.as_ids else self.text(t) + + def build(self, offset: int) -> dict: + out = {"tokens": [], "token_logprobs": [], "top_logprobs": [], "text_offset": []} + for lines in self.seen: + while lines: + t, line = lines.popleft() + f = line.split() + try: # "LP nan": no head logits for that window row + lp = float(f[1]) + alts = [(int(i), float(v)) for i, _, v in (x.partition(":") for x in f[2:])] + except (IndexError, ValueError): + lp, alts = None, [] + if self.allowed and t in self.allowed and lp is not None: + alts = [(t, lp)] + alts # the line never repeats the chosen id among the named ones + out["tokens"].append(self.name(t)) + out["token_logprobs"].append(lp if lp is not None and math.isfinite(lp) else None) + out["top_logprobs"].append({self.name(i): v for i, v in alts if math.isfinite(v)}) + out["text_offset"].append(offset) + offset += 0 if t in self.stop_ids else len(self.text(t)) + return out + + def logprob_span(text: str, part: str, entries: list): """The entries whose tokens spell `part`, a piece of `text`, or None when the tokens do not spell `text` or no token boundary falls at both ends of `part`.""" @@ -7274,7 +7347,7 @@ def _post(self): except EngineDied as e: self._json(503, {"error": {"type": "server_error", "message": str(e)}}) return - if path in ("/v1/chat/completions", "/v1/messages", "/v1/responses"): + if path in ("/v1/chat/completions", "/v1/completions", "/v1/messages", "/v1/responses"): self.record = svc.begin_request(path, req) if path == "/v1/responses": self._responses(req) @@ -7289,6 +7362,8 @@ def _post(self): self._json(code, body) elif path == "/v1/chat/completions": self._openai(req) + elif path == "/v1/completions": + self._completions(req) elif path == "/v1/messages": self._anthropic(req) elif path == "/v1/messages/count_tokens": @@ -7303,7 +7378,7 @@ def _post(self): else {"error": {"type": "server_error", "message": str(e)}}) except ValueError as e: print(f"[strata] 400 invalid request: {e}", flush=True) - if path == "/v1/responses": + if path in ("/v1/responses", "/v1/completions"): self._json(400, responses_error_body(str(e))) else: self._json(400, {"error": {"type": "invalid_request_error", "message": str(e)}}) @@ -7334,9 +7409,19 @@ def _post(self): "message": f"the request is malformed ({type(e).__name__}: {e})"}} self._json(400, responses_error_body(body["error"]["message"]) if path == "/v1/responses" else body) except EngineDied as e: # before the answer started (not streamed) - self._json(503, {"error": {"type": "server_error", "message": f"{e}; the next request restarts it"}}) + body = {"error": {"type": "server_error", "message": f"{e}; the next request restarts it"}} + if getattr(self, "answer_started", False): + self._note(error=body["error"]) + self.wfile.write(b"data: " + json.dumps(body).encode() + b"\n\ndata: [DONE]\n\n") + else: + self._json(503, body) except EngineStuck as e: # an unload or restart that could not end the engine - self._json(503, {"error": {"type": "server_error", "message": str(e)}}) + body = {"error": {"type": "server_error", "message": str(e)}} + if getattr(self, "answer_started", False): + self._note(error=body["error"]) + self.wfile.write(b"data: " + json.dumps(body).encode() + b"\n\ndata: [DONE]\n\n") + else: + self._json(503, body) except ClientGone: # #1999: the client left while the request was prepared self._note(outcome="disconnected") @@ -7563,6 +7648,9 @@ def _captured(self, items, api): delta = item["choices"][0]["delta"] content, reasoning = delta.get("content", ""), delta.get("reasoning_content", "") usage, timings = item.get("usage"), item.get("timings") + elif api == "completions": + content, reasoning = item["choices"][0]["text"], "" + usage, timings = item.get("usage"), None elif api == "responses": kind = item["type"] content = item["delta"] if kind == "response.output_text.delta" else "" @@ -7605,21 +7693,7 @@ def _openai(self, req): # max_tokens (by the rest of the context when it is not sent), which is the only bound OpenAI and vLLM # put on it either. A reply that is not streamed does hold every entry (its token, bytes and up to 20 # alternatives) until it ends, so a long one holds a lot: streaming hands each entry out with its chunk. - if top_logprobs(req) is not None and batch_slots_on(svc.engine) and max_req != 1: - # --parallel/--batch: a slot's "BT " lines carry no LP line, and the batched loop never - # calls emit_lp_line (src/program/generate.cpp); an admission -- a max_tokens 1 request -- is read - # as GEN 1, which does write them. Past that one token a reply would hold entries for a prefix - # only, which a client cannot tell from the model's own answer, so the longer one is refused. - raise ValueError('logprobs is available with "parallel" (batch slots) only for max_tokens 1; ' - 'run the server without "--parallel" for longer replies') - n_vocab = len(getattr(svc.tok, "tokens", ())) or None - if top_logprobs(req) is not None and n_vocab: - # emit_lp_line indexes the logits row with each named id as given, so an id past the vocabulary - # would read out of bounds: refused here, before anything is sent - bad = [i for i in allowed_ids(req) if i >= n_vocab] - if bad: - raise ValueError(f"allowed_token_ids: {bad[0]} is outside the vocabulary " - f"(0..{n_vocab - 1})") + check_logprobs(req, svc, max_req) bias_api.validate_request(req, svc.engine, len(getattr(svc.tok, "tokens", ())) or None) use_mcp = req.get("strata_mcp") is True and svc.mcp is not None # the web app's opt-in (serve/mcp.py) own = {t.get("name") for t in tools or [] if isinstance(t, dict)} # #592: a second line of defence @@ -7660,22 +7734,183 @@ def _openai(self, req): self._note(outcome="disconnected") svc.note_client_disconnect(cancel) # client went away: stop the engine chunks.close() - except EngineDied as e: # mid-stream: say so, then end the stream properly + except (EngineDied, EngineStuck) as e: err = {"error": {"type": "server_error", "message": f"{e}; the next request restarts it"}} self._note(error=err["error"]) - self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + try: + self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") except QueueWaitExpired as e: err = {"error": {"type": "rate_limit_error", "code": "queue_wait_timeout", "message": str(e)}} self._note(error=err["error"]) - self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + try: + self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") except StructuredOutputError as e: err = {"error": {"type": "structured_output_failed", "code": "structured_output_failed", "message": str(e)}} self._note(error=err["error"]) - self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + try: + self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") + except ValueError as e: + err = {"error": {"type": "server_error", "message": str(e)}} + self._note(error=err["error"]) + try: + self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") + + def _completions(self, req): + """OpenAI legacy text completions over a caller-tokenized prompt (no chat template).""" + req = svc.with_shared(req, "openai") + prompt = req.get("prompt") + n_vocab = len(getattr(svc.tok, "tokens", ())) or None + # a string, or one prompt already tokenized (a list of token ids, as vLLM takes it) + given_ids = isinstance(prompt, list) and bool(prompt) and \ + all(isinstance(i, int) and not isinstance(i, bool) and i >= 0 for i in prompt) + if not isinstance(prompt, str) and not given_ids: + raise ValueError("prompt must be a string or a list of token ids (one prompt per request)") + if given_ids and n_vocab and max(prompt) >= n_vocab: + raise ValueError(f"prompt: token id {max(prompt)} is outside the vocabulary (0..{n_vocab - 1})") + for key in ("n", "best_of"): + value = req.get(key, 1) + if value != 1: + raise ValueError(f"{key} must be 1 (the server generates one sequence at a time)") + if req.get("suffix") not in (None, ""): + raise ValueError("suffix is not supported") + if req.get("logprobs") not in (None, False) and top_logprobs(req) is None: + raise ValueError("logprobs must be a whole number 0..20 (or true with top_logprobs)") + if top_logprobs(req) is not None and req.get("stream"): + raise ValueError("logprobs on /v1/completions are not supported with stream: true") + stream_options = req.get("stream_options") + if stream_options is not None: + if not isinstance(stream_options, dict) or set(stream_options) - {"include_usage"}: + raise ValueError("stream_options supports only include_usage") + if not req.get("stream"): + raise ValueError("stream_options are only supported with stream: true") + _, _, template_kwargs = openai_to_messages({**req, "messages": []}) + thinking = template_kwargs.get("enable_thinking", True) is not False + if req.get("enable_thinking") is False or req.get("thinking") is False: + thinking = False + thinking_cfg = req.get("thinking") + if isinstance(thinking_cfg, dict) and thinking_cfg.get("type") == "disabled": + thinking = False + bias_api.normalize(req.get("logit_bias"), len(getattr(svc.tok, "tokens", ())) or None) + ids = list(prompt) if given_ids else svc.tok.encode(prompt, parse_special=True) + if given_ids: + prompt = svc.tok.decode(ids) # for echo and text_offset + if thinking is False: # passthrough stays untouched when thinking is on; asked off, + tail = svc.tok.encode(svc.template.gen_prompt_tail(False), parse_special=True) # the prompt gets the tail + tail_on = svc.tok.encode(svc.template.gen_prompt_tail(True), parse_special=True) # a chat render adds + if tail and ids[-len(tail):] != tail and (not tail_on or ids[-len(tail_on):] != tail_on): + ids = ids + tail + try: + max_new = int(req.get("max_tokens") or 0) + except (TypeError, ValueError): + raise ValueError("max_tokens must be a whole number") + check_logprobs(req, svc, max_new) + bias_api.validate_request(req, svc.engine, len(getattr(svc.tok, "tokens", ())) or None) + svc.reasoning_budget(req) + stop_strings(req) + check_request_sampling(req) + cancel = threading.Event() + self._watch_client(cancel) + lp = CompletionLogprobs(svc, req, cancel) if top_logprobs(req) is not None else None + svc.load() + model = svc.model_for(req) + cid, created = "cmpl-" + uuid.uuid4().hex[:24], int(time.time()) + max_new = max_new if max_new > 0 else None + ctx = svc.engine.max_context + if ctx <= 0: + ctx = getattr(svc.engine, "known_ctx", 0) + if ctx <= 0: + raise EngineStarting("the engine is starting (a minute or two); try again shortly") + room = ctx - CTX_SLACK - len(ids) + if room < 1: + raise ValueError("request exceeds the context window: prompt leaves no room to answer") + if max_new is None: + max_new = room + elif max_new > room: + if not svc.fit_max_tokens: + raise ValueError(f"request exceeds the context window: prompt + max_tokens exceeds context " + f"(at most {room} tokens here)") + max_new = room + + def chunks(): + try: + if req.get("echo"): + yield {"id": cid, "object": "text_completion", "created": created, "model": model, + "choices": [{"text": prompt, "index": 0, "logprobs": None, "finish_reason": None}]} + for kind, value in svc.run(ids, thinking, None, max_new, req, cancel): + if lp is not None: + lp.watch(svc.engine) + if kind == "ping": + yield None + elif kind == "event": + if value.kind in ("content", "reasoning") and value.text: + yield {"id": cid, "object": "text_completion", "created": created, "model": model, + "choices": [{"text": value.text, "index": 0, "logprobs": None, + "finish_reason": None}]} + elif kind == "done": + completion_tokens = value["completion_tokens"] + yield {"id": cid, "object": "text_completion", "created": created, "model": model, + "choices": [{"text": "", "index": 0, "logprobs": None, + "finish_reason": {"cancel": "stop"}.get(value["finish"], value["finish"])}], + "usage": {"prompt_tokens": len(ids), "completion_tokens": completion_tokens, + "total_tokens": len(ids) + completion_tokens}} + finally: + cancel.set() if self.watch_done is None else None + + stream_chunks = self._capture(chunks(), "completions") + stream = bool(req.get("stream")) + if not stream: + text = [] + last = None + for item in stream_chunks: + if item is not None: + if "usage" in item: + last = item + else: + text.append(item["choices"][0]["text"]) + result = last or {"id": cid, "object": "text_completion", "created": created, "model": model, + "choices": [{"index": 0, "text": "", "logprobs": None, "finish_reason": "length"}], + "usage": {"prompt_tokens": len(ids), "completion_tokens": 0, "total_tokens": len(ids)}} + result["choices"][0]["text"] = "".join(text) + if lp is not None: + lp.watch(svc.engine) + result["choices"][0]["logprobs"] = lp.build(len(prompt) if req.get("echo") else 0) + self._json(200, result) + return + self._sse() + try: + for item in stream_chunks: + if item is None: + self.wfile.write(b": keep-alive\n\n") + else: + self.wfile.write(b"data: " + json.dumps(item, ensure_ascii=False).encode() + b"\n\n") + self.wfile.flush() + self.wfile.write(b"data: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") + cancel.set() + stream_chunks.close() + except (EngineDied, EngineStuck) as e: # mid-stream: say so, then end the stream properly + err = {"error": {"type": "server_error", "message": f"{e}; the next request restarts it"}} + self._note(error=err["error"]) + try: + self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") except ValueError as e: # the engine's ERR after the stream started: the err = {"error": {"type": "server_error", "message": str(e)}} # headers are sent, so no 400 now self._note(error=err["error"]) - self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + try: + self.wfile.write(b"data: " + json.dumps(err).encode() + b"\n\ndata: [DONE]\n\n") + except OSError: + self._note(outcome="disconnected") def _responses_ready(self, req, events): """Send a finished Responses result that did not run on the engine.""" diff --git a/serve/test_logprobs.py b/serve/test_logprobs.py index a21237db782..e0ec7b6bd5e 100644 --- a/serve/test_logprobs.py +++ b/serve/test_logprobs.py @@ -522,6 +522,117 @@ def test_slots_behind_replicas(self): self.assertTrue(batch_slots_on(SimpleNamespace(batch=4))) +class Completions(unittest.TestCase): + """Legacy /v1/completions with logprobs: vLLM's `logprobs` object, the shape a decision client + (a typed choice read from one next-token distribution) already reads from `vllm serve`.""" + + def setUp(self): + self.tok = Tokenizer() + self.engine = LpEngine(self.tok, "Six.") + self.svc = Service(self.engine, self.tok, TEMPLATE) + self.httpd = serve(self.svc, port=0) + self.base = f"http://127.0.0.1:{self.httpd.server_address[1]}" + + def tearDown(self): + self.httpd.shutdown() + self.httpd.server_close() + + def post(self, body): + body = {"model": "m", "prompt": "Two times three is", "max_tokens": 8, **body} + req = urllib.request.Request(self.base + "/v1/completions", json.dumps(body).encode(), + {"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=30) as r: + return json.loads(r.read().decode()) + + def refused(self, body): + with self.assertRaises(urllib.error.HTTPError) as cm: + self.post(body) + self.assertEqual(cm.exception.code, 400) + return cm.exception.read().decode() + + def test_the_legacy_shape(self): + choice = self.post({"logprobs": 4})["choices"][0] + lp = choice["logprobs"] + self.assertEqual(choice["text"], "Six.") + # every generated token, the stop token too: a one-token read whose top token is the stop still has its row + self.assertEqual(lp["tokens"], ["S", "i", "x", ".", "<|im_end|>"]) + self.assertEqual(lp["token_logprobs"], [-0.01, -0.02, -0.03, -0.04, -0.05]) + self.assertEqual(lp["top_logprobs"][0], {"A": -0.5, "H": -1.25, "X": -3.0, "[": -4.2}) + self.assertEqual(lp["text_offset"], [0, 1, 2, 3, 4]) + self.assertEqual(self.engine.sampling.get("logprobs"), 4) # logprobs=4 reached the engine + + def test_the_chat_form_of_the_field(self): + lp = self.post({"logprobs": True, "top_logprobs": 2})["choices"][0]["logprobs"] + self.assertEqual("".join(lp["tokens"][:-1]), "Six.") + + def test_echo_offsets_start_after_the_prompt(self): + lp = self.post({"logprobs": 1, "echo": True})["choices"][0]["logprobs"] + self.assertEqual(lp["text_offset"][0], len("Two times three is")) + + def test_allowed_token_ids_are_full_vocabulary_values(self): + # not renormalised (unlike chat): a caller reading a choice can tell how much mass its labels hold + self.engine = self.svc.engine = NamedEngine(self.tok, "A") + lp = self.post({"logprobs": 1, "max_tokens": 1, "allowed_token_ids": [65, 66, 67]})["choices"][0]["logprobs"] + self.assertEqual(lp["tokens"], ["A"]) + self.assertEqual(lp["top_logprobs"][0], {"A": -0.01, "B": -0.7, "C": -0.71}) # as the engine printed them + + def test_tokens_as_token_ids(self): + lp = self.post({"logprobs": 2, "max_tokens": 1, "return_tokens_as_token_ids": True})["choices"][0]["logprobs"] + self.assertEqual(lp["tokens"], ["token_id:83"]) + self.assertEqual(lp["top_logprobs"][0], {"token_id:65": -0.5, "token_id:72": -1.25, + "token_id:88": -3.0, "token_id:91": -4.2}) + + def test_a_tokenized_prompt(self): + out = self.post({"prompt": list(b"Two times three is"), "logprobs": 1, "echo": True}) + self.assertTrue(out["choices"][0]["text"].startswith("Two times three is")) + self.tok.tokens = [str(i) for i in range(300)] + self.assertIn("vocabulary", self.refused({"prompt": [300]})) + self.refused({"prompt": [[1, 2], [3]]}) # one prompt per request + + def test_not_asked_no_logprobs(self): + self.assertIsNone(self.post({})["choices"][0]["logprobs"]) + self.assertIsNone(self.post({"logprobs": False})["choices"][0]["logprobs"]) # as chat: not asked + + def test_a_request_reads_only_its_own_lines(self): + # as Api's: A's last "." is held by its stop string and sent after A has left the queue, when B is + # already generating on the same engine; each reply's tokens are its own + first_out = threading.Event() + engine = NextEngine(self.tok, ["Six.", "Bonjour!"], first_out) + svc = Service(engine, self.tok, TEMPLATE) + svc.fifo = HandOver(engine) + httpd = serve(svc, port=0) + got = {} + try: + self.base = f"http://127.0.0.1:{httpd.server_address[1]}" + a = threading.Thread(target=lambda: got.update(a=self.post({"logprobs": 1, "stop": [".!"]}))) + b = threading.Thread(target=lambda: got.update(b=self.post({"logprobs": 1}))) + a.start() + self.assertTrue(engine.first_in.wait(10)) + b.start() + a.join(30) + first_out.set() + b.join(30) + finally: + first_out.set() + httpd.shutdown() + httpd.server_close() + for key, text in (("a", "Six."), ("b", "Bonjour!")): + choice = got[key]["choices"][0] + self.assertEqual(choice["text"], text) + self.assertEqual("".join(t for t in choice["logprobs"]["tokens"] if t != "<|im_end|>"), text, key) + + def test_refusals(self): + self.assertIn("stream", self.refused({"logprobs": 1, "stream": True})) + self.refused({"logprobs": "yes"}) + self.engine.batch = 2 # batch slots: only a one-token read + self.assertIn("--parallel", self.refused({"logprobs": 1})) + + def test_no_length_cap(self): + # as on chat: max_tokens (or the rest of the context, 0) is the reply's only bound + lp = self.post({"logprobs": 1, "max_tokens": 0})["choices"][0]["logprobs"] + self.assertEqual("".join(lp["tokens"][:-1]), "Six.") + + class Span(unittest.TestCase): def entries(self, *toks): return [{"token": t} for t in toks] diff --git a/serve/test_monitor.py b/serve/test_monitor.py index 6c64b1115f8..bc07a56c7f5 100644 --- a/serve/test_monitor.py +++ b/serve/test_monitor.py @@ -108,6 +108,23 @@ def test_output_and_reasoning_for_both_dialects_and_stream_modes(self): else: self.assertNotIn("response", record) + def test_legacy_completions_capture_output_and_usage(self): + for stream in (False, True): + with self.subTest(stream=stream): + code, reply = self.request("/v1/completions", { + "prompt": "raw prompt", "max_tokens": 128, "enable_thinking": False, "stream": stream}) + self.assertEqual(code, 200) + if stream: + chunks = [json.loads(line[6:]) for line in reply.splitlines() if line.startswith("data: {")] + text = "".join(c["choices"][0]["text"] for c in chunks) + usage = chunks[-1]["usage"] + else: + text, usage = reply["choices"][0]["text"], reply["usage"] + record = self.latest() + self.assertTrue(text) + self.assertEqual(record["output"], text) + self.assertEqual(record["usage"], usage) + def test_error_and_engine_death_are_not_successes(self): # a request the server refuses is an error record: the strict mode (fit_max_tokens off) refuses a max_tokens # that overshoots the context, the default shortens it (serve/test_server.py MaxTokens) diff --git a/serve/test_server.py b/serve/test_server.py index 09c88024865..0c1d8413346 100644 --- a/serve/test_server.py +++ b/serve/test_server.py @@ -99,6 +99,128 @@ def test_count_tokens_is_the_prompt_messages_reads(self): self.assertEqual(s2, 200) self.assertEqual(b["input_tokens"], n_in) + def test_legacy_completions_raw_prompt_nonstream_echo_and_stream(self): + prompt = "raw <|im_start|>user" + status, body = self.post("/v1/completions", {"model": "m", "prompt": prompt, + "max_tokens": 16, "temperature": 0, + "enable_thinking": False}) + self.assertEqual(status, 200) + self.assertEqual(body["object"], "text_completion") + self.assertTrue(body["choices"][0]["text"]) + tail_len = len(self.svc.tok.encode(self.svc.template.gen_prompt_tail(False), parse_special=True)) + self.assertEqual(body["usage"]["prompt_tokens"], # asked off: the tail + len(self.svc.tok.encode(prompt, parse_special=True)) + tail_len) # a chat render adds + status, echoed = self.post("/v1/completions", {"model": "m", "prompt": prompt, "max_tokens": 8, + "temperature": 0, "enable_thinking": False, + "echo": True}) + self.assertEqual(status, 200) + self.assertTrue(echoed["choices"][0]["text"].startswith(prompt)) + req = urllib.request.Request(self.base + "/v1/completions", data=json.dumps({ + "model": "m", "prompt": "stream me", "max_tokens": 8, "temperature": 0, + "enable_thinking": False, "stream": True}).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=30) as response: + stream = response.read().decode() + self.assertIn('"object": "text_completion"', stream) + self.assertIn('"usage"', stream) + self.assertIn("data: [DONE]", stream) + + def test_legacy_completions_stream_text_matches_nonstream(self): + for echo, budget in ((False, 64), (True, 64), (False, 3000), (True, 3000)): + with self.subTest(echo=echo, budget=budget): + body = {"model": "m", "prompt": "raw prompt", "max_tokens": budget, + "temperature": 0, "enable_thinking": False, "echo": echo} + status, reply = self.post("/v1/completions", body) + self.assertEqual(status, 200) + req = urllib.request.Request(self.base + "/v1/completions", data=json.dumps({ + **body, "stream": True}).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=30) as response: + raw = response.read().decode() + chunks = [json.loads(line[6:]) for line in raw.splitlines() if line.startswith("data: {")] + self.assertEqual("".join(c["choices"][0]["text"] for c in chunks), reply["choices"][0]["text"]) + self.assertEqual(chunks[-1]["choices"][0]["text"], "") + self.assertEqual(chunks[-1]["choices"][0]["finish_reason"], reply["choices"][0]["finish_reason"]) + self.assertEqual(chunks[-1]["usage"], reply["usage"]) + if echo: + self.assertEqual(chunks[0]["choices"][0]["text"], body["prompt"]) + self.assertTrue(raw.endswith("data: [DONE]\n\n")) + + def test_legacy_completions_rejects_unsupported_shapes(self): + for body in ({"model": "m", "prompt": ["a"]}, {"model": "m", "prompt": "a", "n": 2}, + {"model": "m", "prompt": "a", "suffix": "b"}): + with self.subTest(body=body): + status, _ = self.post("/v1/completions", body) + self.assertEqual(status, 400) + + def test_legacy_completions_honors_disabled_thinking(self): + cases = ({"reasoning_effort": "none"}, {"chat_template_kwargs": {"enable_thinking": False}}, + {"enable_thinking": False}, {"thinking": {"type": "disabled"}}) + for options in cases: + seen = [] + original = self.svc.run + + def capture(ids, thinking, tools, max_new, sampling, cancel, force=None): + seen.append(thinking) + yield from original(ids, thinking, tools, max_new, sampling, cancel, force=force) + + with self.subTest(options=options), mock.patch.object(self.svc, "run", side_effect=capture): + status, _ = self.post("/v1/completions", {"model": "m", "prompt": "raw", "max_tokens": 8, + "temperature": 0, **options}) + self.assertEqual(status, 200) + self.assertEqual(seen, [False]) + + def test_legacy_completions_stream_options_accepted_when_asked(self): + # include_usage is the only field OpenAI defines for it; the stream's usage chunk is there anyway + req = urllib.request.Request(self.base + "/v1/completions", data=json.dumps({ + "model": "m", "prompt": "hi", "max_tokens": 8, "temperature": 0, "stream": True, + "stream_options": {"include_usage": True}}).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=30) as response: + stream = response.read().decode() + self.assertIn('"usage"', stream) + self.assertIn("data: [DONE]", stream) + # it belongs to the stream: a non-stream request carrying it, or one carrying a field the + # API does not define, stays a 400 + status, _ = self.post("/v1/completions", {"model": "m", "prompt": "hi", "max_tokens": 8, + "temperature": 0, "stream_options": {"include_usage": False}}) + self.assertEqual(status, 400) + status, _ = self.post("/v1/completions", {"model": "m", "prompt": "hi", "max_tokens": 8, + "temperature": 0, "stream": True, + "stream_options": {"include_thinking": True}}) + self.assertEqual(status, 400) + + def test_legacy_completions_thinking_off_appends_the_no_think_tail(self): + # the raw prompt passes through untouched; asked off, it gets the tail a chat render would add + # (and asked off when it already ends with either tail - thinking on or off - nothing is added) + tail_off_text = self.svc.template.gen_prompt_tail(False) + tail_on_text = self.svc.template.gen_prompt_tail(True) + encode = lambda text: list(self.svc.tok.encode(text, parse_special=True)) # noqa: E731 + tail_off, tail_on = encode(tail_off_text), encode(tail_on_text) + + def run_ids(body): + seen = [] + original = self.svc.run + + def capture(ids, thinking, tools, max_new, sampling, cancel, force=None): + seen.append(list(ids)) + yield from original(ids, thinking, tools, max_new, sampling, cancel, force=force) + + with mock.patch.object(self.svc, "run", side_effect=capture): + status, _ = self.post("/v1/completions", {"model": "m", "max_tokens": 8, + "temperature": 0, **body}) + self.assertEqual(status, 200) + return seen[0] + + self.assertEqual(run_ids({"prompt": "hi", "enable_thinking": False}), encode("hi") + tail_off) + self.assertEqual(run_ids({"prompt": "hi"}), encode("hi")) # not asked: untouched + self.assertEqual(run_ids({"prompt": "hi", "reasoning_effort": "none"}), encode("hi") + tail_off) + for tail in (tail_off_text, tail_on_text): # already ends with a tail: no second one + self.assertEqual(run_ids({"prompt": "hi" + tail, "enable_thinking": False}), encode("hi") + encode(tail)) + # Some templates add a tail only when thinking is disabled; an empty prompt still needs it. + with mock.patch.object(self.svc.template, "gen_prompt_tail", side_effect=lambda enabled: "" if enabled else tail_off_text): + for prompt in ("", "hi"): + self.assertEqual(run_ids({"prompt": prompt, "enable_thinking": False}), encode(prompt) + tail_off) + def test_request_line_parses_the_engine_summary(self): line = ("strata serve: prompt 1200 tokens = 1000 reused + 200 read in 50 ms (4000.0 tok/s), 30 generated in " "300 ms (100.0 tok/s), drafts accepted 20 of 28, 2 checkpoints") diff --git a/serve/test_structured.py b/serve/test_structured.py index f0e2c74aede..5a47de8f90f 100644 --- a/serve/test_structured.py +++ b/serve/test_structured.py @@ -95,10 +95,14 @@ def test_invalid_generation_is_an_error_and_never_streams_content(self): code, reply = self.chat(script) self.assertEqual(code, 502) self.assertEqual(reply["error"]["code"], "structured_output_failed") - _, raw = self.chat(script, stream=True) - self.assertIn('"structured_output_failed"', raw) + code, raw = self.chat(script, stream=True) + self.assertEqual(code, 200) + chunks = [json.loads(line[6:]) for line in raw.splitlines() if line.startswith("data: {")] + self.assertEqual(len(chunks), 1) + self.assertEqual(chunks[0]["error"]["code"], "structured_output_failed") + self.assertNotIn("HTTP/", raw) self.assertNotIn('"delta"', raw) - self.assertIn("data: [DONE]", raw) + self.assertTrue(raw.endswith("data: [DONE]\n\n")) def test_object_mode_strict_json_and_truncation(self): fmt = {"type": "json_object"} @@ -190,6 +194,11 @@ def test_without_jsonschema_an_object_is_still_required(self): code, reply = self.chat(script) self.assertEqual(code, 502, script) self.assertEqual(reply["error"]["code"], "structured_output_failed") + code, raw = self.chat(script, stream=True) + self.assertEqual(code, 200) + self.assertIn('data: {"error":', raw) + self.assertNotIn("HTTP/", raw) + self.assertTrue(raw.endswith("data: [DONE]\n\n")) self.assertEqual(self.chat('{"answer":4}', response_format={"type": "json_object"})[0], 200) def test_root_unions_of_object_shapes_are_accepted(self):