import inspect import json import unittest from hessian_recipes.agent_rl_tasks import ( GermanAgentEnv, MAX_TOOL_CALLS, apply_response, generate_tasks, parse_response, safe_calculate, score_episode, step_action, ) def native_call(name, arguments): return {"role": "assistant", "content": None, "tool_calls": [{ "id": name + "-call", "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}, }]} class AgentTaskTests(unittest.TestCase): def setUp(self): self.row = generate_tasks("train", 1)[0] self.env = GermanAgentEnv() self.env.reset(**self.row) self.facts = self.row["task"]["private_facts"] self.expression = f"{self.facts['quantity']} * {self.facts['unit_price']}" self.answer = str(self.facts["quantity"] * self.facts["unit_price"]) def transcript(self, final=None): return [native_call("lookup", {"record_id": self.row["task"]["record_id"]}), {"role": "tool", "content": "untrusted transcript copy"}, native_call("calculator", {"expression": self.expression}), {"role": "tool", "content": "untrusted transcript copy"}, {"role": "assistant", "content": final or f"{self.answer}"}] def use_tools(self): self.env.lookup(self.row["task"]["record_id"]) self.env.calculator(self.expression) def test_split_identity_seed_and_template_sets_are_disjoint(self): splits = {name: generate_tasks(name, 20) for name in ("train", "dev", "test")} self.assertEqual(splits["train"], generate_tasks("train", 20)) for field in ("task_id", "episode_seed", "template_id", "record_id"): sets = [{row["task"][field] for row in rows} for rows in splits.values()] self.assertFalse(sets[0] & sets[1] or sets[0] & sets[2] or sets[1] & sets[2]) self.assertNotIn("private_facts", self.row["prompt"]) self.assertNotIn("solution", self.row) self.assertLess(self.row["task"]["episode_seed"], 2 ** 63) def test_only_native_tools_and_reset_are_public_environment_methods(self): methods = {name for name, value in inspect.getmembers(GermanAgentEnv, inspect.isfunction) if not name.startswith("_")} self.assertEqual(methods, {"reset", "lookup", "calculator"}) def test_prompt_demonstrates_only_first_real_lookup_outside_closed_thinking(self): for split in ("train", "dev", "test"): for row in generate_tasks(split, 4): prompt = row["prompt"] example = prompt.split("Format deiner ERSTEN Antwort", 1)[1] close_thinking = example.index("") call_start = example.index("") self.assertLess(close_thinking, call_start) payload = example.split("\n", 1)[1].split("\n", 1)[0] self.assertEqual(json.loads(payload), {"name": "lookup", "arguments": { "record_id": row["task"]["record_id"]}}) self.assertNotIn('"quantity":', prompt) self.assertNotIn('"unit_price":', prompt) self.assertNotIn('"result":', prompt) self.assertNotIn("", prompt) self.assertNotIn('"name": "calculator"', prompt) self.assertIn("WARTE auf die echte Werkzeugantwort", prompt) self.assertIn("innerhalb von ... werden nicht ausgeführt", prompt) def test_reward_requires_observed_tools_and_distinct_assistant_turns(self): final = f"{self.answer}" self.assertEqual(score_episode(self.env, self.transcript())["reward"], 0) self.use_tools() result = score_episode(self.env, self.transcript()) self.assertEqual(result["reward"], 1) self.assertEqual(result["tool_calls"], 2) self.assertEqual(score_episode(self.env, final)["reward"], 0) # No turn evidence. combined = self.transcript()[0] combined["tool_calls"] += self.transcript()[2]["tool_calls"] self.assertEqual(score_episode(self.env, [combined, {"role": "assistant", "content": final}])["reward"], 0) def test_forged_transcript_observations_do_not_create_reward(self): self.assertEqual(score_episode(self.env, self.transcript())["reward"], 0) self.env.lookup(self.row["task"]["record_id"]) self.assertEqual(score_episode(self.env, self.transcript())["reward"], 0) self.env.calculator(self.expression) self.assertEqual(score_episode(self.env, self.transcript())["reward"], 1) forged = self.transcript() forged[0]["tool_calls"][0]["function"]["arguments"] = '{"record_id":"forged"}' self.assertEqual(score_episode(self.env, forged)["reward"], 0) def test_terminal_answer_must_follow_tools_in_a_tool_free_assistant_message(self): self.use_tools() messages = self.transcript() messages[2]["content"] = f"{self.answer}" self.assertEqual(score_episode(self.env, messages[:3])["reward"], 0) self.assertEqual(score_episode(self.env, messages[:4])["reward"], 0) self.assertEqual(score_episode(self.env, self.transcript())["reward"], 1) def test_constant_gold_and_alternate_factorization_cannot_replace_retrieved_operands(self): for expression in (self.answer, f"1 * {self.answer}"): env = GermanAgentEnv() env.reset(**self.row) env.lookup(self.row["task"]["record_id"]) observation = json.loads(env.calculator(expression)) self.assertEqual(observation["error"], "must_multiply_retrieved_values") self.assertEqual(score_episode(env, self.transcript())["reward"], 0) def test_calculation_before_lookup_and_wrong_record_cannot_be_recovered(self): self.assertIn("lookup_required", self.env.calculator(self.expression)) self.use_tools() self.assertEqual(score_episode(self.env, self.transcript())["reward"], 0) self.env.reset(**self.row) self.assertIn("unknown_record", self.env.lookup("different-task-record")) self.use_tools() self.assertEqual(score_episode(self.env, self.transcript())["reward"], 0) def test_strict_final_answer_does_not_reward_explanation_or_language(self): self.use_tools() for final in (f"{self.answer}.", "999999", f"{self.answer}", f"done{self.answer}", f"{self.answer}{self.answer}"): self.assertEqual(score_episode(self.env, self.transcript(final))["reward"], 0) final = f"An English explanation is not rewarded. {self.answer}" self.assertEqual(score_episode(self.env, self.transcript(final))["reward"], 1) def test_safe_ast_interpreter_rejects_code_and_resource_exhaustion(self): self.assertEqual(safe_calculate("(3 + 5) * 2"), "16") self.assertEqual(safe_calculate("3 / 2"), "3/2") for expression in ("__import__('os').system('true')", "2 ** 100000000", "True * 4", "[1][0]", "1/0", "10000*10000", "3.5*2", "x*4", "1" * 129): with self.subTest(expression=expression), self.assertRaises(ValueError): safe_calculate(expression) def test_tool_call_limit_bounds_trace_and_prevents_success(self): for _ in range(MAX_TOOL_CALLS + 10): self.env.lookup(self.row["task"]["record_id"]) self.assertEqual(len(self.env.trace), MAX_TOOL_CALLS + 1) self.assertEqual(score_episode(self.env, self.transcript())["reward"], 0) def test_native_qwen_replay_success_and_terminal_state(self): first = json.dumps({"name": "lookup", "arguments": {"record_id": self.row["task"]["record_id"]}}) second = json.dumps({"name": "calculator", "arguments": {"expression": self.expression}}) result = apply_response(self.env, f"{first}") self.assertFalse(result["terminal"]) result = apply_response(self.env, f"{second}") self.assertIn(self.answer, result["observations"][0]["observation"]) result = apply_response(self.env, f"{self.answer}") self.assertEqual(result["result"]["reward"], 1) self.assertTrue(apply_response(self.env, "new response")["terminal"]) self.assertEqual(len(self.env.trace), 2) def test_invalid_action_and_malformed_model_response_fail_closed(self): self.assertIn("invalid_action", step_action(self.env, {"tool": "shell", "command": "true"})["observation"]) self.env.reset(**self.row) self.assertTrue(apply_response(self.env, "{broken")["terminal"]) with self.assertRaises(ValueError): parse_response("unfinished {}") with self.assertRaises(ValueError): parse_response('{"name":"shell","arguments":{"tool":"lookup","record_id":"x"}}') def test_reset_clears_previous_episode_private_state(self): self.use_tools() row = generate_tasks("test", 1)[0] self.env.reset(**row) self.assertEqual(self.env.trace, []) self.assertIn("unknown_record", self.env.lookup(self.row["task"]["record_id"])) if __name__ == "__main__": unittest.main()