mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 16:42:24 +02:00
fix: format, lint, tests for snapshot_frequency branch
- Remove unused AggregateChannel compat wrapper from test files; switch all DeltaChannel test aliases to import from channels.delta directly - Drop dict-reducer tests (tested AggregateChannel type-inference, not relevant to this DeltaChannel-only branch) - Add value: Value | Any annotation to AggregateChannel.__slots__ (mypy) - Add isinstance asserts for DeltaChannel before replay_writes calls - Remove now-unused _math_compat import and clean up PostgresSaver import Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Sydney Runkle
co-authored by
Claude Sonnet 4.6
Claude Sonnet 4.6
parent
ba12c8264d
commit
f06eb0f76c
@@ -6,8 +6,8 @@ from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
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
|
||||
@@ -127,7 +127,6 @@ def test_delta_channel_basic_two_steps() -> None:
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
@@ -152,7 +151,6 @@ 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.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
@@ -174,7 +172,6 @@ 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.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
# Old BinaryOperatorAggregate checkpoint: plain list treated as backward compat
|
||||
@@ -188,7 +185,6 @@ def test_delta_channel_overwrite() -> None:
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
@@ -207,7 +203,6 @@ 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.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
@@ -241,7 +236,6 @@ 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.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
spec = DeltaChannel(add_messages)
|
||||
@@ -270,7 +264,6 @@ 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.channels._delta import DeltaChannel
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
ch = DeltaChannel(add_messages).from_checkpoint(MISSING)
|
||||
@@ -290,7 +283,6 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
@@ -325,252 +317,6 @@ def test_delta_channel_inmemory_saver_assembles_writes() -> None:
|
||||
assert len(state.values["messages"]) == 4 # 2 human + 2 AI
|
||||
|
||||
|
||||
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)
|
||||
# Should be available (not raise EmptyChannelError) and start empty
|
||||
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."""
|
||||
from langgraph.checkpoint.base import DELTA_SENTINEL
|
||||
|
||||
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 (deepagents pattern)."""
|
||||
|
||||
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"}])
|
||||
|
||||
# Delete file1, add file3
|
||||
ch.update([{"file1.py": None, "file3.py": "content3"}])
|
||||
|
||||
assert ch.get() == {"file2.py": "content2", "file3.py": "content3"}
|
||||
|
||||
# Confirm writes reconstruction produces the same result
|
||||
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]`.
|
||||
|
||||
This is the shape the deepagents filesystem middleware uses for its
|
||||
`files` field; without unwrapping NotRequired we'd fall through to `list`
|
||||
and blow up on the first dict operator call.
|
||||
"""
|
||||
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.
|
||||
|
||||
Mirrors the deepagents filesystem pattern: `files: Annotated[dict, reducer]`
|
||||
where the reducer merges dicts and treats None values as deletions.
|
||||
"""
|
||||
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)
|
||||
|
||||
# Checkpoint stores only the sentinel — per-step writes live in checkpoint_writes.
|
||||
saved = saver.get_tuple(config)
|
||||
assert saved is not None
|
||||
cv = saved.checkpoint["channel_values"]["files"]
|
||||
assert cv 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",
|
||||
}
|
||||
|
||||
# Deletion path must round-trip through writes replay.
|
||||
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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -29,16 +29,6 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, StateGraph
|
||||
from langgraph.graph.message import add_messages
|
||||
|
||||
try:
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
|
||||
_POSTGRES_AVAILABLE = True
|
||||
_POSTGRES_URI = (
|
||||
"postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
|
||||
)
|
||||
except ImportError:
|
||||
_POSTGRES_AVAILABLE = False
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Realistic message payload (~100 tokens / ~400 chars each)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -251,7 +241,9 @@ def run_baseline_benchmark() -> None:
|
||||
|
||||
print("Storage (blob bytes)")
|
||||
print("-" * W)
|
||||
print(f"{'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}")
|
||||
print(
|
||||
f"{'turns':>6} {'ctx':>10} {'add_msgs':>12} {'delta(inf)':>12} {'savings':>8}"
|
||||
)
|
||||
print("-" * W)
|
||||
for turns, b_bytes, d_bytes, b_rt, d_rt, b_wt, d_wt in rows:
|
||||
if b_bytes is None or b_bytes < 0 or d_bytes is None or d_bytes < 0:
|
||||
@@ -315,7 +307,9 @@ def run_snapshot_freq_benchmark() -> None:
|
||||
# Storage table
|
||||
print()
|
||||
print("Storage (blob bytes) — lower is better")
|
||||
header = f"{'turns':>6} {'ctx':>10}" + "".join(f" {f'freq={l}':>{col_w}}" for l in freq_labels)
|
||||
header = f"{'turns':>6} {'ctx':>10}" + "".join(
|
||||
f" {f'freq={freq_label}':>{col_w}}" for freq_label in freq_labels
|
||||
)
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for turns in SWEEP_TURN_COUNTS:
|
||||
@@ -353,7 +347,9 @@ def run_snapshot_freq_benchmark() -> None:
|
||||
|
||||
print()
|
||||
print("Legend:")
|
||||
print(" freq=1 snapshot every write (full blob always — same as add_messages / BinOp)")
|
||||
print(
|
||||
" freq=1 snapshot every write (full blob always — same as add_messages / BinOp)"
|
||||
)
|
||||
print(" freq=N snapshot every N writes; read walks at most N ancestor writes")
|
||||
print(" freq=inf pure delta; read walks entire ancestor chain")
|
||||
print()
|
||||
|
||||
@@ -51,6 +51,7 @@ from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels._delta import DeltaChannel
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -41,6 +41,7 @@ from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.channels.topic import Topic
|
||||
|
||||
Reference in New Issue
Block a user