agentfork-cli 0.1.0__py3-none-any.whl

agentfork/__init__.py

@@ -0,0 +1,3 @@
+from .recorder import trace, tool, wrap_anthropic, wrap_openai
+
+__all__ = ["trace", "tool", "wrap_openai", "wrap_anthropic"]

agentfork/cli.py

@@ -0,0 +1,259 @@
+import json
+from datetime import datetime
+from typing import List, Optional
+
+import typer
+from rich.console import Console
+from rich.syntax import Syntax
+from rich.table import Table
+from rich.tree import Tree
+
+from . import replay, store
+
+app = typer.Typer(help="agentfork: record, inspect and replay LLM agent runs.")
+console = Console()
+STATUS_COLOR = {"ok": "green", "error": "red", "running": "yellow"}
+
+def _load(text):
+    return json.loads(text) if text else None
+
+def _fmt_time(ts):
+    return datetime.fromtimestamp(ts).strftime("%m-%d %H:%M:%S") if ts else "-"
+
+def _duration(run):
+    if not run.get("ended_at"):
+        return "-"
+    return f"{run['ended_at'] - run['started_at']:.1f}s"
+
+def _summarize(step):
+    inp, out = _load(step["input_json"]), _load(step["output_json"])
+    if step["error"]:
+        return f"[red]{step['error']}[/red]"
+    if step["type"] == "tool" and isinstance(inp, dict):
+        args = inp.get("kwargs") or inp.get("args")
+        return f"{args} -> {out}"
+    if step["type"] == "llm" and isinstance(out, dict):
+        if "choices" in out:
+            msg = (out.get("choices") or [{}])[0].get("message") or {}
+            calls = msg.get("tool_calls") or []
+            if calls:
+                return "calls " + ", ".join(c["function"]["name"] for c in calls)
+            return (msg.get("content") or "")[:120]
+        blocks = [b for b in (out.get("content") or []) if isinstance(b, dict)]
+        tools = [b.get("name") for b in blocks if b.get("type") == "tool_use"]
+        if tools:
+            return "calls " + ", ".join(tools)
+        return " ".join(b.get("text", "") for b in blocks if b.get("type") == "text")[:120]
+    return str(out)[:120]
+
+def _label(s):
+    tokens = f" {s['tokens_in']}->{s['tokens_out']} tok" if s["tokens_in"] is not None else ""
+    kind = "cyan" if s["type"] == "llm" else "magenta"
+    cached = " [green](cached)[/green]" if s.get("cached") else ""
+    return (f"[dim]#{s['seq']}[/dim] [{kind}]{s['type']}[/{kind}] [bold]{s['name']}[/bold] "
+            f"[dim]{s['latency_ms']:.0f}ms{tokens}[/dim]{cached}  {_summarize(s)}")
+
+def _resolve(prefix):
+    matches = store.find_runs(prefix)
+    if not matches:
+        console.print(f"[red]No run found starting with '{prefix}'[/red]")
+        raise typer.Exit(1)
+    if len(matches) > 1:
+        console.print(f"[yellow]'{prefix}' matches {len(matches)} runs; use more characters[/yellow]")
+        raise typer.Exit(1)
+    return matches[0]
+
+@app.command("list")
+def list_runs(limit: int = typer.Option(10, help="How many recent runs to show.")):
+    """List recent runs."""
+    table = Table(title="Recent runs")
+    for col in ["id", "name", "status", "steps", "tokens", "duration", "started", "fork of"]:
+        table.add_column(col)
+    for r in store.list_runs(limit):
+        color = STATUS_COLOR.get(r["status"], "white")
+        fork = f"{r['parent_run_id'][:8]}@{r['fork_step']}" if r.get("parent_run_id") else ""
+        table.add_row(r["id"][:8], r["name"], f"[{color}]{r['status']}[/{color}]",
+                      str(r["n_steps"]), str(r["tokens"]), _duration(r), _fmt_time(r["started_at"]), fork)
+    console.print(table)
+
+@app.command()
+def show(run_id: str = typer.Argument(..., help="Run id or its first few characters."),
+         full: bool = typer.Option(False, "--full", help="Show full inputs and outputs.")):
+    """Show one run as a step tree."""
+    run = _resolve(run_id)
+    color = STATUS_COLOR.get(run["status"], "white")
+    fork = (f" [dim]fork of {run['parent_run_id'][:8]} at step {run['fork_step']}[/dim]"
+            if run.get("parent_run_id") else "")
+    root = Tree(f"[bold]{run['name']}[/bold] {run['id'][:8]} "
+                f"[{color}]{run['status']}[/{color}] {_duration(run)}{fork}")
+    nodes = {}
+    for s in store.get_steps(run["id"]):
+        node = nodes.get(s["parent_id"], root).add(_label(s))
+        if full:
+            node.add(Syntax(json.dumps(_load(s["input_json"]), indent=2), "json", word_wrap=True))
+            node.add(Syntax(json.dumps(_load(s["output_json"]), indent=2), "json", word_wrap=True))
+        nodes[s["id"]] = node
+    console.print(root)
+    console.print(f"[bold]Final output:[/bold] {_load(run['output_json'])}")
+
+@app.command()
+def fork(run_id: str = typer.Argument(..., help="Run to fork."),
+         step: int = typer.Option(..., "--step", help="Step number to change (see `show`)."),
+         output: Optional[str] = typer.Option(None, "--output", help="Fake result for a tool step."),
+         model: Optional[str] = typer.Option(None, "--model", help="Different model for an LLM step."),
+         system: Optional[str] = typer.Option(None, "--system", help="System prompt for an LLM step.")):
+    """Re-run a past run from one step with one change; earlier steps come from the cache."""
+    if output is None and model is None and system is None:
+        console.print("[red]Give one change: --output, --model or --system[/red]")
+        raise typer.Exit(1)
+    run = _resolve(run_id)
+    try:
+        new_id, _ = replay.fork(run, step, output=output, model=model, system=system)
+    except Exception as e:
+        console.print(f"[red]Fork failed: {e}[/red]")
+        raise typer.Exit(1)
+    console.print(f"[green]Forked {run['id'][:8]} at step {step} -> new run {new_id[:8]}[/green]")
+    show(new_id, full=False)
+    console.print(f"[dim]Compare: agentfork diff {run['id'][:8]} {new_id[:8]}[/dim]")
+
+@app.command()
+def diff(run_a: str = typer.Argument(..., help="First run."),
+         run_b: str = typer.Argument(..., help="Second run.")):
+    """Compare two runs step by step."""
+    a, b = _resolve(run_a), _resolve(run_b)
+    sa = {s["seq"]: s for s in store.get_steps(a["id"])}
+    sb = {s["seq"]: s for s in store.get_steps(b["id"])}
+    table = Table(title=f"{a['id'][:8]} vs {b['id'][:8]}", show_lines=True)
+    table.add_column("#")
+    table.add_column(a["id"][:8])
+    table.add_column(b["id"][:8])
+    for seq in sorted(set(sa) | set(sb)):
+        left = _label(sa[seq]) if seq in sa else "[dim]-[/dim]"
+        right = _label(sb[seq]) if seq in sb else "[dim]-[/dim]"
+        same = seq in sa and seq in sb and _summarize(sa[seq]) == _summarize(sb[seq])
+        table.add_row(str(seq) if same else f"[yellow]{seq} *[/yellow]", left, right)
+    console.print(table)
+    console.print(f"[bold]{a['id'][:8]} output:[/bold] {_load(a['output_json'])}")
+    console.print(f"[bold]{b['id'][:8]} output:[/bold] {_load(b['output_json'])}")
+
+TEST_TEMPLATE = '''# Generated by `agentfork save-test` from run {run_id}.
+# Replays the recorded run entirely from cache: no API calls, no cost.
+# If you change prompts, tools or agent logic on purpose, record a new run
+# and run `agentfork save-test <new-run-id>` again.
+import json
+from pathlib import Path
+
+from agentfork import replay
+
+FIXTURE = Path(__file__).with_name("run_{short}.json")
+ROOT = Path(__file__).resolve().parents[{depth}]
+
+
+def test_run_{short}_replays(tmp_path, monkeypatch):
+    monkeypatch.setenv("AGENTFORK_DIR", str(tmp_path))
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-agentfork-replay")
+    replay.check(json.loads(FIXTURE.read_text()), ROOT)
+'''
+
+@app.command("save-test")
+def save_test(run_id: str = typer.Argument(..., help="Run to turn into a test."),
+              out: str = typer.Option("tests/agentfork", "--out", help="Folder for the generated test.")):
+    """Save a run as a pytest regression test that replays it from cache."""
+    import os
+    from pathlib import Path
+    run = _resolve(run_id)
+    if not run.get("entry_file"):
+        console.print("[red]This run has no saved entry point; run your agent again and save the new run.[/red]")
+        raise typer.Exit(1)
+    if run.get("parent_run_id"):
+        console.print("[yellow]This run is a fork, and its change isn't in your code yet. "
+                      "Apply the fix, run your agent again, and save that run.[/yellow]")
+        raise typer.Exit(1)
+    cwd = Path.cwd()
+    out_dir = Path(out)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    short = run["id"][:8]
+    fixture = {"run": dict(run, entry_file=os.path.relpath(run["entry_file"], cwd)),
+               "steps": store.get_steps(run["id"])}
+    fixture_path = out_dir / f"run_{short}.json"
+    test_path = out_dir / f"test_run_{short}.py"
+    fixture_path.write_text(json.dumps(fixture, indent=2))
+    depth = len(Path(os.path.relpath(out_dir.resolve(), cwd)).parts)
+    test_path.write_text(TEST_TEMPLATE.format(run_id=run["id"], short=short, depth=depth))
+    console.print(f"[green]Saved {test_path} and {fixture_path.name}[/green]")
+    console.print(f"[dim]Run it with: pytest {test_path}[/dim]")
+
+
+def _make_labels(variants):
+    labels, counts = {}, {}
+    for v in variants:
+        key, value = next(iter(v.items()))
+        if key == "model":
+            labels[id(v)] = f"model={value}"
+        else:
+            counts[key] = counts.get(key, 0) + 1
+            labels[id(v)] = f"{key} #{counts[key]}"
+    return labels
+
+@app.command()
+def sweep(run_id: str = typer.Argument(..., help="Run to sweep."),
+          step: int = typer.Option(..., "--step", help="Step to vary (see `show`)."),
+          model: Optional[str] = typer.Option(None, "--model", help="Comma-separated models to try."),
+          system: Optional[List[str]] = typer.Option(None, "--system", help="System prompt to try; repeat for more."),
+          output: Optional[List[str]] = typer.Option(None, "--output", help="Tool result to try; repeat for more."),
+          repeat: int = typer.Option(1, "--repeat", help="Runs per variant (LLMs vary run to run)."),
+          expect: Optional[str] = typer.Option(None, "--expect", help="Text a good final output contains.")):
+    """Fork one step across several models, prompts or tool results and compare outcomes."""
+    variants = [{"model": m.strip()} for m in (model or "").split(",") if m.strip()]
+    variants += [{"system": s} for s in (system or [])]
+    variants += [{"output": o} for o in (output or [])]
+    if not variants:
+        console.print("[red]Give at least one variant: --model, --system or --output[/red]")
+        raise typer.Exit(1)
+    run = _resolve(run_id)
+    labels = _make_labels(variants)
+    console.print(f"[bold]Sweeping step {step} of {run['id'][:8]}: "
+                  f"{len(variants)} variants x {repeat} = {len(variants) * repeat} runs[/bold]")
+
+    def progress(r):
+        mark = "[green]ok[/green]" if r["status"] == "ok" else "[red]error[/red]"
+        verdict = "" if r["passed"] is None else (" [green]pass[/green]" if r["passed"] else " [red]fail[/red]")
+        console.print(f"  {mark}{verdict} {labels[id(r['variant'])]} [dim]{r['tokens']} tok[/dim]")
+
+    try:
+        results = replay.sweep(run, step, variants, repeat=repeat, expect=expect, on_result=progress)
+    except ValueError as e:
+        console.print(f"[red]{e}[/red]")
+        raise typer.Exit(1)
+
+    groups = {"original": [replay.baseline(run, step, expect)]}
+    for r in results:
+        groups.setdefault(labels[id(r["variant"])], []).append(r)
+    table = Table(title=f"Sweep of step {step} ({run['id'][:8]})", show_lines=True)
+    table.add_column("variant", no_wrap=True)
+    table.add_column("runs", no_wrap=True, min_width=4)
+    table.add_column("ok", no_wrap=True, min_width=3)
+    if expect:
+        table.add_column("pass", no_wrap=True, min_width=4)
+    table.add_column("avg tokens", no_wrap=True, min_width=6)
+    table.add_column("avg time", no_wrap=True, min_width=6)
+    table.add_column("sample output", max_width=36, overflow="ellipsis", no_wrap=True)
+    table.add_column("run ids", no_wrap=True)
+    for label, rs in groups.items():
+        n = len(rs)
+        ids = [(r["run_id"] or "-")[:8] for r in rs]
+        row = [label, str(n), f"{sum(r['status'] == 'ok' for r in rs)}/{n}"]
+        if expect:
+            passed = sum(bool(r["passed"]) for r in rs)
+            color = "green" if passed == n else ("yellow" if passed else "red")
+            row.append(f"[{color}]{passed}/{n}[/{color}]")
+        row += [f"{sum(r['tokens'] for r in rs) / n:.0f}",
+                f"{sum(r['latency_ms'] for r in rs) / n / 1000:.1f}s",
+                str(rs[0]["result"]),
+                ", ".join(ids[:2]) + (f" +{len(ids) - 2}" if len(ids) > 2 else "")]
+        table.add_row(*row)
+    console.print(table)
+    for v in variants:
+        key, value = next(iter(v.items()))
+        if key != "model":
+            console.print(f"[bold]{labels[id(v)]}[/bold]: {value}")

