import math import random import unittest from llm_toolkit import * LOGITS = [2.0, 0.5, 0.3, 0.2, -0.5, -1.0] SCHEMA = {"type": "object", "required": ["name", "intent"], "properties": { "name": {"type": "string"}, "intent": {"type": "string", "enum": ["meeting", "intro", "follow-up"]}}} class TestDecoding(unittest.TestCase): def test_softmax_sums_to_one_and_temperature_sharpens(self): lo, hi = softmax(LOGITS, 0.5), softmax(LOGITS, 2.0) self.assertAlmostEqual(sum(lo), 1.0) self.assertGreater(lo[0], hi[0]) self.assertLess(entropy(lo), entropy(hi)) def test_greedy_and_temperature_zero(self): self.assertEqual(greedy(softmax(LOGITS)), 0) rng = random.Random(1) self.assertEqual({decode_step(LOGITS, rng, temperature=0) for _ in range(20)}, {0}) def test_top_k_keeps_exactly_k(self): p = top_k_filter(softmax(LOGITS), 2) self.assertEqual(sum(1 for x in p if x > 0), 2) self.assertAlmostEqual(sum(p), 1.0) def test_top_p_keeps_smallest_nucleus(self): p = top_p_filter(softmax(LOGITS), 0.8) self.assertEqual([i for i, x in enumerate(p) if x > 0], [0, 1, 2]) self.assertEqual(sum(1 for x in top_p_filter(softmax(LOGITS), 1.0) if x > 0), len(LOGITS)) self.assertEqual(sum(1 for x in top_p_filter(softmax(LOGITS), 0.01) if x > 0), 1) def test_filtered_sampling_never_picks_removed_tokens(self): rng = random.Random(3) seen = {decode_step(LOGITS, rng, 1.0, top_k=2) for _ in range(300)} self.assertTrue(seen <= {0, 1}) def test_bad_arguments(self): with self.assertRaises(ValueError): softmax(LOGITS, 0) with self.assertRaises(ValueError): top_p_filter([0.5, 0.5], 0) class TestAttention(unittest.TestCase): Q = [[1.0, 0.0], [0.0, 1.0], [2.0, 0.0]] K = [[1.0, 0.0], [0.0, 1.0], [1.0, 0.0]] V = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]] def test_rows_sum_to_one(self): _, w = attention(self.Q, self.K, self.V) for row in w: self.assertAlmostEqual(sum(row), 1.0) def test_causal_mask_blocks_the_future(self): _, w = attention(self.Q, self.K, self.V, causal=True) self.assertEqual(w[0][1:], [0.0, 0.0]) self.assertEqual(w[1][2], 0.0) _, w2 = attention(self.Q, self.K, self.V, causal=False) self.assertGreater(w2[0][2], 0.0) def test_changing_a_future_token_does_not_change_the_past(self): base, _ = attention(self.Q, self.K, self.V) K2 = [r[:] for r in self.K]; K2[2] = [9.0, 9.0] V2 = [r[:] for r in self.V]; V2[2] = [9.0, 9.0] new, _ = attention(self.Q, K2, V2) self.assertEqual(base[:2], new[:2]) def test_linear_attention_parallel_equals_recurrent(self): a = linear_attention_parallel(self.Q, self.K, self.V) b = linear_attention_recurrent(self.Q, self.K, self.V) for r, s in zip(a, b): for x, y in zip(r, s): self.assertAlmostEqual(x, y) def test_kv_cache_arithmetic(self): self.assertEqual(kv_cache_bytes(32, 8, 128, 32768), 4 * 2**30) self.assertEqual(kv_cache_bytes(1, 1, 1, 10, bytes_per_value=2, batch=3), 2 * 10 * 2 * 3) class TestConstrained(unittest.TestCase): def test_validate_catches_each_kind_of_error(self): self.assertEqual(validate({"name": "A", "intent": "intro"}, SCHEMA), []) self.assertIn("missing required key 'intent'", validate({"name": "A"}, SCHEMA)[0]) self.assertIn("not one of", validate({"name": "A", "intent": "chat"}, SCHEMA)[0]) self.assertIn("expected string", validate({"name": 5, "intent": "intro"}, SCHEMA)[0]) self.assertTrue(validate(True, {"type": "integer"})) def test_extract_json_handles_chatter_and_fences(self): self.assertEqual(extract_json('Sure! {"a": {"b": "}"}} thanks'), {"a": {"b": "}"}}) self.assertEqual(extract_json('```json\n{"a": 1}\n```'), {"a": 1}) with self.assertRaises(ValueError): extract_json("no json here") def test_retry_loop_recovers_and_feeds_errors_back(self): m = MockJSONModel(['{"name": "S", "intent": "chat"}', 'oops', '{"name": "S", "intent": "intro"}']) value, attempts = generate_json(m, "task", SCHEMA) self.assertEqual((value["intent"], attempts), ("intro", 3)) self.assertIn("rejected", m.prompts[1]) self.assertIn("not one of", m.prompts[1]) def test_retry_loop_gives_up(self): m = MockJSONModel(["never json"]) with self.assertRaises(RuntimeError): generate_json(m, "task", SCHEMA, max_retries=2) self.assertEqual(m.calls, 3) def test_constrained_decoding_is_always_legal(self): vocab = ["meet", "ing", "in", "tro", "follow", "-up", "chat", "xyz"] table = {"meeting": ["meet", "ing"], "intro": ["in", "tro"], "follow-up": ["follow", "-up"]} cons = TrieConstraint(list(table), lambda s: table[s]) loves_illegal = lambda prefix: [0.0] * 6 + [9.0, 9.0] rng = random.Random(5) outs = {constrained_decode(loves_illegal, vocab, cons, rng) for _ in range(100)} self.assertTrue(outs <= set(table)) self.assertGreater(len(outs), 1) class TestContext(unittest.TestCase): def msgs(self): return [{"role": "system", "text": "be careful always", "pinned": True}, {"role": "user", "text": "fix the test"}, {"role": "tool", "text": "x " * 50}, {"role": "assistant", "text": "reading the file now"}, {"role": "tool", "text": "y " * 40}, {"role": "assistant", "text": "found the bug"}, {"role": "tool", "text": "z " * 5}, {"role": "user", "text": "now add a test"}] def test_no_change_when_it_fits(self): m = self.msgs() new, rep = compact(m, limit=1000) self.assertEqual(new, m) self.assertEqual(rep["steps"], []) def test_clearing_tool_results_keeps_the_last_two(self): out = clear_old_tool_results(self.msgs()) tools = [m["text"] for m in out if m["role"] == "tool"] self.assertEqual(tools[0], "[tool result cleared]") self.assertNotEqual(tools[1], "[tool result cleared]") def test_compaction_fits_and_keeps_pinned_and_recent(self): m = self.msgs() new, rep = compact(m, limit=40, reserve_output=5) self.assertTrue(rep["fits"]) self.assertLess(rep["after"], rep["before"]) self.assertEqual(new[0]["text"], "be careful always") self.assertEqual(new[-1]["text"], "now add a test") self.assertTrue(any(x["text"].startswith("Summary:") for x in new)) def test_edges_first_puts_best_at_the_ends(self): out = edges_first([1, 2, 3, 4, 5]) self.assertEqual((out[0], out[-1]), (1, 2)) self.assertEqual(out[len(out) // 2], 5) self.assertEqual(sorted(out), [1, 2, 3, 4, 5]) class TestSampling(unittest.TestCase): def test_coverage_and_vote(self): self.assertAlmostEqual(coverage(0.2, 1), 0.2) self.assertAlmostEqual(coverage(0.2, 10), 1 - 0.8 ** 10) self.assertEqual(majority_vote(["a", "b", "b", "c"]), "b") self.assertEqual(majority_vote(["x", "y"]), "x") if __name__ == "__main__": unittest.main(verbosity=1)