mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-22 09:35:07 +02:00
Sydney RunkleandGitHub
9a27693c64
When a DeltaChannel receives an Overwrite, force that checkpoint to store a snapshot so sparse replay starts from the post-overwrite value. This also aligns live DeltaChannel overwrite handling with BinaryOperatorAggregate by letting Overwrite bypass other reducer writes in the same superstep.
804 lines
28 KiB
Python
804 lines
28 KiB
Python
import operator
|
|
from collections.abc import Sequence
|
|
from typing import Annotated
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
|
from typing_extensions import NotRequired, TypedDict
|
|
|
|
from langgraph._internal._typing import MISSING
|
|
from langgraph.channels.binop import BinaryOperatorAggregate
|
|
from langgraph.channels.delta import DeltaChannel
|
|
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 import START, StateGraph
|
|
from langgraph.graph.message import _messages_delta_reducer
|
|
from langgraph.graph.state import _get_channel
|
|
from langgraph.types import Overwrite
|
|
|
|
pytestmark = pytest.mark.anyio
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core channel primitives
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_last_value() -> None:
|
|
channel = LastValue(int).from_checkpoint(MISSING)
|
|
assert channel.ValueType is int
|
|
assert channel.UpdateType is int
|
|
|
|
with pytest.raises(EmptyChannelError):
|
|
channel.get()
|
|
with pytest.raises(InvalidUpdateError):
|
|
channel.update([5, 6])
|
|
|
|
channel.update([3])
|
|
assert channel.get() == 3
|
|
channel.update([4])
|
|
assert channel.get() == 4
|
|
checkpoint = channel.checkpoint()
|
|
channel = LastValue(int).from_checkpoint(checkpoint)
|
|
assert channel.get() == 4
|
|
|
|
|
|
def test_topic() -> None:
|
|
channel = Topic(str).from_checkpoint(MISSING)
|
|
assert channel.ValueType == Sequence[str]
|
|
assert channel.UpdateType == str | list[str]
|
|
|
|
assert channel.update(["a", "b"])
|
|
assert channel.get() == ["a", "b"]
|
|
assert channel.update([["c", "d"], "d"])
|
|
assert channel.get() == ["c", "d", "d"]
|
|
assert channel.update([])
|
|
with pytest.raises(EmptyChannelError):
|
|
channel.get()
|
|
assert not channel.update([]), "channel already empty"
|
|
assert channel.update(["e"])
|
|
assert channel.get() == ["e"]
|
|
checkpoint = channel.checkpoint()
|
|
channel = Topic(str).from_checkpoint(checkpoint)
|
|
assert channel.get() == ["e"]
|
|
channel_copy = Topic(str).from_checkpoint(checkpoint)
|
|
channel_copy.update(["f"])
|
|
assert channel_copy.get() == ["f"]
|
|
assert channel.get() == ["e"]
|
|
|
|
|
|
def test_topic_accumulate() -> None:
|
|
channel = Topic(str, accumulate=True).from_checkpoint(MISSING)
|
|
assert channel.ValueType == Sequence[str]
|
|
assert channel.UpdateType == str | list[str]
|
|
|
|
assert channel.update(["a", "b"])
|
|
assert channel.get() == ["a", "b"]
|
|
assert channel.update(["b", ["c", "d"], "d"])
|
|
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
|
|
assert not channel.update([])
|
|
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
|
|
checkpoint = channel.checkpoint()
|
|
channel = Topic(str, accumulate=True).from_checkpoint(checkpoint)
|
|
assert channel.get() == ["a", "b", "b", "c", "d", "d"]
|
|
assert channel.update(["e"])
|
|
assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
|
|
|
|
|
|
def test_binop() -> None:
|
|
channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(MISSING)
|
|
assert channel.ValueType is int
|
|
assert channel.UpdateType is int
|
|
|
|
assert channel.get() == 0
|
|
|
|
channel.update([1, 2, 3])
|
|
assert channel.get() == 6
|
|
channel.update([4])
|
|
assert channel.get() == 10
|
|
checkpoint = channel.checkpoint()
|
|
channel = BinaryOperatorAggregate(int, operator.add).from_checkpoint(checkpoint)
|
|
assert channel.get() == 10
|
|
|
|
|
|
def test_untracked_value() -> None:
|
|
channel = UntrackedValue(dict).from_checkpoint(MISSING)
|
|
assert channel.ValueType is dict
|
|
assert channel.UpdateType is dict
|
|
|
|
with pytest.raises(EmptyChannelError):
|
|
channel.get()
|
|
|
|
test_data = {"session": "test", "temp": "dir"}
|
|
channel.update([test_data])
|
|
assert channel.get() == test_data
|
|
|
|
new_data = {"session": "updated", "temp": "newdir"}
|
|
channel.update([new_data])
|
|
assert channel.get() == new_data
|
|
|
|
checkpoint = channel.checkpoint()
|
|
assert checkpoint is MISSING
|
|
|
|
new_channel = UntrackedValue(dict).from_checkpoint(checkpoint)
|
|
with pytest.raises(EmptyChannelError):
|
|
new_channel.get()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeltaChannel — message reducer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_delta_channel_basic_two_steps() -> None:
|
|
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
|
|
|
ch.update([HumanMessage(content="hi", id="h1")])
|
|
d1 = ch.checkpoint()
|
|
assert d1 is MISSING
|
|
|
|
ch.update([AIMessage(content="hello", id="a1")])
|
|
d2 = ch.checkpoint()
|
|
assert d2 is MISSING
|
|
|
|
assert len(ch.get()) == 2
|
|
assert ch.get()[0].content == "hi"
|
|
assert ch.get()[1].content == "hello"
|
|
|
|
|
|
def test_delta_channel_from_checkpoint_writes_list() -> None:
|
|
"""replay_writes on a fresh channel replays through the operator."""
|
|
spec = DeltaChannel(_messages_delta_reducer, list)
|
|
ch = spec.from_checkpoint(MISSING)
|
|
ch.replay_writes(
|
|
[
|
|
("t0", "messages", HumanMessage(content="hi", id="h1")),
|
|
("t1", "messages", AIMessage(content="hello", id="a1")),
|
|
("t2", "messages", HumanMessage(content="bye", id="h2")),
|
|
]
|
|
)
|
|
msgs = ch.get()
|
|
assert len(msgs) == 3
|
|
assert msgs[0].content == "hi"
|
|
assert msgs[1].content == "hello"
|
|
assert msgs[2].content == "bye"
|
|
|
|
|
|
def test_delta_channel_from_checkpoint_backwards_compat() -> None:
|
|
spec = DeltaChannel(_messages_delta_reducer, list)
|
|
old_value = [HumanMessage(content="old", id="h1")]
|
|
ch = spec.from_checkpoint(old_value)
|
|
assert ch.get() == old_value
|
|
|
|
|
|
def test_delta_channel_overwrite() -> None:
|
|
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
|
ch.update([HumanMessage(content="old", id="h1")])
|
|
|
|
ch.update([Overwrite([HumanMessage(content="new", id="h2")])])
|
|
d = ch.checkpoint()
|
|
assert d is MISSING
|
|
assert len(ch.get()) == 1
|
|
assert ch.get()[0].content == "new"
|
|
|
|
|
|
def test_overwrite_dataclass_form_survives_json_roundtrip() -> None:
|
|
"""`Overwrite` serialised with `orjson` collapses to a plain dict but
|
|
must still be recognised as an overwrite by the channel reducer.
|
|
|
|
Without the `type` discriminator the dataclass-erased shape (`{"value":
|
|
...}`) is indistinguishable from a literal channel value, and downstream
|
|
reducers raise `MESSAGE_COERCION_FAILURE` (or similar) on read.
|
|
"""
|
|
import orjson
|
|
|
|
from langgraph._internal._constants import OVERWRITE
|
|
from langgraph.channels.binop import _get_overwrite
|
|
|
|
ow = Overwrite(value=[HumanMessage(content="new", id="h2")])
|
|
erased = orjson.loads(orjson.dumps(ow, default=lambda o: o.model_dump()))
|
|
|
|
assert erased["type"] == OVERWRITE
|
|
is_overwrite, value = _get_overwrite(erased)
|
|
assert is_overwrite
|
|
assert isinstance(value, list)
|
|
assert value[0]["content"] == "new"
|
|
|
|
|
|
def test_overwrite_sentinel_dict_still_recognised() -> None:
|
|
"""The pre-existing `{"__overwrite__": value}` dict form continues to be
|
|
recognised. This is the canonical sentinel emitted by producers that do
|
|
not have an `Overwrite` dataclass available."""
|
|
from langgraph._internal._constants import OVERWRITE
|
|
from langgraph.channels.binop import _get_overwrite
|
|
|
|
is_overwrite, value = _get_overwrite({OVERWRITE: ["b"]})
|
|
assert is_overwrite
|
|
assert value == ["b"]
|
|
|
|
|
|
def test_overwrite_non_matching_dict_not_recognised() -> None:
|
|
"""Dicts that resemble the erased shape but do not carry the
|
|
`__overwrite__` discriminator must not be misclassified as overwrites."""
|
|
from langgraph.channels.binop import _get_overwrite
|
|
|
|
assert _get_overwrite({"value": ["b"]}) == (False, None)
|
|
assert _get_overwrite({"type": "human", "value": "hi"}) == (False, None)
|
|
|
|
|
|
def test_delta_channel_remove_message_and_replay() -> None:
|
|
"""RemoveMessage must round-trip correctly when writes are replayed."""
|
|
spec = DeltaChannel(_messages_delta_reducer, list)
|
|
ch = spec.from_checkpoint(MISSING)
|
|
|
|
ch.update([HumanMessage(content="hi", id="h1")])
|
|
ch.update([AIMessage(content="hello", id="a1")])
|
|
assert ch.get() == [
|
|
HumanMessage(content="hi", id="h1"),
|
|
AIMessage(content="hello", id="a1"),
|
|
]
|
|
|
|
ch.update([RemoveMessage(id="a1")])
|
|
assert ch.get() == [HumanMessage(content="hi", id="h1")]
|
|
|
|
ch2 = spec.from_checkpoint(MISSING)
|
|
ch2.replay_writes(
|
|
[
|
|
("t0", "messages", HumanMessage(content="hi", id="h1")),
|
|
("t1", "messages", AIMessage(content="hello", id="a1")),
|
|
("t2", "messages", RemoveMessage(id="a1")),
|
|
]
|
|
)
|
|
assert ch2.get() == [HumanMessage(content="hi", id="h1")]
|
|
|
|
|
|
def test_delta_channel_update_by_id_and_replay() -> None:
|
|
"""Updating a message by ID must round-trip correctly through writes replay."""
|
|
spec = DeltaChannel(_messages_delta_reducer, list)
|
|
ch = spec.from_checkpoint(MISSING)
|
|
|
|
ch.update([HumanMessage(content="original", id="h1")])
|
|
ch.update([HumanMessage(content="updated", id="h1")])
|
|
assert ch.get() == [HumanMessage(content="updated", id="h1")]
|
|
|
|
ch2 = spec.from_checkpoint(MISSING)
|
|
ch2.replay_writes(
|
|
[
|
|
("t0", "messages", HumanMessage(content="original", id="h1")),
|
|
("t1", "messages", HumanMessage(content="updated", id="h1")),
|
|
]
|
|
)
|
|
assert len(ch2.get()) == 1
|
|
assert ch2.get()[0].content == "updated"
|
|
|
|
|
|
def test_delta_channel_dict_coercion() -> None:
|
|
"""_messages_delta_reducer coerces dict writes to BaseMessage objects.
|
|
|
|
HTTP-driven input always arrives as JSON dicts. The reducer must coerce
|
|
them (same contract as add_messages) so graphs work without a separate
|
|
coercion step.
|
|
"""
|
|
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
|
|
|
# dict input — simulates what arrives from the HTTP API
|
|
ch.update([{"role": "human", "content": "hello", "id": "h1"}])
|
|
assert len(ch.get()) == 1
|
|
assert isinstance(ch.get()[0], HumanMessage)
|
|
assert ch.get()[0].content == "hello"
|
|
assert ch.get()[0].id == "h1"
|
|
|
|
# update by ID via dict
|
|
ch.update([{"role": "ai", "content": "world", "id": "h1"}])
|
|
assert len(ch.get()) == 1
|
|
assert ch.get()[0].content == "world"
|
|
|
|
# remove via RemoveMessage instance (same contract as add_messages)
|
|
ch.update([RemoveMessage(id="h1")])
|
|
assert ch.get() == []
|
|
|
|
|
|
def test_messages_delta_reducer_coerces_state() -> None:
|
|
"""State (left side) is coerced when raw — supports raw initial input
|
|
and deserialized blobs. The steady-state path (state already typed)
|
|
short-circuits and skips coercion.
|
|
"""
|
|
state = [{"role": "human", "content": "hello", "id": "h1"}]
|
|
writes = [[{"role": "ai", "content": "world", "id": "h1"}]]
|
|
result = _messages_delta_reducer(state, writes) # type: ignore[arg-type]
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], AIMessage)
|
|
assert result[0].content == "world"
|
|
assert result[0].id == "h1"
|
|
|
|
|
|
def test_messages_delta_reducer_tuple_write_is_one_message() -> None:
|
|
"""A top-level tuple write is one message-like, not a sequence to flatten.
|
|
|
|
`("user", "hi")` is a valid `MessageLikeRepresentation`; flattening it
|
|
would produce two HumanMessages ("user", "hi") instead of one.
|
|
"""
|
|
result = _messages_delta_reducer([], [("user", "hi")]) # type: ignore[arg-type]
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], HumanMessage)
|
|
assert result[0].content == "hi"
|
|
|
|
|
|
def test_delta_channel_checkpoint_returns_missing() -> None:
|
|
"""checkpoint() always returns MISSING regardless of state.
|
|
|
|
Pregel writes `_DeltaSnapshot(ch.get())` directly into `channel_values`
|
|
on snapshot steps; the channel itself never participates in snapshot
|
|
serialization, so its `checkpoint()` is always the absence sentinel.
|
|
"""
|
|
ch = DeltaChannel(_messages_delta_reducer, list).from_checkpoint(MISSING)
|
|
assert ch.checkpoint() is MISSING
|
|
|
|
ch.update([HumanMessage(content="hi", id="h1")])
|
|
assert ch.checkpoint() is MISSING
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeltaChannel — snapshot frequency
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_delta_channel_snapshot_version_based() -> None:
|
|
"""Snapshots fire when a channel accumulates `snapshot_frequency` updates.
|
|
|
|
Under the version-delta cadence, every time the channel's
|
|
`current_version - last_snapshot_version >= snapshot_frequency` a
|
|
`_DeltaSnapshot` blob is written. Bounds the ancestor walk to at most
|
|
`snapshot_frequency` steps on any read for that channel.
|
|
"""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[
|
|
list, DeltaChannel(_messages_delta_reducer, snapshot_frequency=5)
|
|
]
|
|
other: str
|
|
|
|
def node_a(state: State) -> dict:
|
|
i = len(state["messages"]) // 2
|
|
return {"messages": [AIMessage(content=f"a{i}", id=f"a{i}")]}
|
|
|
|
def node_b(state: State) -> dict:
|
|
return {"other": "y"}
|
|
|
|
g = StateGraph(State)
|
|
g.add_node("a", node_a)
|
|
g.add_node("b", node_b)
|
|
g.add_edge(START, "a")
|
|
g.add_edge("a", "b")
|
|
saver = InMemorySaver()
|
|
graph = g.compile(checkpointer=saver)
|
|
|
|
config = {"configurable": {"thread_id": "t1"}}
|
|
for i in range(6):
|
|
graph.invoke(
|
|
{"messages": [HumanMessage(content=f"h{i}", id=f"h{i}")], "other": ""},
|
|
config,
|
|
)
|
|
|
|
msg_blob_values = [
|
|
saver.serde.loads_typed((type_tag, blob))
|
|
for k, (type_tag, blob) in saver.blobs.items()
|
|
if k[2] == "messages" and type_tag == "msgpack" and blob
|
|
]
|
|
snapshots = [v for v in msg_blob_values if isinstance(v, _DeltaSnapshot)]
|
|
assert snapshots, "expected at least one _DeltaSnapshot blob for messages"
|
|
|
|
state = graph.get_state(config)
|
|
assert len(state.values["messages"]) == 12 # 6 human + 6 AI
|
|
|
|
|
|
# TODO(delta-channel-cadence): the previous "snapshot fires even when channel
|
|
# was not written" test asserted eager step-based snapshotting; under the new
|
|
# version-delta cadence (`should_snapshot` triggers on per-channel update
|
|
# count, not superstep count), no snapshot fires for an unwritten channel.
|
|
# Replace with a test that exercises the version-delta trigger plus the
|
|
# durability="exit" force-snapshot branch — see
|
|
# `docs/superpowers/specs/2026-05-04-delta-channel-batched-reads-design.md`
|
|
# section "Snapshot cadence".
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeltaChannel — end-to-end (InMemorySaver)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
|
"""InMemorySaver assembles writes from checkpoint_writes inside get_tuple."""
|
|
|
|
class State(TypedDict):
|
|
messages: Annotated[list, DeltaChannel(_messages_delta_reducer, list)]
|
|
|
|
n = {"v": 0}
|
|
|
|
def respond(state: State) -> dict:
|
|
n["v"] += 1
|
|
return {"messages": [AIMessage(content=f"ok{n['v']}", id=f"ai{n['v']}")]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("respond", respond)
|
|
builder.add_edge(START, "respond")
|
|
saver = InMemorySaver()
|
|
graph = builder.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "t1"}}
|
|
|
|
graph.invoke({"messages": [HumanMessage(content="hi", id="h1")]}, config)
|
|
graph.invoke({"messages": [HumanMessage(content="bye", id="h2")]}, config)
|
|
|
|
saved = saver.get_tuple(config)
|
|
assert saved is not None
|
|
assert "messages" not in saved.checkpoint["channel_values"]
|
|
|
|
state = graph.get_state(config)
|
|
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
|
|
|
|
|
|
def test_delta_channel_overwrite_superstep_snapshots() -> None:
|
|
def reducer(state: list[str], writes: Sequence[list[str]]) -> list[str]:
|
|
result = list(state)
|
|
for write in writes:
|
|
result.extend(write)
|
|
return result
|
|
|
|
class State(TypedDict):
|
|
items: Annotated[
|
|
list[str], DeltaChannel(reducer, list, snapshot_frequency=1000)
|
|
]
|
|
|
|
def node_a(state: State) -> dict:
|
|
return {"items": ["a"]}
|
|
|
|
def node_b(state: State) -> dict:
|
|
return {"items": Overwrite(["b"])}
|
|
|
|
def node_c(state: State) -> dict:
|
|
return {"items": ["c"]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node_a", node_a)
|
|
builder.add_node("node_b", node_b)
|
|
builder.add_node("node_c", node_c)
|
|
builder.add_edge(START, "node_a")
|
|
builder.add_edge("node_a", "node_b")
|
|
builder.add_edge("node_a", "node_c")
|
|
|
|
saver = InMemorySaver()
|
|
graph = builder.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "overwrite-snapshot"}}
|
|
|
|
result = graph.invoke({"items": ["START"]}, config)
|
|
assert result == {"items": ["b"]}
|
|
|
|
saved = saver.get_tuple(config)
|
|
assert saved is not None
|
|
snapshot = saved.checkpoint["channel_values"].get("items")
|
|
assert isinstance(snapshot, _DeltaSnapshot)
|
|
assert snapshot.value == ["b"]
|
|
assert saved.metadata.get("counters_since_delta_snapshot", {}).get("items") is None
|
|
|
|
|
|
def test_delta_channel_replay_after_overwrite_snapshot() -> None:
|
|
def reducer(state: list[str], writes: Sequence[list[str]]) -> list[str]:
|
|
result = list(state)
|
|
for write in writes:
|
|
result.extend(write)
|
|
return result
|
|
|
|
class State(TypedDict):
|
|
items: Annotated[
|
|
list[str], DeltaChannel(reducer, list, snapshot_frequency=1000)
|
|
]
|
|
|
|
calls = 0
|
|
|
|
def node(state: State) -> dict:
|
|
nonlocal calls
|
|
calls += 1
|
|
if calls == 1:
|
|
return {"items": Overwrite(["reset"])}
|
|
return {"items": ["after"]}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("node", node)
|
|
builder.add_edge(START, "node")
|
|
|
|
saver = InMemorySaver()
|
|
graph = builder.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "overwrite-replay"}}
|
|
|
|
assert graph.invoke({"items": ["before"]}, config) == {"items": ["reset"]}
|
|
first_saved = saver.get_tuple(config)
|
|
assert first_saved is not None
|
|
assert isinstance(
|
|
first_saved.checkpoint["channel_values"].get("items"), _DeltaSnapshot
|
|
)
|
|
|
|
assert graph.invoke({"items": []}, config) == {"items": ["reset", "after"]}
|
|
second_saved = saver.get_tuple(config)
|
|
assert second_saved is not None
|
|
assert "items" not in second_saved.checkpoint["channel_values"]
|
|
assert graph.get_state(config).values == {"items": ["reset", "after"]}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeltaChannel — dict reducer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _delta_channel_with_type(op, typ):
|
|
"""Build a DeltaChannel with an explicit type via the Annotated injection path."""
|
|
return _get_channel("_test", Annotated[typ, DeltaChannel(op)])
|
|
|
|
|
|
def test_delta_channel_dict_reducer_fresh_channel() -> None:
|
|
"""DeltaChannel with a dict reducer starts as empty dict on MISSING checkpoint."""
|
|
|
|
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()
|
|
assert ch.get() == {}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_basic_updates() -> None:
|
|
"""DeltaChannel with a dict reducer accumulates key/value pairs across steps."""
|
|
|
|
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}])
|
|
d1 = ch.checkpoint()
|
|
assert d1 is MISSING
|
|
|
|
ch.update([{"b": 2}])
|
|
d2 = ch.checkpoint()
|
|
assert d2 is MISSING
|
|
|
|
assert ch.get() == {"a": 1, "b": 2}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_writes_reconstruction() -> None:
|
|
"""replay_writes on a fresh channel replays through a dict merge reducer."""
|
|
|
|
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(MISSING)
|
|
ch.replay_writes(
|
|
[
|
|
("t0", "files", {"a": 1}),
|
|
("t1", "files", {"b": 2}),
|
|
("t2", "files", {"c": 3}),
|
|
]
|
|
)
|
|
assert ch.get() == {"a": 1, "b": 2, "c": 3}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_with_deletions() -> None:
|
|
"""Dict reducer that treats None values as deletions works end-to-end."""
|
|
|
|
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)
|
|
ch.update([{"file1.py": "content1", "file2.py": "content2"}])
|
|
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
|
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
|
|
|
spec = _delta_channel_with_type(merge_files, dict)
|
|
ch2 = spec.from_checkpoint(MISSING)
|
|
ch2.replay_writes(
|
|
[
|
|
("t0", "files", {"file1.py": "content1", "file2.py": "content2"}),
|
|
("t1", "files", {"file1.py": None, "file3.py": "content3"}),
|
|
]
|
|
)
|
|
assert ch2.get() == {"file2.py": "content2", "file3.py": "content3"}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_overwrite_in_update() -> None:
|
|
"""Overwrite(dict) in update() must preserve dict shape, not coerce to list."""
|
|
|
|
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}])
|
|
ch.update([Overwrite({"b": 2, "c": 3})])
|
|
assert ch.get() == {"b": 2, "c": 3}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_overwrite_in_writes_replay() -> None:
|
|
"""Overwrite(dict) embedded in replayed writes must reconstruct as dict."""
|
|
|
|
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(MISSING)
|
|
ch.replay_writes(
|
|
[
|
|
("t0", "files", {"a": 1}),
|
|
("t1", "files", Overwrite({"x": 10, "y": 20})),
|
|
("t2", "files", {"z": 30}),
|
|
]
|
|
)
|
|
assert ch.get() == {"x": 10, "y": 20, "z": 30}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_with_notrequired_annotation() -> None:
|
|
"""DeltaChannel infers dict type through `Annotated[NotRequired[dict[...]], ch]`."""
|
|
|
|
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(merge_dicts)]
|
|
ch = _get_channel("files", annotation).from_checkpoint(MISSING)
|
|
assert ch.get() == {}
|
|
ch.update([{"a": 1}])
|
|
ch.update([{"b": 2}])
|
|
assert ch.get() == {"a": 1, "b": 2}
|
|
|
|
|
|
def test_delta_channel_dict_reducer_end_to_end_filesystem() -> None:
|
|
"""End-to-end: graph with dict-reducer (filesystem-style) channel wrapped in DeltaChannel."""
|
|
|
|
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(merge_files)]
|
|
|
|
turn = {"v": 0}
|
|
|
|
def write_file(state: State) -> dict:
|
|
turn["v"] += 1
|
|
n = turn["v"]
|
|
return {"files": {f"/doc_{n}.txt": f"content for turn {n}"}}
|
|
|
|
builder = StateGraph(State)
|
|
builder.add_node("write_file", write_file)
|
|
builder.add_edge(START, "write_file")
|
|
saver = InMemorySaver()
|
|
graph = builder.compile(checkpointer=saver)
|
|
config = {"configurable": {"thread_id": "fs"}}
|
|
|
|
for _ in range(3):
|
|
graph.invoke({"files": {}}, config)
|
|
|
|
saved = saver.get_tuple(config)
|
|
assert saved is not None
|
|
assert "files" not in saved.checkpoint["channel_values"]
|
|
state = graph.get_state(config)
|
|
assert state.values["files"] == {
|
|
"/doc_1.txt": "content for turn 1",
|
|
"/doc_2.txt": "content for turn 2",
|
|
"/doc_3.txt": "content for turn 3",
|
|
}
|
|
|
|
def delete_file(state: State) -> dict:
|
|
return {"files": {"/doc_1.txt": None}}
|
|
|
|
builder2 = StateGraph(State)
|
|
builder2.add_node("write_file", write_file)
|
|
builder2.add_node("delete_file", delete_file)
|
|
builder2.add_edge(START, "write_file")
|
|
builder2.add_edge("write_file", "delete_file")
|
|
turn["v"] = 0
|
|
saver2 = InMemorySaver()
|
|
graph2 = builder2.compile(checkpointer=saver2)
|
|
config2 = {"configurable": {"thread_id": "fs2"}}
|
|
graph2.invoke({"files": {}}, config2)
|
|
state2 = graph2.get_state(config2)
|
|
assert state2.values["files"] == {}
|
|
|
|
|
|
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(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}
|
|
ch = spec.from_checkpoint(old_value)
|
|
assert ch.get() == {"a": 1, "b": 2}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# DeltaChannel — seed / pre-delta migration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_delta_channel_from_checkpoint_honors_seed() -> None:
|
|
"""A non-sentinel value to from_checkpoint is used as the pre-delta seed.
|
|
|
|
Guards the pre-delta migration path: when the saver's ancestor walk hits
|
|
a pre-DeltaChannel blob it passes it as `seed` so replay reconstructs
|
|
the post-migration state correctly rather than replaying from empty.
|
|
"""
|
|
spec = DeltaChannel(_messages_delta_reducer, list)
|
|
seed = [HumanMessage(content="pre-delta", id="p1")]
|
|
ch = spec.from_checkpoint(seed)
|
|
ch.replay_writes(
|
|
[
|
|
("t0", "messages", AIMessage(content="delta-1", id="d1")),
|
|
("t1", "messages", HumanMessage(content="delta-2", id="d2")),
|
|
]
|
|
)
|
|
msgs = ch.get()
|
|
assert [m.content for m in msgs] == ["pre-delta", "delta-1", "delta-2"]
|
|
|
|
|
|
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(_messages_delta_reducer, list)
|
|
seed = [HumanMessage(content="only-snap", id="s1")]
|
|
ch = spec.from_checkpoint(seed)
|
|
ch.replay_writes([])
|
|
assert ch.get() == seed
|
|
|
|
|
|
def test_delta_channel_from_checkpoint_seed_none_is_distinct_from_sentinel() -> None:
|
|
"""`seed=None` must start replay from None, not from an empty channel.
|
|
|
|
The `MISSING` absence sentinel means 'no seed'; passing `None`
|
|
explicitly should feed None to the reducer as the left operand.
|
|
"""
|
|
|
|
def replace(state, writes):
|
|
return writes[-1] if writes else state
|
|
|
|
spec = DeltaChannel(replace, list)
|
|
ch = spec.from_checkpoint(None)
|
|
ch.replay_writes([("t0", "x", "after")])
|
|
assert ch.get() == "after"
|