import copy import json import statistics import subprocess import pytest from hessian_recipes import agent_rl_evaluation as agent_eval from hessian_recipes import agent_rl_report as reporting from hessian_recipes.agent_rl_pilot import source_hashes from hessian_recipes.agent_rl_tasks import generate_tasks from hessian_recipes.evaluation import ( checkpoint_sha256, compare as math_compare, run as math_run, sha256_file, ) from test_agent_rl_evaluation import FixturePolicy, setup as evaluation_setup def write(path, value): path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(value)) def jsonl(path, rows): path.write_text("".join(json.dumps(row) + "\n" for row in rows)) def fixture_run(tmp_path): root = tmp_path / "pilot" root.mkdir() spec_path = evaluation_setup(root, count=2) spec = json.loads(spec_path.read_text()) (root / "tokenizer").rename(root / "frozen-tokenizer") (root / "tasks.jsonl").rename(root / "test.jsonl") spec["tokenizer"]["path"] = "frozen-tokenizer" spec["tasks"]["path"] = "test.jsonl" train = generate_tasks("train", 2) for row in train: row["task"]["private_facts"] = {"quantity": 2, "unit_price": 3} jsonl(root / "train.jsonl", train) jsonl(root / "dev.jsonl", generate_tasks("dev", 1)) model_sha = checkpoint_sha256(root / "checkpoint") merged = root / "merged-checkpoint" merged.mkdir() (merged / "weights.bin").write_bytes(b"synthetic merged artifact") merged_sha = checkpoint_sha256(merged) for label in ("baseline", "candidate"): current = copy.deepcopy(spec) if label == "candidate": current["model"] = {"path": "merged-checkpoint", "sha256": merged_sha} write(root / (label + "-spec.json"), current) agent_eval.run( root / (label + "-spec.json"), root / label, generator=FixturePolicy("skip" if label == "baseline" else "success"), ) training_config = {"mock": "synthetic fixture"} write(root / "training.json", training_config) adapter = root / "training/checkpoint" adapter.mkdir(parents=True) (adapter / "adapter_model.safetensors").write_bytes( b"synthetic test adapter; not a trained policy" ) records, groups = [], [] for step, row in enumerate(train): rewards = [] for index in range(4): episode = agent_eval._episode( row, index, spec, FixturePolicy("success" if index % 2 == 0 else "skip") ) calls = episode["calls"] initial = calls[0]["prompt_token_ids"] complete = calls[-1]["prompt_token_ids"] + calls[-1]["response_token_ids"] completion, spans = complete[len(initial) :], [] mask = [0] * len(completion) for call_index, call in enumerate(calls): offset = len(call["prompt_token_ids"]) - len(initial) count = len(call["response_token_ids"]) mask[offset : offset + count] = [1] * count spans.append( { "offset": offset, "retained_tokens": count, "prompt_ids": call["prompt_token_ids"], "sampled_completion_ids": call["response_token_ids"], "call_index": call_index, } ) score = episode["outcome"] reward = float(score["success"]) rewards.append(reward) version = {"global_step": step, "base_model_sha256": model_sha} messages = [] for call in calls: messages.extend([call["assistant"], *call["observations"]]) records.append( { "task_id": row["id"], "episode_seed": row["task"]["episode_seed"], "split": "train", "optimizer_step": step, "rollout_index": index, "completion_messages": messages, "score": score, "training_reward": reward, "reward_components": {"training_reward": reward}, "policy_version": version, "policy_audit": { "prompt_ids": initial, "completion_ids": completion, "tool_mask": mask, "assistant_tokens": sum(mask), "external_tokens": len(mask) - sum(mask), "sampled_spans": spans, }, } ) groups.append( { "optimizer_step": step, "task_ids": [row["id"]] * 4, "rewards": rewards, "strict_rewards": rewards, "reward_std": statistics.pstdev(rewards), } ) jsonl(root / "training/rollout-evidence.jsonl", records) protocol = source_hashes() training = { "status": "completed", "model": {"path": str(root / "checkpoint"), "sha256": model_sha}, "dataset": {"sha256": sha256_file(root / "train.jsonl")}, "config_sha256": sha256_file(root / "training.json"), "policy_format": { "chat_template_sha256": FixturePolicy.template_metadata["chat_template_sha256"], "helper_sha256": protocol["agent_rl_training"], }, "helper_sha256": protocol["agent_rl_training"], "environment_sha256": protocol["agent_rl_tasks"], "checkpoint": str(adapter), "checkpoint_sha256": checkpoint_sha256(adapter), "evidence_sha256": sha256_file(root / "training/rollout-evidence.jsonl"), "training": { "max_steps": 2, "num_generations": 4, "max_completion_length": 768, "max_tool_calling_iterations": 3, "temperature": 0.6, "top_p": 0.95, "top_k": 20, "min_p": 0, }, "chat_template_kwargs": {"enable_thinking": True}, "reward_groups": groups, "train_log_history": [{"grad_norm": 1.0}, {"grad_norm": 0.5}], "task_coverage": { "consumed_unique_task_ids": sorted(row["id"] for row in train), "total_rollouts": 8, "training_pool_rows": 2, }, "update_evidence": { "changed_trainable_tensors": 1, "mixed_reward_groups": 2, "nonzero_gradient_steps": 2, "max_abs_parameter_delta": 0.001, }, "lora": {"r": 1}, "reward_protocol": {"scope": "training_only"}, "dependencies": {"mode": "fixture"}, } write(root / "training/training-manifest.json", training) write( root / "merged-checkpoint.merge-manifest.json", { "status": "completed", "base_model": {"sha256": model_sha}, "adapter": {"path": str(adapter), "sha256": checkpoint_sha256(adapter)}, "output": {"sha256": merged_sha}, }, ) prior = tmp_path / "prior-sft" prior.mkdir() jsonl( prior / "evaluation.jsonl", [ { "id": "math1", "category": "math", "prompt": "PRIVATE QUESTION", "scorer": "final_numeric", "answer": "6", } ], ) for destination, revision, answer in ((prior, model_sha, "6"), (root, merged_sha, "wrong")): predictions = destination / "math-predictions.jsonl" jsonl(predictions, [{"id": "math1", "trial": 0, "output": answer}]) math_spec = { "schema_version": 1, "benchmark": { "id": "fixture-math", "version": "1", "status": "development", "count": 1, "path": str(prior / "evaluation.jsonl"), "sha256": sha256_file(prior / "evaluation.jsonl"), }, "model": { "id": "fixture", "revision": revision, "tokenizer_id": "fixture", "tokenizer_revision": spec["tokenizer"]["sha256"], "chat_template": "{{ prompt }}", }, "sampling": {"temperature": 0, "max_new_tokens": 8, "seed": 42}, "trials": 1, "adapter": { "type": "fixture", "path": str(predictions), "sha256": sha256_file(predictions), }, "scorer_version": "1", "harness_version": "1", } name = "candidate-spec.json" if destination == prior else "retention-spec.json" write(destination / name, math_spec) math_run( destination / name, destination / ("candidate" if destination == prior else "retention") ) plan = { "model": str(root / "checkpoint"), "model_sha256": model_sha, "tokenizer_sha256": spec["tokenizer"]["sha256"], "source_sha256": protocol, "training_config_sha256": sha256_file(root / "training.json"), "split_sha256": {k: sha256_file(root / (k + ".jsonl")) for k in ("train", "dev", "test")}, "split_sizes": {"train": 2, "dev": 1, "test": 2}, "seed": 42, "steps": 2, "sampling": spec["sampling"], "chat_template_kwargs": spec["chat_template_kwargs"], "max_completion_length": 768, "max_tool_calling_iterations": 3, "retention_run": str(prior), "retention_baseline_report_sha256": sha256_file(prior / "candidate/report.json"), "retention_baseline_spec_sha256": sha256_file(prior / "candidate-spec.json"), "retention_tasks_sha256": sha256_file(prior / "evaluation.jsonl"), } write(root / "plan.json", plan) commit = subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip() write(root / "launch-provenance.json", {"source_commit": commit, "source_sha256": protocol}) write( root / "run-manifest.json", { "status": "completed", "completed_stages": ["baseline", "grpo", "merge", "candidate", "retention"], "plan_sha256": sha256_file(root / "plan.json"), "comparison": agent_eval.compare( root / "baseline/report.json", root / "candidate/report.json" ), "retention_comparison": math_compare( prior / "candidate/report.json", root / "retention/report.json" ), }, ) return root def test_verified_summary_has_no_raw_outputs_and_keeps_retention_regression(tmp_path): root = fixture_run(tmp_path) summary = reporting.summarize(root, tmp_path / "summary") assert summary["status"] == "ready" and not summary["release_eligible"] assert summary["training"]["consumed_unique_tasks"] == 2 assert summary["training"]["total_rollouts"] == 8 assert summary["training"]["strict_train_successes"] == 4 assert summary["training"]["token_audit"]["external_tokens"] > 0 assert summary["agent_eval"]["candidate"]["successful"] == 2 assert summary["agent_eval"]["baseline"]["final_answer_correct"] == 2 assert not summary["MGSM_retention"]["gate"]["passed"] assert "PRIVATE QUESTION" not in (tmp_path / "summary/summary.json").read_text() with pytest.raises(ValueError, match="fresh"): reporting.summarize(root, tmp_path / "summary") @pytest.mark.parametrize( "failure", ["ancestry", "zero_update", "outcome", "retention", "mask", "template"] ) def test_inconsistent_or_nonlearning_evidence_rejected(tmp_path, failure): root = fixture_run(tmp_path) if failure == "ancestry": path = root / "merged-checkpoint.merge-manifest.json" value = json.loads(path.read_text()) value["base_model"]["sha256"] = "f" * 64 elif failure == "zero_update": path = root / "training/training-manifest.json" value = json.loads(path.read_text()) value["update_evidence"]["changed_trainable_tensors"] = 0 elif failure == "template": path = root / "training/training-manifest.json" value = json.loads(path.read_text()) value["policy_format"]["chat_template_sha256"] = "f" * 64 elif failure == "outcome": path = root / "candidate/report.json" value = json.loads(path.read_text()) value["metrics"]["successful"] = 0 elif failure == "retention": path = root / "retention/report.json" value = json.loads(path.read_text()) value["results"][0]["output"] = "6" else: path = root / "training/rollout-evidence.jsonl" rows = [json.loads(line) for line in path.read_text().splitlines()] rows[0]["policy_audit"]["tool_mask"][0] = 0 jsonl(path, rows) path = root / "training/training-manifest.json" value = json.loads(path.read_text()) value["evidence_sha256"] = sha256_file(root / "training/rollout-evidence.jsonl") write(path, value) with pytest.raises((ValueError, RuntimeError)): reporting.summarize(root, tmp_path / "summary") assert not (tmp_path / "summary").exists()