vitnify 0.1.0__py3-none-any.whl

vitnify/_vendor/pck/cas.py

@@ -0,0 +1,181 @@
+"""Content-addressed store + Merkle ledger — the provenance backbone of PCK.
+
+Every source chunk is addressed by its cryptographic hash. A Merkle tree over those
+hashes yields a single ``root`` that attests the whole corpus. Any single chunk can
+then carry an *inclusion proof*: a short path of sibling hashes that recomputes the
+root. The guarantees this buys, and why they matter for AI answers:
+
+  * **Tamper-evidence** — change one byte of one chunk and the root changes.
+  * **Independent verifiability** — anyone holding the root can check a chunk's proof
+    without trusting the party that produced it (no model, no server, no secret key).
+  * **Fabrication is uncredentialed** — a made-up "source" has no leaf in the tree, so
+    it cannot produce a valid proof. This is the cryptographic half of catching a
+    hallucination: a grounded claim points at a chunk with a valid proof; an invented
+    one cannot.
+
+This mirrors the BLAKE3 + Merkle design of a native runtime's ``cas_store`` (whose real
+Rust code is exercised in ``research/pck_provenance/``); here we keep a dependency-light
+pure-Python implementation that uses ``blake3`` when available and falls back to
+``hashlib.blake2b`` otherwise. The wire format of a proof is identical either way.
+"""
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, field
+from typing import Optional
+
+try:  # BLAKE3 if the user has it (matches the substrate); else a stdlib fallback.
+    import blake3 as _blake3  # type: ignore
+
+    _HASH_NAME = "blake3"
+
+    def _digest(b: bytes) -> str:
+        return _blake3.blake3(b).hexdigest()
+except Exception:  # pragma: no cover - depends on environment
+    _HASH_NAME = "blake2b-256"
+
+    def _digest(b: bytes) -> str:
+        return hashlib.blake2b(b, digest_size=32).hexdigest()
+
+
+HASH_NAME = _HASH_NAME
+
+
+def hash_bytes(b: bytes) -> str:
+    """Content address of raw bytes (hex)."""
+    return _digest(b)
+
+
+def hash_text(s: str) -> str:
+    """Content address of a text chunk (UTF-8, hex)."""
+    return _digest(s.encode("utf-8"))
+
+
+def _hash_pair(left: str, right: str) -> str:
+    """Internal Merkle node = hash of the concatenated child hex digests."""
+    return _digest((left + right).encode("utf-8"))
+
+
+# --- inclusion proof -------------------------------------------------------------
+
+@dataclass(frozen=True)
+class ProofStep:
+    """One level of a Merkle inclusion path."""
+
+    sibling: str      # hex digest of the sibling node
+    on_right: bool    # True if the *current* node is the right child at this level
+
+    def to_json(self) -> dict:
+        return {"sibling": self.sibling, "on_right": self.on_right}
+
+    @staticmethod
+    def from_json(d: dict) -> "ProofStep":
+        return ProofStep(sibling=d["sibling"], on_right=bool(d["on_right"]))
+
+
+@dataclass(frozen=True)
+class InclusionProof:
+    """A self-contained proof that ``leaf`` is committed under ``root``.
+
+    Carries everything a third party needs to re-verify against a *published* root:
+    the leaf hash, the sibling path, and the index. No tree, server, or key required.
+    """
+
+    leaf: str
+    index: int
+    path: list[ProofStep]
+    hash_name: str = HASH_NAME
+
+    def verify(self, root: str) -> bool:
+        return verify_proof(self.leaf, self.path, root)
+
+    def to_json(self) -> dict:
+        return {
+            "leaf": self.leaf,
+            "index": self.index,
+            "hash": self.hash_name,
+            "path": [s.to_json() for s in self.path],
+        }
+
+    @staticmethod
+    def from_json(d: dict) -> "InclusionProof":
+        return InclusionProof(
+            leaf=d["leaf"],
+            index=int(d["index"]),
+            path=[ProofStep.from_json(s) for s in d["path"]],
+            hash_name=d.get("hash", HASH_NAME),
+        )
+
+
+def verify_proof(leaf: str, path: list[ProofStep], root: str) -> bool:
+    """Recompute the root from ``leaf`` and ``path``; True iff it equals ``root``."""
+    h = leaf
+    for step in path:
+        h = _hash_pair(step.sibling, h) if step.on_right else _hash_pair(h, step.sibling)
+    return h == root
+
+
+# --- the tree / store ------------------------------------------------------------
+
+@dataclass
+class MerkleCAS:
+    """Content-addressed store with a Merkle commitment over its chunks.
+
+    Odd layers duplicate the final node (Bitcoin-style), so a proof always exists for
+    every leaf. Construction is deterministic in the chunk order.
+    """
+
+    chunks: list[str]
+    leaves: list[str] = field(default_factory=list)
+    layers: list[list[str]] = field(default_factory=list)
+    _index_of_hash: dict[str, int] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        if not self.chunks:
+            raise ValueError("MerkleCAS requires at least one chunk")
+        self.leaves = [hash_text(c) for c in self.chunks]
+        # first occurrence wins if a chunk repeats
+        for i, h in enumerate(self.leaves):
+            self._index_of_hash.setdefault(h, i)
+        self.layers = self._build(self.leaves)
+
+    @staticmethod
+    def _build(leaves: list[str]) -> list[list[str]]:
+        layers = [leaves[:]]
+        while len(layers[-1]) > 1:
+            cur = layers[-1]
+            nxt = [
+                _hash_pair(cur[i], cur[i + 1] if i + 1 < len(cur) else cur[i])
+                for i in range(0, len(cur), 2)
+            ]
+            layers.append(nxt)
+        return layers
+
+    @property
+    def root(self) -> str:
+        return self.layers[-1][0]
+
+    def index_of(self, chunk_or_hash: str) -> Optional[int]:
+        """Index of a chunk by its raw text or by its hash; None if absent."""
+        if chunk_or_hash in self._index_of_hash:
+            return self._index_of_hash[chunk_or_hash]
+        return self._index_of_hash.get(hash_text(chunk_or_hash))
+
+    def contains(self, chunk_or_hash: str) -> bool:
+        return self.index_of(chunk_or_hash) is not None
+
+    def prove_index(self, idx: int) -> InclusionProof:
+        if not 0 <= idx < len(self.leaves):
+            raise IndexError(idx)
+        path: list[ProofStep] = []
+        i = idx
+        for layer in self.layers[:-1]:
+            sib = i ^ 1 if (i ^ 1) < len(layer) else i  # duplicated node is its own sibling
+            path.append(ProofStep(sibling=layer[sib], on_right=bool(i & 1)))
+            i //= 2
+        return InclusionProof(leaf=self.leaves[idx], index=idx, path=path)
+
+    def prove(self, chunk_or_hash: str) -> Optional[InclusionProof]:
+        """Inclusion proof for a chunk (by text or hash); None if not in the corpus."""
+        idx = self.index_of(chunk_or_hash)
+        return None if idx is None else self.prove_index(idx)

