Keep multimodal content inline in v1 rollouts

tests/v1/test_graph.py

@@ -1,6 +1,7 @@
 import base64
 
 import numpy as np
+import pytest
 
 import verifiers.v1 as vf
 from verifiers.v1 import graph
@@ -426,6 +427,92 @@ def second_turn(trace, prompt_ids):
     ]
 
 
+def test_expanded_prompt_is_attributed_while_bridge_uses_logical_tokens():
+    trace = vf.Trace(
+        agent=vf.AgentInfo(config=vf.AgentConfig()),
+        task=vf.TraceTask(type="Task", data=vf.TaskData(idx=0, prompt="x")),
+    )
+    user = vf.UserMessage(content="image")
+    assistant = vf.AssistantMessage(content="looked")
+    assert graph.MessageNode(message=user, token_ids=[9]).logical_ids == [9]
+    assert (
+        graph.MessageNode(
+            message=user, token_ids=[9], renderer_token_ids=[]
+        ).logical_ids
+        == []
+    )
+    graph.prepare_turn(trace, [user]).commit(
+        vf.Response(
+            id="a",
+            created=0,
+            model="t",
+            message=assistant,
+            finish_reason="stop",
+            tokens=TurnTokens(
+                prompt_ids=[1, 9, 9, 3],
+                renderer_prompt_ids=[1, 9, 3],
+                completion_ids=[4],
+                message_spans=[(0, 2)],
+                is_content=[False, True, False],
+                mm_placeholders=[(1, 2)],
+                mm_token_type_id_map={9: 1},
+            ),
+        )
+    )
+
+    assert trace.branches[0].token_ids == [1, 9, 9, 3, 4]
+    assert trace.branches[0].mm_token_type_ids == [0, 1, 1, 0, 0]
+    assert trace.nodes[0].token_ids == [1, 9, 9]
+    assert trace.nodes[0].is_content == [False, True, True]
+    turn = graph.prepare_turn(trace, [user, assistant, vf.UserMessage(content="next")])
+    assert turn.previous_renderer_token_ids() == ([1, 9, 3], [4])
+
+    with pytest.raises(ValueError, match="exactly extend"):
+        turn.commit(
+            vf.Response(
+                id="b",
+                created=0,
+                model="t",
+                message=vf.AssistantMessage(content="bad"),
+                finish_reason="stop",
+                tokens=TurnTokens(
+                    prompt_ids=[1, 9, 8, 3, 4, 5],
+                    renderer_prompt_ids=[1, 9, 3, 4, 5],
+                    bridged=True,
+                    completion_ids=[6],
+                    message_spans=[None, None, (4, 5)],
+                ),
+            )
+        )
+
+    tool = vf.ToolMessage(content="result", tool_call_id="call_0")
+    graph.prepare_turn(trace, [user, assistant, tool]).commit(
+        vf.Response(
+            id="b",
+            created=0,
+            model="t",
+            message=vf.AssistantMessage(content="done"),
+            finish_reason="stop",
+            tokens=TurnTokens(
+                prompt_ids=[1, 9, 9, 3, 4, 5, 6, 7],
+                renderer_prompt_ids=[1, 9, 3, 4, 5, 6, 7],
+                bridged=True,
+                completion_ids=[8],
+                message_spans=[None, None, (4, 6)],
+                is_content=[False, True, False, True, False, True, False],
+                mm_placeholders=[(1, 2)],
+                mm_token_type_id_map={9: 1},
+            ),
+        )
+    )
+
+    tool_node = trace.nodes[-2]
+    assert tool_node.message == tool
+    assert tool_node.token_ids == [5, 6]
+    assert tool_node.renderer_token_ids == [5, 6]
+    assert tool_node.is_content == [False, True]
+
+
 def test_prompt_supplied_assistant_messages_are_not_sampled_turns():
     task = vf.TaskData(idx=0, prompt="few-shot")
     trace = vf.Trace(

verifiers/v1/clients/train.py

@@ -135,6 +135,16 @@ def response_from_generate(
         message_spans = bridged_turn.prompt_message_spans(attribution)
     else:
         message_spans = attribution.message_token_spans()
+    raw_mm_placeholders = result.get("mm_placeholders")
+    mm_placeholders = (
+        sorted(
+            (placeholder["offset"], placeholder["length"])
+            for ranges in raw_mm_placeholders.values()
+            for placeholder in ranges
+        )
+        if raw_mm_placeholders is not None
+        else None
+    )
     return Response(
         id=result.get("request_id", ""),
         created=0,
@@ -154,11 +164,13 @@ def response_from_generate(
         # million-token contexts synchronously on the event loop.
         tokens=TurnTokens.model_construct(
             prompt_ids=prompt_ids,
+            renderer_prompt_ids=result.get("renderer_prompt_ids"),
+            bridged=bridged_turn is not None,
             completion_ids=completion_ids,
             completion_logprobs=result.get("completion_logprobs") or [],
             message_spans=message_spans,
             is_content=attribution.is_content if attribution is not None else None,
-            multi_modal_data=result.get("multi_modal_data"),
+            mm_placeholders=mm_placeholders,
             mm_token_type_id_map=mm_token_type_id_map,
             routed_experts=result.get("routed_experts"),
             sampling_mask=SamplingMask.from_sampling_mask(mask)
@@ -354,11 +366,9 @@ async def get_response(
         from renderers.client import generate
 
         wire_tools = [tool_to_wire(t) for t in tools] if tools else None
-        wire_messages = (
-            [message_to_wire(m) for m in turn.tail] if turn is not None else []
-        )
+        wire_messages = [message_to_wire(m) for m in prompt]
+        wire_tail = [message_to_wire(m) for m in turn.tail] if turn is not None else []
         prompt_ids: list[int] | None = None
-        multi_modal_data = None
         prompt_attribution: RenderedTokens | None = None
         model = body["model"]
         sampling_params = sampling.wire_args()
@@ -377,48 +387,52 @@ async def get_response(
             mm_token_type_id_map = (
                 renderer.mm_token_type_id_map if is_multimodal(renderer) else 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)
-            )
-            previous_ids = turn.previous_token_ids() if can_bridge else None
+            has_images = _has_multimodal_content(prompt)
+            process_multimodal = not has_images
+            if has_images and not getattr(
+                renderer, "supports_process_multimodal", False
+            ):
+                raise NotImplementedError(
+                    f"{type(renderer).__name__} does not support process_multimodal=False"
+                )
+            render_kwargs = {} if process_multimodal else {"process_multimodal": False}
+            # Only build the O(context) previous token stream for a bridgeable tail.
+            can_bridge = turn is not None and _is_valid_incremental_tail(wire_tail)
+            previous_ids = turn.previous_renderer_token_ids() if can_bridge else None
             if previous_ids is not None:
                 previous_prompt_ids, previous_completion_ids = previous_ids
 
                 def bridge():
                     return renderer.bridge_to_next_turn(
                         previous_prompt_ids,
                         previous_completion_ids,
-                        wire_messages,
+                        wire_tail,
                         tools=wire_tools,
+                        **render_kwargs,
                     )
 
                 bridged = await slot.run(bridge)
                 if bridged is not None:
                     prompt_ids = bridged.token_ids
-                    multi_modal_data = bridged.multi_modal_data
                     prompt_attribution = bridged
                     bridged_turn = turn
                     sampling_params["routed_experts_prompt_start"] = max(
-                        len(previous_prompt_ids) + len(previous_completion_ids) - 1,
-                        0,
+                        turn.path_len - 1, 0
                     )
 
             # Render here (encode-side, so through the slot) rather than inside `generate`:
             # handed prebuilt prompt_ids, generate's own renderer touches are decode-side
             # and stop-id reads, safe on a bare renderer without lock or thread hop.
             if prompt_ids is None:
-                wire_messages = [message_to_wire(m) for m in prompt]
                 rendered = await slot.run(
                     lambda: renderer.render(
-                        wire_messages, tools=wire_tools, add_generation_prompt=True
+                        wire_messages,
+                        tools=wire_tools,
+                        add_generation_prompt=True,
+                        **render_kwargs,
                     )
                 )
                 prompt_ids = rendered.token_ids
-                multi_modal_data = rendered.multi_modal_data
                 prompt_attribution = rendered
 
             try:
@@ -428,10 +442,10 @@ def bridge():
                     messages=wire_messages,
                     model=model,
                     prompt_ids=prompt_ids,
-                    multi_modal_data=multi_modal_data,
                     prompt_attribution=prompt_attribution,
                     tools=wire_tools,
                     sampling_params=sampling_params,
+                    process_multimodal=process_multimodal,
                     cache_salt=cache_salt,
                     extra_headers={SESSION_ID_HEADER: session_id}
                     if session_id