redroot 0.2.0__py3-none-any.whl

redroot/propagation.py

@@ -0,0 +1,447 @@
+"""Forward propagation: what happens to every output when an input changes.
+
+:func:`propagate` re-evaluates only the nodes downstream of the edited inputs,
+in recording order, without running the workflow again. It also re-checks
+every guard on the way. The result says, for each output, whether its new
+value is known exactly or whether the workflow must be re-executed.
+
+Soundness rule: an output is only reported ``UPDATED``/``UNCHANGED`` when the
+graph is guaranteed to agree with a re-execution. If anything that steered
+execution changed (a comparison, a conversion, a container's shape, an
+exception), execution may have taken another path, and a trace only records
+the path that was taken. Writes on the other path are invisible to it, so
+*every* output is then reported ``STALE``.
+"""
+
+from __future__ import annotations
+
+import decimal
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any
+
+from redroot._core import (
+    GUARD,
+    LEAF,
+    NOT_REPLAYABLE,
+    Node,
+    Raised,
+    rebuild_tuple,
+    same_value,
+    unwrap,
+)
+from redroot.ops import Opaque, get_op
+
+if TYPE_CHECKING:
+    from redroot.trace import Trace
+
+__all__ = [
+    "Mismatch",
+    "NotReplayableError",
+    "OutputChange",
+    "Propagation",
+    "Reason",
+    "ReasonKind",
+    "Status",
+    "Verification",
+    "evaluate_node",
+    "evaluate_outcome",
+    "node_function",
+    "propagate",
+]
+
+
+class Status(str, Enum):
+    """What propagation knows about one output."""
+
+    UNCHANGED = "unchanged"
+    """The edit does not affect the output."""
+    UPDATED = "updated"
+    """The output has a new value, computed exactly from the graph."""
+    STALE = "stale"
+    """The new value is unknown: re-execute the workflow to obtain it."""
+    UNLINKED = "unlinked"
+    """The output is not linked to any traced input ("not linked to extraction")."""
+
+
+class ReasonKind(str, Enum):
+    """Why the graph alone could not produce an exact result."""
+
+    GUARD_FLIPPED = "guard_flipped"
+    """A recorded decision or conversion now has a different outcome."""
+    GUARD_UNKNOWN = "guard_unknown"
+    """A guard depends on a value that could not be recomputed."""
+    EVALUATION_ERROR = "evaluation_error"
+    """Re-evaluating an operation raised; the workflow would have too."""
+    NOT_REPLAYABLE = "not_replayable"
+    """An operation that cannot be re-run (e.g. an LLM call) has new inputs."""
+    UNTRACKED_INPUT = "untracked_input"
+    """An edited input was not traced (``None``, ``bool``...)."""
+    TYPE_CHANGED = "type_changed"
+    """An input or intermediate value changed type, which code may branch on unseen."""
+
+
+# Reasons that invalidate the whole run rather than individual nodes.
+RUN_LEVEL = frozenset(
+    {
+        ReasonKind.GUARD_FLIPPED,
+        ReasonKind.GUARD_UNKNOWN,
+        ReasonKind.EVALUATION_ERROR,
+        ReasonKind.UNTRACKED_INPUT,
+        ReasonKind.TYPE_CHANGED,
+    }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class Reason:
+    """One explanation of why re-execution is needed."""
+
+    kind: ReasonKind
+    message: str
+    key: str | None = None
+    node: Node | None = None
+    before: Any = None
+    after: Any = None
+
+
+@dataclass(frozen=True, slots=True)
+class OutputChange:
+    """The effect of an edit on one output.
+
+    ``new`` is ``None`` when ``status`` is :attr:`Status.STALE`.
+    """
+
+    key: str
+    status: Status
+    old: Any
+    new: Any
+
+
+@dataclass(frozen=True, slots=True)
+class Mismatch:
+    """An output whose propagated value disagrees with re-execution."""
+
+    key: str
+    status: Status
+    expected: Any
+    actual: Any
+
+
+@dataclass(frozen=True, slots=True)
+class Verification:
+    """Propagated values compared with the outputs of a real re-execution.
+
+    A mismatch on an output that propagation claimed to know exactly is a
+    tracing gap: something influenced the output without being recorded.
+    """
+
+    mismatches: tuple[Mismatch, ...]
+    missing: tuple[str, ...] = ()
+    """Outputs the trace has but the re-execution did not produce."""
+    unexpected: tuple[str, ...] = ()
+    """Outputs the re-execution produced that the trace does not have."""
+
+    @property
+    def ok(self) -> bool:
+        """Whether re-execution confirmed every claim propagation made."""
+        return not (self.mismatches or self.missing or self.unexpected)
+
+
+@dataclass(frozen=True)
+class Propagation:
+    """The result of :func:`propagate`."""
+
+    trace: Trace
+    edits: Mapping[str, Any]
+    outputs: Mapping[str, OutputChange]
+    reasons: tuple[Reason, ...]
+    _values: Mapping[int, Any] = field(repr=False)
+    _stale: frozenset[int] = field(repr=False)
+
+    @property
+    def needs_reexecution(self) -> bool:
+        """Whether some output's new value can only come from re-execution."""
+        return any(r.kind in RUN_LEVEL for r in self.reasons) or any(
+            change.status is Status.STALE for change in self.outputs.values()
+        )
+
+    @property
+    def exact(self) -> bool:
+        """Whether every output's new value is known from the graph alone."""
+        return not self.needs_reexecution
+
+    @property
+    def flipped_guards(self) -> list[Reason]:
+        """The reasons that are guard flips."""
+        return [r for r in self.reasons if r.kind is ReasonKind.GUARD_FLIPPED]
+
+    def with_status(self, status: Status) -> dict[str, OutputChange]:
+        """Outputs with the given status, by key."""
+        return {k: c for k, c in self.outputs.items() if c.status is status}
+
+    @property
+    def updated(self) -> dict[str, OutputChange]:
+        """Outputs with a new, exactly known value."""
+        return self.with_status(Status.UPDATED)
+
+    @property
+    def stale(self) -> dict[str, OutputChange]:
+        """Outputs whose new value needs re-execution."""
+        return self.with_status(Status.STALE)
+
+    def value(self, ref: Any) -> Any:
+        """The new value of any node (by key, traced value or node).
+
+        Raises:
+            LookupError: if the node's new value is unknown.
+        """
+        node = self.trace.node(ref)
+        if node.id in self._stale or any(r.kind in RUN_LEVEL for r in self.reasons):
+            raise LookupError(f"the new value of node #{node.id} is unknown; re-execute")
+        return self._values.get(node.id, node.value)
+
+    def verify(self, actual: Mapping[str, Any]) -> Verification:
+        """Compare with the outputs of a re-execution on the edited inputs.
+
+        Args:
+            actual: Output values by key, e.g. ``retrace.output_values()``.
+        """
+        actual = {key: unwrap(value) for key, value in actual.items()}
+        mismatches = tuple(
+            Mismatch(key, change.status, change.new, actual[key])
+            for key, change in self.outputs.items()
+            if change.status is not Status.STALE
+            and key in actual
+            and not same_value(change.new, actual[key])
+        )
+        if any(r.kind in RUN_LEVEL for r in self.reasons):
+            # Execution may legitimately take another path and emit other keys.
+            return Verification(mismatches)
+        return Verification(
+            mismatches,
+            missing=tuple(k for k in self.outputs if k not in actual),
+            unexpected=tuple(k for k in actual if k not in self.outputs),
+        )
+
+
+def _resolve(arg: Any, values: Mapping[int, Any]) -> Any:
+    kind = type(arg)
+    if kind is Node:
+        arg = values.get(arg.id, arg.value)
+        kind = type(arg)
+        if kind is not Opaque:
+            return arg
+    if kind is Opaque:
+        raise ValueError(f"{arg!r} was not serializable and cannot be re-used")
+    if kind is list:
+        return [_resolve(a, values) for a in arg]
+    if isinstance(arg, tuple):
+        return rebuild_tuple(arg, [_resolve(a, values) for a in arg])
+    if kind is dict:
+        return {k: _resolve(v, values) for k, v in arg.items()}
+    return arg
+
+
+class NotReplayableError(LookupError):
+    """Raised when evaluating a node whose operation cannot be re-run."""
+
+
+def node_function(node: Node) -> Callable[..., Any] | None:
+    """The function that re-evaluates ``node``, or ``None`` if it cannot be re-run.
+
+    A node's own function (set for ``@traced`` calls) takes precedence over
+    the registry, so equally named functions never get mixed up.
+    """
+    if node.fn is NOT_REPLAYABLE:
+        return None
+    if node.fn is not None:
+        return node.fn
+    spec = get_op(node.op)
+    return spec.fn if spec is not None and spec.replayable else None
+
+
+def evaluate_node(node: Node, values: Mapping[int, Any] | None = None) -> Any:
+    """Recompute ``node`` from its operands.
+
+    Operands take their value from ``values`` (by node id) when present, and
+    their recorded value otherwise. ``Decimal`` operations run under the
+    context they were recorded with.
+
+    Raises:
+        NotReplayableError: if the operation is unknown or not replayable.
+    """
+    fn = node_function(node)
+    if fn is None:
+        raise NotReplayableError(f"operation {node.op!r} cannot be re-evaluated")
+    values = values or {}
+    args = _resolve(node.args, values)
+    kwargs = _resolve(node.kwargs, values) if node.kwargs else {}
+    if node.ctx is not None and node.trace is not None:
+        with decimal.localcontext(node.trace._contexts[node.ctx]):
+            return fn(*args, **kwargs)
+    return fn(*args, **kwargs)
+
+
+def evaluate_outcome(node: Node, values: Mapping[int, Any] | None = None) -> Any:
+    """Like :func:`evaluate_node`, but a guard that raises evaluates to :class:`Raised`.
+
+    For a guard, raising *is* the outcome to compare with the recorded one.
+    """
+    if node.kind != GUARD:
+        return evaluate_node(node, values)
+    if node_function(node) is None:
+        raise NotReplayableError(f"operation {node.op!r} cannot be re-evaluated")
+    try:
+        return evaluate_node(node, values)
+    except Exception as exc:
+        return Raised.of(exc)
+
+
+def _describe(node: Node) -> str:
+    from redroot.trace import _render
+
+    return _render(node, 6, top=True)
+
+
+def propagate(trace: Trace, edits: Mapping[str, Any]) -> Propagation:
+    """Recompute ``trace``'s outputs as if its inputs had the edited values.
+
+    Args:
+        trace: A recorded trace.
+        edits: New values by input key, e.g. ``{"ext:bank.balance": 900}``.
+
+    Raises:
+        KeyError: if an edit names an input the trace never saw.
+    """
+    reasons: list[Reason] = []
+    changed: dict[int, Any] = {}
+    for key, raw in edits.items():
+        new = unwrap(raw)
+        node = trace.inputs.get(key)
+        if node is None:
+            untracked = trace.untracked_inputs
+            if key not in untracked:
+                raise KeyError(f"unknown input {key!r}")
+            if not same_value(new, untracked[key]):
+                reasons.append(
+                    Reason(
+                        ReasonKind.UNTRACKED_INPUT,
+                        f"input {key!r} is not traced; only re-execution can apply its new value",
+                        key=key,
+                        before=untracked[key],
+                        after=new,
+                    )
+                )
+            continue
+        if same_value(new, node.value):
+            continue
+        if type(new) is not type(node.value):
+            reasons.append(
+                Reason(
+                    ReasonKind.TYPE_CHANGED,
+                    f"input {key!r} changed type from {type(node.value).__name__} "
+                    f"to {type(new).__name__}",
+                    key=key,
+                    node=node,
+                    before=node.value,
+                    after=new,
+                )
+            )
+        changed[node.id] = new
+
+    values: dict[int, Any] = dict(changed)
+    stale: set[int] = set()
+    if changed:
+        with decimal.localcontext(trace.decimal_context):
+            for node in trace.nodes[min(changed) :]:
+                if node.kind == LEAF:
+                    continue
+                dirty = blocked = False
+                for operand in node.operands():
+                    if operand.id in stale:
+                        blocked = True
+                        break
+                    if operand.id in values:
+                        dirty = True
+                if not (dirty or blocked):
+                    continue
+                reason = _reevaluate(node, values, blocked)
+                if reason is None:
+                    continue
+                if reason.kind is not ReasonKind.GUARD_FLIPPED:
+                    stale.add(node.id)
+                if not blocked or node.kind == GUARD:
+                    reasons.append(reason)
+
+    run_level = any(r.kind in RUN_LEVEL for r in reasons)
+    outputs: dict[str, OutputChange] = {}
+    for key, out in trace.outputs.items():
+        if run_level:
+            change = OutputChange(key, Status.STALE, out.value, None)
+        elif out.node is None:
+            change = OutputChange(key, Status.UNLINKED, out.value, out.value)
+        elif out.node.id in stale:
+            change = OutputChange(key, Status.STALE, out.value, None)
+        elif out.node.id in values:
+            change = OutputChange(key, Status.UPDATED, out.value, values[out.node.id])
+        else:
+            change = OutputChange(key, Status.UNCHANGED, out.value, out.value)
+        outputs[key] = change
+
+    return Propagation(
+        trace=trace,
+        edits=dict(edits),
+        outputs=outputs,
+        reasons=tuple(reasons),
+        _values=values,
+        _stale=frozenset(stale),
+    )
+
+
+def _reevaluate(node: Node, values: dict[int, Any], blocked: bool) -> Reason | None:
+    """Re-evaluate one dirty node, updating ``values``.
+
+    Returns why the node could not be brought up to date, if it could not.
+    """
+    if blocked:
+        if node.kind == GUARD:
+            return Reason(
+                ReasonKind.GUARD_UNKNOWN,
+                f"cannot re-check {_describe(node)}: an input could not be recomputed",
+                node=node,
+            )
+        return Reason(ReasonKind.NOT_REPLAYABLE, "depends on a stale value", node=node)
+    if node_function(node) is None:
+        kind = ReasonKind.GUARD_UNKNOWN if node.kind == GUARD else ReasonKind.NOT_REPLAYABLE
+        return Reason(kind, f"operation {node.op!r} cannot be re-evaluated", node=node)
+    try:
+        new = evaluate_outcome(node, values)
+    except Exception as exc:
+        return Reason(
+            ReasonKind.EVALUATION_ERROR,
+            f"{_describe(node)} raised {type(exc).__name__}: {exc}",
+            node=node,
+        )
+    if same_value(new, node.value):
+        return None  # early cutoff: dependents are unaffected
+    if node.kind == GUARD:
+        return Reason(
+            ReasonKind.GUARD_FLIPPED,
+            f"{_describe(node)} was {node.value!r}, now {new!r}",
+            node=node,
+            before=node.value,
+            after=new,
+        )
+    if type(new) is not type(node.value):
+        # Code may branch on the type (isinstance, formatting) unseen.
+        return Reason(
+            ReasonKind.TYPE_CHANGED,
+            f"{_describe(node)} changed type from {type(node.value).__name__} to {type(new).__name__}",
+            node=node,
+            before=node.value,
+            after=new,
+        )
+    values[node.id] = new
+    return None

redroot/serialization.py

@@ -0,0 +1,245 @@
+"""JSON (de)serialization of traces.
+
+The format is versioned so traces stored with a run can be reloaded later,
+e.g. to propagate a reviewer's edit on a server long after the run::
+
+    {
+      "format": "redroot.trace",
+      "version": 1,
+      "id": "…", "name": "run-42",
+      "decimal_context": {"prec": 28, "rounding": "ROUND_HALF_EVEN", …},
+      "nodes": [
+        {"id": 0, "op": "leaf", "kind": "leaf", "key": "ext:a", "value": 1200},
+        {"id": 1, "op": "gt", "kind": "guard", "args": [{"$node": 0}, 1000], "value": true},
+        …
+      ],
+      "outputs": {"out:total": {"node": 2}, "out:year": {"value": 2025}},
+      "untracked_inputs": {"ext:married": true}
+    }
+
+Values that JSON cannot express are tagged: ``{"$decimal": "1.10"}``,
+``{"$tuple": [...]}``, ``{"$dict": [[key, value], ...]}`` (only for dicts with
+non-string keys), ``{"$float": "nan"}``, ``{"$raised": "ZeroDivisionError"}``
+for an operation that raised, ``{"$node": 3}`` for a reference to another
+node. A node computed under a non-default ``decimal`` context refers to it by
+index (``"ctx"``) into ``"decimal_contexts"``; ``"replayable": false`` marks
+nodes that must never be re-run. Values of any other type are stored as ``{"$opaque": repr}``;
+they load as :class:`~redroot.ops.Opaque`, and nodes using them cannot be
+re-evaluated.
+"""
+
+from __future__ import annotations
+
+import base64
+import decimal
+import math
+from collections.abc import Mapping
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any
+
+from redroot._core import LEAF, NOT_REPLAYABLE, Node, Raised, Traced
+from redroot.ops import Opaque
+
+if TYPE_CHECKING:
+    from redroot.trace import Trace
+
+__all__ = ["FORMAT", "VERSION", "decode_value", "encode_value", "trace_from_dict", "trace_to_dict"]
+
+FORMAT = "redroot.trace"
+VERSION = 1
+
+
+def encode_value(value: Any) -> Any:
+    """Encode a plain value (or node reference) as JSON-compatible data."""
+    if value is None or value is True or value is False:
+        return value
+    kind = type(value)
+    if kind is int or kind is str:
+        return value
+    if kind is float:
+        return value if math.isfinite(value) else {"$float": repr(value)}
+    if kind is Node:
+        return {"$node": value.id}
+    if kind is Decimal:
+        return {"$decimal": str(value)}
+    if kind is list:
+        return [encode_value(item) for item in value]
+    if isinstance(value, tuple):
+        return {"$tuple": [encode_value(item) for item in value]}
+    if kind is dict:
+        if all(type(k) is str for k in value) and not (
+            len(value) == 1 and next(iter(value)).startswith("$")
+        ):
+            return {k: encode_value(v) for k, v in value.items()}
+        return {"$dict": [[encode_value(k), encode_value(v)] for k, v in value.items()]}
+    if kind is slice:
+        return {
+            "$slice": [
+                encode_value(value.start),
+                encode_value(value.stop),
+                encode_value(value.step),
+            ]
+        }
+    if kind is complex:
+        return {"$complex": [encode_value(value.real), encode_value(value.imag)]}
+    if kind is bytes:
+        return {"$bytes": base64.b64encode(value).decode("ascii")}
+    if isinstance(value, Traced):
+        return encode_value(value._rt_node.value)
+    if isinstance(value, Opaque):
+        return {"$opaque": value.description}
+    if isinstance(value, Raised):
+        return {"$raised": value.type_name}
+    return {"$opaque": f"{type(value).__qualname__}: {repr(value)[:200]}"}
+
+
+def decode_value(data: Any, nodes: list[Node] | None = None) -> Any:
+    """Decode :func:`encode_value` output. ``nodes`` resolves ``$node`` references."""
+    if isinstance(data, list):
+        return [decode_value(item, nodes) for item in data]
+    if not isinstance(data, dict):
+        return data
+    if len(data) == 1:
+        tag, payload = next(iter(data.items()))
+        if tag.startswith("$"):
+            return _decode_tagged(tag, payload, nodes)
+    return {k: decode_value(v, nodes) for k, v in data.items()}
+
+
+def _decode_tagged(tag: str, payload: Any, nodes: list[Node] | None) -> Any:
+    if tag == "$node":
+        if nodes is None or not 0 <= payload < len(nodes):
+            raise ValueError(f"reference to unknown or later node #{payload}")
+        return nodes[payload]
+    if tag == "$decimal":
+        return Decimal(payload)
+    if tag == "$float":
+        return float(payload)
+    if tag == "$tuple":
+        return tuple(decode_value(item, nodes) for item in payload)
+    if tag == "$dict":
+        return {decode_value(k, nodes): decode_value(v, nodes) for k, v in payload}
+    if tag == "$slice":
+        return slice(*(decode_value(item, nodes) for item in payload))
+    if tag == "$complex":
+        return complex(*(decode_value(item, nodes) for item in payload))
+    if tag == "$bytes":
+        return base64.b64decode(payload)
+    if tag == "$opaque":
+        return Opaque(payload)
+    if tag == "$raised":
+        return Raised(payload)
+    raise ValueError(f"unknown value tag {tag!r}")
+
+
+def _context_to_dict(ctx: decimal.Context) -> dict[str, Any]:
+    return {
+        "prec": ctx.prec,
+        "rounding": ctx.rounding,
+        "Emin": ctx.Emin,
+        "Emax": ctx.Emax,
+        "capitals": ctx.capitals,
+        "clamp": ctx.clamp,
+        "traps": sorted(signal.__name__ for signal, on in ctx.traps.items() if on),
+    }
+
+
+def _context_from_dict(data: Mapping[str, Any]) -> decimal.Context:
+    traps = [getattr(decimal, name) for name in data.get("traps", [])]
+    return decimal.Context(
+        prec=data["prec"],
+        rounding=data["rounding"],
+        Emin=data["Emin"],
+        Emax=data["Emax"],
+        capitals=data["capitals"],
+        clamp=data["clamp"],
+        traps=traps,
+        flags=[],
+    )
+
+
+def trace_to_dict(trace: Trace) -> dict[str, Any]:
+    """Serialize ``trace`` to JSON-compatible data."""
+    nodes = []
+    for node in trace.nodes:
+        item: dict[str, Any] = {
+            "id": node.id,
+            "op": node.op,
+            "kind": node.kind,
+            "value": encode_value(node.value),
+        }
+        if node.args:
+            item["args"] = [encode_value(arg) for arg in node.args]
+        if node.kwargs:
+            item["kwargs"] = {k: encode_value(v) for k, v in node.kwargs.items()}
+        if node.key is not None:
+            item["key"] = node.key
+        if node.meta:
+            item["meta"] = encode_value(node.meta)
+        if node.ctx is not None:
+            item["ctx"] = node.ctx
+        if node.fn is NOT_REPLAYABLE:
+            item["replayable"] = False
+        nodes.append(item)
+    return {
+        "format": FORMAT,
+        "version": VERSION,
+        "id": trace.id,
+        "name": trace.name,
+        "decimal_context": _context_to_dict(trace.decimal_context),
+        "decimal_contexts": [_context_to_dict(ctx) for ctx in trace.decimal_contexts],
+        "nodes": nodes,
+        "outputs": {
+            key: {"node": out.node.id}
+            if out.node is not None
+            else {"value": encode_value(out.value)}
+            for key, out in trace.outputs.items()
+        },
+        "untracked_inputs": {key: encode_value(v) for key, v in trace.untracked_inputs.items()},
+    }
+
+
+def trace_from_dict(data: Mapping[str, Any]) -> Trace:
+    """Rebuild a trace from :func:`trace_to_dict` output."""
+    from redroot.trace import Output, Trace
+
+    if data.get("format") != FORMAT:
+        raise ValueError(f"not a RedRoot trace (format={data.get('format')!r})")
+    if data.get("version") != VERSION:
+        raise ValueError(f"unsupported trace version {data.get('version')!r}; expected {VERSION}")
+
+    trace = Trace(data.get("name"), id=data["id"])
+    trace.decimal_context = _context_from_dict(data["decimal_context"])
+    trace._contexts = [_context_from_dict(ctx) for ctx in data.get("decimal_contexts", [])]
+    nodes: list[Node] = trace._nodes
+    for item in data["nodes"]:
+        if item["id"] != len(nodes):
+            raise ValueError(f"node ids must be consecutive; got {item['id']} at {len(nodes)}")
+        node = Node(
+            item["op"],
+            tuple(decode_value(arg, nodes) for arg in item.get("args", ())),
+            {k: decode_value(v, nodes) for k, v in item["kwargs"].items()}
+            if "kwargs" in item
+            else None,
+            decode_value(item["value"]),
+            item["kind"],
+            key=item.get("key"),
+            meta=decode_value(item["meta"]) if "meta" in item else None,
+        )
+        node.id = item["id"]
+        node.trace = trace
+        node.ctx = item.get("ctx")
+        if item.get("replayable") is False:
+            node.fn = NOT_REPLAYABLE
+        nodes.append(node)
+        if node.kind == LEAF and node.key is not None:
+            trace._leaves[node.key] = node
+    for key, out in data.get("outputs", {}).items():
+        if "node" in out:
+            node = nodes[out["node"]]
+            trace._outputs[key] = Output(key, node, node.value)
+        else:
+            trace._outputs[key] = Output(key, None, decode_value(out["value"]))
+    for key, value in data.get("untracked_inputs", {}).items():
+        trace._untracked[key] = decode_value(value)
+    return trace