redroot 0.2.0__py3-none-any.whl

redroot/cli.py

@@ -0,0 +1,98 @@
+"""Command-line tools for traces saved with ``Trace.to_json()``."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections.abc import Sequence
+
+from redroot.trace import Trace
+
+
+def _load(path: str) -> Trace:
+    try:
+        with open(path, encoding="utf-8") as f:
+            return Trace.from_dict(json.load(f))
+    except FileNotFoundError:
+        raise SystemExit(f"error: file '{path}' not found") from None
+    except json.JSONDecodeError:
+        raise SystemExit(f"error: '{path}' is not valid JSON") from None
+    except (KeyError, TypeError, ValueError) as exc:
+        raise SystemExit(f"error: '{path}' is not a RedRoot trace: {exc}") from None
+
+
+def _visualize(args: argparse.Namespace) -> int:
+    from redroot.visualizer.web.server import run_server
+
+    run_server(_load(args.trace), port=args.port, open_browser=not args.no_browser)
+    return 0
+
+
+def _explain(args: argparse.Namespace) -> int:
+    trace = _load(args.trace)
+    keys = args.keys or list(trace.outputs)
+    for key in keys:
+        out = trace.outputs.get(key)
+        if out is not None and out.node is None:
+            print(f"{key} = {out.value!r}  (not linked to any traced input)")
+            continue
+        try:
+            node = trace.node(key)
+        except KeyError:
+            print(f"error: no input or output named {key!r}", file=sys.stderr)
+            return 1
+        print(f"{key} = {node.value!r}")
+        print(f"  computed as: {trace.explain(key)}")
+        sources = ", ".join(s.key or f"#{s.id}" for s in trace.sources(key))
+        print(f"  from inputs: {sources or '-'}")
+    return 0
+
+
+def _coverage(args: argparse.Namespace) -> int:
+    from redroot.validation import check_identity, coverage
+
+    trace = _load(args.trace)
+    print(coverage(trace))
+    report = check_identity(trace)
+    status = "ok" if report.ok else f"{len(report.mismatches)} MISMATCHES"
+    print(
+        f"identity .......... {status} ({report.checked} checked, {report.skipped} not replayable)"
+    )
+    return 0 if report.ok else 1
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """Build the ``redroot`` argument parser."""
+    parser = argparse.ArgumentParser(prog="redroot", description="Inspect RedRoot traces.")
+    sub = parser.add_subparsers(dest="command")
+
+    viz = sub.add_parser("visualize", help="open a trace in the web viewer")
+    viz.add_argument("trace", help="trace JSON file")
+    viz.add_argument("--port", type=int, default=8050, help="port to serve on (default 8050)")
+    viz.add_argument("--no-browser", action="store_true", help="do not open a browser")
+    viz.set_defaults(handler=_visualize)
+
+    explain = sub.add_parser("explain", help="show how outputs were computed")
+    explain.add_argument("trace", help="trace JSON file")
+    explain.add_argument("keys", nargs="*", help="input or output keys (default: every output)")
+    explain.set_defaults(handler=_explain)
+
+    cov = sub.add_parser("coverage", help="report coverage and check the trace re-evaluates")
+    cov.add_argument("trace", help="trace JSON file")
+    cov.set_defaults(handler=_coverage)
+    return parser
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    """Run the CLI and return its exit code."""
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if args.command is None:
+        parser.print_help()
+        return 0
+    return int(args.handler(args))
+
+
+if __name__ == "__main__":
+    sys.exit(main())

redroot/functions.py

@@ -0,0 +1,240 @@
+"""Recording whole function calls as single nodes.
+
+Some steps are better seen as one opaque operation than as the arithmetic
+inside them: a tax-table lookup, a zip-to-county resolver, an LLM call. A
+:func:`traced` function runs on plain values with tracing suspended, and its
+result is recorded as one node whose operands are the call's arguments.
+
+Deterministic functions are re-called during propagation. Non-deterministic
+ones (``deterministic=False``, e.g. web lookups or LLM calls) keep their
+recorded result while their inputs are unchanged; if an input changes, the
+values derived from them are reported stale.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable, Iterable
+from contextvars import ContextVar
+from decimal import Decimal
+from typing import Any, TypeVar, cast, overload
+
+from redroot._core import (
+    NOT_REPLAYABLE,
+    Node,
+    Observer,
+    Traced,
+    _active,
+    active_trace,
+    decimal_context_id,
+    encode,
+    iter_nodes,
+    record_raise,
+    record_result,
+    unwrap,
+    unwrap_deep,
+)
+from redroot.ops import get_op, register_op
+
+__all__ = ["annotate", "derive", "traced", "traced_llm"]
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+_call_meta: ContextVar[dict[str, Any] | None] = ContextVar("redroot_call_meta", default=None)
+
+
+@overload
+def traced(fn: F, /) -> F: ...
+
+
+@overload
+def traced(
+    *,
+    name: str | None = None,
+    deterministic: bool = True,
+    meta: dict[str, Any] | None = None,
+) -> Callable[[F], F]: ...
+
+
+def traced(
+    fn: F | None = None,
+    /,
+    *,
+    name: str | None = None,
+    deterministic: bool = True,
+    meta: dict[str, Any] | None = None,
+) -> F | Callable[[F], F]:
+    """Record each call of the decorated function as one node.
+
+    Args:
+        fn: The function, when used as a bare ``@traced`` decorator.
+        name: Operation name stored on nodes; defaults to
+            ``"<module>.<qualname>"``. To re-evaluate a deserialized trace,
+            the function must be importable and registered under this name.
+        deterministic: Whether calling again with equal arguments returns an
+            equal result. Only deterministic functions are re-evaluated.
+        meta: Metadata stored on every node this function records.
+
+    Usage::
+
+        @redroot.traced
+        def national_standard(household_size: int) -> Decimal: ...
+
+        @redroot.traced(name="geo.county", deterministic=False)
+        def county_for_zip(zip_code: str) -> str: ...
+    """
+
+    def decorate(func: F) -> F:
+        module = getattr(func, "__module__", None)
+        op_name = name or (f"{module}.{func.__qualname__}" if module else func.__qualname__)
+        register_op(op_name, func, replayable=deterministic)
+        # Nodes keep the exact function that computed them, so helpers that
+        # share a name (e.g. in two generated workflows) never get mixed up.
+        node_fn = func if deterministic else NOT_REPLAYABLE
+
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            trace = _active.get()
+            if trace is None:
+                return func(*args, **kwargs)
+            if isinstance(trace, Observer):  # inside another traced call
+                trace.touch(args, kwargs)
+                return func(*args, **kwargs)
+            refs, _, linked = encode(args, trace)
+            kw_refs, _, kw_linked = encode(kwargs, trace)
+            # The function is one opaque step: it runs on the original
+            # arguments (so mutations behave as usual) with recording
+            # suspended, while an observer notes every traced value it reads.
+            observer = Observer(trace)
+            trace_token = _active.set(observer)
+            meta_token = _call_meta.set({})
+            try:
+                result = func(*args, **kwargs)
+            except Exception as exc:
+                if linked or kw_linked or observer.touched:
+                    record_raise(trace, op_name, tuple(refs), kw_refs or None, exc, node_fn)
+                raise
+            finally:
+                annotations = _call_meta.get()
+                _call_meta.reset(meta_token)
+                _active.reset(trace_token)
+
+            operand_ids = {node.id for node in iter_nodes((refs, kw_refs))}
+            hidden = {**_nested_nodes((args, kwargs), trace), **observer.touched}
+            extra = [node for node_id, node in sorted(hidden.items()) if node_id not in operand_ids]
+            if not (linked or kw_linked or extra):
+                return result
+            fn: Callable[..., Any] = node_fn
+            if extra:
+                # Inputs reached the call where they cannot be replayed from
+                # (closures, sets, dict views...): depend on them, never replay.
+                fn = NOT_REPLAYABLE
+                kw_refs = {**kw_refs, "__depends_on__": extra}
+            node_meta = {**(meta or {}), **(annotations or {})} or None
+            ctx = decimal_context_id(trace) if _involves_decimal(refs, result) else None
+            return record_result(
+                trace,
+                op_name,
+                tuple(_sanitize(r) for r in refs),
+                kw_refs or None,
+                result,
+                node_meta,
+                ctx,
+                fn,
+            )
+
+        wrapper.__redroot_op__ = op_name  # type: ignore[attr-defined]
+        return wrapper  # type: ignore[return-value]
+
+    if fn is not None:
+        return decorate(fn)
+    return decorate
+
+
+def traced_llm(model: str, task: str | None = None, **meta: Any) -> Callable[[F], F]:
+    """Record calls to a function that queries an LLM.
+
+    LLM calls are not deterministic, so they are never re-run: if their inputs
+    change, values derived from them are reported stale. The model name and
+    any ``meta`` are stored on the node; call :func:`annotate` inside the
+    function to add per-call data such as token counts.
+
+    Args:
+        model: Model identifier, stored as ``meta["model"]``.
+        task: Operation name; defaults to the function's qualified name.
+        **meta: More metadata stored on every node.
+    """
+    return traced(name=task, deterministic=False, meta={"model": model, **meta})
+
+
+_VIEW_TYPES: tuple[type, ...] = (
+    set,
+    frozenset,
+    type({}.keys()),
+    type({}.values()),
+    type({}.items()),
+)
+
+
+def _nested_nodes(obj: Any, trace: Any, found: dict[int, Node] | None = None) -> dict[int, Node]:
+    """Nodes of ``trace`` inside ``obj``, including inside sets and dict views."""
+    found = {} if found is None else found
+    if isinstance(obj, Traced):
+        node = obj._rt_node
+        if node.trace is trace:
+            found[node.id] = node
+    elif type(obj) is dict:
+        for item in obj.values():
+            _nested_nodes(item, trace, found)
+    elif type(obj) is list or isinstance(obj, (tuple, *_VIEW_TYPES)):
+        for item in cast("Iterable[Any]", obj):
+            _nested_nodes(item, trace, found)
+    return found
+
+
+def _sanitize(ref: Any) -> Any:
+    """Replace a set or dict view argument by a plain snapshot for storage."""
+    if isinstance(ref, _VIEW_TYPES):
+        items = [unwrap_deep(item) for item in cast("Iterable[Any]", ref)]
+        return type(ref)(items) if isinstance(ref, (set, frozenset)) else items
+    return ref
+
+
+def _involves_decimal(refs: Any, result: Any) -> bool:
+    if type(unwrap(result)) is Decimal:
+        return True
+    return any(type(node.value) is Decimal for node in iter_nodes(refs))
+
+
+def annotate(**meta: Any) -> None:
+    """Attach metadata to the node recorded for the current traced call.
+
+    Does nothing outside a recording call (e.g. during re-evaluation).
+    """
+    current = _call_meta.get()
+    if current is not None:
+        current.update(meta)
+
+
+def derive(
+    value: Any, inputs: Iterable[Any], op: str, *, meta: dict[str, Any] | None = None
+) -> Any:
+    """Record that ``value`` was computed from ``inputs`` outside RedRoot's view.
+
+    Use this for steps that cannot be decorated, e.g. a result returned by an
+    external service. Unless ``op`` names a registered replayable operation,
+    the node cannot be re-evaluated and dependants become stale when
+    ``inputs`` change.
+
+    Returns:
+        ``value`` in traced form (or unchanged outside an active trace).
+    """
+    trace = active_trace()
+    if trace is None:
+        return value
+    refs, _, linked = encode(tuple(inputs), trace)
+    if not linked:
+        return value
+    spec = get_op(op)
+    fn = None if spec is not None and spec.replayable else NOT_REPLAYABLE
+    return record_result(trace, op, refs, None, unwrap_deep(value), meta, fn=fn)