Inline raw multimodal images in the v1 train client

tests/test_client_multimodal_types.py

@@ -1,7 +1,6 @@
 from types import SimpleNamespace
 
 import pytest
-
 from verifiers.clients.openai_chat_completions_client import OpenAIChatCompletionsClient
 from verifiers.types import (
     AssistantMessage,
@@ -214,7 +213,6 @@ async def test_anthropic_merges_consecutive_tool_results_into_single_user_messag
 async def test_anthropic_from_native_response_extracts_usage():
     anthropic = pytest.importorskip("anthropic")
     from anthropic.types import Message as AnthropicMessage
-
     from verifiers.clients.anthropic_messages_client import AnthropicMessagesClient
 
     client = AnthropicMessagesClient(object())
@@ -267,7 +265,6 @@ async def test_anthropic_tool_call_round_trips_thinking_blocks():
     pytest.importorskip("anthropic")
     from anthropic.types import Message as AnthropicMessage
     from anthropic.types import Usage as AnthropicUsage
-
     from verifiers.clients.anthropic_messages_client import AnthropicMessagesClient
 
     client = AnthropicMessagesClient(object())
@@ -296,3 +293,4 @@ async def test_anthropic_tool_call_round_trips_thinking_blocks():
         {"type": "thinking", "thinking": "hidden chain", "signature": "sig_1"},
         {"type": "tool_use", "id": "call_1", "name": "lookup", "input": {"q": "x"}},
     ]
+

verifiers/types.py

@@ -214,7 +214,7 @@ class ResponseTokens(CustomBaseModel):
     completion_logprobs: list[float]
     routed_experts: RoutedExpertsPayload | None = None
     # Renderer-emitted multimodal sidecar (renderers.base.MultiModalData)
-    # carrying processed pixel_values / placeholder ranges per modality.
+    # carrying raw image descriptors / placeholder ranges per modality.
     # Populated by the renderer client when the rollout went through a
     # multimodal-aware renderer; ``None`` otherwise. Stored as ``Any`` to
     # avoid a hard import dependency on ``renderers`` at this layer.
@@ -261,7 +261,7 @@ class TrajectoryStepTokens(TypedDict):
     is_truncated: bool
     routed_experts: RoutedExpertsPayload | None
     # Renderer-emitted multimodal sidecar (renderers.base.MultiModalData)
-    # carrying processed pixel_values / placeholder ranges per modality.
+    # carrying raw image descriptors / placeholder ranges per modality.
     # ``NotRequired`` because text-only rollouts (and non-renderer client
     # types) never populate it.
     multi_modal_data: NotRequired[Any]

verifiers/v1/clients/train.py

@@ -13,9 +13,9 @@
 from typing import Any
 
 from openai import AsyncOpenAI, OpenAIError
-from renderers import RenderedTokens
 from renderers import OverlongPromptError as RendererOverlongPromptError
-from renderers import RendererConfig
+from renderers import RenderedTokens, RendererConfig
+from renderers.base import is_multimodal
 
 from verifiers.v1.clients.client import SESSION_ID_HEADER, Client
 from verifiers.v1.dialects import FINISH_REASONS, ChatDialect, Dialect, parse_tools
@@ -170,16 +170,6 @@ def _is_valid_incremental_tail(messages: list[dict[str, Any]]) -> bool:
     return all(role == "tool" for role in roles)
 
 
-def _has_multimodal_content(messages) -> bool:
-    for message in messages:
-        content = getattr(message, "content", None)
-        if not isinstance(content, list):
-            continue
-        if any(getattr(part, "type", None) == "image_url" for part in content):
-            return True
-    return False
-
-
 class TrainClient(Client):
     """Renders prompts to token ids and calls a vLLM `/inference/v1/generate` engine."""
 
@@ -266,23 +256,24 @@ async def get_response(
         )
         bridged_turn: PendingTurn | None = None
 
