loomcache 1.0.0__py3-none-any.whl

loom/__init__.py

@@ -0,0 +1,49 @@
+"""Loom: a content-addressed, deterministic replay engine for agent AI workflows."""
+
+from .cache import Cache, DiskCache, default_cache
+from .diff import diff_runs, first_divergence, format_diff
+from .run import Node, Run, current_run, step
+from .tracing import TracedBox, get_node, attach_node, unwrap, unwrap_recursive
+
+# Optional extensions – if dependencies are missing, these imports are skipped.
+try:
+    from .cache_remote import S3Cache, RedisCache
+except ImportError:
+    S3Cache = RedisCache = None
+
+try:
+    from .langchain import wrap_runnable, CachedRunnable
+except ImportError:
+    wrap_runnable = CachedRunnable = None
+
+try:
+    from .async_run import AsyncRun, async_step, gather
+except ImportError:
+    AsyncRun = async_step = gather = None
+
+__version__ = "1.0.0"
+
+__all__ = [
+    "step",
+    "Run",
+    "Node",
+    "current_run",
+    "Cache",
+    "DiskCache",
+    "default_cache",
+    "diff_runs",
+    "first_divergence",
+    "format_diff",
+    "get_node",
+    "attach_node",
+    "unwrap",
+    "unwrap_recursive",
+    "TracedBox",
+    "S3Cache",
+    "RedisCache",
+    "wrap_runnable",
+    "CachedRunnable",
+    "AsyncRun",
+    "async_step",
+    "gather",
+]

loom/async_run.py

