"""Versioned, framework-neutral agent rollout exchange format. This validates captured data, not eligibility for a particular RL backend. Exact prompt and response token IDs must come from inference, never text retokenization. Tool observations belong only in a call's context/prompt; the response is always the sampled assistant output. See docs/agent-harness.md for the full contract. """ from __future__ import annotations import copy import json import math import os import re import tempfile from pathlib import Path from typing import Any, Iterable SCHEMA_VERSION = "hessian.trajectory.v1" def _object(value: Any, required: set[str], optional: set[str], location: str) -> dict: if not isinstance(value, dict): raise ValueError(f"{location}: expected an object") missing = required - value.keys() extra = value.keys() - required - optional if missing or extra: raise ValueError( f"{location}: missing keys {sorted(missing)}, unknown keys {sorted(map(str, extra))}" ) return value def _text(value: Any, location: str) -> None: if not isinstance(value, str) or not value.strip(): raise ValueError(f"{location}: expected a nonempty string") def _revision(value: Any, location: str) -> None: _text(value, location) if value.lower() in {"main", "master", "latest", "head", "unknown"}: raise ValueError(f"{location}: use an identified checkpoint or revision, not {value!r}") def _tokens(value: Any, location: str) -> None: if not isinstance(value, list) or not value: raise ValueError(f"{location}: expected a nonempty token ID list") if any(type(token) is not int or token < 0 for token in value): raise ValueError(f"{location}: token IDs must be nonnegative integers") def _number(value: Any, location: str) -> None: try: valid = type(value) in (int, float) and math.isfinite(value) except OverflowError: valid = False if not valid: raise ValueError(f"{location}: expected a finite number") def validate_trajectory(record: Any) -> dict: """Return an independent copy of a valid v1 episode, or raise ValueError.""" episode = _object(record, { "schema_version", "episode_id", "group_id", "task_id", "seed", "source_kind", "policy", "harness", "reward", "calls", }, set(), "episode") if episode["schema_version"] != SCHEMA_VERSION: raise ValueError(f"schema_version: expected {SCHEMA_VERSION!r}") for field in ("episode_id", "group_id", "task_id"): _text(episode[field], field) if type(episode["seed"]) is not int or episode["seed"] < 0: raise ValueError("seed: expected a nonnegative integer") if episode["source_kind"] not in ("captured", "fixture"): raise ValueError("source_kind: expected captured or fixture") policy = _object(episode["policy"], { "model_id", "revision", "tokenizer_id", "tokenizer_revision", "chat_template_sha256", }, set(), "policy") for field in ("model_id", "tokenizer_id"): _text(policy[field], f"policy.{field}") for field in ("revision", "tokenizer_revision"): _revision(policy[field], f"policy.{field}") if not isinstance(policy["chat_template_sha256"], str) or not re.fullmatch( r"[0-9a-f]{64}", policy["chat_template_sha256"] ): raise ValueError("policy.chat_template_sha256: expected lowercase SHA-256 hex") harness = _object(episode["harness"], {"name", "version", "revision"}, set(), "harness") for field in ("name", "version"): _text(harness[field], f"harness.{field}") _revision(harness["revision"], "harness.revision") reward = _object(episode["reward"], {"value", "scorer", "revision"}, set(), "reward") _number(reward["value"], "reward.value") _text(reward["scorer"], "reward.scorer") _revision(reward["revision"], "reward.revision") if not isinstance(episode["calls"], list) or not episode["calls"]: raise ValueError("calls: expected at least one inference call") call_ids: set[str] = set() trainable = 0 for index, value in enumerate(episode["calls"]): loc = f"calls[{index}]" call = _object(value, { "call_id", "prompt_token_ids", "response_token_ids", "response_loss_mask", "response_role", "finish_reason", "context", }, {"behavior_logprobs"}, loc) _text(call["call_id"], f"{loc}.call_id") if call["call_id"] in call_ids: raise ValueError(f"{loc}.call_id: duplicate call ID") call_ids.add(call["call_id"]) _tokens(call["prompt_token_ids"], f"{loc}.prompt_token_ids") _tokens(call["response_token_ids"], f"{loc}.response_token_ids") if call["response_role"] != "assistant": raise ValueError(f"{loc}.response_role: only assistant responses can be trainable") if call["finish_reason"] not in ("stop", "length", "tool_calls"): raise ValueError(f"{loc}.finish_reason: expected stop, length, or tool_calls") mask = call["response_loss_mask"] if not isinstance(mask, list) or len(mask) != len(call["response_token_ids"]): raise ValueError(f"{loc}.response_loss_mask: must align with response token IDs") if any(type(bit) is not int or bit not in (0, 1) for bit in mask): raise ValueError(f"{loc}.response_loss_mask: expected integer zeros and ones") trainable += sum(mask) if "behavior_logprobs" in call: probs = call["behavior_logprobs"] if not isinstance(probs, list) or len(probs) != len(mask): raise ValueError(f"{loc}.behavior_logprobs: must align with response token IDs") for prob in probs: _number(prob, f"{loc}.behavior_logprobs") if prob > 0: raise ValueError(f"{loc}.behavior_logprobs: log probabilities cannot be positive") if not isinstance(call["context"], list) or not call["context"]: raise ValueError(f"{loc}.context: expected nonempty audit context") for j, message in enumerate(call["context"]): msg = _object(message, {"role", "content"}, set(), f"{loc}.context[{j}]") if msg["role"] not in ("system", "user", "assistant", "tool"): raise ValueError(f"{loc}.context[{j}].role: unsupported role") if not isinstance(msg["content"], str): raise ValueError(f"{loc}.context[{j}].content: expected a string") if not trainable: raise ValueError("calls: episode must contain at least one trainable assistant token") return copy.deepcopy(episode) def _validate_batch(records: Iterable[Any]) -> list[dict]: result = [] episodes: set[str] = set() groups: dict[str, tuple] = {} for record in records: row = validate_trajectory(record) if row["episode_id"] in episodes: raise ValueError(f"duplicate episode_id: {row['episode_id']}") episodes.add(row["episode_id"]) # A GRPO comparison group cannot silently mix tasks, policies, or scorers. signature = (row["task_id"], json.dumps(row["policy"], sort_keys=True), json.dumps(row["harness"], sort_keys=True), row["reward"]["scorer"], row["reward"]["revision"], row["source_kind"]) previous = groups.setdefault(row["group_id"], signature) if previous != signature: raise ValueError(f"group_id {row['group_id']!r}: task, policy, harness, or scorer mismatch") result.append(row) if not result: raise ValueError("trajectory file must contain at least one episode") return result def _unique_keys(pairs: list[tuple]) -> dict: obj = {} for key, value in pairs: if key in obj: raise ValueError(f"duplicate JSON key: {key}") obj[key] = value return obj def read_trajectories(path: str | Path) -> list[dict]: """Load strict JSONL and validate episodes and group consistency.""" records = [] with Path(path).open(encoding="utf-8") as handle: for line_number, line in enumerate(handle, 1): try: row = json.loads(line, object_pairs_hook=_unique_keys) records.append(validate_trajectory(row)) except (ValueError, TypeError) as exc: raise ValueError(f"{path}:{line_number}: {exc}") from exc return _validate_batch(records) def write_trajectories(path: str | Path, records: Iterable[dict]) -> int: """Validate the whole batch, then atomically write JSONL; return row count.""" rows = _validate_batch(records) destination = Path(path) destination.parent.mkdir(parents=True, exist_ok=True) temporary = None try: with tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", dir=destination.parent, prefix=f".{destination.name}.", delete=False) as handle: temporary = handle.name for row in rows: handle.write(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n") os.replace(temporary, destination) finally: if temporary is not None and os.path.exists(temporary): os.unlink(temporary) return len(rows)