loomcache 1.0.0__py3-none-any.whl

loom/hashing.py

@@ -0,0 +1,83 @@
+"""Content-addressed hashing utilities for Loom.
+
+Every cache key in Loom is derived deterministically from:
+  1. The *source code* of the step function (so editing a prompt
+     invalidates the cache automatically).
+  2. The values (or upstream node hashes) of every argument.
+
+This mirrors how Bazel/Nix key build outputs off content instead of
+timestamps, and is the reason Loom can safely skip re-running a step.
+"""
+from __future__ import annotations
+
+import hashlib
+import inspect
+import json
+import pickle
+from typing import Any, Callable, Optional
+
+
+def _stable_json(value: Any) -> Optional[str]:
+    try:
+        return json.dumps(value, sort_keys=True, default=str)
+    except (TypeError, ValueError):
+        return None
+
+
+def hash_value(value: Any) -> str:
+    """Return a stable content hash for an arbitrary Python value.
+
+    If `value` is itself the traced output of an upstream step, its
+    identity for hashing purposes IS the upstream node's hash — this is
+    what turns a chain of step calls into a real dependency graph.
+
+    Otherwise: try JSON first (fast, stable across processes and
+    machines), fall back to pickle bytes, and finally to repr() for
+    exotic objects that are neither JSON-serializable nor picklable.
+    """
+    node = getattr(value, "_loom_node", None)
+    if node is not None:
+        return f"node:{node.node_hash}"
+
+    as_json = _stable_json(value)
+    if as_json is not None:
+        payload = f"json:{as_json}".encode()
+    else:
+        try:
+            payload = b"pickle:" + pickle.dumps(value)
+        except Exception:
+            payload = f"repr:{value!r}".encode()
+    return hashlib.sha256(payload).hexdigest()
+
+
+def hash_source(func: Callable) -> str:
+    """Hash the source code of a function.
+
+    If the function body changes (e.g. you tweak a prompt template),
+    this hash changes, which changes every node hash downstream of it —
+    Loom will correctly recompute rather than serve a stale cache hit.
+    """
+    try:
+        src = inspect.getsource(func)
+    except (OSError, TypeError):
+        src = f"{func.__module__}.{func.__qualname__}"
+    return hashlib.sha256(src.encode()).hexdigest()
+
+
+def hash_node(
+    step_name: str,
+    source_hash: str,
+    args: tuple,
+    kwargs: dict,
+    extra: str = "",
+) -> str:
+    """Combine step identity + source + argument hashes into one node hash."""
+    parts = [f"step:{step_name}", f"src:{source_hash}"]
+    for a in args:
+        parts.append(f"arg:{hash_value(a)}")
+    for k in sorted(kwargs):
+        parts.append(f"kw:{k}={hash_value(kwargs[k])}")
+    if extra:
+        parts.append(f"extra:{extra}")
+    payload = "|".join(parts).encode()
+    return hashlib.sha256(payload).hexdigest()

loom/langchain.py

