from __future__ import annotations import json from pathlib import Path import tempfile from types import SimpleNamespace import unittest from hessian_recipes.agent_rl_training import (_audited_trainer_class, _raw_preserving_template, audit_policy_tokens, audit_task_group, milestone_training_reward, plan) from hessian_recipes.evaluation import checkpoint_sha256 from hessian_recipes.training import sha256 class PolicyTokenAuditTests(unittest.TestCase): def setUp(self): self.prompt = [1, 2] self.completion = [3, 4, 5, 6, 7, 8, 9, 10, 5] self.mask = [1, 1, 1, 0, 0, 0, 1, 1, 1] self.spans = [{"prompt_ids": [1, 2], "completion_ids": [3, 4, 5], "call_index": 0}, {"prompt_ids": [1, 2, 3, 4, 5, 6, 7, 8], "completion_ids": [9, 10, 5], "call_index": 1}] def audit(self, **kwargs): return audit_policy_tokens([self.prompt], [kwargs.get("completion", self.completion)], [kwargs.get("mask", self.mask)], kwargs.get("spans", self.spans)) def test_actual_ids_and_context_prove_assistant_tokens_only(self): evidence = self.audit()[0] self.assertEqual(evidence["assistant_tokens"], 6) self.assertEqual(evidence["external_tokens"], 3) self.assertTrue(evidence["sampled_token_equality_verified"]) self.assertEqual(evidence["sampled_spans"][1]["offset"], 6) def test_json_normalization_is_not_a_sampled_action(self): spans = [dict(span) for span in self.spans] spans[0]["completion_ids"] = [3, 88, 5] with self.assertRaisesRegex(ValueError, "actually sampled"): self.audit(spans=spans) def test_same_tokens_with_wrong_observation_context_fail(self): spans = [dict(span) for span in self.spans] spans[1]["prompt_ids"] = [1, 2, 3, 4, 5, 6, 99, 8] with self.assertRaisesRegex(ValueError, "context"): self.audit(spans=spans) def test_tool_observation_cannot_contribute_policy_loss(self): mask = list(self.mask) mask[4] = 1 with self.assertRaisesRegex(ValueError, "tool mask"): self.audit(mask=mask) def test_sampled_assistant_tokens_cannot_be_mislabeled_as_tool_feedback(self): mask = list(self.mask) mask[1] = 0 with self.assertRaisesRegex(ValueError, "tool mask"): self.audit(mask=mask) def test_truncated_sampled_prefix_is_valid(self): evidence = audit_policy_tokens([[1, 2]], [[3, 4]], [[1, 1]], [self.spans[0]]) self.assertEqual(evidence[0]["assistant_tokens"], 2) def test_missing_or_misaligned_masks_fail_closed(self): with self.assertRaises(ValueError): audit_policy_tokens([[1]], [[2]], None, []) with self.assertRaises(ValueError): self.audit(mask=[1]) def test_upstream_native_outputs_are_audited_before_return(self): outer = self class FakeNativeTrainer: def _generate_single_turn(self, prompts): return [outer.prompt], [[3, 4, 5]], None, {} def _generate(self, prompts): self._generate_single_turn(prompts) return [outer.prompt], [[3, 4, 5]], [[1, 1, 1]], [], 3, None, {} trainer = _audited_trainer_class(FakeNativeTrainer)() trainer._generate([[]]) self.assertTrue(trainer._last_policy_audit[0]["sampled_token_equality_verified"]) def test_unknown_upstream_template_fails_closed(self): with self.assertRaisesRegex(ValueError, "layout changed"): _raw_preserving_template("unknown template") class AgentRecipeTests(unittest.TestCase): def setUp(self): self.temporary = tempfile.TemporaryDirectory() self.addCleanup(self.temporary.cleanup) self.root = Path(self.temporary.name) model = self.root / "model" model.mkdir() (model / "config.json").write_text('{"model_type":"qwen3"}') (model / "model.safetensors").write_bytes(b"weights") self.dataset = self.root / "train.jsonl" self.row = {"id": "task-1", "prompt": "Hole Daten und berechne den Preis.", "task": {"split": "train"}} self.dataset.write_text(json.dumps(self.row) + "\n") self.config = {"schema_version": 1, "model": {"path": "model", "sha256": checkpoint_sha256(model)}, "dataset": {"path": "train.jsonl", "sha256": sha256(self.dataset)}, "output_dir": "rl-output"} self.path = self.root / "recipe.json" def plan(self): self.path.write_text(json.dumps(self.config)) return plan(self.path) def test_plan_is_online_single_group_and_does_not_create_output(self): report = self.plan() self.assertEqual(report["training"]["num_generations"], 4) self.assertEqual(report["fixed"]["num_iterations"], 1) self.assertEqual(report["fixed"]["steps_per_generation"], 1) self.assertEqual(report["chat_template_kwargs"], {"enable_thinking": True}) self.assertFalse((self.root / "rl-output").exists()) def test_hash_mismatch_is_rejected(self): self.config["dataset"]["sha256"] = "0" * 64 with self.assertRaisesRegex(ValueError, "SHA256 mismatch"): self.plan() def test_test_split_can_never_be_used_for_optimizer_updates(self): self.row["task"]["split"] = "test" self.dataset.write_text(json.dumps(self.row) + "\n") self.config["dataset"]["sha256"] = sha256(self.dataset) with self.assertRaisesRegex(ValueError, "train-split"): self.plan() def test_stale_rollout_accumulation_and_unsupported_settings_fail(self): for settings in ({"gradient_accumulation_steps": 2}, {"per_device_train_batch_size": 8}, {"use_vllm": True}, {"max_steps": 1000}, {"temperature": 0}, {"beta": -1}, {"top_p": 1.5}, {"max_tool_calling_iterations": 9}): with self.subTest(settings=settings): self.config["training"] = settings with self.assertRaises(ValueError): self.plan() def test_policy_dropout_or_uncontrolled_lora_is_rejected(self): self.config["lora"] = {"dropout": 0.1} with self.assertRaisesRegex(ValueError, "dropout=0"): self.plan() class TaskGroupAuditTests(unittest.TestCase): def setUp(self): self.task = {"task_id": "train-1", "episode_seed": 9123, "split": "train", "private_facts": {"quantity": 3, "unit_price": 7}} def test_same_task_group_has_explicit_metadata_without_hidden_facts(self): tasks = [dict(self.task) for _ in range(4)] environments = [SimpleNamespace(_task=dict(task)) for task in tasks] metadata = audit_task_group(tasks, environments, 4) self.assertEqual(metadata, [{"task_id": "train-1", "episode_seed": 9123, "split": "train"}] * 4) self.assertNotIn("private_facts", metadata[0]) def test_mixed_task_groups_fail_before_rewards(self): other = self.task | {"task_id": "train-2"} tasks = [self.task, other] with self.assertRaisesRegex(ValueError, "same task identity"): audit_task_group(tasks, [SimpleNamespace(_task=task) for task in tasks], 2) def test_incorrect_environment_binding_is_rejected(self): with self.assertRaisesRegex(ValueError, "actual reset task"): audit_task_group([self.task], [SimpleNamespace(_task=self.task | {"episode_seed": 4})], 1) def test_missing_or_incomplete_group_is_rejected(self): with self.assertRaises(ValueError): audit_task_group(None, [], 4) with self.assertRaises(ValueError): audit_task_group([self.task], [SimpleNamespace(_task=self.task)], 4) def test_non_train_group_is_rejected(self): task = self.task | {"split": "test"} with self.assertRaisesRegex(ValueError, "train split"): audit_task_group([task], [SimpleNamespace(_task=task)], 1) class MilestoneRewardTests(unittest.TestCase): @staticmethod def score(*, lookup=False, calculation=False, success=False, outcomes=()): return {"lookup_success": lookup, "calculation_success": calculation, "success": success, "reward": float(success), "trace": [{"ok": ok} for ok in outcomes]} def test_no_tools_or_fabricated_final_answer_earn_no_progress(self): score = self.score() | {"final_answer": "42", "explanation": "Perfect German explanation."} self.assertEqual(milestone_training_reward(score)["training_reward"], 0.0) def test_authoritative_milestones_give_bounded_partial_credit(self): self.assertAlmostEqual(milestone_training_reward(self.score(lookup=True))["training_reward"], 0.15) score = self.score(lookup=True, calculation=True) reward = milestone_training_reward(score) self.assertAlmostEqual(reward["training_reward"], 0.40) self.assertFalse(reward["strict_success"]) self.assertEqual(score["reward"], 0.0) def test_actual_error_recovery_creates_learning_signal_without_strict_success(self): # The observed rollout: wrong calculator, valid lookup, valid calculator, # correct final answer. The initial process error still fails strict eval. score = self.score(lookup=True, calculation=True, outcomes=(False, True, True)) reward = milestone_training_reward(score) self.assertAlmostEqual(reward["training_reward"], 0.35) self.assertEqual(reward["trusted_tool_errors"], 1) self.assertEqual(reward["strict_reward"], 0.0) self.assertFalse(score["success"]) def test_full_strict_success_outranks_all_partial_workflows(self): reward = milestone_training_reward(self.score(lookup=True, calculation=True, success=True, outcomes=(True, True))) self.assertEqual(reward["training_reward"], 1.0) self.assertEqual(reward["strict_reward"], 1.0) def test_repeated_successful_tool_calls_do_not_repeat_milestone_bonus(self): one = milestone_training_reward(self.score(lookup=True, outcomes=(True,))) many = milestone_training_reward(self.score(lookup=True, outcomes=(True,) * 5)) self.assertEqual(one["training_reward"], many["training_reward"]) def test_error_penalty_uses_real_trace_and_clips_at_zero(self): reward = milestone_training_reward(self.score(lookup=True, outcomes=(False,) * 10)) self.assertEqual(reward["training_reward"], 0.0) self.assertEqual(reward["tool_error_penalty"], 0.5) def test_inconsistent_or_untrusted_scores_fail_closed(self): invalid = [self.score() | {"lookup_success": 1}, self.score() | {"trace": [{}]}, self.score(calculation=True), self.score(lookup=True, calculation=True, success=True, outcomes=(False,)), self.score() | {"reward": 0.15}] for score in invalid: with self.subTest(score=score), self.assertRaises(ValueError): milestone_training_reward(score) if __name__ == "__main__": unittest.main()