"""frameworks_toolkit.py - a tiny toolkit for Week 3b (Frameworks & Orchestration). Standard library only. No API key. Every "language model" here is a mock. Everything is made up for teaching: the tickets, the tasks, the prices and the slip rates. Parts, in order: 1. A DSPy-like mini library: Signature, Predict, ChainOfThought, BestOfN, a metric, evaluate, and the optimizers LabeledFewShot, BootstrapFewShot and random search. 2. A LangGraph-like StateGraph: typed state with reducers, nodes, conditional edges, checkpoints, and interrupt / resume for a human in the loop. 3. A mini-SWE-agent-style harness: one bash action per step, a temp directory, an allowlist, a linear history, an event log, a stuck detector, and a tree of messages (Pi-style). 4. A Temporal-style durable runtime: workflows, activities, an event history, retries, timers, signals, crash and deterministic replay. 5. Tracing: spans with OpenTelemetry-GenAI-like attributes, token and cost accounting, JSONL export, a text trace viewer, a bootstrap confidence interval and a Pareto front. Run `python frameworks_toolkit.py` for the demo. Tested on Python 3.13 only. """ import contextlib import copy import hashlib import itertools import json import math import operator import os import random import re import shlex import subprocess import tempfile import threading import time from dataclasses import dataclass, field from typing import (Annotated, Any, Callable, Dict, Iterable, List, Literal, Optional, Tuple, get_args, get_origin, get_type_hints) # ===================================================================================== # 1. A DSPy-like mini library # ===================================================================================== _TYPES = {"str": str, "int": int, "float": float, "bool": bool} _IDENT = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") class SignatureError(ValueError): """The signature string is not valid.""" class ParseError(ValueError): """The model's reply could not be read as the output fields.""" @dataclass(frozen=True) class Fld: name: str type: str = "str" class Signature: """Inputs, outputs and instructions. 'question -> answer' is the smallest example.""" def __init__(self, inputs: Tuple[Fld, ...], outputs: Tuple[Fld, ...], instructions: str = ""): self.inputs, self.outputs = tuple(inputs), tuple(outputs) names = [f.name for f in self.inputs + self.outputs] if not self.inputs or not self.outputs: raise SignatureError("a signature needs at least one input and one output") if len(set(names)) != len(names): raise SignatureError("field names must be unique: %r" % names) self.instructions = instructions or "Given the fields %s, produce the fields %s." % ( ", ".join(f.name for f in self.inputs), ", ".join(f.name for f in self.outputs)) @staticmethod def _side(text: str) -> Tuple[Fld, ...]: out = [] for item in text.split(","): item = item.strip() name, _, typ = item.partition(":") name, typ = name.strip(), (typ.strip() or "str") if not _IDENT.match(name): raise SignatureError("bad field name %r" % name) if typ not in _TYPES: raise SignatureError("unknown type %r (use one of %s)" % (typ, sorted(_TYPES))) out.append(Fld(name, typ)) return tuple(out) @classmethod def parse(cls, spec: str, instructions: str = "") -> "Signature": if spec.count("->") != 1: raise SignatureError("a signature needs exactly one '->': %r" % spec) left, right = spec.split("->") return cls(cls._side(left), cls._side(right), instructions) def prepend_output(self, name: str, type: str = "str") -> "Signature": """What ChainOfThought does: put a new output field first.""" return Signature(self.inputs, (Fld(name, type),) + self.outputs, self.instructions) def __str__(self): side = lambda fs: ", ".join(f.name if f.type == "str" else "%s: %s" % (f.name, f.type) for f in fs) return "%s -> %s" % (side(self.inputs), side(self.outputs)) def _coerce(text: str, typ: str): text = text.strip() try: if typ == "int": return int(text) if typ == "float": return float(text) if typ == "bool": low = text.lower() if low in ("true", "yes"): return True if low in ("false", "no"): return False raise ValueError(text) except ValueError: raise ParseError("cannot read %r as %s" % (text, typ)) return text class Example: """One data point. with_inputs() says which keys are inputs; the rest are labels.""" def __init__(self, **kw): self._d, self._inputs = dict(kw), () def with_inputs(self, *keys) -> "Example": e = Example(**self._d) e._inputs = tuple(keys) return e def inputs(self) -> Dict[str, Any]: return {k: self._d[k] for k in self._inputs} def labels(self) -> Dict[str, Any]: return {k: v for k, v in self._d.items() if k not in self._inputs} def __getattr__(self, k): if k.startswith("_"): raise AttributeError(k) try: return self._d[k] except KeyError: raise AttributeError(k) def toDict(self): return dict(self._d) class Prediction(Example): """The result of a module call. Fields are attributes.""" class LM: """The smallest LM interface: text in, text out. Subclasses count their calls.""" def __init__(self): self.calls = 0 self.history: List[str] = [] def __call__(self, prompt: str) -> str: raise NotImplementedError class _Settings: lm: Optional[LM] = None settings = _Settings() def configure(lm: Optional[LM] = None): settings.lm = lm _trace_stack: List[list] = [] # single-threaded on purpose def render_prompt(sig: Signature, demos: List[dict], inputs: Dict[str, Any]) -> str: """The text a Predict sends. The layout is my own; real DSPy adapters differ.""" lines = ["INSTRUCTIONS: " + sig.instructions, "OUTPUT FIELDS: " + ", ".join(f.name for f in sig.outputs), 'Reply with one line per output field, like "name: value".'] for d in demos: lines.append("--- demo") lines += ["%s: %s" % (f.name, d[f.name]) for f in sig.inputs + sig.outputs if f.name in d] lines.append("--- input") lines += ["%s: %s" % (f.name, inputs[f.name]) for f in sig.inputs] return "\n".join(lines) def parse_reply(sig: Signature, text: str) -> Dict[str, Any]: found = {} for line in text.splitlines(): name, sep, value = line.partition(":") if sep and name.strip() in {f.name for f in sig.outputs} and name.strip() not in found: found[name.strip()] = value.strip() missing = [f.name for f in sig.outputs if f.name not in found] if missing: raise ParseError("reply is missing %s: %r" % (missing, text[:60])) return {f.name: _coerce(found[f.name], f.type) for f in sig.outputs} class Module: """Base class. A program is a Module that holds Predict modules and plain Python.""" def named_predictors(self, prefix: str = "") -> List[Tuple[str, "Predict"]]: out = [] for k, v in sorted(vars(self).items()): if isinstance(v, Predict): out.append((prefix + k, v)) elif isinstance(v, Module): out += v.named_predictors(prefix + k + ".") return out def set_lm(self, lm: Optional[LM]): for _, p in self.named_predictors(): p.lm = lm def reset_demos(self): for _, p in self.named_predictors(): p.demos = [] def deepcopy(self): return copy.deepcopy(self) def __deepcopy__(self, memo): new = object.__new__(type(self)) memo[id(self)] = new for k, v in vars(self).items(): setattr(new, k, copy.deepcopy(v, memo)) return new def __call__(self, **kw): return self.forward(**kw) class Predict(Module): """One LM call that implements a signature. Holds demos (few-shot examples).""" def __init__(self, signature, lm: Optional[LM] = None): self.signature = Signature.parse(signature) if isinstance(signature, str) else signature self.lm, self.demos = lm, [] def __deepcopy__(self, memo): # share the LM, copy the demos new = object.__new__(Predict) memo[id(self)] = new new.signature, new.lm, new.demos = self.signature, self.lm, copy.deepcopy(self.demos, memo) return new def named_predictors(self, prefix=""): return [(prefix.rstrip(".") or "self", self)] def forward(self, **kw) -> Prediction: names = [f.name for f in self.signature.inputs] if sorted(kw) != sorted(names): raise TypeError("expected inputs %s, got %s" % (names, sorted(kw))) lm = self.lm or settings.lm if lm is None: raise RuntimeError("no LM: call configure(lm=...) or set_lm(...)") reply = lm(render_prompt(self.signature, self.demos, kw)) out = parse_reply(self.signature, reply) if _trace_stack: _trace_stack[-1].append((self, dict(kw), dict(out))) return Prediction(**kw, **out) class ChainOfThought(Module): """Predict with a 'reasoning' output field put first. The mock LM gets no benefit from it.""" def __init__(self, signature, lm: Optional[LM] = None): sig = Signature.parse(signature) if isinstance(signature, str) else signature self.predict = Predict(sig.prepend_output("reasoning"), lm) def forward(self, **kw) -> Prediction: return self.predict(**kw) class BestOfN(Module): """Run a module up to n times and keep the best reward. Stop early at the threshold.""" def __init__(self, module: Module, n: int, reward_fn: Callable, threshold: float = 1.0): if n < 1: raise ValueError("n must be at least 1") self.module, self.n, self.reward_fn, self.threshold = module, n, reward_fn, threshold def forward(self, **kw) -> Prediction: best, best_r, last_err = None, -math.inf, None for _ in range(self.n): try: pred = self.module(**kw) except ParseError as e: last_err = e continue r = self.reward_fn(kw, pred) if r > best_r: best, best_r = pred, r if best_r >= self.threshold: break if best is None: raise last_err return best # ---- made-up data: support tickets ----------------------------------------------- _TICKETS = { "billing": [("I was charged twice for my {plan} plan this month", "high"), ("please refund invoice {num}, the amount is wrong", "low"), ("my card was billed but I cancelled the {plan} plan", "high"), ("why is there an extra charge on invoice {num}", "low")], "shipping": [("my package {num} has not arrived and tracking is stuck", "high"), ("the delivery of order {num} is late", "low"), ("tracking says delivered but the package is missing", "high"), ("can I change the delivery address for order {num}", "low")], "login": [("I cannot log in, the password reset email never came", "high"), ("my account is locked after too many login tries", "high"), ("the two factor code fails when I log in to the {plan} account", "high"), ("reset my password, I forgot it", "low")], "feedback": [("please add dark mode to the app", "low"), ("the new {plan} dashboard is great, thanks", "low"), ("feature idea: export reports as a PDF", "low"), ("I suggest a shortcut key for search", "low")], } CATEGORIES = sorted(_TICKETS) def make_tickets(n: int, seed: int = 0) -> List[Example]: """Made-up tickets. Train and dev share templates, so dev scores are optimistic.""" rng = random.Random(seed) out = [] for i in range(n): cat = CATEGORIES[i % len(CATEGORIES)] text, urgency = rng.choice(_TICKETS[cat]) text = text.format(plan=rng.choice(["basic", "pro", "team"]), num=rng.choice(["4821", "1177", "9034", "3350"])) out.append(Example(ticket=text, category=cat, urgency=urgency).with_inputs("ticket")) rng.shuffle(out) return out _STOP = set("a an the i my to of for is it in on at and or but not be was can this that with as me we you your are".split()) def _words(text: str) -> set: return {w for w in re.findall(r"[a-z0-9]+", text.lower()) if w not in _STOP} def overlap(a: str, b: str) -> float: """Jaccard overlap of word sets, ignoring stop words.""" wa, wb = _words(a), _words(b) return len(wa & wb) / len(wa | wb) if wa | wb else 0.0 class ImitatorLM(LM): """A mock LM that imitates the demos in its prompt. With no demos it guesses each field from `priors`, and with probability `p_format_slip` it forgets the "name: value" format. With demos it never slips on format. It copies the output of the demo whose inputs overlap most with the new input, if the overlap is at least `min_overlap`. These rules are my own simplification of "demos help a small model". """ def __init__(self, priors: Dict[str, List[str]], seed: int = 0, p_format_slip: float = 0.3, min_overlap: float = 0.2, model: str = "mock-imitator"): super().__init__() self.priors, self.rng, self.p_slip, self.min_overlap, self.model = priors, random.Random(seed), p_format_slip, min_overlap, model def __call__(self, prompt: str) -> str: self.calls += 1 self.history.append(prompt) head, _, tail = prompt.partition("--- input\n") fields = re.search(r"OUTPUT FIELDS: (.*)", head).group(1).split(", ") demos = [] for block in head.split("--- demo\n")[1:]: demos.append(dict(l.split(": ", 1) for l in block.split("\n") if ": " in l and not l.startswith(("OUTPUT", "Reply", "INSTR")))) new = dict(l.split(": ", 1) for l in tail.split("\n") if ": " in l) if not demos and self.rng.random() < self.p_slip: return "I think the answer is probably %s." % self.rng.choice(self.priors.get(fields[-1], ["unknown"])) lines = [] for f in fields: if f == "reasoning": lines.append("reasoning: I looked at the key words and picked the closest known case.") continue best, best_o = None, -1.0 for d in demos: if f not in d: continue o = overlap(" ".join(new.values()), " ".join(v for k, v in d.items() if k in new)) if o > best_o: best, best_o = d, o if best is not None and best_o >= self.min_overlap: lines.append("%s: %s" % (f, best[f])) else: lines.append("%s: %s" % (f, self.rng.choice(self.priors.get(f, ["unknown"])))) return "\n".join(lines) def make_lm(seed: int = 0, p_format_slip: float = 0.3) -> ImitatorLM: return ImitatorLM({"category": CATEGORIES, "urgency": ["low", "high"]}, seed=seed, p_format_slip=p_format_slip) class Classify(Module): """One stage: ticket -> category.""" def __init__(self): self.predict = Predict("ticket -> category", None) def forward(self, ticket): return self.predict(ticket=ticket) class Triage(Module): """Two stages: ticket -> category, then (ticket, category) -> urgency. Plain Python glue between.""" def __init__(self): self.category = Predict("ticket -> category") self.urgency = Predict("ticket, category -> urgency") def forward(self, ticket): c = self.category(ticket=ticket).category u = self.urgency(ticket=ticket, category=c).urgency return Prediction(ticket=ticket, category=c, urgency=u) def category_match(example, pred, trace=None) -> bool: return getattr(pred, "category", None) == example.category def triage_match(example, pred, trace=None) -> bool: return pred.category == example.category and pred.urgency == example.urgency @dataclass class EvalResult: score: float n: int failures: int # examples where the program raised lm_calls: int rows: List[Tuple[bool, str]] def evaluate(program: Module, devset: List[Example], metric: Callable, lm: Optional[LM] = None) -> EvalResult: """Mean metric over a dev set. An exception counts as a miss, like DSPy's Evaluate.""" lm = lm or settings.lm before = lm.calls if lm else 0 hits, fails, rows = 0, 0, [] for ex in devset: try: ok = bool(metric(ex, program(**ex.inputs()))) rows.append((ok, "")) except Exception as e: # ParseError and friends ok = False fails += 1 rows.append((False, type(e).__name__)) hits += ok return EvalResult(hits / len(devset) if devset else 0.0, len(devset), fails, (lm.calls - before) if lm else 0, rows) def _label_demo(pred: Predict, ex: Example) -> Optional[dict]: d = dict(ex.toDict()) need = [f.name for f in pred.signature.inputs + pred.signature.outputs] return {k: d[k] for k in need} if all(k in d for k in need) else None class LabeledFewShot: """Put k labeled training examples into every predictor whose fields the labels cover.""" def __init__(self, k: int = 4, seed: int = 0): self.k, self.seed = k, seed def compile(self, student: Module, trainset: List[Example]) -> Module: student = student.deepcopy() rng = random.Random(self.seed) pool = list(trainset) rng.shuffle(pool) for _, p in student.named_predictors(): demos = [d for d in (_label_demo(p, e) for e in pool) if d] p.demos = demos[: self.k] return student class BootstrapFewShot: """Run a teacher on training inputs. Keep traces whose final output passes the metric. The kept (input, output) pairs of every predictor become demos for that predictor. `max_rounds` retries the examples that failed. `max_labeled_demos` adds labeled examples after the bootstrapped ones. The teacher defaults to the student with no demos. """ def __init__(self, metric: Callable, max_bootstrapped_demos: int = 4, max_labeled_demos: int = 0, max_rounds: int = 1, seed: int = 0): self.metric, self.max_b, self.max_l, self.rounds, self.seed = metric, max_bootstrapped_demos, max_labeled_demos, max_rounds, seed def compile(self, student: Module, trainset: List[Example], teacher: Optional[Module] = None) -> Module: student = student.deepcopy() student.reset_demos() teacher = teacher.deepcopy() if teacher is not None else student.deepcopy() pool = list(trainset) random.Random(self.seed).shuffle(pool) t_preds = dict(teacher.named_predictors()) kept: Dict[str, List[dict]] = {name: [] for name in t_preds} used, done = set(), 0 for _ in range(self.rounds): for i, ex in enumerate(pool): if done >= self.max_b: break if i in used: continue _trace_stack.append([]) try: ok = bool(self.metric(ex, teacher(**ex.inputs()))) except Exception: ok = False finally: trace = _trace_stack.pop() if ok and trace: used.add(i) done += 1 names = {id(p): n for n, p in t_preds.items()} for p, inp, out in trace: kept[names[id(p)]].append({**inp, **out}) for name, p in student.named_predictors(): p.demos = kept.get(name, [])[: self.max_b] if self.max_l: extra = [d for d in (_label_demo(p, pool[i]) for i in range(len(pool)) if i not in used) if d] p.demos += extra[: self.max_l] student._bootstrapped = done # how many training examples produced a demo return student class RandomSearch: """Try candidate programs, score each on valset, keep the best (earliest wins ties). Candidates: the program with no demos, labeled few-shot, then `num_candidates` bootstraps with different shuffles. This is a simplified take on DSPy's BootstrapFewShotWithRandomSearch. """ def __init__(self, metric: Callable, num_candidates: int = 4, max_bootstrapped_demos: int = 4, max_labeled_demos: int = 0, max_rounds: int = 2, seed: int = 0): self.metric, self.n, self.max_b, self.max_l, self.rounds, self.seed = metric, num_candidates, max_bootstrapped_demos, max_labeled_demos, max_rounds, seed self.scores: List[float] = [] def compile(self, student: Module, trainset: List[Example], valset: Optional[List[Example]] = None) -> Module: valset = valset or trainset zero = student.deepcopy() zero.reset_demos() cands = [zero, LabeledFewShot(self.max_b, self.seed).compile(student, trainset)] for i in range(self.n): cands.append(BootstrapFewShot(self.metric, self.max_b, self.max_l, self.rounds, self.seed + i + 1).compile(student, trainset)) self.scores = [evaluate(c, valset, self.metric).score for c in cands] return cands[max(range(len(cands)), key=lambda i: (self.scores[i], -i))] # ===================================================================================== # 2. A LangGraph-like StateGraph # ===================================================================================== START, END = "__start__", "__end__" class GraphRecursionError(RuntimeError): """The graph ran more super-steps than the limit.""" class InvalidUpdateError(ValueError): """A node returned a bad update (unknown key, wrong type, or two writes to a key with no reducer).""" class GraphInterrupt(Exception): """Raised inside a node by interrupt(). The graph catches it and saves a checkpoint.""" def __init__(self, interrupt): super().__init__(interrupt) self.interrupt = interrupt @dataclass class Interrupt: value: Any node: str @dataclass class Command: """Pass Command(resume=value) to invoke() to continue a paused run.""" resume: Any = None class _ICtx: def __init__(self, node, values): self.node, self.values, self.counter = node, list(values), 0 _ictx: List[_ICtx] = [] def interrupt(value: Any) -> Any: """Pause the graph and show `value` to the caller. On resume, return the resume value. The node restarts from its first line on resume, so code before interrupt() runs again. """ if not _ictx: raise RuntimeError("interrupt() can only be called inside a graph node") ctx = _ictx[-1] i, ctx.counter = ctx.counter, ctx.counter + 1 if i < len(ctx.values): return ctx.values[i] raise GraphInterrupt(Interrupt(value, ctx.node)) def _empty(tp): origin = get_origin(tp) or tp return origin() if origin in (list, dict, set, str, int, float) else None class InMemorySaver: """Checkpoints in RAM, per thread. Lost when the process ends.""" def __init__(self): self.data: Dict[str, List[dict]] = {} def put(self, thread: str, ckpt: dict): self.data.setdefault(thread, []).append(copy.deepcopy(ckpt)) def list(self, thread: str) -> List[dict]: return copy.deepcopy(self.data.get(thread, [])) def latest(self, thread: str) -> Optional[dict]: items = self.data.get(thread) return copy.deepcopy(items[-1]) if items else None class FileSaver(InMemorySaver): """Checkpoints appended to a JSON Lines file, so a new process can resume (state must be JSON).""" def __init__(self, path: str): super().__init__() self.path = path def put(self, thread: str, ckpt: dict): with open(self.path, "a", encoding="utf-8") as f: f.write(json.dumps({"thread": thread, "ckpt": ckpt}) + "\n") def list(self, thread: str) -> List[dict]: if not os.path.exists(self.path): return [] with open(self.path, encoding="utf-8") as f: rows = [json.loads(l) for l in f if l.strip()] return [r["ckpt"] for r in rows if r["thread"] == thread] def latest(self, thread: str) -> Optional[dict]: items = self.list(thread) return items[-1] if items else None @dataclass class Snapshot: values: dict next: List[str] step: int status: str # "running", "interrupted" or "done" class StateGraph: """Build a graph of nodes that read and write one shared state. The schema is a TypedDict. A key typed Annotated[list, operator.add] has a reducer that merges writes. Keys with no reducer are replaced. """ def __init__(self, schema): hints = get_type_hints(schema, include_extras=True) self.keys, self.reducers, self.types = list(hints), {}, {} for k, tp in hints.items(): if get_origin(tp) is Annotated: base, *meta = get_args(tp) self.types[k] = base fns = [m for m in meta if callable(m)] if fns: self.reducers[k] = fns[0] else: self.types[k] = tp self.nodes: Dict[str, Callable] = {} self.edges: List[Tuple[str, str]] = [] self.cond: Dict[str, Tuple[Callable, Optional[dict]]] = {} def add_node(self, name: str, fn: Callable) -> "StateGraph": if name in (START, END) or name in self.nodes: raise ValueError("bad or duplicate node name %r" % name) self.nodes[name] = fn return self def _check(self, a, b=None): for n, allowed in ((a, {START}), (b, {END})): if n is not None and n not in self.nodes and n not in allowed: raise ValueError("unknown node %r" % n) def add_edge(self, src: str, dst: str) -> "StateGraph": self._check(src, dst) self.edges.append((src, dst)) return self def add_conditional_edges(self, src: str, router: Callable, path_map: Optional[dict] = None) -> "StateGraph": self._check(src) if src in self.cond: raise ValueError("%r already has conditional edges" % src) self.cond[src] = (router, path_map) return self def compile(self, checkpointer=None, recursion_limit: int = 25) -> "CompiledGraph": if not any(a == START for a, _ in self.edges) and START not in self.cond: raise ValueError("the graph has no edge from START") for n in self.nodes: if not any(a == n for a, _ in self.edges) and n not in self.cond: raise ValueError("node %r has no outgoing edge" % n) return CompiledGraph(self, checkpointer, recursion_limit) class CompiledGraph: def __init__(self, g: StateGraph, checkpointer, recursion_limit): self.g, self.saver, self.limit = g, checkpointer, recursion_limit self.events: List[Tuple[str, str]] = [] # ("run", node) / ("skip", node): for tests and demos # -- state updates ---------------------------------------------------------------- def _apply(self, state: dict, writes: List[Tuple[str, dict]]) -> dict: state = dict(state) per_key: Dict[str, List[Any]] = {} for node, upd in writes: for k, v in upd.items(): per_key.setdefault(k, []).append(v) for k, vals in per_key.items(): red = self.g.reducers.get(k) if red is None and len(vals) > 1: raise InvalidUpdateError("key %r got %d writes in one step and has no reducer" % (k, len(vals))) for v in vals: if red is None: state[k] = v else: base = state[k] if k in state else _empty(self.g.types[k]) state[k] = v if base is None else red(base, v) return state def _check_update(self, node: str, upd) -> dict: if upd is None: return {} if not isinstance(upd, dict): raise InvalidUpdateError("node %r must return a dict or None, got %s" % (node, type(upd).__name__)) bad = [k for k in upd if k not in self.g.keys] if bad: raise InvalidUpdateError("node %r wrote unknown keys %s" % (node, bad)) return upd def _targets(self, ran: Iterable[str], state: dict) -> List[str]: out: List[str] = [] for n in ran: dsts = [b for a, b in self.g.edges if a == n] if n in self.g.cond: router, pmap = self.g.cond[n] r = router(copy.deepcopy(state)) for t in (r if isinstance(r, (list, tuple)) else [r]): t = pmap[t] if pmap and t in pmap else t if t != END and t not in self.g.nodes: raise ValueError("router of %r returned unknown target %r" % (n, t)) dsts.append(t) for d in dsts: if d != END and d not in out: out.append(d) return out # -- running ---------------------------------------------------------------------- def _save(self, thread, ckpt): if self.saver and thread is not None: self.saver.put(thread, ckpt) def invoke(self, input: Any = None, config: Optional[dict] = None) -> dict: cfg = config or {} thread = (cfg.get("configurable") or {}).get("thread_id") limit = cfg.get("recursion_limit", self.limit) latest = self.saver.latest(thread) if (self.saver and thread is not None) else None if isinstance(input, Command): if not latest or latest["status"] != "interrupted": raise ValueError("nothing to resume on thread %r" % thread) ck = latest node = ck["interrupt"]["node"] ck["resume"].setdefault(node, []).append(input.resume) ck["status"], ck["interrupt"] = "running", None elif input is None: if not latest or latest["status"] == "done": raise ValueError("no unfinished run on thread %r" % thread) if latest["status"] == "interrupted": raise ValueError("thread %r is waiting: pass Command(resume=...)" % thread) ck = latest else: if latest and latest["status"] == "interrupted": raise ValueError("thread %r is waiting: pass Command(resume=...)" % thread) base = latest["state"] if latest else {} state = self._apply(base, [(START, self._check_update(START, dict(input)))]) ck = {"step": latest["step"] if latest else 0, "state": state, "next": [], "writes": {}, "resume": {}, "status": "running", "interrupt": None} ck["next"] = self._targets([START], state) self._save(thread, ck) return self._loop(thread, ck, limit) def _loop(self, thread, ck, limit) -> dict: run_steps = 0 while ck["next"]: if run_steps >= limit: raise GraphRecursionError("hit the limit of %d super-steps" % limit) writes: List[Tuple[str, dict]] = [] for node in sorted(ck["next"]): if node in ck["writes"]: writes.append((node, ck["writes"][node])) self.events.append(("skip", node)) continue _ictx.append(_ICtx(node, ck["resume"].get(node, []))) try: self.events.append(("run", node)) upd = self._check_update(node, self.g.nodes[node](copy.deepcopy(ck["state"]))) except GraphInterrupt as gi: ck["status"] = "interrupted" ck["interrupt"] = {"node": node, "value": gi.interrupt.value} self._save(thread, ck) return {**ck["state"], "__interrupt__": [gi.interrupt]} except Exception: self._save(thread, ck) # keep the writes of nodes that already finished raise finally: _ictx.pop() ck["writes"][node] = upd writes.append((node, upd)) state = self._apply(ck["state"], writes) ran = sorted(ck["next"]) ck = {"step": ck["step"] + 1, "state": state, "next": self._targets(ran, state), "writes": {}, "resume": {}, "status": "running", "interrupt": None} run_steps += 1 if not ck["next"]: ck["status"] = "done" self._save(thread, ck) ck["status"] = "done" return ck["state"] # -- inspection --------------------------------------------------------------------- def get_state(self, config: dict) -> Snapshot: ck = self.saver.latest(config["configurable"]["thread_id"]) if ck is None: raise KeyError("no checkpoint for that thread") return Snapshot(ck["state"], ck["next"], ck["step"], ck["status"]) def history(self, config: dict) -> List[dict]: return self.saver.list(config["configurable"]["thread_id"]) def to_mermaid(self) -> str: lines = ["graph TD"] for a, b in self.g.edges: lines.append(" %s --> %s" % (a, b)) for a, (router, pmap) in sorted(self.g.cond.items()): targets = list(pmap.values()) if pmap else [] if not targets: lit = get_type_hints(router).get("return") targets = list(get_args(lit)) if lit is not None and get_origin(lit) is not None else [] for t in targets: lines.append(" %s -.-> %s" % (a, t)) return "\n".join(lines) # ===================================================================================== # 3. A mini-SWE-agent-style harness # ===================================================================================== DEFAULT_ALLOW = ("ls", "cat", "echo", "grep", "head", "tail", "wc", "mkdir", "touch", "cp", "mv", "sort", "pwd", "printf", "diff") DONE_MARK = "TASK_COMPLETE" # my own marker for "I am finished" _FENCE = re.compile(r"```bash\n(.*?)\n```", re.S) class SandboxError(Exception): """A command was refused before it ran.""" @dataclass class Observation: returncode: int output: str truncated: int = 0 # number of characters cut def split_commands(cmd: str) -> List[List[str]]: """Split a command line into simple commands. Refuses things the toy allowlist cannot judge.""" if "\n" in cmd or "`" in cmd or "$(" in cmd or "<<" in cmd or "<(" in cmd or ">(" in cmd: raise SandboxError("newlines, backticks, $( ), heredocs and process substitution are not allowed") lex = shlex.shlex(cmd, posix=True, punctuation_chars=True) lex.whitespace_split = True try: toks = list(lex) except ValueError as e: raise SandboxError("cannot parse command: %s" % e) cmds, cur = [], [] for t in toks: if t in (";", "&&", "||", "|", "&"): if cur: cmds.append(cur) cur = [] else: cur.append(t) if cur: cmds.append(cur) if not cmds: raise SandboxError("empty command") return cmds class Sandbox: """Runs one bash command at a time in a directory, after an allowlist check. This is a TOY. It refuses obvious escapes, but a real harness uses a container or a VM. Each command is a fresh process, so `cd` does not carry over to the next command. """ def __init__(self, workdir: str, allow=DEFAULT_ALLOW, timeout: float = 5.0, max_output: int = 2000): self.workdir, self.allow, self.timeout, self.max_output = os.path.realpath(workdir), tuple(allow), timeout, max_output def check(self, cmd: str): for argv in split_commands(cmd): toks = [t for t in argv if t not in (">", ">>", "<", "2>", "2>>", ">&", "2>&1")] if not toks: raise SandboxError("empty command") if toks[0] not in self.allow: raise SandboxError("%r is not on the allowlist %s" % (toks[0], list(self.allow))) for t in toks[1:]: pathlike = "/" in t or t.startswith(".") if (t.startswith("/") or t.startswith("~") or ".." in t.split("/") or "$" in t or "`" in t or (pathlike and any(c in t for c in "[{?"))): raise SandboxError("argument %r looks like a path or expansion outside the work directory" % t) def run(self, cmd: str) -> Observation: self.check(cmd) p = subprocess.Popen(["bash", "-c", cmd], cwd=self.workdir, env={"PATH": "/usr/bin:/bin", "LC_ALL": "C"}, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, start_new_session=True) try: out, _ = p.communicate(timeout=self.timeout) rc = p.returncode except subprocess.TimeoutExpired: try: os.killpg(p.pid, 9) except ProcessLookupError: pass out, _ = p.communicate() return Observation(124, "timed out after %ss" % self.timeout) cut = max(0, len(out) - self.max_output) return Observation(rc, out[: self.max_output] + ("\n[truncated %d characters]" % cut if cut else ""), cut) @dataclass class Event: id: int kind: str # system, user, assistant, observation, format_error, refusal text: str _ROLE = {"system": "system", "user": "user", "assistant": "assistant", "observation": "user", "format_error": "user", "refusal": "user"} def rebuild_messages(events: List[Event]) -> List[dict]: """The model's messages are rebuilt from the event log on every step (an OpenHands-style idea).""" return [{"role": _ROLE[e.kind], "content": e.text} for e in events] def events_to_jsonl(events: List[Event]) -> str: return "\n".join(json.dumps({"id": e.id, "kind": e.kind, "text": e.text}) for e in events) def events_from_jsonl(text: str) -> List[Event]: return [Event(**json.loads(l)) for l in text.splitlines() if l.strip()] @dataclass class RunResult: status: str # submitted, step_limit, stuck, model_error steps: int events: List[Event] class ScriptedModel: """A mock model that replays a list of replies. It keeps every message list it was shown.""" def __init__(self, replies: List[str]): self.replies, self.seen = list(replies), [] def __call__(self, messages: List[dict]) -> str: self.seen.append(copy.deepcopy(messages)) if not self.replies: raise RuntimeError("script ended") return self.replies.pop(0) class _NullSpan: def __init__(self): self.attributes: Dict[str, Any] = {} @contextlib.contextmanager def _span(tracer, name: str, **attrs): """A tracer span if there is a tracer, else a throwaway object (so callers need no ifs).""" if tracer is None: yield _NullSpan() else: with tracer.span(name, **attrs) as sp: yield sp class BashAgent: """A bash-only agent loop with a linear history. Each step: show the history to the model, take its one ```bash block, run it, add the output. Ends on the done marker, the step limit, or `stuck` identical (action, output) pairs. """ def __init__(self, model: Callable, sandbox: Sandbox, system: str = "You fix files. Reply with one ```bash block.", step_limit: int = 10, stuck: int = 3, tracer=None): self.model, self.sandbox, self.system, self.step_limit, self.stuck = model, sandbox, system, step_limit, stuck self.tracer = tracer self.events: List[Event] = [] def _add(self, kind: str, text: str): self.events.append(Event(len(self.events), kind, text)) def run(self, task: str) -> RunResult: with _span(self.tracer, "invoke_agent bash-agent", **{"gen_ai.operation.name": "invoke_agent", "gen_ai.agent.name": "bash-agent"}) as sp: res = self._run(task) sp.attributes["app.status"] = res.status sp.attributes["app.steps"] = res.steps return res def _run(self, task: str) -> RunResult: self.events = [] self._add("system", self.system) self._add("user", task) recent: List[Tuple[str, str]] = [] for step in range(1, self.step_limit + 1): try: msgs = rebuild_messages(self.events) with _span(self.tracer, "chat mock-model", **{"gen_ai.operation.name": "chat", "gen_ai.provider.name": "mock", "gen_ai.request.model": "mock-model"}) as sp: reply = self.model(msgs) _usage(sp, "mock-model", json.dumps(msgs), reply) except Exception as e: self._add("format_error", "model error: %s" % e) return RunResult("model_error", step, self.events) self._add("assistant", reply) blocks = _FENCE.findall(reply) if len(blocks) != 1: self._add("format_error", "Reply with exactly one ```bash block. I found %d." % len(blocks)) continue cmd = blocks[0].strip() if cmd == "echo " + DONE_MARK: return RunResult("submitted", step, self.events) try: with _span(self.tracer, "execute_tool bash", **{"gen_ai.operation.name": "execute_tool", "gen_ai.tool.name": "bash"}) as sp: try: obs = self.sandbox.run(cmd) except SandboxError: sp.attributes["app.refused"] = True raise text = "exit %d\n%s" % (obs.returncode, obs.output) self._add("observation", text) except SandboxError as e: text = "refused: %s" % e self._add("refusal", text) recent.append((cmd, text)) if len(recent) >= self.stuck and len(set(recent[-self.stuck:])) == 1: return RunResult("stuck", step, self.events) return RunResult("step_limit", self.step_limit, self.events) class SessionTree: """History as a tree (a Pi-style idea): go back to an earlier message and branch from there.""" def __init__(self, system: str = "system"): self.nodes: Dict[int, dict] = {0: {"parent": None, "role": "system", "content": system}} self.head = 0 def add(self, role: str, content: str) -> int: nid = len(self.nodes) self.nodes[nid] = {"parent": self.head, "role": role, "content": content} self.head = nid return nid def branch_from(self, node_id: int): if node_id not in self.nodes: raise KeyError(node_id) self.head = node_id def path(self, node_id: Optional[int] = None) -> List[dict]: nid, out = self.head if node_id is None else node_id, [] while nid is not None: n = self.nodes[nid] out.append({"role": n["role"], "content": n["content"]}) nid = n["parent"] return out[::-1] def leaves(self) -> List[int]: parents = {n["parent"] for n in self.nodes.values()} return sorted(i for i in self.nodes if i not in parents) class FileTools: """Four tools in a directory: read, write, edit, bash (the surface Pi's README lists). `edit` replaces exactly one occurrence, so an ambiguous edit fails loudly. """ def __init__(self, workdir: str, sandbox: Optional[Sandbox] = None): self.root = os.path.realpath(workdir) self.sandbox = sandbox or Sandbox(workdir) def _path(self, rel: str) -> str: p = os.path.realpath(os.path.join(self.root, rel)) if os.path.commonpath([p, self.root]) != self.root: raise SandboxError("path %r is outside the work directory" % rel) return p def read(self, path: str) -> str: with open(self._path(path), encoding="utf-8") as f: return f.read() def write(self, path: str, text: str) -> str: p = self._path(path) os.makedirs(os.path.dirname(p), exist_ok=True) with open(p, "w", encoding="utf-8") as f: f.write(text) return "wrote %d characters" % len(text) def edit(self, path: str, old: str, new: str) -> str: text = self.read(path) n = text.count(old) if n != 1: raise ValueError("edit needs exactly one match for %r, found %d" % (old, n)) self.write(path, text.replace(old, new)) return "edited" def bash(self, cmd: str) -> Observation: return self.sandbox.run(cmd) def call(self, name: str, **kw): if name not in ("read", "write", "edit", "bash"): raise KeyError("no tool named %r" % name) return getattr(self, name)(**kw) # ===================================================================================== # 4. A Temporal-style durable runtime # ===================================================================================== class SimulatedCrash(Exception): """The worker 'died' at a chosen point. Used to test replay.""" class NonDeterminismError(RuntimeError): """On replay the workflow code did something different from its recorded history.""" class WorkflowSuspended(Exception): """The workflow is waiting for a timer or a signal.""" class NonRetryable(Exception): """An activity raises this to say: do not retry.""" class ActivityError(Exception): """An activity failed for good (out of attempts, or not retryable).""" @dataclass class RetryPolicy: """Backoff between attempts. Temporal's defaults are 1 s, x2, a cap of 100 x the first interval, and unlimited attempts [from its docs]. My default of 3 attempts is my own.""" initial_interval: float = 1.0 backoff: float = 2.0 max_interval: float = 100.0 max_attempts: int = 3 non_retryable: Tuple[type, ...] = () def delay(self, failed_attempt: int) -> float: return min(self.initial_interval * self.backoff ** (failed_attempt - 1), self.max_interval) class HistoryStore: """An append-only event history per workflow. Optionally mirrored to a JSON Lines file.""" def __init__(self, path: Optional[str] = None): self.path, self.events = path, {} if path and os.path.exists(path): with open(path, encoding="utf-8") as f: for line in f: if line.strip(): row = json.loads(line) self.events.setdefault(row["wf"], []).append(row["event"]) def append(self, wf: str, event: dict) -> dict: evs = self.events.setdefault(wf, []) event = dict(event, id=len(evs) + 1) evs.append(event) if self.path: with open(self.path, "a", encoding="utf-8") as f: f.write(json.dumps({"wf": wf, "event": event}) + "\n") return event def for_workflow(self, wf: str) -> List[dict]: return list(self.events.get(wf, [])) _activity_local = threading.local() def activity_info() -> dict: """Inside an activity: {'key': a stable idempotency key, 'attempt': 1-based attempt number}.""" return dict(getattr(_activity_local, "info", {})) @dataclass class RunStatus: status: str # completed, failed, suspended result: Any = None detail: str = "" def _jsonable(x): return json.loads(json.dumps(x)) def _hash(args) -> str: return hashlib.sha1(json.dumps(args, sort_keys=True).encode()).hexdigest()[:10] class Runtime: """Runs workflows by REPLAY. The workflow function is plain Python that calls ctx.activity, ctx.sleep, ctx.wait_for_signal and ctx.side_effect. Each call gets a sequence number. If the history already has the answer for that number, it is returned and nothing runs again. Time here is virtual: sleeping costs nothing, so tests are fast.""" def __init__(self, store: Optional[HistoryStore] = None): self.store = store or HistoryStore() self.activities: Dict[str, Tuple[Callable, RetryPolicy]] = {} self.workflows: Dict[str, Callable] = {} self.now = 0.0 self.crash_at: Optional[Tuple[int, str]] = None # (event id, "before" | "after") self.executions: List[Tuple[str, int]] = [] # every real activity attempt: (name, attempt) def register_activity(self, name: str, fn: Callable, retry: Optional[RetryPolicy] = None): self.activities[name] = (fn, retry or RetryPolicy()) def register_workflow(self, name: str, fn: Callable): self.workflows[name] = fn def _record(self, wf: str, type_: str, **data) -> dict: nxt = len(self.store.for_workflow(wf)) + 1 if self.crash_at and self.crash_at == (nxt, "before"): self.crash_at = None raise SimulatedCrash("crash before event %d (%s)" % (nxt, type_)) ev = self.store.append(wf, dict(type=type_, t=self.now, **data)) if self.crash_at and self.crash_at == (nxt, "after"): self.crash_at = None raise SimulatedCrash("crash after event %d (%s)" % (nxt, type_)) return ev def start(self, wf: str, name: str, *args) -> RunStatus: if self.store.for_workflow(wf): raise ValueError("workflow id %r already exists" % wf) if name not in self.workflows: raise KeyError("no workflow named %r" % name) self._record(wf, "WorkflowStarted", name=name, args=_jsonable(list(args))) return self.run(wf) def signal(self, wf: str, name: str, payload: Any = None): self._record(wf, "SignalReceived", name=name, payload=_jsonable(payload)) def advance(self, seconds: float): self.now += seconds def status(self, wf: str) -> str: evs = self.store.for_workflow(wf) if not evs: raise KeyError(wf) return {"WorkflowCompleted": "completed", "WorkflowFailed": "failed"}.get(evs[-1]["type"], "running") def run(self, wf: str) -> RunStatus: """Run or replay a workflow until it finishes, fails or must wait.""" evs = self.store.for_workflow(wf) if not evs: raise KeyError(wf) last = evs[-1] if last["type"] == "WorkflowCompleted": return RunStatus("completed", last["result"]) if last["type"] == "WorkflowFailed": return RunStatus("failed", None, last["error"]) started = evs[0] ctx = Ctx(self, wf) try: result = self.workflows[started["name"]](ctx, *started["args"]) except WorkflowSuspended as s: return RunStatus("suspended", None, str(s)) except ActivityError as e: self._record(wf, "WorkflowFailed", error=str(e)) return RunStatus("failed", None, str(e)) self._record(wf, "WorkflowCompleted", result=_jsonable(result)) return RunStatus("completed", _jsonable(result)) class Ctx: """What a workflow function sees. Every method here is a 'command' with a sequence number.""" def __init__(self, rt: Runtime, wf: str): self.rt, self.wf, self.seq = rt, wf, -1 def _events(self, seq: int) -> List[dict]: return [e for e in self.rt.store.for_workflow(self.wf) if e.get("seq") == seq] def _next(self) -> int: self.seq += 1 return self.seq def activity(self, name: str, *args) -> Any: seq = self._next() fn, policy = self.rt.activities[name] evs = self._events(seq) sched = [e for e in evs if e["type"] == "ActivityScheduled"] if sched: if sched[0]["name"] != name or sched[0]["args_hash"] != _hash(_jsonable(list(args))): raise NonDeterminismError("command %d was %s%s in history but the code now calls %s%s" % ( seq, sched[0]["name"], "(...)", name, "(...)")) else: self.rt._record(self.wf, "ActivityScheduled", seq=seq, name=name, args_hash=_hash(_jsonable(list(args)))) for e in evs: if e["type"] == "ActivityCompleted": return e["result"] if e["type"] == "ActivityFailed": raise ActivityError(e["error"]) attempt = 1 + sum(1 for e in evs if e["type"] == "ActivityAttemptFailed") while True: _activity_local.info = {"key": "%s:%d" % (self.wf, seq), "attempt": attempt} self.rt.executions.append((name, attempt)) try: result = _jsonable(fn(*args)) except Exception as ex: final = isinstance(ex, NonRetryable) or isinstance(ex, policy.non_retryable) or attempt >= policy.max_attempts if final: self.rt._record(self.wf, "ActivityFailed", seq=seq, error="%s: %s" % (type(ex).__name__, ex), attempts=attempt) raise ActivityError("%s failed after %d attempt(s): %s" % (name, attempt, ex)) d = policy.delay(attempt) self.rt._record(self.wf, "ActivityAttemptFailed", seq=seq, attempt=attempt, error=str(ex), retry_in=d) self.rt.advance(d) attempt += 1 continue finally: _activity_local.info = {} self.rt._record(self.wf, "ActivityCompleted", seq=seq, result=result, attempts=attempt) return result def sleep(self, seconds: float): seq = self._next() evs = self._events(seq) if any(e["type"] == "TimerFired" for e in evs): return started = [e for e in evs if e["type"] == "TimerStarted"] started = started[0] if started else self.rt._record(self.wf, "TimerStarted", seq=seq, fire_at=self.rt.now + seconds) if self.rt.now >= started["fire_at"]: self.rt._record(self.wf, "TimerFired", seq=seq) return raise WorkflowSuspended("sleeping until t=%g" % started["fire_at"]) def wait_for_signal(self, name: str) -> Any: seq = self._next() evs = self._events(seq) done = [e for e in evs if e["type"] == "SignalConsumed"] all_ev = self.rt.store.for_workflow(self.wf) if done: return next(e["payload"] for e in all_ev if e["id"] == done[0]["signal_id"]) taken = {e["signal_id"] for e in all_ev if e["type"] == "SignalConsumed"} for e in all_ev: if e["type"] == "SignalReceived" and e["name"] == name and e["id"] not in taken: self.rt._record(self.wf, "SignalConsumed", seq=seq, signal_id=e["id"]) return e["payload"] raise WorkflowSuspended("waiting for signal %r" % name) def side_effect(self, fn: Callable) -> Any: seq = self._next() evs = [e for e in self._events(seq) if e["type"] == "SideEffect"] if evs: return evs[0]["value"] value = _jsonable(fn()) self.rt._record(self.wf, "SideEffect", seq=seq, value=value) return value def now(self) -> float: """The virtual time, recorded once, so a replay sees the same value.""" return self.side_effect(lambda: self.rt.now) def format_history(events: List[dict]) -> str: """One line per event, like a short event-history view.""" lines = [] for e in events: extra = {k: v for k, v in e.items() if k not in ("id", "type", "t")} lines.append("%2d t=%-7g %-22s %s" % (e["id"], e["t"], e["type"], json.dumps(extra, sort_keys=True)[:70])) return "\n".join(lines) # ===================================================================================== # 5. Tracing, cost, and a little statistics # ===================================================================================== PRICES = {"mock-imitator": (0.15, 0.60), "mock-model": (0.15, 0.60)} # made-up USD per 1M input / output tokens def approx_tokens(text: str) -> int: """A rough estimate: characters / 4, rounded up. Real tokenizers differ.""" return max(1, math.ceil(len(text) / 4)) def cost_usd(model: str, in_tokens: int, out_tokens: int, prices: Dict[str, Tuple[float, float]] = PRICES) -> float: pin, pout = prices[model] return in_tokens * pin / 1e6 + out_tokens * pout / 1e6 def _usage(sp, model: str, prompt: str, reply: str): i, o = approx_tokens(prompt), approx_tokens(reply) sp.attributes.update({"gen_ai.usage.input_tokens": i, "gen_ai.usage.output_tokens": o, "gen_ai.response.finish_reasons": ["stop"], "app.cost_usd": cost_usd(model, i, o)}) @dataclass class Span: trace_id: str span_id: str parent_id: Optional[str] name: str start: float end: float = 0.0 attributes: Dict[str, Any] = field(default_factory=dict) status: str = "OK" events: List[dict] = field(default_factory=list) @property def duration(self) -> float: return self.end - self.start class Tracer: """Collects spans. Parent and child come from `with` nesting. The clock can be injected.""" def __init__(self, clock: Callable[[], float] = time.time, seed: int = 0, capture_content: bool = False): self.clock, self.rng, self.capture_content = clock, random.Random(seed), capture_content self.spans: List[Span] = [] self._local = threading.local() def _stack(self) -> List[Span]: if not hasattr(self._local, "stack"): self._local.stack = [] return self._local.stack @contextlib.contextmanager def span(self, name: str, **attributes): st = self._stack() parent = st[-1] if st else None sp = Span(parent.trace_id if parent else "%032x" % self.rng.getrandbits(128), "%016x" % self.rng.getrandbits(64), parent.span_id if parent else None, name, self.clock(), attributes=dict(attributes)) st.append(sp) try: yield sp except Exception as e: sp.status = "ERROR" sp.attributes["error.type"] = type(e).__name__ sp.events.append({"name": "exception", "message": str(e)[:200]}) raise finally: sp.end = self.clock() st.pop() self.spans.append(sp) def content(self, sp: Span, key: str, value: Any): """Record prompt or argument text. Off by default: it can hold private data (the spec makes it opt-in).""" if self.capture_content: sp.attributes[key] = value class TracedLM: """Wrap any LM so each call becomes a span named 'chat ' with token and cost attributes.""" def __init__(self, lm, tracer: Tracer, model: str = "mock-imitator"): self.lm, self.tracer, self.model = lm, tracer, model @property def calls(self) -> int: return self.lm.calls def __call__(self, prompt: str) -> str: with self.tracer.span("chat " + self.model, **{"gen_ai.operation.name": "chat", "gen_ai.provider.name": "mock", "gen_ai.request.model": self.model}) as sp: self.tracer.content(sp, "gen_ai.input.messages", prompt) reply = self.lm(prompt) _usage(sp, self.model, prompt, reply) return reply def traced_tool(tracer: Tracer, name: str, fn: Callable, *args, call_id: str = "", **kw): """Run a tool inside an 'execute_tool ' span.""" with tracer.span("execute_tool " + name, **{"gen_ai.operation.name": "execute_tool", "gen_ai.tool.name": name, "gen_ai.tool.call.id": call_id}) as sp: tracer.content(sp, "gen_ai.tool.call.arguments", {"args": list(args), "kw": kw}) return fn(*args, **kw) def traced_node(tracer: Tracer, name: str, fn: Callable) -> Callable: """Wrap a graph node. The span name 'node ' is my own; the GenAI spec has no graph-node span.""" def wrapper(state): with tracer.span("node " + name, **{"graph.node": name}): return fn(state) return wrapper def export_jsonl(spans: List[Span]) -> str: """One JSON object per line. Field names follow OTLP loosely, but this is not OTLP.""" rows = [] for s in spans: rows.append(json.dumps({"trace_id": s.trace_id, "span_id": s.span_id, "parent_span_id": s.parent_id, "name": s.name, "start_time_unix_nano": int(s.start * 1e9), "end_time_unix_nano": int(s.end * 1e9), "attributes": s.attributes, "status": {"code": s.status}, "events": s.events}, sort_keys=True)) return "\n".join(rows) def load_jsonl(text: str) -> List[Span]: out = [] for line in text.splitlines(): if line.strip(): d = json.loads(line) out.append(Span(d["trace_id"], d["span_id"], d["parent_span_id"], d["name"], d["start_time_unix_nano"] / 1e9, d["end_time_unix_nano"] / 1e9, d["attributes"], d["status"]["code"], d["events"])) return out def summarize(spans: List[Span]) -> dict: by_op: Dict[str, dict] = {} for s in spans: op = s.attributes.get("gen_ai.operation.name", "other") row = by_op.setdefault(op, {"count": 0, "seconds": 0.0, "errors": 0}) row["count"] += 1 row["seconds"] += s.duration row["errors"] += s.status == "ERROR" return {"spans": len(spans), "errors": sum(s.status == "ERROR" for s in spans), "input_tokens": sum(s.attributes.get("gen_ai.usage.input_tokens", 0) for s in spans), "output_tokens": sum(s.attributes.get("gen_ai.usage.output_tokens", 0) for s in spans), "cost_usd": sum(s.attributes.get("app.cost_usd", 0.0) for s in spans), "by_operation": by_op} def render_trace(spans: List[Span]) -> str: """A text tree: name, duration in ms, tokens, cost, and errors.""" kids: Dict[Optional[str], List[Span]] = {} ids = {s.span_id for s in spans} for s in sorted(spans, key=lambda s: (s.start, s.end)): kids.setdefault(s.parent_id if s.parent_id in ids else None, []).append(s) lines: List[str] = [] def walk(parent, depth): for s in kids.get(parent, []): bits = ["%.1f ms" % (s.duration * 1000)] a = s.attributes if "gen_ai.usage.input_tokens" in a: bits.append("tok %d/%d" % (a["gen_ai.usage.input_tokens"], a["gen_ai.usage.output_tokens"])) if "app.cost_usd" in a: bits.append("$%.6f" % a["app.cost_usd"]) if s.status == "ERROR": bits.append("ERROR " + str(a.get("error.type", ""))) if a.get("app.refused"): bits.append("refused") lines.append("%s%s [%s]" % (" " * depth, s.name, ", ".join(bits))) walk(s.span_id, depth + 1) walk(None, 0) return "\n".join(lines) def bootstrap_ci(values: List[float], n_resamples: int = 1000, alpha: float = 0.05, seed: int = 0) -> Tuple[float, float, float]: """Mean and a percentile bootstrap interval. Resample with replacement, take the mean each time.""" if not values: raise ValueError("no values") rng = random.Random(seed) means = sorted(sum(rng.choice(values) for _ in values) / len(values) for _ in range(n_resamples)) lo = means[int((alpha / 2) * n_resamples)] hi = means[min(n_resamples - 1, int((1 - alpha / 2) * n_resamples))] return sum(values) / len(values), lo, hi def pareto_front(points: List[Tuple[str, float, float]]) -> List[str]: """points are (name, cost, score). Keep a point unless another has cost <= and score >=, and is better on one.""" keep = [] for n, c, s in points: dominated = any((c2 <= c and s2 >= s) and (c2 < c or s2 > s) for n2, c2, s2 in points if n2 != n) if not dominated: keep.append(n) return keep # ===================================================================================== # Demo # ===================================================================================== class FakeClock: """A clock that moves 12.5 ms on every read, so trace output is the same on every run.""" def __init__(self, step: float = 0.0125): self.t, self.step = 1000.0, step def __call__(self) -> float: self.t += self.step return self.t def _h(title: str): print("\n== %s ==" % title) def demo(): from typing import TypedDict train, dev = make_tickets(40, seed=1), make_tickets(40, seed=2) _h("1. DSPy-like: signature, prompt, and optimizers (made-up tickets, mock LM)") sig = Signature.parse("ticket -> category") print("signature:", sig, "| with CoT:", sig.prepend_output("reasoning")) p = Predict(sig, make_lm(0)) p.demos = [{"ticket": "refund invoice 1177, the amount is wrong", "category": "billing"}] print(render_prompt(sig, p.demos, {"ticket": "why is there an extra charge"})) k, seeds, table = 8, range(10), {} makers = [("no demos", lambda s: None), ("labeled few-shot", lambda s: LabeledFewShot(k, s)), ("bootstrap, 1 round", lambda s: BootstrapFewShot(category_match, k, 0, 1, s)), ("bootstrap, 3 rounds", lambda s: BootstrapFewShot(category_match, k, 0, 3, s)), ("bootstrap + labeled", lambda s: BootstrapFewShot(category_match, k, k, 3, s)), ("random search", lambda s: RandomSearch(category_match, 4, k, 0, 3, s))] for name, mk in makers: for s in seeds: lm = make_lm(s) configure(lm) opt, prog = mk(s), Classify() if opt: prog = opt.compile(prog, train) configure(make_lm(100 + s)) table.setdefault(name, []).append(evaluate(prog, dev, category_match).score) print("dev accuracy, mean over 10 seeds (k=%d demos, 40 dev tickets):" % k) for name, v in table.items(): print(" %-22s %.3f" % (name, sum(v) / len(v))) lm = make_lm(3) configure(lm) tri = BootstrapFewShot(triage_match, 6, 0, 3, 0).compile(Triage(), train) print("two-stage Triage: bootstrapped %d training examples; demos per predictor: %s" % ( tri._bootstrapped, {n: len(q.demos) for n, q in tri.named_predictors()})) best = BootstrapFewShot(category_match, k, k, 3, 0).compile(Classify(), train) _h("2. LangGraph-like: state, conditional edge, checkpoint, interrupt, resume") class TicketState(TypedDict): ticket: str category: str reply: str log: Annotated[list, operator.add] runs = {"approve": 0} classify_lm = make_lm(7) configure(classify_lm) def classify(s): return {"category": best(ticket=s["ticket"]).category, "log": ["classify"]} def draft(s): return {"reply": "Hi, about your %s question: we are on it." % s["category"], "log": ["draft"]} def approve(s): runs["approve"] += 1 # this line runs again on resume answer = interrupt({"question": "Send this reply?", "reply": s["reply"]}) return {"log": ["approved" if answer else "rejected"]} def send(s): return {"log": ["send"]} def route(s) -> Literal["draft", "__end__"]: return END if s["category"] == "feedback" else "draft" b = StateGraph(TicketState) for n, f in [("classify", classify), ("draft", draft), ("approve", approve), ("send", send)]: b.add_node(n, f) b.add_edge(START, "classify").add_conditional_edges("classify", route) b.add_edge("draft", "approve").add_edge("approve", "send").add_edge("send", END) path = os.path.join(tempfile.mkdtemp(), "checkpoints.jsonl") cfg = {"configurable": {"thread_id": "ticket-1"}} g1 = b.compile(checkpointer=FileSaver(path)) r = g1.invoke({"ticket": "I was charged twice for my pro plan this month"}, cfg) print("first call stops at:", r["__interrupt__"][0]) g2 = b.compile(checkpointer=FileSaver(path)) # a new compiled graph, as after a restart print("snapshot after restart:", g2.get_state(cfg).status, g2.get_state(cfg).next) r = g2.invoke(Command(resume=True), cfg) print("final log:", r["log"], "| approve node ran", runs["approve"], "times | checkpoints:", len(g2.history(cfg))) print(g1.to_mermaid()) _h("3. mini-SWE-agent-style harness: one bash action per step, in a temp directory") d = tempfile.mkdtemp() with open(os.path.join(d, "config.ini"), "w") as f: f.write("[server]\nhost = localhost\ntimout = 30\nport = 8080\n") fence = lambda c, t="Next.": "%s\n```bash\n%s\n```" % (t, c) script = [fence("ls"), fence("grep -n timout config.ini"), fence("printf '[server]\\nhost = localhost\\ntimeout = 30\\nport = 8080\\n' > config.ini"), fence("grep -n timeout config.ini"), fence("echo " + DONE_MARK)] agent = BashAgent(ScriptedModel(script), Sandbox(d), step_limit=8) res = agent.run("The key 'timout' in config.ini is a typo. Fix it.") for e in res.events: if e.kind == "assistant": print(" action %s" % _FENCE.findall(e.text)[0][:66]) elif e.kind in ("observation", "refusal"): print(" %-11s %s" % (e.kind, e.text.splitlines()[0][:66])) print("status:", res.status, "| steps:", res.steps, "| file fixed:", "timeout = 30" in open(os.path.join(d, "config.ini")).read()) bad = [fence(c) for c in ("cat ../../etc/passwd", "rm -rf .", "cat $HOME/.profile", "curl http://example.com")] + [fence("echo " + DONE_MARK)] res2 = BashAgent(ScriptedModel(bad), Sandbox(d)).run("misbehave") print("adversarial script: %d of 4 commands refused; status %s" % (sum(e.kind == "refusal" for e in res2.events), res2.status)) st = BashAgent(ScriptedModel([fence("ls")] * 6), Sandbox(d), stuck=3).run("loop") print("a model that repeats itself: status %s after %d steps" % (st.status, st.steps)) tree = SessionTree() a = tree.add("user", "fix the bug") tree.add("assistant", "try plan A") tree.branch_from(a) tree.add("assistant", "try plan B") print("session tree: %d nodes, leaves %s, current path: %s" % (len(tree.nodes), tree.leaves(), [m["content"] for m in tree.path()])) _h("4. Temporal-style durable run: crash, replay, and the at-least-once catch") notified = [] def classify_activity(text): return {"category": best(ticket=text).category} def notify_naive(category): notified.append(("naive", category)) return "sent" def notify_keyed(category): key = activity_info()["key"] if key not in [k for k, _ in notified]: notified.append((key, category)) return "sent" def remind(ctx, text): found = ctx.activity("classify", text) # an LM call: recorded once, never repeated on replay ctx.sleep(3 * 86400) # a timer: costs nothing while it waits return ctx.activity("notify", found["category"]) def make_rt(path, notify): rt = Runtime(HistoryStore(path)) rt.register_activity("classify", classify_activity) rt.register_activity("notify", notify) rt.register_workflow("remind", remind) return rt for label, notify in (("without a key", notify_naive), ("with an idempotency key", notify_keyed)): del notified[:] hp = os.path.join(tempfile.mkdtemp(), "history.jsonl") classify_lm.calls = 0 rt = make_rt(hp, notify) s1 = rt.start("wf-1", "remind", "I cannot log in, the password reset email never came") rt.advance(3 * 86400) rt.crash_at = (7, "before") # die after notify ran, before its result is recorded try: rt.run("wf-1") except SimulatedCrash as e: crash = str(e) rt2 = make_rt(hp, notify) # a new worker reads the same history file rt2.now = rt.now s2 = rt2.run("wf-1") print("%-24s first run: %s | crash: %s | after replay: %s | LM calls: %d | notify ran: %d, effect applied: %d" % ( label, s1.status, crash, s2.status, classify_lm.calls, sum(1 for n, _ in rt.executions + rt2.executions if n == "notify"), len(notified))) print(format_history(rt2.store.for_workflow("wf-1"))) _h("5. Tracing: spans, tokens, cost (made-up prices), JSONL, and a text viewer") tr = Tracer(clock=FakeClock(), seed=1) tlm = TracedLM(make_lm(9), tr) configure(tlm) d2 = tempfile.mkdtemp() with open(os.path.join(d2, "config.ini"), "w") as f: f.write("[server]\ntimout = 30\n") with tr.span("invoke_workflow ticket-flow", **{"gen_ai.operation.name": "invoke_workflow"}): best(ticket="the delivery of order 1177 is late") BashAgent(ScriptedModel([fence("grep -n timout config.ini"), fence("cat ../x"), fence("echo " + DONE_MARK)]), Sandbox(d2), tracer=tr).run("find the typo") traced_tool(tr, "lookup_order", lambda n: {"order": n, "status": "late"}, "1177", call_id="call_1") print(render_trace(tr.spans)) sm = summarize(tr.spans) print("totals: %d spans, %d errors, tokens in/out %d/%d, cost $%.6f" % (sm["spans"], sm["errors"], sm["input_tokens"], sm["output_tokens"], sm["cost_usd"])) print("by operation:", {k: v["count"] for k, v in sorted(sm["by_operation"].items())}) print("one JSONL line:", export_jsonl(tr.spans[:1])[:150] + " ...") _h("6. Cost and confidence: made-up numbers for two harness settings (not real measurements)") per_task = {"lean": ([1, 1, 0, 1, 1, 1, 0, 1, 1, 1], [0.6, 0.7, 0.9, 0.5, 0.6, 0.7, 1.1, 0.6, 0.5, 0.7]), "heavy": ([1, 1, 1, 1, 1, 1, 0, 0, 1, 1], [1.4, 1.5, 1.6, 1.2, 1.3, 1.5, 1.9, 1.3, 1.2, 1.4])} pts = [] for name, (ok, cost) in per_task.items(): m, lo, hi = bootstrap_ci(ok, 1000, 0.05, 0) c = sum(cost) / len(cost) pts.append((name, c, m)) print("%-6s success %.2f (95%% interval %.2f to %.2f), mean cost $%.2f per task" % (name, m, lo, hi, c)) print("Pareto front (lower cost, higher success):", pareto_front(pts)) def main(argv: List[str]) -> int: """`python frameworks_toolkit.py` runs the demo. `... view spans.jsonl` shows a saved trace as a tree.""" if len(argv) == 3 and argv[1] == "view": with open(argv[2], encoding="utf-8") as f: spans = load_jsonl(f.read()) print(render_trace(spans)) sm = summarize(spans) print("totals: %d spans, %d errors, tokens in/out %d/%d, cost $%.6f" % ( sm["spans"], sm["errors"], sm["input_tokens"], sm["output_tokens"], sm["cost_usd"])) return 0 if len(argv) == 1: demo() return 0 print("usage: python frameworks_toolkit.py [view spans.jsonl]") return 2 if __name__ == "__main__": import sys sys.exit(main(sys.argv))