@@ -0,0 +1,133 @@
+"""LangChain adapter: cache any Runnable via Loom.
+
+    from loom.langchain import wrap_runnable
+    from langchain.chains import LLMChain
+
+    chain = LLMChain(...)
+    cached_chain = wrap_runnable(chain, name="my_chain")
+
+    with loom.Run("run") as run:
+        result = cached_chain.invoke({"input": "Hello"})
+"""
+
+from __future__ import annotations
+
+import hashlib
+import inspect
+import time
+from typing import Any, Optional, TypeVar
+
+from loom import current_run, default_cache, attach_node, hashing
+from loom.run import Node
+from loom.cache import Cache
+
+try:
+    from langchain_core.runnables import Runnable, RunnableConfig
+except ImportError:
+    # Stub for optional dependency
+    class Runnable:  # type: ignore[no-redef]
+        pass
+
+    RunnableConfig = Any
+
+T = TypeVar("T", bound=Runnable)
+
+
+class CachedRunnable(Runnable):
+    """A Runnable that delegates to the original and caches invoke results."""
+
+    def __init__(
+        self,
+        runnable: Runnable,
+        *,
+        name: Optional[str] = None,
+        cache: Optional[Cache] = None,
+        version: Optional[str] = "",
+        source_hash: Optional[str] = None,
+    ):
+        self._runnable = runnable
+        self._cache = cache or default_cache()
+        self._name = name or getattr(runnable, "name", None) or runnable.__class__.__name__
+        self._version = version or ""
+
+        if source_hash is None:
+            try:
+                src = inspect.getsource(runnable.invoke)
+                source_hash = hashlib.sha256(src.encode()).hexdigest()
+            except (OSError, TypeError):
+                source_hash = "unknown"
+        self._source_hash = source_hash
+
+    def _cache_key(self, input: Any, config: Optional[RunnableConfig] = None) -> str:
+        identity = f"langchain:{self._name}:{self._source_hash}:{self._version}"
+        input_hash = hashing.hash_value(input)
+        return hashing.hash_node(
+            step_name=self._name,
+            source_hash=identity,
+            args=(input_hash,),
+            kwargs={},
+            extra="langchain",
+        )
+
+    def invoke(self, input: Any, config: Optional[RunnableConfig] = None, **kwargs) -> Any:
+        run = current_run()
+        if run is None:
+            return self._runnable.invoke(input, config, **kwargs)
+
+        key = self._cache_key(input, config)
+        entry = self._cache.get(key)
+
+        if entry is not None:
+            output = entry.output
+            cache_hit = True
+            duration = 0.0
+            time_saved = entry.metadata.get("duration_s", 0.0)
+        else:
+            start = time.time()
+            output = self._runnable.invoke(input, config, **kwargs)
+            duration = time.time() - start
+            self._cache.put(
+                key,
+                output,
+                metadata={
+                    "step_name": self._name,
+                    "timestamp": time.time(),
+                    "duration_s": duration,
+                },
+            )
+            cache_hit = False
+            time_saved = 0.0
+
+        node = Node(
+            node_hash=key,
+            step_name=self._name,
+            parents=[],
+            cache_hit=cache_hit,
+            duration_s=duration,
+            time_saved_s=time_saved,
+            timestamp=time.time(),
+            args_repr=repr(input),
+            kwargs_repr=repr(config),
+            output_repr=repr(output),
+        )
+        run.record(node)
+        return attach_node(output, node)
+
+    # Async and batch methods are not yet implemented (v1.0 roadmap).
+    async def ainvoke(self, input: Any, config: Optional[RunnableConfig] = None, **kwargs) -> Any:
+        raise NotImplementedError("Async caching is not yet supported for Runnables.")
+
+    def batch(self, inputs, config=None, **kwargs):
+        raise NotImplementedError("Batch caching is not yet supported.")
+
+
+def wrap_runnable(
+    runnable: T,
+    *,
+    name: Optional[str] = None,
+    cache: Optional[Cache] = None,
+    version: Optional[str] = "",
+    source_hash: Optional[str] = None,
+) -> T:
+    """Wrap a LangChain Runnable with Loom caching."""
+    return CachedRunnable(runnable, name=name, cache=cache, version=version, source_hash=source_hash)

loom/run.py

