mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-20 23:11:43 +02:00
Merge branch 'main' into byhow/langgraph/finish-siblings-on-invalid-resume
This commit is contained in:
@@ -14,8 +14,8 @@ async def memory_checkpointer():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_validate_memory_base():
|
||||
"""InMemorySaver passes all base capability tests."""
|
||||
async def test_validate_memory():
|
||||
"""InMemorySaver passes the tests of every capability it implements."""
|
||||
report = await validate(memory_checkpointer)
|
||||
report.print_report()
|
||||
assert report.passed_all_base(), f"Base tests failed: {report.to_dict()}"
|
||||
assert report.passed_all(), f"Capability tests failed: {report.to_dict()}"
|
||||
|
||||
@@ -469,7 +469,7 @@ class PostgresSaver(BasePostgresSaver):
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
if not checkpoint_id:
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
|
||||
@@ -417,7 +417,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
|
||||
thread_id = config["configurable"]["thread_id"]
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
if not checkpoint_id:
|
||||
target = await self.aget_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
|
||||
@@ -164,6 +164,40 @@ def test_sync_walk_reads_nothing_newer_than_the_target(
|
||||
assert read == _ids_from_target_down(configs)
|
||||
|
||||
|
||||
def _empty_checkpoint_id(config: dict) -> dict:
|
||||
return {"configurable": {**config["configurable"], "checkpoint_id": ""}}
|
||||
|
||||
|
||||
async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
await saver.setup()
|
||||
configs = await _abuild_chain(saver)
|
||||
latest = await saver.aget_delta_channel_history(
|
||||
config=configs[-1], channels=[CHANNEL]
|
||||
)
|
||||
|
||||
result = await saver.aget_delta_channel_history(
|
||||
config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert latest[CHANNEL]["writes"]
|
||||
assert result == latest
|
||||
|
||||
|
||||
def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
|
||||
saver.setup()
|
||||
configs = _build_chain(saver)
|
||||
latest = saver.get_delta_channel_history(config=configs[-1], channels=[CHANNEL])
|
||||
|
||||
result = saver.get_delta_channel_history(
|
||||
config=_empty_checkpoint_id(configs[-1]), channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert latest[CHANNEL]["writes"]
|
||||
assert result == latest
|
||||
|
||||
|
||||
async def test_root_target_has_no_history_and_still_terminates(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
@@ -539,7 +539,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
if not checkpoint_id:
|
||||
target = self.get_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
|
||||
@@ -653,7 +653,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
||||
thread_id = str(config["configurable"]["thread_id"])
|
||||
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
||||
checkpoint_id = get_checkpoint_id(config)
|
||||
if checkpoint_id is None:
|
||||
if not checkpoint_id:
|
||||
target = await self.aget_tuple(config)
|
||||
if target is None:
|
||||
return {ch: {"writes": []} for ch in channels}
|
||||
|
||||
@@ -65,6 +65,39 @@ async def test_async_walk_reaches_parent_whatever_the_id_order(
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
EMPTY_CHECKPOINT_ID: dict[str, Any] = {
|
||||
"configurable": {**CONFIG["configurable"], "checkpoint_id": ""}
|
||||
}
|
||||
|
||||
|
||||
def test_sync_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = saver.put(CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {})
|
||||
saver.put_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
saver.put(root, _checkpoint("z-newer", {}), {}, {})
|
||||
|
||||
got = saver.get_delta_channel_history(
|
||||
config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
async def test_async_empty_checkpoint_id_reads_from_the_latest_checkpoint() -> None:
|
||||
async with AsyncSqliteSaver.from_conn_string(":memory:") as saver:
|
||||
root = await saver.aput(
|
||||
CONFIG, _checkpoint("a-older", {CHANNEL: "seed"}), {}, {}
|
||||
)
|
||||
await saver.aput_writes(root, [(CHANNEL, "write-root")], "task")
|
||||
await saver.aput(root, _checkpoint("z-newer", {}), {}, {})
|
||||
|
||||
got = await saver.aget_delta_channel_history(
|
||||
config=EMPTY_CHECKPOINT_ID, channels=[CHANNEL]
|
||||
)
|
||||
|
||||
assert got[CHANNEL] == EXPECTED
|
||||
|
||||
|
||||
def test_walk_reaches_root_of_long_chain_with_descending_ids() -> None:
|
||||
steps = 40
|
||||
with SqliteSaver.from_conn_string(":memory:") as saver:
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
"""Run delta-channel conformance capabilities against InMemorySaver."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
conformance = pytest.importorskip(
|
||||
"langgraph.checkpoint.conformance",
|
||||
reason="langgraph-checkpoint-conformance not installed",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delta_channel_conformance():
|
||||
# Imported inside the test: the module-level importorskip above is what
|
||||
# makes these safe, so they cannot move to the top of the file.
|
||||
from langgraph.checkpoint.conformance import validate # noqa: PLC0415
|
||||
from langgraph.checkpoint.conformance.initializer import ( # noqa: PLC0415
|
||||
checkpointer_test,
|
||||
)
|
||||
|
||||
from langgraph.checkpoint.memory import InMemorySaver # noqa: PLC0415
|
||||
|
||||
@checkpointer_test(name="InMemorySaver")
|
||||
async def mem_saver():
|
||||
yield InMemorySaver()
|
||||
|
||||
report = await validate(
|
||||
mem_saver,
|
||||
capabilities={
|
||||
"delta_channel_history",
|
||||
},
|
||||
)
|
||||
for cap, result in report.results.items():
|
||||
if result.passed is False:
|
||||
details = "\n".join(result.failures or [])
|
||||
pytest.fail(f"Capability {cap} failed:\n{details}")
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable, Mapping
|
||||
from collections.abc import Callable, Iterable, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from inspect import signature
|
||||
from typing import Any, Literal, cast
|
||||
@@ -32,6 +32,7 @@ from langgraph._internal._constants import (
|
||||
)
|
||||
from langgraph._internal._typing import MISSING
|
||||
from langgraph.channels.base import BaseChannel
|
||||
from langgraph.channels.binop import _get_overwrite
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.managed.base import ManagedValueMapping, ManagedValueSpec
|
||||
|
||||
@@ -152,6 +153,23 @@ def delta_channels_with_pending_writes(
|
||||
}
|
||||
|
||||
|
||||
def delta_channels_overwritten(
|
||||
specs: Mapping[str, Any], writes: Iterable[tuple[str, Any]]
|
||||
) -> set[str]:
|
||||
"""Return the names of the DeltaChannels that `writes` set with an `Overwrite`.
|
||||
|
||||
`update_state` saves a full snapshot of these channels in the checkpoint it
|
||||
creates, like the loop does when a node returns an `Overwrite`. Otherwise,
|
||||
reading the channel later starts from an older snapshot and replays the
|
||||
writes the `Overwrite` threw away.
|
||||
"""
|
||||
return {
|
||||
ch
|
||||
for ch, value in writes
|
||||
if isinstance(specs.get(ch), DeltaChannel) and _get_overwrite(value)[0]
|
||||
}
|
||||
|
||||
|
||||
def checkpoint_superseded(
|
||||
saver: BaseCheckpointSaver, config: RunnableConfig, saved: CheckpointTuple
|
||||
) -> bool:
|
||||
@@ -256,6 +274,7 @@ def create_checkpoint(
|
||||
get_next_version: GetNextVersion | None = None,
|
||||
channels_to_snapshot: set[str] | None = None,
|
||||
stored_versions: ChannelVersions | None = None,
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]] | None = None,
|
||||
) -> Checkpoint:
|
||||
"""Build a new Checkpoint from the previous one and live channel state.
|
||||
|
||||
@@ -314,7 +333,9 @@ def create_checkpoint(
|
||||
id=id or str(uuid6(clock_seq=step)),
|
||||
channel_values=values,
|
||||
channel_versions=channel_versions,
|
||||
versions_seen=_mark_bumps_seen(checkpoint["versions_seen"], bumped),
|
||||
versions_seen=_mark_bumps_seen(
|
||||
checkpoint["versions_seen"], bumped, trigger_to_nodes or {}
|
||||
),
|
||||
updated_channels=None if updated_channels is None else sorted(updated_channels),
|
||||
)
|
||||
|
||||
@@ -322,19 +343,27 @@ def create_checkpoint(
|
||||
def _mark_bumps_seen(
|
||||
versions_seen: dict[str, ChannelVersions],
|
||||
bumped: Mapping[str, tuple[Any, Any]],
|
||||
trigger_to_nodes: Mapping[str, Sequence[str]],
|
||||
) -> dict[str, ChannelVersions]:
|
||||
"""Advance whoever had seen a bumped channel's old version to the new one.
|
||||
|
||||
A bump that only stores a snapshot is not a write. Left unseen, it would
|
||||
re-fire `interrupt_before` and rerun the channel's subscribers. For each
|
||||
entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the one the
|
||||
node really read, so `versions_seen_without_bumps` can put the read back.
|
||||
re-fire `interrupt_before` and rerun the channel's subscribers. A channel
|
||||
bumped from no version was never written, so it also goes to the
|
||||
subscribers that never ran: they have no entry, and would start on the bump.
|
||||
For each entry it advances, `SNAPSHOT_BUMPS` keeps the new version and the
|
||||
one the node really read, so `versions_seen_without_bumps` can put the read
|
||||
back.
|
||||
"""
|
||||
if not bumped:
|
||||
return versions_seen
|
||||
out = dict(versions_seen)
|
||||
for k, (old, _) in bumped.items():
|
||||
if old is None:
|
||||
for node in trigger_to_nodes.get(k, ()):
|
||||
out.setdefault(node, {})
|
||||
marks = dict(versions_seen.get(SNAPSHOT_BUMPS, {}))
|
||||
for node, seen in versions_seen.items():
|
||||
for node, seen in out.items():
|
||||
if node == SNAPSHOT_BUMPS:
|
||||
continue
|
||||
for k, (old, new) in bumped.items():
|
||||
|
||||
@@ -19,6 +19,7 @@ from typing import (
|
||||
TypeVar,
|
||||
cast,
|
||||
)
|
||||
from uuid import UUID, uuid5
|
||||
|
||||
from langchain_core.callbacks import AsyncParentRunManager, ParentRunManager
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
@@ -269,6 +270,9 @@ class PregelLoop:
|
||||
# `_put_exit_delta_writes` uses this to decide between anchoring on
|
||||
# the existing parent (True) or creating a lazy stub (False).
|
||||
_has_persisted_parent: bool = False
|
||||
# True iff `__enter__` loaded the thread's latest checkpoint, not one a
|
||||
# `checkpoint_id` addressed, so nothing has been built on it yet.
|
||||
_loaded_latest: bool = False
|
||||
|
||||
managed: ManagedValueMapping
|
||||
checkpoint: Checkpoint
|
||||
@@ -1117,16 +1121,26 @@ class PregelLoop:
|
||||
self._exit_delta_writes.append(
|
||||
(self.step, NULL_TASK_ID, "", c, v)
|
||||
)
|
||||
# Persist delta-channel input writes so sub-freq inputs are
|
||||
# recoverable via ancestor walk (mirrors the Command input path).
|
||||
# A DeltaChannel reads its input from the writes stored on the
|
||||
# checkpoint this run starts from, under a task id of their own:
|
||||
# readers apply a checkpoint's NULL_TASK_ID writes as its own state.
|
||||
# A new thread has no checkpoint to store them on, and one a
|
||||
# `checkpoint_id` addressed may have children that would read them,
|
||||
# so then the input checkpoint snapshots the channel instead.
|
||||
if self.durability != "exit":
|
||||
delta_input = [
|
||||
(c, v)
|
||||
for c, v in input_writes
|
||||
if isinstance(self.specs.get(c), DeltaChannel)
|
||||
]
|
||||
if delta_input:
|
||||
self.put_writes(NULL_TASK_ID, delta_input)
|
||||
if delta_input and self._has_persisted_parent and self._loaded_latest:
|
||||
self.put_writes(
|
||||
str(uuid5(UUID(self.checkpoint["id"]), INPUT)), delta_input
|
||||
)
|
||||
else:
|
||||
self._delta_channels_forced_snapshot.update(
|
||||
c for c, _ in delta_input
|
||||
)
|
||||
# save input checkpoint
|
||||
self.updated_channels = updated_channels
|
||||
self._put_checkpoint({"source": "input"})
|
||||
@@ -1256,6 +1270,7 @@ class PregelLoop:
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=self.checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
for k in channels_to_snapshot:
|
||||
new_counters[k] = (0, 0)
|
||||
@@ -1774,6 +1789,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
|
||||
# Normal case: fetch the most recent checkpoint for this
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = self.checkpointer.get_tuple(self.checkpoint_config)
|
||||
self._loaded_latest = True
|
||||
|
||||
# Capture before the synthetic-empty fallback below overwrites `saved`.
|
||||
# `_put_exit_delta_writes` uses this on first run (no persisted parent)
|
||||
@@ -2030,6 +2046,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
|
||||
# Normal case: fetch the most recent checkpoint for this
|
||||
# graph/thread. Returns None on first invocation.
|
||||
saved = await self.checkpointer.aget_tuple(self.checkpoint_config)
|
||||
self._loaded_latest = True
|
||||
|
||||
# Capture before the synthetic-empty fallback below overwrites `saved`.
|
||||
# `_put_exit_delta_writes` uses this on first run (no persisted parent)
|
||||
|
||||
@@ -137,6 +137,7 @@ from langgraph.pregel._checkpoint import (
|
||||
copy_checkpoint,
|
||||
create_checkpoint,
|
||||
create_checkpoint_plan_for_update_state_api,
|
||||
delta_channels_overwritten,
|
||||
delta_channels_with_pending_writes,
|
||||
empty_checkpoint,
|
||||
get_updated_channels_from_tasks,
|
||||
@@ -1766,6 +1767,7 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
@@ -1795,6 +1797,9 @@ class Pregel(
|
||||
checkpointer.get_next_version,
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, input_writes
|
||||
)
|
||||
|
||||
# apply input write to channels
|
||||
next_step = (
|
||||
@@ -1822,6 +1827,7 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
@@ -2033,6 +2039,9 @@ class Pregel(
|
||||
),
|
||||
)
|
||||
updated_channels = get_updated_channels_from_tasks(run_tasks)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, (w for t in run_tasks for w in t.writes)
|
||||
)
|
||||
# The base's other children replay whatever is stored on it, so an
|
||||
# edit of an older checkpoint stores none of its writes there: the
|
||||
# checkpoint written here carries them, its delta channels
|
||||
@@ -2086,6 +2095,7 @@ class Pregel(
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = checkpointer.put(
|
||||
checkpoint_config,
|
||||
@@ -2258,6 +2268,7 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
@@ -2287,6 +2298,9 @@ class Pregel(
|
||||
checkpointer.get_next_version,
|
||||
self.trigger_to_nodes,
|
||||
)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, input_writes
|
||||
)
|
||||
|
||||
# apply input write to channels
|
||||
next_step = (
|
||||
@@ -2314,6 +2328,7 @@ class Pregel(
|
||||
get_next_version=checkpointer.get_next_version,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
@@ -2524,6 +2539,9 @@ class Pregel(
|
||||
),
|
||||
)
|
||||
updated_channels = get_updated_channels_from_tasks(run_tasks)
|
||||
fork_pending |= delta_channels_overwritten(
|
||||
self.channels, (w for t in run_tasks for w in t.writes)
|
||||
)
|
||||
# The base's other children replay whatever is stored on it, so an
|
||||
# edit of an older checkpoint stores none of its writes there: the
|
||||
# checkpoint written here carries them, its delta channels
|
||||
@@ -2577,6 +2595,7 @@ class Pregel(
|
||||
else None,
|
||||
channels_to_snapshot=channels_to_snapshot,
|
||||
stored_versions=checkpoint_previous_versions,
|
||||
trigger_to_nodes=self.trigger_to_nodes,
|
||||
)
|
||||
next_config = await checkpointer.aput(
|
||||
checkpoint_config,
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
"""A run's input to a DeltaChannel input channel reads back on the checkpoints
|
||||
built from it, and on no others."""
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Durability
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
def _sorted_extend(current: list, writes: list) -> list:
|
||||
return sorted([*current, *(item for write in writes for item in write)])
|
||||
|
||||
|
||||
def _delta_input_graph() -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend),
|
||||
"plain": BinaryOperatorAggregate(list, lambda a, b: sorted(a + b)),
|
||||
"go": LastValue(int),
|
||||
},
|
||||
input_channels=["log", "plain", "go"],
|
||||
output_channels=["log", "plain"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
def test_each_run_input_reads_back_on_its_own_checkpoints(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
graph.invoke({"log": [0], "plain": [0], "go": 1}, config, durability=durability)
|
||||
graph.invoke({"log": [5], "plain": [5], "go": 1}, config, durability=durability)
|
||||
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
async def test_each_run_input_reads_back_on_its_own_checkpoints_async(
|
||||
durability: Durability,
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
await graph.ainvoke(
|
||||
{"log": [0], "plain": [0], "go": 1}, config, durability=durability
|
||||
)
|
||||
await graph.ainvoke(
|
||||
{"log": [5], "plain": [5], "go": 1}, config, durability=durability
|
||||
)
|
||||
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
OTHER_BRANCH_INPUTS = pytest.mark.parametrize(
|
||||
"other_branch_input",
|
||||
[{"go": 1}, {"log": [5], "plain": [5], "go": 1}],
|
||||
ids=["other-branch-without-delta-input", "other-branch-with-delta-input"],
|
||||
)
|
||||
|
||||
|
||||
@OTHER_BRANCH_INPUTS
|
||||
def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch(
|
||||
durability: Durability, other_branch_input: dict
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"go": 1}, config, durability=durability)
|
||||
older = graph.get_state(config).config
|
||||
graph.invoke(other_branch_input, config, durability=durability)
|
||||
|
||||
graph.invoke({"log": [7], "plain": [7], "go": 1}, older, durability=durability)
|
||||
|
||||
for state in graph.get_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
|
||||
|
||||
@OTHER_BRANCH_INPUTS
|
||||
async def test_run_input_from_an_older_checkpoint_stays_out_of_its_other_branch_async(
|
||||
durability: Durability, other_branch_input: dict
|
||||
) -> None:
|
||||
graph = _delta_input_graph()
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"go": 1}, config, durability=durability)
|
||||
older = (await graph.aget_state(config)).config
|
||||
await graph.ainvoke(other_branch_input, config, durability=durability)
|
||||
|
||||
await graph.ainvoke(
|
||||
{"log": [7], "plain": [7], "go": 1}, older, durability=durability
|
||||
)
|
||||
|
||||
async for state in graph.aget_state_history(config):
|
||||
assert state.values.get("log", []) == state.values.get("plain", [])
|
||||
@@ -0,0 +1,78 @@
|
||||
import pytest
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
CONFIG = {"configurable": {"thread_id": "t"}}
|
||||
|
||||
|
||||
def _extend(current: list, writes: list) -> list:
|
||||
return [*current, *(item for write in writes for item in write)]
|
||||
|
||||
|
||||
def _graph(reads: list) -> Pregel:
|
||||
writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d")
|
||||
reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d)))
|
||||
return Pregel(
|
||||
nodes={"writer": writer, "reader": reader},
|
||||
channels={"a": LastValue(str), "d": DeltaChannel(_extend)},
|
||||
input_channels="a",
|
||||
output_channels=["d"],
|
||||
checkpointer=InMemorySaver(),
|
||||
)
|
||||
|
||||
|
||||
def test_an_input_update_from_before_the_first_write_starts_only_the_writer() -> None:
|
||||
reads: list = []
|
||||
graph = _graph(reads)
|
||||
graph.invoke("go", CONFIG)
|
||||
first = next(
|
||||
s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
|
||||
)
|
||||
|
||||
fork = graph.update_state(first, {"a": "go"}, as_node="__input__")
|
||||
|
||||
assert graph.get_state(fork).next == ("writer",)
|
||||
graph.invoke(None, fork)
|
||||
assert reads == [[1], [1]]
|
||||
|
||||
|
||||
def test_a_replay_from_before_the_first_write_forks_with_only_the_writer_next() -> None:
|
||||
reads: list = []
|
||||
graph = _graph(reads)
|
||||
graph.invoke("go", CONFIG)
|
||||
first = next(
|
||||
s.config for s in graph.get_state_history(CONFIG) if s.metadata["step"] == -1
|
||||
)
|
||||
|
||||
graph.invoke(None, first, durability="sync")
|
||||
|
||||
fork = next(
|
||||
s for s in graph.get_state_history(CONFIG) if s.metadata["source"] == "fork"
|
||||
)
|
||||
assert fork.next == ("writer",)
|
||||
graph.invoke(None, fork.config)
|
||||
assert reads == [[1], [1], [1]]
|
||||
|
||||
|
||||
async def test_an_ainput_update_from_before_the_first_write_starts_only_the_writer() -> (
|
||||
None
|
||||
):
|
||||
reads: list = []
|
||||
graph = _graph(reads)
|
||||
await graph.ainvoke("go", CONFIG)
|
||||
first = [
|
||||
s.config
|
||||
async for s in graph.aget_state_history(CONFIG)
|
||||
if s.metadata["step"] == -1
|
||||
][0]
|
||||
|
||||
fork = await graph.aupdate_state(first, {"a": "go"}, as_node="__input__")
|
||||
|
||||
assert (await graph.aget_state(fork)).next == ("writer",)
|
||||
await graph.ainvoke(None, fork)
|
||||
assert reads == [[1], [1]]
|
||||
@@ -0,0 +1,110 @@
|
||||
"""An `Overwrite` through `update_state` snapshots its DeltaChannel on the
|
||||
checkpoint the update saves, as a node's `Overwrite` does on the loop's."""
|
||||
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import HumanMessage
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.checkpoint.serde.types import _DeltaSnapshot
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Overwrite
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
|
||||
class _State(TypedDict):
|
||||
messages: Annotated[list, DeltaChannel(_messages_delta_reducer)]
|
||||
|
||||
|
||||
def _messages_graph(saver: InMemorySaver) -> Any:
|
||||
builder = StateGraph(_State)
|
||||
builder.add_node("model", lambda state: {})
|
||||
builder.add_edge(START, "model")
|
||||
return builder.compile(checkpointer=saver)
|
||||
|
||||
|
||||
def _extend(current: list, writes: list) -> list:
|
||||
return [*current, *(item for write in writes for item in write)]
|
||||
|
||||
|
||||
def _delta_input_graph(saver: InMemorySaver) -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log")
|
||||
return Pregel(
|
||||
nodes={"n": node},
|
||||
channels={"log": DeltaChannel(_extend), "go": LastValue(int)},
|
||||
input_channels=["log", "go"],
|
||||
output_channels=["log"],
|
||||
checkpointer=saver,
|
||||
)
|
||||
|
||||
|
||||
def test_update_state_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _messages_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"messages": [HumanMessage(content="a", id="1")]}, config)
|
||||
|
||||
graph.update_state(
|
||||
config,
|
||||
{"messages": Overwrite([HumanMessage(content="b", id="2")])},
|
||||
as_node="model",
|
||||
)
|
||||
|
||||
head = saver.get_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
|
||||
assert [m.content for m in graph.get_state(config).values["messages"]] == ["b"]
|
||||
|
||||
|
||||
async def test_aupdate_state_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _messages_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"messages": [HumanMessage(content="a", id="1")]}, config)
|
||||
|
||||
await graph.aupdate_state(
|
||||
config,
|
||||
{"messages": Overwrite([HumanMessage(content="b", id="2")])},
|
||||
as_node="model",
|
||||
)
|
||||
|
||||
head = await saver.aget_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("messages"), _DeltaSnapshot)
|
||||
values = (await graph.aget_state(config)).values
|
||||
assert [m.content for m in values["messages"]] == ["b"]
|
||||
|
||||
|
||||
def test_update_state_as_input_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _delta_input_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
graph.invoke({"log": [0], "go": 1}, config)
|
||||
|
||||
graph.update_state(config, {"log": Overwrite([1])}, as_node="__input__")
|
||||
|
||||
head = saver.get_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
|
||||
assert graph.get_state(config).values["log"] == [1]
|
||||
|
||||
|
||||
async def test_aupdate_state_as_input_with_an_overwrite_snapshots_the_channel() -> None:
|
||||
saver = InMemorySaver()
|
||||
graph = _delta_input_graph(saver)
|
||||
config = {"configurable": {"thread_id": "t"}}
|
||||
await graph.ainvoke({"log": [0], "go": 1}, config)
|
||||
|
||||
await graph.aupdate_state(config, {"log": Overwrite([1])}, as_node="__input__")
|
||||
|
||||
head = await saver.aget_tuple(config)
|
||||
assert head is not None
|
||||
assert isinstance(head.checkpoint["channel_values"].get("log"), _DeltaSnapshot)
|
||||
assert (await graph.aget_state(config)).values["log"] == [1]
|
||||
Reference in New Issue
Block a user