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
@@ -232,6 +233,53 @@ def second_turn(trace, prompt_ids):
     ]
 
 
+def test_expanded_prompt_is_canonical_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")
+    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, 2, 3],
+                completion_ids=[4],
+                message_spans=[(0, 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]
+    turn = graph.prepare_turn(trace, [user, assistant, vf.UserMessage(content="next")])
+    assert turn.previous_token_ids() == ([1, 2, 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, 2, 3, 4, 5],
+                    completion_ids=[6],
+                    message_spans=[None, None, (4, 5)],
+                ),
+            )
+        )
+
+
 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

@@ -152,11 +152,11 @@ 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"),
             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_token_type_id_map=mm_token_type_id_map,
             routed_experts=result.get("routed_experts"),
             kept_tokens=KeptTokens(**kept)
@@ -352,11 +352,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()
@@ -374,13 +372,15 @@ 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)
-            )
+            raw_multimodal = _has_multimodal_content(prompt)
+            if raw_multimodal and not getattr(
+                renderer, "supports_raw_multimodal", False
+            ):
+                raise NotImplementedError(
+                    f"{type(renderer).__name__} does not support raw multimodal inference"
+                )
+            # 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_token_ids() if can_bridge else None
             if previous_ids is not None:
                 previous_prompt_ids, previous_completion_ids = previous_ids
@@ -389,33 +389,33 @@ def bridge():
                     return renderer.bridge_to_next_turn(
                         previous_prompt_ids,
                         previous_completion_ids,
-                        wire_messages,
+                        wire_tail,
                         tools=wire_tools,
+                        **({"process_multimodal": False} if raw_multimodal else {}),
                     )
 
                 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,
+                        **({"process_multimodal": False} if raw_multimodal else {}),
                     )
                 )
                 prompt_ids = rendered.token_ids
-                multi_modal_data = rendered.multi_modal_data
                 prompt_attribution = rendered
 
             try:
@@ -425,10 +425,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,
+                    raw_multimodal=raw_multimodal,
                     extra_headers={SESSION_ID_HEADER: session_id}
                     if session_id
                     else None,