From 6f866584fad03d6457c922461c34d9e01831eb88 Mon Sep 17 00:00:00 2001 From: Tisha Chawla Date: Wed, 29 Jul 2026 00:35:33 +0530 Subject: [PATCH] integrations: Chronicle record/replay + unbiased overhead benchmark Attaches Chronicle to the two neutral MAS from the outside, via the LLMClient seam. Nothing in src/testbench imports Chronicle; the code lives under integrations/ so the tool-agnostic CI rule (tests/test_agnostic.py) stays green. - integrations/chronicle/chronicle_client.py: a ChronicleClient that wraps any LLMClient and records each model call as an envelope, replaying it on stub. - demo.py / demo_loop.py: record -> replay -> cut-point on the orchestrator and the looping evaluator-optimizer MAS. - benchmark.py: unbiased recording-overhead + determinism harness. Because the workload cannot import Chronicle, this overhead cannot be tuned to flatter the tool. Recording adds ~28-300 us per crossing (median ~163 us, under 0.1% of a 300 ms model call); replay is deterministic with zero real model calls. - test_replay_regression.py + test_overhead_benchmark.py: assert detection, determinism, and zero model calls for both MAS. - workers.py: doc comments naming each agent and the model-call seam (tool neutral, no imports). Co-Authored-By: Claude Opus 4.8 Signed-off-by: Tisha Chawla --- .gitignore | 3 + integrations/chronicle/README.md | 66 ++++++ integrations/chronicle/benchmark.py | 209 ++++++++++++++++++ integrations/chronicle/chronicle_client.py | 120 ++++++++++ integrations/chronicle/demo.py | 85 +++++++ integrations/chronicle/demo_loop.py | 77 +++++++ .../chronicle/test_overhead_benchmark.py | 28 +++ .../chronicle/test_replay_regression.py | 110 +++++++++ src/testbench/orchestrator_workers/workers.py | 11 + 9 files changed, 709 insertions(+) create mode 100644 integrations/chronicle/README.md create mode 100644 integrations/chronicle/benchmark.py create mode 100644 integrations/chronicle/chronicle_client.py create mode 100644 integrations/chronicle/demo.py create mode 100644 integrations/chronicle/demo_loop.py create mode 100644 integrations/chronicle/test_overhead_benchmark.py create mode 100644 integrations/chronicle/test_replay_regression.py diff --git a/.gitignore b/.gitignore index 9332be4..c3d7d5e 100644 --- a/.gitignore +++ b/.gitignore @@ -19,3 +19,6 @@ env/ # Local *.log .env + +# Chronicle integration demo output (regenerated) +integrations/chronicle/.fixtures/ diff --git a/integrations/chronicle/README.md b/integrations/chronicle/README.md new file mode 100644 index 0000000..5564383 --- /dev/null +++ b/integrations/chronicle/README.md @@ -0,0 +1,66 @@ +# Chronicle integration (record and replay a testbench MAS) + +This folder attaches [chronicle](https://github.com/theagentplane/chronicle) to the +testbench **from the outside**. Nothing in `src/testbench` imports chronicle (a CI test +enforces that), so this code lives here under `integrations/` instead, where it is free +to import both. + +## The one idea + +Every testbench agent depends on a single model seam: `LLMClient.complete(model, messages)`. +Chronicle wraps that seam. On a live run it writes down each call; on replay it hands the +recorded answer straight back instead of calling the model. + +``` +agent -> ChronicleClient.complete(...) -> real client (RECORD: also writes an Envelope) +agent -> ChronicleClient.complete(...) -> recorded answer (REPLAY: no real client call) +``` + +We wrap the client, not the event sink, because replay has to **return** a recorded value +into the agent. The event sink only observes; it cannot feed a value back. + +## Terms (same as chronicle) + +- **Boundary**: one model call. Here the boundary name is the agent role (`supervisor`, + `researcher`, `analyst`, `writer`), read from the `ROLE:` system message. +- **Envelope**: the record of one boundary crossing: its input messages, its output text, + token usage, and model. It records that I/O, not any side effects. +- **Fixture**: the committed folder of Envelopes for one run (`graph.json` + one JSON per + boundary). +- **Stub / live**: on replay a *stubbed* boundary returns its recorded output without + running; a *live* boundary runs the real client. A **cut-point** is the one boundary you + set live to test a change while the rest stays stubbed. + +## Run the demos + +Both MAS use the *same* `ChronicleClient` with no changes. + +```bash +python integrations/chronicle/demo.py # orchestrator-workers (linear, 4 calls) +python integrations/chronicle/demo_loop.py # evaluator-optimizer (loops, 6 calls) +``` + +Each prints three phases and how many times the real model client was actually called: + +``` +1) RECORD real model calls: 4/6 (fixture written) +2) REPLAY (stub all) real model calls: 0 (identical result) +3) CUT-POINT real model calls: 1 (one boundary ran live) +``` + +The loop demo cut-points `generator@2` (the second draft only), which shows how chronicle +tells repeated calls apart by **invocation index**. + +## Run the tests + +```bash +python -m pytest integrations/chronicle/test_replay_regression.py -q +``` + +## Files + +- `chronicle_client.py` : the `ChronicleClient` adapter (wraps any `LLMClient`). +- `demo.py` : orchestrator-workers: record, replay, cut-point. +- `demo_loop.py` : evaluator-optimizer: same, on a looping MAS. +- `test_replay_regression.py` : the phases above as assertions, for both MAS. +- `.fixtures/` : demo output, gitignored (regenerated on each run). diff --git a/integrations/chronicle/benchmark.py b/integrations/chronicle/benchmark.py new file mode 100644 index 0000000..e5ad311 --- /dev/null +++ b/integrations/chronicle/benchmark.py @@ -0,0 +1,209 @@ +"""Unbiased recording-overhead and determinism measurement for Chronicle, run on the +neutral testbench MAS (which imports nothing from Chronicle). + +Because the workload does not know Chronicle exists (a CI test enforces that), the +overhead measured here cannot be tuned to flatter the tool. This is the number to cite +for recording overhead on a realistic multi-agent system, alongside the fault-detection +numbers from Chronicle's own incident benchmark. + + python integrations/chronicle/benchmark.py + python integrations/chronicle/benchmark.py --json out.json --tex table.tex + +Reports, for each MAS: +- recording overhead (in-memory) in microseconds per crossing and as a fraction of a + typical model call, +- deterministic replay (N replays reproduce the run) with zero real model calls. +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import sys +import tempfile +from dataclasses import asdict, dataclass +from pathlib import Path +from time import perf_counter + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import chronicle +from chronicle import ReplayPlan + +from chronicle_client import ChronicleClient +from testbench.core import LLMClient, ModelResponse, RunConfig, make_client +from testbench.evaluator_optimizer.loop import refine +from testbench.orchestrator_workers.graph import build_graph, initial_state + +MODEL_LATENCY_MS = 300.0 +REPS = 150 +ROUNDS = 7 +REPLAYS = 20 + +ORCH_TASK = "What is token governance and why does it matter for agents?" +LOOP_TASK = "Explain token governance to a new engineer in two sentences." + + +class _Spy: + """Counts how many calls actually reach the real client.""" + + def __init__(self, inner: LLMClient) -> None: + self.inner = inner + self.calls = 0 + + def complete(self, model, messages, *, max_output_tokens=None) -> ModelResponse: + self.calls += 1 + return self.inner.complete(model, messages, max_output_tokens=max_output_tokens) + + +def _prep_orch(client: LLMClient): + # Compile the LangGraph once, so timing measures the run, not graph construction. + graph = build_graph(client, sink=None, cfg=RunConfig()) + return lambda: graph.invoke(initial_state(ORCH_TASK))["brief"] + + +def _prep_loop(client: LLMClient): + return lambda: ( + lambda r: f"{r.answer}|{r.score}|{r.iterations}" + )(refine(LOOP_TASK, client, sink=None, cfg=RunConfig())) + + +# name -> (prepare(client) -> zero-arg runnable returning a stable outcome, crossings) +WORKLOADS = { + "Orchestrator": (_prep_orch, 4), + "Loop": (_prep_loop, 6), +} + + +def _best_time(fn, reps: int = REPS, rounds: int = ROUNDS) -> float: + for _ in range(min(reps, 30)): + fn() + best = float("inf") + for _ in range(rounds): + start = perf_counter() + for _ in range(reps): + fn() + best = min(best, (perf_counter() - start) / reps) + return best + + +@dataclass +class Result: + mas: str + crossings: int + t_base_us: float + t_rec_us: float + overhead_us_per_crossing: float + overhead_pct_at_model: float + deterministic: bool + replay_model_calls: int + + +def measure(name: str, prepare, crossings: int) -> Result: + # Baseline: the plain client, no Chronicle in the path. Graph built once. + base_run = prepare(make_client(RunConfig())) + t_base = _best_time(base_run) + + # Recording: ChronicleClient, in-memory (store=None) to isolate compute from disk. + rec_run = prepare(ChronicleClient(make_client(RunConfig()))) + + def rec() -> None: + with chronicle.record(name, store=None): + rec_run() + + t_rec = _best_time(rec) + + # Determinism + zero model calls: record once, replay REPLAYS times. + with tempfile.TemporaryDirectory() as tmp: + fixture = str(Path(tmp) / name) + spy = _Spy(make_client(RunConfig())) + run = prepare(ChronicleClient(spy)) + with chronicle.record(name, export=fixture): + first = run() + outcomes, replay_calls = [], 0 + for i in range(REPLAYS): + spy.calls = 0 + with chronicle.replay_trace(fixture, ReplayPlan()): + outcomes.append(run()) + if i == 0: + replay_calls = spy.calls + deterministic = all(o == first for o in outcomes) + + per_crossing_us = (t_rec - t_base) * 1e6 / crossings if crossings else 0.0 + return Result( + mas=name, + crossings=crossings, + t_base_us=t_base * 1e6, + t_rec_us=t_rec * 1e6, + overhead_us_per_crossing=per_crossing_us, + overhead_pct_at_model=per_crossing_us / (MODEL_LATENCY_MS * 1000.0) * 100.0, + deterministic=deterministic, + replay_model_calls=replay_calls, + ) + + +def run_all() -> list[Result]: + return [measure(name, prep, crossings) for name, (prep, crossings) in WORKLOADS.items()] + + +def print_table(results: list[Result]) -> None: + print() + print(f" {'MAS':<13} {'cross':>5} {'t_base':>9} {'t_rec':>9} {'rec/cross':>10} " + f"{'%@model':>8} {'determ':>7} {'replay calls':>13}") + print(" " + "-" * 82) + for r in results: + print(f" {r.mas:<13} {r.crossings:>5} {r.t_base_us:>8.1f}u {r.t_rec_us:>8.1f}u " + f"{r.overhead_us_per_crossing:>9.1f}u {r.overhead_pct_at_model:>7.3f}% " + f"{'yes' if r.deterministic else 'NO':>7} {r.replay_model_calls:>13}") + print(" " + "-" * 82) + med = statistics.median(r.overhead_us_per_crossing for r in results) + med_pct = statistics.median(r.overhead_pct_at_model for r in results) + calls = sum(r.replay_model_calls for r in results) + print(f"\n unbiased recording overhead : {med:.1f} us/crossing " + f"(= {med_pct:.3f}% of a {MODEL_LATENCY_MS:.0f} ms model call)") + print(f" replay model calls : {calls} (across all workloads)") + print(f" deterministic replay : {'yes' if all(r.deterministic for r in results) else 'NO'}" + f" ({REPLAYS} replays each)\n") + + +def to_latex(results: list[Result]) -> str: + med = statistics.median(r.overhead_us_per_crossing for r in results) + med_pct = statistics.median(r.overhead_pct_at_model for r in results) + lines = [ + "% Auto-generated by integrations/chronicle/benchmark.py (neutral testbench)", + "\\newcommand{\\unbiasedoverheadus}{%.1f}" % med, + "\\newcommand{\\unbiasedoverheadpct}{%.3f}" % med_pct, + "", + "\\begin{tabular}{lrrr}", + "\\toprule", + "MAS & Crossings & Rec./cross. (\\textmu s) & Overhead \\\\", + "\\midrule", + ] + for r in results: + lines.append( + f"{r.mas} & {r.crossings} & {r.overhead_us_per_crossing:.1f} & " + f"{r.overhead_pct_at_model:.3f}\\% \\\\" + ) + lines += ["\\bottomrule", "\\end{tabular}"] + return "\n".join(lines) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Unbiased Chronicle overhead on the testbench") + parser.add_argument("--json", type=Path) + parser.add_argument("--tex", type=Path) + args = parser.parse_args() + + results = run_all() + print_table(results) + if args.json: + args.json.write_text(json.dumps([asdict(r) for r in results], indent=2), encoding="utf-8") + print(f" wrote {args.json}") + if args.tex: + args.tex.write_text(to_latex(results), encoding="utf-8") + print(f" wrote {args.tex}") + + +if __name__ == "__main__": + main() diff --git a/integrations/chronicle/chronicle_client.py b/integrations/chronicle/chronicle_client.py new file mode 100644 index 0000000..bca8320 --- /dev/null +++ b/integrations/chronicle/chronicle_client.py @@ -0,0 +1,120 @@ +"""A chronicle recorder that plugs into the testbench's LLMClient seam. + +The testbench imports nothing from chronicle (a CI test enforces that). Instead, +chronicle attaches from the *outside* by wrapping the one dependency every agent has: +the ``LLMClient``. Each ``complete(...)`` call becomes a chronicle Envelope on the way +out (record), and on replay the recorded Envelope is handed straight back without ever +calling the real client. + +Why wrap the client and not the event sink? Because replay has to *substitute the return +value* of a model call, and the client's ``complete`` is the only seam that returns a +value the agent then uses. The event sink is fire-and-forget, good for observing but not +for feeding a recorded answer back in. + +This file lives outside ``src/testbench`` on purpose, so it can import chronicle without +breaking the agnostic rule. +""" + +from __future__ import annotations + +from chronicle import get_session +from chronicle.envelope.schema import ActionResult, InputState +from chronicle.session import SessionMode + +from testbench.core import LLMClient, Message, ModelResponse, Usage + + +def _boundary_id(messages: list[Message]) -> str: + """Use the agent's role as the boundary name (supervisor, researcher, ...). + + Every testbench agent puts a leading ``ROLE: `` system message on its call, + so we can name each boundary without the MAS telling us who is calling. + """ + for m in messages: + if m.role == "system" and m.content.startswith("ROLE:"): + return m.content.split("\n", 1)[0][len("ROLE:") :].strip() + return "model" + + +class ChronicleClient: + """Wraps any ``LLMClient`` and records/replays every ``complete`` call. + + - Outside a chronicle ``record`` / ``replay_trace`` block it is transparent: it just + calls the inner client (nothing is recorded). + - Inside ``record(...)``: it calls the inner client and writes one Envelope per call. + - Inside ``replay_trace(...)``: a *stubbed* boundary returns the recorded response and + never touches the inner client; a *live* boundary (a cut-point) runs the inner client + and its result is captured for asserts. + """ + + def __init__(self, inner: LLMClient) -> None: + self.inner = inner + + def complete( + self, + model: str, + messages: list[Message], + *, + max_output_tokens: int | None = None, + ) -> ModelResponse: + session = get_session() + boundary_id = _boundary_id(messages) + input_state = InputState( + messages=[{"role": m.role, "content": m.content} for m in messages], + ) + + # REPLAY + this boundary is stubbed: hand back the recorded answer, no model call. + if session.mode is SessionMode.REPLAY and _should_stub(session, boundary_id): + envelope = session._fixture_for(boundary_id) + return _response_from_envelope(envelope, model) + + # Transparent when chronicle is not active (no trace begun in LIVE mode). + if session.mode is SessionMode.LIVE and not getattr(session, "trace_id", None): + return self.inner.complete(model, messages, max_output_tokens=max_output_tokens) + + # Run the real client (recording live, or running a live cut-point on replay). + resp = self.inner.complete(model, messages, max_output_tokens=max_output_tokens) + + if session.mode is SessionMode.REPLAY: + # Cut-point: this boundary ran live. Capture its real output so a test can + # assert on it, and advance the per-boundary cursor like chronicle does. + idx = session._replay_cursor.get(boundary_id, 0) + 1 + session.capture_live_input(boundary_id, idx, input_state) + session.capture_live_result(boundary_id, idx, resp) + session.next_invocation(boundary_id) + session._replay_cursor[boundary_id] = idx + else: + # Recording: one Envelope for this crossing (input + output + usage + model). + session.record_envelope( + boundary_id, + "llm", + input_state, + ActionResult( + completion=resp.text, + token_usage={ + "input_tokens": resp.usage.input_tokens, + "output_tokens": resp.usage.output_tokens, + }, + ), + model_version=resp.model, + ) + return resp + + +def _should_stub(session, boundary_id: str) -> bool: + invocation_index = session._replay_cursor.get(boundary_id, 0) + 1 + return session.replay_plan.should_stub(boundary_id, invocation_index) + + +def _response_from_envelope(envelope, model: str) -> ModelResponse: + """Rebuild the testbench's ModelResponse from a recorded Envelope.""" + action = envelope.action_result + usage = action.token_usage or {} + return ModelResponse( + text=action.completion or "", + usage=Usage( + input_tokens=int(usage.get("input_tokens", 0)), + output_tokens=int(usage.get("output_tokens", 0)), + ), + model=envelope.metadata.model_version or model, + ) diff --git a/integrations/chronicle/demo.py b/integrations/chronicle/demo.py new file mode 100644 index 0000000..321a498 --- /dev/null +++ b/integrations/chronicle/demo.py @@ -0,0 +1,85 @@ +"""Record a testbench MAS run with chronicle, then replay it and cut-point test a fix. + +Run it: + + python integrations/chronicle/demo.py + +Nothing in ``src/testbench`` imports chronicle. This script wires the two together from +the outside: it wraps the testbench's ``LLMClient`` with ``ChronicleClient`` and drives +the unchanged orchestrator-workers MAS. + +The three phases print how many times the *real* model client was actually called, which +is the whole point: replay reproduces the run with zero model calls. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +# Make "chronicle_client" importable whether run from the repo root or this folder. +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import chronicle +from chronicle import ReplayPlan + +from chronicle_client import ChronicleClient +from testbench.core import LLMClient, Message, ModelResponse, RunConfig, make_client +from testbench.orchestrator_workers.graph import build_graph, initial_state + +TASK = "What is token governance and why does it matter for agents?" +FIXTURE = str(Path(__file__).resolve().parent / ".fixtures" / "brief-run") + + +class CountingClient: + """A see-through spy: forwards every call, but counts how many really happen.""" + + def __init__(self, inner: LLMClient) -> None: + self.inner = inner + self.calls = 0 + + def complete(self, model, messages, *, max_output_tokens=None) -> ModelResponse: + self.calls += 1 + return self.inner.complete(model, messages, max_output_tokens=max_output_tokens) + + +def run_mas(client: LLMClient) -> dict: + graph = build_graph(client, sink=None, cfg=RunConfig()) + return graph.invoke(initial_state(TASK)) + + +def main() -> None: + # Real (offline, deterministic) client -> counter -> chronicle recorder. + spy = CountingClient(make_client(RunConfig())) + client = ChronicleClient(spy) + + # ---- Phase 1: RECORD ------------------------------------------------------------- + spy.calls = 0 + with chronicle.record("brief-run", export=FIXTURE): + recorded = run_mas(client) + print("1) RECORD") + print(f" real model calls: {spy.calls}") + print(f" fixture written : {FIXTURE}") + print(f" brief starts : {recorded['brief'].splitlines()[0]!r}\n") + + # ---- Phase 2: REPLAY (stub everything) ------------------------------------------- + spy.calls = 0 + with chronicle.replay_trace(FIXTURE, ReplayPlan()): # default: stub all + replayed = run_mas(client) + print("2) REPLAY (all boundaries stubbed)") + print(f" real model calls: {spy.calls} <- zero: served from the fixture") + print(f" identical brief : {replayed['brief'] == recorded['brief']}\n") + + # ---- Phase 3: CUT-POINT (run only the writer live) ------------------------------- + spy.calls = 0 + plan = ReplayPlan().live("writer", 1) # writer runs live; everything upstream stubbed + with chronicle.replay_trace(FIXTURE, plan) as session: + run_mas(client) + writer_out = session.captured_result("writer", 1) + print("3) CUT-POINT (writer live, upstream stubbed)") + print(f" real model calls: {spy.calls} <- only the writer") + print(f" writer live out : {writer_out.text.splitlines()[0]!r}") + + +if __name__ == "__main__": + main() diff --git a/integrations/chronicle/demo_loop.py b/integrations/chronicle/demo_loop.py new file mode 100644 index 0000000..cb06bf6 --- /dev/null +++ b/integrations/chronicle/demo_loop.py @@ -0,0 +1,77 @@ +"""Record + replay the *looping* MAS (evaluator-optimizer) with chronicle. + + python integrations/chronicle/demo_loop.py + +This MAS retries: generate a draft, score it, and if the score is too low feed the +critique back and generate again (up to 3 rounds). So the ``generator`` and ``evaluator`` +boundaries are each hit several times. Chronicle keeps them apart by *invocation index*: +generator@1, generator@2, generator@3 are distinct recorded calls. + +The same ``ChronicleClient`` from the first demo is reused with no changes: it does not +care which MAS is running, only that model calls go through ``LLMClient.complete``. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import chronicle +from chronicle import ReplayPlan + +from chronicle_client import ChronicleClient +from testbench.core import LLMClient, ModelResponse, RunConfig, make_client +from testbench.evaluator_optimizer.loop import refine + +TASK = "Explain token governance to a new engineer in two sentences." +FIXTURE = str(Path(__file__).resolve().parent / ".fixtures" / "refine-run") + + +class CountingClient: + """Forwards every call, counting how many really reach the model.""" + + def __init__(self, inner: LLMClient) -> None: + self.inner = inner + self.calls = 0 + + def complete(self, model, messages, *, max_output_tokens=None) -> ModelResponse: + self.calls += 1 + return self.inner.complete(model, messages, max_output_tokens=max_output_tokens) + + +def main() -> None: + spy = CountingClient(make_client(RunConfig())) + client = ChronicleClient(spy) + + # ---- Phase 1: RECORD (the loop runs until the score passes 0.8) ------------------- + spy.calls = 0 + with chronicle.record("refine-run", export=FIXTURE): + result = refine(TASK, client, sink=None, cfg=RunConfig()) + print("1) RECORD") + print(f" real model calls: {spy.calls} (generator x3 + evaluator x3)") + print(f" iterations : {result.iterations} final score: {result.score:.2f}") + print(f" fixture written : {FIXTURE}\n") + + # ---- Phase 2: REPLAY (stub everything) ------------------------------------------- + spy.calls = 0 + with chronicle.replay_trace(FIXTURE, ReplayPlan()): + replayed = refine(TASK, client, sink=None, cfg=RunConfig()) + print("2) REPLAY (all boundaries stubbed)") + print(f" real model calls: {spy.calls} <- zero: the whole loop is served from the fixture") + print(f" same result : {replayed.answer == result.answer and replayed.score == result.score}\n") + + # ---- Phase 3: CUT-POINT on ONE invocation (the 2nd draft only) ------------------- + spy.calls = 0 + plan = ReplayPlan().live("generator", 2) # only generator's 2nd attempt runs live + with chronicle.replay_trace(FIXTURE, plan) as session: + refine(TASK, client, sink=None, cfg=RunConfig()) + second_draft = session.captured_result("generator", 2) + print("3) CUT-POINT (generator@2 live, everything else stubbed)") + print(f" real model calls: {spy.calls} <- only the 2nd draft") + print(f" generator@2 out : {second_draft.text!r}") + + +if __name__ == "__main__": + main() diff --git a/integrations/chronicle/test_overhead_benchmark.py b/integrations/chronicle/test_overhead_benchmark.py new file mode 100644 index 0000000..63ab066 --- /dev/null +++ b/integrations/chronicle/test_overhead_benchmark.py @@ -0,0 +1,28 @@ +"""The unbiased overhead benchmark runs, replays deterministically, and makes zero real +model calls. Timing itself is not asserted (it is machine-dependent); correctness is. + +Run: python -m pytest integrations/chronicle/test_overhead_benchmark.py -q +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import benchmark + + +def test_benchmark_runs_deterministically_with_zero_model_calls(): + # Shrink the timing loops so the test is fast; correctness is independent of them. + benchmark.REPS = 5 + benchmark.ROUNDS = 2 + benchmark.REPLAYS = 3 + + results = benchmark.run_all() + + assert {r.mas for r in results} == set(benchmark.WORKLOADS) + assert all(r.crossings > 0 for r in results) + assert all(r.deterministic for r in results) + assert all(r.replay_model_calls == 0 for r in results) diff --git a/integrations/chronicle/test_replay_regression.py b/integrations/chronicle/test_replay_regression.py new file mode 100644 index 0000000..7403806 --- /dev/null +++ b/integrations/chronicle/test_replay_regression.py @@ -0,0 +1,110 @@ +"""Self-contained proof that chronicle records and replays a testbench MAS. + +This is not part of the testbench's own suite (its pytest ``testpaths`` is ``tests`` +only, and that suite must stay tool-agnostic). Run it explicitly: + + python -m pytest integrations/chronicle/test_replay_regression.py -q + +It records the orchestrator-workers run to a temp fixture, then proves replay reproduces +it with zero model calls, and that a cut-point runs exactly one boundary live. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) + +import chronicle +from chronicle import ReplayPlan + +from chronicle_client import ChronicleClient +from testbench.core import LLMClient, ModelResponse, RunConfig, make_client +from testbench.evaluator_optimizer.loop import refine +from testbench.orchestrator_workers.graph import build_graph, initial_state + +TASK = "What is token governance and why does it matter for agents?" +LOOP_TASK = "Explain token governance to a new engineer in two sentences." + + +class _Spy: + def __init__(self, inner: LLMClient) -> None: + self.inner = inner + self.calls = 0 + + def complete(self, model, messages, *, max_output_tokens=None) -> ModelResponse: + self.calls += 1 + return self.inner.complete(model, messages, max_output_tokens=max_output_tokens) + + +def _run(client: LLMClient) -> dict: + return build_graph(client, sink=None, cfg=RunConfig()).invoke(initial_state(TASK)) + + +def test_record_then_replay_makes_no_model_calls(tmp_path): + spy = _Spy(make_client(RunConfig())) + client = ChronicleClient(spy) + fixture = str(tmp_path / "brief-run") + + with chronicle.record("brief-run", export=fixture): + recorded = _run(client) + assert spy.calls == 4 # supervisor, researcher, analyst, writer + + spy.calls = 0 + with chronicle.replay_trace(fixture, ReplayPlan()): # stub everything + replayed = _run(client) + assert spy.calls == 0 # served entirely from the fixture + assert replayed["brief"] == recorded["brief"] # deterministic reproduction + + +def test_cut_point_runs_only_the_target_boundary_live(tmp_path): + spy = _Spy(make_client(RunConfig())) + client = ChronicleClient(spy) + fixture = str(tmp_path / "brief-run") + + with chronicle.record("brief-run", export=fixture): + _run(client) + + spy.calls = 0 + with chronicle.replay_trace(fixture, ReplayPlan().live("writer", 1)) as session: + _run(client) + writer_out = session.captured_result("writer", 1) + assert spy.calls == 1 # only the writer ran live + assert writer_out.text.startswith("Brief:") + + +def test_looping_mas_replays_each_invocation_index(tmp_path): + spy = _Spy(make_client(RunConfig())) + client = ChronicleClient(spy) + fixture = str(tmp_path / "refine-run") + + with chronicle.record("refine-run", export=fixture): + recorded = refine(LOOP_TASK, client, sink=None, cfg=RunConfig()) + # 3 rounds: generator + evaluator each called 3 times. + assert spy.calls == 6 + assert recorded.iterations == 3 + + spy.calls = 0 + with chronicle.replay_trace(fixture, ReplayPlan()): + replayed = refine(LOOP_TASK, client, sink=None, cfg=RunConfig()) + assert spy.calls == 0 + assert replayed.answer == recorded.answer + assert replayed.score == recorded.score + + +def test_cut_point_targets_a_single_invocation(tmp_path): + spy = _Spy(make_client(RunConfig())) + client = ChronicleClient(spy) + fixture = str(tmp_path / "refine-run") + + with chronicle.record("refine-run", export=fixture): + refine(LOOP_TASK, client, sink=None, cfg=RunConfig()) + + # Run only the generator's 2nd attempt live; the other five calls stay stubbed. + spy.calls = 0 + with chronicle.replay_trace(fixture, ReplayPlan().live("generator", 2)) as session: + refine(LOOP_TASK, client, sink=None, cfg=RunConfig()) + second_draft = session.captured_result("generator", 2) + assert spy.calls == 1 + assert second_draft is not None diff --git a/src/testbench/orchestrator_workers/workers.py b/src/testbench/orchestrator_workers/workers.py index e51cb2b..4fac00b 100644 --- a/src/testbench/orchestrator_workers/workers.py +++ b/src/testbench/orchestrator_workers/workers.py @@ -23,6 +23,9 @@ def _call( extra: list[tuple[str, str]] | None = None, kind: str = "model", ) -> str: + # Shared helper: every agent makes its one model call here. This single + # `client.complete(...)` line is the one decision point an external observer can + # wrap to record or meter the call, without the agent knowing about it. messages = [Message("system", f"ROLE: {role}"), Message("user", user)] for label, value in extra or []: messages.append(Message("user", f"{label}\n{value}")) @@ -42,6 +45,9 @@ def _call( def supervisor_node( state: ResearchState, client: LLMClient, sink: EventSink, cfg: RunConfig ) -> StateUpdate: + """Supervisor: the router. On its first visit it asks the model for a plan, then on + every visit it picks the next unfinished worker (researcher -> analyst -> writer) + and finally FINISH. It makes its one model call only once (to plan).""" plan = state.get("plan") or "" if not plan: plan = _call( @@ -65,6 +71,7 @@ def supervisor_node( def researcher_node( state: ResearchState, client: LLMClient, sink: EventSink, cfg: RunConfig ) -> StateUpdate: + """Researcher: gathers raw facts about the task from the model. Writes `research`.""" text = _call(client, sink, role="researcher", model=cfg.model_worker, user=state["task"]) return {"research": text} @@ -72,6 +79,8 @@ def researcher_node( def analyst_node( state: ResearchState, client: LLMClient, sink: EventSink, cfg: RunConfig ) -> StateUpdate: + """Analyst: turns the researcher's raw facts (passed in as SOURCE) into a few key + points. Writes `analysis`.""" text = _call( client, sink, @@ -86,6 +95,8 @@ def analyst_node( def writer_node( state: ResearchState, client: LLMClient, sink: EventSink, cfg: RunConfig ) -> StateUpdate: + """Writer: composes the final brief from the analyst's key points (passed in as + ANALYSIS). Writes `brief`, the MAS output.""" text = _call( client, sink,