vitnify/capability.py

@@ -0,0 +1,38 @@
+"""Userspace, kernel-free capability broker (adapted from the PoC).
+Containment by construction: every side-effecting tool call routes through here,
+and a tool the run wasn't granted is simply unreachable -- a jailbreak that reaches
+an ungranted tool hits a wall, not a classifier. Each call is recorded to the event
+log (with its result, so replay can re-inject it). No ML, no score, no bypass.
+"""
+from __future__ import annotations
+import hashlib
+from .events import EventLog, Kind
+
+def _rh(x) -> str:
+    return hashlib.blake2b(repr(x).encode(), digest_size=16).hexdigest()
+
+class Broker:
+    def __init__(self, capabilities, tools: dict, log: EventLog, replay=None):
+        self.caps = set(capabilities)
+        self.tools = tools
+        self.log = log
+        self.replay = replay        # if set: list of recorded tool results to re-inject
+
+    def call(self, tool: str, *args):
+        allowed = tool in self.caps
+        if not allowed:
+            self.log.append(Kind.TOOL_CALL, {"tool": tool, "args": list(args), "decision": "DENY"})
+            return False, None
+        if self.replay is not None:
+            # external effect: DO NOT re-run the live tool; serve the recorded result
+            result = self.replay.pop(0)
+        else:
+            result = self.tools[tool](*args)
+        self.log.append(Kind.TOOL_CALL, {"tool": tool, "args": list(args),
+                                         "decision": "ALLOW", "result": result, "result_hash": _rh(result)})
+        return True, result
+
+def recorded_tool_results(log: EventLog):
+    """Extract the ordered list of ALLOW tool results, to feed a replay Broker."""
+    return [e.payload["result"] for e in log.events
+            if e.kind == Kind.TOOL_CALL.value and e.payload.get("decision") == "ALLOW"]