@@ -0,0 +1,250 @@
+"""Core execution-tracking primitives: steps, nodes, and runs.
+
+    @loom.step
+    def plan(query: str) -> str:
+        return llm_call(f"Plan: {query}")
+
+    with loom.Run("my-pipeline") as run:
+        result = plan("hello world")
+    run.save()
+
+Every `@loom.step` call made while a `Run` is active is hashed,
+looked up in the cache, executed only on a miss, and recorded onto
+`run.nodes` in order.
+"""
+from __future__ import annotations
+
+import contextvars
+import functools
+import json
+import time
+import uuid
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
+
+from . import hashing
+from .cache import Cache, default_cache
+from .tracing import attach_node, get_node, TracedBox, unwrap_recursive
+
+_current_run: contextvars.ContextVar = contextvars.ContextVar(
+    "loom_current_run", default=None
+)
+
+
+def current_run() -> Optional["Run"]:
+    """The `Run` currently active via a `with loom.Run(...)` block, if any."""
+    return _current_run.get()
+
+
+@dataclass
+class Node:
+    """One recorded execution (or cache hit) of a single `@step` call."""
+
+    node_hash: str
+    step_name: str
+    parents: list
+    cache_hit: bool
+    duration_s: float
+    time_saved_s: float
+    timestamp: float
+    args_repr: str
+    kwargs_repr: str
+    output_repr: str
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+def _safe_repr(value: Any, limit: int = 300) -> str:
+    try:
+        text = repr(value)
+    except Exception:
+        text = f"<unrepr-able {type(value).__name__}>"
+    return text if len(text) <= limit else text[: limit - 3] + "..."
+
+
+def _unwrap(value: Any) -> Any:
+    if isinstance(value, TracedBox):
+        return value.unwrap()
+    return value
+
+
+class Run:
+    """A single execution of a pipeline.
+
+    Use as a context manager. Every `@step` call made while this Run is
+    active gets recorded, in order, into `run.nodes`.
+    """
+
+    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
+
+    # -- context manager ------------------------------------------------
+    def __enter__(self) -> "Run":
+        self._token = _current_run.set(self)
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        _current_run.reset(self._token)
+        return False
+
+    # -- recording --------------------------------------------------------
+    def record(self, node: Node) -> None:
+        self.nodes.append(node)
+
+    # -- persistence --------------------------------------------------------
+    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
+
+    @classmethod
+    def load(cls, path: Union[str, Path]) -> "Run":
+        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
+
+    # -- introspection --------------------------------------------------------
+    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),
+        }
+
+    def fork(self, pipeline_fn: Callable, name: Optional[str] = None, **new_kwargs) -> "Run":
+        """Re-run `pipeline_fn` (typically with one input changed).
+
+        Because every `@step` call is content-addressed, any step whose
+        hash is unaffected by the change is served from cache instantly;
+        only the changed step and everything downstream of it actually
+        re-executes. This mirrors how `bazel build` or `make` only
+        rebuild the targets whose inputs changed.
+        """
+        forked = Run(name or f"{self.name}-fork", cache=self.cache, root=self.root)
+        with forked:
+            pipeline_fn(**new_kwargs)
+        return forked
+
+    def __repr__(self) -> str:
+        return f"<Run {self.run_id} nodes={len(self.nodes)}>"
+
+
+def step(func: Optional[Callable] = None, *, cache: Optional[Cache] = None):
+    """Decorator that turns a plain function into a cached, tracked step.
+
+    Every call is:
+      1. Hashed from (source code of the function + argument values, or
+         upstream node hashes for arguments that are themselves the
+         traced output of another step).
+      2. Looked up in the cache — on a hit, the cached output is
+         returned immediately with zero re-execution.
+      3. On a miss, executed normally, cached, and recorded.
+    """
+
+    def decorator(f: Callable) -> Callable:
+        source_hash = hashing.hash_source(f)
+
+        @functools.wraps(f)
+        def wrapper(*args, **kwargs):
+            run = current_run()
+            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)
+            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 (unwrapped)
+                cache_hit = True
+                duration = 0.0
+                time_saved = float(entry.metadata.get("duration_s", 0.0))
+                clean_output = output
+            else:
+                # Unwrap arguments before calling the function
+                call_args = tuple(_unwrap(a) for a in args)
+                call_kwargs = {k: _unwrap(v) for k, v in kwargs.items()}
+                t0 = time.time()
+                raw_output = f(*call_args, **call_kwargs)
+                duration = time.time() - t0
+
+                # Recursively unwrap any traced objects to get a clean, pickleable value
+                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
+
+            # Create the Node object (for both cache hit and miss)
+            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=_safe_repr(args),
+                kwargs_repr=_safe_repr(kwargs),
+                output_repr=_safe_repr(clean_output),
+            )
+
+            if run is not None:
+                run.record(node)
+
+            # Attach the node to the clean output and return traced version
+            return attach_node(clean_output, node)
+
+        wrapper._loom_step = True  # type: ignore[attr-defined]
+        wrapper._loom_source_hash = source_hash  # type: ignore[attr-defined]
+        wrapper.__wrapped_func__ = f  # type: ignore[attr-defined]
+        return wrapper
+
+    if func is not None:
+        return decorator(func)
+    return decorator

loom/trace.py

@@ -0,0 +1,15 @@
+"""Alias module: `from loom import trace` then use `trace.step`, `trace.Run`.
+
+Purely a namespacing convenience some people prefer stylistically —
+identical in behavior to importing `step` / `Run` from `loom` directly.
+
+    from loom import trace
+
+    @trace.step
+    def plan(query: str) -> str:
+        ...
+
+    with trace.Run("my-pipeline") as run:
+        ...
+"""
+from .run import Node, Run, current_run, step  # noqa: F401

loom/tracing.py

