"""Pipelines of named steps with swappable implementations and recorded traces.""" from __future__ import annotations import inspect from dataclasses import dataclass, field from typing import Any, Callable, Hashable, Mapping, Sequence from .noise import Draw, NoiseSource __all__ = ["AS_RUN", "Case", "CostMeter", "Pipeline", "Step", "Trace", "step"] AS_RUN = "as_run" class Step: """A named pipeline step with an as-run implementation and alternatives. The step's output is stored under ``output`` (default: its name). Its dependencies ``deps`` are names of case inputs or upstream outputs; by default they are read from the function's parameter names (excluding ``noise``). If the function has a ``noise`` parameter it receives a :class:`~agentblame.noise.NoiseHandle`. Register alternatives with ``@my_step.variant("name")``. The first registered variant is the default "fix" used for blame unless ``fix=`` is passed to :meth:`variant` or a ``fixes`` mapping is given at analysis time. """ def __init__(self, fn: Callable, name: str | None = None, deps: Sequence[str] | None = None, output: str | None = None, cost: float = 1.0): self.name = name or fn.__name__ self.output = output or self.name if deps is None: deps = [p for p in inspect.signature(fn).parameters if p != "noise"] self.deps = list(deps) self.variants: dict[str, tuple[Callable, float, bool]] = {} self.default_fix: str | None = None self._register(AS_RUN, fn, cost) self.cost = cost def _register(self, name: str, fn: Callable, cost: float) -> None: params = inspect.signature(fn).parameters wants_noise = "noise" in params or any( p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) self.variants[name] = (fn, cost, wants_noise) def variant(self, name: str, cost: float | None = None, fix: bool | None = None) -> Callable[[Callable], Callable]: """Decorator registering an alternative implementation.""" if name == AS_RUN: raise ValueError(f"'{AS_RUN}' is reserved for the original implementation") def deco(fn: Callable) -> Callable: self._register(name, fn, self.cost if cost is None else cost) if fix or (fix is None and self.default_fix is None): self.default_fix = name return fn return deco def call(self, variant: str, args: Mapping[str, Any], handle) -> Any: try: fn, _, wants_noise = self.variants[variant] except KeyError: raise KeyError(f"step '{self.name}' has no variant '{variant}' " f"(has {list(self.variants)})") from None if wants_noise: return fn(**args, noise=handle) return fn(**args) def __repr__(self) -> str: return f"Step({self.name!r}, variants={list(self.variants)})" def step(name: str | None = None, *, deps: Sequence[str] | None = None, output: str | None = None, cost: float = 1.0) -> Callable[[Callable], Step]: """Decorator form of :class:`Step`: ``@step("planner", cost=2.0)``.""" if callable(name): # used bare: @step return Step(name) def deco(fn: Callable) -> Step: return Step(fn, name=name, deps=deps, output=output, cost=cost) return deco @dataclass class Case: """A test case: inputs to the pipeline and the seed of its factual world.""" id: str inputs: dict seed: Hashable = 0 @dataclass class Trace: """Everything recorded about one pipeline execution.""" outputs: dict # case inputs + every step output variants: dict # step -> variant used draws: list[Draw] # all random draws, in order score: float # metric(outputs) cost: float # cost actually spent (cache misses only) world_id: str world_seed: Hashable step_order: list[str] = field(default_factory=list) def output_of(self, step_name: str, pipeline: "Pipeline") -> Any: return self.outputs[pipeline.steps_by_name[step_name].output] class CostMeter: """Accumulates the cost of step executions (cache hits are free).""" def __init__(self) -> None: self.reset() def reset(self) -> None: self.total = 0.0 self.executions = 0 self.cache_hits = 0 self.by_step: dict[str, float] = {} def charge(self, step: str, cost: float) -> None: self.total += cost self.executions += 1 self.by_step[step] = self.by_step.get(step, 0.0) + cost def snapshot(self) -> tuple[float, int, int]: return self.total, self.executions, self.cache_hits def _fingerprint(x: Any) -> str: """Cache key for step inputs. Outputs should have a stable, faithful repr.""" if isinstance(x, dict): return "{" + ",".join(f"{k!r}:{_fingerprint(v)}" for k, v in sorted(x.items(), key=lambda kv: repr(kv[0]))) + "}" return repr(x) class Pipeline: """An ordered DAG of steps plus a success metric. ``metric(outputs) -> float`` receives the dict of case inputs and all step outputs. With ``cache=True`` each step execution is memoized on ``(step, variant, world, inputs)``: counterfactual runs that only differ downstream re-use upstream results, and their cost is not charged again. """ def __init__(self, steps: Sequence[Step], metric: Callable[[dict], float], name: str = "pipeline", cache: bool = True): self.steps = list(steps) self.steps_by_name = {s.name: s for s in self.steps} if len(self.steps_by_name) != len(self.steps): raise ValueError("step names must be unique") produced: set[str] = set() for s in self.steps: for d in s.deps: if d in self.steps_by_name and self.steps_by_name[d].output not in produced: raise ValueError(f"step '{s.name}' depends on '{d}' which runs later") produced.add(s.output) self.metric = metric self.name = name self.use_cache = cache self._cache: dict[tuple, tuple[Any, list[Draw]]] = {} self.meter = CostMeter() @property def step_names(self) -> list[str]: return [s.name for s in self.steps] def clear_cache(self) -> None: self._cache.clear() def default_fixes(self) -> dict[str, str]: return {s.name: s.default_fix for s in self.steps if s.default_fix is not None} def run(self, inputs: Mapping[str, Any], intervention: Mapping[str, str] | None = None, noise: NoiseSource | None = None) -> Trace: """Execute the pipeline under ``intervention`` (step -> variant) in world ``noise``.""" intervention = dict(intervention or {}) unknown = set(intervention) - set(self.steps_by_name) if unknown: raise KeyError(f"unknown steps in intervention: {unknown}") noise = noise if noise is not None else NoiseSource(0) values = dict(inputs) draws: list[Draw] = [] variants: dict[str, str] = {} spent = 0.0 for s in self.steps: v = intervention.get(s.name, AS_RUN) variants[s.name] = v try: args = {d: values[d] for d in s.deps} except KeyError as e: raise KeyError(f"step '{s.name}' needs '{e.args[0]}', not provided") from None key = (s.name, v, noise.world_id, _fingerprint(args)) hit = self._cache.get(key) if self.use_cache else None if hit is not None: out, step_draws = hit self.meter.cache_hits += 1 else: handle = noise.handle(s.name) out = s.call(v, args, handle) step_draws = handle.draws cost = s.variants[v][1] self.meter.charge(s.name, cost) spent += cost if self.use_cache: self._cache[key] = (out, step_draws) values[s.output] = out draws.extend(step_draws) score = float(self.metric(values)) return Trace(values, variants, draws, score, spent, noise.world_id, noise.seed, self.step_names)