vitnify/certificate.py

@@ -0,0 +1,139 @@
+"""vitnify-receipt v1 -- the signed, self-verifying execution receipt.
+
+One receipt binds an agent run into a single object:
+
+    program hash | granted capabilities | hash-chained event log (Merkle root + head)
+    | the engine's per-step model-computation digests
+
+so an action can never be detached from the computation and context that produced
+it. The receipt digest is BLAKE3 over a canonical body (matching the vitni-tensor
+engine and the event log); the receipt is signed with ed25519 and embeds its own
+public key, so ANYONE can verify it offline with no model, no network, and no secret.
+
+Two verification levels:
+  * level 1 (integrity): recompute the Merkle root, chain head, and event count from
+    the raw events, and check the ed25519 signature. No model required.
+  * level 2 (recomputation): re-run each bound model step through vitni-tensor and
+    confirm every model_digest reproduces bit-for-bit.
+
+verify_certificate() implements level 1 and is fully independent. HMAC signing is
+retained only as a fallback for environments without an asymmetric key.
+"""
+from __future__ import annotations
+import os, sys, json, hmac
+from dataclasses import dataclass, field, asdict
+from .events import EventLog
+
+from ._vendor.pck.cas import MerkleCAS  # noqa: E402
+
+FORMAT = "vitnify-receipt v1"
+
+try:
+    import blake3 as _blake3
+    def _digest32(b: bytes) -> str:
+        return _blake3.blake3(b).hexdigest()
+except ImportError:
+    import hashlib
+    def _digest32(b: bytes) -> str:
+        return hashlib.blake2b(b, digest_size=32).hexdigest()
+
+# cryptography is imported lazily inside the ed25519 helpers so the fallback path
+# still runs on a stock Python without the dependency.
+
+
+def _canon(o) -> str:
+    return json.dumps(o, sort_keys=True, separators=(",", ":"))
+
+
+def gen_ed25519():
+    """Return (private_key, public_key_hex). Persist the private key to sign future runs."""
+    from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+    priv = Ed25519PrivateKey.generate()
+    return priv, priv.public_key().public_bytes_raw().hex()
+
+
+@dataclass
+class ExecutionCertificate:
+    program_hash: str
+    capabilities: list
+    event_root: str
+    n_events: int
+    head_hash: str
+    model_digests: list = field(default_factory=list)   # engine per-step digests, bound by the log
+    v: str = FORMAT
+    sig: str | None = None
+    sig_alg: str = "none"
+    pubkey: str | None = None       # ed25519 public key (hex) -- verifier needs no secret
+
+    def body(self) -> dict:
+        return {"v": self.v, "program_hash": self.program_hash,
+                "capabilities": sorted(self.capabilities), "event_root": self.event_root,
+                "n_events": self.n_events, "head_hash": self.head_hash,
+                "model_digests": list(self.model_digests)}
+
+    def digest(self) -> str:
+        return _digest32(_canon(self.body()).encode())
+
+    def sign_ed25519(self, priv) -> "ExecutionCertificate":
+        self.sig = priv.sign(bytes.fromhex(self.digest())).hex()
+        self.pubkey = priv.public_key().public_bytes_raw().hex()
+        self.sig_alg = "ed25519"
+        return self
+
+    def sign_hmac(self, key: bytes) -> "ExecutionCertificate":
+        import hashlib
+        self.sig = hmac.new(key, self.digest().encode(), hashlib.blake2b).hexdigest()
+        self.sig_alg = "hmac-blake2b"
+        return self
+
+    def to_json(self) -> str:
+        return _canon(asdict(self))
+
+
+def issue_certificate(program_hash, capabilities, log: EventLog, priv=None, key: bytes | None = None):
+    """Build and sign a vitnify-receipt for a completed run.
+
+    The model-computation digests are read from the log's llm_call events, so the
+    receipt binds exactly the model steps that were recorded. ed25519 (self-verifying)
+    is the canonical signing path; pass `priv`.
+    """
+    cas = MerkleCAS(log.chunks())
+    cert = ExecutionCertificate(
+        program_hash, sorted(set(capabilities)), cas.root, len(log), log.head(),
+        model_digests=log.model_digests(),
+    )
+    if priv is not None:
+        cert.sign_ed25519(priv)
+    elif key is not None:
+        cert.sign_hmac(key)
+    return cert, cas
+
+
+def verify_certificate(cert: ExecutionCertificate, log: EventLog, key: bytes | None = None) -> dict:
+    """Level-1 verification: recompute everything from the raw events; trust nothing.
+
+    ed25519 receipts verify against their OWN embedded public key -- no secret required.
+    (Level 2 -- recomputing each model_digest through vitni-tensor -- is a separate step
+    that additionally needs the model weights.)
+    """
+    cas = MerkleCAS(log.chunks())
+    checks = {
+        "format": cert.v == FORMAT,
+        "root_matches": cas.root == cert.event_root,
+        "head_matches": log.head() == cert.head_hash,
+        "count_matches": len(log) == cert.n_events,
+        "model_digests_match": log.model_digests() == list(cert.model_digests),
+    }
+    if cert.sig_alg == "ed25519":
+        try:
+            from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
+            Ed25519PublicKey.from_public_bytes(bytes.fromhex(cert.pubkey)).verify(
+                bytes.fromhex(cert.sig), bytes.fromhex(cert.digest()))
+            checks["sig_valid"] = True
+        except Exception:
+            checks["sig_valid"] = False
+    elif cert.sig_alg == "hmac-blake2b" and key is not None:
+        import hashlib
+        checks["sig_valid"] = cert.sig == hmac.new(key, cert.digest().encode(), hashlib.blake2b).hexdigest()
+    checks["ok"] = all(checks.values())
+    return checks

