diff --git a/verifiers/v1/README.md b/verifiers/v1/README.md index e04b6f52a5..79f3662229 100644 --- a/verifiers/v1/README.md +++ b/verifiers/v1/README.md @@ -16,7 +16,10 @@ abstractions and on-disk output. Everything is pydantic-typed; `import verifiers - **Minimal & pythonic** — the high-level abstractions without the implementation bulk; plain classes + decorators (`@vf.reward` / `@vf.metric` / ...). - **Training-ready traces** — exact token ids + logprobs straight from an agentic rollout - (renderer client), with branching recovered for compaction / subagents. + (renderer client); one training sample per branch, recovered for compaction / subagents. +- **Delta-native trace graph** — each message is stored once as a node linked to its + predecessor, so a trace's size is linear in turns, not quadratic; branches fall out of + walking the graph, and a training sample is a cheap concat of node tokens along a path. - **Hub-native + v0-compatible** — ids install on demand from the Environments Hub, and classic v0 envs run through the same CLIs via a bridge. @@ -145,9 +148,9 @@ uv run eval code-golf-v1 -n 1 -r 2 # group rewards: a @vf.group_reward scores N A rollout isn't always linear. The `compact` harness rewrites its context every turn — a fresh `[system, user]` carrying its running notes plus the last tool output — so each turn -is its own *branch*. `branching` recovers them from the flat trajectory and -`trace.branches` / `num_branches` expose it (a linear harness is one branch; the compact -harness is one per turn — it also handles subagents): +is its own *branch*. Branches fall out of the message graph — each leaf's root→leaf path is +one branch, exposed by `trace.branches` / `num_branches` (a linear harness is one branch; +the compact harness is one per turn — it also handles subagents): ```bash uv run eval wiki-search-v1 -n 1 --harness.id compact # fresh prompt each turn → num_branches == turns @@ -164,10 +167,11 @@ uv run eval gsm8k-v1 -n 1 --client.type renderers \ # renderers: client-side to --client.base-url http://localhost:8000/v1 # token-in/out traces (needs a vLLM engine) ``` -With `renderers`, each `trace.trajectory[i].tokens` carries the exact `prompt_ids` / -`completion_ids` / `completion_logprobs` the engine saw — training-ready token data -straight from an agentic rollout, with zero agent changes. (When the engine returns ids on -the response itself, the openai client picks them up too — no renderer required.) +With `renderers`, each graph node carries the exact tokens the engine saw — `token_ids` +plus a per-token trainable `mask` and `logprobs` — so concatenating a branch's nodes is a +ready training sample, straight from an agentic rollout with zero agent changes. (When the +engine returns ids on the response itself, the openai client picks them up too — no +renderer required.) ### Limits & retries diff --git a/verifiers/v1/__init__.py b/verifiers/v1/__init__.py index 0c12a35a21..6f54584eab 100644 --- a/verifiers/v1/__init__.py +++ b/verifiers/v1/__init__.py @@ -43,14 +43,13 @@ from verifiers.v1.task import Resources, Task, WireTask from verifiers.v1.taskset import Taskset, TasksetConfig, ToolsConfig from verifiers.v1.tools import Tools, run_mcp_server +from verifiers.v1.graph import MessageNode from verifiers.v1.trace import ( Branch, Error, TimeSpan, Timing, Trace, - Turn, - TurnTokens, ) from verifiers.v1.types import ( AssistantMessage, @@ -62,6 +61,7 @@ SystemMessage, Tool, ToolCall, + TurnTokens, ToolMessage, Usage, UserMessage, @@ -90,7 +90,7 @@ "WireTask", "Resources", "Trace", - "Turn", + "MessageNode", "Branch", "TurnTokens", "Timing", diff --git a/verifiers/v1/branching.py b/verifiers/v1/branching.py deleted file mode 100644 index c8e91b0984..0000000000 --- a/verifiers/v1/branching.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Branch detection: split a flat trajectory into the linear histories that made it. - -A rollout records every model call in one flat `trajectory`, but the conversation -isn't always linear. An harness may **compact** its context (a turn's prompt drops -earlier history) or run **subagents** (several independent histories run concurrently -in one trajectory). Each maximal linear history is a **branch**. - -We recover branches with one rule — *a turn extends a branch when that branch's last -prompt+response is a prefix of the turn's prompt* — applied with the more reliable -signal available: - - 1. **tokens** — prefix match on token ids (prompt_ids/completion_ids), when every - turn carries them (renderer client). Robust to message reformatting. - 2. **messages** — prefix match on the messages themselves. Assumes the harness echoes - prior turns **byte-identically** (`reasoning_content` excepted — it never - round-trips through a prompt; `None`/`""` content compare equal). Our built-in - harnesses do; some external ones may not. - -Every branch stays active and a turn extends the **longest** matching branch, so -compaction overlap and concurrent subagent branches resolve correctly. These functions -are pure (no Trace import at runtime) and test in isolation. -""" - -from __future__ import annotations - -import operator -from collections.abc import Callable -from typing import TYPE_CHECKING, Any - -from verifiers.v1.types import AssistantMessage, Message, ToolMessage - -if TYPE_CHECKING: - from verifiers.v1.trace import Turn - - -def same_message(a: Message, b: Message) -> bool: - """Equality on the fields that round-trip through a prompt: role, content, tool - calls (assistant) / call id (tool). Ignores `reasoning_content`, and treats `None` - and `""` content as equal — an external harness may echo a tool-call-only assistant - message with either, and that shouldn't split a branch.""" - if type(a) is not type(b) or (a.content or "") != (b.content or ""): - return False - if isinstance(a, AssistantMessage): - assert isinstance(b, AssistantMessage) - calls_a, calls_b = a.tool_calls or [], b.tool_calls or [] - return len(calls_a) == len(calls_b) and all( - x.id == y.id and x.name == y.name and x.arguments == y.arguments - for x, y in zip(calls_a, calls_b) - ) - if isinstance(a, ToolMessage): - assert isinstance(b, ToolMessage) - return a.tool_call_id == b.tool_call_id - return True # system / user: content already matched - - -def segment(turns: list[Turn]) -> list[list[int]]: - """Group turn indices into branches — each a list of indices in order, possibly - non-contiguous (concurrent subagent branches). Uses token-id prefixes when every turn - carries them, else message prefixes. It builds, per turn, the `head` (the prompt that - must extend a branch) and the `seq` (prompt+response, a branch's running prefix), then - matches.""" - if turns and all(t.tokens and t.tokens.prompt_ids for t in turns): - heads = [list(t.tokens.prompt_ids) for t in turns] - seqs = [[*t.tokens.prompt_ids, *t.tokens.completion_ids] for t in turns] - eq: Callable[[Any, Any], bool] = operator.eq - else: - heads = [list(t.prompt) for t in turns] - seqs = [[*t.prompt, t.response.message] for t in turns] - eq = same_message - return _forest(heads, seqs, eq) - - -def _forest( - heads: list[list], seqs: list[list], eq: Callable[[Any, Any], bool] -) -> list[list[int]]: - """Longest-prefix multi-match: a turn joins the active branch whose running `seq` is - the longest prefix of the turn's `head`, else starts a new branch. Keeping every - branch active + longest-match is what resolves compaction overlap and concurrent - subagent branches.""" - branches: list[list[int]] = [] - prefixes: list[list] = [] # parallel to branches: each branch's running seq - for i, head in enumerate(heads): - best, best_len = -1, -1 - for b, prefix in enumerate(prefixes): - if ( - len(prefix) > best_len - and len(head) >= len(prefix) - and all(eq(p, h) for p, h in zip(prefix, head)) - ): - best, best_len = b, len(prefix) - if best >= 0: - branches[best].append(i) - prefixes[best] = seqs[i] - else: - branches.append([i]) - prefixes.append(seqs[i]) - return branches diff --git a/verifiers/v1/cli/dashboard.py b/verifiers/v1/cli/dashboard.py index a8351876c1..45e51362c3 100644 --- a/verifiers/v1/cli/dashboard.py +++ b/verifiers/v1/cli/dashboard.py @@ -81,15 +81,10 @@ def _tokens(trace: Trace) -> str: final context the model saw. (Output can exceed the final context — reasoning tokens count toward completions but aren't re-fed — so it's not derived by subtraction.)""" branches = trace.branches - if not branches or not branches[0].turns: + if not branches or not branches[0].nodes: return "" - usages = [ - t.response.usage for t in branches[0].turns if t.response.usage is not None - ] - if not usages: - return "" - output = sum(u.completion_tokens for u in usages) - return f"{format_count(usages[-1].prompt_tokens)}/{format_count(output)} tokens" + b = branches[0] + return f"{format_count(b.prompt_len)}/{format_count(b.completion_len)} tokens" def _groups(rollouts: list[Rollout]) -> list[list[Rollout]]: diff --git a/verifiers/v1/clients/renderer.py b/verifiers/v1/clients/renderer.py index b2d114e82a..ce7db26f43 100644 --- a/verifiers/v1/clients/renderer.py +++ b/verifiers/v1/clients/renderer.py @@ -62,6 +62,12 @@ def response_from_generate(result: dict, model: str) -> Response: ] or None prompt_ids = result.get("prompt_ids") or [] completion_ids = result.get("completion_ids") or [] + # Per-message token spans (the renderer's attribution) let the trace graph store each + # message's tokens once; carried transiently on TurnTokens and consumed by `graph.add_turn`. + attribution = result.get("prompt_attribution") + message_spans = ( + attribution.message_token_spans() if attribution is not None else None + ) return Response( id=result.get("request_id", ""), created=0, @@ -79,6 +85,7 @@ def response_from_generate(result: dict, model: str) -> Response: prompt_ids=prompt_ids, completion_ids=completion_ids, completion_logprobs=result.get("completion_logprobs") or [], + message_spans=message_spans, ), ) diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py new file mode 100644 index 0000000000..639fca5e18 --- /dev/null +++ b/verifiers/v1/graph.py @@ -0,0 +1,179 @@ +"""Message-graph trajectory: store each message once, recover branches by walking. + +A rollout is a graph of `MessageNode`s — one per distinct message, each linked to its +predecessor. The conversation is a path from a root to a leaf; branches (compaction, +subagents) are simply multiple leaves, so branching falls out of the walk. Each node stores +only the tokens it *adds* to the cumulative sequence, keeping size linear in turns and +making a branch's training sample a cheap concat of node `token_ids`/`mask`/`logprobs` along +its path. + +Token attribution (renderer client): the renderer reports, per prompt, each message's token +span (`RenderedTokens.message_token_spans()`, carried on `TurnTokens.message_spans`). A new +input message's node gets its span plus the leading template scaffold since the previous +message; the trailing scaffold (the generation prompt) goes on the assistant node, prefixed +to its sampled completion. By construction `concat(node.token_ids along a path)` reproduces +the exact `prompt_ids + completion_ids` the model saw. +""" + +from __future__ import annotations + +import hashlib +from typing import TYPE_CHECKING + +from pydantic import Field + +from verifiers.v1.types import ( + AssistantMessage, + FinishReason, + Message, + Response, + StrictBaseModel, + ToolMessage, +) + +if TYPE_CHECKING: + from verifiers.v1.trace import Branch, Trace + + +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 + sequence; the mask/logprobs make it a training sample.""" + + parent: int | None = None + """Index into `Trace.nodes` of the predecessor message; None for a root.""" + message: Message + """The message this node carries (system / user / assistant / tool).""" + token_ids: list[int] = Field(default_factory=list) + """This message's delta contribution to the cumulative token sequence: its leading + template scaffold + its own tokens — for an assistant, the generation-prompt scaffold + followed by the sampled completion. Concatenated along a path, these reproduce the exact + `prompt_ids + completion_ids` the model saw.""" + mask: list[bool] = Field(default_factory=list) + """Per-token, parallel to `token_ids`: True for trainable, model-sampled tokens (only an + assistant node's completion span); False for template scaffold and every input-message + token.""" + 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.""" + + +def message_hash(message: Message) -> str: + """Stable content hash on the fields that round-trip through a prompt — role, content + (None and "" equal), assistant tool calls, tool call id; `reasoning_content` ignored. + Two messages hash equal iff they're the same conversational message, so a re-stated + prefix message dedups to one node. The dedup key for sharing a prefix across + turns/branches; salt-free so it is identical across processes and after deserialization.""" + parts: list[str] = [type(message).__name__, message.content or ""] + if isinstance(message, AssistantMessage): + for tc in message.tool_calls or []: + parts += [tc.id, tc.name, tc.arguments] + elif isinstance(message, ToolMessage): + parts.append(message.tool_call_id) + return hashlib.blake2b("\x00".join(parts).encode(), digest_size=16).hexdigest() + + +def _head_index(trace: "Trace") -> dict[tuple[int | None, str], int]: + """`(parent, msg_hash) -> node_id`, rebuilt lazily from `nodes` after deserialization.""" + if not trace._head_index and trace.nodes: + trace._head_index = { + (node.parent, message_hash(node.message)): nid + for nid, node in enumerate(trace.nodes) + } + return trace._head_index + + +def add_turn(trace: "Trace", prompt: "list[Message]", response: Response) -> None: + """Insert one model turn (its prompt messages + its response) into the graph. Reuses any + existing prefix nodes (by `(parent, hash)`), creates a node per new message attributing + its tokens, and appends a fresh assistant node holding the generation-prompt scaffold + + the sampled completion. + + Token attribution anchors new tokens to the cumulative *stored* length of the reused + prefix (`path_len`), not message spans — the previous assistant's closing scaffold lives + in its later input-form span but not its stored generation form, so anchoring on spans + would drop it. The new tokens (`prompt_ids[path_len:]`) are split among the new input + messages by span (leading template scaffold folds into the following message), and the + trailing generation prompt goes on the assistant node before its sampled completion. By + construction `concat(node.token_ids along the path) == prompt_ids + completion_ids`.""" + tokens = response.tokens + prompt_ids = list(tokens.prompt_ids) if tokens else [] + spans = tokens.message_spans if tokens else None + idx = _head_index(trace) + + parent: int | None = None + path_len = 0 # cumulative stored token length of the reused prefix + # cursor: in prompt_ids, the end of the previous *new* message's tokens + cursor: int | None = None + for i, msg in enumerate(prompt): + key = (parent, message_hash(msg)) + existing = idx.get(key) + if cursor is None and existing is not None: # still extending the shared prefix + parent = existing + path_len += len(trace.nodes[existing].token_ids) + continue + start = path_len if cursor is None else cursor + span = spans[i] if spans and i < len(spans) else None + end = span[1] if span else start + node_tokens = prompt_ids[start:end] + trace.nodes.append( + MessageNode( + parent=parent, + message=msg, + token_ids=node_tokens, + mask=[False] * len(node_tokens), + ) + ) + parent = len(trace.nodes) - 1 + idx[key] = parent + cursor = end + + # Assistant node: trailing scaffold (the generation prompt) + the sampled completion. + comp_ids = list(tokens.completion_ids) if tokens else [] + gen_start = path_len if cursor is None else cursor + gen_prompt = prompt_ids[gen_start:] + trace.nodes.append( + MessageNode( + parent=parent, + message=response.message, + token_ids=[*gen_prompt, *comp_ids], + mask=[False] * len(gen_prompt) + [True] * len(comp_ids), + logprobs=list(tokens.completion_logprobs) if tokens else [], + finish_reason=response.finish_reason, + ) + ) + # Register the assistant so the next turn's prompt (which restates it) reuses this node. + idx[(parent, message_hash(response.message))] = len(trace.nodes) - 1 + + +# --- walking the graph (views) --------------------------------------------------------- + + +def _path_to(trace: "Trace", leaf: int) -> list[int]: + """Node ids from the root down to `leaf` (inclusive), in order.""" + path: list[int] = [] + nid: int | None = leaf + while nid is not None: + path.append(nid) + nid = trace.nodes[nid].parent + path.reverse() + return path + + +def leaves(trace: "Trace") -> list[int]: + """Node ids that are no node's parent — one per branch (the last node of each).""" + has_child = {n.parent for n in trace.nodes if n.parent is not None} + return [i for i in range(len(trace.nodes)) if i not in has_child] + + +def branches_from_nodes(trace: "Trace") -> list["Branch"]: + """Each leaf's root→leaf node path becomes a `Branch` — one per leaf (one branch when + linear, several under compaction or subagents).""" + from verifiers.v1.trace import Branch + + return [ + Branch(index=i, nodes=[trace.nodes[nid] for nid in _path_to(trace, leaf)]) + for i, leaf in enumerate(leaves(trace)) + ] diff --git a/verifiers/v1/interception/server.py b/verifiers/v1/interception/server.py index aad6500c33..1ffd637572 100644 --- a/verifiers/v1/interception/server.py +++ b/verifiers/v1/interception/server.py @@ -1,9 +1,10 @@ """The interception server: harness chat-completions, caught and proxied. Every rollout runs an harness program whose OpenAI-style calls are caught here: a small -localhost server routes each `POST /v1/chat/completions` to our `Client`, records a -`Turn`, and returns the result in OpenAI shape. We inject `OPENAI_BASE_URL`/`OPENAI_API_KEY` -so the program's SDK talks to us. Chat completions only, no streaming. +localhost server routes each `POST /v1/chat/completions` to our `Client`, records the turn +into the trace's message graph, and returns the result in OpenAI shape. We inject +`OPENAI_BASE_URL`/`OPENAI_API_KEY` so the program's SDK talks to us. Chat completions only, +no streaming. One server multiplexes many rollouts: each rollout registers a `RolloutSession` under its own secret (the bearer token the harness already sends), and the server routes by that @@ -25,7 +26,8 @@ from aiohttp import web from verifiers.v1.clients import RolloutContext -from verifiers.v1.trace import Trace, Turn +from verifiers.v1 import graph +from verifiers.v1.trace import Trace from verifiers.v1.types import ( AssistantMessage, Message, @@ -265,9 +267,8 @@ async def handle_chat(self, request: web.Request) -> web.Response: except Exception as e: # surface to the program as an API error logger.warning("model call failed: id=%s %s", session.trace.id, e) return web.json_response({"error": str(e)}, status=502) - session.trace.trajectory.append( - Turn(prompt=prompt, response=response, tokens=response.tokens) - ) # branches are derived from the trajectory (see Trace.branches) + graph.add_turn(session.trace, prompt, response) # one node per new message; + # branches fall out of walking the graph (see Trace.branches / verifiers.v1.graph) last = response # Hand back to the program when the model wants a tool (the program runs it) or # when there's no user simulator to keep the conversation going. diff --git a/verifiers/v1/legacy.py b/verifiers/v1/legacy.py index 2705ab3c14..9f3fb170c9 100644 --- a/verifiers/v1/legacy.py +++ b/verifiers/v1/legacy.py @@ -30,7 +30,8 @@ RunRolloutResponse, ) from verifiers.v1.task import WireTask -from verifiers.v1.trace import Error, TimeSpan, Timing, Trace, Turn +from verifiers.v1 import graph +from verifiers.v1.trace import Error, TimeSpan, Timing, Trace from verifiers.v1.types import ( AssistantMessage, Response, @@ -186,20 +187,6 @@ def rollout_output_to_trace(out: dict, task_idx: int) -> Trace: system prompt / instruction / answer. ``is_truncated`` is a computed v1 field derived from the final turn's ``finish_reason`` and the stop condition.""" model = str(out.get("model") or "") - trajectory: list[Turn] = [] - for step in out.get("trajectory") or []: - if not isinstance(step, dict): - continue - # The renderer records tokens on both the turn (training reads these) and the - # response, mirroring the native v1 client; keep both so the trace is identical. - tokens = _to_v1_tokens(step.get("tokens")) - trajectory.append( - Turn( - prompt=_to_v1_messages(step.get("prompt")), - response=_to_v1_response(step.get("response"), model, tokens), - tokens=tokens, - ) - ) error = None raw_error = out.get("error") @@ -215,7 +202,6 @@ def rollout_output_to_trace(out: dict, task_idx: int) -> Trace: trace: Trace = Trace[WireTask]( task=_to_wire_task(task_idx, out.get("prompt"), out.get("answer")), - trajectory=trajectory, rewards={"reward": float(out.get("reward") or 0.0)}, metrics={k: float(v) for k, v in (out.get("metrics") or {}).items()}, is_completed=bool(out.get("is_completed", True)), @@ -223,6 +209,17 @@ def rollout_output_to_trace(out: dict, task_idx: int) -> Trace: errors=[error] if error else [], timing=_timing(out.get("timing")), ) + # Rebuild the message graph from the v0 steps. v0 tokens carry no per-message spans, so + # attribution is coarse (the per-turn delta lands on the assistant node) — still linear. + for step in out.get("trajectory") or []: + if not isinstance(step, dict): + continue + tokens = _to_v1_tokens(step.get("tokens")) + graph.add_turn( + trace, + _to_v1_messages(step.get("prompt")), + _to_v1_response(step.get("response"), model, tokens), + ) return trace diff --git a/verifiers/v1/rollout.py b/verifiers/v1/rollout.py index 67b1dafa74..d5b0f35eb7 100644 --- a/verifiers/v1/rollout.py +++ b/verifiers/v1/rollout.py @@ -193,7 +193,7 @@ async def run( trace.id, self.task.idx, trace.reward, - len(trace.trajectory), + trace.num_turns, trace.error.type if trace.error else trace.stop_condition, ) return trace diff --git a/verifiers/v1/trace.py b/verifiers/v1/trace.py index e19cf429e2..3ca5d0e5e9 100644 --- a/verifiers/v1/trace.py +++ b/verifiers/v1/trace.py @@ -17,15 +17,14 @@ from pydantic import Field, PrivateAttr, computed_field -from verifiers.v1 import branching +from verifiers.v1 import graph +from verifiers.v1.graph import MessageNode from verifiers.v1.task import TaskT from verifiers.v1.types import ( AssistantMessage, Messages, - Response, StrictBaseModel, ToolMessage, - TurnTokens, ) logger = logging.getLogger(__name__) @@ -61,62 +60,70 @@ class Error(StrictBaseModel): ) -class Turn(StrictBaseModel): - """One model turn: the prompt sent, the response, and optional token encoding.""" +class Branch(StrictBaseModel): + """A linear run of messages whose context grew without being rewritten — a root→leaf + path in the message graph. A conversation that compacts (or runs subagents) splits into + several branches; a linear one is a single branch. `messages` is the full conversation; + one training sample is built per branch.""" - prompt: Messages - response: Response - tokens: TurnTokens | None = None + index: int + nodes: list[MessageNode] + @computed_field @property - def num_prompt_tokens(self) -> int: - if self.tokens and self.tokens.prompt_ids: - return len(self.tokens.prompt_ids) - return self.response.usage.prompt_tokens if self.response.usage else 0 + def num_turns(self) -> int: + """Model turns (assistant messages) in this branch.""" + return sum(1 for n in self.nodes if isinstance(n.message, AssistantMessage)) @property - def num_completion_tokens(self) -> int: - if self.tokens and self.tokens.completion_ids: - return len(self.tokens.completion_ids) - return self.response.usage.completion_tokens if self.response.usage else 0 - - -class Branch(StrictBaseModel): - """A linear run of turns whose context grew without being rewritten. A trajectory - that compacts splits into several branches (see `branching`); a linear one is a - single branch. `messages` is the branch's full conversation — its last turn's - prompt plus that turn's response (the last prompt already holds the branch's - earlier turns).""" + def messages(self) -> Messages: + """The branch's full conversation, in order.""" + return [n.message for n in self.nodes] - index: int - turns: list[Turn] + @property + def token_ids(self) -> list[int]: + """The branch's full token sequence — every node's tokens concatenated in order + (final-turn prompt + every completion). The training sample's input ids.""" + return [t for node in self.nodes for t in node.token_ids] - @computed_field @property - def num_turns(self) -> int: - """Model turns in this branch.""" - return len(self.turns) + def sampled_mask(self) -> list[bool]: + """Per-token trainable flag aligned to `token_ids`: True for the model-sampled + (completion) tokens, False for prompt/template scaffold.""" + return [m for node in self.nodes for m in node.mask] @property - def prompt_len(self) -> int: - """Input context size: this branch's final-turn prompt (the full last context).""" - return self.turns[-1].num_prompt_tokens if self.turns else 0 + def logprobs(self) -> list[float]: + """Per-token sampling logprobs aligned to `token_ids` — the node logprobs spread onto + their sampled positions, 0.0 on every non-sampled token.""" + out: list[float] = [] + for node in self.nodes: + li = 0 + for sampled in node.mask: + if sampled: + out.append(node.logprobs[li] if li < len(node.logprobs) else 0.0) + li += 1 + else: + out.append(0.0) + return out @property def completion_len(self) -> int: - """All assistant-generated tokens across this branch's turns.""" - return sum(turn.num_completion_tokens for turn in self.turns) + """All assistant-generated (model-sampled) tokens across this branch.""" + return sum(sum(n.mask) for n in self.nodes) @property def total_tokens(self) -> int: - """This branch's final-turn sequence length (prompt + completion).""" - last = self.turns[-1] if self.turns else None - return last.num_prompt_tokens + last.num_completion_tokens if last else 0 + """This branch's full sequence length (final-turn prompt + every completion).""" + return sum(len(n.token_ids) for n in self.nodes) @property - def messages(self) -> Messages: - last = self.turns[-1] - return [*last.prompt, last.response.message] + def prompt_len(self) -> int: + """Input context size: the final-turn prompt = full sequence minus the last completion.""" + last_completion = next( + (sum(n.mask) for n in reversed(self.nodes) if any(n.mask)), 0 + ) + return self.total_tokens - last_completion class Trace(StrictBaseModel, Generic[TaskT]): @@ -126,8 +133,10 @@ class Trace(StrictBaseModel, Generic[TaskT]): """Unique id for this rollout, auto-generated per trace.""" task: TaskT """The (immutable) task being solved — fully typed, flows into scoring.""" - trajectory: list[Turn] = Field(default_factory=list) - """Every model turn in order — the ground truth.""" + nodes: list[MessageNode] = Field(default_factory=list) + """The message graph — one node per distinct message, linked to its predecessor (see + `graph`). The ground truth; `trajectory` and `branches` are views over it. Stores each + message once, so size is linear (not quadratic) in turns.""" rewards: dict[str, float] = Field(default_factory=dict) """Per-`@reward`-function contributions, with each function's weight applied.""" @@ -141,9 +150,9 @@ class Trace(StrictBaseModel, Generic[TaskT]): rollout was retried). `error` exposes the most recent.""" timing: Timing = Field(default_factory=Timing) - _branch_cache: dict[int, list[list[int]]] = PrivateAttr(default_factory=dict) - """Branch segmentation (index-groups, no turn copies) cached by trajectory length — - the trajectory only grows, so its length is a sufficient invalidation key.""" + _head_index: dict = PrivateAttr(default_factory=dict) + """`(parent, msg_hash) -> node_id` for the graph builder (`graph.add_turn`); rebuilt + lazily from `nodes` after deserialization.""" @computed_field @property @@ -160,10 +169,16 @@ def error(self) -> Error | None: def has_error(self) -> bool: return bool(self.errors) - @property - def last_turn(self) -> Turn | None: - """The final model turn, or None for an empty trajectory.""" - return self.trajectory[-1] if self.trajectory else None + def _last_assistant(self) -> "MessageNode | None": + """The most recent assistant node, or None for a trace with no responses.""" + return next( + ( + n + for n in reversed(self.nodes) + if isinstance(n.message, AssistantMessage) + ), + None, + ) @property def prompt_len(self) -> int: @@ -185,43 +200,26 @@ def total_tokens(self) -> int: @property def has_response(self) -> bool: - """Whether the final assistant turn produced non-empty content.""" - last = self.last_turn - return bool(last and last.response.message.content) - - def _branch_groups(self) -> list[list[int]]: - """Branch index-groups for the current trajectory, recomputed only when a turn - is added (the trajectory is append-only, so its length is a sufficient key). - Caches the groups (ints), never copies of the turns.""" - n = len(self.trajectory) - if n not in self._branch_cache: - self._branch_cache = {n: branching.segment(self.trajectory)} - return self._branch_cache[n] + """Whether the most recent assistant message produced non-empty content.""" + last = self._last_assistant() + return bool(last and last.message.content) - @computed_field @property def branches(self) -> list[Branch]: - """The trajectory segmented into linear branches (see `branching`): one branch - when linear, several under compaction or subagents. The structured view of what - the harness saw, replacing a flat message list. Branches hold turn references, not - copies.""" - return [ - Branch(index=i, turns=[self.trajectory[j] for j in group]) - for i, group in enumerate(self._branch_groups()) - ] + """The conversation segmented into linear branches — a view over the graph: each + leaf's root→leaf path is a branch (one when linear, several under compaction or + subagents). Branching falls out of the walk; see `graph.branches_from_nodes`.""" + return graph.branches_from_nodes(self) - @computed_field @property def num_branches(self) -> int: - """How many branches the trajectory has (1 = linear; >1 = compaction/subagents).""" - return len(self._branch_groups()) + """How many branches (1 = linear; >1 = compaction/subagents).""" + return len(graph.leaves(self)) - @computed_field @property def num_turns(self) -> int: - """Total model turns across the whole trajectory (all branches); per-branch - counts are on each `Branch.num_turns`.""" - return len(self.trajectory) + """Total model turns (assistant nodes) across all branches.""" + return sum(1 for n in self.nodes if isinstance(n.message, AssistantMessage)) @computed_field @property @@ -238,21 +236,23 @@ def is_truncated(self) -> bool: "harness_timeout", ): return True - last = self.last_turn - return bool(last and last.response.finish_reason == "length") + last = self._last_assistant() + return bool(last and last.finish_reason == "length") @property def assistant_messages(self) -> list[AssistantMessage]: """Every model response, in order — one per turn, branch-independent.""" - return [turn.response.message for turn in self.trajectory] + return [ + n.message for n in self.nodes if isinstance(n.message, AssistantMessage) + ] @property def tool_messages(self) -> list[ToolMessage]: - """The tool results in the latest full context — the last turn's prompt. (For a - linear rollout that's every tool result; computed straight off the trajectory, - like `assistant_messages`, with no branch reconstruction.)""" - last = self.trajectory[-1].prompt if self.trajectory else [] - return [m for m in last if isinstance(m, ToolMessage)] + """The tool results in the latest full context — the main (last) branch's + conversation. For a linear rollout that's every tool result.""" + branches = self.branches + messages = branches[-1].messages if branches else [] + return [m for m in messages if isinstance(m, ToolMessage)] def record_metric(self, name: str, value: float) -> None: """Record a single `@metric` result under `name`. Warns if it overrides an diff --git a/verifiers/v1/types.py b/verifiers/v1/types.py index 64321e1574..b37dcf694d 100644 --- a/verifiers/v1/types.py +++ b/verifiers/v1/types.py @@ -105,6 +105,14 @@ class TurnTokens(StrictBaseModel): completion_ids: list[int] = Field(default_factory=list) completion_logprobs: list[float] = Field(default_factory=list) + # Transient build carrier: per-prompt-message token spans into `prompt_ids`, produced by + # the renderer (`RenderedTokens.message_token_spans()`) and consumed by the graph builder + # to attribute tokens per message, then dropped — never persisted (it would reintroduce + # the quadratic data the graph removes). + message_spans: list[tuple[int, int] | None] | None = Field( + default=None, exclude=True + ) + class Response(StrictBaseModel): """One model completion, provider-agnostic."""