@@ -0,0 +1,182 @@
+"""Async execution support for Loom.
+
+Use `AsyncRun` in an async context manager with async steps.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextvars
+import functools
+import json
+import time
+import uuid
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
+
+from loom import hashing, default_cache
+from loom.run import Node, current_run
+from loom.tracing import attach_node, get_node, unwrap_recursive
+from loom.cache import Cache
+
+_current_async_run: contextvars.ContextVar = contextvars.ContextVar(
+    "loom_current_async_run", default=None
+)
+
+
+class AsyncRun:
+    """Async context manager for recording async step executions."""
+
+    def __init__(
+        self,
+        name: str,
+        cache: Optional[Cache] = None,
+        root: Union[str, Path] = ".loom_runs",
+    ):
+        self.name = name
+        self.run_id = f"{name}-{uuid.uuid4().hex[:8]}"
+        self.cache = cache or default_cache()
+        self.root = Path(root)
+        self.nodes: list = []
+        self.created_at = time.time()
+        self._token = None
+
+    async def __aenter__(self) -> AsyncRun:
+        self._token = _current_async_run.set(self)
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        _current_async_run.reset(self._token)
+
+    def record(self, node: Node) -> None:
+        self.nodes.append(node)
+
+    def to_dict(self) -> dict:
+        return {
+            "run_id": self.run_id,
+            "name": self.name,
+            "created_at": self.created_at,
+            "nodes": [n.to_dict() for n in self.nodes],
+        }
+
+    def save(self, path: Optional[Union[str, Path]] = None) -> Path:
+        self.root.mkdir(parents=True, exist_ok=True)
+        out_path = Path(path) if path else self.root / f"{self.run_id}.json"
+        out_path.write_text(json.dumps(self.to_dict(), indent=2, default=str))
+        return out_path
+
+    def stats(self) -> dict:
+        hits = sum(1 for n in self.nodes if n.cache_hit)
+        total = len(self.nodes)
+        return {
+            "total_nodes": total,
+            "cache_hits": hits,
+            "cache_misses": total - hits,
+            "hit_rate": (hits / total) if total else 0.0,
+            "wall_time_s": round(sum(n.duration_s for n in self.nodes), 3),
+            "time_saved_s": round(sum(n.time_saved_s for n in self.nodes), 3),
+        }
+
+    @classmethod
+    def load(cls, path: Union[str, Path]) -> AsyncRun:
+        data = json.loads(Path(path).read_text())
+        run = cls.__new__(cls)
+        run.name = data["name"]
+        run.run_id = data["run_id"]
+        run.created_at = data["created_at"]
+        run.cache = default_cache()
+        run.root = Path(path).parent
+        run.nodes = [Node(**n) for n in data["nodes"]]
+        run._token = None
+        return run
+
+
+def async_step(func: Optional[Callable] = None, *, cache: Optional[Cache] = None):
+    """Decorator that works for async functions.
+
+    If the function is async, it returns a coroutine that will be awaited
+    by the caller. For sync functions, use `@loom.step` instead.
+    """
+
+    def decorator(f: Callable) -> Callable:
+        source_hash = hashing.hash_source(f)
+        is_coroutine = asyncio.iscoroutinefunction(f)
+
+        @functools.wraps(f)
+        async def async_wrapper(*args, **kwargs):
+            run = _current_async_run.get()
+            if run is None:
+                # Fallback to calling the function directly (could be sync or async)
+                if is_coroutine:
+                    return await f(*args, **kwargs)
+                else:
+                    return f(*args, **kwargs)
+
+            active_cache = cache or (run.cache if run is not None else default_cache())
+            node_hash = hashing.hash_node(f.__qualname__, source_hash, args, kwargs)
+
+            # Determine parents from traced arguments
+            parents = [
+                n.node_hash
+                for n in (get_node(a) for a in list(args) + list(kwargs.values()))
+                if n is not None
+            ]
+
+            entry = active_cache.get(node_hash)
+            if entry is not None:
+                output = entry.output          # already clean
+                cache_hit = True
+                duration = 0.0
+                time_saved = float(entry.metadata.get("duration_s", 0.0))
+                clean_output = output
+            else:
+                # Unwrap arguments
+                from loom.run import _unwrap
+                call_args = tuple(_unwrap(a) for a in args)
+                call_kwargs = {k: _unwrap(v) for k, v in kwargs.items()}
+                t0 = time.time()
+                if is_coroutine:
+                    raw_output = await f(*call_args, **call_kwargs)
+                else:
+                    raw_output = f(*call_args, **call_kwargs)
+                duration = time.time() - t0
+
+                clean_output = unwrap_recursive(raw_output)
+
+                active_cache.put(
+                    node_hash,
+                    clean_output,
+                    metadata={
+                        "step_name": f.__qualname__,
+                        "timestamp": time.time(),
+                        "duration_s": duration,
+                    },
+                )
+                cache_hit = False
+                time_saved = 0.0
+
+            node = Node(
+                node_hash=node_hash,
+                step_name=f.__qualname__,
+                parents=parents,
+                cache_hit=cache_hit,
+                duration_s=duration,
+                time_saved_s=time_saved,
+                timestamp=time.time(),
+                args_repr=repr(args),
+                kwargs_repr=repr(kwargs),
+                output_repr=repr(clean_output),
+            )
+            run.record(node)
+            return attach_node(clean_output, node)
+
+        return async_wrapper
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+async def gather(*coros, return_exceptions=False):
+    """Run multiple async steps concurrently, with Loom tracking."""
+    return await asyncio.gather(*coros, return_exceptions=return_exceptions)

loom/cache.py

@@ -0,0 +1,114 @@
+"""Content-addressed cache backends.
+
+Loom stores every step's output under a key derived purely from its
+content hash — the same idea as Git's object store or Bazel's action
+cache. Two runs (even in different processes, on different days, on
+different machines sharing this cache) that produce the same hash will
+transparently share the same cached output.
+
+`Cache` is a tiny abstract interface so alternate backends (Redis, S3,
+a shared network drive for a team) can be dropped in later without
+touching the rest of Loom.
+"""
+from __future__ import annotations
+
+import json
+import pickle
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional, Union
+
+
+@dataclass
+class CacheEntry:
+    output: Any
+    metadata: dict
+
+
+class Cache:
+    """Abstract cache interface. Subclass to plug in Redis/S3/etc."""
+
+    def get(self, key: str) -> Optional[CacheEntry]:
+        raise NotImplementedError
+
+    def put(self, key: str, output: Any, metadata: dict) -> None:
+        raise NotImplementedError
+
+    def stats(self) -> dict:
+        raise NotImplementedError
+
+
+class DiskCache(Cache):
+    """Default local cache: `<root>/objects/<hash[:2]>/<hash>.pkl`.
+
+    Sharding by the first two hex characters keeps any one directory
+    from accumulating too many files, the same trick Git uses for its
+    object store.
+    """
+
+    def __init__(self, root: Union[str, Path] = ".loom_cache"):
+        self.root = Path(root)
+        self.objects_dir = self.root / "objects"
+        self.objects_dir.mkdir(parents=True, exist_ok=True)
+        self._hits = 0
+        self._misses = 0
+
+    def _paths(self, key: str):
+        shard = self.objects_dir / key[:2]
+        shard.mkdir(parents=True, exist_ok=True)
+        return shard / f"{key}.pkl", shard / f"{key}.json"
+
+    def get(self, key: str) -> Optional[CacheEntry]:
+        data_path, meta_path = self._paths(key)
+        if not data_path.exists():
+            self._misses += 1
+            return None
+
+        # Remove empty or corrupted files
+        if data_path.stat().st_size == 0:
+            data_path.unlink()
+            if meta_path.exists():
+                meta_path.unlink()
+            self._misses += 1
+            return None
+
+        try:
+            with open(data_path, "rb") as f:
+                output = pickle.load(f)
+        except (EOFError, pickle.UnpicklingError, Exception):
+            # Corrupt file – delete it and treat as a miss
+            data_path.unlink()
+            if meta_path.exists():
+                meta_path.unlink()
+            self._misses += 1
+            return None
+
+        self._hits += 1
+        metadata = {}
+        if meta_path.exists():
+            try:
+                metadata = json.loads(meta_path.read_text())
+            except Exception:
+                metadata = {}
+        return CacheEntry(output=output, metadata=metadata)
+
+    def put(self, key: str, output: Any, metadata: dict) -> None:
+        data_path, meta_path = self._paths(key)
+        with open(data_path, "wb") as f:
+            pickle.dump(output, f)
+        meta_path.write_text(json.dumps(metadata, indent=2, default=str))
+
+    def stats(self) -> dict:
+        total = self._hits + self._misses
+        hit_rate = self._hits / total if total else 0.0
+        return {"hits": self._hits, "misses": self._misses, "hit_rate": hit_rate}
+
+
+_default_cache: Optional[DiskCache] = None
+
+
+def default_cache() -> DiskCache:
+    global _default_cache
+    if _default_cache is None:
+        _default_cache = DiskCache()
+    return _default_cache

loom/cache_remote.py

@@ -0,0 +1,143 @@
+"""Remote cache backends: S3 and Redis.
+
+To use, install the corresponding extra:
+
+    pip install loomtrace[s3]   # for S3
+    pip install loomtrace[redis]  # for Redis
+"""
+
+from __future__ import annotations
+
+import json
+import pickle
+from typing import Any, Optional
+
+from .cache import Cache, CacheEntry
+
+
+class S3Cache(Cache):
+    """Amazon S3 cache backend.
+
+    Args:
+        bucket: S3 bucket name.
+        prefix: Key prefix (e.g., "loom_cache/").
+        region_name: AWS region, optional.
+        aws_access_key_id, aws_secret_access_key: Optional credentials.
+    """
+
+    def __init__(
+        self,
+        bucket: str,
+        prefix: str = "loom_cache/",
+        region_name: Optional[str] = None,
+        aws_access_key_id: Optional[str] = None,
+        aws_secret_access_key: Optional[str] = None,
+    ):
+        try:
+            import boto3
+        except ImportError as exc:
+            raise ImportError("Install boto3: pip install loomtrace[s3]") from exc
+
+        self.bucket = bucket
+        self.prefix = prefix.rstrip("/") + "/"
+        self.s3 = boto3.client(
+            "s3",
+            region_name=region_name,
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+        )
+        self._hits = 0
+        self._misses = 0
+
+    def _key(self, key: str) -> str:
+        return f"{self.prefix}{key[:2]}/{key}.pkl"
+
+    def _meta_key(self, key: str) -> str:
+        return f"{self.prefix}{key[:2]}/{key}.json"
+
+    def get(self, key: str) -> Optional[CacheEntry]:
+        try:
+            obj = self.s3.get_object(Bucket=self.bucket, Key=self._key(key))
+            output = pickle.loads(obj["Body"].read())
+        except self.s3.exceptions.NoSuchKey:
+            self._misses += 1
+            return None
+
+        self._hits += 1
+        metadata = {}
+        try:
+            meta_obj = self.s3.get_object(Bucket=self.bucket, Key=self._meta_key(key))
+            metadata = json.loads(meta_obj["Body"].read().decode())
+        except self.s3.exceptions.NoSuchKey:
+            pass
+
+        return CacheEntry(output=output, metadata=metadata)
+
+    def put(self, key: str, output: Any, metadata: dict) -> None:
+        self.s3.put_object(
+            Bucket=self.bucket,
+            Key=self._key(key),
+            Body=pickle.dumps(output),
+        )
+        self.s3.put_object(
+            Bucket=self.bucket,
+            Key=self._meta_key(key),
+            Body=json.dumps(metadata).encode(),
+        )
+
+    def stats(self) -> dict:
+        total = self._hits + self._misses
+        return {
+            "hits": self._hits,
+            "misses": self._misses,
+            "hit_rate": self._hits / total if total else 0.0,
+        }
+
+
+class RedisCache(Cache):
+    """Redis cache backend.
+
+    Args:
+        url: Redis URL (e.g., redis://localhost:6379/0).
+        key_prefix: Prefix for all keys.
+    """
+
+    def __init__(self, url: str = "redis://localhost:6379/0", key_prefix: str = "loom:"):
+        try:
+            import redis
+        except ImportError as exc:
+            raise ImportError("Install redis-py: pip install loomtrace[redis]") from exc
+
+        self.redis = redis.from_url(url)
+        self.key_prefix = key_prefix
+        self._hits = 0
+        self._misses = 0
+
+    def _data_key(self, key: str) -> str:
+        return f"{self.key_prefix}{key}:data"
+
+    def _meta_key(self, key: str) -> str:
+        return f"{self.key_prefix}{key}:meta"
+
+    def get(self, key: str) -> Optional[CacheEntry]:
+        data = self.redis.get(self._data_key(key))
+        if data is None:
+            self._misses += 1
+            return None
+        self._hits += 1
+        output = pickle.loads(data)
+        meta = self.redis.get(self._meta_key(key))
+        metadata = json.loads(meta) if meta else {}
+        return CacheEntry(output=output, metadata=metadata)
+
+    def put(self, key: str, output: Any, metadata: dict) -> None:
+        self.redis.set(self._data_key(key), pickle.dumps(output))
+        self.redis.set(self._meta_key(key), json.dumps(metadata))
+
+    def stats(self) -> dict:
+        total = self._hits + self._misses
+        return {
+            "hits": self._hits,
+            "misses": self._misses,
+            "hit_rate": self._hits / total if total else 0.0,
+        }

loom/cli.py

@@ -0,0 +1,92 @@
+"""Command-line interface for Loom.
+
+    loom show <run.json>
+    loom diff <run_a.json> <run_b.json>
+    loom stats <run.json>
+    loom web [--runs-dir DIR] [--host HOST] [--port PORT]
+"""
+from __future__ import annotations
+
+import argparse
+
+from .diff import diff_runs, format_diff
+from .run import Run
+
+
+def cmd_show(args):
+    run = Run.load(args.run_path)
+    print(f"Run: {run.name}  ({run.run_id})")
+    print(f"Nodes: {len(run.nodes)}\n")
+    for i, node in enumerate(run.nodes):
+        tag = "HIT " if node.cache_hit else "MISS"
+        print(f"[{i}] {tag}  {node.step_name}  ({node.node_hash[:10]}...)")
+        print(f"      -> {node.output_repr}")
+    stats = run.stats()
+    print(
+        f"\nCache hit rate: {stats['hit_rate']:.0%}  "
+        f"({stats['cache_hits']}/{stats['total_nodes']})"
+    )
+    print(f"Wall time: {stats['wall_time_s']}s   Time saved by cache: {stats['time_saved_s']}s")
+
+
+def cmd_diff(args):
+    run_a = Run.load(args.run_a)
+    run_b = Run.load(args.run_b)
+    diffs = diff_runs(run_a, run_b)
+    print(format_diff(diffs))
+    changed = [d for d in diffs if d.status != "same"]
+    if changed:
+        print(
+            f"\n{len(changed)} node(s) differ. First divergence: "
+            f"[{changed[0].index}] {changed[0].step_name}"
+        )
+    else:
+        print("\nRuns are identical.")
+
+
+def cmd_stats(args):
+    run = Run.load(args.run_path)
+    for k, v in run.stats().items():
+        print(f"{k}: {v}")
+
+
+def cmd_web(args):
+    try:
+        from loom.web import serve
+    except ImportError as exc:
+        raise SystemExit("Install loomtrace[web] to use the web UI.") from exc
+    serve(runs_dir=args.runs_dir, host=args.host, port=args.port)
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(
+        prog="loom",
+        description="Loom: content-addressed execution engine for agent workflows.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_show = sub.add_parser("show", help="Show all nodes in a saved run")
+    p_show.add_argument("run_path")
+    p_show.set_defaults(func=cmd_show)
+
+    p_diff = sub.add_parser("diff", help="Diff two saved runs node-by-node")
+    p_diff.add_argument("run_a")
+    p_diff.add_argument("run_b")
+    p_diff.set_defaults(func=cmd_diff)
+
+    p_stats = sub.add_parser("stats", help="Show cache hit-rate stats for a run")
+    p_stats.add_argument("run_path")
+    p_stats.set_defaults(func=cmd_stats)
+
+    p_web = sub.add_parser("web", help="Launch web UI")
+    p_web.add_argument("--runs-dir", default=".loom_runs", help="Directory containing run JSONs")
+    p_web.add_argument("--host", default="127.0.0.1", help="Host to bind")
+    p_web.add_argument("--port", type=int, default=5000, help="Port to bind")
+    p_web.set_defaults(func=cmd_web)
+
+    args = parser.parse_args(argv)
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

loom/diff.py

@@ -0,0 +1,58 @@
+"""Node-level diffing between two runs.
+
+See exactly *where* two pipeline executions first diverge — not just
+whether their final outputs differ.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+from .run import Run, Node
+
+
+@dataclass
+class NodeDiff:
+    index: int
+    step_name: str
+    status: str  # "same" | "changed" | "only_in_a" | "only_in_b"
+    node_a: Optional[Node]
+    node_b: Optional[Node]
+
+
+def diff_runs(run_a: Run, run_b: Run) -> List[NodeDiff]:
+    """Compare two runs node-by-node, in call order."""
+    diffs: List[NodeDiff] = []
+    max_len = max(len(run_a.nodes), len(run_b.nodes))
+    for i in range(max_len):
+        na = run_a.nodes[i] if i < len(run_a.nodes) else None
+        nb = run_b.nodes[i] if i < len(run_b.nodes) else None
+        if na is None:
+            diffs.append(NodeDiff(i, nb.step_name, "only_in_b", na, nb))
+        elif nb is None:
+            diffs.append(NodeDiff(i, na.step_name, "only_in_a", na, nb))
+        elif na.node_hash == nb.node_hash:
+            diffs.append(NodeDiff(i, na.step_name, "same", na, nb))
+        else:
+            diffs.append(NodeDiff(i, na.step_name, "changed", na, nb))
+    return diffs
+
+
+def first_divergence(run_a: Run, run_b: Run) -> Optional[NodeDiff]:
+    """The earliest node (in call order) where the two runs differ, or
+    None if they are identical."""
+    for d in diff_runs(run_a, run_b):
+        if d.status != "same":
+            return d
+    return None
+
+
+def format_diff(diffs: List[NodeDiff]) -> str:
+    symbols = {"same": " ", "changed": "~", "only_in_a": "-", "only_in_b": "+"}
+    lines = []
+    for d in diffs:
+        lines.append(f"{symbols[d.status]} [{d.index}] {d.step_name}  ({d.status})")
+        if d.status == "changed":
+            lines.append(f"      a: {d.node_a.output_repr}")
+            lines.append(f"      b: {d.node_b.output_repr}")
+    return "\n".join(lines)