vitnify/deterministic_llm.py

@@ -0,0 +1,64 @@
+"""Deterministic local inference -- the piece that makes a replay CERTIFIABLE.
+Greedy decode over TinyLlama, with an optional batch-invariant linear (your bi_linear
+work). The certificate commits to a hash of the LOGITS at each step, not just the token:
+tokens rarely flip under batch load, but logits always shift -- so a bit-identical
+certificate can only be reproduced when inference is batch-invariant.
+"""
+from __future__ import annotations
+import hashlib
+from contextlib import nullcontext
+import torch, torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+_ORIG_LINEAR = F.linear
+
+def _bi_linear(x, w, b=None):
+    """Fixed-chunk fp32 reduction -> output independent of batch size (bit-identical)."""
+    d, CH = x.shape[-1], 128
+    xf, wf, acc = x.float(), w.float(), None
+    for c in range(0, d, CH):
+        p = xf[..., c:c+CH] @ wf[:, c:c+CH].t()
+        acc = p if acc is None else acc + p
+    out = (acc + b.float()) if b is not None else acc
+    return out.to(x.dtype)
+
+class batch_invariant:
+    def __enter__(self): F.linear = _bi_linear
+    def __exit__(self, *a): F.linear = _ORIG_LINEAR
+
+def hash_tensor(t: torch.Tensor) -> str:
+    return hashlib.blake2b(t.detach().to("cpu").contiguous().numpy().tobytes(), digest_size=16).hexdigest()
+
+class DeterministicLM:
+    def __init__(self, name="TinyLlama/TinyLlama-1.1B-Chat-v1.0", device=None):
+        self.dev = device or ("mps" if torch.backends.mps.is_available() else "cpu")
+        self.tok = AutoTokenizer.from_pretrained(name, local_files_only=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            name, torch_dtype=torch.float16, local_files_only=True).eval().to(self.dev)
+        self.V = self.model.config.vocab_size
+
+    def chat_prompt(self, messages) -> str:
+        """Render messages with the model's own chat template (Qwen, Llama, etc.)."""
+        return self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+    @torch.no_grad()
+    def generate(self, prompt: str, n_new=8, batch_load=0, invariant=False):
+        """Greedy decode. batch_load simulates production co-batching (padded neighbors).
+        Returns (tokens, per_step_logit_hashes). Our request is always row 0."""
+        ids = self.tok(prompt, return_tensors="pt").input_ids.to(self.dev)
+        ctx = batch_invariant() if invariant else nullcontext()
+        toks, hashes = [], []
+        with ctx:
+            cur = ids
+            for _ in range(n_new):
+                if batch_load:
+                    pad = torch.randint(0, self.V, (batch_load, cur.shape[1]), device=self.dev)
+                    batch = torch.cat([cur, pad], 0)
+                else:
+                    batch = cur
+                logits = self.model(batch).logits[0, -1]
+                hashes.append(hash_tensor(logits))
+                nxt = int(logits.argmax())
+                toks.append(nxt)
+                cur = torch.cat([cur, torch.tensor([[nxt]], device=self.dev)], 1)
+        return toks, hashes

