renderers 0.1.8.dev2__tar.gz → 0.1.8.dev4__tar.gz

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: renderers
-Version: 0.1.8.dev2
+Version: 0.1.8.dev4
 Summary: Chat template renderers — deterministic message-to-token conversion for LLM training
 License-Expression: Apache-2.0
 License-File: LICENSE

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/__init__.py

@@ -36,6 +36,7 @@ from renderers.base import (
     reject_assistant_in_extension,
     trim_to_turn_close,
 )
+from renderers.client import OverlongPromptError
 from renderers.deepseek_v3 import DeepSeekV3Renderer
 from renderers.default import DefaultRenderer
 from renderers.glm5 import GLM5Renderer
@@ -69,6 +70,7 @@ __all__ = [
     "MultiModalData",
     "MultimodalRenderer",
     "Nemotron3Renderer",
+    "OverlongPromptError",
     "ParsedResponse",
     "ParsedToolCall",
     "PlaceholderRange",

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.8.dev2'
-__version_tuple__ = version_tuple = (0, 1, 8, 'dev2')
+__version__ = version = '0.1.8.dev4'
+__version_tuple__ = version_tuple = (0, 1, 8, 'dev4')
 
 __commit_id__ = commit_id = None

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/base.py

@@ -177,8 +177,162 @@ class RenderedTokens:
     token_ids: list[int] = field(default_factory=list)
     message_indices: list[int] = field(default_factory=list)
     sampled_mask: list[bool] = field(default_factory=list)
+    message_roles: list[str] = field(default_factory=list)
     multi_modal_data: "MultiModalData | None" = None
 
+    def tokens_per_message(
+        self, n_messages: int | None = None, *, sampled_only: bool = False
+    ) -> list[int]:
+        """Count rendered tokens attributed to each caller-relative message.
+
+        ``out[i]`` is the number of tokens with ``message_indices[k] == i``,
+        i.e. tokens the renderer attributed to ``messages[i]``. This
+        includes template scaffolding the renderer wraps around the
+        message — the ``<|im_start|>role\\n`` opener, the closing
+        ``<|im_end|>\\n``, etc. — because those are the renderer's own
+        attribution decision and are preserved verbatim here. Tokens with
+        ``message_indices[k] == -1`` (scaffolding outside any single
+        message, e.g. the trailing generation prompt) are not counted.
+
+        With ``sampled_only=True``, counts only tokens the model would
+        have emitted at inference (``sampled_mask[k] is True``). For
+        example, length-penalty signals in RL: the template wraps each
+        assistant turn in scaffolding tokens (e.g. ``<|im_start|>assistant\\n``,
+        ``<|im_end|>\\n``) that are constant-size and not chosen by the
+        model, so they shouldn't enter the penalty. For roles the model
+        never samples (``user``, ``tool``, ``system``), the
+        ``sampled_only`` count is zero by construction. Renderers that
+        don't populate ``sampled_mask`` (``DefaultRenderer`` — the Jinja
+        template is opaque) return all zeros under ``sampled_only=True``.
+
+        ``n_messages`` defaults to ``len(self.message_roles)``, which
+        every Renderer populates with the caller-relative message list
+        (caller's ``messages`` for ``render()``; ``new_messages`` for
+        ``bridge_to_next_turn()``). Pass it explicitly only to truncate
+        — indices outside ``[0, n_messages)`` are ignored, so passing a
+        smaller value won't raise; it just drops the tail. Values larger
+        than ``len(self.message_roles)`` are clamped, so the returned
+        list never claims more messages than the renderer attributed.
+
+        Works on results from both :meth:`Renderer.render` and
+        :meth:`Renderer.bridge_to_next_turn`. For a bridge result the
+        indices are relative to the new messages the bridge added, not
+        the full conversation history; the prior portion is uniformly
+        ``-1`` (and ``sampled_mask`` uniformly ``False``), so it
+        contributes nothing to either count.
+        """
+        if n_messages is None:
+            n_messages = len(self.message_roles)
+        else:
+            n_messages = min(n_messages, len(self.message_roles))
+        out = [0] * n_messages
+        if sampled_only:
+            if len(self.sampled_mask) != len(self.token_ids):
+                return out
+            for idx, sampled in zip(self.message_indices, self.sampled_mask):
+                if sampled and 0 <= idx < n_messages:
+                    out[idx] += 1
+        else:
+            for idx in self.message_indices:
+                if 0 <= idx < n_messages:
+                    out[idx] += 1
+        return out
+
+    def message_token_spans(self) -> list[tuple[int, int] | None]:
+        """Per-message ``(start, end)`` slices into :attr:`token_ids`.
+
+        ``out[i]`` is the half-open span ``[start, end)`` such that
+        ``token_ids[start:end]`` are the tokens attributed to
+        ``messages[i]`` (or ``new_messages[i]`` for a bridge result).
+        Messages that contributed no tokens get ``None``. Renderer
+        scaffolding outside any message (``message_indices[k] == -1``)
+        is not represented.
+
+        Hand-coded renderers emit each message's tokens contiguously,
+        so the span is well-defined. The implementation tolerates
+        non-contiguous attribution by returning the outer span
+        ``(first_k, last_k + 1)``; if you suspect interleaving, slice
+        ``message_indices`` yourself to verify.
+
+        Returns ``len(self.message_roles)`` entries when ``message_roles``
+        is populated. Otherwise infers the count from
+        ``max(message_indices) + 1`` — useful for manually-constructed
+        ``RenderedTokens`` in tests but only correct when the last
+        message contributed at least one token.
+
+        Cheap to call: single pass over ``message_indices``. Re-call
+        rather than caching the result if you mutate the dataclass.
+        """
+        if self.message_roles:
+            n_messages = len(self.message_roles)
+        else:
+            max_idx = -1
+            for idx in self.message_indices:
+                if idx > max_idx:
+                    max_idx = idx
+            n_messages = max_idx + 1
+
+        firsts: list[int] = [-1] * n_messages
+        lasts: list[int] = [-1] * n_messages
+        for k, idx in enumerate(self.message_indices):
+            if 0 <= idx < n_messages:
+                if firsts[idx] == -1:
+                    firsts[idx] = k
+                lasts[idx] = k
+
+        out: list[tuple[int, int] | None] = []
+        for i in range(n_messages):
+            if firsts[i] == -1:
+                out.append(None)
+            else:
+                out.append((firsts[i], lasts[i] + 1))
+        return out
+
+    def role_token_spans(self) -> dict[str, list[tuple[int, int]]]:
+        """:meth:`message_token_spans` regrouped by ``message_roles``.
+
+        Maps each role appearing in :attr:`message_roles` to a list of
+        ``(start, end)`` spans — one per occurrence of that role, in
+        message order. Messages with no contributed tokens are skipped.
+        Returns an empty dict if :attr:`message_roles` is empty.
+
+        Intended for per-role statistics that operate on per-token
+        signals — e.g. ``logprobs[start:end]`` for each assistant span
+        to compute per-turn perplexity, or
+        ``attention[start:end]`` for tool-response attention analysis.
+        """
+        spans = self.message_token_spans()
+        out: dict[str, list[tuple[int, int]]] = {}
+        for role, span in zip(self.message_roles, spans):
+            if span is None:
+                out.setdefault(role, [])
+                continue
+            out.setdefault(role, []).append(span)
+        return out
+
+    def tokens_by_role(self, *, sampled_only: bool = False) -> dict[str, int]:
+        """Sum :meth:`tokens_per_message` grouped by ``message_roles``.
+
+        Convenience for length-penalty bookkeeping in RL trainers:
+        ``rendered.tokens_by_role(sampled_only=True)["assistant"]`` is
+        the count of tokens the model actually emitted across all
+        assistant turns — template scaffolding excluded.
+        ``rendered.tokens_by_role()["tool"]`` is the raw count of
+        tool-response tokens (``sampled_only`` is zero for ``tool`` by
+        construction since the model never samples those).
+
+        Roles present in :attr:`message_roles` always appear in the
+        returned dict, even with post-filter count ``0``, so callers
+        can index directly without ``KeyError`` on conversations that
+        happen to lack a role. Returns an empty dict if
+        :attr:`message_roles` is empty.
+        """
+        counts = self.tokens_per_message(sampled_only=sampled_only)
+        out: dict[str, int] = {}
+        for role, n in zip(self.message_roles, counts):
+            out[role] = out.get(role, 0) + n
+        return out
+
 
 class ToolCallParseStatus(str, enum.Enum):
     """Per-attempt outcome of parsing a single ``<tool_call>`` block.
@@ -358,6 +512,25 @@ class Renderer(Protocol):
         list so far with ``add_generation_prompt=True`` — except prev
         sampled tokens are kept verbatim rather than re-rendered).
 
+        Attribution on the returned ``RenderedTokens``:
+
+        - ``message_indices`` is ``-1`` over the entire prior portion
+          (length ``len(previous_ids)`` after :func:`trim_to_turn_close`)
+          because the bridge gets the prior as raw token lists with no
+          attribution. Over the bridge-added portion, indices are
+          relative to ``new_messages``: a token rendered as part of
+          ``new_messages[i]`` carries ``i``, and inter-turn separators /
+          the trailing generation prompt carry ``-1``. So
+          ``bridge.tokens_per_message(len(new_messages))`` gives the
+          per-new-message token count for length-penalty bookkeeping.
+        - ``sampled_mask`` is uniformly ``False`` across the entire
+          returned sequence. The bridge output is consumed as the next
+          turn's prompt; nothing it emits was model-sampled, and the
+          bridge has no way to recover which prior tokens were. If the
+          caller needs that distinction for the prior portion, they
+          have it directly: every token in ``prev_completion_ids`` was
+          sampled; every token in ``prev_prompt_ids`` was not.
+
         Text-only renderers return :class:`RenderedTokens` with
         ``multi_modal_data=None``. Multimodal renderers (see
         :class:`MultimodalRenderer`) populate ``multi_modal_data`` so
@@ -593,6 +766,8 @@ MODEL_RENDERER_MAP: dict[str, str] = {
     "Qwen/Qwen3-14B": "qwen3",
     "Qwen/Qwen3-32B": "qwen3",
     "Qwen/Qwen3-30B-A3B": "qwen3",
+    "Qwen/Qwen3-30B-A3B-Instruct-2507": "qwen3",
+    "Qwen/Qwen3-30B-A3B-Thinking-2507": "qwen3",
     "Qwen/Qwen3-235B-A22B": "qwen3",
     # Qwen3.5. All seven sizes share the same renderer. The 4B / 9B /
     # 35B-A3B / 122B-A10B / 397B-A17B chat template defaults
@@ -619,6 +794,7 @@ MODEL_RENDERER_MAP: dict[str, str] = {
     "Qwen/Qwen3-VL-30B-A3B-Instruct": "qwen3-vl",
     # GLM-5 family (GLM-4.7 reuses the GLM-5 template).
     "zai-org/GLM-5": "glm-5",
+    "zai-org/GLM-5-FP8": "glm-5",
     "zai-org/GLM-4.7-Flash": "glm-5",
     "zai-org/GLM-5.1": "glm-5.1",
     # GLM-4.5.

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/client.py

@@ -14,10 +14,11 @@ from __future__ import annotations
 import asyncio
 import base64
 import logging
+from collections.abc import Mapping
 from typing import Any, cast
 
 import numpy as np
-from openai import AsyncOpenAI, BadRequestError
+from openai import AsyncOpenAI
 
 from renderers.base import (
     Message,
@@ -31,6 +32,79 @@ from renderers.base import (
 _request_logger = logging.getLogger("renderers.client")
 
 
+class OverlongPromptError(Exception):
+    """The rendered prompt exceeds the engine's context window.
+
+    Raised by :func:`generate` when the rendered token sequence is strictly
+    longer than the resolved cap — either an explicit ``max_prompt_len`` the
+    caller passed in, or the engine's ``max_model_len`` discovered via
+    ``GET /v1/models``. Caught client-side before the engine ever sees the
+    request, so callers route the failure to a deterministic policy (skip /
+    truncate / count) instead of round-tripping through an engine 4xx.
+
+    Named after the corresponding ``verifiers.errors.OverlongPromptError``;
+    the two are distinct classes (different package hierarchies) but the
+    concept is the same and downstream clients translate one to the other.
+    """
+
+    def __init__(self, *, prompt_len: int, max_prompt_len: int) -> None:
+        self.prompt_len = prompt_len
+        self.max_prompt_len = max_prompt_len
+        super().__init__(
+            f"Prompt length ({prompt_len}) exceeds maximum "
+            f"context length ({max_prompt_len})."
+        )
+
+
+# Per-process cache of resolved engine context-length caps, keyed by
+# ``(base_url, model)``. ``None`` is the "we asked the engine and it didn't
+# tell us" sentinel — distinct from "key missing" (haven't asked yet). The
+# lock serializes the first lookup per key; cache hits avoid the lock.
+_max_prompt_len_cache: dict[tuple[str, str], int | None] = {}
+_max_prompt_len_lock = asyncio.Lock()
+
+
+async def _resolve_max_prompt_len(client: AsyncOpenAI, model: str) -> int | None:
+    """Discover ``max_model_len`` from the engine via ``GET /v1/models``.
+
+    OpenAI-API-compatible engines expose model metadata at this endpoint;
+    vLLM extends its ``ModelCard`` with a ``max_model_len`` field. Engines
+    that don't (SGLang as of this writing, third-party gateways, etc.) get
+    a cached ``None`` and the pre-flight overflow check silently disables —
+    callers fall back to whatever reactive handling they have for engine
+    4xx, which the verifiers ``@handle_openai_overlong_prompt`` decorator
+    already supplies for the prime-rl path.
+
+    Any exception during lookup (network error, non-JSON body, attribute
+    miss on a mock client in tests) is treated as "unknown cap": cached
+    ``None`` so we don't retry on every call.
+    """
+    key = (str(getattr(client, "base_url", "")), model)
+    if key in _max_prompt_len_cache:
+        return _max_prompt_len_cache[key]
+    async with _max_prompt_len_lock:
+        if key in _max_prompt_len_cache:
+            return _max_prompt_len_cache[key]
+        try:
+            payload = await client.get("/models", cast_to=cast(Any, dict[str, Any]))
+        except Exception as exc:
+            _request_logger.debug("max_prompt_len lookup failed: %s", exc)
+            _max_prompt_len_cache[key] = None
+            return None
+        value: int | None = None
+        for card in payload.get("data") or []:
+            if not isinstance(card, Mapping):
+                continue
+            if card.get("id") != model:
+                continue
+            raw = card.get("max_model_len")
+            if isinstance(raw, int) and raw > 0:
+                value = raw
+            break
+        _max_prompt_len_cache[key] = value
+        return value
+
+
 async def _maybe_offload(renderer: Renderer | RendererPool, fn):
     """Run sync renderer work on a thread iff ``renderer`` is a pool.
 
@@ -58,6 +132,7 @@ async def generate(
     cache_salt: str | None = None,
     priority: int | None = None,
     extra_headers: dict[str, str] | None = None,
+    max_prompt_len: int | None = None,
 ) -> dict[str, Any]:
     """Tokenize messages, call vLLM /inference/v1/generate, parse the response.
 
@@ -74,6 +149,16 @@ async def generate(
     mm_placeholders, kwargs_data) before POSTing. The serializer imports
     ``vllm.*`` lazily so text-only consumers never pay for the import.
 
+    ``max_prompt_len`` controls the pre-flight overflow check. When the
+    rendered prompt is strictly longer than the cap, the request is never
+    sent and ``OverlongPromptError`` is raised. If ``max_prompt_len`` is
+    ``None`` (the default), the cap is auto-discovered once per
+    ``(base_url, model)`` via ``GET /v1/models`` (vLLM's
+    ``ModelCard.max_model_len`` extension); engines that don't expose it
+    cache a ``None`` cap and the pre-flight silently disables. Engine 4xx
+    that still slip through propagate raw — converting them into a domain
+    error is the calling client's job (its error shape is engine-specific).
+
     Returns a dict with: request_id, prompt_ids, completion_ids,
     completion_logprobs, content, reasoning_content, tool_calls,
     finish_reason, routed_experts.
@@ -96,6 +181,13 @@ async def generate(
 
     prompt_ids, stop_token_ids, mm_data = await _maybe_offload(renderer, _prepare)
 
+    if max_prompt_len is None:
+        max_prompt_len = await _resolve_max_prompt_len(client, model)
+    if max_prompt_len is not None and len(prompt_ids) > max_prompt_len:
+        raise OverlongPromptError(
+            prompt_len=len(prompt_ids), max_prompt_len=max_prompt_len
+        )
+
     sp: dict[str, Any] = dict(sampling_params or {})
     sp["stop_token_ids"] = stop_token_ids
     sp["logprobs"] = 1
@@ -135,16 +227,7 @@ async def generate(
     }
     if extra_headers:
         post_kwargs["options"] = cast(Any, {"headers": extra_headers})
-    try:
-        data = await client.post(endpoint, **post_kwargs)
-    except BadRequestError as exc:
-        _log_overlong_prompt_diagnostic(
-            prompt_ids=prompt_ids,
-            messages=messages,
-            max_tokens=sp.get("max_tokens"),
-            exc=exc,
-        )
-        raise
+    data = await client.post(endpoint, **post_kwargs)
 
     choice = (data.get("choices") or [{}])[0]
     completion_ids = choice.get("token_ids") or []
@@ -225,6 +308,7 @@ def _build_mm_features(
     to change. Don't pre-build the abstraction with one engine in tree.
     """
     from renderers.qwen3_vl import Qwen3VLRenderer
+    from renderers.qwen35 import Qwen35Renderer
 
     # Type dispatch only needs the renderer class. Pools expose
     # ``renderer_cls`` as a snapshot attribute, so we don't have to check
@@ -233,7 +317,10 @@ def _build_mm_features(
         renderer.renderer_cls if isinstance(renderer, RendererPool) else type(renderer)
     )
 
-    if issubclass(renderer_cls, Qwen3VLRenderer):
+    # Qwen3-VL and Qwen3.5 both ship ``pixel_values`` + ``image_grid_thw``
+    # via the shared Qwen2-VL field factory. ``spatial_merge_size=2`` is
+    # the family default and matches every Qwen-VL processor in tree.
+    if issubclass(renderer_cls, (Qwen3VLRenderer, Qwen35Renderer)):
         return _build_qwen_vl_features(mm_data, spatial_merge_size=2)
 
     raise NotImplementedError(
@@ -305,44 +392,3 @@ def _build_qwen_vl_features(
         out["kwargs_data"] = None
 
     return out
-
-
-def _log_overlong_prompt_diagnostic(
-    *,
-    prompt_ids: list[int],
-    messages: list[Message],
-    max_tokens: int | None,
-    exc: BadRequestError,
-) -> None:
-    """Log a structured snapshot when vLLM rejects with 4xx — usually overlong.
-
-    Captures total prompt length, per-message role + character count, and
-    the first chunk of the response body.
-    """
-    body_text = ""
-    response = getattr(exc, "response", None)
-    if response is not None:
-        body_text = (response.text or "")[:500].replace("\n", " ")
-    msg_summary = []
-    for i, m in enumerate(messages):
-        role = m.get("role", "?")
-        content = m.get("content")
-        if isinstance(content, str):
-            content_len = len(content)
-        elif isinstance(content, list):
-            content_len = sum(
-                len(p.get("text", "")) if isinstance(p, dict) else 0 for p in content
-            )
-        else:
-            content_len = 0
-        tool_calls = m.get("tool_calls")
-        tc_count = len(tool_calls) if tool_calls else 0
-        msg_summary.append(f"[{i}]{role}(c={content_len},tc={tc_count})")
-    _request_logger.warning(
-        "vllm 4xx prompt_len=%d messages=%d max_tokens=%s per_msg=%s response_body=%s",
-        len(prompt_ids),
-        len(messages),
-        max_tokens,
-        " ".join(msg_summary),
-        body_text,
-    )

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/deepseek_v3.py

@@ -210,7 +210,10 @@ class DeepSeekV3Renderer:
                 emit_text("<think>\n", -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -271,22 +274,29 @@ class DeepSeekV3Renderer:
             return None
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the
-        # caller blanket-masks it via ``prompt_mask=[False]*N``, so we
-        # don't track sampled_mask here. Local helpers accept the kwarg
-        # for signature compatibility with ``_render_tool`` and ignore
-        # it; the returned ``RenderedTokens`` leaves ``sampled_mask``
-        # empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -329,7 +339,13 @@ class DeepSeekV3Renderer:
         if self._enable_thinking:
             emit_text("<think>\n", -1)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     # ------------------------------------------------------------------
     # Assistant rendering

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/default.py

@@ -143,7 +143,12 @@ class DefaultRenderer:
             token_ids = full_ids
             message_indices.extend([-1] * len(gen_tokens))
 
-        return RenderedTokens(token_ids=token_ids, message_indices=message_indices)
+        message_roles = [m.get("role") or "" for m in messages]
+        return RenderedTokens(
+            token_ids=token_ids,
+            message_indices=message_indices,
+            message_roles=message_roles,
+        )
 
     def _apply(self, messages, *, tools=None, add_generation_prompt=False) -> list[int]:
         kwargs = dict(self._chat_template_kwargs)

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/glm45.py

@@ -203,7 +203,10 @@ class GLM45Renderer:
                 emit_special(self._think_end, -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -271,22 +274,29 @@ class GLM45Renderer:
         last_prev = previous_ids[-1]
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the
-        # caller blanket-masks it via ``prompt_mask=[False]*N``, so we
-        # don't track sampled_mask here. Local helpers accept the kwarg
-        # for signature compatibility with ``_render_tool`` and ignore
-        # it; the returned ``RenderedTokens`` leaves ``sampled_mask``
-        # empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -318,7 +328,13 @@ class GLM45Renderer:
             emit_special(self._think, -1)
             emit_special(self._think_end, -1)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     def _render_assistant(
         self,

{renderers-0.1.8.dev2 → renderers-0.1.8.dev4}/renderers/glm5.py

@@ -220,7 +220,10 @@ class GLM5Renderer:
                 emit_special(self._think_end, -1, is_sampled=False)
 
         return RenderedTokens(
-            token_ids=tokens, message_indices=indices, sampled_mask=sampled
+            token_ids=tokens,
+            message_indices=indices,
+            sampled_mask=sampled,
+            message_roles=[m.get("role") or "" for m in messages],
         )
 
     def render_ids(
@@ -292,22 +295,29 @@ class GLM5Renderer:
         last_prev = previous_ids[-1]
 
         ext: list[int] = []
-
-        # Bridge output is consumed as the next turn's prompt — the
-        # caller blanket-masks it via ``prompt_mask=[False]*N``, so we
-        # don't track sampled_mask here. Local helpers accept the kwarg
-        # for signature compatibility with ``_render_assistant`` /
-        # ``_render_tool`` and ignore it; the returned ``RenderedTokens``
-        # leaves ``sampled_mask`` empty.
+        ext_indices: list[int] = []
+        ext_sampled: list[bool] = []
+
+        # Bridge populates ``message_indices`` (relative to ``new_messages``)
+        # and ``sampled_mask`` (uniformly ``False`` — every token the
+        # bridge emits is template scaffolding for the next prompt, not
+        # something the model sampled). Downstream consumers can run
+        # :meth:`RenderedTokens.tokens_per_message` on the bridge output
+        # to get per-new-message token counts without re-rendering.
         def emit_special(
-            token_id: int, _msg_idx: int = -1, *, is_sampled: bool = False
+            token_id: int, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
             ext.append(token_id)
+            ext_indices.append(msg_idx)
+            ext_sampled.append(is_sampled)
 
         def emit_text(
-            text: str, _msg_idx: int = -1, *, is_sampled: bool = False
+            text: str, msg_idx: int = -1, *, is_sampled: bool = False
         ) -> None:
-            ext.extend(self._encode(text))
+            ids = self._encode(text)
+            ext.extend(ids)
+            ext_indices.extend([msg_idx] * len(ids))
+            ext_sampled.extend([is_sampled] * len(ids))
 
         for i, msg in enumerate(new_messages):
             role = msg.get("role")
@@ -340,7 +350,13 @@ class GLM5Renderer:
         else:
             emit_special(self._think_end, -1)
 
-        return RenderedTokens(token_ids=previous_ids + ext)
+        total_len = len(previous_ids) + len(ext)
+        return RenderedTokens(
+            token_ids=previous_ids + ext,
+            message_indices=[-1] * len(previous_ids) + ext_indices,
+            sampled_mask=[False] * total_len,
+            message_roles=[m.get("role") or "" for m in new_messages],
+        )
 
     def _render_assistant(
         self,