feat(v1): store engine weight-version spans on assistant nodes

verifiers/v1/clients/train.py

@@ -30,6 +30,7 @@
     ToolCall,
     TurnTokens,
     Usage,
+    WeightVersionSpan,
 )
 
 logger = logging.getLogger(__name__)
@@ -181,6 +182,9 @@ def response_from_generate(
             sampling_mask=SamplingMask.from_sampling_mask(mask)
             if (mask := result.get("sampling_mask"))
             else None,
+            weight_versions=[WeightVersionSpan(**span) for span in spans]
+            if (spans := result.get("weight_versions"))
+            else None,
         ),
     )
 

verifiers/v1/graph.py

@@ -49,6 +49,7 @@
     TextContentPart,
     Tool,
     ToolMessage,
+    WeightVersionSpan,
 )
 
 if TYPE_CHECKING:
@@ -164,6 +165,10 @@ class MessageNode(BaseModel):
     `ids` stores the flat token ids and `counts` stores each token's row size. Assistant
     nodes only. The arrays serialize as raw-byte `__nd__` dictionaries.
     """
+    weight_versions: list[WeightVersionSpan] | None = None
+    """Engine weight-version spans over this node's sampled tokens (indexes into the
+    completion, as vLLM returns them). Assistant nodes only; None when the engine labels
+    none."""
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
@@ -997,6 +1002,7 @@ def _commit_turn(turn: PendingTurn, response: Response) -> int:
             else [],
             # TurnTokens is discarded after commit, so transfer its logprobs without copying.
             logprobs=tokens.completion_logprobs if tokens else [],
+            weight_versions=tokens.weight_versions if tokens else None,
         )
     )
     assistant_id = len(trace.nodes) - 1

verifiers/v1/types.py

@@ -201,6 +201,14 @@ class RoutedExperts(TypedDict):
     start: int
 
 
+class WeightVersionSpan(BaseModel):
+    """Completion tokens `[start, end)` were sampled by the engine weights labelled `version`."""
+
+    version: str
+    start: int
+    end: int
+
+
 @dataclass
 class SamplingMask:
     """Sampling masks stored as flat int32 `ids` and `counts` arrays.
@@ -276,6 +284,9 @@ class TurnTokens(BaseModel):
     # Transient carrier (excluded): per-completion-token sampling masks,
     # attributed to the assistant node by the turn's `commit`, then dropped.
     sampling_mask: SamplingMask | None = Field(default=None, exclude=True)
+    # Transient carrier (excluded): the engine's weight-version spans over the completion,
+    # attributed to the assistant node by the turn's `commit`, then dropped.
+    weight_versions: list[WeightVersionSpan] | None = Field(default=None, exclude=True)
 
 
 class Response(BaseModel):