vitnify/engine.py

@@ -0,0 +1,46 @@
+"""Bridge to the vitni-tensor engine.
+
+Runs a model step through the deterministic engine and returns its
+`vitnify-receipt v1` model-computation digest, so the SDK can bind what the model
+actually computed into the signed execution receipt (`EventLog.append_llm_call`).
+
+The engine is the `vitni-receipt` binary (built from the vitni-tensor crate). Point
+`VITNI_RECEIPT_BIN` at it, or put it on PATH.
+"""
+from __future__ import annotations
+import os, json, subprocess
+
+try:
+    import blake3 as _blake3
+    def _h32(b: bytes) -> str:
+        return _blake3.blake3(b).hexdigest()
+except ImportError:
+    import hashlib
+    def _h32(b: bytes) -> str:
+        return hashlib.blake2b(b, digest_size=32).hexdigest()
+
+DEFAULT_BIN = os.environ.get("VITNI_RECEIPT_BIN", "vitni-receipt")
+
+
+def prompt_hash(tokens: list[int]) -> str:
+    return _h32(json.dumps(list(tokens), separators=(",", ":")).encode())
+
+
+class Engine:
+    """Deterministic model backend. Every `run` is reproducible bit-for-bit and yields
+    a cross-vendor model-computation digest to bind into the receipt."""
+
+    def __init__(self, gguf: str, model_id: str, binary: str = DEFAULT_BIN):
+        self.gguf = gguf
+        self.model_id = model_id
+        self.binary = binary
+
+    def run(self, prompt_tokens: list[int], n_new: int = 16) -> dict:
+        """Return {"tokens": [...], "model_digest": hex, "weights_hash": hex}."""
+        proc = subprocess.run(
+            [self.binary, "--gguf", self.gguf,
+             "--prompt", ",".join(str(t) for t in prompt_tokens),
+             "--n", str(n_new), "--model-id", self.model_id],
+            capture_output=True, text=True, check=True,
+        )
+        return json.loads(proc.stdout)

vitnify/events.py

