diff --git a/.gitignore b/.gitignore index b154d3b56..d34e545b9 100644 --- a/.gitignore +++ b/.gitignore @@ -101,3 +101,4 @@ dmypy.json .editorconfig .scratch .worktrees/ +new_pr_desc.md diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 1097757e9..fd42146cf 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -385,11 +385,11 @@ async def test_delta_channel_chain_reconstruction(saver_name: str) -> None: from langchain_core.messages import AIMessage, HumanMessage from langgraph.channels.delta import DeltaChannel from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer from typing_extensions import TypedDict class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] def respond(state: State) -> dict: n = len(state["messages"]) diff --git a/libs/langgraph/langgraph/channels/delta.py b/libs/langgraph/langgraph/channels/delta.py index f109e32c2..2ca944e55 100644 --- a/libs/langgraph/langgraph/channels/delta.py +++ b/libs/langgraph/langgraph/channels/delta.py @@ -23,38 +23,43 @@ __all__ = ("DeltaChannel",) class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): - """Fold-reducer channel with configurable snapshot cadence. + """Reducer channel that stores only a sentinel in checkpoint blobs and + reconstructs state by replaying ancestor writes through the reducer. - `snapshot_frequency=None` (default): pure delta — stores only - `DELTA_SENTINEL` in checkpoint blobs; reads replay all ancestor writes. + The reducer receives the current accumulated value and the full list of + new writes for that step in one call: + ``reducer(state, [write1, write2, ...]) -> new_state``. - `snapshot_frequency=N`: pregel's `create_checkpoint` writes a full - `_DeltaSnapshot` blob every N steps (eagerly, even if the channel had - no write that step). Reads walk at most N ancestor checkpoints before - hitting the snapshot, bounding replay depth to N regardless of thread - length. + ``snapshot_frequency=None`` (default): pure delta — stores only + ``DELTA_SENTINEL`` in checkpoint blobs; reads replay all ancestor writes. + + ``snapshot_frequency=N``: ``create_checkpoint`` writes a full + ``_DeltaSnapshot`` blob every N steps, bounding replay depth to N. Parameters: - typ: The value type (e.g. `list`, `dict`). - operator: Binary reducer `(Value, Value) -> Value`. + reducer: ``(state, list[writes]) -> new_state``. Receives the current + accumulated value and the list of all writes for this step. + typ: The value type (e.g. ``list``, ``dict``). Inferred automatically + from the outer type when used inside ``Annotated[T, DeltaChannel(...)]``. snapshot_frequency: Every Nth pregel step writes a snapshot blob. - `None` (default) = pure delta, never snapshot. + ``None`` (default) = pure delta, never snapshot. """ - __slots__ = ("value", "operator", "snapshot_frequency") + __slots__ = ("value", "reducer", "snapshot_frequency") value: Value | Any def __init__( self, - typ: type[Value], - operator: Callable[[Any, Any], Any], + reducer: Callable[[Any, Sequence[Any]], Any], + typ: type[Value] | None = None, *, snapshot_frequency: int | None = None, ) -> None: + if typ is None: + typ = list # type: ignore[assignment] # placeholder; overridden by _is_field_channel super().__init__(typ) - self.operator = operator + self.reducer = reducer self.snapshot_frequency = snapshot_frequency - # Normalize abstract / parameterized types to their concrete counterparts. typ = _strip_extras(typ) if typ in (collections.abc.Sequence, collections.abc.MutableSequence): typ = list @@ -70,7 +75,7 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return False if self.snapshot_frequency != other.snapshot_frequency: return False - return _operators_equal(self.operator, other.operator) + return _operators_equal(self.reducer, other.reducer) @property def ValueType(self) -> Any: @@ -90,33 +95,22 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): def copy(self) -> Self: new = self.__class__( - self.typ, self.operator, snapshot_frequency=self.snapshot_frequency + self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency ) new.key = self.key new.value = self.value if self.value is MISSING else _copy.copy(self.value) return new - def _apply_write(self, value: Any, write: Any) -> Any: - is_overwrite, overwrite_value = _get_overwrite(write) - if is_overwrite: - return ( - _copy.copy(overwrite_value) - if overwrite_value is not None - else self.typ() - ) - base = self.typ() if value is MISSING else value - return self.operator(base, write) - def from_checkpoint(self, checkpoint: Any) -> Self: """Initialize from a stored blob or sentinel. Blob types (dispatched via serde ext code, not dict key inspection): - * `DELTA_SENTINEL` / `MISSING`: start empty; caller replays writes. - * `_DeltaSnapshot(value)`: restore value directly from snapshot. + * ``DELTA_SENTINEL`` / ``MISSING``: start empty; caller replays writes. + * ``_DeltaSnapshot(value)``: restore value directly from snapshot. * plain value (migration from old BinOp blobs): use directly. """ new = self.__class__( - self.typ, self.operator, snapshot_frequency=self.snapshot_frequency + self.reducer, self.typ, snapshot_frequency=self.snapshot_frequency ) new.key = self.key if checkpoint is MISSING or checkpoint is DELTA_SENTINEL: @@ -128,9 +122,24 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return new def replay_writes(self, writes: Sequence[PendingWrite]) -> None: - """Fold ancestor writes oldest→newest into current value.""" - for _, _, value in writes: - self.value = self._apply_write(self.value, value) + """Apply ancestor writes oldest→newest via a single reducer call. + + If any write is an Overwrite, the last one in the sequence acts as + the reset point: its value becomes the new base and only writes + after it are passed to the reducer. + """ + values = [v for _, _, v in writes] + if not values: + return + base = self.value + start = 0 + for i, v in enumerate(values): + is_ow, ow_value = _get_overwrite(v) + if is_ow: + base = _copy.copy(ow_value) if ow_value is not None else self.typ() + start = i + 1 + remaining = values[start:] + self.value = self.reducer(base, remaining) if remaining else base def update(self, values: Sequence[Any]) -> bool: if not values: @@ -147,10 +156,17 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): raise InvalidUpdateError(msg) overwrite_idx = i if overwrite_idx is not None: - self.value = self._apply_write(self.value, values[overwrite_idx]) + _, overwrite_value = _get_overwrite(values[overwrite_idx]) + base = ( + _copy.copy(overwrite_value) + if overwrite_value is not None + else self.typ() + ) + remaining = [v for i, v in enumerate(values) if i != overwrite_idx] + self.value = self.reducer(base, remaining) if remaining else base return True - for value in values: - self.value = self._apply_write(self.value, value) + base = self.typ() if self.value is MISSING else self.value + self.value = self.reducer(base, list(values)) return True def get(self) -> Any: @@ -162,10 +178,10 @@ class DeltaChannel(Generic[Value], BaseChannel[Any, Any, Any]): return self.value is not MISSING def checkpoint(self) -> Any: - """Return stored representation: always `DELTA_SENTINEL`. + """Return stored representation: always ``DELTA_SENTINEL``. - Snapshot decisions are made by `create_checkpoint` in pregel (which - has the step number) via `is_snapshot_step`. `checkpoint()` is only + Snapshot decisions are made by ``create_checkpoint`` in pregel (which + has the step number) via ``is_snapshot_step``. ``checkpoint()`` is only called for non-snapshot steps or when no checkpointer is available. """ if self.value is MISSING: diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 18c96d436..53388b812 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -244,6 +244,49 @@ def add_messages( return merged +def _messages_delta_reducer( + state: list[AnyMessage], writes: list[list[AnyMessage]] +) -> list[AnyMessage]: + """**Experimental.** Batch reducer for use with ``DeltaChannel``. + + Processes all writes for a step in one pass — dedup by ID, ``RemoveMessage`` + tombstoning — without calling ``add_messages``. Assumes writes contain + already-typed ``BaseMessage`` objects (no raw-dict coercion). + + Use ``add_messages`` as the reducer for ``BinaryOperatorAggregate`` or + anywhere raw message dicts / strings need to be coerced first. + + Example:: + + from typing import Annotated + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.message import _messages_delta_reducer + + class State(TypedDict): + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] + """ + from itertools import chain + + index: dict[str, int] = {m.id: i for i, m in enumerate(state) if m.id is not None} + result: list[AnyMessage | None] = list(state) + for msg in chain.from_iterable( + [w] if isinstance(w, BaseMessage) else w for w in writes + ): + mid = msg.id + if mid is None: + result.append(msg) + elif isinstance(msg, RemoveMessage): + if mid in index: + result[index[mid]] = None + del index[mid] + elif mid in index: + result[index[mid]] = msg + else: + index[mid] = len(result) + result.append(msg) + return [m for m in result if m is not None] + + @deprecated( "MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.", category=None, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 4f49210e1..d2ee7527a 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1688,8 +1688,8 @@ def _is_field_channel(typ: type[Any]) -> BaseChannel | None: ): origin = origin.__args__[0] item = item.__class__( + item.reducer, origin, - item.operator, snapshot_frequency=item.snapshot_frequency, ) return item diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index cfeacda15..2a5aa5446 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -1312,6 +1312,8 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): + _delta_write_futs: list[asyncio.Future[Any]] + def __init__( self, input: Any | None, @@ -1490,7 +1492,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): if saved.pending_writes is not None else [] ) - self._delta_write_futs: list[asyncio.Future[Any]] = [] + self._delta_write_futs = [] self.submit = await self.stack.enter_async_context( AsyncBackgroundExecutor(self.config) ) diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index 562a45001..2f1cf6d15 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -12,7 +12,7 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.errors import EmptyChannelError, InvalidUpdateError -from langgraph.graph.message import add_messages +from langgraph.graph.message import _messages_delta_reducer pytestmark = pytest.mark.anyio @@ -127,9 +127,9 @@ def test_delta_channel_basic_two_steps() -> None: from langchain_core.messages import AIMessage, HumanMessage from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer - ch = DeltaChannel(list, add_messages).from_checkpoint(MISSING) + ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) # Step 1: one message added ch.update([HumanMessage(content="hi", id="h1")]) @@ -151,9 +151,9 @@ def test_delta_channel_from_checkpoint_writes_list() -> None: """replay_writes on a fresh channel replays through the operator.""" from langchain_core.messages import AIMessage, HumanMessage - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer - spec = DeltaChannel(list, add_messages) + spec = DeltaChannel(_messages_delta_reducer, list) ch = spec.from_checkpoint(DELTA_SENTINEL) ch.replay_writes( [ @@ -172,10 +172,10 @@ def test_delta_channel_from_checkpoint_writes_list() -> None: def test_delta_channel_from_checkpoint_backwards_compat() -> None: from langchain_core.messages import HumanMessage - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer # Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat - spec = DeltaChannel(list, add_messages) + spec = DeltaChannel(_messages_delta_reducer, list) old_value = [HumanMessage(content="old", id="h1")] ch = spec.from_checkpoint(old_value) assert ch.get() == old_value @@ -185,10 +185,10 @@ def test_delta_channel_overwrite() -> None: from langchain_core.messages import HumanMessage from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer from langgraph.types import Overwrite - ch = DeltaChannel(list, add_messages).from_checkpoint(MISSING) + ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) ch.update([HumanMessage(content="old", id="h1")]) ch.update([Overwrite([HumanMessage(content="new", id="h2")])]) @@ -203,9 +203,9 @@ def test_delta_channel_remove_message_and_replay() -> None: """RemoveMessage must round-trip correctly when writes are replayed.""" from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer - spec = DeltaChannel(list, add_messages) + spec = DeltaChannel(_messages_delta_reducer, list) ch = spec.from_checkpoint(MISSING) # Step 1: add two messages @@ -236,9 +236,9 @@ def test_delta_channel_update_by_id_and_replay() -> None: """Updating a message by ID must round-trip correctly through writes replay.""" from langchain_core.messages import HumanMessage - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer - spec = DeltaChannel(list, add_messages) + spec = DeltaChannel(_messages_delta_reducer, list) ch = spec.from_checkpoint(MISSING) # Step 1: add a message @@ -264,9 +264,9 @@ def test_delta_channel_checkpoint_returns_sentinel() -> None: """checkpoint() always returns DELTA_SENTINEL regardless of state.""" from langgraph.checkpoint.base import DELTA_SENTINEL - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer - ch = DeltaChannel(list, add_messages).from_checkpoint(MISSING) + ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING) assert ch.checkpoint() is DELTA_SENTINEL from langchain_core.messages import HumanMessage @@ -290,12 +290,12 @@ def test_delta_channel_snapshot_step_based() -> None: from typing_extensions import TypedDict from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer # snapshot_frequency=5: snapshot every 5 pregel steps class State(TypedDict): messages: Annotated[ - list, DeltaChannel(list, add_messages, snapshot_frequency=5) + list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=5) ] other: str @@ -349,11 +349,11 @@ def test_delta_channel_snapshot_fires_even_when_not_written() -> None: from typing_extensions import TypedDict from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): messages: Annotated[ - list, DeltaChannel(list, add_messages, snapshot_frequency=3) + list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=3) ] tick: int @@ -406,10 +406,10 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None: from typing_extensions import TypedDict from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer, list)] n = {"v": 0} @@ -451,14 +451,17 @@ def _delta_channel_with_type(operator, typ): from langgraph.channels.delta import DeltaChannel from langgraph.graph.state import _get_channel - return _get_channel("_test", Annotated[typ, DeltaChannel(typ, operator)]) + return _get_channel("_test", Annotated[typ, DeltaChannel(operator)]) def test_delta_channel_dict_reducer_fresh_channel() -> None: """DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint.""" - def merge_dicts(left: dict, right: dict) -> dict: - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) assert ch.is_available() @@ -468,8 +471,11 @@ def test_delta_channel_dict_reducer_fresh_channel() -> None: def test_delta_channel_dict_reducer_basic_updates() -> None: """DeltaChannel with a dict reducer accumulates key/value pairs across steps.""" - def merge_dicts(left: dict, right: dict) -> dict: - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) @@ -487,8 +493,11 @@ def test_delta_channel_dict_reducer_basic_updates() -> None: def test_delta_channel_dict_reducer_writes_reconstruction() -> None: """replay_writes on a fresh channel replays through a dict merge reducer.""" - def merge_dicts(left: dict, right: dict) -> dict: - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result spec = _delta_channel_with_type(merge_dicts, dict) ch = spec.from_checkpoint(DELTA_SENTINEL) @@ -505,15 +514,14 @@ def test_delta_channel_dict_reducer_writes_reconstruction() -> None: def test_delta_channel_dict_reducer_with_deletions() -> None: """Dict reducer that treats None values as deletions works end-to-end.""" - def merge_files(left: dict | None, right: dict) -> dict: - if left is None: - return {k: v for k, v in right.items() if v is not None} - result = {**left} - for k, v in right.items(): - if v is None: - result.pop(k, None) - else: - result[k] = v + def merge_files(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + for k, v in w.items(): + if v is None: + result.pop(k, None) + else: + result[k] = v return result ch = _delta_channel_with_type(merge_files, dict).from_checkpoint(MISSING) @@ -536,8 +544,11 @@ def test_delta_channel_dict_reducer_overwrite_in_update() -> None: """Overwrite(dict) in update() must preserve dict shape, not coerce to list.""" from langgraph.types import Overwrite - def merge_dicts(left: dict, right: dict) -> dict: - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) ch.update([{"a": 1}]) @@ -549,8 +560,11 @@ def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None: """Overwrite(dict) embedded in replayed writes must reconstruct as dict.""" from langgraph.types import Overwrite - def merge_dicts(left: dict, right: dict) -> dict: - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result spec = _delta_channel_with_type(merge_dicts, dict) ch = spec.from_checkpoint(DELTA_SENTINEL) @@ -573,12 +587,13 @@ def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None: from langgraph.channels.delta import DeltaChannel from langgraph.graph.state import _get_channel - def merge_dicts(left: dict | None, right: dict) -> dict: - if left is None: - return dict(right) - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result - annotation = Annotated[NotRequired[dict[str, int]], DeltaChannel(dict, merge_dicts)] + annotation = Annotated[NotRequired[dict[str, int]], DeltaChannel(merge_dicts)] ch = _get_channel("files", annotation).from_checkpoint(MISSING) assert ch.get() == {} ch.update([{"a": 1}]) @@ -596,19 +611,18 @@ def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None: from langgraph.channels.delta import DeltaChannel from langgraph.graph import START, StateGraph - def merge_files(left: dict | None, right: dict) -> dict: - if left is None: - return {k: v for k, v in right.items() if v is not None} - result = {**left} - for k, v in right.items(): - if v is None: - result.pop(k, None) - else: - result[k] = v + def merge_files(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + for k, v in w.items(): + if v is None: + result.pop(k, None) + else: + result[k] = v return result class State(TypedDict): - files: Annotated[dict[str, str], DeltaChannel(dict, merge_files)] + files: Annotated[dict[str, str], DeltaChannel(merge_files)] turn = {"v": 0} @@ -657,8 +671,11 @@ def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None: def test_delta_channel_dict_reducer_backwards_compat() -> None: """A pre-DeltaChannel dict checkpoint must load as a dict, not be listified.""" - def merge_dicts(left: dict, right: dict) -> dict: - return {**left, **right} + def merge_dicts(state: dict, writes: list) -> dict: + result = dict(state) + for w in writes: + result.update(w) + return result spec = _delta_channel_with_type(merge_dicts, dict) old_value = {"a": 1, "b": 2} @@ -678,7 +695,7 @@ def test_delta_channel_from_checkpoint_honors_seed() -> None: a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs the post-migration state correctly rather than replaying from empty. """ - spec = DeltaChannel(list, add_messages) + spec = DeltaChannel(_messages_delta_reducer, list) seed = [HumanMessage(content="pre-delta", id="p1")] ch = spec.from_checkpoint(seed) ch.replay_writes( @@ -694,7 +711,7 @@ def test_delta_channel_from_checkpoint_honors_seed() -> None: def test_delta_channel_from_checkpoint_seed_without_writes() -> None: """Reconstruction at a pre-delta ancestor with no newer deltas returns just the seed — the saver's terminator fired immediately.""" - spec = DeltaChannel(list, add_messages) + spec = DeltaChannel(_messages_delta_reducer, list) seed = [HumanMessage(content="only-snap", id="s1")] ch = spec.from_checkpoint(seed) ch.replay_writes([]) @@ -708,10 +725,10 @@ def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_sentinel() -> explicitly should feed None to the reducer as the left operand. """ - def replace(left, right): - return right + def replace(state, writes): + return writes[-1] if writes else state - spec = DeltaChannel(list, replace) + spec = DeltaChannel(replace, list) ch = spec.from_checkpoint(None) ch.replay_writes([("t0", "x", "after")]) # Reducer replaces; seed=None → first write produces "after". diff --git a/libs/langgraph/tests/test_delta_channel_benchmark.py b/libs/langgraph/tests/test_delta_channel_benchmark.py index 41bc68c25..8136ed717 100644 --- a/libs/langgraph/tests/test_delta_channel_benchmark.py +++ b/libs/langgraph/tests/test_delta_channel_benchmark.py @@ -28,7 +28,7 @@ from typing_extensions import TypedDict from langgraph.channels.delta import DeltaChannel from langgraph.graph import END, StateGraph -from langgraph.graph.message import add_messages +from langgraph.graph.message import _messages_delta_reducer, add_messages try: from langgraph.checkpoint.postgres import PostgresSaver @@ -114,12 +114,14 @@ class BinaryState(TypedDict): class DeltaState(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] def _make_delta_state(snapshot_frequency: int | float) -> type: """Create a TypedDict with DeltaChannel at the given snapshot_frequency.""" - channel = DeltaChannel(list, add_messages, snapshot_frequency=snapshot_frequency) + channel = DeltaChannel( + _messages_delta_reducer, snapshot_frequency=snapshot_frequency + ) # Use the functional TypedDict form so the Annotated type is stored as an # already-evaluated object rather than a forward-reference string (which # would fail when get_type_hints tries to resolve 'snapshot_frequency'). diff --git a/libs/langgraph/tests/test_delta_channel_migration.py b/libs/langgraph/tests/test_delta_channel_migration.py index e74a713da..9793fd831 100644 --- a/libs/langgraph/tests/test_delta_channel_migration.py +++ b/libs/langgraph/tests/test_delta_channel_migration.py @@ -70,6 +70,13 @@ def _noop(_state: Any) -> dict: return {} +def _list_concat(state: list, writes: list) -> list: + result = list(state) + for w in writes: + result.extend(w if isinstance(w, list) else [w]) + return result + + def _binop_graph(checkpointer: Any) -> Any: class BinopState(TypedDict): items: Annotated[list, BinaryOperatorAggregate(list, operator.add)] @@ -85,7 +92,7 @@ def _binop_graph(checkpointer: Any) -> Any: def _delta_graph(checkpointer: Any) -> Any: class DeltaState(TypedDict): - items: Annotated[list, DeltaChannel(list, operator.add)] + items: Annotated[list, DeltaChannel(_list_concat)] return ( StateGraph(DeltaState) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f898ac4a1..ee95fdfd7 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9409,10 +9409,10 @@ async def test_delta_channel_end_to_end_inmemory() -> None: from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] def respond(state: State) -> dict: n = len(state["messages"]) @@ -9450,10 +9450,10 @@ async def test_delta_channel_time_travel() -> None: from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] counter = {"n": 0} @@ -9507,10 +9507,10 @@ async def test_delta_channel_remove_message_end_to_end() -> None: from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] def respond(state: State) -> dict: return {"messages": [AIMessage(content="reply", id="ai-1")]} @@ -9553,10 +9553,10 @@ async def test_delta_channel_update_by_id_end_to_end() -> None: from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] def update_msg(state: State) -> dict: # re-send h1 with updated content @@ -9596,10 +9596,10 @@ async def test_delta_channel_async_write_ordering() -> None: from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import START, StateGraph - from langgraph.graph.message import add_messages + from langgraph.graph.message import _messages_delta_reducer class State(TypedDict): - messages: Annotated[list, DeltaChannel(list, add_messages)] + messages: Annotated[list, DeltaChannel(_messages_delta_reducer)] def respond(state: State) -> dict: i = len(state["messages"]) @@ -9616,8 +9616,7 @@ async def test_delta_channel_async_write_ordering() -> None: async def tracked_aput(self, config, checkpoint, metadata, new_versions): has_sentinel = any( - v is DELTA_SENTINEL - for v in checkpoint.get("channel_values", {}).values() + v is DELTA_SENTINEL for v in checkpoint.get("channel_values", {}).values() ) order.append("aput_sentinel" if has_sentinel else "aput_other") return await original_aput(self, config, checkpoint, metadata, new_versions)