#!/usr/bin/env python3 """Tests for patterns_toolkit.py. Standard library only. Run: python test_patterns_toolkit.py""" import contextlib import io import json import os import sys import threading import unittest sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import patterns_toolkit as pt # noqa: E402 def toks(n): """A text of exactly n tokens.""" return " ".join("w%d" % i for i in range(n)) class TestMocksAndMeter(unittest.TestCase): def test_count_tokens_counts_words_and_punctuation(self): self.assertEqual(pt.count_tokens(""), 0) self.assertEqual(pt.count_tokens("hello world"), 2) self.assertEqual(pt.count_tokens("a, b."), 4) def test_scripted_lm_replays_in_order_and_records_prompts(self): lm = pt.ScriptedLM(["x", "y"]) self.assertEqual((lm("p1"), lm("p2")), ("x", "y")) self.assertEqual(lm.prompts, ["p1", "p2"]) def test_scripted_lm_raises_when_script_runs_out(self): lm = pt.ScriptedLM(["only"]) lm("a") with self.assertRaises(pt.ScriptExhausted): lm("b") def test_rule_lm_records_prompts(self): lm = pt.RuleLM(lambda p: p.upper()) self.assertEqual(lm("abc"), "ABC") self.assertEqual(lm.prompts, ["abc"]) def test_meter_counts_calls_and_tokens(self): m = pt.Meter() llm = m.wrap(pt.RuleLM(lambda p: "one two")) llm("a b c") llm("d") self.assertEqual((m.calls, m.tokens_in, m.tokens_out), (2, 4, 4)) def test_meter_is_safe_across_threads(self): m = pt.Meter() llm = m.wrap(pt.RuleLM(lambda p: "x")) threads = [threading.Thread(target=lambda: [llm("a b") for _ in range(50)]) for _ in range(4)] [t.start() for t in threads] [t.join() for t in threads] self.assertEqual((m.calls, m.tokens_in), (200, 400)) class TestCalcAndSearch(unittest.TestCase): def test_safe_eval_arithmetic(self): self.assertEqual(pt.safe_eval("2 + 3 * 4"), 14) self.assertEqual(pt.safe_eval("-(2 - 5) ** 2"), -9) self.assertEqual(pt.safe_eval("7 // 2 + 7 % 4"), 6) def test_safe_eval_rejects_names_calls_and_attributes(self): for bad in ("__import__('os')", "x + 1", "(1).real", "[1]", "'a' * 3", "True + 1"): with self.assertRaises((ValueError, SyntaxError), msg=bad): pt.safe_eval(bad) def test_safe_eval_blocks_huge_powers(self): with self.assertRaises(ValueError): pt.safe_eval("2 ** 1000") def test_calc_returns_text_and_never_raises(self): self.assertEqual(pt.calc("2009 - 1971"), "38") self.assertEqual(pt.calc("1 / 4"), "0.25") self.assertTrue(pt.calc("1 / 0").startswith("error")) self.assertTrue(pt.calc("1 +").startswith("error")) self.assertTrue(pt.calc("abc").startswith("error")) def test_search_picks_best_overlap(self): self.assertEqual(pt.search("who founded Harbor Labs"), "Mira Okoye") self.assertEqual(pt.search("year Harbor Labs was founded"), "2009") self.assertEqual(pt.search("birth year of Mira Okoye"), "1971") def test_search_needs_three_shared_words(self): self.assertEqual(pt.search("harbor"), "no result") self.assertEqual(pt.search("the weather in Paris"), "no result") class TestPromptChain(unittest.TestCase): def test_output_of_one_step_feeds_the_next(self): lm = pt.RuleLM(lambda p: "[" + p.split(":")[-1].strip() + "]") r = pt.prompt_chain(lm, ["A: {input}", "B: {input}"], "x") self.assertTrue(r.ok) self.assertEqual(r.outputs, ["[x]", "[[x]]"]) self.assertEqual(lm.prompts, ["A: x", "B: [x]"]) def test_gate_stops_the_chain_at_the_failing_step(self): lm = pt.ScriptedLM(["ok", "toolongtoolong", "never"]) r = pt.prompt_chain(lm, ["{input}"] * 3, "t", gate=lambda i, out: len(out) < 5) self.assertEqual((r.ok, r.stopped_at, len(r.outputs)), (False, 1, 2)) self.assertEqual(len(lm.prompts), 2) def test_no_steps_is_an_ok_empty_chain(self): r = pt.prompt_chain(pt.ScriptedLM([]), [], "t") self.assertEqual((r.ok, r.outputs), (True, [])) def test_outputs_are_stripped(self): r = pt.prompt_chain(pt.ScriptedLM([" hi \n"]), ["{input}"], "t") self.assertEqual(r.outputs, ["hi"]) class TestRouting(unittest.TestCase): def setUp(self): self.routes = {"password": (None, "P: {input}"), "refund": (None, "R: {input}")} def test_routes_to_the_matching_label(self): lm = pt.ScriptedLM(["refund", "done"]) self.assertEqual(pt.route(lm, "money back", self.routes), ("refund", "done")) self.assertEqual(lm.prompts[1], "R: money back") def test_label_match_is_case_insensitive_and_tolerates_extra_words(self): lm = pt.ScriptedLM(["Label: PASSWORD.", "ok"]) self.assertEqual(pt.route(lm, "x", self.routes)[0], "password") def test_unknown_label_goes_to_default(self): lm = pt.ScriptedLM(["weather", "ok"]) self.assertEqual(pt.route(lm, "x", self.routes, default="refund")[0], "refund") def test_unknown_label_without_default_raises(self): with self.assertRaises(pt.RoutingError): pt.route(pt.ScriptedLM(["weather"]), "x", self.routes) def test_a_route_can_use_a_different_model(self): cheap = pt.ScriptedLM(["cheap answer"]) routes = {"easy": (cheap, "E: {input}"), "hard": (None, "H: {input}")} clf = pt.ScriptedLM(["easy"]) self.assertEqual(pt.route(clf, "q", routes), ("easy", "cheap answer")) self.assertEqual(len(clf.prompts), 1) # the classifier was not used for the answer self.assertEqual(cheap.prompts, ["E: q"]) def test_label_must_match_a_whole_word(self): with self.assertRaises(pt.RoutingError): pt.route(pt.ScriptedLM(["passwords"]), "x", {"password": (None, "{input}")}) class TestSectioningAndVoting(unittest.TestCase): def test_sectioning_keeps_template_order(self): lm = pt.RuleLM(lambda p: p.split(":")[0]) self.assertEqual(pt.sectioning(lm, "t", ["a: {input}", "b: {input}", "c: {input}"]), ["a", "b", "c"]) def test_sectioning_really_runs_calls_at_the_same_time(self): barrier = threading.Barrier(2, timeout=3) def fn(prompt): barrier.wait() # both calls must be inside fn together return prompt self.assertEqual(pt.sectioning(pt.RuleLM(fn), "t", ["a {input}", "b {input}"], max_workers=2), ["a t", "b t"]) def test_voting_majority(self): v = pt.voting(pt.ScriptedLM(["B", "A", "B"]), "q", 3) self.assertEqual((v.winner, v.counts, round(v.agreement, 3)), ("B", {"B": 2, "A": 1}, 0.667)) def test_voting_tie_goes_to_first_seen(self): self.assertEqual(pt.voting(pt.ScriptedLM(["A", "B", "A", "B"]), "q", 4).winner, "A") def test_voting_applies_parse(self): v = pt.voting(pt.ScriptedLM(["Answer: 4", "answer: 4 ", "Answer: 5"]), "q", 3, parse=lambda s: s.lower().replace("answer:", "").strip()) self.assertEqual((v.winner, v.counts["4"]), ("4", 2)) def test_voting_needs_at_least_one_sample(self): with self.assertRaises(ValueError): pt.voting(pt.ScriptedLM([]), "q", 0) class TestOrchestrator(unittest.TestCase): def test_parse_subtasks_accepts_fenced_json(self): self.assertEqual(pt.parse_subtasks('Here:\n```json\n["a", " b "]\n```'), ["a", "b"]) def test_parse_subtasks_rejects_bad_input(self): for bad in ("no list", "[1, 2]", "[]", '[""]', "[not json]"): with self.assertRaises(pt.PlanError, msg=bad): pt.parse_subtasks(bad) def test_plan_work_and_combine(self): lm = pt.ScriptedLM(['["s1", "s2"]', "r1", "r2", "final text"]) r = pt.orchestrate(lm, "task") self.assertEqual((r.subtasks, r.final, r.plan_attempts), (["s1", "s2"], "final text", 1)) self.assertEqual(r.results, [("s1", "r1"), ("s2", "r2")]) self.assertIn("- s1: r1", lm.prompts[-1]) def test_a_bad_plan_is_retried_once(self): lm = pt.ScriptedLM(["I think we should...", '["only"]', "r", "done"]) r = pt.orchestrate(lm, "t") self.assertEqual(r.plan_attempts, 2) self.assertIn("not valid", lm.prompts[1]) def test_two_bad_plans_raise(self): with self.assertRaises(pt.PlanError): pt.orchestrate(pt.ScriptedLM(["nope", "still nope"]), "t") def test_too_many_subtasks_are_cut_and_flagged(self): lm = pt.ScriptedLM(['["a","b","c","d"]', "1", "2", "end"]) r = pt.orchestrate(lm, "t", max_subtasks=2) self.assertEqual((r.subtasks, r.truncated), (["a", "b"], True)) def test_custom_worker_is_used(self): lm = pt.ScriptedLM(['["a"]', "end"]) r = pt.orchestrate(lm, "t", worker=lambda s: "W:" + s) self.assertEqual(r.results, [("a", "W:a")]) self.assertEqual(len(lm.prompts), 2) class TestRefine(unittest.TestCase): def test_stops_when_the_evaluator_accepts(self): r = pt.refine(lambda t, fb: "x" * (len(fb) + 1), lambda d: (len(d) >= 3, "longer"), "t", max_rounds=5) self.assertEqual((r.converged, len(r.rounds), r.final), (True, 3, "xxx")) def test_feedback_is_passed_to_the_generator(self): seen = [] pt.refine(lambda t, fb: seen.append(list(fb)) or "d", lambda d: (False, "fix%d" % len(seen)), "t", max_rounds=3) self.assertEqual(seen, [[], ["fix1"], ["fix1", "fix2"]]) def test_gives_up_after_max_rounds_and_returns_last_draft(self): r = pt.refine(lambda t, fb: "d%d" % len(fb), lambda d: (False, "no"), "t", max_rounds=2) self.assertEqual((r.converged, r.final, len(r.rounds)), (False, "d1", 2)) def test_max_rounds_must_be_positive(self): with self.assertRaises(ValueError): pt.refine(lambda t, fb: "", lambda d: (True, ""), "t", max_rounds=0) class TestReAct(unittest.TestCase): def test_parse_react(self): self.assertEqual(pt.parse_react("Thought: look\nAction: search[a b]"), ("look", "search", "a b")) self.assertEqual(pt.parse_react("Action: finish[42]"), ("", "finish", "42")) self.assertEqual(pt.parse_react("just text")[1], None) self.assertEqual(pt.parse_react("Thought: t\nAction: search[x]\nAction: calc[1]")[1:], ("search", "x")) def test_parse_react_allows_brackets_inside_the_argument(self): self.assertEqual(pt.parse_react("Action: calc[(1+2)*[3]]")[1:], ("calc", "(1+2)*[3]")) def test_happy_path(self): r = pt.react(pt.RuleLM(pt.age_policy("Harbor Labs").fn), pt.QUESTION, pt.TOOLS) self.assertEqual((r.answer, r.status, len(r.steps)), ("38", "finished", 5)) self.assertEqual([s.tool for s in r.steps], ["search", "search", "search", "calc", "finish"]) self.assertEqual(r.steps[1].arg, "birth year of Mira Okoye") def test_unknown_tool_becomes_an_observation(self): lm = pt.ScriptedLM(["Action: fly[x]", "Action: finish[ok]"]) r = pt.react(lm, "q", pt.TOOLS) self.assertIn("unknown tool 'fly'", r.steps[0].observation) self.assertEqual(r.answer, "ok") def test_tool_exception_becomes_an_observation(self): def boom(arg): raise RuntimeError("disk full") r = pt.react(pt.ScriptedLM(["Action: boom[x]", "Action: finish[done]"]), "q", {"boom": boom}) self.assertEqual(r.steps[0].observation, "error: RuntimeError: disk full") def test_format_error_is_reported_and_the_loop_continues(self): lm = pt.ScriptedLM(["I will just answer.", "Action: finish[7]"]) r = pt.react(lm, "q", pt.TOOLS) self.assertIn("format error", r.steps[0].observation) self.assertEqual(r.answer, "7") self.assertIn("format error", lm.prompts[1]) # the model is told what went wrong def test_step_limit(self): lm = pt.RuleLM(lambda p: "Action: calc[%d + 1]" % len(p)) r = pt.react(lm, "q", pt.TOOLS, max_steps=4) self.assertEqual((r.status, r.answer, len(r.steps)), ("step_limit", None, 4)) def test_stuck_detector_stops_exact_repeats(self): lm = pt.RuleLM(lambda p: "Thought: again\nAction: search[nothing here]") r = pt.react(lm, "q", pt.TOOLS, max_steps=20, stuck_after=3) self.assertEqual((r.status, len(r.steps)), ("stuck", 3)) def test_a_changed_observation_is_not_stuck(self): counter = iter(range(100)) r = pt.react(pt.RuleLM(lambda p: "Action: tick[x]"), "q", {"tick": lambda a: str(next(counter))}, max_steps=6, stuck_after=2) self.assertEqual(r.status, "step_limit") def test_the_prompt_grows_every_turn(self): lm = pt.age_policy("Harbor Labs") pt.react(lm, pt.QUESTION, pt.TOOLS) sizes = [pt.count_tokens(p) for p in lm.prompts] self.assertEqual(sizes, sorted(sizes)) self.assertEqual(len(set(sizes)), len(sizes)) self.assertEqual(sum(sizes), 483) # my own count for the demo run def test_lessons_go_into_the_header_before_the_question(self): lm = pt.ScriptedLM(["Action: finish[1]"]) pt.react(lm, "the question", pt.TOOLS, memory=["lesson one", "lesson two"]) head, tail = lm.prompts[0].split("Question:") self.assertIn("- lesson one", head) self.assertIn("- lesson two", head) self.assertEqual(tail.strip(), "the question") def test_tool_names_are_listed_in_the_header(self): lm = pt.ScriptedLM(["Action: finish[1]"]) pt.react(lm, "q", pt.TOOLS) self.assertIn("calc, search", lm.prompts[0]) class TestReflexion(unittest.TestCase): def setUp(self): self.q = "How old was the founder of Ridge Corp when Ridge Corp was founded?" def run_it(self, buggy=True, **kw): def trial(mem): return pt.react(pt.age_policy("Ridge Corp", buggy=buggy), self.q, pt.TOOLS, memory=mem) def evaluate(res): return (True, "ok") if res.answer == "40" else (False, "Wrong: the answer %s is negative." % res.answer) reflect = lambda res, fb: "I subtracted in the wrong order. " + fb # noqa: E731 return pt.reflexion(trial, evaluate, reflect, **kw) def test_a_lesson_fixes_the_second_trial(self): r = self.run_it() self.assertEqual([t.ok for t in r.trials], [False, True]) self.assertEqual([t.answer for t in r.trials], ["-40", "40"]) self.assertEqual(r.trials[0].memory_used, []) self.assertEqual(len(r.trials[1].memory_used), 1) def test_no_reflection_is_needed_when_the_first_trial_passes(self): r = self.run_it(buggy=False) self.assertEqual((r.success, len(r.trials), r.memory), (True, 1, [])) def test_gives_up_after_max_trials(self): n = [] out = pt.reflexion(lambda mem: pt.AgentResult("x", "finished", []), lambda res: (False, "no"), lambda res, fb: n.append(1) or "lesson%d" % len(n), max_trials=3) self.assertEqual((out.success, len(out.trials)), (False, 3)) def test_memory_is_a_sliding_window(self): counter = iter(range(10)) out = pt.reflexion(lambda mem: pt.AgentResult("x", "finished", []), lambda res: (False, "no"), lambda res, fb: "L%d" % next(counter), max_trials=5, memory_size=2) self.assertEqual(out.memory, ["L3", "L4"]) self.assertEqual(len(out.trials[-1].memory_used), 2) def test_reflect_gets_the_trajectory_and_the_feedback(self): got = [] pt.reflexion(lambda mem: pt.AgentResult("a", "finished", [pt.Step("t", "x", "y", "z")]), lambda res: (False, "FB"), lambda res, fb: got.append((res.answer, fb)) or "L", max_trials=1) self.assertEqual(got, [("a", "FB")]) def test_lessons_reach_the_actors_prompt(self): r = self.run_it() self.assertTrue(r.success) lm = pt.age_policy("Ridge Corp", buggy=True) pt.react(lm, self.q, pt.TOOLS, memory=r.memory) self.assertIn("wrong order", lm.prompts[0]) class TestCodeSession(unittest.TestCase): def test_print_is_captured(self): self.assertEqual(pt.CodeSession().run("print(1 + 1)"), pt.Obs(True, "2")) def test_variables_persist_between_runs(self): s = pt.CodeSession() s.run("x = 41") self.assertEqual(s.run("print(x + 1)").text, "42") def test_tools_are_functions_in_the_namespace(self): s = pt.CodeSession({"double": lambda n: n * 2}) self.assertEqual(s.run("print(double(4))").text, "8") def test_errors_become_text_and_the_session_survives(self): s = pt.CodeSession() o = s.run("y = 5\nprint(1 / 0)") self.assertEqual((o.ok, o.text), (False, "ZeroDivisionError: division by zero")) self.assertEqual(s.run("print(y)").text, "5") def test_partial_output_is_kept_with_the_error(self): o = pt.CodeSession().run("print('before')\nraise ValueError('bad')") self.assertEqual(o.text, "before\nValueError: bad") def test_syntax_error(self): o = pt.CodeSession().run("def (:") self.assertFalse(o.ok) self.assertTrue(o.text.startswith("SyntaxError")) def test_imports_and_files_are_not_available(self): s = pt.CodeSession() for code in ("import os", "open('/etc/passwd')", "eval('1')", "exec('1')", "__import__('os')"): o = s.run(code) self.assertFalse(o.ok, code) def test_runaway_loop_is_stopped(self): s = pt.CodeSession(max_lines=200) o = s.run("while True:\n pass") self.assertEqual((o.ok, o.text), (False, "StepLimit: more than 200 lines executed")) self.assertEqual(s.run("print('alive')").text, "alive") def test_line_budget_counts_inside_functions_too(self): o = pt.CodeSession(max_lines=100).run("def f():\n while True:\n pass\nf()") self.assertIn("StepLimit", o.text) def test_tracing_is_restored_after_a_run(self): before = sys.gettrace() pt.CodeSession().run("print(1)") pt.CodeSession(max_lines=5).run("while True:\n pass") self.assertIs(sys.gettrace(), before) def test_long_output_is_truncated(self): o = pt.CodeSession(max_output=50).run("print('x' * 500)") self.assertIn("[truncated", o.text) self.assertLess(len(o.text), 100) def test_this_is_not_a_security_sandbox(self): # The docstring says so, and this documents why: attribute walks still work. o = pt.CodeSession().run("print(().__class__.__name__)") self.assertEqual(o.text, "tuple") class TestCodeActAndJsonActions(unittest.TestCase): def test_codeact_recovers_from_a_tool_error(self): r = pt.codeact(pt.RuleLM(pt.codeact_phone_policy), pt.PHONE_TASK, pt.phone_tools()) self.assertEqual((r.answer, r.status, len(r.turns)), ("India", "finished", 3)) self.assertEqual(r.turns[0].obs.text, "ValueError: unknown country 'Deutschland'") self.assertTrue(r.turns[1].obs.ok) def test_json_actions_need_one_turn_per_tool_call(self): r = pt.json_actions(pt.RuleLM(pt.json_phone_policy), pt.PHONE_TASK, pt.phone_tools()) self.assertEqual((r.answer, len(r.turns)), ("India", 9)) def test_codeact_uses_fewer_turns_and_tokens_on_this_task(self): mj, mc = pt.Meter(), pt.Meter() pt.json_actions(mj.wrap(pt.RuleLM(pt.json_phone_policy)), pt.PHONE_TASK, pt.phone_tools()) pt.codeact(mc.wrap(pt.RuleLM(pt.codeact_phone_policy)), pt.PHONE_TASK, pt.phone_tools()) self.assertEqual((mj.calls, mc.calls), (9, 3)) self.assertLess(mc.tokens_in * 3, mj.tokens_in) def test_the_phone_prices_hand_check(self): # my own arithmetic: 799*1 = 799; 118000*0.0068 = 802.4; 749*1.08 = 808.92; 66000*0.012 = 792 t = pt.phone_tools() usd = {c: t["get_price"]("X-1", c) * t["get_rate"](c) for c in pt.COUNTRIES} self.assertEqual(min(usd, key=usd.get), "India") self.assertAlmostEqual(usd["Japan"], 802.4, places=6) self.assertAlmostEqual(usd["Germany"], 808.92, places=6) def test_final_without_code_ends_the_run(self): r = pt.codeact(pt.ScriptedLM(["FINAL: 42"]), "q", {}) self.assertEqual((r.answer, r.status), ("42", "finished")) def test_reply_without_code_or_final_is_a_format_error(self): lm = pt.ScriptedLM(["hmm", "FINAL: ok"]) r = pt.codeact(lm, "q", {}) self.assertIn("format error", r.turns[0].obs.text) self.assertEqual(r.answer, "ok") def test_turn_limit(self): r = pt.codeact(pt.RuleLM(lambda p: "```python\nprint(1)\n```"), "q", {}, max_turns=3) self.assertEqual((r.status, r.answer, len(r.turns)), ("turn_limit", None, 3)) def test_code_wins_if_a_reply_has_both_code_and_final(self): lm = pt.ScriptedLM(["```python\nprint(5)\n```\nFINAL: guess", "FINAL: 5"]) r = pt.codeact(lm, "q", {}) self.assertEqual(r.turns[0].obs.text, "5") def test_json_bad_reply_is_an_observation(self): r = pt.json_actions(pt.ScriptedLM(["not json", '{"final": 1}']), "q", {}) self.assertIn("JSONDecodeError", r.turns[0].obs.text) def test_json_unknown_tool_is_an_observation(self): r = pt.json_actions(pt.ScriptedLM(['{"tool": "nope", "args": {}}', '{"final": 1}']), "q", {}) self.assertIn("KeyError", r.turns[0].obs.text) PLAN = ("Plan: find the founder, then the numbers.\n#E1 = search[who founded Harbor Labs]\n" "#E2 = search[birth year of #E1]\n#E3 = search[year Harbor Labs was founded]\n#E4 = calc[#E3 - #E2]") class TestPlanParsing(unittest.TestCase): def test_parses_steps_and_ignores_other_lines(self): plan = pt.parse_plan(PLAN, ["search", "calc"]) self.assertEqual([s.id for s in plan], ["E1", "E2", "E3", "E4"]) self.assertEqual(plan[1], pt.PlanStep("E2", "search", "birth year of #E1")) def test_unknown_tool_is_rejected(self): with self.assertRaises(pt.PlanParseError): pt.parse_plan("#E1 = shell[rm -rf /]", ["search"]) def test_duplicate_id_is_rejected(self): with self.assertRaises(pt.PlanParseError): pt.parse_plan("#E1 = calc[1]\n#E1 = calc[2]", ["calc"]) def test_reference_to_a_later_or_missing_step_is_rejected(self): for bad in ("#E1 = calc[#E2]\n#E2 = calc[1]", "#E1 = calc[#E9]", "#E1 = calc[#E1 + 1]"): with self.assertRaises(pt.PlanParseError, msg=bad): pt.parse_plan(bad, ["calc"]) def test_empty_plan_is_rejected(self): with self.assertRaises(pt.PlanParseError): pt.parse_plan("I will think about it.", ["calc"]) def test_waves_group_independent_steps(self): plan = pt.parse_plan(PLAN, ["search", "calc"]) self.assertEqual([[s.id for s in w] for w in pt.waves(plan)], [["E1", "E3"], ["E2"], ["E4"]]) def test_waves_of_a_straight_line_and_of_independent_steps(self): line = pt.parse_plan("#E1 = calc[1]\n#E2 = calc[#E1]\n#E3 = calc[#E2]", ["calc"]) flat = pt.parse_plan("#E1 = calc[1]\n#E2 = calc[2]\n#E3 = calc[3]", ["calc"]) self.assertEqual([len(w) for w in pt.waves(line)], [1, 1, 1]) self.assertEqual([len(w) for w in pt.waves(flat)], [3]) class TestExecutePlan(unittest.TestCase): def test_runs_the_plan_without_any_model(self): ev = pt.execute_plan(pt.parse_plan(PLAN, ["search", "calc"]), pt.TOOLS) self.assertEqual(ev, {"E1": "Mira Okoye", "E2": "1971", "E3": "2009", "E4": "38"}) def test_parallel_and_sequential_give_the_same_evidence(self): plan = pt.parse_plan(PLAN, ["search", "calc"]) self.assertEqual(pt.execute_plan(plan, pt.TOOLS, parallel=True), pt.execute_plan(plan, pt.TOOLS)) def test_independent_steps_really_run_together_in_parallel_mode(self): barrier = threading.Barrier(2, timeout=3) def meet(arg): barrier.wait() return arg plan = pt.parse_plan("#E1 = meet[a]\n#E2 = meet[b]", ["meet"]) self.assertEqual(pt.execute_plan(plan, {"meet": meet}, parallel=True), {"E1": "a", "E2": "b"}) def test_a_failing_tool_becomes_evidence(self): def boom(arg): raise KeyError("x") ev = pt.execute_plan(pt.parse_plan("#E1 = boom[a]\n#E2 = calc[1 + 1]", ["boom", "calc"]), {"boom": boom, "calc": pt.calc}) self.assertTrue(ev["E1"].startswith("error: KeyError")) self.assertEqual(ev["E2"], "2") def test_references_are_replaced_by_evidence_text(self): seen = [] pt.execute_plan(pt.parse_plan("#E1 = t[a]\n#E2 = t[x #E1 y]", ["t"]), {"t": lambda a: seen.append(a) or "R"}) self.assertEqual(seen, ["a", "x R y"]) class TestRewoo(unittest.TestCase): def solver(self): return pt.RuleLM(lambda p: "Answer: " + p.split("#E4: ")[1].splitlines()[0]) def test_two_model_calls_however_many_tools_run(self): m = pt.Meter() r = pt.rewoo(m.wrap(pt.ScriptedLM([PLAN])), m.wrap(self.solver()), pt.QUESTION, pt.TOOLS) self.assertEqual((r.answer, m.calls, len(r.plan), r.waves), ("Answer: 38", 2, 4, 3)) def test_the_solver_sees_the_plan_and_the_evidence(self): solver = self.solver() pt.rewoo(pt.ScriptedLM([PLAN]), solver, pt.QUESTION, pt.TOOLS) self.assertIn("#E2 = search[birth year of #E1]", solver.prompts[0]) self.assertIn("#E1: Mira Okoye", solver.prompts[0]) def test_uses_fewer_prompt_tokens_than_react_on_the_same_question(self): mr, ma = pt.Meter(), pt.Meter() pt.rewoo(mr.wrap(pt.ScriptedLM([PLAN])), mr.wrap(self.solver()), pt.QUESTION, pt.TOOLS) pt.react(ma.wrap(pt.age_policy("Harbor Labs")), pt.QUESTION, pt.TOOLS) self.assertEqual((mr.calls, ma.calls), (2, 5)) self.assertLess(mr.tokens_in * 3, ma.tokens_in) def test_a_bad_plan_raises_before_any_tool_runs(self): ran = [] with self.assertRaises(pt.PlanParseError): pt.rewoo(pt.ScriptedLM(["#E1 = nope[x]"]), pt.ScriptedLM([]), "q", {"search": lambda a: ran.append(a)}) self.assertEqual(ran, []) def test_the_plan_cannot_adapt_to_a_surprise(self): # Honest limit: if E1 returns 'no result', the plan still runs E2 with that text. plan = "#E1 = search[who founded Nowhere Inc]\n#E2 = search[birth year of #E1]" ev = pt.execute_plan(pt.parse_plan(plan, ["search"]), pt.TOOLS) self.assertEqual(ev, {"E1": "no result", "E2": "no result"}) class TestSession(unittest.TestCase): def make(self): s = pt.Session() s.append("system", "sys") s.append("user", "task") s.append("assistant", "run: ls") s.append("observation", "a.py b.py") return s def test_append_links_events_into_a_chain(self): s = self.make() self.assertEqual([(e.id, e.parent) for e in s.events], [(0, None), (1, 0), (2, 1), (3, 2)]) self.assertEqual(s.leaf, 3) def test_path_returns_root_to_leaf(self): s = self.make() self.assertEqual([e.id for e in s.path()], [0, 1, 2, 3]) self.assertEqual([e.id for e in s.path(1)], [0, 1]) def test_a_branch_is_a_second_child(self): s = self.make() s.append("assistant", "alt", parent=1) self.assertEqual([e.id for e in s.path()], [0, 1, 4]) self.assertEqual([e.id for e in s.path(3)], [0, 1, 2, 3]) self.assertEqual(len(s.events), 5) # nothing was removed def test_bad_parent_and_bad_kind_are_rejected(self): s = self.make() with self.assertRaises(ValueError): s.append("assistant", "x", parent=99) with self.assertRaises(ValueError): s.append("thought", "x") def test_events_are_immutable(self): e = self.make().get(1) with self.assertRaises(Exception): e.text = "changed" def test_observers_see_every_event(self): seen = [] s = pt.Session() s.observers.append(lambda e: seen.append(e.kind)) s.append("user", "a") s.append("assistant", "b") self.assertEqual(seen, ["user", "assistant"]) def test_jsonl_round_trip(self): s = self.make() s.append("compaction", "sum", first_kept=2) s2 = pt.Session.from_jsonl(s.to_jsonl()) self.assertEqual(s2.events, s.events) self.assertEqual(s2.leaf, s.leaf) def test_jsonl_with_out_of_order_ids_is_rejected(self): bad = '{"first_kept": null, "id": 1, "kind": "user", "parent": null, "text": "x"}' with self.assertRaises(ValueError): pt.Session.from_jsonl(bad) def chat(n_turns, obs_tokens=50, system="sys"): """A session with a system event, a user event, and n_turns of (assistant, observation).""" s = pt.Session() s.append("system", system) s.append("user", "Fix the bug in parser.py") for t in range(n_turns): s.append("assistant", "run: cat src/mod%d.py" % t) s.append("observation", toks(obs_tokens) + " t%d" % t) return s class TestBuildContext(unittest.TestCase): def test_roles_follow_event_kinds(self): msgs = pt.build_context(chat(1)) self.assertEqual([m.role for m in msgs], ["system", "user", "assistant", "tool"]) self.assertEqual([m.src for m in msgs], [0, 1, 2, 3]) def test_without_masking_everything_is_shown(self): s = chat(5) self.assertEqual(len(pt.build_context(s)), 12) def test_masking_keeps_the_last_m_observations_in_full(self): s = chat(5) msgs = pt.build_context(s, mask_window=2) tool = [m for m in msgs if m.role == "tool"] self.assertEqual([m.text.startswith("[output omitted") for m in tool], [True, True, True, False, False]) def test_masking_never_touches_actions_or_the_task(self): s = chat(5) full = [m for m in pt.build_context(s) if m.role != "tool"] masked = [m for m in pt.build_context(s, mask_window=1) if m.role != "tool"] self.assertEqual(full, masked) def test_the_placeholder_names_the_size_and_the_recall_id(self): s = chat(3, obs_tokens=50) ph = pt.build_context(s, mask_window=1)[3].text self.assertEqual(ph, "[output omitted: 51 tokens; recall id 3]") self.assertEqual(pt.recall(s, 3), s.get(3).text) def test_short_observations_are_not_replaced_by_a_longer_placeholder(self): s = chat(3, obs_tokens=2) self.assertEqual(pt.build_context(s, mask_window=1), pt.build_context(s)) def test_mask_window_zero_masks_every_observation(self): msgs = pt.build_context(chat(2), mask_window=0) self.assertTrue(all(m.text.startswith("[output omitted") for m in msgs if m.role == "tool")) def test_window_larger_than_history_masks_nothing(self): s = chat(3) self.assertEqual(pt.build_context(s, mask_window=10), pt.build_context(s)) def test_stride_moves_the_masked_boundary_in_steps(self): counts = [] for turns in range(1, 16): msgs = pt.build_context(chat(turns), mask_window=3, mask_stride=4) counts.append(sum(m.text.startswith("[output omitted") for m in msgs)) self.assertEqual(counts, [0, 0, 0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12]) def test_stride_one_is_the_plain_rolling_window(self): s = chat(9) self.assertEqual(pt.build_context(s, mask_window=3, mask_stride=1), pt.build_context(s, mask_window=3)) def test_bad_stride_is_rejected(self): with self.assertRaises(ValueError): pt.build_context(chat(1), mask_window=1, mask_stride=0) def test_building_a_context_never_changes_the_log(self): s = chat(6) before = s.to_jsonl() pt.build_context(s, mask_window=1) pt.compact(s, keep_recent_tokens=100) self.assertTrue(s.to_jsonl().startswith(before)) def test_context_follows_the_chosen_branch(self): s = chat(2) s.append("assistant", "alt action", parent=1) self.assertEqual(pt.build_context(s)[-1].text, "alt action") self.assertEqual(pt.build_context(s, leaf=5)[-1].role, "tool") def test_context_contains_and_tokens(self): msgs = pt.build_context(chat(1, obs_tokens=10)) self.assertTrue(pt.context_contains(msgs, "w9")) self.assertEqual(pt.context_tokens(msgs), 1 + 7 + 8 + 11) class TestCompaction(unittest.TestCase): def test_should_compact_threshold(self): self.assertFalse(pt.should_compact(3000, 4000, 1000)) self.assertTrue(pt.should_compact(3001, 4000, 1000)) def test_compaction_shrinks_the_context_and_keeps_the_log(self): s = chat(10, obs_tokens=100) n_events = len(s.events) before = pt.context_tokens(pt.build_context(s)) c = pt.compact(s, keep_recent_tokens=300) after = pt.build_context(s) self.assertLess(pt.context_tokens(after), before / 2) self.assertEqual(len(s.events), n_events + 1) self.assertEqual(c.kind, "compaction") def test_the_view_is_system_then_summary_then_the_kept_tail(self): s = chat(10, obs_tokens=100) c = pt.compact(s, keep_recent_tokens=300) ids = [m.src for m in pt.build_context(s)] self.assertEqual(ids[:2], [0, c.id]) self.assertEqual(ids[2], c.first_kept) self.assertEqual(ids[2:], list(range(c.first_kept, c.id))) def test_the_cut_never_falls_on_an_observation(self): for keep in range(50, 700, 37): s = chat(10, obs_tokens=100) c = pt.compact(s, keep_recent_tokens=keep) if c: self.assertNotEqual(s.get(c.first_kept).kind, "observation", keep) def test_kept_tail_has_at_least_the_requested_tokens_or_ends_at_a_boundary(self): s = chat(10, obs_tokens=100) c = pt.compact(s, keep_recent_tokens=300) kept = pt.context_tokens([m for m in pt.build_context(s) if m.src >= c.first_kept and m.src != c.id]) self.assertGreaterEqual(kept, 100) def test_nothing_to_compact_returns_none(self): self.assertIsNone(pt.compact(chat(1, obs_tokens=10), keep_recent_tokens=10000)) def test_summary_has_goal_actions_and_files(self): s = chat(6) s.append("assistant", "sed -i 's/a/b/' src/fix.py") s.append("observation", "ok") c = pt.compact(s, keep_recent_tokens=120) self.assertIn("## Goal\nFix the bug in parser.py", c.text) self.assertIn("run: cat src/mod0.py", c.text) self.assertIn("read: ", c.text) def test_errors_survive_in_the_summary_unless_dropped(self): s = pt.Session() s.append("system", "sys") s.append("user", "goal") s.append("assistant", "run: pytest") s.append("observation", "FAILED\nError: bad token at line 3") for t in range(6): s.append("assistant", "run: ls %d" % t) s.append("observation", toks(60)) c = pt.compact(s, keep_recent_tokens=100) self.assertIn("Error: bad token at line 3", c.text) lost = pt.extractive_summary([pt.Message("tool", "Error: x", 0)], keep_errors=False) self.assertNotIn("Critical", lost) def test_summary_respects_its_token_budget(self): msgs = [pt.Message("user", "goal text", 0)] + [pt.Message("assistant", "run: cat f%d.py" % i, i) for i in range(80)] self.assertLessEqual(pt.count_tokens(pt.extractive_summary(msgs, max_tokens=60)), 60) def test_a_second_compaction_carries_files_and_errors_forward(self): s = pt.Session() s.append("system", "sys") s.append("user", "goal") s.append("assistant", "run: cat first.py") s.append("observation", "Error: boom in first.py") for t in range(8): s.append("assistant", "run: cat a%d.py" % t) s.append("observation", toks(80)) c1 = pt.compact(s, keep_recent_tokens=200) for t in range(8): s.append("assistant", "run: cat b%d.py" % t) s.append("observation", toks(80)) c2 = pt.compact(s, keep_recent_tokens=200) self.assertGreater(c2.id, c1.id) self.assertIn("first.py", c2.text) self.assertIn("Error: boom", c2.text) ids = [m.src for m in pt.build_context(s)] self.assertNotIn(c1.id, ids) # only the newest summary is shown self.assertEqual(ids[:2], [0, c2.id]) def test_compaction_on_a_branch_uses_that_branchs_path(self): s = chat(8, obs_tokens=100) s.leaf = 5 for t in range(6): s.append("assistant", "alt %d" % t) s.append("observation", toks(100)) c = pt.compact(s, keep_recent_tokens=250) self.assertIn(c.first_kept, [e.id for e in s.path()]) class TestBranchesAndIsolation(unittest.TestCase): def test_switch_branch_summarizes_only_the_abandoned_part(self): s = chat(3) # ids 0..7 s.append("assistant", "run: git log", parent=3) s.append("observation", "abc fix") b = pt.switch_branch(s, 7) self.assertEqual(b.kind, "branch_summary") self.assertEqual(b.parent, 7) self.assertIn("git log", b.text) self.assertNotIn("mod2", b.text) # shared history is not repeated self.assertEqual(s.leaf, b.id) def test_both_branches_stay_in_the_log(self): s = chat(2) s.append("assistant", "side", parent=3) n = len(s.events) pt.switch_branch(s, 5) self.assertEqual(len(s.events), n + 1) self.assertEqual(s.get(6).text, "side") def test_the_summary_is_part_of_the_new_branchs_context(self): s = chat(2) s.append("assistant", "run: git status", parent=3) b = pt.switch_branch(s, 5) msgs = pt.build_context(s) self.assertEqual(msgs[-1].src, b.id) self.assertIn("git status", msgs[-1].text) def test_switching_to_the_same_branch_leaves_nothing_behind(self): s = chat(2) b = pt.switch_branch(s, s.leaf) self.assertEqual(b.text, "(nothing left behind)") def test_isolation_gives_the_parent_only_the_report(self): parent = pt.Session() parent.append("system", "sys") parent.append("user", "task") def worker(child): for i in range(5): child.append("assistant", "run: cat f%d" % i) child.append("observation", toks(200)) return "the parser is in src/parse.py" r = pt.isolate(parent, "find the parser", worker) self.assertEqual([e.kind for e in parent.events], ["system", "user", "assistant", "observation"]) self.assertEqual(parent.get(3).text, "the parser is in src/parse.py") self.assertGreater(r.child_tokens, 1000) self.assertLess(r.parent_added_tokens, 20) self.assertEqual(len(r.child.events), 12) # the child's full record is kept for audit def test_the_child_starts_clean(self): parent = pt.Session() parent.append("user", "SECRET parent history") seen = [] pt.isolate(parent, "t", lambda child: seen.append(child.to_jsonl()) or "r") self.assertNotIn("SECRET", seen[0]) def expected_full_history(turns=40): """Independent arithmetic: tokens sent per call when nothing is removed, from the raw trajectory.""" system, traj = pt.make_trajectory(turns) base = pt.count_tokens(system) + pt.count_tokens("Fix the failing parser test in the repository.") used, run = [], base for action, obs in traj: used.append(run) run += pt.count_tokens(action) + pt.count_tokens(obs) return used class TestTrajectoryAndSimulation(unittest.TestCase): def test_trajectory_is_deterministic_and_has_the_needle_and_an_error(self): a, b = pt.make_trajectory(), pt.make_trajectory() self.assertEqual(a, b) self.assertIn(pt.NEEDLE, a[1][3][1]) self.assertIn("Error: test_parse failed", a[1][4][1]) self.assertEqual(len(a[1]), 40) def test_full_history_matches_hand_arithmetic(self): used = expected_full_history() r = pt.simulate(pt.Strategy("full")) self.assertEqual(r.input_tokens, sum(used)) self.assertEqual(r.peak_context, max(used)) self.assertEqual(r.calls, 40) def test_full_history_is_append_only_so_each_call_reuses_the_previous_input(self): used = expected_full_history() r = pt.simulate(pt.Strategy("full")) self.assertEqual(r.cached_tokens, sum(used[:-1])) self.assertAlmostEqual(r.cost, ((sum(used) - sum(used[:-1])) * 3.0 + sum(used[:-1]) * 0.30) / 1e6, places=12) def test_masking_matches_a_brute_force_count(self): system, traj = pt.make_trajectory() base = pt.count_tokens(system) + pt.count_tokens("Fix the failing parser test in the repository.") total = 0 for t in range(40): used = base for k, (a, o) in enumerate(traj[:t]): n = pt.count_tokens(o) recent = k >= t - 10 used += pt.count_tokens(a) + (n if recent or n <= 11 else 11) # a placeholder is 11 tokens total += used self.assertEqual(pt.simulate(pt.Strategy("m", mask_window=10)).input_tokens, total) def test_masking_saves_input_tokens_but_loses_the_needle_from_view(self): full, m = pt.simulate(pt.Strategy("f")), pt.simulate(pt.Strategy("m", mask_window=10)) self.assertLess(m.input_tokens, full.input_tokens * 0.6) self.assertTrue(full.needle_in_context) self.assertFalse(m.needle_in_context) def test_a_masked_needle_can_be_recalled_from_the_log(self): system, traj = pt.make_trajectory() s = pt.Session() s.append("system", system) s.append("user", "task") for a, o in traj: s.append("assistant", a) s.append("observation", o) view = pt.build_context(s, mask_window=10) self.assertFalse(pt.context_contains(view, pt.NEEDLE)) placeholder = next(m for m in view if m.text.startswith("[output omitted") and m.src == 9) self.assertIn("recall id 9", placeholder.text) self.assertIn(pt.NEEDLE, pt.recall(s, 9)) def test_toy_cache_model_a_rolling_mask_breaks_the_prefix_every_turn(self): full = pt.simulate(pt.Strategy("f")) roll = pt.simulate(pt.Strategy("m", mask_window=10)) step = pt.simulate(pt.Strategy("s", mask_window=10, mask_stride=10)) frac = lambda r: r.cached_tokens / r.input_tokens # noqa: E731 self.assertGreater(frac(full), 0.9) self.assertLess(frac(roll), 0.3) self.assertGreater(frac(step), 0.8) self.assertLess(step.cost, roll.cost) # same masking idea, cheaper under this price model self.assertLess(full.cost, roll.cost) # with 10x cheaper cache hits, an intact prefix wins here def test_compaction_keeps_the_context_under_the_window(self): r = pt.simulate(pt.Strategy("c", compact=True), window=4000, reserve=1000) full = pt.simulate(pt.Strategy("f")) self.assertGreater(r.compactions, 3) self.assertLessEqual(r.peak_context, 4000) self.assertLess(r.input_tokens, full.input_tokens * 0.2) self.assertFalse(r.needle_in_context) def test_isolation_cuts_what_the_parent_pays_and_loses_what_the_report_omits(self): full = pt.simulate(pt.Strategy("f")) iso = pt.simulate(pt.Strategy("i", isolate_first=12)) self.assertEqual(iso.calls, 40) # 12 in the child plus 28 in the parent self.assertLess(iso.input_tokens, full.input_tokens) self.assertLess(iso.peak_context, full.peak_context) self.assertFalse(iso.needle_in_context) def test_per_call_series_adds_up_and_matches_the_hand_count(self): r = pt.simulate(pt.Strategy("f")) self.assertEqual(r.per_call, expected_full_history()) self.assertEqual(sum(pt.simulate(pt.Strategy("m", mask_window=10)).per_call), pt.simulate(pt.Strategy("m", mask_window=10)).input_tokens) def test_cost_formula(self): r = pt.simulate(pt.Strategy("m", mask_window=10)) expect = ((r.input_tokens - r.cached_tokens) * pt.PRICE_IN + r.cached_tokens * pt.PRICE_CACHED) / 1e6 self.assertAlmostEqual(r.cost, expect, places=12) self.assertEqual(pt.PRICE_IN / pt.PRICE_CACHED, 10) def test_simulation_is_deterministic(self): s = pt.Strategy("x", mask_window=5, compact=True) self.assertEqual(pt.simulate(s), pt.simulate(s)) def test_fewer_turns_fewer_tokens(self): self.assertLess(pt.simulate(pt.Strategy("f"), turns=10).input_tokens, pt.simulate(pt.Strategy("f"), turns=20).input_tokens) class TestScriptedPolicies(unittest.TestCase): def test_buggy_actor_subtracts_in_the_wrong_order_until_it_sees_the_lesson(self): q = "How old was the founder of Ridge Corp when Ridge Corp was founded?" bad = pt.react(pt.age_policy("Ridge Corp", buggy=True), q, pt.TOOLS) good = pt.react(pt.age_policy("Ridge Corp", buggy=True), q, pt.TOOLS, memory=["mind the order"]) self.assertEqual((bad.answer, good.answer), ("-40", "40")) def test_phone_tools_reject_unknown_countries(self): with self.assertRaises(ValueError): pt.phone_tools()["get_rate"]("Deutschland") class TestDemoAndCli(unittest.TestCase): def run_main(self, argv): out, err = io.StringIO(), io.StringIO() with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err): code = pt.main(argv) return code, out.getvalue(), err.getvalue() def test_full_demo_runs_and_prints_seven_sections(self): code, out, _ = self.run_main([]) self.assertEqual(code, 0) for n in range(1, 8): self.assertIn("== %d. " % n, out) def test_demo_is_deterministic(self): self.assertEqual(self.run_main([])[1], self.run_main(["demo"])[1]) def test_one_section(self): code, out, _ = self.run_main(["demo", "3"]) self.assertEqual(code, 0) self.assertIn("== 3. Reflexion", out) self.assertNotIn("== 2. ", out) def test_demo_numbers_match_the_notes(self): out = self.run_main([])[1] self.assertIn("answer=38 status=finished model calls=5 prompt tokens=483", out) self.assertIn("JSON actions: answer=India model calls=9", out) self.assertIn("CodeAct: answer=India model calls=3", out) self.assertIn("ReWOO: Answer: 38 | model calls=2", out) def test_bad_usage_returns_2(self): for argv in (["bogus"], ["demo", "9"], ["demo", "x"], ["demo", "1", "2"]): code, _, err = self.run_main(argv) self.assertEqual(code, 2, argv) self.assertIn("usage", err) if __name__ == "__main__": unittest.main()