"""Tests for frameworks_toolkit.py. Run: python test_frameworks_toolkit.py -v Standard library only. No network and no API key.""" import contextlib import io import json import operator import os import shutil import tempfile import time import unittest from typing import Annotated, TypedDict, Literal import frameworks_toolkit as ft def scripted_lm(replies): """An LM that returns the given replies in order.""" class S(ft.LM): def __call__(self, prompt): self.calls += 1 self.history.append(prompt) return replies[(self.calls - 1) % len(replies)] return S() class TmpDirCase(unittest.TestCase): def setUp(self): self.dir = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, self.dir, True) ft.configure(None) def path(self, name): return os.path.join(self.dir, name) # ----------------------------------------------------------------------------- 1. DSPy-like class TestSignature(unittest.TestCase): def test_parse_simple(self): s = ft.Signature.parse("question -> answer") self.assertEqual([f.name for f in s.inputs], ["question"]) self.assertEqual([f.name for f in s.outputs], ["answer"]) self.assertEqual(str(s), "question -> answer") def test_parse_typed_and_multiple(self): s = ft.Signature.parse("context, question: str -> answer: int, ok: bool") self.assertEqual([(f.name, f.type) for f in s.inputs], [("context", "str"), ("question", "str")]) self.assertEqual([(f.name, f.type) for f in s.outputs], [("answer", "int"), ("ok", "bool")]) self.assertEqual(str(s), "context, question -> answer: int, ok: bool") def test_bad_signatures(self): for bad in ["question answer", "a -> b -> c", " -> b", "a -> ", "1a -> b", "a -> b: dict", "a, a -> b", "a -> a"]: with self.assertRaises(ft.SignatureError, msg=bad): ft.Signature.parse(bad) def test_prepend_output_keeps_original(self): s = ft.Signature.parse("q -> a") c = s.prepend_output("reasoning") self.assertEqual([f.name for f in c.outputs], ["reasoning", "a"]) self.assertEqual([f.name for f in s.outputs], ["a"]) def test_default_and_custom_instructions(self): self.assertIn("produce the fields a", ft.Signature.parse("q -> a").instructions) self.assertEqual(ft.Signature.parse("q -> a", "Be brief.").instructions, "Be brief.") class TestPromptAndParse(unittest.TestCase): def setUp(self): self.sig = ft.Signature.parse("q -> a: int") def test_render_has_demos_inputs_and_output_fields(self): p = ft.render_prompt(self.sig, [{"q": "1+1", "a": 2}], {"q": "2+2"}) self.assertIn("OUTPUT FIELDS: a", p) self.assertIn("--- demo\nq: 1+1\na: 2", p) self.assertTrue(p.endswith("--- input\nq: 2+2")) def test_parse_reply_typed_and_ignores_noise(self): self.assertEqual(ft.parse_reply(self.sig, "thinking...\na: 42\nextra: x"), {"a": 42}) def test_parse_reply_missing_field(self): with self.assertRaises(ft.ParseError): ft.parse_reply(self.sig, "the answer is 4") def test_coerce(self): self.assertEqual(ft._coerce("3.5", "float"), 3.5) self.assertIs(ft._coerce("Yes", "bool"), True) self.assertIs(ft._coerce("no", "bool"), False) for typ, bad in [("int", "x"), ("float", "x"), ("bool", "maybe")]: with self.assertRaises(ft.ParseError): ft._coerce(bad, typ) class TestExample(unittest.TestCase): def test_inputs_labels_and_attrs(self): e = ft.Example(q="x", a="y").with_inputs("q") self.assertEqual(e.inputs(), {"q": "x"}) self.assertEqual(e.labels(), {"a": "y"}) self.assertEqual(e.a, "y") with self.assertRaises(AttributeError): e.nope def test_with_inputs_returns_a_copy(self): e = ft.Example(q="x", a="y") self.assertEqual(e.inputs(), {}) self.assertEqual(e.with_inputs("q").inputs(), {"q": "x"}) self.assertEqual(e.inputs(), {}) class TestPredict(TmpDirCase): def test_predict_reads_typed_output(self): p = ft.Predict("q -> a: int", scripted_lm(["a: 7"])) out = p(q="x") self.assertEqual((out.q, out.a), ("x", 7)) def test_wrong_inputs_raise(self): p = ft.Predict("q -> a", scripted_lm(["a: 1"])) with self.assertRaises(TypeError): p(q="x", extra="y") with self.assertRaises(TypeError): p() def test_needs_an_lm(self): with self.assertRaises(RuntimeError): ft.Predict("q -> a")(q="x") def test_global_lm_and_own_lm(self): ft.configure(scripted_lm(["a: global"])) self.assertEqual(ft.Predict("q -> a")(q="x").a, "global") self.assertEqual(ft.Predict("q -> a", scripted_lm(["a: own"]))(q="x").a, "own") def test_demos_go_into_the_prompt(self): lm = scripted_lm(["a: 1"]) p = ft.Predict("q -> a", lm) p.demos = [{"q": "one", "a": "uno"}] p(q="two") self.assertIn("q: one\na: uno", lm.history[0]) def test_trace_records_predictor_calls(self): p = ft.Predict("q -> a", scripted_lm(["a: 1"])) ft._trace_stack.append([]) try: p(q="x") finally: trace = ft._trace_stack.pop() self.assertEqual(len(trace), 1) self.assertIs(trace[0][0], p) self.assertEqual(trace[0][2], {"a": "1"}) def test_chain_of_thought_adds_reasoning_first(self): lm = scripted_lm(["reasoning: because\na: 5"]) cot = ft.ChainOfThought("q -> a", lm) out = cot(q="x") self.assertEqual((out.reasoning, out.a), ("because", "5")) self.assertIn("OUTPUT FIELDS: reasoning, a", lm.history[0]) def test_named_predictors_found_in_nested_modules(self): t = ft.Triage() self.assertEqual([n for n, _ in t.named_predictors()], ["category", "urgency"]) c = ft.Classify() self.assertEqual([n for n, _ in c.named_predictors()], ["predict"]) self.assertEqual([n for n, _ in ft.Predict("q -> a").named_predictors()], ["self"]) class TestBestOfN(TmpDirCase): def test_keeps_best_reward_and_stops_early(self): lm = scripted_lm(["a: 1", "a: 9", "a: 5"]) b = ft.BestOfN(ft.Predict("q -> a: int", lm), 3, lambda kw, p: p.a, threshold=9) self.assertEqual(b(q="x").a, 9) self.assertEqual(lm.calls, 2) # stopped when the threshold was met def test_runs_all_n_when_threshold_is_not_met(self): lm = scripted_lm(["a: 1", "a: 2", "a: 3"]) b = ft.BestOfN(ft.Predict("q -> a: int", lm), 3, lambda kw, p: p.a, threshold=100) self.assertEqual(b(q="x").a, 3) self.assertEqual(lm.calls, 3) def test_skips_bad_replies_and_raises_if_all_bad(self): lm = scripted_lm(["junk", "a: 4"]) self.assertEqual(ft.BestOfN(ft.Predict("q -> a: int", lm), 2, lambda kw, p: p.a)(q="x").a, 4) with self.assertRaises(ft.ParseError): ft.BestOfN(ft.Predict("q -> a: int", scripted_lm(["junk"])), 2, lambda kw, p: 0)(q="x") def test_n_must_be_positive(self): with self.assertRaises(ValueError): ft.BestOfN(ft.Predict("q -> a"), 0, lambda kw, p: 0) class TestImitatorLM(TmpDirCase): def test_copies_the_closest_demo(self): lm = ft.make_lm(0, p_format_slip=0.0) p = ft.Predict("ticket -> category", lm) p.demos = [{"ticket": "refund invoice 1177", "category": "billing"}, {"ticket": "package tracking is stuck", "category": "shipping"}] self.assertEqual(p(ticket="the tracking of my package is late").category, "shipping") self.assertEqual(p(ticket="refund for invoice 9034").category, "billing") def test_same_seed_same_replies(self): a, b = ft.make_lm(5), ft.make_lm(5) prompt = ft.render_prompt(ft.Signature.parse("ticket -> category"), [], {"ticket": "hello"}) self.assertEqual([a(prompt) for _ in range(6)], [b(prompt) for _ in range(6)]) def test_slip_breaks_format_only_without_demos(self): lm = ft.make_lm(0, p_format_slip=1.0) p = ft.Predict("ticket -> category", lm) with self.assertRaises(ft.ParseError): p(ticket="hello there") p.demos = [{"ticket": "hello there", "category": "feedback"}] self.assertEqual(p(ticket="hello there").category, "feedback") def test_counts_calls_and_keeps_prompts(self): lm = ft.make_lm(0) p = ft.Predict("ticket -> category", lm) p.demos = [{"ticket": "x", "category": "login"}] p(ticket="a") p(ticket="b") self.assertEqual((lm.calls, len(lm.history)), (2, 2)) def test_overlap_ignores_stop_words(self): self.assertEqual(ft.overlap("the a of", "to is it"), 0.0) self.assertAlmostEqual(ft.overlap("refund invoice", "refund invoice"), 1.0) self.assertAlmostEqual(ft.overlap("refund invoice", "refund order"), 1 / 3) class TestTicketsAndEvaluate(TmpDirCase): def test_make_tickets_is_deterministic_and_balanced(self): a, b = ft.make_tickets(40, 3), ft.make_tickets(40, 3) self.assertEqual([e.toDict() for e in a], [e.toDict() for e in b]) counts = {c: sum(e.category == c for e in a) for c in ft.CATEGORIES} self.assertEqual(set(counts.values()), {10}) self.assertEqual(a[0].inputs().keys(), {"ticket"}) self.assertIn(a[0].urgency, ("low", "high")) def test_evaluate_scores_and_counts_calls(self): lm = scripted_lm(["category: billing"]) ft.configure(lm) dev = [ft.Example(ticket="x", category="billing").with_inputs("ticket"), ft.Example(ticket="y", category="login").with_inputs("ticket")] r = ft.evaluate(ft.Classify(), dev, ft.category_match) self.assertEqual((r.score, r.n, r.failures, r.lm_calls), (0.5, 2, 0, 2)) self.assertEqual([ok for ok, _ in r.rows], [True, False]) def test_exception_counts_as_a_miss(self): ft.configure(scripted_lm(["no format here"])) dev = ft.make_tickets(5, 1) r = ft.evaluate(ft.Classify(), dev, ft.category_match) self.assertEqual((r.score, r.failures), (0.0, 5)) self.assertEqual(r.rows[0][1], "ParseError") def test_empty_devset(self): ft.configure(scripted_lm(["category: x"])) self.assertEqual(ft.evaluate(ft.Classify(), [], ft.category_match).score, 0.0) class TestOptimizers(TmpDirCase): def setUp(self): super().setUp() self.train, self.dev = ft.make_tickets(40, 1), ft.make_tickets(40, 2) def test_labeled_few_shot_puts_k_demos_and_leaves_original(self): prog = ft.Classify() out = ft.LabeledFewShot(5, seed=1).compile(prog, self.train) self.assertEqual(len(out.predict.demos), 5) self.assertEqual(prog.predict.demos, []) self.assertEqual(set(out.predict.demos[0]), {"ticket", "category"}) def test_labeled_few_shot_fills_every_stage_of_a_program(self): out = ft.LabeledFewShot(3).compile(ft.Triage(), self.train) self.assertEqual({n: len(p.demos) for n, p in out.named_predictors()}, {"category": 3, "urgency": 3}) self.assertEqual(set(out.urgency.demos[0]), {"ticket", "category", "urgency"}) def test_bootstrap_with_always_false_metric_keeps_nothing(self): ft.configure(ft.make_lm(0)) out = ft.BootstrapFewShot(lambda e, p, trace=None: False, 4, 0, 2).compile(ft.Classify(), self.train) self.assertEqual(out.predict.demos, []) self.assertEqual(out._bootstrapped, 0) def test_bootstrap_with_always_true_metric_stops_at_max(self): ft.configure(ft.make_lm(0, p_format_slip=0.0)) out = ft.BootstrapFewShot(lambda e, p, trace=None: True, 4, 0, 1).compile(ft.Classify(), self.train) self.assertEqual(len(out.predict.demos), 4) self.assertEqual(out._bootstrapped, 4) def test_bootstrapped_demos_come_from_passing_runs_only(self): ft.configure(ft.make_lm(2)) out = ft.BootstrapFewShot(ft.category_match, 6, 0, 3).compile(ft.Classify(), self.train) gold = {e.ticket: e.category for e in self.train} self.assertGreater(len(out.predict.demos), 0) for d in out.predict.demos: self.assertEqual(d["category"], gold[d["ticket"]]) def test_bootstrap_traces_every_stage_of_a_two_stage_program(self): ft.configure(ft.make_lm(3)) out = ft.BootstrapFewShot(ft.triage_match, 5, 0, 4).compile(ft.Triage(), self.train) n = out._bootstrapped self.assertGreater(n, 0) self.assertEqual(len(out.category.demos), n) self.assertEqual(len(out.urgency.demos), n) self.assertEqual(set(out.urgency.demos[0]), {"ticket", "category", "urgency"}) def test_bootstrap_does_not_touch_the_student_or_its_lm(self): lm = ft.make_lm(0) prog = ft.Classify() prog.set_lm(lm) out = ft.BootstrapFewShot(ft.category_match, 4, 0, 2).compile(prog, self.train) self.assertEqual(prog.predict.demos, []) self.assertIs(out.predict.lm, lm) # the LM is shared, not copied self.assertGreater(lm.calls, 0) def test_more_rounds_never_bootstrap_fewer_examples(self): got = [] for rounds in (1, 3): ft.configure(ft.make_lm(4)) got.append(ft.BootstrapFewShot(ft.category_match, 8, 0, rounds, 1).compile(ft.Classify(), self.train)._bootstrapped) self.assertGreaterEqual(got[1], got[0]) def test_max_labeled_demos_are_added_after_bootstrapped_ones(self): ft.configure(ft.make_lm(1)) out = ft.BootstrapFewShot(ft.category_match, 3, 4, 2).compile(ft.Classify(), self.train) self.assertEqual(len(out.predict.demos), out._bootstrapped + 4) def test_a_teacher_can_be_given(self): good = ft.Classify() good.set_lm(scripted_lm(["category: billing"])) student = ft.Classify() student.set_lm(scripted_lm(["category: login"])) train = [ft.Example(ticket="t%d" % i, category="billing").with_inputs("ticket") for i in range(6)] out = ft.BootstrapFewShot(ft.category_match, 3, 0, 1).compile(student, train, teacher=good) self.assertEqual(len(out.predict.demos), 3) self.assertTrue(all(d["category"] == "billing" for d in out.predict.demos)) def test_optimizing_beats_no_demos_on_average(self): base, tuned = [], [] for s in range(5): ft.configure(ft.make_lm(s)) prog = ft.BootstrapFewShot(ft.category_match, 8, 8, 3, s).compile(ft.Classify(), self.train) ft.configure(ft.make_lm(100 + s)) tuned.append(ft.evaluate(prog, self.dev, ft.category_match).score) base.append(ft.evaluate(ft.Classify(), self.dev, ft.category_match).score) self.assertGreater(sum(tuned) / 5 - sum(base) / 5, 0.4) def test_random_search_returns_the_best_candidate(self): ft.configure(ft.make_lm(0)) rs = ft.RandomSearch(ft.category_match, 3, 6, 0, 2) best = rs.compile(ft.Classify(), self.train, self.dev) self.assertEqual(len(rs.scores), 5) # no demos, labeled, and 3 bootstraps i = rs.scores.index(max(rs.scores)) self.assertEqual(rs.scores[i], max(rs.scores)) self.assertTrue(all(isinstance(x, float) for x in rs.scores)) if i == 0: self.assertEqual(best.predict.demos, []) else: self.assertGreater(len(best.predict.demos), 0) def test_random_search_ties_pick_the_earliest_candidate(self): ft.configure(ft.make_lm(0, p_format_slip=0.0)) rs = ft.RandomSearch(lambda e, p, trace=None: True, 2, 3, 0, 1) best = rs.compile(ft.Classify(), self.train) self.assertEqual(best.predict.demos, []) # candidate 0 is the program with no demos # ----------------------------------------------------------------------------- 2. StateGraph class S1(TypedDict): x: int log: Annotated[list, operator.add] class S2(TypedDict): a: str note: str def build_linear(): b = ft.StateGraph(S1) b.add_node("inc", lambda s: {"x": s["x"] + 1, "log": ["inc"]}) b.add_node("dbl", lambda s: {"x": s["x"] * 2, "log": ["dbl"]}) b.add_edge(ft.START, "inc").add_edge("inc", "dbl").add_edge("dbl", ft.END) return b class TestGraphBasics(unittest.TestCase): def test_reducers_are_read_from_annotations(self): g = ft.StateGraph(S1) self.assertEqual(g.keys, ["x", "log"]) self.assertIn("log", g.reducers) self.assertNotIn("x", g.reducers) def test_linear_run(self): out = build_linear().compile().invoke({"x": 3, "log": []}) self.assertEqual(out, {"x": 8, "log": ["inc", "dbl"]}) def test_reducer_starts_from_an_empty_list(self): out = build_linear().compile().invoke({"x": 1}) self.assertEqual(out["log"], ["inc", "dbl"]) def test_node_returning_none_changes_nothing(self): b = ft.StateGraph(S2) b.add_node("n", lambda s: None).add_edge(ft.START, "n").add_edge("n", ft.END) self.assertEqual(b.compile().invoke({"a": "q"}), {"a": "q"}) def test_nodes_get_a_copy_of_the_state(self): def sneaky(s): s["log"].append("mutated") return {"x": 1} b = ft.StateGraph(S1) b.add_node("n", sneaky).add_edge(ft.START, "n").add_edge("n", ft.END) self.assertEqual(b.compile().invoke({"x": 0, "log": ["a"]})["log"], ["a"]) def test_bad_updates(self): for bad, exc in [({"nope": 1}, ft.InvalidUpdateError), ("text", ft.InvalidUpdateError)]: b = ft.StateGraph(S2) b.add_node("n", lambda s, bad=bad: bad).add_edge(ft.START, "n").add_edge("n", ft.END) with self.assertRaises(exc): b.compile().invoke({"a": "q"}) def test_compile_checks(self): with self.assertRaises(ValueError): ft.StateGraph(S2).compile() # no START edge b = ft.StateGraph(S2) b.add_node("n", lambda s: None).add_edge(ft.START, "n") with self.assertRaises(ValueError): b.compile() # n has no way out with self.assertRaises(ValueError): b.add_edge("n", "ghost") with self.assertRaises(ValueError): b.add_node("n", lambda s: None) # duplicate with self.assertRaises(ValueError): b.add_node(ft.END, lambda s: None) # reserved class TestGraphRouting(unittest.TestCase): def make(self, router, pmap=None): b = ft.StateGraph(S1) b.add_node("a", lambda s: {"log": ["a"]}).add_node("b", lambda s: {"log": ["b"]}).add_node("c", lambda s: {"log": ["c"]}) b.add_edge(ft.START, "a").add_conditional_edges("a", router, pmap) b.add_edge("b", ft.END).add_edge("c", ft.END) return b def test_conditional_edge_picks_a_branch(self): b = self.make(lambda s: "b" if s["x"] > 0 else "c") self.assertEqual(b.compile().invoke({"x": 1})["log"], ["a", "b"]) self.assertEqual(b.compile().invoke({"x": 0})["log"], ["a", "c"]) def test_path_map_translates_labels(self): b = self.make(lambda s: "yes", {"yes": "b", "no": "c"}) self.assertEqual(b.compile().invoke({"x": 0})["log"], ["a", "b"]) def test_router_can_return_end(self): b = self.make(lambda s: ft.END) self.assertEqual(b.compile().invoke({"x": 0})["log"], ["a"]) def test_router_to_unknown_node_fails(self): with self.assertRaises(ValueError): self.make(lambda s: "ghost").compile().invoke({"x": 0}) def test_fan_out_runs_in_one_super_step_and_merges(self): b = self.make(lambda s: ["c", "b"]) g = b.compile() out = g.invoke({"x": 0}) self.assertEqual(out["log"], ["a", "b", "c"]) # merged in sorted node order self.assertEqual(g.events, [("run", "a"), ("run", "b"), ("run", "c")]) def test_two_writes_to_a_plain_key_are_refused(self): b = ft.StateGraph(S2) b.add_node("p", lambda s: {"note": "p"}).add_node("q", lambda s: {"note": "q"}) b.add_edge(ft.START, "p").add_edge(ft.START, "q").add_edge("p", ft.END).add_edge("q", ft.END) with self.assertRaises(ft.InvalidUpdateError): b.compile().invoke({"a": "x"}) def test_a_loop_hits_the_recursion_limit(self): b = ft.StateGraph(S1) b.add_node("n", lambda s: {"x": s["x"] + 1}).add_edge(ft.START, "n") b.add_conditional_edges("n", lambda s: "n" if s["x"] < 100 else ft.END) with self.assertRaises(ft.GraphRecursionError): b.compile(recursion_limit=5).invoke({"x": 0}) self.assertEqual(b.compile().invoke({"x": 0}, {"recursion_limit": 200})["x"], 100) self.assertEqual(b.compile(recursion_limit=30).invoke({"x": 80})["x"], 100) def test_mermaid_shows_static_and_conditional_edges(self): def r(s) -> Literal["b", "c"]: return "b" text = self.make(r).compile().to_mermaid() self.assertIn("__start__ --> a", text) self.assertIn("a -.-> b", text) self.assertIn("a -.-> c", text) self.assertIn("b --> __end__", text) class S3(TypedDict): draft: str approved: bool msgs: Annotated[list, operator.add] def build_review(counter=None): """draft -> review (asks a human) -> send or discard.""" counter = counter if counter is not None else [] def draft(s): return {"draft": "refund 20 USD", "msgs": ["drafted"]} def review(s): counter.append("review started") # runs again on resume ok = ft.interrupt({"question": "send this?", "draft": s["draft"]}) return {"approved": bool(ok), "msgs": ["reviewed"]} b = ft.StateGraph(S3) b.add_node("draft", draft).add_node("review", review) b.add_node("send", lambda s: {"msgs": ["sent"]}).add_node("discard", lambda s: {"msgs": ["discarded"]}) b.add_edge(ft.START, "draft").add_edge("draft", "review") b.add_conditional_edges("review", lambda s: "send" if s["approved"] else "discard") b.add_edge("send", ft.END).add_edge("discard", ft.END) return b, counter CFG = {"configurable": {"thread_id": "t1"}} class TestInterruptResume(TmpDirCase): def test_pause_shows_the_question_and_saves_a_checkpoint(self): b, _ = build_review() g = b.compile(ft.InMemorySaver()) out = g.invoke({"draft": "", "approved": False}, CFG) self.assertEqual(out["__interrupt__"][0].value["question"], "send this?") self.assertEqual(out["__interrupt__"][0].node, "review") snap = g.get_state(CFG) self.assertEqual((snap.status, snap.next), ("interrupted", ["review"])) self.assertEqual(snap.values["msgs"], ["drafted"]) def test_resume_finishes_the_run_and_the_node_restarts(self): b, counter = build_review() g = b.compile(ft.InMemorySaver()) g.invoke({"draft": "", "approved": False}, CFG) out = g.invoke(ft.Command(resume=True), CFG) self.assertEqual(out["msgs"], ["drafted", "reviewed", "sent"]) self.assertEqual(counter, ["review started", "review started"]) # the node ran twice self.assertEqual(g.events.count(("run", "draft")), 1) # earlier nodes did not rerun self.assertEqual(g.get_state(CFG).status, "done") def test_the_human_can_say_no(self): b, _ = build_review() g = b.compile(ft.InMemorySaver()) g.invoke({"draft": "", "approved": False}, CFG) self.assertEqual(g.invoke(ft.Command(resume=False), CFG)["msgs"][-1], "discarded") def test_resume_needs_a_waiting_thread(self): b, _ = build_review() g = b.compile(ft.InMemorySaver()) with self.assertRaises(ValueError): g.invoke(ft.Command(resume=True), CFG) def test_new_input_while_waiting_is_refused(self): b, _ = build_review() g = b.compile(ft.InMemorySaver()) g.invoke({"draft": "", "approved": False}, CFG) with self.assertRaises(ValueError): g.invoke({"draft": "other"}, CFG) with self.assertRaises(ValueError): g.invoke(None, CFG) def test_interrupt_outside_a_graph_is_an_error(self): with self.assertRaises(RuntimeError): ft.interrupt("hello") def test_two_interrupts_in_one_node_resume_in_order(self): def node(s): a = ft.interrupt("first?") b = ft.interrupt("second?") return {"msgs": ["%s+%s" % (a, b)]} b = ft.StateGraph(S3) b.add_node("n", node).add_edge(ft.START, "n").add_edge("n", ft.END) g = b.compile(ft.InMemorySaver()) self.assertEqual(g.invoke({"draft": ""}, CFG)["__interrupt__"][0].value, "first?") self.assertEqual(g.invoke(ft.Command(resume="A"), CFG)["__interrupt__"][0].value, "second?") self.assertEqual(g.invoke(ft.Command(resume="B"), CFG)["msgs"], ["A+B"]) def test_interrupt_needs_a_checkpointer_to_resume(self): b, _ = build_review() g = b.compile() # no saver out = g.invoke({"draft": "", "approved": False}, CFG) self.assertIn("__interrupt__", out) with self.assertRaises(ValueError): g.invoke(ft.Command(resume=True), CFG) class TestCheckpoints(TmpDirCase): def test_one_checkpoint_per_super_step_plus_the_input(self): g = build_linear().compile(ft.InMemorySaver()) g.invoke({"x": 1}, CFG) hist = g.history(CFG) self.assertEqual(len(hist), 3) # input, after inc, after dbl (done) self.assertEqual([h["step"] for h in hist], [0, 1, 2]) self.assertEqual(hist[-1]["status"], "done") self.assertEqual(hist[1]["state"]["x"], 2) # time travel: the state after the first node def test_threads_are_separate(self): g = build_linear().compile(ft.InMemorySaver()) g.invoke({"x": 1}, {"configurable": {"thread_id": "a"}}) g.invoke({"x": 10}, {"configurable": {"thread_id": "b"}}) self.assertEqual(g.get_state({"configurable": {"thread_id": "a"}}).values["x"], 4) self.assertEqual(g.get_state({"configurable": {"thread_id": "b"}}).values["x"], 22) with self.assertRaises(KeyError): g.get_state({"configurable": {"thread_id": "zzz"}}) def test_a_second_input_on_the_same_thread_continues_the_state(self): g = build_linear().compile(ft.InMemorySaver()) g.invoke({"x": 1}, CFG) out = g.invoke({"x": 5}, CFG) # x is replaced, log is appended to self.assertEqual(out["log"], ["inc", "dbl", "inc", "dbl"]) def test_file_saver_survives_a_new_process(self): path = self.path("ckpt.jsonl") b, _ = build_review() b.compile(ft.FileSaver(path)).invoke({"draft": "", "approved": False}, CFG) b2, counter2 = build_review() # a "restart": new objects, same file g2 = b2.compile(ft.FileSaver(path)) self.assertEqual(g2.get_state(CFG).status, "interrupted") out = g2.invoke(ft.Command(resume=True), CFG) self.assertEqual(out["msgs"], ["drafted", "reviewed", "sent"]) with open(path) as f: self.assertTrue(all(json.loads(line) for line in f)) # every line is JSON def test_in_memory_saver_forgets_on_restart(self): b, _ = build_review() b.compile(ft.InMemorySaver()).invoke({"draft": "", "approved": False}, CFG) b2, _ = build_review() with self.assertRaises(ValueError): b2.compile(ft.InMemorySaver()).invoke(ft.Command(resume=True), CFG) def test_crash_in_a_node_then_continue_without_redoing_finished_work(self): calls = {"ok": 0, "flaky": 0} def ok(s): calls["ok"] += 1 return {"msgs": ["ok"]} def flaky(s): calls["flaky"] += 1 if calls["flaky"] == 1: raise RuntimeError("boom") return {"msgs": ["flaky"]} b = ft.StateGraph(S3) b.add_node("a_ok", ok).add_node("b_flaky", flaky).add_node("end", lambda s: {"msgs": ["end"]}) b.add_edge(ft.START, "a_ok").add_edge(ft.START, "b_flaky").add_edge("a_ok", "end").add_edge("b_flaky", "end").add_edge("end", ft.END) g = b.compile(ft.InMemorySaver()) with self.assertRaises(RuntimeError): g.invoke({"draft": ""}, CFG) out = g.invoke(None, CFG) # continue from the last checkpoint self.assertEqual(out["msgs"], ["ok", "flaky", "end"]) self.assertEqual(calls, {"ok": 1, "flaky": 2}) # "ok" finished before the crash and was not redone def test_continue_on_a_finished_thread_is_an_error(self): g = build_linear().compile(ft.InMemorySaver()) g.invoke({"x": 1}, CFG) with self.assertRaises(ValueError): g.invoke(None, CFG) # ----------------------------------------------------------------------------- 3. Harness def fence(cmd): return "I will run this.\n```bash\n%s\n```" % cmd class TestSplitAndSandbox(TmpDirCase): def test_split_commands(self): self.assertEqual(ft.split_commands("ls -l | wc -l"), [["ls", "-l"], ["wc", "-l"]]) self.assertEqual(ft.split_commands("echo 'a b' && cat x; ls"), [["echo", "a b"], ["cat", "x"], ["ls"]]) self.assertEqual(ft.split_commands("echo hi > out.txt"), [["echo", "hi", ">", "out.txt"]]) def test_split_refuses_what_it_cannot_judge(self): for bad in ["echo `id`", "echo $(id)", "cat < a.txt").returncode, 0) obs = sb.run("cat a.txt") self.assertEqual((obs.returncode, obs.output), (0, "hello\n")) self.assertTrue(os.path.exists(self.path("a.txt"))) def test_not_on_the_allowlist(self): for cmd in ["rm a.txt", "python3 -c 1", "curl http://example.com", "ls && rm x", "echo hi | sh"]: with self.assertRaises(ft.SandboxError, msg=cmd): ft.Sandbox(self.dir).run(cmd) def test_escape_attempts_are_refused(self): sb = ft.Sandbox(self.dir) for cmd in ["cat /etc/passwd", "cat ../x", "cat a/../../x", "ls ~", "echo $HOME", "echo hi > /tmp/x", "echo hi >> ../x", "cat ./.[.]/x", "ls ..", "cat ${HOME}", "echo $PATH", "cp x /tmp/y"]: with self.assertRaises(ft.SandboxError, msg=cmd): sb.run(cmd) def test_brackets_in_quoted_text_are_fine(self): sb = ft.Sandbox(self.dir) sb.run("printf '[server]\\nport = 80\\n' > c.ini") self.assertEqual(sb.run("cat c.ini").output, "[server]\nport = 80\n") def test_every_command_is_a_fresh_process(self): sb = ft.Sandbox(self.dir, allow=("cd", "pwd", "mkdir")) sb.run("mkdir sub") sb.run("cd sub") self.assertEqual(sb.run("pwd").output.strip(), os.path.realpath(self.dir)) def test_environment_is_minimal(self): out = ft.Sandbox(self.dir, allow=("env",)).run("env").output self.assertIn("PATH=", out) self.assertNotIn("HOME=", out) os.environ["SECRET_TOKEN_FOR_TEST"] = "x" self.addCleanup(os.environ.pop, "SECRET_TOKEN_FOR_TEST", None) self.assertNotIn("SECRET_TOKEN", ft.Sandbox(self.dir, allow=("env",)).run("env").output) def test_timeout_kills_a_blocked_command_quickly(self): os.mkfifo(self.path("pipe")) sb = ft.Sandbox(self.dir, timeout=0.5) t0 = time.monotonic() obs = sb.run("cat pipe") # blocks forever: nobody writes to the FIFO took = time.monotonic() - t0 self.assertEqual(obs.returncode, 124) self.assertIn("timed out", obs.output) self.assertLess(took, 3.0) self.assertGreaterEqual(took, 0.4) def test_long_output_is_cut_with_a_marker(self): with open(self.path("big.txt"), "w") as f: f.write("x" * 5000) obs = ft.Sandbox(self.dir, max_output=100).run("cat big.txt") self.assertEqual(obs.truncated, 4900) self.assertIn("[truncated 4900 characters]", obs.output) self.assertLess(len(obs.output), 160) def test_nonzero_exit_is_reported_not_raised(self): obs = ft.Sandbox(self.dir).run("cat missing.txt") self.assertNotEqual(obs.returncode, 0) self.assertIn("No such file", obs.output) class TestEventsAndAgent(TmpDirCase): def write_cfg(self, text="timout = 30\n"): with open(self.path("config.ini"), "w") as f: f.write(text) def test_messages_are_rebuilt_from_events(self): evs = [ft.Event(0, "system", "s"), ft.Event(1, "user", "u"), ft.Event(2, "assistant", "a"), ft.Event(3, "observation", "o"), ft.Event(4, "refusal", "r")] self.assertEqual([m["role"] for m in ft.rebuild_messages(evs)], ["system", "user", "assistant", "user", "user"]) def test_events_round_trip_through_jsonl(self): evs = [ft.Event(0, "system", "line1\nline2"), ft.Event(1, "user", "q\"uote")] self.assertEqual(ft.events_from_jsonl(ft.events_to_jsonl(evs)), evs) def test_scripted_model_replays_and_records(self): m = ft.ScriptedModel(["one", "two"]) self.assertEqual((m([{"role": "user", "content": "a"}]), m([])), ("one", "two")) self.assertEqual(len(m.seen[0]), 1) with self.assertRaises(RuntimeError): m([]) def test_agent_fixes_a_file_and_submits(self): self.write_cfg() model = ft.ScriptedModel([fence("cat config.ini"), fence("echo 'timeout = 30' > config.ini"), fence("echo TASK_COMPLETE")]) res = ft.BashAgent(model, ft.Sandbox(self.dir)).run("fix the typo") self.assertEqual((res.status, res.steps), ("submitted", 3)) self.assertEqual(open(self.path("config.ini")).read(), "timeout = 30\n") def test_history_is_linear_each_prompt_extends_the_last(self): self.write_cfg() model = ft.ScriptedModel([fence("cat config.ini"), fence("ls"), fence("echo TASK_COMPLETE")]) ft.BashAgent(model, ft.Sandbox(self.dir)).run("look around") for a, b in zip(model.seen, model.seen[1:]): self.assertEqual(b[:len(a)], a) self.assertEqual(len(b), len(a) + 2) # the model's reply and the observation def test_wrong_format_is_reported_back_and_the_loop_continues(self): model = ft.ScriptedModel(["no block here", fence("ls") + "\n" + fence("ls"), fence("echo TASK_COMPLETE")]) res = ft.BashAgent(model, ft.Sandbox(self.dir)).run("t") self.assertEqual(res.status, "submitted") self.assertEqual([e.kind for e in res.events].count("format_error"), 2) self.assertIn("found 0", [e.text for e in res.events if e.kind == "format_error"][0]) def test_refusals_are_shown_to_the_model(self): model = ft.ScriptedModel([fence("rm -rf x"), fence("echo TASK_COMPLETE")]) res = ft.BashAgent(model, ft.Sandbox(self.dir)).run("t") self.assertEqual(res.status, "submitted") self.assertIn("refused", [e.text for e in res.events if e.kind == "refusal"][0]) self.assertIn("refused", model.seen[1][-1]["content"]) def test_step_limit(self): model = ft.ScriptedModel([fence("echo %d" % i) for i in range(10)]) res = ft.BashAgent(model, ft.Sandbox(self.dir), step_limit=4).run("t") self.assertEqual((res.status, res.steps), ("step_limit", 4)) def test_stuck_detector(self): model = ft.ScriptedModel([fence("ls")] * 10) res = ft.BashAgent(model, ft.Sandbox(self.dir), stuck=3).run("t") self.assertEqual((res.status, res.steps), ("stuck", 3)) def test_different_commands_are_not_stuck(self): model = ft.ScriptedModel([fence("ls"), fence("pwd"), fence("ls"), fence("echo TASK_COMPLETE")]) self.assertEqual(ft.BashAgent(model, ft.Sandbox(self.dir), stuck=2).run("t").status, "submitted") def test_model_failure_ends_the_run_cleanly(self): res = ft.BashAgent(ft.ScriptedModel([]), ft.Sandbox(self.dir)).run("t") self.assertEqual(res.status, "model_error") self.assertEqual(res.events[-1].kind, "format_error") class TestSessionTreeAndTools(TmpDirCase): def test_branching(self): t = ft.SessionTree("sys") a = t.add("user", "task") b = t.add("assistant", "plan A") t.branch_from(a) c = t.add("assistant", "plan B") self.assertEqual([m["content"] for m in t.path()], ["sys", "task", "plan B"]) self.assertEqual([m["content"] for m in t.path(b)], ["sys", "task", "plan A"]) self.assertEqual(t.leaves(), [b, c]) with self.assertRaises(KeyError): t.branch_from(99) def test_four_tools(self): ft_ = ft.FileTools(self.dir) self.assertEqual(ft_.call("write", path="sub/a.txt", text="one two one"), "wrote 11 characters") self.assertEqual(ft_.call("read", path="sub/a.txt"), "one two one") self.assertEqual(ft_.call("edit", path="sub/a.txt", old="two", new="2"), "edited") self.assertEqual(ft_.call("read", path="sub/a.txt"), "one 2 one") self.assertEqual(ft_.call("bash", cmd="cat sub/a.txt").output, "one 2 one") def test_edit_needs_exactly_one_match(self): t = ft.FileTools(self.dir) t.write("a.txt", "x x y") for old in ("x", "zzz"): with self.assertRaises(ValueError): t.edit("a.txt", old, "q") self.assertEqual(t.read("a.txt"), "x x y") def test_paths_and_names_are_checked(self): t = ft.FileTools(self.dir) for rel in ("../x.txt", "/etc/passwd", "a/../../x"): with self.assertRaises(ft.SandboxError, msg=rel): t.read(rel) with self.assertRaises(KeyError): t.call("delete", path="x") os.symlink("/etc", self.path("link")) with self.assertRaises(ft.SandboxError): t.read("link/hostname") # ----------------------------------------------------------------------------- 4. Durable execution def make_runtime(store=None, flaky_fails=0, notes=None): """A runtime with: lookup (pure), flaky (fails N times), notify (side effect, logs to notes).""" rt = ft.Runtime(store) state = {"fails": flaky_fails} notes = notes if notes is not None else [] def lookup(x): return {"value": x * 2, "pair": (x, x)} def flaky(x): if state["fails"] > 0: state["fails"] -= 1 raise IOError("temporary failure") return x + 100 def notify(msg): notes.append((msg, ft.activity_info()["key"], ft.activity_info()["attempt"])) return "sent" def broken(x): raise ft.NonRetryable("bad input") rt.register_activity("lookup", lookup) rt.register_activity("flaky", flaky, ft.RetryPolicy(initial_interval=1, backoff=2, max_attempts=3)) rt.register_activity("notify", notify) rt.register_activity("broken", broken, ft.RetryPolicy(max_attempts=5)) return rt, notes def wf_simple(ctx, x): return ctx.activity("lookup", x)["value"] + 1 def wf_timer(ctx, x): a = ctx.activity("lookup", x) ctx.sleep(100) b = ctx.activity("notify", "done %d" % a["value"]) return [a["value"], b] class TestRetryPolicy(unittest.TestCase): def test_delays_double_and_are_capped(self): p = ft.RetryPolicy(initial_interval=1, backoff=2, max_interval=5) self.assertEqual([p.delay(i) for i in (1, 2, 3, 4)], [1, 2, 4, 5]) class TestDurable(unittest.TestCase): def test_simple_workflow_and_its_history(self): rt, _ = make_runtime() rt.register_workflow("w", wf_simple) st = rt.start("a", "w", 5) self.assertEqual((st.status, st.result), ("completed", 11)) self.assertEqual([e["type"] for e in rt.store.for_workflow("a")], ["WorkflowStarted", "ActivityScheduled", "ActivityCompleted", "WorkflowCompleted"]) self.assertEqual([e["id"] for e in rt.store.for_workflow("a")], [1, 2, 3, 4]) self.assertEqual(rt.status("a"), "completed") def test_results_go_through_json(self): rt, _ = make_runtime() rt.register_workflow("w", lambda ctx: ctx.activity("lookup", 2)) self.assertEqual(rt.start("a", "w").result, {"value": 4, "pair": [2, 2]}) # the tuple became a list def test_running_a_finished_workflow_runs_nothing_again(self): rt, _ = make_runtime() rt.register_workflow("w", wf_simple) rt.start("a", "w", 5) n = len(rt.executions) self.assertEqual(rt.run("a").result, 11) self.assertEqual(len(rt.executions), n) def test_start_errors(self): rt, _ = make_runtime() rt.register_workflow("w", wf_simple) rt.start("a", "w", 1) with self.assertRaises(ValueError): rt.start("a", "w", 1) with self.assertRaises(KeyError): rt.start("b", "nope") with self.assertRaises(KeyError): rt.status("zzz") def test_retry_then_success_records_each_failed_attempt(self): rt, _ = make_runtime(flaky_fails=2) rt.register_workflow("w", lambda ctx: ctx.activity("flaky", 1)) st = rt.start("a", "w") self.assertEqual((st.status, st.result), ("completed", 101)) types = [e["type"] for e in rt.store.for_workflow("a")] self.assertEqual(types.count("ActivityAttemptFailed"), 2) self.assertEqual(rt.store.for_workflow("a")[-2]["attempts"], 3) self.assertEqual(rt.now, 3.0) # waited 1 s then 2 s (virtual time) self.assertEqual(rt.executions, [("flaky", 1), ("flaky", 2), ("flaky", 3)]) def test_out_of_attempts_fails_the_workflow(self): rt, _ = make_runtime(flaky_fails=10) rt.register_workflow("w", lambda ctx: ctx.activity("flaky", 1)) st = rt.start("a", "w") self.assertEqual(st.status, "failed") self.assertIn("after 3 attempt", st.detail) self.assertEqual(rt.status("a"), "failed") self.assertEqual(rt.store.for_workflow("a")[-2]["type"], "ActivityFailed") self.assertEqual(rt.run("a").status, "failed") # stays failed on replay def test_non_retryable_stops_at_once(self): rt, _ = make_runtime() rt.register_workflow("w", lambda ctx: ctx.activity("broken", 1)) st = rt.start("a", "w") self.assertEqual(st.status, "failed") self.assertEqual(rt.executions, [("broken", 1)]) def test_policy_can_name_non_retryable_types(self): rt = ft.Runtime() def f(): raise KeyError("x") rt.register_activity("f", f, ft.RetryPolicy(max_attempts=5, non_retryable=(KeyError,))) rt.register_workflow("w", lambda ctx: ctx.activity("f")) self.assertEqual(rt.start("a", "w").status, "failed") self.assertEqual(len(rt.executions), 1) def test_timer_suspends_then_completes_without_redoing_activities(self): rt, notes = make_runtime() rt.register_workflow("w", wf_timer) self.assertEqual(rt.start("a", "w", 4).status, "suspended") self.assertEqual(rt.run("a").status, "suspended") # still not time rt.advance(99) self.assertEqual(rt.run("a").status, "suspended") rt.advance(1) st = rt.run("a") self.assertEqual((st.status, st.result), ("completed", [8, "sent"])) self.assertEqual([n for n, _ in rt.executions].count("lookup"), 1) # lookup ran once in total self.assertEqual(len([e for e in rt.store.for_workflow("a") if e["type"] == "TimerStarted"]), 1) def test_crash_after_a_recorded_result_does_not_repeat_the_activity(self): store = ft.HistoryStore() rt, notes = make_runtime(store) rt.register_workflow("w", wf_timer) rt.crash_at = (3, "after") # right after lookup's ActivityCompleted with self.assertRaises(ft.SimulatedCrash): rt.start("a", "w", 4) rt2, _ = make_runtime(store, notes=notes) # a new worker with the same history rt2.register_workflow("w", wf_timer) self.assertEqual(rt2.run("a").status, "suspended") # replay reached the timer rt2.advance(100) st = rt2.run("a") self.assertEqual((st.status, st.result), ("completed", [8, "sent"])) self.assertEqual([n for n, _ in rt.executions + rt2.executions].count("lookup"), 1) def test_crash_before_the_result_is_recorded_runs_the_activity_again(self): store = ft.HistoryStore() notes = [] rt, _ = make_runtime(store, notes=notes) rt.register_workflow("w", wf_timer) rt.start("a", "w", 4) rt.advance(100) rt.crash_at = (7, "before") # notify ran, but its ActivityCompleted was lost with self.assertRaises(ft.SimulatedCrash): rt.run("a") rt2, _ = make_runtime(store, notes=notes) rt2.register_workflow("w", wf_timer) rt2.now = rt.now self.assertEqual(rt2.run("a").status, "completed") self.assertEqual(len(notes), 2) # at-least-once: the side effect happened twice self.assertEqual(notes[0][1], notes[1][1]) # ...but both calls carried the same idempotency key self.assertEqual(notes[0][1], "a:2") def test_an_idempotent_effect_applied_once_despite_two_runs(self): store, applied = ft.HistoryStore(), set() rt = ft.Runtime(store) runs = [] def charge(amount): runs.append(amount) applied.add(ft.activity_info()["key"]) # the "payment provider" dedupes by key return "ok" rt.register_activity("charge", charge) rt.register_workflow("w", lambda ctx: ctx.activity("charge", 5)) rt.crash_at = (3, "before") with self.assertRaises(ft.SimulatedCrash): rt.start("a", "w") rt2 = ft.Runtime(store) rt2.register_activity("charge", charge) rt2.register_workflow("w", lambda ctx: ctx.activity("charge", 5)) rt2.run("a") self.assertEqual((len(runs), len(applied)), (2, 1)) def test_attempt_numbers_continue_after_a_crash_in_the_retry_loop(self): store = ft.HistoryStore() rt, _ = make_runtime(store, flaky_fails=1) rt.register_workflow("w", lambda ctx: ctx.activity("flaky", 1)) rt.crash_at = (3, "after") # after the first AttemptFailed with self.assertRaises(ft.SimulatedCrash): rt.start("a", "w") rt2, _ = make_runtime(store, flaky_fails=0) rt2.register_workflow("w", lambda ctx: ctx.activity("flaky", 1)) self.assertEqual(rt2.run("a").result, 101) self.assertEqual(rt2.executions, [("flaky", 2)]) def test_history_survives_in_a_file(self): path = os.path.join(tempfile.mkdtemp(), "h.jsonl") self.addCleanup(shutil.rmtree, os.path.dirname(path), True) rt, _ = make_runtime(ft.HistoryStore(path)) rt.register_workflow("w", wf_simple) rt.start("a", "w", 5) again = ft.HistoryStore(path) self.assertEqual(again.for_workflow("a"), rt.store.for_workflow("a")) self.assertEqual(again.append("a", {"type": "X", "t": 0})["id"], 5) class TestDeterminism(unittest.TestCase): def test_changing_the_order_of_commands_is_caught(self): store = ft.HistoryStore() rt, _ = make_runtime(store) rt.register_workflow("w", lambda ctx: (ctx.activity("lookup", 1), ctx.sleep(10), ctx.activity("lookup", 2))[2]) rt.start("a", "w") rt2, _ = make_runtime(store) rt2.register_workflow("w", lambda ctx: (ctx.activity("notify", "x"), ctx.sleep(10))[0]) # a different program with self.assertRaises(ft.NonDeterminismError): rt2.run("a") def test_changing_the_arguments_is_caught(self): store = ft.HistoryStore() rt, _ = make_runtime(store) rt.register_workflow("w", lambda ctx: (ctx.activity("lookup", 1), ctx.sleep(10))[0]) rt.start("a", "w") rt2, _ = make_runtime(store) rt2.register_workflow("w", lambda ctx: (ctx.activity("lookup", 2), ctx.sleep(10))[0]) with self.assertRaises(ft.NonDeterminismError): rt2.run("a") def test_random_in_workflow_code_breaks_replay_but_side_effect_does_not(self): import random store = ft.HistoryStore() rt, _ = make_runtime(store) picked = [] def bad(ctx): n = random.choice([1, 2, 3, 4, 5, 6, 7, 8]) # not deterministic picked.append(n) ctx.activity("lookup", n) ctx.sleep(10) return n rt.register_workflow("w", bad) random.seed(123) rt.start("a", "w") rt.advance(10) # pick a seed whose first choice differs from the one in the history other = next(s for s in range(100) if random.Random(s).choice([1, 2, 3, 4, 5, 6, 7, 8]) != picked[0]) random.seed(other) with self.assertRaises(ft.NonDeterminismError): rt.run("a") def good(ctx): n = ctx.side_effect(lambda: random.choice([1, 2, 3, 4, 5, 6, 7, 8])) ctx.activity("lookup", n) ctx.sleep(10) return n rt3, _ = make_runtime(ft.HistoryStore()) rt3.register_workflow("w", good) first = rt3.start("g", "w") rt3.advance(10) results = set() for seed in range(10): random.seed(seed) results.add(rt3.run("g").result) self.assertEqual(len(results), 1) self.assertEqual(first.status, "suspended") def test_now_is_recorded_once(self): rt, _ = make_runtime() def w(ctx): t = ctx.now() ctx.sleep(50) return t rt.register_workflow("w", w) rt.now = 7 rt.start("a", "w") rt.advance(50) self.assertEqual(rt.run("a").result, 7) class TestSignals(unittest.TestCase): def wf(self, ctx): a = ctx.wait_for_signal("approve") b = ctx.wait_for_signal("approve") return [a, b] def test_waits_until_a_signal_arrives(self): rt, _ = make_runtime() rt.register_workflow("w", self.wf) self.assertEqual(rt.start("a", "w").status, "suspended") rt.signal("a", "approve", "yes") self.assertEqual(rt.run("a").status, "suspended") # the second wait is still open rt.signal("a", "approve", "again") self.assertEqual(rt.run("a").result, ["yes", "again"]) def test_a_signal_that_arrives_early_is_kept(self): rt, _ = make_runtime() rt.register_workflow("w", self.wf) rt.register_workflow("v", lambda ctx: (ctx.sleep(5), ctx.wait_for_signal("go"))[1]) rt.start("a", "v") rt.signal("a", "go", {"n": 1}) rt.advance(5) self.assertEqual(rt.run("a").result, {"n": 1}) def test_other_signal_names_do_not_match_and_each_signal_is_used_once(self): rt, _ = make_runtime() rt.register_workflow("w", lambda ctx: ctx.wait_for_signal("go")) rt.start("a", "w") rt.signal("a", "stop", 1) self.assertEqual(rt.run("a").status, "suspended") rt.signal("a", "go", 2) self.assertEqual(rt.run("a").result, 2) self.assertEqual(rt.run("a").result, 2) # replay gives the same answer def test_format_history(self): rt, _ = make_runtime() rt.register_workflow("w", wf_simple) rt.start("a", "w", 1) text = ft.format_history(rt.store.for_workflow("a")) self.assertEqual(len(text.splitlines()), 4) self.assertIn("ActivityCompleted", text) # ----------------------------------------------------------------------------- 5. Tracing and statistics class TickClock: """Moves by 1 second on every read, so durations are easy to compute by hand.""" def __init__(self): self.t = 0.0 def __call__(self): self.t += 1.0 return self.t class TestCost(unittest.TestCase): def test_approx_tokens(self): self.assertEqual((ft.approx_tokens(""), ft.approx_tokens("abcd"), ft.approx_tokens("abcde")), (1, 1, 2)) def test_cost_by_hand(self): # 100 input tokens at 0.15 USD per million, plus 10 output tokens at 0.60 USD per million self.assertAlmostEqual(ft.cost_usd("mock-model", 100, 10), 100 * 0.15 / 1e6 + 10 * 0.60 / 1e6) self.assertAlmostEqual(ft.cost_usd("mock-model", 100, 10), 2.1e-5) self.assertEqual(ft.cost_usd("mock-model", 0, 0), 0.0) with self.assertRaises(KeyError): ft.cost_usd("unknown", 1, 1) class TestTracer(unittest.TestCase): def test_nesting_gives_parent_ids_and_one_trace_id(self): tr = ft.Tracer(TickClock()) with tr.span("root") as r: with tr.span("child"): pass with tr.span("child2"): pass by = {s.name: s for s in tr.spans} self.assertEqual([s.name for s in tr.spans], ["child", "child2", "root"]) # spans close inner first self.assertEqual(by["child"].parent_id, by["root"].span_id) self.assertIsNone(by["root"].parent_id) self.assertEqual(len({s.trace_id for s in tr.spans}), 1) def test_separate_roots_get_separate_trace_ids(self): tr = ft.Tracer(TickClock()) with tr.span("a"): pass with tr.span("b"): pass self.assertNotEqual(tr.spans[0].trace_id, tr.spans[1].trace_id) def test_ids_depend_on_the_seed_only(self): def ids(seed): tr = ft.Tracer(TickClock(), seed=seed) with tr.span("a"): with tr.span("b"): pass return [(s.trace_id, s.span_id) for s in tr.spans] self.assertEqual(ids(1), ids(1)) self.assertNotEqual(ids(1), ids(2)) self.assertEqual(len(ids(1)[0][0]), 32) self.assertEqual(len(ids(1)[0][1]), 16) def test_durations_with_a_fake_clock(self): tr = ft.Tracer(TickClock()) with tr.span("outer"): # start=1 with tr.span("inner"): # start=2, end=3 pass by = {s.name: s for s in tr.spans} # outer end=4 self.assertEqual((by["inner"].duration, by["outer"].duration), (1.0, 3.0)) def test_an_exception_marks_the_span_and_is_re_raised(self): tr = ft.Tracer(TickClock()) with self.assertRaises(ZeroDivisionError): with tr.span("outer"): with tr.span("bad"): 1 / 0 by = {s.name: s for s in tr.spans} self.assertEqual((by["bad"].status, by["bad"].attributes["error.type"]), ("ERROR", "ZeroDivisionError")) self.assertEqual(by["outer"].status, "ERROR") # the error passes through the parent too self.assertEqual(len(tr.spans), 2) # both spans were still recorded with tr.span("after"): # and the stack is clean pass self.assertIsNone(tr.spans[-1].parent_id) def test_content_is_off_unless_asked_for(self): for capture, expect in [(False, False), (True, True)]: tr = ft.Tracer(TickClock(), capture_content=capture) lm = ft.TracedLM(scripted_lm(["a: 1"]), tr, "mock-model") lm("secret prompt") self.assertEqual("gen_ai.input.messages" in tr.spans[0].attributes, expect) def test_traced_lm_attributes_and_cost(self): tr = ft.Tracer(TickClock()) lm = ft.TracedLM(scripted_lm(["x" * 40]), tr, "mock-model") self.assertEqual(lm("p" * 400), "x" * 40) sp = tr.spans[0] self.assertEqual(sp.name, "chat mock-model") a = sp.attributes self.assertEqual((a["gen_ai.operation.name"], a["gen_ai.request.model"]), ("chat", "mock-model")) self.assertEqual((a["gen_ai.usage.input_tokens"], a["gen_ai.usage.output_tokens"]), (100, 10)) self.assertAlmostEqual(a["app.cost_usd"], 2.1e-5) self.assertEqual(lm.calls, 1) def test_traced_tool(self): tr = ft.Tracer(TickClock(), capture_content=True) self.assertEqual(ft.traced_tool(tr, "add", lambda a, b: a + b, 2, 3, call_id="c1"), 5) sp = tr.spans[0] self.assertEqual(sp.name, "execute_tool add") self.assertEqual((sp.attributes["gen_ai.tool.name"], sp.attributes["gen_ai.tool.call.id"]), ("add", "c1")) self.assertEqual(sp.attributes["gen_ai.tool.call.arguments"], {"args": [2, 3], "kw": {}}) def test_traced_node_inside_a_graph(self): tr = ft.Tracer(TickClock()) b = ft.StateGraph(S1) b.add_node("inc", ft.traced_node(tr, "inc", lambda s: {"x": s["x"] + 1})) b.add_edge(ft.START, "inc").add_edge("inc", ft.END) with tr.span("invoke_workflow demo"): b.compile().invoke({"x": 0}) self.assertEqual([s.name for s in tr.spans], ["node inc", "invoke_workflow demo"]) self.assertEqual(tr.spans[0].parent_id, tr.spans[1].span_id) def test_agent_run_produces_an_agent_span_with_chat_and_tool_children(self): d = tempfile.mkdtemp() self.addCleanup(shutil.rmtree, d, True) tr = ft.Tracer(TickClock()) model = ft.ScriptedModel([fence("rm x"), fence("ls"), fence("echo TASK_COMPLETE")]) ft.BashAgent(model, ft.Sandbox(d), tracer=tr).run("t") names = [s.name for s in tr.spans] self.assertEqual(names.count("chat mock-model"), 3) self.assertEqual(names.count("execute_tool bash"), 2) root = tr.spans[-1] self.assertEqual((root.name, root.attributes["app.status"], root.attributes["app.steps"]), ("invoke_agent bash-agent", "submitted", 3)) self.assertTrue(all(s.parent_id == root.span_id for s in tr.spans[:-1])) refused = [s for s in tr.spans if s.attributes.get("app.refused")] self.assertEqual((len(refused), refused[0].status), (1, "ERROR")) class TestExportAndView(unittest.TestCase): def make(self): tr = ft.Tracer(TickClock(), capture_content=True) lm = ft.TracedLM(scripted_lm(["x" * 40, "y" * 8]), tr, "mock-model") with tr.span("invoke_agent demo", **{"gen_ai.operation.name": "invoke_agent"}): lm("p" * 400) lm("q" * 40) try: with tr.span("execute_tool boom", **{"gen_ai.operation.name": "execute_tool"}): raise ValueError("bad") except ValueError: pass return tr def test_jsonl_round_trip(self): tr = self.make() text = ft.export_jsonl(tr.spans) rows = [json.loads(l) for l in text.splitlines()] self.assertEqual(len(rows), 4) self.assertEqual(set(rows[0]), {"trace_id", "span_id", "parent_span_id", "name", "start_time_unix_nano", "end_time_unix_nano", "attributes", "status", "events"}) back = ft.load_jsonl(text) self.assertEqual([(s.name, s.span_id, s.parent_id, s.status, s.attributes) for s in back], [(s.name, s.span_id, s.parent_id, s.status, s.attributes) for s in tr.spans]) self.assertAlmostEqual(back[0].duration, tr.spans[0].duration) self.assertEqual(ft.load_jsonl(""), []) def test_summary_totals_by_hand(self): s = ft.summarize(self.make().spans) self.assertEqual(s["spans"], 4) self.assertEqual(s["errors"], 1) self.assertEqual((s["input_tokens"], s["output_tokens"]), (100 + 10, 10 + 2)) self.assertAlmostEqual(s["cost_usd"], ft.cost_usd("mock-model", 100, 10) + ft.cost_usd("mock-model", 10, 2)) self.assertEqual(s["by_operation"]["chat"]["count"], 2) self.assertEqual(s["by_operation"]["execute_tool"]["errors"], 1) def test_viewer_shows_a_tree(self): text = ft.render_trace(self.make().spans) lines = text.splitlines() self.assertTrue(lines[0].startswith("invoke_agent demo")) self.assertTrue(lines[1].startswith(" chat mock-model")) self.assertIn("tok 100/10", lines[1]) self.assertIn("$0.000021", lines[1]) self.assertIn("ERROR ValueError", lines[3]) self.assertEqual(len(lines), 4) def test_viewer_handles_a_partial_trace(self): spans = self.make().spans[1:] # the root is missing: children show as roots self.assertEqual(len(ft.render_trace(spans).splitlines()), 3) class TestStats(unittest.TestCase): def test_bootstrap_ci_basic_properties(self): m, lo, hi = ft.bootstrap_ci([1, 1, 0, 1, 0, 1, 1, 1, 0, 1]) self.assertAlmostEqual(m, 0.7) self.assertLessEqual(lo, m) self.assertGreaterEqual(hi, m) self.assertEqual(ft.bootstrap_ci([1, 1, 0, 1]), ft.bootstrap_ci([1, 1, 0, 1])) # same seed, same answer def test_constant_data_gives_a_point_interval(self): self.assertEqual(ft.bootstrap_ci([0.5] * 8), (0.5, 0.5, 0.5)) def test_more_data_gives_a_tighter_interval(self): small = [1, 1, 1, 0, 1, 1, 0, 1, 1, 1] big = small * 10 _, l1, h1 = ft.bootstrap_ci(small) _, l2, h2 = ft.bootstrap_ci(big) self.assertLess(h2 - l2, h1 - l1) def test_empty_input(self): with self.assertRaises(ValueError): ft.bootstrap_ci([]) def test_pareto_front(self): pts = [("a", 1.0, 0.8), ("b", 2.0, 0.8), ("c", 2.0, 0.9), ("d", 0.5, 0.5), ("e", 3.0, 0.7)] self.assertEqual(ft.pareto_front(pts), ["a", "c", "d"]) # b and e are beaten by a (cheaper, no worse) def test_pareto_ties_are_both_kept(self): self.assertEqual(ft.pareto_front([("a", 1.0, 0.5), ("b", 1.0, 0.5)]), ["a", "b"]) self.assertEqual(ft.pareto_front([]), []) class TestDemo(unittest.TestCase): def run_demo(self): buf = io.StringIO() with contextlib.redirect_stdout(buf): ft.demo() return buf.getvalue() def test_demo_runs_and_is_deterministic(self): a, b = self.run_demo(), self.run_demo() self.assertEqual(a, b) for n in range(1, 7): self.assertIn("== %d. " % n, a) def test_demo_makes_the_teaching_points(self): out = self.run_demo() self.assertIn("file fixed: True", out) self.assertIn("notify ran: 2, effect applied: 2", out) self.assertIn("notify ran: 2, effect applied: 1", out) self.assertIn("Pareto front (lower cost, higher success): ['lean']", out) class TestCommandLine(TmpDirCase): def run_cli(self, *args): import subprocess import sys here = os.path.dirname(os.path.abspath(ft.__file__)) return subprocess.run([sys.executable, os.path.join(here, "frameworks_toolkit.py"), *args], capture_output=True, text=True, timeout=60) def test_view_prints_a_saved_trace(self): tr = ft.Tracer(TickClock()) lm = ft.TracedLM(scripted_lm(["x" * 40]), tr, "mock-model") with tr.span("invoke_agent demo", **{"gen_ai.operation.name": "invoke_agent"}): lm("p" * 400) path = self.path("spans.jsonl") with open(path, "w") as f: f.write(ft.export_jsonl(tr.spans)) r = self.run_cli("view", path) self.assertEqual(r.returncode, 0) self.assertIn("invoke_agent demo", r.stdout) self.assertIn(" chat mock-model", r.stdout) self.assertIn("totals: 2 spans, 0 errors, tokens in/out 100/10", r.stdout) def test_bad_usage_gives_a_message(self): r = self.run_cli("nonsense") self.assertEqual(r.returncode, 2) self.assertIn("usage", r.stdout) if __name__ == "__main__": unittest.main()