@@ -0,0 +1,132 @@
+"""Transparent value wrapping so Loom can track data lineage.
+
+When a @step function returns a value, Loom attaches an invisible
+`_loom_node` marker to it. If that value is later passed into another
+@step call, Loom sees the marker and knows there is a dependency edge
+between the two nodes — without requiring the user to declare any
+dependencies manually. This is what lets a plain chain of Python
+function calls become a real, hashable DAG.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+_TRACED_SUBCLASS_CACHE: dict = {}
+
+
+def _traced_subclass(base: type) -> type:
+    if base in _TRACED_SUBCLASS_CACHE:
+        return _TRACED_SUBCLASS_CACHE[base]
+
+    class _Traced(base):  # type: ignore[misc,valid-type]
+        __slots__ = ("_loom_node",)
+
+    _Traced.__name__ = f"Traced{base.__name__.capitalize()}"
+    _TRACED_SUBCLASS_CACHE[base] = _Traced
+    return _Traced
+
+
+class TracedBox:
+    """Fallback wrapper for values that can't be subclassed or tagged
+    directly (e.g. `bool`, or some C-extension objects).
+
+    Behaves like the wrapped value for repr/equality/truthiness/hash,
+    and transparently forwards attribute access, so downstream code can
+    usually treat it exactly like the original object. Call `.unwrap()`
+    to get the raw value back explicitly.
+    """
+
+    __slots__ = ("_loom_value", "_loom_node")
+
+    def __init__(self, value: Any, node: Any):
+        object.__setattr__(self, "_loom_value", value)
+        object.__setattr__(self, "_loom_node", node)
+
+    def __getattr__(self, name):
+        return getattr(object.__getattribute__(self, "_loom_value"), name)
+
+    def __repr__(self):
+        return repr(object.__getattribute__(self, "_loom_value"))
+
+    def __str__(self):
+        return str(object.__getattribute__(self, "_loom_value"))
+
+    def __eq__(self, other):
+        return object.__getattribute__(self, "_loom_value") == other
+
+    def __hash__(self):
+        return hash(object.__getattribute__(self, "_loom_value"))
+
+    def __bool__(self):
+        return bool(object.__getattribute__(self, "_loom_value"))
+
+    def unwrap(self):
+        return object.__getattribute__(self, "_loom_value")
+
+
+# CPython allows adding `__slots__` to subclasses of `str`, `tuple`, and
+# `frozenset` (their C layout has room for it), but NOT to subclasses of
+# `int` or `float` (fixed-size numeric layout, no room for extra slots) --
+# those go through `TracedBox` instead, same as `bool`.
+_DIRECTLY_TAGGABLE_IMMUTABLES = (str, tuple, frozenset)
+_DIRECTLY_TAGGABLE_MUTABLES = (list, dict, set)
+
+
+def attach_node(value: Any, node: Any) -> Any:
+    """Attach a lineage marker to `value`, returning a traced version of it.
+
+    Supports directly: str, tuple, frozenset, list, dict, set, and any
+    plain object with a `__dict__`. Everything else (notably `int`,
+    `float`, and `bool`, whose fixed C layout doesn't allow extra
+    attributes) falls back to the transparent `TracedBox` wrapper.
+    """
+    if getattr(value, "_loom_node", None) is not None:
+        return value  # already traced (e.g. a step just returns its input)
+
+    cls = type(value)
+
+    if cls in _DIRECTLY_TAGGABLE_IMMUTABLES or cls in _DIRECTLY_TAGGABLE_MUTABLES:
+        traced_cls = _traced_subclass(cls)
+        obj = traced_cls(value)
+        object.__setattr__(obj, "_loom_node", node)
+        return obj
+
+    try:
+        object.__setattr__(value, "_loom_node", node)
+        return value
+    except (AttributeError, TypeError):
+        pass
+
+    return TracedBox(value, node)
+
+
+def get_node(value: Any):
+    return getattr(value, "_loom_node", None)
+
+
+def unwrap(value: Any) -> Any:
+    if isinstance(value, TracedBox):
+        return value.unwrap()
+    return value
+
+
+def unwrap_recursive(value: Any) -> Any:
+    """Recursively remove all Loom tracing from a value, returning a plain Python object."""
+    if hasattr(value, "_loom_node"):
+        # It's a traced object. Convert to its base type.
+        if isinstance(value, TracedBox):
+            return unwrap_recursive(value.unwrap())
+        else:
+            # Subclass of built‑in (e.g., TracedStr). Get the base class and convert.
+            base = value.__class__.__base__
+            try:
+                return unwrap_recursive(base(value))
+            except Exception:
+                # Fallback – return the value as‑is (might still break pickle)
+                return value
+    elif isinstance(value, dict):
+        return {unwrap_recursive(k): unwrap_recursive(v) for k, v in value.items()}
+    elif isinstance(value, (list, tuple, set, frozenset)):
+        return type(value)(unwrap_recursive(v) for v in value)
+    else:
+        return value