diff --git a/libs/langgraph/tests/test_channels.py b/libs/langgraph/tests/test_channels.py index bb528371c..0b2d4ce47 100644 --- a/libs/langgraph/tests/test_channels.py +++ b/libs/langgraph/tests/test_channels.py @@ -317,6 +317,233 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None: assert len(state.values["messages"]) == 4 # 2 human + 2 AI +# --------------------------------------------------------------------------- +# Dict-reducer tests +# --------------------------------------------------------------------------- + + +def _delta_channel_with_type(operator, typ): + """Build a DeltaChannel with an explicit type via the Annotated injection path.""" + from typing import Annotated + + from langgraph.channels.delta import DeltaChannel + from langgraph.graph.state import _get_channel + + 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} + + 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(left: dict, right: dict) -> dict: + return {**left, **right} + + ch = _delta_channel_with_type(merge_dicts, dict).from_checkpoint(MISSING) + + ch.update([{"a": 1}]) + d1 = ch.checkpoint() + assert d1 is DELTA_SENTINEL + + ch.update([{"b": 2}]) + d2 = ch.checkpoint() + assert d2 is DELTA_SENTINEL + + 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(left: dict, right: dict) -> dict: + return {**left, **right} + + spec = _delta_channel_with_type(merge_dicts, dict) + ch = spec.from_checkpoint(DELTA_SENTINEL) + 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(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 + 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(DELTA_SENTINEL) + 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.""" + from langgraph.types import Overwrite + + def merge_dicts(left: dict, right: dict) -> dict: + return {**left, **right} + + 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.""" + from langgraph.types import Overwrite + + def merge_dicts(left: dict, right: dict) -> dict: + return {**left, **right} + + spec = _delta_channel_with_type(merge_dicts, dict) + ch = spec.from_checkpoint(DELTA_SENTINEL) + 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]`.""" + from typing import Annotated + + from typing_extensions import NotRequired + + 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} + + 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.""" + from typing import Annotated + + from langgraph.checkpoint.memory import InMemorySaver + from typing_extensions import TypedDict + + 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 + 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 saved.checkpoint["channel_values"]["files"] is DELTA_SENTINEL + 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(left: dict, right: dict) -> dict: + return {**left, **right} + + 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} + + # --------------------------------------------------------------------------- # seed / pre-delta migration # ---------------------------------------------------------------------------