feat(langgraph): append-only TranscriptChannel; edits and removals fork

libs/langgraph/langgraph/channels/transcript.py

@@ -0,0 +1,300 @@
+from __future__ import annotations
+
+import uuid
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+from typing_extensions import Self
+
+from langgraph._internal._typing import MISSING
+from langgraph.channels.base import BaseChannel
+from langgraph.channels.binop import _get_overwrite
+from langgraph.errors import InvalidUpdateError
+
+__all__ = ("TranscriptChannel", "TranscriptEntry", "TranscriptRef")
+
+EntryKind = Literal["append", "reset"]
+
+
+@dataclass(frozen=True, slots=True)
+class TranscriptEntry:
+    """One log entry. `append` carries a message, `reset` the full list."""
+
+    seq: int
+    kind: EntryKind
+    item_id: str = ""
+    value: Any = None
+
+
+@dataclass(slots=True)
+class TranscriptRef:
+    """Checkpoint value of a `TranscriptChannel`: a pointer into a log.
+
+    Written: `entries` holds entries with seq in (base_head, head].
+    Read: `entries` holds every entry in [1, head], fork ancestors included.
+    """
+
+    transcript_id: str
+    head: int
+    base_head: int = 0
+    entries: list[TranscriptEntry] = field(default_factory=list)
+    parent_transcript_id: str = ""
+    fork_seq: int = 0
+    is_tip: bool = True
+
+
+def _coerce(update: Any) -> list[Any]:
+    from langchain_core.messages import (
+        BaseMessageChunk,
+        convert_to_messages,
+        message_chunk_to_message,
+    )
+
+    items = update if isinstance(update, list) else [update]
+    out = []
+    for m in convert_to_messages(items):
+        m = message_chunk_to_message(m) if isinstance(m, BaseMessageChunk) else m
+        if m.id is None:
+            m.id = str(uuid.uuid4())
+        out.append(m)
+    return out
+
+
+def _is_remove(m: Any) -> bool:
+    from langchain_core.messages import RemoveMessage
+
+    return isinstance(m, RemoveMessage)
+
+
+def _remove_all_id() -> str:
+    from langgraph.graph.message import REMOVE_ALL_MESSAGES
+
+    return REMOVE_ALL_MESSAGES
+
+
+class TranscriptChannel(BaseChannel[list, Any, TranscriptRef]):
+    """Message list persisted as an append-only log.
+
+    The checkpoint stores `(transcript_id, head)`. A checkpointer that
+    supports transcripts persists new entries atomically with the checkpoint
+    and returns them on read, so storage grows linearly and loading any
+    checkpoint costs O(len(transcript)).
+
+    Writes contain only new or changed messages:
+
+    * a new id appends;
+    * an existing id (edit) or a `RemoveMessage` forks: a new transcript
+      shares the prefix before the edited message and replays what follows;
+    * `REMOVE_ALL_MESSAGES` or `Overwrite` writes one `reset` entry.
+
+    Loading a checkpoint that is not its transcript's tip also forks.
+    """
+
+    __slots__ = (
+        "value",
+        "pos",
+        "anchor",
+        "resets",
+        "transcript_id",
+        "head",
+        "base_head",
+        "pending",
+        "parent_transcript_id",
+        "fork_seq",
+    )
+
+    def __init__(self, typ: Any = list) -> None:
+        super().__init__(typ)
+        self.value: list[Any] = []
+        # id -> index in value
+        self.pos: dict[str, int] = {}
+        # id -> seq of the entry that introduced it (an append or a reset)
+        self.anchor: dict[str, int] = {}
+        # seqs of reset entries
+        self.resets: set[int] = set()
+        self.transcript_id = str(uuid.uuid4())
+        self.head = 0
+        # Highest seq known committed for this channel's lineage.
+        self.base_head = 0
+        self.pending: list[TranscriptEntry] = []
+        # Set only for a fork this channel created and has not seen committed.
+        self.parent_transcript_id = ""
+        self.fork_seq = 0
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, TranscriptChannel)
+
+    @property
+    def ValueType(self) -> Any:
+        return list
+
+    @property
+    def UpdateType(self) -> Any:
+        return Any
+
+    def _new(self) -> Self:
+        new = self.__class__(self.typ)
+        new.key = self.key
+        return new
+
+    def copy(self) -> Self:
+        new = self._new()
+        new.value = list(self.value)
+        new.pos = dict(self.pos)
+        new.anchor = dict(self.anchor)
+        new.resets = set(self.resets)
+        new.transcript_id = self.transcript_id
+        new.head = self.head
+        new.base_head = self.base_head
+        new.pending = list(self.pending)
+        new.parent_transcript_id = self.parent_transcript_id
+        new.fork_seq = self.fork_seq
+        return new
+
+    def from_checkpoint(self, checkpoint: Any) -> Self:
+        new = self._new()
+        if checkpoint is MISSING:
+            return new
+        if not isinstance(checkpoint, TranscriptRef):
+            # Plain list from a regular channel: start a transcript from it.
+            new._reset(_coerce(list(checkpoint)))
+            return new
+        for e in checkpoint.entries:
+            if e.kind == "reset":
+                new._load_reset(e.seq, e.value)
+            else:
+                new._load_append(e.seq, e.value)
+        new.head = new.base_head = checkpoint.head
+        if checkpoint.is_tip:
+            new.transcript_id = checkpoint.transcript_id
+        else:
+            new.parent_transcript_id = checkpoint.transcript_id
+            new.fork_seq = checkpoint.head
+        return new
+
+    def get(self) -> list[Any]:
+        return self.value
+
+    def is_available(self) -> bool:
+        return True
+
+    def update(self, values: Sequence[Any]) -> bool:
+        if not values:
+            return False
+        for v in values:
+            is_overwrite, overwrite = _get_overwrite(v)
+            if is_overwrite:
+                self._reset(_coerce(list(overwrite)))
+                continue
+            msgs = _coerce(v)
+            remove_all = [
+                i
+                for i, m in enumerate(msgs)
+                if _is_remove(m) and m.id == _remove_all_id()
+            ]
+            if remove_all:
+                self._reset(msgs[remove_all[-1] + 1 :])
+                continue
+            for m in msgs:
+                self._write(m)
+        return True
+
+    def _write(self, m: Any) -> None:
+        i = self.pos.get(m.id)
+        if i is None:
+            if _is_remove(m):
+                raise InvalidUpdateError(
+                    f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
+                )
+            self._append(m)
+            return
+        current = self.value[i]
+        if current is m or (not _is_remove(m) and current == m):
+            # Resend of a message already in the transcript.
+            return
+        replacement = [] if _is_remove(m) else [m]
+        self._fork_at(m.id, replacement)
+
+    def _fork_at(self, item_id: str, replacement: list[Any]) -> None:
+        """Fork before `item_id`, swap it for `replacement`, replay the rest."""
+        # The fork point must already be committed, so it is at most base_head.
+        fork = min(self.anchor[item_id] - 1, self.base_head)
+        prefix = [m for m in self.value if self.anchor[m.id] <= fork]
+        # Entries in (fork, head] that brought in the tail, in order. After a
+        # reset no older item survives, so a reset group is the whole list.
+        groups: list[tuple[int, list[Any]]] = []
+        for m in self.value:
+            a = self.anchor[m.id]
+            if a <= fork:
+                continue
+            if groups and groups[-1][0] == a:
+                groups[-1][1].append(m)
+            else:
+                groups.append((a, [m]))
+
+        # Items surviving the latest reset after the fork point all anchor at
+        # or after it, so replay starts with that reset, even when empty.
+        latest_reset = max((r for r in self.resets if r > fork), default=None)
+        if latest_reset is not None and (not groups or groups[0][0] != latest_reset):
+            groups.insert(0, (latest_reset, []))
+        old_resets = self.resets
+
+        if self.parent_transcript_id:
+            # Fork created this run: base_head == fork_seq, so branch off its parent.
+            parent = self.parent_transcript_id
+        elif fork == 0:
+            parent = ""
+        else:
+            parent = self.transcript_id
+        resets = {r for r in self.resets if r <= fork}
+        anchors = {m.id: self.anchor[m.id] for m in prefix}
+
+        self.transcript_id = str(uuid.uuid4())
+        self.parent_transcript_id = parent
+        self.fork_seq = self.head = self.base_head = fork
+        self.pending = []
+        self.value = prefix
+        self.pos = {m.id: i for i, m in enumerate(prefix)}
+        self.anchor = anchors
+        self.resets = resets
+        for a, group in groups:
+            swapped = [
+                n for m in group for n in (replacement if m.id == item_id else [m])
+            ]
+            if a in old_resets:
+                self._reset(swapped)
+            else:
+                for m in swapped:
+                    self._append(m)
+
+    def _append(self, m: Any) -> None:
+        self.head += 1
+        self.pending.append(TranscriptEntry(self.head, "append", m.id, m))
+        self._load_append(self.head, m)
+
+    def _reset(self, msgs: list[Any]) -> None:
+        self.head += 1
+        self.pending.append(TranscriptEntry(self.head, "reset", "", list(msgs)))
+        self._load_reset(self.head, msgs)
+
+    def _load_append(self, seq: int, m: Any) -> None:
+        self.pos[m.id] = len(self.value)
+        self.anchor[m.id] = seq
+        self.value.append(m)
+
+    def _load_reset(self, seq: int, msgs: list[Any]) -> None:
+        self.value = list(msgs)
+        self.pos = {m.id: i for i, m in enumerate(self.value)}
+        self.anchor = dict.fromkeys(self.pos, seq)
+        self.resets.add(seq)
+
+    def checkpoint(self) -> TranscriptRef:
+        return TranscriptRef(
+            transcript_id=self.transcript_id,
+            head=self.head,
+            base_head=self.base_head,
+            entries=list(self.pending),
+            parent_transcript_id=self.parent_transcript_id,
+            fork_seq=self.fork_seq,
+        )