agentfork/recorder.py

@@ -0,0 +1,201 @@
+import contextvars, difflib, functools, hashlib, inspect, json, time, uuid, warnings
+from . import store
+
+_run_id = contextvars.ContextVar("run_id", default=None)
+_parent_step = contextvars.ContextVar("parent_step", default=None)
+_seq = contextvars.ContextVar("seq", default=None)
+_replay = contextvars.ContextVar("replay", default=None)
+
+_LLM_DECODERS = {}
+last_run_id = None
+last_run_seq = 0
+
+class ReplayMismatch(AssertionError):
+    """The agent did something different from the recorded run."""
+
+def _decode_llm(data):
+    if isinstance(data, dict):
+        if data.get("object") == "chat.completion" and "openai" in _LLM_DECODERS:
+            return _LLM_DECODERS["openai"](data)
+        if data.get("type") == "message" and "anthropic" in _LLM_DECODERS:
+            return _LLM_DECODERS["anthropic"](data)
+    return data
+
+DECODERS = {"llm": _decode_llm}
+
+def to_jsonable(obj):
+    if hasattr(obj, "model_dump"):
+        return to_jsonable(obj.model_dump())
+    if isinstance(obj, dict):
+        return {str(k): to_jsonable(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [to_jsonable(v) for v in obj]
+    if obj is None or isinstance(obj, (str, int, float, bool)):
+        return obj
+    return repr(obj)
+
+def dumps(obj):
+    return json.dumps(to_jsonable(obj), sort_keys=True)
+
+def make_cache_key(step_type, name, input_json):
+    return hashlib.sha256(f"{step_type}|{name}|{input_json}".encode()).hexdigest()
+
+def _tokens(usage, *names):
+    for n in names:
+        value = getattr(usage, n, None)
+        if value is not None:
+            return value
+    return None
+
+def _apply_override(step_type, kwargs, override, provider=None):
+    kwargs = dict(kwargs)
+    if step_type == "llm":
+        if override.get("model"):
+            kwargs["model"] = override["model"]
+        if override.get("system"):
+            if provider == "anthropic":
+                kwargs["system"] = override["system"]
+            else:
+                msgs = [m for m in kwargs.get("messages", [])
+                        if not (isinstance(m, dict) and m.get("role") == "system")]
+                kwargs["messages"] = [{"role": "system", "content": override["system"]}] + msgs
+    return kwargs
+
+def _last_descendant_seq(steps, step):
+    ids, last = {step["id"]}, step["seq"]
+    for s in sorted(steps.values(), key=lambda s: s["seq"]):
+        if s["seq"] > step["seq"] and s["parent_id"] in ids:
+            ids.add(s["id"])
+            last = s["seq"]
+    return last
+
+def _explain_mismatch(seq, step_type, name, cached_step, input_json):
+    if cached_step["type"] != step_type or cached_step["name"] != name:
+        return (f"step {seq} was {cached_step['type']} {cached_step['name']} in the recording "
+                f"but is now {step_type} {name}")
+    old = json.dumps(json.loads(cached_step["input_json"]), indent=2, sort_keys=True).splitlines()
+    new = json.dumps(json.loads(input_json), indent=2, sort_keys=True).splitlines()
+    diff = list(difflib.unified_diff(old, new, "recorded", "now", lineterm="", n=1))[:40]
+    return f"step {seq} ({step_type} {name}) got different inputs than the recorded run\n" + "\n".join(diff)
+
+def record_step(step_type, name, fn, args, kwargs, provider=None):
+    run_id = _run_id.get()
+    if run_id is None:
+        return fn(*args, **kwargs)
+    counter = _seq.get()
+    counter[0] += 1
+    seq = counter[0]
+    replay = _replay.get()
+    strict = bool(replay and replay.get("strict"))
+    cached_step, use_forced = None, False
+    if replay and seq < replay["fork_step"]:
+        cached_step = replay["steps"].get(seq)
+        if strict and cached_step is None:
+            raise ReplayMismatch(f"step {seq} ({step_type} {name}) is new: "
+                                 f"the recorded run only had {len(replay['steps'])} steps")
+    elif replay and seq == replay["fork_step"]:
+        if "output" in replay["override"]:
+            use_forced = True
+        else:
+            kwargs = _apply_override(step_type, kwargs, replay["override"], provider)
+    if (replay and not strict and seq > replay["fork_step"] and step_type == "llm"
+            and replay["override"].get("model")):
+        kwargs = dict(kwargs, model=replay["override"]["model"])
+    if step_type == "llm" and kwargs.get("model") and cached_step is None:
+        name = kwargs["model"]
+    input_json = dumps({"args": args, "kwargs": kwargs})
+    key = make_cache_key(step_type, name, input_json)
+    if cached_step is not None and cached_step["cache_key"] != key:
+        if strict:
+            raise ReplayMismatch(_explain_mismatch(seq, step_type, name, cached_step, input_json))
+        warnings.warn(f"agentfork: step {seq} inputs differ from the original run; reusing its cached output")
+    step_id, parent = uuid.uuid4().hex, _parent_step.get()
+    token = _parent_step.set(step_id)
+    started = time.time()
+    output, error, cached = None, None, 0
+    try:
+        if cached_step is not None:
+            decode = DECODERS.get(step_type, lambda x: x)
+            output = decode(json.loads(cached_step["output_json"]))
+            counter[0] = _last_descendant_seq(replay["steps"], cached_step)
+            cached = 1
+        elif use_forced:
+            output = replay["override"]["output"]
+        else:
+            output = fn(*args, **kwargs)
+        return output
+    except Exception as e:
+        error = repr(e)
+        raise
+    finally:
+        _parent_step.reset(token)
+        usage = None if cached else getattr(output, "usage", None)
+        store.insert("steps", {
+            "id": step_id, "run_id": run_id, "parent_id": parent, "seq": seq,
+            "type": step_type, "name": name,
+            "input_json": input_json, "output_json": dumps(output), "cache_key": key,
+            "tokens_in": _tokens(usage, "prompt_tokens", "input_tokens"),
+            "tokens_out": _tokens(usage, "completion_tokens", "output_tokens"),
+            "latency_ms": (time.time() - started) * 1000,
+            "error": error, "started_at": started, "cached": cached,
+        })
+
+def trace(fn=None, *, name=None):
+    def deco(f):
+        try:
+            entry_file = inspect.getsourcefile(f)
+        except TypeError:
+            entry_file = None
+
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            global last_run_id, last_run_seq
+            if _run_id.get() is not None:
+                return f(*args, **kwargs)
+            run_id = uuid.uuid4().hex
+            replay = _replay.get() or {}
+            parent_run = replay.get("parent_run_id")
+            store.insert("runs", {
+                "id": run_id, "name": name or f.__name__, "started_at": time.time(),
+                "status": "running", "input_json": dumps({"args": args, "kwargs": kwargs}),
+                "entry_file": entry_file, "entry_name": f.__qualname__,
+                "parent_run_id": parent_run, "fork_step": replay.get("fork_step") if parent_run else None,
+            })
+            last_run_id = run_id
+            t_run, t_seq, t_parent = _run_id.set(run_id), _seq.set([0]), _parent_step.set(None)
+            status, output = "error", None
+            try:
+                output = f(*args, **kwargs)
+                status = "ok"
+                return output
+            finally:
+                last_run_seq = _seq.get()[0]
+                _run_id.reset(t_run); _seq.reset(t_seq); _parent_step.reset(t_parent)
+                store.update_run(run_id, ended_at=time.time(), status=status, output_json=dumps(output))
+                print(f"[agentfork] saved run {run_id[:8]} ({status})")
+        return wrapper
+    return deco(fn) if fn else deco
+
+def tool(fn):
+    @functools.wraps(fn)
+    def wrapper(*args, **kwargs):
+        return record_step("tool", fn.__name__, fn, args, kwargs)
+    return wrapper
+
+def wrap_openai(client):
+    from openai.types.chat import ChatCompletion
+    _LLM_DECODERS["openai"] = ChatCompletion.model_validate
+    original = client.chat.completions.create
+    def create(*args, **kwargs):
+        return record_step("llm", kwargs.get("model", "llm"), original, args, kwargs, provider="openai")
+    client.chat.completions.create = create
+    return client
+
+def wrap_anthropic(client):
+    from anthropic.types import Message
+    _LLM_DECODERS["anthropic"] = Message.model_validate
+    original = client.messages.create
+    def create(*args, **kwargs):
+        return record_step("llm", kwargs.get("model", "llm"), original, args, kwargs, provider="anthropic")
+    client.messages.create = create
+    return client

agentfork/replay.py

@@ -0,0 +1,109 @@
+import importlib.util, json, sys
+from pathlib import Path
+from . import recorder, store
+
+def load_entry(run):
+    path, qualname = run.get("entry_file"), run.get("entry_name")
+    if not path or not Path(path).exists():
+        raise RuntimeError("this run has no saved entry point; run your agent again and fork the new run")
+    folder = str(Path(path).parent)
+    if folder not in sys.path:
+        sys.path.insert(0, folder)
+    spec = importlib.util.spec_from_file_location("agentfork_entry", path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    obj = module
+    for part in qualname.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+def fork(run, step, output=None, model=None, system=None):
+    steps = {s["seq"]: s for s in store.get_steps(run["id"])}
+    if step not in steps:
+        raise ValueError(f"run {run['id'][:8]} has no step {step}")
+    kind = steps[step]["type"]
+    override = {}
+    if output is not None:
+        if kind != "tool":
+            raise ValueError("--output only works on tool steps")
+        override["output"] = output
+    if model or system:
+        if kind != "llm":
+            raise ValueError("--model and --system only work on llm steps")
+        override.update({"model": model, "system": system})
+    fn = load_entry(run)
+    inputs = json.loads(run["input_json"])
+    token = recorder._replay.set({"parent_run_id": run["id"], "fork_step": step,
+                                  "steps": steps, "override": override})
+    try:
+        result = fn(*inputs["args"], **inputs["kwargs"])
+    finally:
+        recorder._replay.reset(token)
+    return recorder.last_run_id, result
+
+def check(fixture, root="."):
+    run = fixture["run"]
+    steps = {s["seq"]: s for s in fixture["steps"]}
+    fn = load_entry(dict(run, entry_file=str(Path(root) / run["entry_file"])))
+    inputs = json.loads(run["input_json"])
+    token = recorder._replay.set({"parent_run_id": None, "fork_step": float("inf"),
+                                  "steps": steps, "override": {}, "strict": True})
+    try:
+        result = fn(*inputs["args"], **inputs["kwargs"])
+    finally:
+        recorder._replay.reset(token)
+    if recorder.last_run_seq != max(steps, default=0):
+        raise recorder.ReplayMismatch(
+            f"the recorded run had {max(steps, default=0)} steps but the replay made {recorder.last_run_seq}")
+    expected = json.loads(run["output_json"])
+    if recorder.to_jsonable(result) != expected:
+        raise recorder.ReplayMismatch(f"final output changed\n  recorded: {expected!r}\n  now:      {result!r}")
+    return result
+
+def _check_variants(run, step, variants):
+    steps = {s["seq"]: s for s in store.get_steps(run["id"])}
+    if step not in steps:
+        raise ValueError(f"run {run['id'][:8]} has no step {step}")
+    kind = steps[step]["type"]
+    for v in variants:
+        if "output" in v and kind != "tool":
+            raise ValueError("--output only works on tool steps")
+        if ("model" in v or "system" in v) and kind != "llm":
+            raise ValueError("--model and --system only work on llm steps")
+
+def _summary(steps, result, expect):
+    return {
+        "tokens": sum((s["tokens_in"] or 0) + (s["tokens_out"] or 0) for s in steps),
+        "latency_ms": sum(s["latency_ms"] or 0 for s in steps),
+        "passed": None if expect is None else expect.lower() in str(result).lower(),
+    }
+
+def baseline(run, step, expect=None):
+    steps = [s for s in store.get_steps(run["id"]) if s["seq"] >= step]
+    result = json.loads(run["output_json"]) if run["output_json"] else None
+    return dict({"variant": {"original": True}, "run_id": run["id"],
+                 "status": run["status"], "result": result}, **_summary(steps, result, expect))
+
+def sweep(run, step, variants, repeat=1, expect=None, on_result=None):
+    import contextlib, io
+    _check_variants(run, step, variants)
+    results = []
+    for variant in variants:
+        for _ in range(repeat):
+            before = recorder.last_run_id
+            status, run_id, result = "error", None, None
+            try:
+                with contextlib.redirect_stdout(io.StringIO()):
+                    run_id, result = fork(run, step, **variant)
+                status = "ok"
+            except Exception as e:
+                result = repr(e)
+                if recorder.last_run_id != before:
+                    run_id = recorder.last_run_id
+            live = [s for s in (store.get_steps(run_id) if run_id else []) if not s["cached"]]
+            record = dict({"variant": variant, "run_id": run_id, "status": status,
+                           "result": result}, **_summary(live, result, expect))
+            results.append(record)
+            if on_result:
+                on_result(record)
+    return results

agentfork/store.py

@@ -0,0 +1,70 @@
+import os, sqlite3
+from contextlib import closing
+from pathlib import Path
+
+SCHEMA = """
+CREATE TABLE IF NOT EXISTS runs (
+  id TEXT PRIMARY KEY, name TEXT, started_at REAL, ended_at REAL,
+  status TEXT, input_json TEXT, output_json TEXT,
+  parent_run_id TEXT, fork_step INTEGER, entry_file TEXT, entry_name TEXT
+);
+CREATE TABLE IF NOT EXISTS steps (
+  id TEXT PRIMARY KEY, run_id TEXT, parent_id TEXT, seq INTEGER,
+  type TEXT, name TEXT, input_json TEXT, output_json TEXT,
+  cache_key TEXT, tokens_in INTEGER, tokens_out INTEGER,
+  latency_ms REAL, error TEXT, started_at REAL, cached INTEGER DEFAULT 0
+);
+CREATE INDEX IF NOT EXISTS idx_steps_run ON steps(run_id, seq);
+"""
+
+MIGRATIONS = [
+    "ALTER TABLE runs ADD COLUMN entry_file TEXT",
+    "ALTER TABLE runs ADD COLUMN entry_name TEXT",
+    "ALTER TABLE steps ADD COLUMN cached INTEGER DEFAULT 0",
+]
+
+def db_path():
+    return Path(os.environ.get("AGENTFORK_DIR", ".agentfork")) / "traces.db"
+
+def connect():
+    path = db_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(path)
+    conn.executescript(SCHEMA)
+    for sql in MIGRATIONS:
+        try:
+            conn.execute(sql)
+        except sqlite3.OperationalError:
+            pass
+    return conn
+
+def _exec(sql, params):
+    with closing(connect()) as conn, conn:
+        conn.execute(sql, params)
+
+def insert(table, row):
+    cols = ", ".join(row)
+    marks = ", ".join("?" for _ in row)
+    _exec(f"INSERT INTO {table} ({cols}) VALUES ({marks})", list(row.values()))
+
+def update_run(run_id, **fields):
+    sets = ", ".join(f"{k}=?" for k in fields)
+    _exec(f"UPDATE runs SET {sets} WHERE id=?", [*fields.values(), run_id])
+
+def query(sql, params=()):
+    with closing(connect()) as conn:
+        conn.row_factory = sqlite3.Row
+        return [dict(r) for r in conn.execute(sql, params)]
+
+def list_runs(limit=10):
+    return query("""
+        SELECT r.*, COUNT(s.id) AS n_steps,
+               COALESCE(SUM(s.tokens_in), 0) + COALESCE(SUM(s.tokens_out), 0) AS tokens
+        FROM runs r LEFT JOIN steps s ON s.run_id = r.id
+        GROUP BY r.id ORDER BY r.started_at DESC LIMIT ?""", (limit,))
+
+def find_runs(prefix):
+    return query("SELECT * FROM runs WHERE id LIKE ? ORDER BY started_at DESC", (prefix + "%",))
+
+def get_steps(run_id):
+    return query("SELECT * FROM steps WHERE run_id = ? ORDER BY seq", (run_id,))

agentfork_cli-0.1.0.dist-info/METADATA

@@ -0,0 +1,125 @@
+Metadata-Version: 2.4
+Name: agentfork-cli
+Version: 0.1.0
+Summary: Fork your LLM agent at any step and sweep prompts, models or tool results to find what fixes it.
+Project-URL: Homepage, https://github.com/bimal1023/agentfork
+Project-URL: Issues, https://github.com/bimal1023/agentfork/issues
+Author: Bimal Kumal
+License-Expression: MIT
+License-File: LICENSE
+Keywords: agents,anthropic,debugging,llm,openai,replay,testing
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Software Development :: Debuggers
+Classifier: Topic :: Software Development :: Testing
+Requires-Python: >=3.9
+Requires-Dist: rich
+Requires-Dist: typer
+Provides-Extra: anthropic
+Requires-Dist: anthropic; extra == 'anthropic'
+Provides-Extra: dev
+Requires-Dist: anthropic; extra == 'dev'
+Requires-Dist: openai; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Provides-Extra: openai
+Requires-Dist: openai; extra == 'openai'
+Description-Content-Type: text/markdown
+
+# agentfork
+
+[![tests](https://github.com/bimal1023/agentfork/actions/workflows/tests.yml/badge.svg)](https://github.com/bimal1023/agentfork/actions/workflows/tests.yml)
+
+**Fork your LLM agent at any step, then sweep prompts, models or tool results to find what actually fixes it.**
+
+![agentfork demo](https://raw.githubusercontent.com/bimal1023/agentfork/main/docs/demo.gif)
+
+When an agent misbehaves at step 3, the usual loop is: tweak the prompt, re-run everything, eyeball one output, repeat. That is slow, costs tokens on steps that didn't change, and one lucky run proves nothing. agentfork records every LLM and tool call, replays earlier steps from cache, and runs **many variants of one step, several times each**, so you get pass rates instead of anecdotes.
+
+Everything runs locally. No account, no server, no data leaves your machine.
+
+## The headline example
+
+A refund agent reads `"amount": 4999, "currency": "USD"` from an order API. The amount is in cents, and the agent refunds **$4,999** for a $49.99 laptop ([examples/refund_agent_cents.py](examples/refund_agent_cents.py)). Is the fix a better model or a better prompt?
+
+```text
+$ agentfork sweep de994fe6 --step 3 --repeat 3 --expect "49.99" --model gpt-4o-mini,gpt-4o
+
+variant             runs  pass  avg tokens  sample output
+original            1     0/1   445         I have processed a refund of $4,999.00...
+model=gpt-4o-mini   3     0/3   444         I have processed a refund of $4,999.00...
+model=gpt-4o        3     0/3   292         I apologize, it seems there was a mis...
+```
+
+A bigger model doesn't fix it. One sentence in the prompt does:
+
+```text
+$ agentfork sweep de994fe6 --step 3 --repeat 3 --expect "49.99" \
+    --system "You are a support agent. Look up the order, then refund the full order amount in USD." \
+    --system "... Order amounts from the API are in cents; divide by 100 before refunding."
+
+variant      runs  pass  avg tokens  sample output
+original     1     0/1   445         I have processed a refund of $4,999.00...
+system #1    3     0/3   445         A full refund of $4,999.00 has ...
+system #2    3     3/3   460         I have processed a refund of $49.99...
+```
+
+Steps 1 and 2 replayed from cache in every run, so only the step that mattered was paid for.
+
+## Quickstart
+
+```bash
+pip install "agentfork-cli[openai]"      # or agentfork[anthropic]
+```
+
+Add three things to your agent:
+
+```python
+from openai import OpenAI
+from agentfork import trace, tool, wrap_openai
+
+client = wrap_openai(OpenAI())      # records every LLM call (wrap_anthropic for Claude)
+
+@tool                               # records every tool call
+def lookup_order(order_id): ...
+
+@trace                              # one run = one call of your agent
+def run_agent(question): ...
+```
+
+Run your agent as usual. Each run is saved to `.agentfork/traces.db`.
+
+## Commands
+
+| Command | What it does |
+| --- | --- |
+| `agentfork list` | Recent runs with steps, tokens, duration and fork lineage |
+| `agentfork show <id> [--full]` | A run as a step tree, with optional full JSON |
+| `agentfork fork <id> --step N --output/--model/--system ...` | Re-run from step N with one change; earlier steps come from cache |
+| `agentfork sweep <id> --step N --model a,b --system ... --repeat K --expect TEXT` | Try many variants of step N, K times each, and compare pass rate, tokens and time |
+| `agentfork diff <a> <b>` | Compare two runs step by step |
+| `agentfork save-test <id>` | Save a run as a pytest test that replays from cache in CI with no API key |
+
+## How it works
+
+1. **Record.** `@trace`, `@tool` and `wrap_openai`/`wrap_anthropic` capture each step's inputs, outputs, tokens and latency, plus a hash of its inputs.
+2. **Fork.** agentfork re-imports your agent and re-runs it. Steps before N are served from the cache, step N gets your change, and everything after runs live. A model change applies to every LLM call from step N on.
+3. **Sweep.** Runs fork for each variant `--repeat` times and checks each final output for your `--expect` text.
+4. **Test.** `save-test` writes a JSON snapshot plus a pytest file. The test replays every step from cache and fails with a diff if any step's inputs, the step count, or the final output changes.
+
+## Related work
+
+agentfork overlaps with record/replay tools like [agent-rewind](https://pypi.org/project/agent-rewind/), which supports many more providers, and [agentstep](https://pypi.org/project/agentstep/), a LangGraph time-travel debugger with a web UI. agentfork's focus is the step after replay: sweeping many variants of one step, with repeats and pass rates, to choose a fix with evidence.
+
+## Roadmap
+
+- [x] Recorder, CLI viewer, fork, diff, snapshot tests
+- [x] OpenAI and Anthropic support
+- [x] Sweep with repeats and pass rates
+- [ ] Cost estimates per variant
+- [ ] Parallel sweeps
+- [ ] LangGraph integration
+
+## License
+
+MIT

agentfork_cli-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+agentfork/__init__.py,sha256=UhSuHu7h_AokisRv3K_gsu4UgxaCoarehiLutha5c7E,125
+agentfork/cli.py,sha256=J5Vr_OxX-OeAzpcnHbfP_xB9j9wINw3HNqwvTfr0AzA,12479
+agentfork/recorder.py,sha256=G-dnnxOy9ZcggK-yeADY_V0Nb5tY_unDscEgLh4zuH8,8528
+agentfork/replay.py,sha256=G3btKJRhtrK02M7ltDp1DFK1_JDUkMe-5pa7RqWe9hw,4802
+agentfork/store.py,sha256=8eLc6WXJGDT0aUY1rpwm5L9aSHllPk22Rg8Ze4WCnFM,2374
+agentfork_cli-0.1.0.dist-info/METADATA,sha256=XgKdjBM7S1heJmRXLkWWgwYrmRq6b4RlkN_xZ-_Wgxg,5918
+agentfork_cli-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+agentfork_cli-0.1.0.dist-info/entry_points.txt,sha256=o3dJQFfMtuIfx3v1dA0_raUOmVSMmcMsrLR6Sm0aRJ4,48
+agentfork_cli-0.1.0.dist-info/licenses/LICENSE,sha256=VugILMcFud9WqPBP_SD92VRzT1ClG_0C6hhfJNIjeAs,1068
+agentfork_cli-0.1.0.dist-info/RECORD,,

agentfork_cli-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

agentfork_cli-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+agentfork = agentfork.cli:app

agentfork_cli-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Bimal Kumal
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.