wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
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"
|