diff --git a/docs/server-config.md b/docs/server-config.md index 5f692ac..0703561 100644 --- a/docs/server-config.md +++ b/docs/server-config.md @@ -1319,6 +1319,7 @@ generation dialects. The rest: | `presence_penalty` | honored | honored | honored | | | `frequency_penalty` | honored | honored | honored | | | `stream_options` | honored | ignored | ignored | `include_usage` adds the final usage chunk (chat + `/v1/completions`) | +| `timings_per_token` | honored | ignored | ignored | streamed chat chunks carry `timings.predicted_n`, the exact cumulative output-token count (llama.cpp convention) | | `response_format` | honored | honored | honored | `json_schema` / `json_object`; unknown types are rejected (see below) | | `logprobs` | honored | ignored | ignored | chat-only; `/v1/completions` never returns logprobs | | `top_logprobs` | honored | ignored | ignored | capped by `TOP_LOGPROBS_K` (below) | diff --git a/gmlx/assistant_brain.py b/gmlx/assistant_brain.py index dd2fab9..75fe72c 100644 --- a/gmlx/assistant_brain.py +++ b/gmlx/assistant_brain.py @@ -158,6 +158,7 @@ def turn(self, user_text: str) -> Iterator[BrainEvent]: spoken: list = [] # all answer text this turn (memory) text_parts: list = [] # current round's uncommitted text stats: dict = {} + timings: dict | None = None # last server timings seen (gmlx) completed = False committed_tool_round = False # any assistant+tool round appended? try: @@ -182,6 +183,12 @@ def turn(self, user_text: str) -> Iterator[BrainEvent]: if "_finish" in delta: finish = delta["_finish"] continue + if "_timings" in delta: + timings = delta["_timings"] or {} + n = timings.get("predicted_n") + if n: + yield ("count", int(n)) + continue if delta.get("reasoning"): yield ("status", "thinking") continue @@ -269,6 +276,8 @@ def turn(self, user_text: str) -> Iterator[BrainEvent]: except Exception: # noqa: BLE001 - best-effort pass if completed: + if timings: + stats = {**stats, "timings": timings} yield ("done", stats) def _execute(self, call: dict) -> str: diff --git a/gmlx/chat.py b/gmlx/chat.py index 10d7f3f..e34964a 100644 --- a/gmlx/chat.py +++ b/gmlx/chat.py @@ -836,6 +836,10 @@ def _fmt_stat_line(stats: dict, ctx_used, ctx_max) -> str: if stats.get("prompt_tps"): p += f" @ {stats['prompt_tps']:.0f} tok/s" parts.append(p) + # Queue wait + prefill, as one number: everything before the first token + # arrived. The gen rate beside it is decode-only, so nothing is blended. + if stats.get("ttft_s", 0.0) >= 0.1: + parts.append(f"ttft {stats['ttft_s']:.1f}s") if stats.get("gen_tokens"): parts.append( f"gen {_fmt_k(stats['gen_tokens'])} tok @ {stats.get('gen_tps', 0.0):.1f} tok/s" @@ -1761,6 +1765,58 @@ def __exit__(self, *exc): termios.tcsetattr(self.fd, termios.TCSADRAIN, self._saved) +class _RateTicker: + """Windowed live tok/s for the streaming ticker, from exact cumulative + token counts only: local generation chunks carry ``generation_tokens``, + and gmlx servers stream the same count per chunk when the client asks + for ``timings_per_token`` (adapted into the same field). Chunks without + a count are ignored - no estimated rates. Returns a new status string + at most a few times a second, else None. + + A server turn can span several requests (the assistant tool loop), each + restarting its count at 0; a count below the last one seen rolls the + previous request's total into a base offset so the ticker stays + cumulative across rounds.""" + + WINDOW_S = 2.0 + PUSH_S = 0.3 + + def __init__(self, clock=time.monotonic): + self._clock = clock + self._window: list[tuple[float, float]] = [] + self._base = 0 + self._last_n = 0 + self._last_push = 0.0 + self._last_text = None + + def push(self, r) -> str | None: + n = int(getattr(r, "generation_tokens", 0) or 0) + if not n: + return None + if n < self._last_n: + self._base += self._last_n + self._last_n = n + total = float(self._base + n) + t = self._clock() + self._window.append((t, total)) + while self._window and t - self._window[0][0] > self.WINDOW_S: + self._window.pop(0) + if t - self._last_push < self.PUSH_S or len(self._window) < 2: + return None + dt = self._window[-1][0] - self._window[0][0] + if dt < 0.5: + return None + rate = (self._window[-1][1] - self._window[0][1]) / dt + text = f"{rate:.0f} tok/s" + if text == self._last_text: + self._last_push = t + return None + self._last_push = t + self._last_text = text + return text + + + def _stream_reply( chunks, state: ChatState, @@ -1795,6 +1851,7 @@ def _stream_reply( theme=theme, answer_sink=renderer.feed if renderer else None, ) + ticker = _RateTicker() if renderer is not None else None def _toggle() -> None: # Ctrl-O: collapse<->expand thinking, live for this reply and persisted as @@ -1833,6 +1890,10 @@ def _accept(text: str) -> str: if out: _accept(out) printer.tick() + if ticker is not None: + status = ticker.push(r) + if status is not None: + renderer.set_status(status) last = r if stopped or esc.pressed(): canceled = not stopped @@ -1858,6 +1919,9 @@ def _accept(text: str) -> str: "gen_tokens": int(getattr(last, "generation_tokens", 0) or 0), "gen_tps": float(getattr(last, "generation_tps", 0.0) or 0.0), } + ttft = float(getattr(last, "ttft_s", 0.0) or 0.0) + if ttft: + stats["ttft_s"] = ttft accepts = list(getattr(drafter, "accept_lens", None) or []) drafts = list(getattr(drafter, "draft_lens", None) or []) if drafts: @@ -2159,8 +2223,13 @@ def _setup_assistant(args): max_items=a.memory.max_items) # Usage chunks are gated on stream_options server-side; sampling knobs - # join this dict per turn (see _sync_assistant_extra). + # join this dict per turn (see _sync_assistant_extra). Per-chunk stream + # timings feed the live tok/s ticker, but only gmlx servers know the + # field (strict OpenAI backends reject unknown body params), so it rides + # only when the models probe saw gmlx entries. extra: dict = {"stream_options": {"include_usage": True}} + if caps.get("gmlx"): + extra["timings_per_token"] = True def seam(burl, *, model, messages, max_tokens, api_key=None, tools=None, timeout=600.0): @@ -2210,6 +2279,7 @@ def _assistant_reply(brain, user_text: str, state: ChatState) -> tuple[str, bool status_shown = [False] t0 = time.monotonic() + t_first = [None] # first generated-token event def _clear_status(): if status_shown[0]: @@ -2217,27 +2287,49 @@ def _clear_status(): sys.stdout.flush() status_shown[0] = False + def _mark_first(): + if t_first[0] is None: + t_first[0] = time.monotonic() + def chunks(): try: for kind, payload in brain.turn(user_text): if kind == "say": + _mark_first() _clear_status() yield SimpleNamespace(text=payload) elif kind == "status": + _mark_first() # thinking/tool events ride on tokens _clear_status() sys.stdout.write(f"[assistant] {payload}...") sys.stdout.flush() status_shown[0] = True + elif kind == "count": + # Exact cumulative output tokens for this request round + # (gmlx per-chunk stream timings) - feeds the tok/s ticker. + _mark_first() + yield SimpleNamespace(text="", generation_tokens=payload) elif kind == "done": u = payload or {} n = int(u.get("completion_tokens") or 0) - el = time.monotonic() - t0 + end = time.monotonic() + # ttft covers everything before the first token reached + # this client: queue wait plus prefill. + ttft = (t_first[0] - t0) if t_first[0] is not None else 0.0 + tm = u.get("timings") or {} + gen_tps = float(tm.get("predicted_per_second") or 0.0) + if not gen_tps: + # No server timings (foreign server): decode-only + # wall-clock rate, so the wait never dilutes it. + dec = end - t0 - ttft + gen_tps = (n / dec) if dec > 0 and n else 0.0 yield SimpleNamespace( text="", prompt_tokens=int(u.get("prompt_tokens") or 0), - prompt_tps=0.0, + prompt_tps=float(tm.get("prompt_per_second") or 0.0), generation_tokens=n, - generation_tps=(n / el) if el > 0 and n else 0.0) + generation_tps=gen_tps, + ttft_s=ttft) except TalkClientError as e: _clear_status() print(f"\n[chat] server error: {e}", file=sys.stderr) diff --git a/gmlx/render.py b/gmlx/render.py index 911172b..7ab7e84 100644 --- a/gmlx/render.py +++ b/gmlx/render.py @@ -402,9 +402,28 @@ def _term_size(): _CONTROL_CHARS = re.compile(r"[\x00-\x08\x0b-\x1f\x7f]") +def _find_last(haystack: list[str], needle: list[str]) -> int: + """Last index where ``needle`` appears as a contiguous run in + ``haystack``; -1 when absent or empty.""" + if not needle or len(needle) > len(haystack): + return -1 + for i in range(len(haystack) - len(needle), -1, -1): + if haystack[i : i + len(needle)] == needle: + return i + return -1 + + class StreamRenderer: """Answer-channel sink: repaint the current block in place, print completed - blocks permanently. Plain mode passes text straight through.""" + blocks permanently. Plain mode passes text straight through. + + Repaints are diff-aware: the painted live region is tracked line-by-line + and only rows that changed since the last paint are rewritten (a token + append usually costs one row, not a screenful), so a fast stream through + a slow transport (tmux, ssh, a recorder) does not backpressure the reader + loop. The paint interval also adapts to measured paint cost, so an + expensive backend (rich re-renders the whole block) degrades to fewer, + larger updates instead of falling behind the stream.""" def __init__( self, @@ -415,6 +434,7 @@ def __init__( size_fn=None, clock=None, min_repaint_s: float = 0.04, + max_repaint_s: float = 0.25, ): self.mode = mode self.theme = theme or resolve_theme(color=(mode != "plain")) @@ -422,7 +442,10 @@ def __init__( self._size = size_fn or _term_size self._clock = clock or time.monotonic self._min_repaint = min_repaint_s + self._max_repaint = max_repaint_s + self._interval = min_repaint_s self._buf = BlockBuffer() + self._screen: list[str] = [] # painted rows of the live region (for diff) self._painted = 0 # terminal lines the live block occupies self._committed = 0 # rendered lines of the live block already # scroll-committed: a block taller than the @@ -431,8 +454,16 @@ def __init__( # the last screenful self._frozen = False # resized mid-block: appends raw to block end self._raw_emitted = 0 # chars of the current block already raw-written + self._src_skip = 0 # fence body source lines dropped from live + # renders: once a fence block scrolls far + # past the viewport, its committed body + # lines are immutable, so re-rendering them + # every paint only burns time (rich re-parse + # is O(block)); the live render input stays + # viewport-sized instead self._last_paint = 0.0 self._width = None + self._status = None # live one-line ticker under the block self._backend = _RichBackend(self.theme) if mode == "rich" else None # -- public ------------------------------------------------------------------ @@ -448,6 +479,17 @@ def feed(self, text: str) -> None: self._finish_block(src) self._paint_live() + def set_status(self, text: str | None) -> None: + """Live one-line ticker under the streaming block (e.g. a tok/s + readout). It lives on the cursor's resting line below the painted + region, so it needs no diff or commit bookkeeping; block boundaries + and finalize clear it. Ignored in plain mode, while frozen, and + before the first paint of a block (no resting line yet).""" + self._status = text + if self.mode == "plain" or self._frozen or not self._painted: + return + self._paint_status() + def finalize(self) -> None: """End of reply: last unthrottled paint of the tail; cursor ends on a fresh line below the rendered output.""" @@ -461,9 +503,14 @@ def finalize(self) -> None: self._reset_block() return if tail.strip(): - self._repaint(self._render(tail)[self._committed :]) + self._repaint(self._render(self._trimmed(tail))[self._committed :]) + if self._status: + self._w("\x1b[2K") + self._status = None self._painted = 0 self._committed = 0 + self._src_skip = 0 + self._screen = [] # -- internals ---------------------------------------------------------------- @@ -473,6 +520,42 @@ def _render(self, src: str) -> list[str]: return self._backend.render(src, width) return _render_lite(src, width, self.theme) + def _trimmed(self, src: str) -> str: + """Apply the live trim to a full block source: opening fence line plus + the untrimmed body tail. Identity when no trim is active.""" + if not self._src_skip: + return src + lines = src.splitlines(keepends=True) + return lines[0] + "".join(lines[1 + self._src_skip :]) + + def _paint_source(self, budget: int) -> list[str]: + """Render the live block for painting. A fence block far taller than + the viewport gets its committed body dropped from the render input so + per-paint cost stays viewport-sized: the rebase relocates + ``_committed`` by finding the currently painted rows inside the + trimmed render, and skips the trim entirely on any mismatch (always + correct, only slower).""" + buf = self._buf + body = len(buf._lines) - 1 + if ( + buf._fence is not None + and len(self._screen) > 4 + and body - self._src_skip > 2 * budget + 16 + ): + new_skip = body - (budget + 16) + candidate = ( + buf._lines[0] + + "".join(buf._lines[1 + new_skip :]) + + buf._partial + ) + rendered = self._render(candidate) + i = _find_last(rendered, self._screen[:-1]) + if i >= 0: + self._src_skip = new_skip + self._committed = i + return rendered + return self._render(self._trimmed(buf.current)) + def _finish_block(self, src: str) -> None: if self._frozen: self._w(src[self._raw_emitted :]) @@ -480,11 +563,13 @@ def _finish_block(self, src: str) -> None: self._reset_block() return if src.strip(): - self._repaint(self._render(src)[self._committed :]) - self._w("\n") + self._repaint(self._render(self._trimmed(src))[self._committed :]) + self._w("\x1b[2K\n" if self._status else "\n") self._painted = 0 self._raw_emitted = 0 self._committed = 0 + self._src_skip = 0 + self._screen = [] def _paint_live(self) -> None: cur = self._buf.current @@ -495,7 +580,7 @@ def _paint_live(self) -> None: if not cur.strip(): return now = self._clock() - if now - self._last_paint < self._min_repaint: + if now - self._last_paint < self._interval: return self._last_paint = now size = self._size() @@ -504,8 +589,8 @@ def _paint_live(self) -> None: self._freeze(cur) return self._width = width - lines = self._render(cur)[self._committed :] budget = max(4, size.lines - 2) + lines = self._paint_source(budget)[self._committed :] if len(lines) >= budget: # Taller than the viewport: rows above the screen edge cannot be # repainted in place. Commit everything but the last screenful @@ -513,23 +598,52 @@ def _paint_live(self) -> None: # keep only the tail live. Rendered prefixes are append-stable # for fences, paragraphs and lists (greedy wrap), so committed # lines never need rewriting. - self._committed += len(lines) - (budget - 1) + slide = len(lines) - (budget - 1) + self._committed += slide + self._repaint(lines[slide:], slide=slide) + else: self._repaint(lines) - self._painted = budget - 1 - return - self._repaint(lines) - - def _repaint(self, lines: list[str]) -> None: + if self._status: + self._paint_status() + # Adapt the paint cadence to what a paint actually costs (render + + # write, which blocks when the transport is saturated): painting may + # use about a third of wall time, between the min and max. Non-paint + # work per chunk is sub-millisecond, so this cannot starve the reader + # loop. The interval runs from paint END, so a slow paint never eats + # the following quiet period. + end = self._clock() + self._interval = min( + self._max_repaint, max(self._min_repaint, (end - now) * 3.0) + ) + self._last_paint = end + + def _repaint(self, lines: list[str], slide: int = 0) -> None: + """Bring the painted live region to ``lines``, rewriting only rows + that differ from what is on screen. ``slide`` is how many top rows of + the previous region were scroll-committed by this paint; they stay on + screen as-is and drop out of the diff.""" + keep = self._screen[slide:] if slide else self._screen + lim = min(len(keep), len(lines)) + d = 0 + while d < lim and keep[d] == lines[d]: + d += 1 out = [] - if self._painted: - out.append(f"\x1b[{self._painted}A") - for ln in lines: + up = len(keep) - d + if up > 0: + out.append(f"\x1b[{up}A") + for ln in lines[d:]: out.append("\x1b[2K" + ln + "\n") - extra = self._painted - len(lines) + extra = len(keep) - len(lines) if extra > 0: - out.append("\x1b[2K\n" * extra) - out.append(f"\x1b[{extra}A") - self._w("".join(out)) + # A shrink vacates rows below the new region; the ticker's old + # resting line sits one row past the vacated span, so clear it + # too while we are down there. + span = extra + (1 if self._status else 0) + out.append("\x1b[2K\n" * span) + out.append(f"\x1b[{span}A") + if out: + self._w("".join(out)) + self._screen = list(lines) self._painted = len(lines) def _freeze(self, cur: str) -> None: @@ -538,6 +652,9 @@ def _freeze(self, cur: str) -> None: of this block. With part of the block already scroll-committed the painted tail just stays as-is (stale width) and the stream continues raw below it.""" + if self._status: + self._w("\x1b[2K") + self._status = None if self._committed: self._w("\n") else: @@ -548,10 +665,19 @@ def _freeze(self, cur: str) -> None: self._frozen = True self._raw_emitted = len(cur) self._painted = 0 + self._screen = [] + + def _paint_status(self) -> None: + sgr = getattr(self.theme, "stat", "") + text = self._status or "" + painted = f"{sgr}{text}{self.theme.reset}" if sgr and text else text + self._w("\x1b[2K" + painted + "\r") def _reset_block(self) -> None: self._frozen = False self._raw_emitted = 0 self._painted = 0 self._committed = 0 + self._screen = [] + self._status = None self._width = None diff --git a/gmlx/server_patches/__init__.py b/gmlx/server_patches/__init__.py index 9d42f77..12fd236 100644 --- a/gmlx/server_patches/__init__.py +++ b/gmlx/server_patches/__init__.py @@ -75,6 +75,7 @@ install_ignore_eos, install_openai_stop_sequences, install_stream_thinking_seed, + install_stream_timings, install_thinking_budget_fix, install_vanilla_stream_chunks, ) @@ -156,6 +157,7 @@ "install_server_patches", "install_sse_keepalive", "install_stream_thinking_seed", + "install_stream_timings", "install_thinking_budget_fix", "install_vanilla_stream_chunks", "install_xtc_sampling", @@ -212,6 +214,7 @@ def install_server_patches(cfg, *, reload_fn=None) -> None: install_admit_headroom_gate() install_chat_template_kwargs() install_thinking_budget_fix() + install_stream_timings() install_openai_stop_sequences() install_api_contract() # Before the load-offload / profile-capture / keepalive wrappers so they diff --git a/gmlx/server_patches/api_contract.py b/gmlx/server_patches/api_contract.py index 298b69c..b8e7f4b 100644 --- a/gmlx/server_patches/api_contract.py +++ b/gmlx/server_patches/api_contract.py @@ -58,10 +58,12 @@ # /v1/chat/completions (openai.py chat_completions_endpoint + the gmlx stop # filter). ``tool_choice`` is consumed here: "none" is enforced below, other -# values are documented as template-dependent. +# values are documented as template-dependent. ``timings_per_token`` is the +# per-chunk stream-timings switch (install_stream_timings). CHAT_CONSUMED = _GEN_ARGS_CONSUMED | _GMLX_CONSUMED | frozenset({ "model", "messages", "stream", "stream_options", "adapter_path", "resize_shape", "tools", "tool_choice", "top_logprobs", "stop", + "timings_per_token", }) # /v1/responses (openai.py responses_endpoint). No ``stop`` here: the gmlx diff --git a/gmlx/server_patches/chat_behavior.py b/gmlx/server_patches/chat_behavior.py index dbd1780..a104f92 100644 --- a/gmlx/server_patches/chat_behavior.py +++ b/gmlx/server_patches/chat_behavior.py @@ -624,3 +624,114 @@ def model_dump_json(self, **kwargs): model_dump_json.__dict__[_PATCH_FLAG] = True chunk_cls.model_dump_json = model_dump_json + + +# Per-chunk stream timings +# A live client display needs the server's exact cumulative output-token count +# while a reply streams; the OpenAI shape only carries it in the final usage +# chunk, and counting SSE events undercounts whenever one chunk carries several +# tokens (an MTP verify round) or a token's delta is suppressed (thinking +# markers, tool-call markup). The engine's token items do carry exact counts +# (``token_count``), and each passes through ``GenerationMetrics.record_chunk`` +# right before its chunk is yielded, so the running total is available at the +# route without touching the upstream generator. This patch mirrors +# llama.cpp's convention: a request with ``timings_per_token: true`` gets a +# ``timings`` object on each streamed content chunk, here with the running +# ``predicted_n``. The count crosses from ``record_chunk`` to the route's SSE +# rewrite through a per-request contextvar cell: the route wrapper sets it in +# the request task, and starlette's response task inherits that context. +# Off by default; without the request field the stream is byte-identical. +_STREAM_TIMINGS_FLAG = "_kq_gguf_stream_timings_patch" + +_stream_count_cell: contextvars.ContextVar = contextvars.ContextVar( + "gmlx_stream_token_count", default=None) + + +def _count_record_chunk(original): + def record_chunk(self, chunk): + original(self, chunk) + cell = _stream_count_cell.get() + if cell is None: + return + n = getattr(chunk, "generation_tokens", None) + if n: # stream_generate results: cumulative + cell["n"] = int(n) + else: # engine tokens: per-item count + cell["n"] += int(getattr(chunk, "token_count", 1) or 1) + record_chunk.__dict__[_STREAM_TIMINGS_FLAG] = True + return record_chunk + + +async def _timings_sse(body, cell): + """Wrap a chat-completions SSE body iterator: each content chunk gains a + ``timings`` object carrying the running ``predicted_n``. Usage-only and + role-only chunks pass through untouched.""" + import json + + pending = "" + try: + async for raw in body: + pending += raw.decode() if isinstance(raw, (bytes, bytearray)) else raw + while "\n\n" in pending: + event, pending = pending.split("\n\n", 1) + yield _stamp_timings_event(event, cell, json) + "\n\n" + if pending: + yield pending + finally: + aclose = getattr(body, "aclose", None) + if aclose is not None: + try: + await aclose() + except Exception: + pass + + +def _stamp_timings_event(event: str, cell, json) -> str: + if not event.startswith("data: "): + return event + payload = event[len("data: "):] + if payload.strip() == "[DONE]": + return event + try: + obj = json.loads(payload) + except ValueError: + return event + choices = obj.get("choices") or [] + if not choices: + return event # final usage chunk + delta = choices[0].get("delta") or {} + if not (delta.get("content") or delta.get("reasoning") + or delta.get("tool_calls")): + return event # role-only or empty delta + obj["timings"] = {"predicted_n": cell["n"]} + return "data: " + json.dumps(obj) + + +def install_stream_timings() -> None: + """Honour ``timings_per_token`` on streamed chat completions: each content + chunk carries ``timings.predicted_n``, the exact cumulative output-token + count (MTP verify rounds included). Idempotent per route.""" + from starlette.responses import StreamingResponse + + generation = importlib.import_module("mlx_vlm.server.generation") + metrics_cls = generation.GenerationMetrics + if not getattr(metrics_cls.record_chunk, _STREAM_TIMINGS_FLAG, False): + metrics_cls.record_chunk = _count_record_chunk(metrics_cls.record_chunk) + + app = importlib.import_module("mlx_vlm.server.app").app + + def _make(original): + async def endpoint(request, http_request): + want = bool(getattr(request, "timings_per_token", False)) \ + and bool(getattr(request, "stream", False)) + if not want: + return await original(request, http_request) + cell = {"n": 0} + _stream_count_cell.set(cell) + result = await original(request, http_request) + if isinstance(result, StreamingResponse): + result.body_iterator = _timings_sse(result.body_iterator, cell) + return result + return endpoint + + _wrap_post_routes(app, _CHAT_PATHS, _STREAM_TIMINGS_FLAG, _make) diff --git a/gmlx/talk_client.py b/gmlx/talk_client.py index c80fb42..0a8b7d9 100644 --- a/gmlx/talk_client.py +++ b/gmlx/talk_client.py @@ -155,6 +155,7 @@ def _caps(): "tts": any(e.get("tts") for e in entries), "chat_ids": chat_ids, "default": default, + "gmlx": any(e.get("owned_by") == "gmlx" for e in entries), } return _decode_body("capability probe", _caps) @@ -278,11 +279,13 @@ def stream_chat(base_url: str, *, model: str, messages: list, timeout: float = 600.0, extra: dict | None = None) -> Iterator[dict]: """Stream ``/v1/chat/completions`` -> the raw ``delta`` dict per SSE chunk - (plus ``{"_finish": ...}``/``{"_usage": ...}`` markers). ``tools`` is an - OpenAI function-spec list (the assistant brain's loop); ``tool_calls`` - deltas pass through verbatim. ``extra`` merges additional payload fields - (sampling passthrough, stream_options). Closing the generator closes the - HTTP response - that is the cancellation path.""" + (plus ``{"_finish": ...}``/``{"_usage": ...}``/``{"_timings": ...}`` + markers; ``_timings`` relays a chunk's ``timings`` object, sent per content + chunk by gmlx servers when the request carries ``timings_per_token``). + ``tools`` is an OpenAI function-spec list (the assistant brain's loop); + ``tool_calls`` deltas pass through verbatim. ``extra`` merges additional + payload fields (sampling passthrough, stream_options). Closing the + generator closes the HTTP response - that is the cancellation path.""" payload = {"model": model, "messages": messages, "stream": True} if max_tokens is not None: # None = server default (uncapped chat) payload["max_tokens"] = max_tokens @@ -311,6 +314,8 @@ def stream_chat(base_url: str, *, model: str, messages: list, yield delta if choice.get("finish_reason"): yield {"_finish": choice["finish_reason"]} + if chunk.get("timings"): + yield {"_timings": chunk["timings"]} if chunk.get("usage"): yield {"_usage": chunk["usage"]} finally: @@ -400,6 +405,7 @@ def _is_sentence_end(self, i: int) -> bool: # Brain protocol (the phase-2 seam) BrainEvent = tuple # ("say", text) | ("status", label) | ("done", stats: dict) + # | ("count", n) - cumulative output tokens this round class Brain(Protocol): diff --git a/tests/test_assistant_brain.py b/tests/test_assistant_brain.py index a10c3ac..81a74d7 100644 --- a/tests/test_assistant_brain.py +++ b/tests/test_assistant_brain.py @@ -102,6 +102,23 @@ def test_reasoning_deltas_become_status(): assert ("say", "It is ") in events +def test_timings_markers_become_count_events(): + final_t = {"predicted_n": 5, "predicted_per_second": 50.0, + "prompt_per_second": 600.0} + stream = _stream_script([{"content": "It is "}, {"_timings": {"predicted_n": 3}}, + {"content": "5pm."}, {"_timings": {"predicted_n": 5}}, + {"_finish": "stop"}, + {"_usage": {"completion_tokens": 5}}, + {"_timings": final_t}]) + events = list(_brain(stream).turn("hi")) + assert [e for e in events if e[0] == "count"] == [ + ("count", 3), ("count", 5), ("count", 5)] + assert ("say", "It is ") in events and ("say", "5pm.") in events + # the last timings seen ride on the done stats (decode/prefill rates) + assert events[-1] == ("done", {"completion_tokens": 5, + "timings": final_t}) + + # -- the tool loop ------------------------------------------------------------- def test_tool_round_executes_and_recalls(): stream = _stream_script(_tool_round(text="Checking. "), _PROSE) diff --git a/tests/test_chat.py b/tests/test_chat.py index 56deafe..b1bfd26 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -967,3 +967,80 @@ def test_xtc_rides_along_to_the_assistant_server(state, capsys): chat._sync_assistant_extra(state) assert state.assistant_extra["xtc_probability"] == 0.5 assert "xtc_threshold" not in state.assistant_extra + + +# -- _RateTicker ----------------------------------------------------------------- + + +def test_rate_ticker_exact_from_generation_tokens(): + from types import SimpleNamespace + from gmlx.chat import _RateTicker + + now = [0.0] + tk = _RateTicker(clock=lambda: now[0]) + status = None + for i in range(1, 31): + now[0] = i * 0.1 + s = tk.push(SimpleNamespace(text="x", generation_tokens=i * 15)) + if s: + status = s + assert status is not None and status.endswith(" tok/s") + assert not status.startswith("~") + assert abs(int(status.split()[0]) - 150) <= 2 + + +def test_stat_line_shows_ttft_beside_decode_rate(): + from gmlx.chat import _fmt_stat_line + + line = _fmt_stat_line({"prompt_tokens": 2296, "prompt_tps": 670.0, + "gen_tokens": 607, "gen_tps": 50.3, + "ttft_s": 7.2}, 2900, 0) + assert "ttft 7.2s" in line + assert "@ 50.3 tok/s" in line + # no wait recorded (local pipeline): segment absent + line = _fmt_stat_line({"gen_tokens": 607, "gen_tps": 50.3}, 0, 0) + assert "ttft" not in line + + +def test_rate_ticker_silent_without_counts(): + from types import SimpleNamespace + from gmlx.chat import _RateTicker + + now = [0.0] + tk = _RateTicker(clock=lambda: now[0]) + for i in range(1, 31): + now[0] = i * 0.1 + assert tk.push(SimpleNamespace(text="abcd" * 14)) is None + + +def test_rate_ticker_accumulates_across_round_restarts(): + from types import SimpleNamespace + from gmlx.chat import _RateTicker + + now = [0.0] + tk = _RateTicker(clock=lambda: now[0]) + status = None + for i in range(1, 31): + now[0] = i * 0.1 + # Two server rounds at 150 tok/s: the count restarts at chunk 16 + # (a new request in the tool loop) but the rate must not glitch. + n = i * 15 if i <= 15 else (i - 15) * 15 + s = tk.push(SimpleNamespace(text="x", generation_tokens=n)) + if s: + status = s + assert status is not None + assert abs(int(status.split()[0]) - 150) <= 2 + + +def test_rate_ticker_throttles_pushes(): + from types import SimpleNamespace + from gmlx.chat import _RateTicker + + now = [0.0] + tk = _RateTicker(clock=lambda: now[0]) + pushes = 0 + for i in range(1, 101): + now[0] = i * 0.05 # 20 chunks/s for 5s + if tk.push(SimpleNamespace(text="hello world here", generation_tokens=i * 5)): + pushes += 1 + assert 0 < pushes <= 18 diff --git a/tests/test_render.py b/tests/test_render.py index 02ee320..e5388f3 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -326,3 +326,137 @@ def test_feed_strips_model_ansi_and_control_chars(): assert "\x1b" not in text and "\x00" not in text and "\x07" not in text assert "\r" not in text assert "hi " in text and "red" in text and "ok" in text + + +# -- diff-aware repaint -------------------------------------------------------- + + +def test_stream_diff_append_rewrites_only_changed_rows(): + t = _Term() + r = t.renderer() + r.feed("```\nalpha\nbravo\n") + t.out.clear() + r.feed("charlie\n") # append: alpha/bravo rows are unchanged + paint = _strip(t.text()) + assert "charlie" in paint + assert "alpha" not in paint and "bravo" not in paint + + +def test_stream_diff_oversize_append_stays_small(): + t = _Term(columns=41, lines=8) # budget 6: block taller than the viewport + r = t.renderer() + r.feed("```\n") + for i in range(8): + r.feed(f"line {i}\n") + t.out.clear() + r.feed("line 8\n") # slides the live window by one + out = t.text() + assert "line 8" in _strip(out) + # Only the boundary rows repaint; committed rows never rewrite, so the + # cursor-up span stays a couple of rows, not the whole window. + assert all(int(n) <= 3 for n in re.findall(r"\x1b\[(\d+)A", out)) + assert "line 5" not in _strip(out) + + +def test_stream_diff_unchanged_finalize_writes_nothing(): + t = _Term() + r = t.renderer() + r.feed("steady text\n") + t.out.clear() + r.finalize() # tail render equals the painted screen + assert t.text() == "" + + +def test_stream_adaptive_interval_backs_off(): + ticks = iter([1.0, 2.0, 2.1, 2.2, 2.4, 3.4, 3.45]) + t = _Term() + r = rd.StreamRenderer( + "lite", _THEME, write=t.write, size_fn=lambda: t.size, + clock=lambda: next(ticks), + ) + r.feed("a") # paint start 1.0, end 2.0 -> interval 0.25 + n = len(t.out) + r.feed("b") # now=2.1: 0.1 since paint end, skipped + assert len(t.out) == n + r.feed("c") # now=2.2: still inside the interval + assert len(t.out) == n + r.feed("d") # now=2.4: past the interval, paints again + assert len(t.out) > n + + +# -- fence live-render trimming ------------------------------------------------- + + +def test_stream_fence_trim_bounds_render_input(): + t = _Term(columns=41, lines=8) # budget 6: trim kicks in past 28 body lines + r = t.renderer() + r.feed("```\n") + for i in range(40): + r.feed(f"row {i}\n") + assert r._src_skip > 0 + live = r._trimmed(r._buf.current) + assert len(live.splitlines()) <= 2 * 6 + 16 + 2 + r.feed("```\n") + r.feed("\nafter\n") + r.finalize() + stripped = _strip(t.text()) + for i in range(40): + assert f"row {i}" in stripped, i + assert "after" in stripped + assert "```" not in stripped + + +def test_stream_fence_trim_resets_between_blocks(): + t = _Term(columns=41, lines=8) + r = t.renderer() + r.feed("```\n") + for i in range(40): + r.feed(f"row {i}\n") + r.feed("```\n") + assert r._src_skip == 0 # completion resets the trim + r.feed("\nplain paragraph text\n") + r.finalize() + assert "plain paragraph text" in _strip(t.text()) + + +def test_stream_paragraph_never_trims(): + t = _Term(columns=41, lines=8) + r = t.renderer() + for i in range(40): + r.feed(f"word{i} ") + assert r._src_skip == 0 + r.finalize() + assert "word39" in _strip(t.text()) + + +# -- live status ticker ---------------------------------------------------------- + + +def test_stream_status_paints_on_resting_line(): + t = _Term() + r = t.renderer() + r.feed("some streaming text\n") + r.set_status("148 tok/s") + out = t.text() + assert "148 tok/s" in _strip(out) + assert out.endswith("\r") # cursor parked back at line start + + +def test_stream_status_ignored_before_first_paint(): + t = _Term() + r = t.renderer() + r.set_status("99 tok/s") + assert "99 tok/s" not in t.text() + + +def test_stream_status_cleared_by_block_end_and_finalize(): + t = _Term() + r = t.renderer() + r.feed("first block\n") + r.set_status("120 tok/s") + r.feed("\nsecond block\n") # separator must clear the ticker line + n = t.text().count("\x1b[2K\n") + assert n >= 1 + r.set_status("125 tok/s") + r.finalize() + assert t.text().rstrip("\n").endswith("\x1b[2K") diff --git a/tests/test_server_patches.py b/tests/test_server_patches.py index bf17ac6..0d6834d 100644 --- a/tests/test_server_patches.py +++ b/tests/test_server_patches.py @@ -1737,6 +1737,76 @@ def test_install_openai_stop_wraps_routes_idempotent(): assert len(_APP.app.router.routes) == n +# Per-chunk stream timings: exact cumulative output-token counts stamped onto +# streamed content chunks when the request asks for timings_per_token. +def test_count_record_chunk_engine_tokens_and_cumulative_results(): + recorded = [] + counted = sp_chat._count_record_chunk(lambda self, c: recorded.append(c)) + cell = {"n": 0} + token = sp_chat._stream_count_cell.set(cell) + try: + counted(None, types.SimpleNamespace(token_count=3)) # MTP verify round + counted(None, types.SimpleNamespace()) # plain token + assert cell["n"] == 4 + counted(None, types.SimpleNamespace(generation_tokens=10)) # cumulative + assert cell["n"] == 10 + finally: + sp_chat._stream_count_cell.reset(token) + assert len(recorded) == 3 # original still ran + + +def test_count_record_chunk_noop_without_cell(): + counted = sp_chat._count_record_chunk(lambda self, c: None) + counted(None, types.SimpleNamespace(token_count=5)) # no cell: no error + + +def test_timings_sse_stamps_content_chunks_only(): + import asyncio + import json + + cell = {"n": 0} + role = _sse({"id": "c1", "object": "chat.completion.chunk", "created": 1, + "model": "m", + "choices": [{"index": 0, "delta": {"role": "assistant"}, + "finish_reason": None}]}) + usage = _sse({"id": "c1", "choices": [], "usage": {"total_tokens": 9}}) + feed = [(0, role), (4, _sse(_chunk(content="hi"))), + (9, _sse(_chunk(content=" there"))), + (9, _sse(_chunk(finish="stop"))), (9, usage), + (9, "data: [DONE]\n\n")] + + async def upstream(): + for n, e in feed: # the count advances as the stream does + cell["n"] = n + yield e + + async def collect(): + return [e async for e in sp_chat._timings_sse(upstream(), cell)] + + out = asyncio.run(collect()) + assert out[-1] == "data: [DONE]\n\n" + objs = [json.loads(e[len("data: "):]) for e in out + if e.startswith("data: ") and "[DONE]" not in e] + assert [o["timings"]["predicted_n"] for o in objs if "timings" in o] == [4, 9] + assert "timings" not in objs[0] # role chunk + assert "timings" not in objs[3] # finish chunk + assert "timings" not in objs[4] # usage chunk + + +def test_install_stream_timings_wraps_routes_idempotent(): + sp.install_stream_timings() + generation = importlib.import_module("mlx_vlm.server.generation") + assert getattr(generation.GenerationMetrics.record_chunk, + sp_chat._STREAM_TIMINGS_FLAG, False) + routes = {getattr(r, "path", None): r for r in _APP.app.router.routes} + for path in sp_common._CHAT_PATHS: + assert getattr(routes[path].endpoint, + sp_chat._STREAM_TIMINGS_FLAG, False) + n = len(_APP.app.router.routes) + sp.install_stream_timings() # idempotent + assert len(_APP.app.router.routes) == n + + def test_openai_stop_endpoint_e2e_with_stub(): """POST through FastAPI with a stub original handler: proves the signature propagation parses the body and the wrapper trims non-stream responses.""" diff --git a/tests/test_talk_client.py b/tests/test_talk_client.py index d890c7a..9a72d53 100644 --- a/tests/test_talk_client.py +++ b/tests/test_talk_client.py @@ -210,6 +210,31 @@ def test_stream_chat_yields_deltas_and_markers(monkeypatch): assert resp.closed +def test_stream_chat_relays_timings_marker(monkeypatch): + with_timings = _chunk({"content": "Hel"}) + with_timings["timings"] = {"predicted_n": 2} + final = {"choices": [], "usage": {"total_tokens": 5}, + "timings": {"predicted_n": 5, "predicted_per_second": 50.0}} + resp = _FakeResp(_sse(with_timings, final)) + monkeypatch.setattr(tc, "_open_stream", + lambda url, payload, api_key, timeout: resp) + got = list(tc.stream_chat("http://h:1/v1", model="m", messages=[], + max_tokens=10)) + assert {"_timings": {"predicted_n": 2}} in got + # the final chunk's full timings relay too (rates for the stat line) + assert {"_timings": final["timings"]} in got + assert {"_usage": {"total_tokens": 5}} in got + + +def test_probe_capabilities_marks_gmlx_server(monkeypatch): + payload = _models_payload() + monkeypatch.setattr(tc, "_http_get_json", + lambda url, timeout=5.0, api_key=None: payload) + assert not tc.probe_capabilities("http://h:1/v1")["gmlx"] + payload["data"][0]["owned_by"] = "gmlx" + assert tc.probe_capabilities("http://h:1/v1")["gmlx"] + + def test_stream_chat_close_on_break(monkeypatch): resp = _FakeResp(_sse(_chunk({"content": "a"}), _chunk({"content": "b"}))) monkeypatch.setattr(tc, "_open_stream",