libs/langgraph/tests/test_channels_transcript.py

@@ -0,0 +1,119 @@
+from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+
+from langgraph._internal._typing import MISSING
+from langgraph.channels.transcript import TranscriptChannel, TranscriptRef
+from langgraph.graph.message import REMOVE_ALL_MESSAGES
+
+
+class FakeStore:
+    """Commits refs like a transcript-aware checkpointer and reads them back."""
+
+    def __init__(self) -> None:
+        self.transcripts: dict[str, dict] = {}
+
+    def commit(self, ch: TranscriptChannel) -> TranscriptRef:
+        ref = ch.checkpoint()
+        t = self.transcripts.setdefault(
+            ref.transcript_id,
+            {"parent": ref.parent_transcript_id, "fork": ref.fork_seq, "entries": {}},
+        )
+        for e in ref.entries:
+            t["entries"][e.seq] = e
+        return ref
+
+    def _read(self, tid: str, head: int) -> list:
+        t = self.transcripts[tid]
+        out = self._read(t["parent"], min(head, t["fork"])) if t["parent"] else []
+        return out + [t["entries"][s] for s in sorted(t["entries"]) if s <= head]
+
+    def load(self, ref: TranscriptRef, *, tip: bool = True) -> TranscriptChannel:
+        return TranscriptChannel().from_checkpoint(
+            TranscriptRef(
+                ref.transcript_id,
+                ref.head,
+                entries=self._read(ref.transcript_id, ref.head),
+                is_tip=tip,
+            )
+        )
+
+
+def contents(ch: TranscriptChannel) -> list[str]:
+    return [m.content for m in ch.get()]
+
+
+def test_appends_and_resends() -> None:
+    store = FakeStore()
+    ch = TranscriptChannel().from_checkpoint(MISSING)
+    ch.update([[HumanMessage("h0", id="h0"), AIMessage("a0", id="a0")]])
+    ch = store.load(store.commit(ch))
+    ch.update([[HumanMessage("h1", id="h1")]])
+    ref = store.commit(ch)
+    assert [e.seq for e in ref.entries] == [3]
+    ch = store.load(ref)
+    assert contents(ch) == ["h0", "a0", "h1"]
+
+    ch.update([[AIMessage("a0", id="a0"), HumanMessage("h1", id="h1")]])
+    assert ch.checkpoint().entries == []
+
+
+def test_edit_forks_and_keeps_original() -> None:
+    store = FakeStore()
+    ch = TranscriptChannel().from_checkpoint(MISSING)
+    ch.update([[HumanMessage("h0", id="h0"), AIMessage("a0", id="a0")]])
+    ch = store.load(store.commit(ch))
+    ch.update([[HumanMessage("h1", id="h1"), AIMessage("a1", id="a1")]])
+    original = store.commit(ch)
+    ch = store.load(original)
+
+    ch.update([[AIMessage("a0-edited", id="a0")]])
+    edited = store.commit(ch)
+    assert edited.parent_transcript_id == original.transcript_id
+    assert edited.fork_seq == 1
+    assert contents(store.load(edited)) == ["h0", "a0-edited", "h1", "a1"]
+    assert contents(store.load(original, tip=False)) == ["h0", "a0", "h1", "a1"]
+
+    # A second fork before the first is committed branches off the same parent.
+    ch.update([[RemoveMessage(id="h1")]])
+    removed = store.commit(ch)
+    assert removed.parent_transcript_id == original.transcript_id
+    assert contents(store.load(removed)) == ["h0", "a0-edited", "a1"]
+
+
+def test_edit_before_any_commit_has_no_parent() -> None:
+    ch = TranscriptChannel().from_checkpoint(MISSING)
+    ch.update([[HumanMessage("x", id="x"), AIMessage("y", id="y")]])
+    ch.update([[AIMessage("y2", id="y")]])
+    ref = ch.checkpoint()
+    assert (ref.parent_transcript_id, ref.fork_seq) == ("", 0)
+    assert contents(ch) == ["x", "y2"]
+
+
+def test_edit_inside_reset_replays_reset() -> None:
+    store = FakeStore()
+    ch = TranscriptChannel().from_checkpoint(MISSING)
+    ch.update([[HumanMessage("h0", id="h0")]])
+    ch = store.load(store.commit(ch))
+    ch.update(
+        [[RemoveMessage(id=REMOVE_ALL_MESSAGES), HumanMessage("p", id="p"), AIMessage("q", id="q")]]
+    )
+    ch = store.load(store.commit(ch))
+    ch.update([[AIMessage("q2", id="q")]])
+    ref = store.commit(ch)
+    assert [(e.seq, e.kind) for e in ref.entries] == [(2, "reset")]
+    assert contents(store.load(ref)) == ["p", "q2"]
+
+
+def test_load_non_tip_forks() -> None:
+    store = FakeStore()
+    ch = TranscriptChannel().from_checkpoint(MISSING)
+    ch.update([[HumanMessage("h0", id="h0")]])
+    first = store.commit(ch)
+    ch = store.load(first)
+    ch.update([[AIMessage("a0", id="a0")]])
+    store.commit(ch)
+
+    branch = store.load(first, tip=False)
+    branch.update([[AIMessage("other", id="o")]])
+    ref = store.commit(branch)
+    assert ref.parent_transcript_id == first.transcript_id
+    assert contents(store.load(ref)) == ["h0", "other"]