@@ -0,0 +1,95 @@
+"""Canonical, hash-chained agent event log.
+
+Every nondeterministic input/output of an agent run becomes one canonical,
+hash-linked Event. The chain gives append-only tamper-evidence during recording;
+certificate.py Merkle-commits the same chunks for inclusion proofs and binds the
+root into the signed vitnify-receipt.
+
+Hashing is BLAKE3 throughout (matching the vitni-tensor engine and the
+`vitnify-receipt v1` format), so the whole receipt is one hash family end to end.
+"""
+from __future__ import annotations
+import json
+from dataclasses import dataclass, asdict
+from enum import Enum
+
+try:
+    import blake3 as _blake3
+    def _h32(b: bytes) -> str:
+        return _blake3.blake3(b).hexdigest()
+except ImportError:  # stdlib fallback; wire format is identical either way
+    import hashlib
+    def _h32(b: bytes) -> str:
+        return hashlib.blake2b(b, digest_size=32).hexdigest()
+
+
+def canon(obj) -> str:
+    return json.dumps(obj, sort_keys=True, separators=(",", ":"))
+
+
+def h(s: str) -> str:
+    return _h32(s.encode())
+
+
+class Kind(str, Enum):
+    LLM_CALL   = "llm_call"     # (prompt_hash, tokens, seed, model_digest) -- the model step
+    TOOL_CALL  = "tool_call"    # (tool, args, decision, result) -- external, re-injected on replay
+    ENTROPY    = "entropy"      # (source, value) -- RNG/clock draw, from-log on replay
+    AGENT_STEP = "agent_step"   # (state) -- orchestration checkpoint
+
+
+@dataclass
+class Event:
+    i: int
+    kind: str
+    payload: dict
+    prev: str
+    def canonical(self) -> str:
+        return canon({"i": self.i, "kind": self.kind, "payload": self.payload, "prev": self.prev})
+    @property
+    def hash(self) -> str:
+        return h(self.canonical())
+
+
+class EventLog:
+    def __init__(self):
+        self.events: list[Event] = []
+
+    def append(self, kind, payload: dict) -> Event:
+        prev = self.events[-1].hash if self.events else "genesis"
+        ev = Event(len(self.events), kind.value if isinstance(kind, Kind) else str(kind), payload, prev)
+        self.events.append(ev)
+        return ev
+
+    def append_llm_call(self, prompt_hash: str, tokens: list, seed, model_digest: str) -> Event:
+        """Record a model step and bind its vitni-tensor model-computation digest.
+
+        `model_digest` is the engine's `vitnify-receipt v1` digest for this forward
+        pass. Committing it here (and thus in the receipt's Merkle root) is what binds
+        the deterministic model recomputation to the agent run.
+        """
+        return self.append(Kind.LLM_CALL, {
+            "prompt_hash": prompt_hash, "tokens": list(tokens),
+            "seed": seed, "model_digest": model_digest,
+        })
+
+    def model_digests(self) -> list[str]:
+        """Ordered engine model-computation digests bound by this log (for L2 recompute)."""
+        return [e.payload["model_digest"] for e in self.events
+                if e.kind == Kind.LLM_CALL.value and "model_digest" in e.payload]
+
+    def chunks(self) -> list[str]:
+        return [e.canonical() for e in self.events]
+
+    def head(self) -> str:
+        return self.events[-1].hash if self.events else "genesis"
+
+    def to_json(self) -> str:
+        return canon([asdict(e) for e in self.events])
+
+    @classmethod
+    def from_events(cls, events: list[Event]) -> "EventLog":
+        log = cls(); log.events = events; return log
+
+    def __len__(self):
+        return len(self.events)

vitnify/langgraph_adapter.py

