"""Tests for agent_data_toolkit.py. Run: python3 test_agent_data_toolkit.py""" import contextlib import io import json import math import random import unittest import agent_data_toolkit as ad from agent_data_toolkit import Event, Trace def mk(events=None, reward=1.0, tid="t1", user="u00", prompt="Add the numbers: 1, 2", answer="Answer: 3", train=True, started=0.0, ttype="sum", harness="react", scrubbed=False, verifier=True): """A small valid trace. Pass events to replace the default tool-free flow.""" evs = events if events is not None else [ Event("user", prompt, started + 0.5), Event("assistant", answer, started + 1.0)] return Trace(tid, "s-" + tid, user, ttype, harness, "m", started, evs, reward, {"type": "exact", "expected": "3"} if verifier else None, {"log": True, "train": train, "share": False}, scrubbed) def fb(kind, text="", ts=2.0): return Event("feedback", text, ts, {"type": kind}) class TestShared(unittest.TestCase): def test_stable_hash_is_stable(self): self.assertEqual(ad.stable_hash("a", 1), ad.stable_hash("a", 1)) self.assertEqual(ad.stable_hash("a"), 0xca978112ca1bbdca) def test_stable_hash_separates_parts(self): self.assertNotEqual(ad.stable_hash("a", "bc"), ad.stable_hash("ab", "c")) def test_stable_hash_bits(self): self.assertLess(ad.stable_hash("x", bits=32), 2 ** 32) self.assertLess(ad.stable_hash("x", bits=64), 2 ** 64) self.assertNotEqual(ad.stable_hash("x", bits=32), ad.stable_hash("x", bits=64)) def test_rng_for_repeats(self): a, b = ad.rng_for("k", 1), ad.rng_for("k", 1) self.assertEqual([a.random() for _ in range(3)], [b.random() for _ in range(3)]) self.assertNotEqual(ad.rng_for("k", 1).random(), ad.rng_for("k", 2).random()) def test_words(self): self.assertEqual(ad.words("Don't STOP, 42 times!"), ["don't", "stop", "42", "times"]) self.assertEqual(ad.words(""), []) def test_mean(self): self.assertEqual(ad.mean([1, 2, 3]), 2.0) self.assertEqual(ad.mean([]), 0.0) def test_entropy(self): self.assertAlmostEqual(ad.entropy([1, 1]), 1.0) self.assertAlmostEqual(ad.entropy([5, 5, 5, 5]), 2.0) self.assertEqual(ad.entropy([7]), 0.0) self.assertEqual(ad.entropy([]), 0.0) self.assertEqual(ad.entropy([0, 0]), 0.0) self.assertAlmostEqual(ad.entropy({"a": 1, "b": 1, "c": 0}), 1.0) def test_count_by(self): self.assertEqual(ad.count_by([{"k": "a"}, {"k": "b"}, {"k": "a"}], "k"), {"a": 2, "b": 1}) class TestTraceModel(unittest.TestCase): def test_event_round_trip(self): e = Event("user", "hi", 1.5, {"a": 1}) d = e.to_dict() self.assertEqual(Event.from_dict(d), e) d["meta"]["a"] = 2 self.assertEqual(e.meta["a"], 1) def test_event_from_dict_without_meta(self): self.assertEqual(Event.from_dict({"kind": "user", "text": "x", "ts": 0}).meta, {}) def test_trace_round_trip(self): t = mk([Event("user", "p", 0.5), Event("assistant", "a", 1.0)]) self.assertEqual(Trace.from_dict(t.to_dict()), t) def test_scrubbed_flag_round_trips(self): t = mk(scrubbed=True) self.assertTrue(Trace.from_dict(t.to_dict()).scrubbed) self.assertEqual(ad.loads_jsonl(ad.dumps_jsonl([t]))[0], t) def test_trace_from_dict_defaults(self): d = mk().to_dict() del d["consent"], d["scrubbed"], d["reward"], d["verifier"] t = Trace.from_dict(d) self.assertEqual(t.consent, ad.DEFAULT_CONSENT) self.assertFalse(t.scrubbed) self.assertIsNone(t.reward) def test_default_consent_is_private(self): self.assertEqual(ad.DEFAULT_CONSENT, {"log": True, "train": False, "share": False}) a, b = mk(), mk() a.consent["train"] = False a.consent["x"] = 1 self.assertNotIn("x", b.consent) def test_prompt_and_final_text(self): t = mk([Event("system", "s", 0), Event("user", "first", 1), Event("assistant", "a1", 2), Event("user", "second", 3), Event("assistant", "a2", 4)]) self.assertEqual(t.prompt(), "first") self.assertEqual(t.final_text(), "a2") e = mk([Event("system", "s", 0)]) self.assertEqual((e.prompt(), e.final_text()), ("", "")) def test_feedback_list(self): t = mk([Event("user", "p", 1), fb("copy", ts=2), Event("assistant", "a", 3), fb("edit", "x", 4)]) self.assertEqual([e.meta["type"] for e in t.feedback()], ["copy", "edit"]) class TestLogger(unittest.TestCase): def setUp(self): self.lg = ad.TraceLogger() self.tid = self.lg.start("s1", "u1", "sum", "react", "m", 10.0) def test_ids_count_up(self): self.assertEqual(self.tid, "t00001") self.assertEqual(self.lg.start("s2", "u", "sum", "react", "m", 0), "t00002") self.assertEqual(ad.TraceLogger("x-").start("s", "u", "sum", "r", "m", 0), "x-00001") def test_consent_merge(self): t = self.lg.all()[0] self.assertEqual(t.consent, ad.DEFAULT_CONSENT) tid = self.lg.start("s", "u", "sum", "r", "m", 0, {"train": True}) self.assertEqual(self.lg._traces[tid].consent, {"log": True, "train": True, "share": False}) def test_log_appends(self): self.lg.log(self.tid, "user", "hello", 11.0, a=1) e = self.lg.all()[0].events[0] self.assertEqual((e.kind, e.text, e.ts, e.meta), ("user", "hello", 11.0, {"a": 1})) def test_log_converts_text_to_str(self): self.lg.log(self.tid, "user", 42, 11.0) self.assertEqual(self.lg.all()[0].events[0].text, "42") def test_log_unknown_trace(self): with self.assertRaises(ad.TraceError): self.lg.log("nope", "user", "x", 1) def test_log_unknown_kind(self): with self.assertRaises(ad.TraceError): self.lg.log(self.tid, "banana", "x", 11) def test_log_time_must_not_go_back(self): self.lg.log(self.tid, "user", "x", 12.0) with self.assertRaises(ad.TraceError): self.lg.log(self.tid, "assistant", "y", 11.9) self.lg.log(self.tid, "assistant", "y", 12.0) def test_first_event_not_before_start(self): with self.assertRaises(ad.TraceError): self.lg.log(self.tid, "user", "x", 9.0) def test_end_sets_fields_and_closes(self): t = self.lg.end(self.tid, 1.0, {"type": "exact", "expected": "1"}) self.assertEqual((t.reward, t.verifier["expected"]), (1.0, "1")) with self.assertRaises(ad.TraceError): self.lg.log(self.tid, "user", "x", 20) with self.assertRaises(ad.TraceError): self.lg.end(self.tid, 1.0) def test_end_checks(self): with self.assertRaises(ad.TraceError): self.lg.end("nope") with self.assertRaises(ad.TraceError): self.lg.end(self.tid, 1.5) with self.assertRaises(ad.TraceError): self.lg.end(self.tid, -0.1) self.assertIsNone(self.lg.end(self.tid).reward) def test_open_ids(self): t2 = self.lg.start("s", "u", "sum", "r", "m", 0) self.lg.end(self.tid) self.assertEqual(self.lg.open_ids(), [t2]) class TestValidate(unittest.TestCase): def test_good(self): self.assertEqual(ad.validate_trace(mk()), []) def test_no_events(self): self.assertEqual(ad.validate_trace(mk([])), ["no events"]) def test_first_must_be_user(self): p = ad.validate_trace(mk([Event("assistant", "hi", 1)])) self.assertIn("first message is not from the user", p) def test_system_may_come_first(self): self.assertEqual(ad.validate_trace(mk([Event("system", "s", 0), Event("user", "u", 1)])), []) def test_tool_result_without_call(self): p = ad.validate_trace(mk([Event("user", "u", 1), Event("tool_result", "r", 2)])) self.assertTrue(any("tool_result without tool_call" in x for x in p)) def test_unanswered_call(self): p = ad.validate_trace(mk([Event("user", "u", 1), Event("tool_call", "c", 2)])) self.assertIn("ends with an unanswered tool_call", p) def test_two_calls_two_results(self): evs = [Event("user", "u", 1), Event("tool_call", "c", 2), Event("tool_call", "c", 3), Event("tool_result", "r", 4), Event("tool_result", "r", 5)] self.assertEqual(ad.validate_trace(mk(evs)), []) def test_time_backwards(self): p = ad.validate_trace(mk([Event("user", "u", 2), Event("assistant", "a", 1)])) self.assertTrue(any("timestamps go backwards" in x for x in p)) def test_event_before_start(self): p = ad.validate_trace(mk([Event("user", "u", 5)], started=10.0)) self.assertTrue(any("timestamps go backwards" in x for x in p)) def test_unknown_kind(self): p = ad.validate_trace(mk([Event("user", "u", 1), Event("banana", "x", 2)])) self.assertTrue(any("unknown kind" in x for x in p)) def test_bad_feedback_type(self): p = ad.validate_trace(mk([Event("user", "u", 1), Event("feedback", "", 2, {"type": "love"})])) self.assertTrue(any("bad feedback type" in x for x in p)) p = ad.validate_trace(mk([Event("user", "u", 1), Event("feedback", "", 2)])) self.assertTrue(any("bad feedback type" in x for x in p)) def test_reward_range(self): self.assertIn("reward outside 0..1", ad.validate_trace(mk(reward=2.0))) self.assertEqual(ad.validate_trace(mk(reward=None)), []) class TestJsonlAndSpans(unittest.TestCase): def test_round_trip(self): ts = ad.simulate_traffic(15, seed=2, user_sim=ad.UserSim(2)) self.assertEqual(ad.loads_jsonl(ad.dumps_jsonl(ts)), ts) def test_one_line_each_and_sorted_keys(self): ts = [mk(tid="a"), mk(tid="b")] txt = ad.dumps_jsonl(ts) lines = txt.splitlines() self.assertEqual(len(lines), 2) self.assertTrue(txt.endswith("\n")) self.assertEqual(list(json.loads(lines[0])), sorted(json.loads(lines[0]))) def test_loads_skips_blank_lines(self): txt = ad.dumps_jsonl([mk(tid="a")]) self.assertEqual(len(ad.loads_jsonl("\n" + txt + "\n\n")), 1) def test_empty(self): self.assertEqual(ad.dumps_jsonl([]), "") self.assertEqual(ad.loads_jsonl(""), []) def test_span_tree(self): evs = [Event("user", "u", 1.0), Event("assistant", "think", 2.0), Event("tool_call", "c", 3.0), Event("tool_result", "r", 5.0), Event("assistant", "Answer: 1", 6.0)] root = ad.to_span_tree(mk(evs, started=0.5)) self.assertEqual((root["name"], root["start"], root["end"]), ("session", 0.5, 6.0)) self.assertEqual([k["name"] for k in root["children"]], ["llm_call", "tool_call", "llm_call"]) self.assertEqual([(k["start"], k["end"]) for k in root["children"]], [(2.0, 3.0), (3.0, 5.0), (6.0, 6.0)]) self.assertEqual(root["children"][0]["attributes"]["chars"], 5) self.assertEqual(ad.span_summary(root), {"llm_call": 2, "tool_call": 1, "seconds": 5.5}) def test_span_tree_no_events(self): root = ad.to_span_tree(mk([], started=3.0)) self.assertEqual((root["start"], root["end"], root["children"]), (3.0, 3.0, [])) class TestWorld(unittest.TestCase): def test_solve_each_type(self): self.assertEqual(ad.solve("sum", "Add the numbers: 3, 9, 4"), "16") self.assertEqual(ad.solve("reverse", "Reverse the order of these words: red green blue"), "blue green red") self.assertEqual(ad.solve("upper", "Write this text in uppercase: hello world"), "HELLO WORLD") self.assertEqual(ad.solve("vowels", "Count the vowels in the word: banana"), "3") self.assertEqual(ad.solve("vowels", "Count the vowels in the word: rhythm"), "0") self.assertEqual(ad.solve("vowels", "Count the vowels in the word: Queue"), "4") self.assertEqual(ad.solve("vowels", "Count the vowels in the word: APPLE"), "2") self.assertEqual(ad.solve("upper", "Write this text in uppercase: a: b"), "A: B") self.assertEqual(ad.solve("sort", "Sort these numbers from smallest to largest: 5, 20, 9"), "5, 9, 20") self.assertEqual(ad.solve("domain", "What is the domain part of this email address: ada@example.org"), "example.org") def test_solve_errors(self): with self.assertRaises(ValueError): ad.solve("sum", "Add the numbers") with self.assertRaises(ValueError): ad.solve("nope", "x: y") def test_make_prompt_unknown(self): with self.assertRaises(ValueError): ad.make_prompt("nope", random.Random(0)) def test_make_task_is_consistent_and_deterministic(self): for k in ad.TYPES: a = ad.make_task(k, ad.rng_for("t", k), "id1") b = ad.make_task(k, ad.rng_for("t", k), "id1") self.assertEqual(a, b) self.assertEqual(a.answer, ad.solve(k, a.prompt)) self.assertEqual(a.type, k) def test_prompt_shapes(self): rng = random.Random(1) for _ in range(30): n = len(ad.make_prompt("sum", rng).split(": ")[1].split(",")) self.assertIn(n, (2, 3, 4)) n = len(ad.make_prompt("sort", rng).split(": ")[1].split(",")) self.assertIn(n, (3, 4, 5)) n = len(ad.make_prompt("reverse", rng).split(": ")[1].split()) self.assertIn(n, (3, 4)) self.assertEqual(len(ad.make_prompt("upper", rng).split(": ")[1].split()), 2) self.assertIn("@", ad.make_prompt("domain", rng)) def test_verifier(self): self.assertEqual(ad.run_verifier({"type": "exact", "expected": "3"}, " 3\n"), 1.0) self.assertEqual(ad.run_verifier({"type": "exact", "expected": "3"}, "4"), 0.0) self.assertEqual(ad.run_verifier({"type": "exact", "expected": 3}, "3"), 1.0) with self.assertRaises(ValueError): ad.run_verifier({"type": "fuzzy"}, "x") t = ad.WorldTask("i", "sum", "p: 1", "1") self.assertEqual(ad.verifier_spec(t), {"type": "exact", "expected": "1"}) def test_corrupt_never_returns_the_answer(self): for seed in range(50): rng = random.Random(seed) for a in ("0", "16", "-3", "blue green", "x", "5, 9, 20"): self.assertNotEqual(ad.corrupt("sum", a, rng), a) self.assertEqual(ad.corrupt("upper", "ab", random.Random(0)), "a") self.assertEqual(ad.corrupt("upper", "a", random.Random(0)), "a?") def test_prompt_key(self): self.assertEqual(len(ad.prompt_key("x")), 12) self.assertEqual(ad.prompt_key("x"), ad.prompt_key("x")) self.assertNotEqual(ad.prompt_key("x"), ad.prompt_key("y")) class TestMockAgent(unittest.TestCase): def task(self, k="sum"): return ad.make_task(k, ad.rng_for("w", k), "tid-" + k) def test_effective_skill(self): t = self.task() self.assertEqual(ad.MockAgent({"sum": 0.5}).effective_skill(t), 0.5) self.assertAlmostEqual(ad.MockAgent({"sum": 0.5}, "json", trained_harnesses=["react"]).effective_skill(t), 0.5 * ad.HARNESS_TRANSFER) self.assertEqual(ad.MockAgent({"sum": 0.5}, "json", trained_harnesses=["react", "json"]).effective_skill(t), 0.5) self.assertEqual(ad.MockAgent({"sum": 0.5}, "json").effective_skill(t), 0.5) def test_memory_floor(self): t = self.task() key = ad.prompt_key(t.prompt) self.assertEqual(ad.MockAgent({"sum": 0.5}, memorized={key}).effective_skill(t), ad.MEMORY_SKILL) self.assertEqual(ad.MockAgent({"sum": 0.99}, memorized={key}).effective_skill(t), 0.99) self.assertEqual(ad.MockAgent({"sum": 0.5}, memorized={"other"}).effective_skill(t), 0.5) def test_memory_and_harness_together(self): t = self.task() a = ad.MockAgent({"sum": 0.5}, "json", trained_harnesses=["react"], memorized={ad.prompt_key(t.prompt)}) self.assertEqual(a.effective_skill(t), ad.MEMORY_SKILL) def test_attempts_perfect_and_hopeless(self): t = self.task() self.assertEqual(ad.MockAgent({"sum": 1.0}).attempts(t), [t.answer]) out = ad.MockAgent({"sum": 0.0}).attempts(t) self.assertEqual(len(out), 2) self.assertTrue(all(x != t.answer for x in out)) def test_attempts_deterministic_and_seeded(self): t = self.task() a = ad.MockAgent({"sum": 0.5}, seed=1) self.assertEqual(a.attempts(t), ad.MockAgent({"sum": 0.5}, seed=1).attempts(t)) outs = {tuple(ad.MockAgent({"sum": 0.5}, seed=s).attempts(t)) for s in range(40)} self.assertGreater(len(outs), 1) def test_second_attempt_has_half_the_skill(self): a = ad.MockAgent({"sum": 0.6}) second = [a.attempts(ad.WorldTask(f"t{i}", "sum", "Add the numbers: 1, 2", "3")) for i in range(3000)] retried = [x for x in second if len(x) == 2] self.assertAlmostEqual(len(retried) / 3000, 0.4, delta=0.04) self.assertAlmostEqual(sum(x[1] == "3" for x in retried) / len(retried), 0.3, delta=0.04) def test_attempt_rates(self): wins = sum(ad.MockAgent({"sum": 0.7}, seed=s).attempts(self.task())[0] == self.task().answer for s in range(1000)) self.assertAlmostEqual(wins / 1000, 0.7, delta=0.05) def test_run_perfect_react(self): lg = ad.TraceLogger() t = self.task() tr = ad.MockAgent({"sum": 1.0}).run(t, lg, "u1", "s1", 100.0) self.assertEqual([e.kind for e in tr.events], ["system", "user", "assistant", "tool_call", "tool_result", "assistant"]) self.assertEqual(tr.reward, 1.0) self.assertEqual(tr.verifier, {"type": "exact", "expected": t.answer}) self.assertEqual(tr.events[3].text, f"solve[{t.prompt.split(': ', 1)[1]}]") self.assertEqual(tr.final_text(), f"Answer: {t.answer}") self.assertEqual(tr.events[1].meta["task_id"], t.task_id) self.assertEqual(ad.validate_trace(tr), []) self.assertEqual((tr.started, tr.user_id, tr.session_id, tr.harness, tr.task_type), (100.0, "u1", "s1", "react", "sum")) def test_run_json_harness(self): lg = ad.TraceLogger() t = self.task() tr = ad.MockAgent({"sum": 1.0}, "json").run(t, lg, "u", "s", 0) call = json.loads(tr.events[3].text) self.assertEqual(call["name"], "solve") self.assertEqual(call["arguments"]["input"], t.prompt.split(": ", 1)[1]) self.assertIn("thought", json.loads(tr.events[2].text)) self.assertIn("Harness: json", tr.events[0].text) def test_run_hopeless_retries_once(self): lg = ad.TraceLogger() tr = ad.MockAgent({"sum": 0.0}).run(self.task(), lg, "u", "s", 0) self.assertEqual([e.kind for e in tr.events].count("tool_call"), 2) self.assertEqual(tr.reward, 0.0) self.assertIn("try again", tr.events[5].text) self.assertEqual(ad.validate_trace(tr), []) def test_run_times_increase(self): tr = ad.MockAgent({"sum": 0.0}).run(self.task(), ad.TraceLogger(), "u", "s", 5.0) ts = [e.ts for e in tr.events] self.assertEqual(ts, sorted(ts)) self.assertEqual(ts[0], 5.0) def test_run_consent_passed_through(self): tr = ad.MockAgent().run(self.task(), ad.TraceLogger(), "u", "s", 0, consent={"train": True}) self.assertTrue(tr.consent["train"]) class TestUserSim(unittest.TestCase): def sessions(self, skill, n=600, seed=0): return ad.simulate_traffic(n, seed=seed, agent=ad.MockAgent({k: skill for k in ad.TYPES}, seed=seed), user_sim=ad.UserSim(seed)) def test_types_by_outcome(self): good = {"copy", "thumbs_up", "thumbs_down"} bad = {"regenerate", "edit", "thumbs_down", "abandon", "copy"} for t in self.sessions(1.0, 300): self.assertTrue({e.meta["type"] for e in t.feedback()} <= good) for t in self.sessions(0.0, 300): self.assertTrue({e.meta["type"] for e in t.feedback()} <= bad) def test_rates_on_success(self): ts = self.sessions(1.0, 1500) n = len(ts) copy = sum(any(e.meta["type"] == "copy" for e in t.feedback()) for t in ts) / n up = sum(any(e.meta["type"] == "thumbs_up" for e in t.feedback()) for t in ts) / n none = sum(not t.feedback() for t in ts) / n self.assertAlmostEqual(copy, 0.45, delta=0.05) self.assertAlmostEqual(up, 0.20, delta=0.05) self.assertAlmostEqual(none, 0.30, delta=0.05) def test_rates_on_failure(self): ts = self.sessions(0.0, 1500) n = len(ts) regen = sum(any(e.meta["type"] == "regenerate" for e in t.feedback()) for t in ts) / n reph = sum(sum(e.kind == "user" for e in t.events) == 2 for t in ts) / n none = sum(not t.feedback() and sum(e.kind == "user" for e in t.events) == 1 for t in ts) / n self.assertAlmostEqual(regen, 0.25, delta=0.05) self.assertAlmostEqual(reph, 0.10, delta=0.04) self.assertAlmostEqual(none, 0.10, delta=0.04) def test_edit_carries_the_right_answer(self): seen = 0 for t in self.sessions(0.0, 200): for e in t.feedback(): if e.meta["type"] == "edit": seen += 1 self.assertEqual(e.text, "Answer: " + t.verifier["expected"]) self.assertGreater(seen, 10) def test_rephrase_repeats_the_prompt(self): seen = 0 for t in self.sessions(0.0, 200): us = [e for e in t.events if e.kind == "user"] if len(us) == 2: seen += 1 self.assertEqual(us[1].text, "again: " + us[0].text) self.assertGreater(seen, 5) def test_deterministic(self): a, b = self.sessions(0.5, 50), self.sessions(0.5, 50) self.assertEqual(a, b) def test_no_sim_no_feedback(self): ts = ad.simulate_traffic(30, seed=1) self.assertEqual(sum(len(t.feedback()) for t in ts), 0) class TestTraffic(unittest.TestCase): def test_count_ids_and_logger(self): lg = ad.TraceLogger() ts = ad.simulate_traffic(40, seed=3, logger=lg) self.assertEqual(len(ts), 40) self.assertEqual(len({t.trace_id for t in ts}), 40) self.assertEqual(lg.all(), ts) self.assertEqual(lg.open_ids(), []) def test_weights(self): ts = ad.simulate_traffic(60, seed=1, weights={"sum": 1.0}) self.assertEqual({t.task_type for t in ts}, {"sum"}) ts = ad.simulate_traffic(200, seed=1, weights={"sum": 3.0, "sort": 1.0}) c = {k: sum(t.task_type == k for t in ts) for k in ("sum", "sort")} self.assertEqual(c["sum"] + c["sort"], 200) self.assertGreater(c["sum"], c["sort"] * 2) def test_deterministic_and_seeded(self): self.assertEqual(ad.simulate_traffic(20, seed=1), ad.simulate_traffic(20, seed=1)) self.assertNotEqual([t.prompt() for t in ad.simulate_traffic(20, seed=1)], [t.prompt() for t in ad.simulate_traffic(20, seed=2)]) def test_timing_users_and_consent(self): ts = ad.simulate_traffic(30, seed=1, start=1000.0, gap=60.0, n_users=3) self.assertEqual([t.started for t in ts], [1000.0 + 60.0 * i for i in range(30)]) self.assertTrue({t.user_id for t in ts} <= {"u00", "u01", "u02"}) for t in ts: self.assertEqual(t.consent, ad.user_consent(t.user_id)) def test_session_prefix_and_task_id(self): ts = ad.simulate_traffic(3, seed=4, session_prefix="zz") self.assertEqual([t.session_id for t in ts], ["zz4-0000", "zz4-0001", "zz4-0002"]) def test_user_consent(self): self.assertEqual(ad.user_consent("u1"), ad.user_consent("u1")) cs = [ad.user_consent(f"user{i}") for i in range(800)] self.assertTrue(all(c["log"] for c in cs)) self.assertAlmostEqual(sum(c["train"] for c in cs) / 800, 0.75, delta=0.05) self.assertAlmostEqual(sum(c["share"] for c in cs) / 800, 0.25, delta=0.05) self.assertTrue(all(c["train"] for c in cs if c["share"])) def full(tid, prompt, final, reward, user="u00", train=True, started=0.0, n_calls=1, ttype="sum", harness="react", extra=()): evs = [Event("user", prompt, started + 0.5)] t = started + 1.0 for _ in range(n_calls): evs += [Event("assistant", "Thought: go", t), Event("tool_call", "solve[x]", t + 0.1), Event("tool_result", "r", t + 0.2)] t += 1.0 evs.append(Event("assistant", final, t)) evs += list(extra) return mk(evs, reward, tid, user, train=train, started=started, ttype=ttype, harness=harness) class TestSft(unittest.TestCase): def test_only_passed_and_consented(self): ts = [full("a", "p1", "Answer: 3", 1.0), full("b", "p2", "Answer: 4", 0.0), full("c", "p3", "Answer: 5", 1.0, train=False)] out = ad.traces_to_sft(ts) self.assertEqual([e["trace_id"] for e in out], ["a"]) self.assertEqual([e["trace_id"] for e in ad.traces_to_sft(ts, require_consent=False)], ["a", "c"]) self.assertEqual([e["trace_id"] for e in ad.traces_to_sft(ts, min_reward=0.0)], ["a", "b"]) def test_unlabelled_and_invalid_dropped(self): unl = full("u", "p", "Answer: 1", None) bad = mk([Event("assistant", "x", 1.0)], 1.0, "bad") self.assertEqual(ad.traces_to_sft([unl, bad], min_reward=0.0), []) def test_partial_reward_threshold(self): t = full("h", "p", "Answer: 1", 0.5) self.assertEqual(len(ad.traces_to_sft([t], min_reward=0.5)), 1) self.assertEqual(len(ad.traces_to_sft([t], min_reward=0.6)), 0) def test_max_events(self): short, long_ = full("s", "p", "Answer: 1", 1.0), full("l", "p", "Answer: 1", 1.0, n_calls=3) self.assertEqual([e["trace_id"] for e in ad.traces_to_sft([short, long_], max_events=6)], ["s"]) self.assertEqual(len(ad.traces_to_sft([short, long_], max_events=None)), 2) self.assertEqual(len(ad.traces_to_sft([short], max_events=5)), 1) self.assertEqual(len(ad.traces_to_sft([short], max_events=4)), 0) def test_roles_and_loss_mask(self): ex = ad.traces_to_sft([full("a", "p1", "Answer: 3", 1.0)])[0] self.assertEqual([(m["role"], m["train"]) for m in ex["messages"]], [("user", False), ("assistant", True), ("assistant", True), ("tool", False), ("assistant", True)]) sysd = mk([Event("system", "s", 0.1), Event("user", "u", 0.5), Event("assistant", "a", 1.0)], 1.0) m = ad.traces_to_sft([sysd])[0]["messages"] self.assertEqual((m[0]["role"], m[0]["train"]), ("system", False)) def test_feedback_not_in_messages(self): ex = ad.traces_to_sft([full("a", "p", "Answer: 3", 1.0, extra=[fb("copy", ts=9.0)])])[0] self.assertEqual(len(ex["messages"]), 5) def test_fields(self): ex = ad.traces_to_sft([full("a", "p1", "Answer: 3", 1.0, user="u7", started=12.0, ttype="sort", harness="json")])[0] self.assertEqual((ex["task_type"], ex["harness"], ex["reward"], ex["started"], ex["user_id"], ex["prompt"]), ("sort", "json", 1.0, 12.0, "u7", "p1")) def test_loss_tokens_and_harness_mix(self): ex = ad.traces_to_sft([full("a", "one two three", "Answer: 3", 1.0)])[0] learned, ctx = ad.loss_tokens(ex) self.assertEqual(ctx, 3 + 1) self.assertEqual(learned, 2 + 2 + 2) mix = ad.harness_mix([{"harness": "react"}, {"harness": "json"}, {"harness": "react"}]) self.assertEqual(mix, {"json": 1, "react": 2}) self.assertEqual(list(mix), ["json", "react"]) class TestPreferences(unittest.TestCase): def test_pass_beats_fail(self): ts = [full("a", "p", "Answer: 3", 1.0), full("b", "p", "Answer: 4", 0.0)] out = ad.traces_to_preferences(ts) self.assertEqual(len(out), 1) self.assertEqual((out[0]["chosen"], out[0]["rejected"], out[0]["source"]), ("Answer: 3", "Answer: 4", "outcome")) self.assertEqual((out[0]["chosen_trace"], out[0]["rejected_trace"]), ("a", "b")) def test_no_gap_no_pair(self): self.assertEqual(ad.traces_to_preferences([full("a", "p", "x", 1.0), full("b", "p", "y", 1.0)]), []) self.assertEqual(ad.traces_to_preferences([full("a", "p", "x", 0.0), full("b", "p", "y", 0.0)]), []) def test_single_trace_no_pair(self): self.assertEqual(ad.traces_to_preferences([full("a", "p", "x", 1.0)]), []) def test_different_prompts_do_not_pair(self): self.assertEqual(ad.traces_to_preferences([full("a", "p1", "x", 1.0), full("b", "p2", "y", 0.0)]), []) def test_shorter_wins_ties_among_passes(self): ts = [full("long", "p", "Answer: 3", 1.0, n_calls=2), full("short", "p", "Answer: 3", 1.0, n_calls=1), full("bad", "p", "Answer: 9", 0.0)] self.assertEqual(ad.traces_to_preferences(ts)[0]["chosen_trace"], "short") def test_longer_trace_is_rejected_among_fails(self): ts = [full("ok", "p", "Answer: 3", 1.0), full("f1", "p", "Answer: 8", 0.0, n_calls=1), full("f2", "p", "Answer: 9", 0.0, n_calls=2)] self.assertEqual(ad.traces_to_preferences(ts)[0]["rejected_trace"], "f2") def test_graded_rewards(self): ts = [full("a", "p", "x", 0.5), full("b", "p", "y", 0.25), full("c", "p", "z", 0.75)] o = ad.traces_to_preferences(ts)[0] self.assertEqual((o["chosen_trace"], o["rejected_trace"]), ("c", "b")) def test_consent_and_unlabelled(self): ts = [full("a", "p", "x", 1.0), full("b", "p", "y", 0.0, train=False)] self.assertEqual(ad.traces_to_preferences(ts), []) self.assertEqual(len(ad.traces_to_preferences(ts, require_consent=False)), 1) self.assertEqual(ad.traces_to_preferences([full("a", "p", "x", 1.0), full("b", "p", "y", None)]), []) def test_invalid_trace_ignored(self): bad = mk([Event("assistant", "x", 1.0)], 0.0, "bad", prompt="p") self.assertEqual(ad.traces_to_preferences([full("a", "p", "x", 1.0), bad]), []) def test_sorted_by_prompt(self): ts = [full("a", "zz", "x", 1.0), full("b", "zz", "y", 0.0), full("c", "aa", "x", 1.0), full("d", "aa", "y", 0.0)] self.assertEqual([p["prompt"] for p in ad.traces_to_preferences(ts)], ["aa", "zz"]) class TestEditPairs(unittest.TestCase): def test_edit_becomes_chosen(self): t = full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", "Answer: 3", 9.0)]) out = ad.edit_pairs([t]) self.assertEqual((out[0]["chosen"], out[0]["rejected"], out[0]["source"]), ("Answer: 3", "Answer: 4", "edit")) def test_same_text_or_empty_ignored(self): self.assertEqual(ad.edit_pairs([full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", "Answer: 4", 9.0)])]), []) self.assertEqual(ad.edit_pairs([full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", " ", 9.0)])]), []) self.assertEqual(ad.edit_pairs([full("a", "p", "Answer: 4", 0.0, extra=[fb("copy", "", 9.0)])]), []) def test_whitespace_only_difference_is_not_an_edit(self): self.assertEqual(ad.edit_pairs([full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", " Answer: 4 ", 9.0)])]), []) def test_consent(self): t = full("a", "p", "Answer: 4", 0.0, train=False, extra=[fb("edit", "Answer: 3", 9.0)]) self.assertEqual(ad.edit_pairs([t]), []) self.assertEqual(len(ad.edit_pairs([t], require_consent=False)), 1) def test_one_pair_per_trace(self): t = full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", "Answer: 3", 9.0), fb("edit", "Answer: 2", 10.0)]) out = ad.edit_pairs([t]) self.assertEqual((len(out), out[0]["chosen"]), (1, "Answer: 3")) class TestEvalCases(unittest.TestCase): def test_failures_only_by_default(self): ts = [full("a", "p1", "x", 1.0), full("b", "p2", "y", 0.0)] self.assertEqual([c["prompt"] for c in ad.traces_to_eval_cases(ts)], ["p2"]) self.assertEqual([c["prompt"] for c in ad.traces_to_eval_cases(ts, only_failed=False)], ["p1", "p2"]) def test_dedupe_by_prompt(self): ts = [full("a", "p", "x", 0.0), full("b", "p", "y", 0.0)] c = ad.traces_to_eval_cases(ts) self.assertEqual((len(c), c[0]["from_trace"]), (1, "a")) def test_max_per_type(self): ts = [full(f"s{i}", f"p{i}", "x", 0.0, ttype="sum") for i in range(5)] + [full("o", "q", "x", 0.0, ttype="sort")] c = ad.traces_to_eval_cases(ts, max_per_type=2) self.assertEqual(sorted(x["task_type"] for x in c), ["sort", "sum", "sum"]) def test_needs_verifier_and_label(self): nov = mk(None, 0.0, "n", verifier=False) unl = full("u", "pp", "x", None) self.assertEqual(ad.traces_to_eval_cases([nov, unl]), []) def test_fields(self): c = ad.traces_to_eval_cases([full("a", "p", "x", 0.0, user="u4", started=7.0)])[0] self.assertEqual(c["case_id"], "c-" + ad.prompt_key("p")) self.assertEqual((c["verifier"]["type"], c["user_id"], c["started"], c["task_type"]), ("exact", "u4", 7.0, "sum")) ts = [full("a", "p", "x", 0.0)] ad.traces_to_eval_cases(ts)[0]["verifier"]["type"] = "changed" self.assertEqual(ts[0].verifier["type"], "exact") def test_invalid_trace_skipped(self): bad = mk([Event("assistant", "x", 1.0)], 0.0, "bad") self.assertEqual(ad.traces_to_eval_cases([bad]), []) class TestShingles(unittest.TestCase): def test_shingles(self): self.assertEqual(ad.shingles("a b c d e"), {"a b c", "b c d", "c d e"}) self.assertEqual(ad.shingles("a b c d e", k=2), {"a b", "b c", "c d", "d e"}) self.assertEqual(ad.shingles("A, b! C"), {"a b c"}) self.assertEqual(ad.shingles("a b"), {"a b"}) self.assertEqual(ad.shingles(""), set()) def test_jaccard(self): self.assertEqual(ad.jaccard({1, 2}, {1, 2}), 1.0) self.assertEqual(ad.jaccard({1}, {2}), 0.0) self.assertAlmostEqual(ad.jaccard({1, 2}, {2, 3}), 1 / 3) self.assertEqual(ad.jaccard(set(), set()), 1.0) self.assertEqual(ad.jaccard(set(), {1}), 0.0) def test_exact_key(self): self.assertEqual(ad.exact_key("Hello, World"), ad.exact_key("hello world")) self.assertNotEqual(ad.exact_key("hello world"), ad.exact_key("hello worlds")) self.assertEqual(len(ad.exact_key("x")), 16) class TestMinHash(unittest.TestCase): def test_shape_and_determinism(self): s = {"a", "b", "c"} sig = ad.minhash(s, 32, 0) self.assertEqual(len(sig), 32) self.assertEqual(sig, ad.minhash(s, 32, 0)) self.assertNotEqual(sig, ad.minhash(s, 32, 1)) self.assertEqual(ad.minhash(set(), 4), (2 ** 64 - 1,) * 4) def test_similarity(self): a = ad.minhash({str(i) for i in range(30)}, 256) self.assertEqual(ad.minhash_similarity(a, a), 1.0) b = ad.minhash({str(i) for i in range(15, 45)}, 256) self.assertAlmostEqual(ad.minhash_similarity(a, b), 1 / 3, delta=0.12) c = ad.minhash({"x" + str(i) for i in range(30)}, 256) self.assertLess(ad.minhash_similarity(a, c), 0.1) with self.assertRaises(ValueError): ad.minhash_similarity(a, ad.minhash({"1"}, 8)) def corpus(n=60, seed=0): rng = random.Random(seed) vocab = [f"w{i}" for i in range(40)] base = [" ".join(rng.choice(vocab) for _ in range(30)) for _ in range(n // 2)] docs = {f"d{i}": t for i, t in enumerate(base)} for i, t in enumerate(base[:10]): w = t.split() w[15] = "changed" docs[f"n{i}"] = " ".join(w) return docs class TestNearDuplicates(unittest.TestCase): def test_identical_found(self): out = ad.near_duplicates({"a": "the quick brown fox jumps", "b": "The quick brown fox jumps!", "c": "something else entirely here"}) self.assertEqual(out, [("a", "b", 1.0)]) def test_one_word_changed_in_long_text(self): t = " ".join(f"w{i}" for i in range(40)) u = t.replace("w20", "zzz") out = ad.near_duplicates({"a": t, "b": u}, threshold=0.8) self.assertEqual(len(out), 1) self.assertAlmostEqual(out[0][2], 35 / 41, places=3) self.assertEqual(ad.near_duplicates({"a": t, "b": u}, threshold=0.9), []) def test_matches_brute_force(self): docs = corpus() fast = ad.near_duplicates(docs, 0.7) slow = ad.brute_force_near_duplicates(docs, 0.7) self.assertEqual(fast, slow) self.assertGreaterEqual(len(slow), 10) def test_threshold_is_inclusive(self): out = ad.near_duplicates({"a": "one two three four", "b": "one two three four"}, threshold=1.0) self.assertEqual(out, [("a", "b", 1.0)]) def test_candidates_are_a_small_superset(self): docs = corpus(60, 1) cands = ad.lsh_candidates(docs) pairs = {(a, b) for a, b, _ in ad.near_duplicates(docs, 0.7)} self.assertTrue(pairs <= set(cands)) self.assertLess(len(cands), 0.1 * (len(docs) * (len(docs) - 1) / 2)) self.assertEqual(cands, sorted(cands)) with self.assertRaises(ValueError): ad.lsh_candidates(docs, num_perm=64, bands=10) def test_partly_similar_texts_are_rarely_candidates(self): head = " ".join(f"s{i}" for i in range(10)) docs = {f"d{i}": head + " " + " ".join(f"u{i}x{j}" for j in range(20)) for i in range(20)} self.assertAlmostEqual(ad.jaccard(ad.shingles(docs["d0"]), ad.shingles(docs["d1"])), 8 / 48, places=3) self.assertLess(len(ad.lsh_candidates(docs)), 20) def test_output_sorted_pairs(self): out = ad.near_duplicates(corpus(), 0.7) self.assertEqual(out, sorted(out)) self.assertTrue(all(a < b for a, b, _ in out)) def test_bands_must_divide(self): with self.assertRaises(ValueError): ad.near_duplicates({"a": "x y z"}, num_perm=64, bands=10) def test_empty_and_single(self): self.assertEqual(ad.near_duplicates({}), []) self.assertEqual(ad.near_duplicates({"a": "x y z w"}), []) def test_brute_force_threshold(self): d = {"a": "a b c d", "b": "a b c e"} self.assertEqual(ad.brute_force_near_duplicates(d, 0.3), [("a", "b", 0.3333)]) self.assertEqual(ad.brute_force_near_duplicates(d, 0.4), []) ITEMS = [{"id": i, "text": f"prompt number {i} about topic {i % 7}", "group": f"g{i % 25}", "ts": float(i)} for i in range(200)] class TestSplits(unittest.TestCase): def test_group_split_keeps_groups_together(self): sp = ad.split_by_group(ITEMS, lambda x: x["group"]) where = {} for name, items in sp.items(): for it in items: where.setdefault(it["group"], set()).add(name) self.assertTrue(all(len(v) == 1 for v in where.values())) self.assertEqual(sum(len(v) for v in sp.values()), 200) def test_group_split_proportions_and_salt(self): items = [{"g": i} for i in range(3000)] sp = ad.split_by_group(items, lambda x: x["g"]) self.assertAlmostEqual(len(sp["train"]) / 3000, 0.8, delta=0.03) self.assertAlmostEqual(len(sp["dev"]) / 3000, 0.1, delta=0.03) sp2 = ad.split_by_group(items, lambda x: x["g"], salt="other") self.assertNotEqual([x["g"] for x in sp["test"]], [x["g"] for x in sp2["test"]]) self.assertEqual(sp, ad.split_by_group(items, lambda x: x["g"])) def test_group_split_boundary_goes_to_the_next_split(self): from unittest import mock with mock.patch.object(ad, "stable_hash", return_value=800000): sp = ad.split_by_group([{"g": 1}], lambda x: x["g"]) self.assertEqual({k: len(v) for k, v in sp.items()}, {"train": 0, "dev": 1, "test": 0}) with mock.patch.object(ad, "stable_hash", return_value=999999): sp = ad.split_by_group([{"g": 1}], lambda x: x["g"]) self.assertEqual(len(sp["test"]), 1) def test_group_split_names_and_validation(self): sp = ad.split_by_group(ITEMS, lambda x: x["id"], (0.5, 0.5), ("a", "b")) self.assertEqual(set(sp), {"a", "b"}) self.assertTrue(sp["a"] and sp["b"]) with self.assertRaises(ValueError): ad.split_by_group(ITEMS, lambda x: x["id"], (0.5, 0.4)) with self.assertRaises(ValueError): ad.split_by_group(ITEMS, lambda x: x["id"], (0.5, 0.5), ("a", "b", "c")) def test_random_split(self): sp = ad.split_random(ITEMS) self.assertEqual({k: len(v) for k, v in sp.items()}, {"train": 160, "dev": 20, "test": 20}) ids = [x["id"] for v in sp.values() for x in v] self.assertEqual(sorted(ids), list(range(200))) self.assertEqual(sp, ad.split_random(ITEMS)) self.assertNotEqual(sp["test"], ad.split_random(ITEMS, seed=1)["test"]) self.assertEqual([x["id"] for x in sp["train"]], sorted(x["id"] for x in sp["train"])) with self.assertRaises(ValueError): ad.split_random(ITEMS, (0.9, 0.2, 0.1)) def test_random_split_loses_nothing_to_rounding(self): sp = ad.split_random(list(range(10)), (0.33, 0.33, 0.34)) self.assertEqual(sum(len(v) for v in sp.values()), 10) self.assertEqual(len(sp["test"]), 4) def test_random_split_last_gets_remainder(self): sp = ad.split_random(list(range(7)), (0.5, 0.25, 0.25)) self.assertEqual(sum(len(v) for v in sp.values()), 7) def test_time_split_boundaries(self): sp = ad.time_split(ITEMS, lambda x: x["ts"], 100.0, 150.0) self.assertEqual({k: len(v) for k, v in sp.items()}, {"train": 100, "dev": 50, "test": 50}) self.assertEqual(max(x["ts"] for x in sp["train"]), 99.0) self.assertEqual(min(x["ts"] for x in sp["dev"]), 100.0) self.assertEqual(min(x["ts"] for x in sp["test"]), 150.0) with self.assertRaises(ValueError): ad.time_split(ITEMS, lambda x: x["ts"], 150.0, 100.0) ad.time_split(ITEMS, lambda x: x["ts"], 100.0, 100.0) def it(text, group="g", ts=0.0): return {"text": text, "group": group, "ts": ts} class TestLeakage(unittest.TestCase): key = staticmethod(lambda x: x["text"]) def rep(self, splits, **kw): return ad.leakage_report(splits, self.key, **kw) def test_clean(self): r = self.rep({"train": [it("alpha beta gamma delta")], "test": [it("one two three four five")]}) self.assertTrue(r["ok"]) self.assertEqual((r["exact"], r["near"], r["groups"], r["time_violations"]), ([], [], [], 0)) def test_exact_copy_ignores_case_and_punctuation(self): r = self.rep({"train": [it("Alpha beta gamma delta.")], "test": [it("alpha beta gamma delta")]}) self.assertEqual(r["exact"], [("train:0", "test:0")]) self.assertEqual(r["near"], []) self.assertFalse(r["ok"]) def test_near_copy(self): t = " ".join(f"w{i}" for i in range(40)) r = self.rep({"train": [it(t)], "dev": [it(t.replace("w20", "zzz"))]}) self.assertEqual(len(r["near"]), 1) self.assertEqual(r["near"][0][:2], ("train:0", "dev:0")) self.assertEqual(r["exact"], []) def test_near_threshold(self): t = " ".join(f"w{i}" for i in range(40)) r = ad.leakage_report({"train": [it(t)], "dev": [it(t.replace("w20", "zzz"))]}, self.key, threshold=0.95) self.assertEqual(r["near"], []) def test_near_copies_inside_train_are_not_leaks(self): r = self.rep({"train": [it("a b c d e f"), it("a b c d e f")], "test": [it("totally different text here")]}) self.assertTrue(r["ok"]) def test_shared_group(self): r = self.rep({"train": [it("aaa bbb ccc ddd", "u1")], "test": [it("eee fff ggg hhh", "u1")]}, group_key=lambda x: x["group"]) self.assertEqual(r["groups"], [("test", "u1")]) self.assertFalse(r["ok"]) r = self.rep({"train": [it("aaa bbb ccc ddd", "u1")], "test": [it("eee fff ggg hhh", "u2")]}, group_key=lambda x: x["group"]) self.assertTrue(r["ok"]) def test_time_violation(self): sp = {"train": [it("aaa bbb ccc ddd", ts=10.0)], "dev": [it("eee fff ggg hhh", ts=10.0)], "test": [it("iii jjj kkk lll", ts=11.0)]} r = self.rep(sp, ts_key=lambda x: x["ts"]) self.assertEqual(r["time_violations"], 1) self.assertFalse(r["ok"]) sp["dev"][0]["ts"] = 10.5 self.assertEqual(self.rep(sp, ts_key=lambda x: x["ts"])["time_violations"], 0) def test_without_train_nothing_to_check(self): self.assertTrue(self.rep({"dev": [it("a b c d")], "test": [it("a b c d")]})["ok"]) def test_multiple_splits_are_each_checked(self): r = self.rep({"train": [it("alpha beta gamma delta")], "dev": [it("alpha beta gamma delta")], "test": [it("alpha beta gamma delta")]}) self.assertEqual(len(r["exact"]), 2) def test_group_split_by_text_passes_the_check(self): sp = ad.split_by_group(ITEMS, lambda x: x["text"]) self.assertTrue(self.rep(sp, group_key=lambda x: x["text"])["ok"]) def test_random_split_of_duplicated_prompts_fails_the_check(self): items = [it(f"repeat this prompt number {i % 10} please", f"g{i}", float(i)) for i in range(100)] r = self.rep(ad.split_random(items), group_key=lambda x: x["text"]) self.assertFalse(r["ok"]) self.assertGreater(len(r["exact"]), 0) class TestDecontaminate(unittest.TestCase): key = staticmethod(lambda x: x["text"]) def test_drops_exact_and_near(self): t = " ".join(f"w{i}" for i in range(40)) train = [it(t), it(t.replace("w20", "zzz")), it("keep me please now yes"), it("another distinct row here")] test = [it(t)] kept, dropped = ad.decontaminate(train, test, self.key) self.assertEqual([x["text"] for x in kept], ["keep me please now yes", "another distinct row here"]) self.assertEqual(len(dropped), 2) def test_exact_with_case(self): kept, dropped = ad.decontaminate([it("Hello there my friend")], [it("hello there my FRIEND")], self.key) self.assertEqual((kept, len(dropped)), ([], 1)) def test_exact_copies_go_even_when_near_matching_is_off(self): kept, dropped = ad.decontaminate([it("alpha beta gamma delta"), it("other words here now")], [it("Alpha beta gamma delta")], self.key, threshold=1.01) self.assertEqual(([x["text"] for x in kept], len(dropped)), (["other words here now"], 1)) def test_nothing_to_drop(self): kept, dropped = ad.decontaminate([it("aaa bbb ccc ddd")], [it("eee fff ggg hhh")], self.key) self.assertEqual((len(kept), dropped), (1, [])) self.assertEqual(ad.decontaminate([], [it("x y z w")], self.key), ([], [])) def test_near_dups_inside_train_are_kept(self): kept, _ = ad.decontaminate([it("a b c d e f"), it("a b c d e f")], [it("zzz yyy xxx www")], self.key) self.assertEqual(len(kept), 2) def test_threshold(self): t = " ".join(f"w{i}" for i in range(40)) kept, _ = ad.decontaminate([it(t.replace("w20", "zzz"))], [it(t)], self.key, threshold=0.95) self.assertEqual(len(kept), 1) def ex(kind, reward, prompt="p", harness="react"): return {"task_type": kind, "reward": reward, "prompt": prompt, "harness": harness} class TestTrainOn(unittest.TestCase): def test_no_examples_no_change(self): s = {"sum": 0.5, "sort": 0.4} self.assertEqual(ad.train_on(s, []), s) self.assertIsNot(ad.train_on(s, []), s) def test_step_formula(self): s = {"sum": 0.5, "sort": 0.4} out = ad.train_on(s, [ex("sum", 1.0)] * 15) self.assertAlmostEqual(out["sum"], 0.75) self.assertEqual(out["sort"], 0.4) def test_scale_caps_the_step(self): s = {"sum": 0.5} self.assertAlmostEqual(ad.train_on(s, [ex("sum", 1.0)] * 30)["sum"], 0.75) self.assertAlmostEqual(ad.train_on(s, [ex("sum", 1.0)] * 3)["sum"], 0.5 + 0.5 * (3 / 15) * 0.5) self.assertAlmostEqual(ad.train_on(s, [ex("sum", 1.0)] * 10, scale=10)["sum"], 0.75) def test_lr(self): self.assertAlmostEqual(ad.train_on({"sum": 0.5}, [ex("sum", 1.0)] * 15, lr=1.0)["sum"], 0.99) self.assertAlmostEqual(ad.train_on({"sum": 0.5}, [ex("sum", 1.0)] * 15, lr=0.2)["sum"], 0.5 + 0.2 * 0.5) def test_bad_data_hurts_and_even_data_does_nothing(self): s = {"sum": 0.5} self.assertAlmostEqual(ad.train_on(s, [ex("sum", 0.0)] * 15)["sum"], 0.5 - 0.5 * 0.5) self.assertAlmostEqual(ad.train_on(s, [ex("sum", 1.0), ex("sum", 0.0)] * 8)["sum"], 0.5) def test_bad_data_moves_a_strong_skill_down_proportionally(self): self.assertAlmostEqual(ad.train_on({"sum": 0.8}, [ex("sum", 0.0)] * 15)["sum"], 0.8 - 0.5 * 0.8) def test_clamped(self): self.assertEqual(ad.train_on({"sum": 0.98}, [ex("sum", 1.0)] * 15, lr=5.0)["sum"], 0.99) self.assertEqual(ad.train_on({"sum": 0.02}, [ex("sum", 0.0)] * 15, lr=5.0)["sum"], 0.01) def test_input_not_mutated(self): s = {"sum": 0.5} ad.train_on(s, [ex("sum", 1.0)] * 15) self.assertEqual(s, {"sum": 0.5}) class TestEvalAndLabels(unittest.TestCase): def test_make_eval_set(self): cs = ad.make_eval_set() self.assertEqual(len(cs), 60) self.assertEqual(len({c.task_id for c in cs}), 60) self.assertEqual(cs, ad.make_eval_set()) self.assertEqual(len(ad.make_eval_set(per_type=3)), 18) self.assertNotEqual([c.prompt for c in cs], [c.prompt for c in ad.make_eval_set(seed=1)]) self.assertEqual({c.type for c in cs}, set(ad.TYPES)) def test_expected_score(self): cs = ad.make_eval_set(per_type=2) self.assertAlmostEqual(ad.expected_score(ad.MockAgent(), cs), sum(ad.INIT_SKILL.values()) / 6) self.assertEqual(ad.expected_score(ad.MockAgent({k: 1.0 for k in ad.TYPES}), cs), 1.0) drop = {ad.prompt_key(c.prompt) for c in cs if c.type == "sum"} a = ad.MockAgent({k: (0.0 if k == "sum" else 1.0) for k in ad.TYPES}) self.assertEqual(ad.expected_score(a, cs, drop), 1.0) self.assertEqual(ad.expected_score(a, cs, {ad.prompt_key(c.prompt) for c in cs}), 0.0) def test_label_traces(self): ts = ad.simulate_traffic(30, seed=1) rng = random.Random(0) lab = ad.label_traces(ts, 10, 0.0, rng) self.assertEqual(len(lab), 10) by = {t.trace_id: t for t in ts} self.assertTrue(all(v == by[k].reward for k, v in lab.items())) self.assertEqual(ad.label_traces(ts, 0, 0.0, random.Random(0)), {}) self.assertEqual(ad.label_traces(ts, -5, 0.0, random.Random(0)), {}) self.assertEqual(len(ad.label_traces(ts, 99, 0.0, random.Random(0))), 30) def test_label_noise(self): ts = ad.simulate_traffic(30, seed=1) by = {t.trace_id: t for t in ts} lab = ad.label_traces(ts, 30, 1.0, random.Random(0)) self.assertTrue(all(v == 1.0 - by[k].reward for k, v in lab.items())) ts = ad.simulate_traffic(600, seed=1) by = {t.trace_id: t for t in ts} lab = ad.label_traces(ts, 600, 0.3, random.Random(0)) flipped = sum(v != by[k].reward for k, v in lab.items()) / len(lab) self.assertAlmostEqual(flipped, 0.3, delta=0.08) def test_label_skips_unlabelled(self): t = mk(reward=None, tid="x") self.assertEqual(ad.label_traces([t], 5, 0.0, random.Random(0)), {}) def test_label_deterministic(self): ts = ad.simulate_traffic(30, seed=1) self.assertEqual(ad.label_traces(ts, 10, 0.2, random.Random(5)), ad.label_traces(ts, 10, 0.2, random.Random(5))) class TestFlywheel(unittest.TestCase): @classmethod def setUpClass(cls): cls.base = ad.run_flywheel(ad.FlywheelConfig()) def test_history_shape(self): h = self.base["history"] self.assertEqual([r["round"] for r in h], list(range(6))) self.assertTrue(all(r["logged"] == 80 and r["labelled"] == 50 for r in h)) self.assertEqual(len(self.base["traces"]), 480) def test_deterministic(self): self.assertEqual(ad.run_flywheel(ad.FlywheelConfig())["history"], self.base["history"]) self.assertNotEqual(ad.run_flywheel(ad.FlywheelConfig(seed=1))["history"], self.base["history"]) def test_it_improves(self): h = self.base["history"] self.assertGreater(h[-1]["eval_clean"], h[0]["eval_clean"] + 0.15) self.assertGreater(h[-1]["online"], h[0]["online"] + 0.1) self.assertTrue(all(self.base["skill"][k] >= ad.INIT_SKILL[k] for k in ad.TYPES)) self.assertTrue(all(1 <= h[i]["round"] for i in range(1, 6))) def test_leaky_score_is_higher_than_clean(self): h = self.base["history"] self.assertGreater(h[-1]["eval_all"], h[-1]["eval_clean"] + 0.01) self.assertEqual(h[0]["eval_overlap"] > 0, h[0]["eval_all"] != h[0]["eval_clean"]) self.assertGreater(h[-1]["eval_overlap"], h[0]["eval_overlap"]) def test_overlap_counts_prompts_seen_in_training(self): mem = self.base["memorized"] n = sum(ad.prompt_key(c.prompt) in mem for c in ad.make_eval_set()) self.assertEqual(n, self.base["history"][-1]["eval_overlap"]) def test_no_labels_means_no_learning(self): r = ad.run_flywheel(ad.FlywheelConfig(label_budget=0)) self.assertTrue(all(h["sft"] == 0 for h in r["history"])) self.assertEqual(len({h["eval_clean"] for h in r["history"]}), 1) self.assertEqual(r["skill"], ad.INIT_SKILL) def test_zero_rounds(self): r = ad.run_flywheel(ad.FlywheelConfig(rounds=0)) self.assertEqual((r["history"], r["traces"]), ([], [])) def test_filter_failures_matters(self): raw = ad.run_flywheel(ad.FlywheelConfig(filter_failures=False)) self.assertGreaterEqual(raw["history"][0]["sft"], self.base["history"][0]["sft"]) self.assertLess(raw["history"][-1]["eval_clean"], self.base["history"][-1]["eval_clean"]) def test_consent_limits_the_data(self): free = ad.run_flywheel(ad.FlywheelConfig(use_consent=False)) self.assertGreater(sum(h["sft"] for h in free["history"]), sum(h["sft"] for h in self.base["history"])) def test_label_noise_hurts(self): noisy = ad.run_flywheel(ad.FlywheelConfig(label_noise=0.5)) self.assertLess(noisy["history"][-1]["eval_clean"], self.base["history"][-1]["eval_clean"]) def test_harness_mismatch(self): same = ad.run_flywheel(ad.FlywheelConfig())["history"][-1]["eval_clean"] diff = ad.run_flywheel(ad.FlywheelConfig(deploy_harness="json"))["history"][-1]["eval_clean"] mix = ad.run_flywheel(ad.FlywheelConfig(harnesses=("react", "json"), deploy_harness="json"))["history"][-1]["eval_clean"] self.assertLess(diff, same - 0.1) self.assertGreater(mix, diff + 0.1) def test_init_skill_used_and_not_mutated(self): init = {k: 0.2 for k in ad.TYPES} r = ad.run_flywheel(ad.FlywheelConfig(rounds=1, label_budget=0), init) self.assertEqual(r["skill"], init) self.assertEqual(init, {k: 0.2 for k in ad.TYPES}) class TestCollapse(unittest.TestCase): P = {"a": 0.5, "b": 0.25, "c": 0.15, "d": 0.06, "e": 0.03, "f": 0.01} def test_rows_and_start(self): rows = ad.collapse_demo(self.P, 40, 5, 1) self.assertEqual([r["generation"] for r in rows], list(range(6))) self.assertEqual(rows[0]["support"], 6) self.assertEqual(rows[0]["p_rarest"], 0.01) self.assertAlmostEqual(rows[0]["entropy"], ad.entropy(list(self.P.values())), places=3) def test_unnormalised_input(self): a = ad.collapse_demo({k: v * 10 for k, v in self.P.items()}, 40, 3, 1) self.assertEqual(a, ad.collapse_demo(self.P, 40, 3, 1)) def test_tails_vanish_and_never_return(self): rows = ad.collapse_demo(self.P, 30, 25, 4) sup = [r["support"] for r in rows] self.assertEqual(sup, sorted(sup, reverse=True)) self.assertLess(sup[-1], sup[0]) self.assertLess(rows[-1]["entropy"], rows[0]["entropy"]) def test_real_data_keeps_the_tails(self): rows = ad.collapse_demo(self.P, 30, 25, 4, real_mix=0.2) self.assertTrue(all(r["support"] == 6 for r in rows)) self.assertGreaterEqual(min(r["p_rarest"] for r in rows), 0.2 * 0.01 - 1e-4) def test_a_single_item_stays_a_probability(self): rows = ad.collapse_demo({"a": 1.0}, 10, 3, 1) self.assertEqual([r["p_rarest"] for r in rows], [1.0] * 4) self.assertEqual([r["support"] for r in rows], [1] * 4) def test_full_mix_never_moves(self): rows = ad.collapse_demo(self.P, 30, 5, 4, real_mix=1.0) self.assertEqual(len({r["entropy"] for r in rows}), 1) def test_bigger_samples_collapse_slower(self): small = ad.collapse_demo(self.P, 20, 20, 2)[-1]["support"] big = ad.collapse_demo(self.P, 5000, 20, 2)[-1]["support"] self.assertGreaterEqual(big, small) self.assertEqual(big, 6) def test_deterministic(self): self.assertEqual(ad.collapse_demo(self.P, 30, 6, 9), ad.collapse_demo(self.P, 30, 6, 9)) self.assertNotEqual(ad.collapse_demo(self.P, 30, 6, 9), ad.collapse_demo(self.P, 30, 6, 10)) def cand(i, kind, prompt, answer, origin="clean"): return {"id": f"c{i}", "type": kind, "prompt": prompt, "answer": answer, "origin": origin} class TestSyntheticTasks(unittest.TestCase): def test_paraphrase_light(self): p = "Add the numbers: 3, 9, 4" q = ad.paraphrase_light(p) self.assertEqual(q, "Sum these numbers: 3, 9, 4") self.assertEqual(ad.solve("sum", q), ad.solve("sum", p)) self.assertEqual(ad.paraphrase_light("Odd head: x y"), "Odd head: x y") for k in ad.TYPES: pr = ad.make_prompt(k, random.Random(1)) self.assertNotEqual(ad.paraphrase_light(pr), pr) self.assertEqual(ad.solve(k, ad.paraphrase_light(pr)), ad.solve(k, pr)) def test_propose_basics(self): cs = ad.propose_candidates(300, 1) self.assertEqual(len(cs), 300) self.assertEqual(len({c["id"] for c in cs}), 300) self.assertEqual(cs, ad.propose_candidates(300, 1)) self.assertNotEqual(cs, ad.propose_candidates(300, 2)) self.assertEqual({c["origin"] for c in cs}, {"clean", "bad_answer", "invalid", "near_dup", "dup"}) def test_propose_without_noise_is_all_clean(self): cs = ad.propose_candidates(100, 1, 0, 0, 0, 0) self.assertEqual({c["origin"] for c in cs}, {"clean"}) self.assertTrue(all(ad.filter_valid(c) is None and ad.filter_verified(c) is None for c in cs)) def test_propose_origins_do_what_they_say(self): cs = ad.propose_candidates(600, 3) by_prompt = {} for c in cs: if c["origin"] == "bad_answer": self.assertEqual(ad.filter_verified(c), "wrong_answer") elif c["origin"] == "invalid": self.assertIsNotNone(ad.filter_valid(c)) elif c["origin"] == "near_dup": self.assertIsNone(ad.filter_verified(c)) earlier = cs[:cs.index(c)] self.assertTrue(any(o["origin"] == "clean" and o["answer"] == c["answer"] and o["prompt"] != c["prompt"] and ad.payload_words(o["prompt"]) == ad.payload_words(c["prompt"]) for o in earlier)) by_prompt.setdefault(c["prompt"], []).append(c) self.assertTrue(any(c["origin"] == "dup" for c in cs)) def test_propose_rates(self): cs = ad.propose_candidates(3000, 5) frac = lambda o: sum(c["origin"] == o for c in cs) / 3000 self.assertAlmostEqual(frac("bad_answer"), 0.10, delta=0.02) self.assertAlmostEqual(frac("invalid"), 0.08, delta=0.02) self.assertAlmostEqual(frac("near_dup"), 0.12, delta=0.02) self.assertAlmostEqual(frac("dup"), 0.05, delta=0.02) def test_first_candidate_cannot_be_a_copy(self): for seed in range(30): self.assertNotIn(ad.propose_candidates(1, seed, 0, 0, 1.0, 0)[0]["origin"], ("near_dup", "dup")) def test_filter_valid(self): ok = cand(0, "sum", "Add the numbers: 3, 9, 4", "16") self.assertIsNone(ad.filter_valid(ok)) self.assertEqual(ad.filter_valid(dict(ok, prompt=" ")), "empty") self.assertEqual(ad.filter_valid(dict(ok, prompt="")), "empty") self.assertEqual(ad.filter_valid(dict(ok, prompt="Do it")), "length") self.assertEqual(ad.filter_valid(dict(ok, prompt="Add the numbers: " + "1, " * 100)), "length") self.assertEqual(ad.filter_valid(dict(ok, prompt="Add the numbers: {a}, {b}")), "placeholder") self.assertEqual(ad.filter_valid(dict(ok, prompt="Add the numbers: TODO fill")), "placeholder") self.assertIsNone(ad.filter_valid(dict(ok, prompt="x" * 15))) self.assertEqual(ad.filter_valid(dict(ok, prompt="x" * 14)), "length") self.assertIsNone(ad.filter_valid(dict(ok, prompt="x" * 200))) self.assertEqual(ad.filter_valid(dict(ok, prompt="x" * 201)), "length") def test_filter_verified(self): ok = cand(0, "sum", "Add the numbers: 3, 9, 4", "16") self.assertIsNone(ad.filter_verified(ok)) self.assertIsNone(ad.filter_verified(dict(ok, answer=16))) self.assertIsNone(ad.filter_verified(dict(ok, answer=" 16 "))) self.assertEqual(ad.filter_verified(dict(ok, answer="17")), "wrong_answer") self.assertEqual(ad.filter_verified(dict(ok, prompt="Add the numbers")), "unparseable") self.assertEqual(ad.filter_verified(dict(ok, prompt="Add the numbers: x, y")), "unparseable") self.assertEqual(ad.filter_verified(dict(ok, type="nope")), "unparseable") self.assertEqual(ad.filter_verified(cand(1, "domain", "Give the domain: nobody", "x")), "unparseable") def test_payload_words(self): self.assertEqual(ad.payload_words("Add the numbers: 3, 9, 4"), {"3", "9", "4"}) self.assertEqual(ad.payload_words("no colon here"), {"no", "colon", "here"}) def test_dedupe_exact_and_paraphrase(self): a = cand(0, "sum", "Add the numbers: 3, 9, 4", "16") b = cand(1, "sum", "Sum these numbers: 3, 9, 4", "16", "near_dup") c = cand(2, "sum", "Add the numbers: 3, 9, 7", "19") d = cand(3, "sum", "Add the numbers: 3, 9, 4", "16", "dup") kept, rej = ad.filter_dedupe([a, b, c, d]) self.assertEqual([x["id"] for x in kept], ["c0", "c2"]) self.assertEqual([x["id"] for x in rej], ["c1", "c3"]) def test_dedupe_is_per_type(self): a = cand(0, "sum", "Add the numbers: 3, 9", "12") b = cand(1, "sort", "Sort these numbers from smallest to largest: 3, 9", "3, 9") self.assertEqual(len(ad.filter_dedupe([a, b])[0]), 2) def test_dedupe_threshold_edge(self): a = cand(0, "sum", "Add the numbers: 1, 2, 3, 4", "10") b = cand(1, "sum", "Add the numbers: 1, 2, 3, 4, 5", "15") self.assertEqual(len(ad.filter_dedupe([a, b], threshold=0.8)[0]), 1) self.assertEqual(len(ad.filter_dedupe([a, b], threshold=0.81)[0]), 2) def test_dedupe_against_existing(self): a = cand(0, "sum", "Add the numbers: 3, 9, 4", "16") kept, rej = ad.filter_dedupe([a], existing=[cand(9, "sum", "Sum these numbers: 3, 9, 4", "16")]) self.assertEqual((kept, len(rej)), ([], 1)) def test_pass_count(self): c = cand(0, "sum", "Add the numbers: 3, 9, 4", "16") self.assertEqual(ad.pass_count(c, ad.MockAgent({"sum": 1.0}), 8), 8) self.assertEqual(ad.pass_count(c, ad.MockAgent({"sum": 0.0}), 8), 0) self.assertEqual(ad.pass_count(c, ad.MockAgent({"sum": 0.5}), 0), 0) a = ad.MockAgent({"sum": 0.5}) self.assertEqual(ad.pass_count(c, a, 8, 1), ad.pass_count(c, a, 8, 1)) mean_ = sum(ad.pass_count(dict(c, id=f"x{i}"), a, 8) for i in range(500)) / 500 self.assertAlmostEqual(mean_, 4.0, delta=0.3) def test_pass_count_depends_on_the_seed(self): a = ad.MockAgent({"sum": 0.5}) cs = [cand(i, "sum", "Add the numbers: 1, 2", "3") for i in range(40)] self.assertNotEqual([ad.pass_count(c, a, 8, 1) for c in cs], [ad.pass_count(c, a, 8, 2) for c in cs]) def test_difficulty_filter(self): cs = [cand(i, "sum", f"Add the numbers: {i}, 1", str(i + 1)) for i in range(60)] kept, rej = ad.filter_difficulty(cs, ad.MockAgent({"sum": 0.5}), 8) self.assertTrue(all(1 <= c["passes"] <= 7 for c in kept)) self.assertTrue(all(c["passes"] in (0, 8) for c in rej)) self.assertGreater(len(kept), 50) self.assertEqual(len(kept) + len(rej), 60) self.assertEqual(ad.filter_difficulty(cs, ad.MockAgent({"sum": 1.0}), 8)[0], []) self.assertEqual(ad.filter_difficulty(cs, ad.MockAgent({"sum": 0.0}), 8)[0], []) self.assertNotIn("passes", cs[0]) def test_difficulty_is_noisy_for_easy_tasks(self): cs = [cand(i, "sum", f"Add the numbers: {i}, 1", str(i + 1)) for i in range(400)] _, rej = ad.filter_difficulty(cs, ad.MockAgent({"sum": 0.9}), 8) self.assertAlmostEqual(len(rej) / 400, 0.9 ** 8, delta=0.06) def test_pipeline(self): r = ad.generate_pipeline(300, ad.MockAgent(seed=1), seed=2) f = [c for _, c in r["funnel"]] self.assertEqual(f[0], 300) self.assertEqual(f, sorted(f, reverse=True)) self.assertEqual([n for n, _ in r["funnel"]], ["proposed", "valid and verified", "after dedupe", "after difficulty"]) self.assertEqual(sum(r["rejects"].values()) + len(r["kept"]), 300) self.assertEqual(f[-1], len(r["kept"])) self.assertEqual(sum(r["origin_proposed"].values()), 300) self.assertEqual(sum(r["origin_kept"].values()), len(r["kept"])) for bad in ("bad_answer", "invalid", "dup"): self.assertNotIn(bad, r["origin_kept"]) self.assertEqual(r, ad.generate_pipeline(300, ad.MockAgent(seed=1), seed=2)) def test_pipeline_rejects_match_hidden_origin(self): r = ad.generate_pipeline(400, ad.MockAgent(seed=1), seed=2) self.assertEqual(r["rejects"]["wrong_answer"], r["origin_proposed"]["bad_answer"]) self.assertEqual(r["rejects"].get("empty", 0) + r["rejects"].get("length", 0) + r["rejects"].get("placeholder", 0), r["origin_proposed"]["invalid"]) def test_pipeline_existing_pool_shrinks_output(self): a = ad.MockAgent(seed=1) first = ad.generate_pipeline(200, a, seed=2) again = ad.generate_pipeline(200, a, seed=3, existing=first["kept"]) alone = ad.generate_pipeline(200, a, seed=3) self.assertLess(len(again["kept"]), len(alone["kept"])) def test_diversity(self): ts = [cand(0, "sum", "Add the numbers: 1, 2", "3"), cand(1, "sum", "Add the numbers: 2, 3", "5"), cand(2, "sort", "Sort these numbers: 3, 1", "1, 3")] d = ad.diversity(ts) self.assertEqual((d["n"], d["types"], d["openings"]), (3, 2, 2)) self.assertAlmostEqual(d["type_entropy"], 0.918, places=3) self.assertEqual(ad.diversity([])["n"], 0) two = [cand(0, "sum", "Add the numbers: 1, 2", "3"), cand(1, "vowels", "Add the vowels now: ab", "1")] self.assertEqual(ad.diversity(two)["openings"], 2) class TestSignals(unittest.TestCase): def sig(self, **kw): base = {"thumbs_up": 0, "thumbs_down": 0, "regenerate": 0, "edit": 0, "copy": 0, "abandon": 0, "edit_ratio": 0.0, "rephrase": False} base.update(kw) return base def test_extract_counts(self): t = full("a", "p", "Answer: 4", 0.0, extra=[fb("regenerate", ts=8), fb("regenerate", ts=9), fb("thumbs_down", ts=10), fb("abandon", ts=11)]) s = ad.extract_signals(t) self.assertEqual((s["regenerate"], s["thumbs_down"], s["abandon"], s["copy"], s["n_feedback"]), (2, 1, 1, 0, 4)) self.assertFalse(s["rephrase"]) def test_extract_edit_ratio(self): t = full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", "Answer: 3", 9.0)]) s = ad.extract_signals(t) self.assertEqual(s["edit"], 1) self.assertAlmostEqual(s["edit_ratio"], 1 - 8 / 9, places=3) t = full("a", "p", "Answer: 4", 0.0, extra=[fb("edit", "Something else entirely", 9.0)]) self.assertGreater(ad.extract_signals(t)["edit_ratio"], 0.7) def test_extract_ignores_non_feedback_events_and_unknown_types(self): t = full("a", "p", "Answer: 4", 1.0, extra=[Event("feedback", "", 9.0, {"type": "love"})]) s = ad.extract_signals(t) self.assertEqual(s["n_feedback"], 0) def test_extract_rephrase(self): same = mk([Event("user", "add the numbers: 3, 9, 4", 1), Event("assistant", "x", 2), Event("user", "again: add the numbers: 3, 9, 4", 3)]) self.assertTrue(ad.extract_signals(same)["rephrase"]) diff = mk([Event("user", "add the numbers: 3, 9, 4", 1), Event("assistant", "x", 2), Event("user", "thanks a lot, now sort these", 3)]) self.assertFalse(ad.extract_signals(diff)["rephrase"]) self.assertFalse(ad.extract_signals(mk())["rephrase"]) def test_score_single_signals(self): self.assertEqual(ad.signals_to_score(self.sig()), 0.0) self.assertAlmostEqual(ad.signals_to_score(self.sig(copy=1)), 0.6) self.assertAlmostEqual(ad.signals_to_score(self.sig(copy=5)), 0.6) self.assertEqual(ad.signals_to_score(self.sig(thumbs_up=1)), 1.0) self.assertEqual(ad.signals_to_score(self.sig(thumbs_down=1)), -1.0) self.assertAlmostEqual(ad.signals_to_score(self.sig(regenerate=1)), -0.5) self.assertEqual(ad.signals_to_score(self.sig(regenerate=3)), -1.0) self.assertAlmostEqual(ad.signals_to_score(self.sig(abandon=1)), -0.8) self.assertAlmostEqual(ad.signals_to_score(self.sig(rephrase=True)), -0.6) def test_score_edit(self): self.assertAlmostEqual(ad.signals_to_score(self.sig(edit=1, edit_ratio=0.0)), -0.4) self.assertAlmostEqual(ad.signals_to_score(self.sig(edit=1, edit_ratio=1.0)), -0.8) self.assertEqual(ad.signals_to_score(self.sig(edit=0, edit_ratio=1.0)), 0.0) def test_score_combinations_and_clipping(self): self.assertAlmostEqual(ad.signals_to_score(self.sig(copy=1, thumbs_down=1)), -0.4) self.assertEqual(ad.signals_to_score(self.sig(thumbs_down=1, abandon=1, regenerate=2)), -1.0) self.assertEqual(ad.signals_to_score(self.sig(thumbs_up=1, copy=1)), 1.0) def test_weak_label_edges(self): self.assertEqual(ad.weak_label(0.4), 1) self.assertEqual(ad.weak_label(0.39), None) self.assertEqual(ad.weak_label(-0.4), 0) self.assertEqual(ad.weak_label(-0.39), None) self.assertEqual(ad.weak_label(0.0), None) self.assertEqual(ad.weak_label(0.5, hi=0.6), None) self.assertEqual(ad.weak_label(-0.5, lo=-0.3), 0) def test_signal_quality(self): good_ok = full("a", "p1", "x", 1.0, extra=[fb("thumbs_up", ts=9)]) good_bad = full("b", "p2", "x", 0.0, extra=[fb("thumbs_up", ts=9)]) bad_ok = full("c", "p3", "x", 1.0, extra=[fb("thumbs_down", ts=9)]) bad_bad = full("d", "p4", "x", 0.0, extra=[fb("thumbs_down", ts=9)]) silent = full("e", "p5", "x", 1.0) q = ad.signal_quality([good_ok, good_ok, good_bad, good_bad, bad_ok, bad_bad, silent]) self.assertEqual((q["tp"], q["fp"], q["fn"], q["tn"]), (2, 2, 1, 1)) self.assertAlmostEqual(q["coverage"], 6 / 7, places=3) self.assertAlmostEqual(q["precision_good"], 2 / 4, places=3) self.assertAlmostEqual(q["recall_good"], 2 / 3, places=3) self.assertAlmostEqual(q["accuracy"], 3 / 6, places=3) def test_partial_reward_is_not_a_pass(self): half = full("h", "p", "x", 0.5, extra=[fb("thumbs_up", ts=9)]) q = ad.signal_quality([half]) self.assertEqual((q["tp"], q["fp"]), (0, 1)) def test_signal_quality_empty(self): q = ad.signal_quality([]) self.assertEqual((q["n"], q["coverage"], q["accuracy"], q["precision_good"], q["recall_good"]), (0, 0.0, 0.0, 0.0, 0.0)) def test_signal_quality_on_simulated_traffic(self): ts = ad.simulate_traffic(600, seed=2, user_sim=ad.UserSim(2)) q = ad.signal_quality(ts) self.assertGreater(q["precision_good"], 0.9) self.assertGreater(q["accuracy"], 0.8) self.assertLess(q["coverage"], 0.95) class TestBradleyTerry(unittest.TestCase): def test_two_models(self): bt = ad.bradley_terry([("a", "b")] * 3 + [("b", "a")]) self.assertAlmostEqual(bt["a"] / bt["b"], 3.0, places=6) self.assertAlmostEqual(bt["a"] * bt["b"], 1.0, places=6) self.assertAlmostEqual(bt["a"] / (bt["a"] + bt["b"]), 0.75, places=6) def test_equal_votes_equal_strength(self): bt = ad.bradley_terry([("a", "b"), ("b", "a")]) self.assertAlmostEqual(bt["a"], bt["b"]) self.assertAlmostEqual(bt["a"], 1.0) def test_three_models_order_and_scale(self): votes = [("a", "b")] * 6 + [("b", "a")] * 2 + [("b", "c")] * 5 + [("c", "b")] * 3 + [("a", "c")] * 7 + [("c", "a")] * 1 bt = ad.bradley_terry(votes) self.assertGreater(bt["a"], bt["b"]) self.assertGreater(bt["b"], bt["c"]) self.assertAlmostEqual(bt["a"] * bt["b"] * bt["c"], 1.0, places=6) def test_mle_conditions(self): votes = [("a", "b")] * 6 + [("b", "a")] * 2 + [("b", "c")] * 5 + [("c", "b")] * 3 + [("a", "c")] * 7 + [("c", "a")] * 1 bt = ad.bradley_terry(votes) for i in "abc": w = sum(1 for x in votes if x[0] == i) e = sum((sum(1 for x in votes if set(x) == {i, j})) * bt[i] / (bt[i] + bt[j]) for j in "abc" if j != i) self.assertAlmostEqual(w, e, places=4) def test_chain_of_models(self): votes = [("a", "b")] * 4 + [("b", "a")] * 4 + [("b", "c")] * 2 + [("c", "b")] * 6 bt = ad.bradley_terry(votes) self.assertAlmostEqual(bt["a"], bt["b"], places=6) self.assertGreater(bt["c"], bt["b"]) self.assertAlmostEqual(bt["c"] / bt["b"], 3.0, places=4) class TestLuhn(unittest.TestCase): def test_luhn(self): self.assertTrue(ad.luhn_ok("4111111111111111")) self.assertTrue(ad.luhn_ok("4111 1111 1111 1111")) self.assertTrue(ad.luhn_ok("4111-1111-1111-1111")) self.assertFalse(ad.luhn_ok("4111111111111112")) self.assertFalse(ad.luhn_ok("411111111111")) self.assertFalse(ad.luhn_ok("41111111111111111111")) self.assertTrue(ad.luhn_ok("378282246310005")) self.assertTrue(ad.luhn_ok("6011111111111117")) self.assertFalse(ad.luhn_ok("0000 0000 0000")) self.assertTrue(ad.luhn_ok("0000 0000 0000 0")) self.assertTrue(ad.luhn_ok("0" * 19)) self.assertFalse(ad.luhn_ok("0" * 20)) class TestScrubber(unittest.TestCase): def s(self, text, names=()): sc = ad.Scrubber(names) return sc.scrub(text), sc def test_email_same_value_same_placeholder(self): out, sc = self.s("Write to ada@example.org and ADA@example.org, or bob@mail.test.") self.assertEqual(out, "Write to and , or .") self.assertEqual(sc.counts, {"EMAIL": 3}) def test_email_with_several_dots(self): self.assertEqual(self.s("x@a.b.c and y@mail.example.org")[0], " and ") def test_phone_formats(self): for p in ("(415) 555-0132", "415-555-0132", "415.555.0132", "+1 415 555 0132", "1-415-555-0132"): out, _ = self.s(f"call {p} now") self.assertEqual(out, "call now", p) def test_ssn(self): self.assertEqual(self.s("ssn 123-45-6789 ok")[0], "ssn ok") self.assertEqual(self.s("not 12-345-6789")[0], "not 12-345-6789") def test_card_needs_luhn(self): self.assertEqual(self.s("card 4111 1111 1111 1111.")[0], "card .") self.assertEqual(self.s("order 4111 1111 1111 1112.")[0], "order 4111 1111 1111 1112.") self.assertEqual(self.s("card 4111-1111-1111-1111")[0], "card ") def test_ip(self): self.assertEqual(self.s("host 10.0.0.12 up")[0], "host up") self.assertEqual(self.s("host 255.255.255.255")[0], "host ") self.assertEqual(self.s("version 999.1.1.1")[0], "version 999.1.1.1") self.assertEqual(self.s("version 1.2.3")[0], "version 1.2.3") def test_key(self): k = "sk-" + "A1" * 10 self.assertEqual(self.s(f"key {k} end")[0], "key end") self.assertEqual(self.s("key sk-abc end")[0], "key sk-abc end") self.assertEqual(self.s("AKIA" + "Z9" * 10)[0], "") def test_names(self): out, sc = self.s("Ada Lovelace met ada lovelace and Adam.", ["Ada Lovelace"]) self.assertEqual(out, " met and Adam.") out, _ = self.s("Ada and Adam", ["Ada", "Adam"]) self.assertEqual(out, " and ") self.assertEqual(self.s("hello Ada", [""])[0], "hello Ada") self.assertEqual(self.s("Radar", ["Ada"])[0], "Radar") def test_name_with_regex_characters(self): self.assertEqual(self.s("call me a.b (ok)", ["a.b"])[0], "call me (ok)") self.assertEqual(self.s("call me axb", ["a.b"])[0], "call me axb") def test_numbering_is_per_kind(self): out, _ = self.s("a@x.org 415-555-0132 b@x.org 415-555-0199") self.assertEqual(out, " ") def test_known_misses(self): self.assertEqual(self.s("I live at 12 Oak Street, Springfield")[0], "I live at 12 Oak Street, Springfield") self.assertEqual(self.s("my name is Grace Hopper")[0], "my name is Grace Hopper") def test_state_is_kept_between_calls(self): sc = ad.Scrubber() self.assertEqual(sc.scrub("a@x.org"), "") self.assertEqual(sc.scrub("b@x.org then a@x.org"), " then ") def test_non_string_input(self): self.assertEqual(ad.Scrubber().scrub(42), "42") def test_residual(self): self.assertEqual(ad.residual_pii("a@x.org and 123-45-6789"), ["EMAIL", "SSN"]) self.assertEqual(ad.residual_pii("nothing here"), []) self.assertEqual(ad.residual_pii("order 4111 1111 1111 1112"), []) self.assertEqual(ad.residual_pii("card 4111 1111 1111 1111"), ["CARD"]) self.assertEqual(ad.residual_pii(" "), []) def test_pseudonymize(self): a = ad.pseudonymize_user("u1", "salt") self.assertEqual(a, ad.pseudonymize_user("u1", "salt")) self.assertNotEqual(a, ad.pseudonymize_user("u1", "other")) self.assertNotEqual(a, ad.pseudonymize_user("u2", "salt")) self.assertTrue(a.startswith("p-") and len(a) == 12) self.assertNotIn("u1", a) class TestScrubTrace(unittest.TestCase): def trace(self): evs = [Event("user", "Mail ada@example.org please", 1.0, {"note": "from ada@example.org", "n": 3}), Event("assistant", "Sent to ada@example.org and bob@x.org", 2.0)] t = mk(evs, 1.0, "t9", user="u7") t.verifier = {"type": "exact", "expected": "ada@example.org"} return t def test_scrubbed_copy(self): t = self.trace() s = ad.scrub_trace(t, "salt") self.assertEqual(s.events[0].text, "Mail please") self.assertEqual(s.events[1].text, "Sent to and ") self.assertEqual(s.events[0].meta, {"note": "from ", "n": 3}) self.assertEqual(s.verifier, {"type": "exact", "expected": ""}) self.assertTrue(s.scrubbed) self.assertEqual(s.user_id, ad.pseudonymize_user("u7", "salt")) def test_original_is_untouched(self): t = self.trace() ad.scrub_trace(t, "salt") self.assertIn("ada@example.org", t.events[0].text) self.assertEqual(t.user_id, "u7") self.assertFalse(t.scrubbed) self.assertEqual(t.verifier["expected"], "ada@example.org") def test_names_and_no_verifier(self): t = mk([Event("user", "Hi Grace", 1.0)], None, verifier=False) s = ad.scrub_trace(t, "s", names=["Grace"]) self.assertEqual((s.events[0].text, s.verifier), ("Hi ", None)) def test_other_fields_kept(self): t = self.trace() s = ad.scrub_trace(t, "salt") self.assertEqual((s.trace_id, s.reward, s.consent, s.harness, s.started), (t.trace_id, t.reward, t.consent, t.harness, t.started)) def test_trace_residual_pii(self): t = self.trace() self.assertEqual(ad.trace_residual_pii(t), ["EMAIL"]) self.assertEqual(ad.trace_residual_pii(ad.scrub_trace(t, "s")), []) t2 = mk([Event("user", "x", 1.0, {"k": "123-45-6789"})]) self.assertEqual(ad.trace_residual_pii(t2), ["SSN"]) def test_simulated_domain_tasks(self): ts = ad.simulate_traffic(80, seed=2) self.assertTrue(any(ad.trace_residual_pii(t) for t in ts)) self.assertFalse(any(ad.trace_residual_pii(ad.scrub_trace(t, "s")) for t in ts)) class TestConsent(unittest.TestCase): def traces(self): return [mk(tid="a", train=True, scrubbed=True), mk(tid="b", train=False, scrubbed=True), mk(tid="c", train=True, scrubbed=False)] def test_train(self): kept, why = ad.consent_filter(self.traces(), "train") self.assertEqual([t.trace_id for t in kept], ["a"]) self.assertEqual(why, {"no_consent": 1, "not_scrubbed": 1}) def test_log_needs_no_scrubbing(self): kept, why = ad.consent_filter(self.traces(), "log") self.assertEqual(len(kept), 3) self.assertEqual(why, {"no_consent": 0, "not_scrubbed": 0}) def test_share(self): ts = self.traces() ts[0].consent["share"] = True ts[2].consent["share"] = True kept, why = ad.consent_filter(ts, "share") self.assertEqual([t.trace_id for t in kept], ["a"]) self.assertEqual(why, {"no_consent": 1, "not_scrubbed": 1}) def test_unknown_purpose(self): with self.assertRaises(ValueError): ad.consent_filter(self.traces(), "sell") def test_default_consent_blocks_training(self): t = mk(scrubbed=True) t.consent = dict(ad.DEFAULT_CONSENT) self.assertEqual(ad.consent_filter([t], "train")[0], []) t.consent = {} self.assertEqual(ad.consent_filter([t], "log")[0], []) class TestRetentionAndDeletion(unittest.TestCase): def test_retention_boundary(self): day = 86400.0 ts = [mk(tid="old", started=0.0), mk(tid="edge", started=10 * day), mk(tid="new", started=19 * day)] kept, n = ad.retention_filter(ts, now=40 * day, max_age_days=30) self.assertEqual(([t.trace_id for t in kept], n), (["edge", "new"], 1)) kept, n = ad.retention_filter(ts, now=40 * day, max_age_days=31) self.assertEqual(([t.trace_id for t in kept], n), (["edge", "new"], 1)) kept, n = ad.retention_filter(ts, now=40 * day, max_age_days=41) self.assertEqual((len(kept), n), (3, 0)) kept, n = ad.retention_filter(ts, now=40 * day, max_age_days=25) self.assertEqual(([t.trace_id for t in kept], n), (["new"], 2)) def test_retention_none_left(self): kept, n = ad.retention_filter([mk(started=0.0)], now=1e9, max_age_days=1) self.assertEqual((kept, n), ([], 1)) def test_delete_user(self): ts = [mk(tid="a", user="u1"), mk(tid="b", user="u2"), mk(tid="c", user="u1")] kept, n = ad.delete_user(ts, "u1") self.assertEqual(([t.trace_id for t in kept], n), (["b"], 2)) self.assertEqual(ad.delete_user(ts, "nobody")[1], 0) def test_delete_pseudonymized_user_needs_the_salt(self): ts = [ad.scrub_trace(mk(tid="a", user="u1"), "salt"), ad.scrub_trace(mk(tid="b", user="u2"), "salt")] self.assertEqual(ad.delete_user(ts, "u1")[1], 0) kept, n = ad.delete_user(ts, "u1", "salt") self.assertEqual((len(kept), n), (1, 1)) self.assertEqual(ad.delete_user(ts, "u1", "wrong-salt")[1], 0) def card_args(traces=None, **over): ts = traces if traces is not None else [mk(tid="a", ttype="sum", started=5.0, scrubbed=True, user="u1"), mk(tid="b", ttype="sort", started=9.0, scrubbed=True, user="u1", reward=0.0, harness="json")] kw = dict(name="demo", version="1.0", traces=ts, purpose="Testing.", sources=["a simulator"], filters=["scrubbed"], known_limits=["tiny"], not_for=["training real models"]) kw.update(over) return kw class TestDataCard(unittest.TestCase): def test_counts_come_from_the_data(self): c = ad.make_data_card(**card_args()) self.assertEqual(c["counts"], {"traces": 2, "by_task_type": {"sort": 1, "sum": 1}, "users": 1}) self.assertEqual(c["harnesses"], {"json": 1, "react": 1}) self.assertEqual(c["labels"], {"labelled": 2, "mean_reward": 0.5}) self.assertEqual(c["time_range"], [5.0, 9.0]) self.assertEqual(c["privacy"]["scrubbed"], 2) self.assertEqual(c["privacy"]["unscrubbed"], 0) self.assertEqual(c["privacy"]["train_consent"], 2) self.assertEqual(c["privacy"]["share_consent"], 0) self.assertEqual(c["privacy"]["residual_pii_traces"], 0) def test_unlabelled_and_empty(self): c = ad.make_data_card(**card_args([mk(reward=None, scrubbed=True)])) self.assertEqual(c["labels"], {"labelled": 0, "mean_reward": None}) e = ad.make_data_card(**card_args([])) self.assertEqual((e["time_range"], e["counts"]["traces"], e["labels"]["mean_reward"]), (None, 0, None)) def test_residual_pii_counted(self): t = mk([Event("user", "mail ada@example.org", 1.0)], 1.0, scrubbed=True) c = ad.make_data_card(**card_args([t, mk(tid="z")])) self.assertEqual(c["privacy"]["residual_pii_traces"], 1) def test_splits_and_leakage(self): rep = {"ok": False, "exact": [1, 2], "near": [3], "groups": [], "time_violations": 4} c = ad.make_data_card(**card_args(splits={"train": [1, 2, 3], "test": [4]}, leakage=rep)) self.assertEqual(c["splits"], {"train": 3, "test": 1}) self.assertEqual(c["leakage"], {"ok": False, "exact": 2, "near": 1, "shared_groups": 0, "time_violations": 4}) self.assertNotIn("leakage", ad.make_data_card(**card_args())) def test_inputs_are_copied(self): src = ["a"] c = ad.make_data_card(**card_args(sources=src)) src.append("b") self.assertEqual(c["sources"], ["a"]) def test_validate_ok(self): self.assertEqual(ad.validate_card(ad.make_data_card(**card_args())), []) def test_validate_missing_and_empty(self): c = ad.make_data_card(**card_args()) del c["version"] self.assertEqual(ad.validate_card(c), ["missing version"]) for f in ("purpose", "sources", "filters", "known_limits", "not_for"): c = ad.make_data_card(**card_args()) c[f] = [] if f != "purpose" else "" self.assertEqual(ad.validate_card(c), [f"empty {f}"]) self.assertEqual(len(ad.validate_card({})), len(ad.CARD_FIELDS)) def test_validate_contradictions(self): c = ad.make_data_card(**card_args(splits={"train": [1] * 5})) self.assertEqual(ad.validate_card(c), ["splits hold more items than the card counts"]) c = ad.make_data_card(**card_args(traces=[mk(tid="a", scrubbed=False)])) c["privacy"]["share_consent"] = 1 self.assertEqual(ad.validate_card(c), ["unscrubbed traces in a card that lists sharing"]) c = ad.make_data_card(**card_args(leakage={"ok": False, "exact": [1], "near": [], "groups": [], "time_violations": 0})) self.assertEqual(ad.validate_card(c), ["leakage report is not clean"]) def test_render(self): rep = {"ok": True, "exact": [], "near": [], "groups": [], "time_violations": 0} c = ad.make_data_card(**card_args(splits={"train": [1, 2], "test": []}, leakage=rep, synthetic_note="All mock.")) md = ad.render_card(c) for s in ("# Data card: demo (version 1.0)", "**Purpose.** Testing.", "- Traces: 2 from 1 users", "- By task type: sort 1, sum 1", "- Harnesses: json 1, react 1", "- Labelled: 2, mean reward 0.5", "Time range (seconds from start): 5 to 9", "- Scrubbed: 2, not scrubbed: 0", "- train 2, test 0", "Leakage checks: clean", "- Source: a simulator", "- Filter: scrubbed", "**Synthetic data.** All mock.", "## Known limits", "- tiny", "## Not for", "- training real models"): self.assertIn(s, md) self.assertTrue(md.endswith("\n")) self.assertEqual(md, ad.render_card(c)) def test_render_optional_parts(self): md = ad.render_card(ad.make_data_card(**card_args())) self.assertNotIn("Leakage", md) self.assertNotIn("Synthetic data", md) self.assertNotIn("## Splits", md) md = ad.render_card(ad.make_data_card(**card_args([]))) self.assertNotIn("Time range", md) rep = {"ok": False, "exact": [1], "near": [], "groups": [], "time_violations": 0} self.assertIn("NOT clean", ad.render_card(ad.make_data_card(**card_args(leakage=rep)))) class TestEndToEnd(unittest.TestCase): def test_safe_pipeline(self): """traffic -> scrub -> consent -> group split -> leakage clean -> card valid""" ts = ad.simulate_traffic(200, seed=7, user_sim=ad.UserSim(7)) sc = [ad.scrub_trace(t, "salt") for t in ts] keep, _ = ad.consent_filter(sc, "train") self.assertTrue(keep) self.assertFalse(any(ad.trace_residual_pii(t) for t in keep)) items = [{"t": t, "text": t.prompt()} for t in keep] sp = ad.split_by_group(items, lambda x: x["text"]) rep = ad.leakage_report(sp, lambda x: x["text"], lambda x: x["text"]) self.assertTrue(rep["ok"]) card = ad.make_data_card("e2e", "1", keep, "x", ["sim"], ["f"], ["l"], ["n"], splits=sp, leakage=rep) self.assertEqual(ad.validate_card(card), []) sft = ad.traces_to_sft(keep) self.assertTrue(sft) self.assertTrue(all(not m["train"] or m["role"] == "assistant" for e in sft for m in e["messages"])) def test_raw_traces_fail_the_gates(self): ts = ad.simulate_traffic(100, seed=7) self.assertEqual(ad.consent_filter(ts, "train")[0], []) self.assertTrue(any(ad.trace_residual_pii(t) for t in ts)) def test_traces_to_three_kinds_of_data(self): ts = ad.simulate_traffic(300, seed=8, user_sim=ad.UserSim(8)) self.assertTrue(ad.traces_to_sft(ts)) self.assertTrue(ad.edit_pairs(ts)) self.assertTrue(ad.traces_to_eval_cases(ts)) for c in ad.traces_to_eval_cases(ts): self.assertEqual(ad.run_verifier(c["verifier"], c["verifier"]["expected"]), 1.0) def test_generated_tasks_feed_the_agent(self): r = ad.generate_pipeline(100, ad.MockAgent(seed=1), seed=4) lg = ad.TraceLogger() agent = ad.MockAgent(seed=2) for c in r["kept"][:20]: t = ad.WorldTask(c["id"], c["type"], c["prompt"], c["answer"]) tr = agent.run(t, lg, "u", c["id"], 0.0) self.assertEqual(ad.validate_trace(tr), []) class TestDemo(unittest.TestCase): def run_demo(self): buf = io.StringIO() with contextlib.redirect_stdout(buf): ad.demo() return buf.getvalue() def test_demo_is_deterministic_and_complete(self): a = self.run_demo() self.assertEqual(a, self.run_demo()) for n in range(1, 10): self.assertIn(f"== {n}. ", a) self.assertIn("same: True", a) self.assertIn("card problems: []", a) self.assertIn("JSON Lines round trip equal: True", a) def test_main_guard(self): with open(ad.__file__) as fh: src = fh.read() self.assertIn('if __name__ == "__main__":', src) if __name__ == "__main__": unittest.main(verbosity=2)