-        # Only build the (O(context)) previous-turn token ids once the cheap guards pass — a
-        # multimodal prompt or a tail that isn't a clean `[tool*, user?]` extension can't bridge.
-        can_bridge = (
-            turn is not None
-            and not _has_multimodal_content(prompt)
-            and _is_valid_incremental_tail(wire_messages)
-        )
+        # Only build the (O(context)) previous-turn token ids once the cheap guards pass: a
+        # tail that isn't a clean `[tool*, user?]` extension can't bridge.
+        can_bridge = turn is not None and _is_valid_incremental_tail(wire_messages)
         previous_ids = turn.previous_token_ids() if can_bridge else None
         if previous_ids is not None:
             previous_prompt_ids, previous_completion_ids = previous_ids
 
             def bridge():
+                kwargs: dict[str, Any] = {"tools": wire_tools}
+                if is_multimodal(renderer):
+                    kwargs["previous_multi_modal_data"] = (
+                        turn.previous_multi_modal_data()
+                    )
                 return renderer.bridge_to_next_turn(
                     previous_prompt_ids,
                     previous_completion_ids,
                     wire_messages,
-                    tools=wire_tools,
+                    **kwargs,
                 )
 
             bridged = await _maybe_offload(renderer, bridge)

verifiers/v1/graph.py

 MessageNode.model_construct( 
 MessageNode.model_construct( 
@@ -31,13 +31,15 @@
 
 from verifiers.v1.types import (
     AssistantMessage,
+    FinishReason,
     KeptTokens,
     Message,
     Response,
     StrictBaseModel,
     TextContentPart,
     Tool,
     ToolMessage,
+    Usage,
 )
 
 if TYPE_CHECKING:
@@ -62,6 +64,46 @@ def _decode_ndarray(d: dict) -> np.ndarray:
     return np.frombuffer(d["data"], dtype=np.dtype(d["dtype"])).reshape(d["shape"])
 
 
+_PROCESSED_MM_KEYS = frozenset({"pixel_values", "image_embeds", "image_features"})
+
+
+def _contains_processed_mm_key(value: Any) -> bool:
+    if isinstance(value, dict):
+        return bool(_PROCESSED_MM_KEYS.intersection(value)) or any(
+            _contains_processed_mm_key(v) for v in value.values()
+        )
+    if isinstance(value, (list, tuple)):
+        return any(_contains_processed_mm_key(v) for v in value)
+    return False
+
+
+def _validate_raw_mm_item(item: Any) -> dict[str, Any]:
+    if not isinstance(item, dict):
+        raise TypeError(
+            "v1 multimodal sidecars must be raw image descriptor dicts, "
+            f"got {type(item).__name__}"
+        )
+    if _contains_processed_mm_key(item):
+        raise TypeError(
+            "v1 multimodal sidecars must be raw image descriptors, "
+            "not processed multimodal payloads"
+        )
+    if not isinstance(item.get("raw_image_data"), str) or not item["raw_image_data"]:
+        raise ValueError("v1 multimodal sidecars require inline raw_image_data")
+    return dict(item)
+
+
+def _validate_raw_mm_data(mmd: MultiModalData) -> MultiModalData:
+    return MultiModalData(
+        mm_hashes={k: list(v) for k, v in mmd.mm_hashes.items()},
+        mm_placeholders={k: list(v) for k, v in mmd.mm_placeholders.items()},
+        mm_items={
+            modality: [_validate_raw_mm_item(item) for item in items]
+            for modality, items in mmd.mm_items.items()
+        },
+    )
+
+
 class MessageNode(StrictBaseModel):
     """One message in the graph: a message plus the tokens it adds to the cumulative
     sequence. Concatenating a root→leaf path's nodes reconstructs that branch's full token
@@ -102,12 +144,19 @@ class MessageNode(StrictBaseModel):
     logprobs: list[float] = Field(default_factory=list)
     """Sampling logprobs for the sampled tokens — length equals the number of True entries in
     `mask`; empty for input messages."""
+    finish_reason: FinishReason = None
+    """The response's finish reason (assistant nodes only) — kept for truncation detection."""
     multi_modal_data: SkipJsonSchema[MultiModalData | None] = None
-    """The renderer items for the images this message's content introduces (pixel tensors,
-    grids, hashes, placeholders) — the only carrier of the pixels from the env server to the
-    trainer. `Branch.multi_modal_data` concatenates them along the path into the training
-    `mm_kwargs`. Rides the wire as raw bytes (msgpack `bin`) since pydantic can't JSON the numpy;
-    kept off disk by the dump-site `exclude` in prime-rl (the tensors bloat the rollout jsonl)."""
+    """The renderer items for images this message introduces.
+
+    With the raw-image path, items are descriptors (hashes, grid metadata, and the inline
+    base64 image source), not image processor tensors. `Branch.multi_modal_data` concatenates
+    them along the path for the trainer. Old processed-payload sidecars are rejected.
+    """
+    usage: Usage | None = None
+    """Provider-reported token usage for this message's response (assistant nodes). Preserved on
+    the wire and on disk so dashboards can show token counts and cost even when the endpoint
+    returns no token ids."""
     routed_experts: SkipJsonSchema[np.ndarray | None] = None
     """This node's slice of the MoE expert-routing array — uint8 `[len(token_ids), layers,
     top_k]`, the expert ids inference selected for exactly this node's tokens. Attributed from
@@ -125,46 +174,48 @@ class MessageNode(StrictBaseModel):
 
     @field_serializer("multi_modal_data")
     def serialize_multi_modal_data(self, mmd: MultiModalData | None) -> dict | None:
-        """`MultiModalData` -> msgpack-safe dict so the pixel tensors ride the wire; numpy
-        `mm_items` values become raw-bytes `__nd__` dicts (every renderer emits `return_tensors="np"`)."""
+        """`MultiModalData` -> msgpack-safe raw descriptor dict."""
         if mmd is None:
             return None
+        mmd = _validate_raw_mm_data(mmd)
         return {
             "mm_hashes": {k: list(v) for k, v in mmd.mm_hashes.items()},
             "mm_placeholders": {
                 modality: [{"offset": p.offset, "length": p.length} for p in ranges]
                 for modality, ranges in mmd.mm_placeholders.items()
             },
             "mm_items": {
-                modality: [
-                    {k: _encode_ndarray(v) for k, v in item.items()} for item in items
-                ]
+                modality: [dict(item) for item in items]
                 for modality, items in mmd.mm_items.items()
             },
         }
 
     @field_validator("multi_modal_data", mode="before")
     @classmethod
     def deserialize_multi_modal_data(cls, value: Any) -> MultiModalData | None:
-        if value is None or isinstance(value, MultiModalData):
+        if value is None:
             return value
+        if isinstance(value, MultiModalData):
+            return _validate_raw_mm_data(value)
         if not isinstance(value, dict):
             raise TypeError(f"cannot build MultiModalData from {type(value).__name__}")
-        return MultiModalData(
-            mm_hashes={k: list(v) for k, v in (value.get("mm_hashes") or {}).items()},
-            mm_placeholders={
-                modality: [
-                    PlaceholderRange(offset=p["offset"], length=p["length"])
-                    for p in ranges
-                ]
-                for modality, ranges in (value.get("mm_placeholders") or {}).items()
-            },
-            mm_items={
-                modality: [
-                    {k: _decode_ndarray(v) for k, v in item.items()} for item in items
-                ]
-                for modality, items in (value.get("mm_items") or {}).items()
-            },
+        return _validate_raw_mm_data(
+            MultiModalData(
+                mm_hashes={
+                    k: list(v) for k, v in (value.get("mm_hashes") or {}).items()
+                },
+                mm_placeholders={
+                    modality: [
+                        PlaceholderRange(offset=p["offset"], length=p["length"])
+                        for p in ranges
+                    ]
+                    for modality, ranges in (value.get("mm_placeholders") or {}).items()
+                },
+                mm_items={
+                    modality: list(items)
+                    for modality, items in (value.get("mm_items") or {}).items()
+                },
+            )
         )
 
     @field_serializer("routed_experts")
@@ -332,6 +383,23 @@ def prompt_message_spans(
             for span in tail_spans
         ]
 
+    def previous_multi_modal_data(self) -> MultiModalData | None:
+        """Concatenate multimodal sidecars attached to the reusable prefix."""
+        merged = MultiModalData()
+        found = False
+        for nid in self.prefix_node_ids:
+            mmd = self.trace.nodes[nid].multi_modal_data
+            if mmd is None or mmd.is_empty():
+                continue
+            found = True
+            for modality, items in mmd.mm_items.items():
+                merged.mm_items.setdefault(modality, []).extend(items)
+            for modality, hashes in mmd.mm_hashes.items():
+                merged.mm_hashes.setdefault(modality, []).extend(hashes)
+            for modality, placeholders in mmd.mm_placeholders.items():
+                merged.mm_placeholders.setdefault(modality, []).extend(placeholders)
+        return merged if found else None
+
     def commit(self, response: Response, tools: list[Tool] | None = None) -> int:
         """Add this turn to the graph; returns the committed assistant node's id."""
         assistant_id = _commit_turn(self, response)
@@ -392,8 +460,9 @@ def _attribute_mm(
     renderer emits items per modality in prompt order (message order, then content-part order),
     so we walk the path advancing a per-modality cursor over every message's media but write
     only the nodes created this turn — `path[:num_reused]` is the reused prefix, already
-    attributed when first created. Item order is all training needs; placeholder offsets aren't
-    carried."""
+    attributed when first created. Each node gets the hashes/items/placeholders for exactly the
+    media it introduced, preserving vLLM multimodal-list alignment when those node sidecars are
+    later merged for bridge or training."""
     if mmd is None or mmd.is_empty():
         return
     cursors: dict[str, int] = {}
@@ -403,6 +472,7 @@ def _attribute_mm(
             continue
         node_items: dict[str, list] = {}
         node_hashes: dict[str, list] = {}
+        node_placeholders: dict[str, list[PlaceholderRange]] = {}
         for part in content:
             modality = _part_modality(part)
             if modality is None:
@@ -414,13 +484,20 @@ def _attribute_mm(
                 continue
             items = mmd.mm_items.get(modality) or []
             hashes = mmd.mm_hashes.get(modality) or []
+            placeholders = mmd.mm_placeholders.get(modality) or []
             if k < len(items):
                 node_items.setdefault(modality, []).append(items[k])
             if k < len(hashes):
                 node_hashes.setdefault(modality, []).append(hashes[k])
-        if node_items:
-            trace.nodes[node_id].multi_modal_data = MultiModalData(
-                mm_items=node_items, mm_hashes=node_hashes
+            if k < len(placeholders):
+                node_placeholders.setdefault(modality, []).append(placeholders[k])
+        if node_items or node_hashes or node_placeholders:
+            trace.nodes[node_id].multi_modal_data = _validate_raw_mm_data(
+                MultiModalData(
+                    mm_items=node_items,
+                    mm_hashes=node_hashes,
+                    mm_placeholders=node_placeholders,
+                )
             )
 
 

verifiers/v1/interception/server.py

@@ -33,9 +33,9 @@
 from pydantic import TypeAdapter, ValidationError
 from pydantic_core import PydanticSerializationError, from_json, to_json
 
+from verifiers.v1 import graph
 from verifiers.v1.dialects import DIALECTS, Dialect
 from verifiers.v1.dialects.base import is_sse_done_event
-from verifiers.v1 import graph
 from verifiers.v1.errors import (
     OverlongPromptError,
     ProviderError,