@@ -0,0 +1,101 @@
+"""LangGraph adapter — make a real LangGraph agent recordable, replayable, certifiable,
+and contained, with no change to the agent's logic.
+
+Two integration points:
+  1. build_react_graph(session, ...) -> a compiled LangGraph StateGraph whose *agent*
+     node calls session.llm() and whose *tool* node calls session.tool() (the capability
+     wall). Because everything flows through the VitniReplay Session, the whole graph is
+     recorded, deterministically replayable, and certified — for free.
+  2. VitniReplayCallback(session) -> a LangChain BaseCallbackHandler that records LLM/tool
+     events from ANY LangChain/LangGraph agent (observability + certificate). This path is
+     record-only: callbacks observe but cannot re-inject, so it doesn't drive replay.
+"""
+from __future__ import annotations
+from typing import TypedDict, Optional
+from .react import parse_action, REACT_SYS
+from .events import Kind
+
+try:  # optional dependency: only needed to actually build/run the graph
+    from langgraph.graph import StateGraph, START, END
+    from langchain_core.callbacks import BaseCallbackHandler
+    _IMPORT_ERROR: Exception | None = None
+except ImportError as exc:  # keep the module importable without the extra installed
+    StateGraph = START = END = None  # type: ignore[assignment]
+    BaseCallbackHandler = object  # type: ignore[assignment,misc]
+    _IMPORT_ERROR = exc
+
+
+def _require_langgraph() -> None:
+    """Raise a clear, actionable error if the LangGraph extra is not installed."""
+    if _IMPORT_ERROR is not None:
+        raise ImportError(
+            "vitnify's LangGraph adapter requires 'langgraph' and 'langchain-core'. "
+            'Install them with:  pip install "vitnify[langgraph]"'
+        ) from _IMPORT_ERROR
+
+
+class VRState(TypedDict, total=False):
+    prompt: str
+    steps: int
+    last_text: str
+    last_act: Optional[dict]
+
+
+def initial_prompt(task: str, tools_desc: str) -> str:
+    return (f"<|system|>\n{REACT_SYS.format(tools=tools_desc)}</s>\n"
+            f"<|user|>\n{task}</s>\n<|assistant|>\n")
+
+
+def build_react_graph(session, tools_desc: str, tool_names, max_steps: int = 3, n_new: int = 40):
+    """Compile a LangGraph StateGraph wired to a VitniReplay Session."""
+    _require_langgraph()
+
+    def agent_node(state: VRState) -> VRState:
+        text = session.llm(state["prompt"], n_new=n_new)          # recorded + recomputable
+        act = parse_action(text, tool_names)
+        return {"last_text": text, "last_act": act, "steps": state.get("steps", 0) + 1}
+
+    def route(state: VRState) -> str:
+        act = state.get("last_act")
+        if act is None or "final" in act or state.get("steps", 0) >= max_steps:
+            return "end"
+        return "tools"
+
+    def tool_node(state: VRState) -> VRState:
+        act = state["last_act"]
+        ok, obs = session.tool(act["tool"], *act["args"])          # capability wall
+        observation = obs if ok else f"DENIED: agent has no capability '{act['tool']}'"
+        return {"prompt": state["prompt"] + state["last_text"] + f"\nObservation: {observation}\n"}
+
+    g = StateGraph(VRState)
+    g.add_node("agent", agent_node)
+    g.add_node("tools", tool_node)
+    g.add_edge(START, "agent")
+    g.add_conditional_edges("agent", route, {"tools": "tools", "end": END})
+    g.add_edge("tools", "agent")
+    return g.compile()
+
+
+class VitniReplayCallback(BaseCallbackHandler):
+    """Observability adapter for ANY LangChain/LangGraph agent: records LLM + tool events
+    into the session's event log. Record-only (callbacks can't re-inject for replay)."""
+    def __init__(self, session):
+        _require_langgraph()
+        self.session = session
+        self._pending = ("?", "")
+
+    def on_llm_end(self, response, **kwargs):
+        try:
+            text = response.generations[0][0].text
+        except Exception:
+            text = str(response)[:200]
+        self.session.log.append(Kind.LLM_CALL, {"text": text, "via": "callback"})
+
+    def on_tool_start(self, serialized, input_str, **kwargs):
+        name = serialized.get("name", "?") if isinstance(serialized, dict) else "?"
+        self._pending = (name, str(input_str))
+
+    def on_tool_end(self, output, **kwargs):
+        name, inp = self._pending
+        self.session.log.append(Kind.TOOL_CALL,
+            {"tool": name, "args": [inp], "decision": "OBSERVED", "result": str(output)[:200]})