From 4dda404da37ec889e428699a0148e97f46f381e2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 2 Apr 2025 12:44:59 -0700 Subject: [PATCH] Implement checkpoint migration - Migrate start:{node} channels to branch:to:{node} - Migrate {node} channels to branch:to:{node} --- .../langgraph/checkpoint/base/__init__.py | 1 + libs/langgraph/langgraph/graph/state.py | 90 +- libs/langgraph/langgraph/pregel/__init__.py | 22 +- libs/langgraph/langgraph/pregel/checkpoint.py | 32 +- libs/langgraph/langgraph/pregel/loop.py | 14 + libs/langgraph/tests/any_str.py | 5 + .../tests/test_checkpoint_migration.py | 1374 +++++++++++++++++ 7 files changed, 1500 insertions(+), 38 deletions(-) create mode 100644 libs/langgraph/tests/test_checkpoint_migration.py diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index fbaa45fe4..9df070b18 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -455,6 +455,7 @@ def get_checkpoint_metadata( ) -> CheckpointMetadata: """Get checkpoint metadata in a backwards-compatible manner.""" metadata = metadata.copy() + print("metadata", metadata) for obj in (config.get("metadata"), config.get("configurable")): if not obj: continue diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 93fd65052..7d47d4b52 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -2,6 +2,7 @@ import inspect import logging import typing import warnings +from collections import defaultdict from functools import partial from inspect import isclass, isfunction, ismethod, signature from types import FunctionType @@ -35,7 +36,16 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue -from langgraph.constants import EMPTY_SEQ, MISSING, NS_END, NS_SEP, SELF, TAG_HIDDEN +from langgraph.checkpoint.base import Checkpoint +from langgraph.constants import ( + EMPTY_SEQ, + INTERRUPT, + MISSING, + NS_END, + NS_SEP, + SELF, + TAG_HIDDEN, +) from langgraph.errors import ( ErrorCode, InvalidUpdateError, @@ -922,6 +932,84 @@ class CompiledStateGraph(CompiledGraph): ) ) + def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None: + """Migrate a checkpoint to new channel layout.""" + + values = checkpoint["channel_values"] + versions = checkpoint["channel_versions"] + seen = checkpoint["versions_seen"] + + # empty checkpoints do not need migration + if not versions: + return + + # current version + if checkpoint["v"] >= 3: + return + + # migrate from v2 to v3 + if any(k.startswith("start:") for k in versions): + # Migrate from start:node to branch:to:node + for k in list(versions): + if k.startswith("start:"): + # confirm node is present + node = k.split(":")[1] + if node not in self.nodes: + continue + # get next version + new_k = f"branch:to:{node}" + new_v = ( + max(versions[new_k], versions.pop(k)) + if new_k in versions + else versions.pop(k) + ) + # update seen + for ss in (seen.get(node, {}), seen.get(INTERRUPT, {})): + if k in ss: + s = ss.pop(k) + if new_k in ss: + ss[new_k] = max(s, ss[new_k]) + else: + ss[new_k] = s + # update value + if new_k not in values and k in values: + values[new_k] = values.pop(k) + # update version + versions[new_k] = new_v + + if not set(self.nodes).isdisjoint(versions): + # Migrate from "node" to "branch:to:node" + source_to_target = defaultdict(list) + for start, end in self.builder.edges: + if start != START and end != END: + source_to_target[start].append(end) + for k in list(versions): + if k == START: + continue + if k in self.nodes: + v = versions.pop(k) + c = values.pop(k, MISSING) + for end in source_to_target[k]: + # get next version + new_k = f"branch:to:{end}" + new_v = max(versions[new_k], v) if new_k in versions else v + # update seen + for ss in (seen.get(end, {}), seen.get(INTERRUPT, {})): + if k in ss: + s = ss.pop(k) + if new_k in ss: + ss[new_k] = max(s, ss[new_k]) + else: + ss[new_k] = s + # update value + if new_k not in values and c is not MISSING: + values[new_k] = c + # update version + versions[new_k] = new_v + # pop interrupt seen + if INTERRUPT in seen: + seen[INTERRUPT].pop(k, MISSING) + def _get_state_reader( builder: StateGraph, schema: Type[Any] diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a5a500e01..8e2aa812f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -48,6 +48,7 @@ from langgraph.channels.base import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, + Checkpoint, CheckpointTuple, copy_checkpoint, ) @@ -766,6 +767,10 @@ class Pregel(PregelProtocol): for name, node in self.get_subgraphs(namespace=namespace, recurse=recurse): yield name, node + def _migrate_checkpoint(self, checkpoint: Checkpoint) -> None: + """Migrate a saved checkpoint to new channel layout.""" + pass + def _prepare_state_snapshot( self, config: RunnableConfig, @@ -784,6 +789,9 @@ class Pregel(PregelProtocol): tasks=(), ) + # migrate checkpoint if needed + self._migrate_checkpoint(saved.checkpoint) + with ChannelsManager( self.channels, saved.checkpoint, @@ -897,6 +905,9 @@ class Pregel(PregelProtocol): tasks=(), ) + # migrate checkpoint if needed + self._migrate_checkpoint(saved.checkpoint) + async with AsyncChannelsManager( self.channels, saved.checkpoint, @@ -1222,6 +1233,7 @@ class Pregel(PregelProtocol): # get last checkpoint config = ensure_config(self.config, input_config) saved = checkpointer.get_tuple(config) + self._migrate_checkpoint(saved.checkpoint) checkpoint = ( copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() ) @@ -1632,6 +1644,7 @@ class Pregel(PregelProtocol): # get last checkpoint config = ensure_config(self.config, input_config) saved = await checkpointer.aget_tuple(config) + self._migrate_checkpoint(saved.checkpoint) checkpoint = ( copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() ) @@ -2277,6 +2290,7 @@ class Pregel(PregelProtocol): manager=run_manager, debug=debug, trigger_to_nodes=self.trigger_to_nodes, + migrate_checkpoint=self._migrate_checkpoint, ) as loop: # create runner runner = PregelRunner( @@ -2570,12 +2584,8 @@ class Pregel(PregelProtocol): interrupt_after=interrupt_after_, manager=run_manager, debug=debug, - # `self.nodes` can be modified after creation of `Pregel`. For example, - # that's how StateGraph compilation currently works. - # For now, we recompute the trigger_to_nodes mapping every time the - # loop is created. We could potentially memoize this if it becomes a - # performance issue. - trigger_to_nodes=_trigger_to_nodes(self.nodes), + trigger_to_nodes=self.trigger_to_nodes, + migrate_checkpoint=self._migrate_checkpoint, ) as loop: # create runner runner = PregelRunner( diff --git a/libs/langgraph/langgraph/pregel/checkpoint.py b/libs/langgraph/langgraph/pregel/checkpoint.py index ee1d2bca8..b4d96091e 100644 --- a/libs/langgraph/langgraph/pregel/checkpoint.py +++ b/libs/langgraph/langgraph/pregel/checkpoint.py @@ -5,9 +5,8 @@ from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint from langgraph.checkpoint.base.id import uuid6 from langgraph.constants import MISSING -from langgraph.pregel.read import PregelNode -LATEST_VERSION = 2 +LATEST_VERSION = 3 def empty_checkpoint() -> Checkpoint: @@ -50,32 +49,3 @@ def create_checkpoint( versions_seen=checkpoint["versions_seen"], pending_sends=checkpoint.get("pending_sends", []), ) - - -def migrate_checkpoint( - checkpoint: Checkpoint, - channels: Mapping[str, BaseChannel], - nodes: Mapping[str, PregelNode], -) -> None: - """Migrate a checkpoint to new channel layout.""" - - values = checkpoint["channel_values"] - versions = checkpoint["channel_versions"] - seen = checkpoint["versions_seen"] - - if any(k.startswith("start:") for k in versions): - # Migrate from start:node to branch:to:node - for k in values: - if k.startswith("start:"): - node = k.split(":")[1] - new_k = f"branch:to:{node}" - if node not in nodes: - continue - v = versions.pop(k) - s = seen.get(node, {}).pop(k, None) - if s is None or s < v: - values[new_k] = values.pop(k) - # TODO handle s == v - # TODO handle new_k already in values - - # TODO Migrate from "node" to "branch:to:node" diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index cacdafe06..4c4ed6dde 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -30,6 +30,7 @@ from typing_extensions import ParamSpec, Self from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import ( + EXCLUDED_METADATA_KEYS, WRITES_IDX_MAP, BaseCheckpointSaver, ChannelVersions, @@ -174,6 +175,7 @@ class PregelLoop(LoopProtocol): Any, ] ] + _migrate_checkpoint: Optional[Callable[[Checkpoint], None]] submit: Submit channels: Mapping[str, BaseChannel] managed: ManagedValueMapping @@ -211,6 +213,7 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] = None, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, + migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, checkpoint_every_step: bool = True, ) -> None: @@ -236,6 +239,7 @@ class PregelLoop(LoopProtocol): CONFIG_KEY_CHECKPOINT_ID not in config[CONF] or CONFIG_KEY_DEDUPE_TASKS in config[CONF] ) + self._migrate_checkpoint = migrate_checkpoint self.trigger_to_nodes = trigger_to_nodes self.checkpoint_every_step = checkpoint_every_step self.debug = debug @@ -723,6 +727,8 @@ class PregelLoop(LoopProtocol): # bail if no checkpointer if self._checkpointer_put_after_previous is not None: for k, v in self.config["metadata"].items(): + if k in EXCLUDED_METADATA_KEYS: + continue metadata.setdefault(k, v) # type: ignore # create new checkpoint @@ -899,6 +905,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, + migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( @@ -916,6 +923,7 @@ class SyncPregelLoop(PregelLoop, ContextManager): interrupt_before=interrupt_before, manager=manager, debug=debug, + migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, ) self.stack = ExitStack() @@ -984,6 +992,8 @@ class SyncPregelLoop(PregelLoop, ContextManager): saved = CheckpointTuple( self.config, empty_checkpoint(), {"step": -2}, None, [] ) + elif self._migrate_checkpoint is not None: + self._migrate_checkpoint(saved.checkpoint) self.checkpoint_config = { **self.config, **saved.config, @@ -1042,6 +1052,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): stream_keys: Union[str, Sequence[str]] = EMPTY_SEQ, input_model: Optional[Type[BaseModel]] = None, debug: bool = False, + migrate_checkpoint: Optional[Callable[[Checkpoint], None]] = None, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, ) -> None: super().__init__( @@ -1059,6 +1070,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): interrupt_before=interrupt_before, manager=manager, debug=debug, + migrate_checkpoint=migrate_checkpoint, trigger_to_nodes=trigger_to_nodes, ) self.stack = AsyncExitStack() @@ -1127,6 +1139,8 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): saved = CheckpointTuple( self.config, empty_checkpoint(), {"step": -2}, None, [] ) + elif self._migrate_checkpoint is not None: + self._migrate_checkpoint(saved.checkpoint) self.checkpoint_config = { **self.config, **saved.config, diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index 5643a00fb..7f63ea801 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -4,6 +4,11 @@ from typing import Any, Sequence, Union from typing_extensions import Self +class AnyObject: + def __eq__(self, value): + return True + + class FloatBetween(float): def __new__(cls, min_value: float, max_value: float) -> Self: return super().__new__(cls, min_value) diff --git a/libs/langgraph/tests/test_checkpoint_migration.py b/libs/langgraph/tests/test_checkpoint_migration.py new file mode 100644 index 000000000..32423aabf --- /dev/null +++ b/libs/langgraph/tests/test_checkpoint_migration.py @@ -0,0 +1,1374 @@ +import operator +import time +from collections import defaultdict +from typing import Annotated, Literal, Optional, Union + +import pytest +from typing_extensions import TypedDict + +from langgraph.checkpoint.base import ( + BaseCheckpointSaver, + CheckpointTuple, + copy_checkpoint, +) +from langgraph.graph.state import StateGraph +from langgraph.types import Command, Interrupt, PregelTask, StateSnapshot, interrupt +from langgraph.utils.config import patch_configurable +from tests.any_str import AnyDict, AnyObject, AnyStr +from tests.conftest import ( + REGULAR_CHECKPOINTERS_ASYNC, + REGULAR_CHECKPOINTERS_SYNC, + awith_checkpointer, +) + +pytestmark = pytest.mark.anyio + + +def get_expected_history(*, exc_task_results: bool = False) -> list[StateSnapshot]: + return [ + StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "answer": "doc1,doc2,doc3,doc4", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + "step": 4, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(), + ), + StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + next=("qa",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"retriever_one": {"docs": ["doc1", "doc2"]}}, + "step": 3, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="qa", + path=("__pregel_pull", "qa"), + error=None, + interrupts=() + if exc_task_results + else ( + Interrupt( + value="", + resumable=True, + ns=[AnyStr("qa:")], + ), + ), + state=None, + result=None + if exc_task_results + else {"answer": "doc1,doc2,doc3,doc4"}, + ), + ), + ), + StateSnapshot( + values={ + "query": "analyzed: query: what is weather in sf", + "docs": ["doc3", "doc4"], + }, + next=("retriever_one",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + "step": 2, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="retriever_one", + path=("__pregel_pull", "retriever_one"), + error=None, + interrupts=(), + state=None, + result=None if exc_task_results else {"docs": ["doc1", "doc2"]}, + ), + ), + ), + StateSnapshot( + values={"query": "query: what is weather in sf", "docs": []}, + next=("analyzer_one", "retriever_two"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"rewrite_query": {"query": "query: what is weather in sf"}}, + "step": 1, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="analyzer_one", + path=("__pregel_pull", "analyzer_one"), + error=None, + interrupts=(), + state=None, + result=None + if exc_task_results + else {"query": "analyzed: query: what is weather in sf"}, + ), + PregelTask( + id=AnyStr(), + name="retriever_two", + path=("__pregel_pull", "retriever_two"), + error=None, + interrupts=(), + state=None, + result={"docs": ["doc3", "doc4"]}, + ), + ), + ), + StateSnapshot( + values={"query": "what is weather in sf", "docs": []}, + next=("rewrite_query",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=( + PregelTask( + id=AnyStr(), + name="rewrite_query", + path=("__pregel_pull", "rewrite_query"), + error=None, + interrupts=(), + state=None, + result=None + if exc_task_results + else {"query": "query: what is weather in sf"}, + ), + ), + ), + StateSnapshot( + values={"docs": []}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"__start__": {"query": "what is weather in sf"}}, + "step": -1, + "parents": {}, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + tasks=( + PregelTask( + id=AnyStr(), + name="__start__", + path=("__pregel_pull", "__start__"), + error=None, + interrupts=(), + state=None, + result={"query": "what is weather in sf"}, + ), + ), + ), + ] + + +SAVED_CHECKPOINTS = { + "3": [ + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2149-6faa-8004-9d848038f10a", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.237381+00:00", + "id": "1f00fd5f-2149-6faa-8004-9d848038f10a", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000004.0.18727156933289513", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:retriever_one": "00000000000000000000000000000005.0.7710812646219019", + "docs": "00000000000000000000000000000005.0.7916507770116351", + "branch:to:qa": "00000000000000000000000000000006.0.6375257096095945", + "answer": "00000000000000000000000000000006.0.9100669543952636", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + "analyzer_one": { + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583" + }, + "retriever_one": { + "branch:to:retriever_one": "00000000000000000000000000000004.0.5907938097782264" + }, + "__interrupt__": { + "query": "00000000000000000000000000000004.0.18727156933289513", + "docs": "00000000000000000000000000000005.0.7916507770116351", + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_one": "00000000000000000000000000000005.0.7710812646219019", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:qa": "00000000000000000000000000000005.0.5602643794940962", + }, + "qa": { + "branch:to:qa": "00000000000000000000000000000005.0.5602643794940962" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + "step": 4, + "parents": {}, + "thread_id": "1", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2140-6fd6-8003-2051ce36b79c", + } + }, + pending_writes=[], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2140-6fd6-8003-2051ce36b79c", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.233695+00:00", + "id": "1f00fd5f-2140-6fd6-8003-2051ce36b79c", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000004.0.18727156933289513", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:retriever_one": "00000000000000000000000000000005.0.7710812646219019", + "docs": "00000000000000000000000000000005.0.7916507770116351", + "branch:to:qa": "00000000000000000000000000000005.0.5602643794940962", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + "analyzer_one": { + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583" + }, + "retriever_one": { + "branch:to:retriever_one": "00000000000000000000000000000004.0.5907938097782264" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "branch:to:qa": None, + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": {"retriever_one": {"docs": ["doc1", "doc2"]}}, + "step": 3, + "parents": {}, + "thread_id": "1", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-213c-6940-8002-a28f475a6478", + } + }, + pending_writes=[ + ( + "2430f303-da9f-2e3e-738c-2e8ea28e8973", + "__interrupt__", + [ + Interrupt( + value="", + resumable=True, + ns=["qa:2430f303-da9f-2e3e-738c-2e8ea28e8973"], + ) + ], + ), + ("00000000-0000-0000-0000-000000000000", "__resume__", ""), + ("2430f303-da9f-2e3e-738c-2e8ea28e8973", "__resume__", [""]), + ( + "2430f303-da9f-2e3e-738c-2e8ea28e8973", + "answer", + "doc1,doc2,doc3,doc4", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-213c-6940-8002-a28f475a6478", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.231890+00:00", + "id": "1f00fd5f-213c-6940-8002-a28f475a6478", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000004.0.18727156933289513", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000004.0.15766851053750708", + "branch:to:retriever_two": "00000000000000000000000000000004.0.04821745244115927", + "branch:to:retriever_one": "00000000000000000000000000000004.0.5907938097782264", + "docs": "00000000000000000000000000000004.0.972701399851098", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + "analyzer_one": { + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "branch:to:retriever_one": None, + "docs": ["doc3", "doc4"], + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + "step": 2, + "parents": {}, + "thread_id": "1", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2039-6354-8001-2c508c8dffd9", + } + }, + pending_writes=[ + ("a5602426-85f2-1fe4-c9e4-bd0127e8e53e", "docs", ["doc1", "doc2"]), + ("a5602426-85f2-1fe4-c9e4-bd0127e8e53e", "branch:to:qa", None), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2039-6354-8001-2c508c8dffd9", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.125661+00:00", + "id": "1f00fd5f-2039-6354-8001-2c508c8dffd9", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000003.0.04057405566428263", + "branch:to:rewrite_query": "00000000000000000000000000000003.0.14126716107927562", + "branch:to:analyzer_one": "00000000000000000000000000000003.0.7165779439892241", + "branch:to:retriever_two": "00000000000000000000000000000003.0.7762711252277583", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "rewrite_query": { + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252" + }, + }, + "channel_values": { + "query": "query: what is weather in sf", + "branch:to:analyzer_one": None, + "branch:to:retriever_two": None, + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": {"rewrite_query": {"query": "query: what is weather in sf"}}, + "step": 1, + "parents": {}, + "thread_id": "1", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2038-613e-8000-ce5ebe65eb97", + } + }, + pending_writes=[ + ( + "4e7cb70b-7e0f-52d0-d8aa-5439bd3f84de", + "query", + "analyzed: query: what is weather in sf", + ), + ( + "4e7cb70b-7e0f-52d0-d8aa-5439bd3f84de", + "branch:to:retriever_one", + None, + ), + ("abcbc448-cfba-ac2b-2e39-346808f20add", "docs", ["doc3", "doc4"]), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2038-613e-8000-ce5ebe65eb97", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.125200+00:00", + "id": "1f00fd5f-2038-613e-8000-ce5ebe65eb97", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.6697367414225304", + "query": "00000000000000000000000000000002.0.3399249312096154", + "branch:to:rewrite_query": "00000000000000000000000000000002.0.05597832024496252", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + }, + "channel_values": { + "query": "what is weather in sf", + "branch:to:rewrite_query": None, + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": None, + "step": 0, + "parents": {}, + "thread_id": "1", + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2036-6ce4-bfff-ac42e9890362", + } + }, + pending_writes=[ + ( + "d1c3a2d6-5ca2-d4c5-5217-35d86cce48a4", + "query", + "query: what is weather in sf", + ), + ( + "d1c3a2d6-5ca2-d4c5-5217-35d86cce48a4", + "branch:to:analyzer_one", + None, + ), + ( + "d1c3a2d6-5ca2-d4c5-5217-35d86cce48a4", + "branch:to:retriever_two", + None, + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fd5f-2036-6ce4-bfff-ac42e9890362", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T15:20:01.124678+00:00", + "id": "1f00fd5f-2036-6ce4-bfff-ac42e9890362", + "channel_versions": { + "__start__": "00000000000000000000000000000001.0.7234984738744598" + }, + "versions_seen": {"__input__": {}}, + "channel_values": {"__start__": {"query": "what is weather in sf"}}, + "pending_sends": [], + }, + metadata={ + "source": "input", + "writes": {"__start__": {"query": "what is weather in sf"}}, + "step": -1, + "parents": {}, + "thread_id": "1", + }, + parent_config=None, + pending_writes=[ + ( + "a9e2a749-9870-1952-0a6c-b23b6729ffda", + "query", + "what is weather in sf", + ), + ( + "a9e2a749-9870-1952-0a6c-b23b6729ffda", + "branch:to:rewrite_query", + None, + ), + ], + ), + ], + "2-start:*": [ + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515f-6b88-8004-6fffb69dd465", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.825576+00:00", + "id": "1f00fe48-515f-6b88-8004-6fffb69dd465", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000004.0.05732679770452498", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "analyzer_one": "00000000000000000000000000000005.0.648286705356163", + "docs": "00000000000000000000000000000005.0.19918575623485935", + "retriever_two": "00000000000000000000000000000005.0.46629341414062697", + "retriever_one": "00000000000000000000000000000006.0.9577453764095437", + "answer": "00000000000000000000000000000006.0.27361287406148327", + "qa": "00000000000000000000000000000006.0.24260043089701677", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884" + }, + "retriever_one": { + "analyzer_one": "00000000000000000000000000000004.0.9322215406936268" + }, + "__interrupt__": { + "query": "00000000000000000000000000000004.0.05732679770452498", + "docs": "00000000000000000000000000000005.0.19918575623485935", + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "analyzer_one": "00000000000000000000000000000005.0.648286705356163", + "retriever_one": "00000000000000000000000000000005.0.0523757506060204", + "retriever_two": "00000000000000000000000000000005.0.46629341414062697", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + }, + "qa": { + "retriever_one": "00000000000000000000000000000005.0.0523757506060204" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + "qa": "qa", + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + "thread_id": "1", + "step": 4, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515c-679e-8003-5f85a56d5dba", + } + }, + pending_writes=[], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515c-679e-8003-5f85a56d5dba", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.824251+00:00", + "id": "1f00fe48-515c-679e-8003-5f85a56d5dba", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000004.0.05732679770452498", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "analyzer_one": "00000000000000000000000000000005.0.648286705356163", + "docs": "00000000000000000000000000000005.0.19918575623485935", + "retriever_two": "00000000000000000000000000000005.0.46629341414062697", + "retriever_one": "00000000000000000000000000000005.0.0523757506060204", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884" + }, + "retriever_one": { + "analyzer_one": "00000000000000000000000000000004.0.9322215406936268" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "retriever_one": "retriever_one", + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": {"retriever_one": {"docs": ["doc1", "doc2"]}}, + "thread_id": "1", + "step": 3, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515b-6d12-8002-82a8f9213eae", + } + }, + pending_writes=[ + ( + "4ee8637e-0a95-285e-75bc-4da721c0beab", + "__interrupt__", + [ + Interrupt( + value="", + resumable=True, + ns=["qa:4ee8637e-0a95-285e-75bc-4da721c0beab"], + ) + ], + ), + ("00000000-0000-0000-0000-000000000000", "__resume__", ""), + ("4ee8637e-0a95-285e-75bc-4da721c0beab", "__resume__", [""]), + ( + "4ee8637e-0a95-285e-75bc-4da721c0beab", + "answer", + "doc1,doc2,doc3,doc4", + ), + ("4ee8637e-0a95-285e-75bc-4da721c0beab", "qa", "qa"), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-515b-6d12-8002-82a8f9213eae", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.823978+00:00", + "id": "1f00fe48-515b-6d12-8002-82a8f9213eae", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000004.0.05732679770452498", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000004.0.2372002638794427", + "branch:to:retriever_two": "00000000000000000000000000000004.0.8860781568140047", + "analyzer_one": "00000000000000000000000000000004.0.9322215406936268", + "docs": "00000000000000000000000000000004.0.49012772235571145", + "retriever_two": "00000000000000000000000000000004.0.9223450775254257", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + "analyzer_one": { + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955" + }, + "retriever_two": { + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884" + }, + }, + "channel_values": { + "query": "analyzed: query: what is weather in sf", + "analyzer_one": "analyzer_one", + "docs": ["doc3", "doc4"], + "retriever_two": "retriever_two", + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": { + "analyzer_one": {"query": "analyzed: query: what is weather in sf"}, + "retriever_two": {"docs": ["doc3", "doc4"]}, + }, + "thread_id": "1", + "step": 2, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5059-6b30-8001-2a9ab4ca7d82", + } + }, + pending_writes=[ + ("16295c56-f44e-31fa-8fad-fff3f9022629", "docs", ["doc1", "doc2"]), + ( + "16295c56-f44e-31fa-8fad-fff3f9022629", + "retriever_one", + "retriever_one", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5059-6b30-8001-2a9ab4ca7d82", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.718258+00:00", + "id": "1f00fe48-5059-6b30-8001-2a9ab4ca7d82", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000003.0.10748450241039154", + "start:rewrite_query": "00000000000000000000000000000003.0.2916637829964738", + "rewrite_query": "00000000000000000000000000000003.0.9534854313752955", + "branch:to:retriever_two": "00000000000000000000000000000003.0.29217346538810884", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "rewrite_query": { + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763" + }, + }, + "channel_values": { + "query": "query: what is weather in sf", + "rewrite_query": "rewrite_query", + "branch:to:retriever_two": "rewrite_query", + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": {"rewrite_query": {"query": "query: what is weather in sf"}}, + "thread_id": "1", + "step": 1, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5058-6a46-8000-086bffc73797", + } + }, + pending_writes=[ + ( + "baecc0e3-ea00-0e00-9436-e33cd2527faf", + "query", + "analyzed: query: what is weather in sf", + ), + ( + "baecc0e3-ea00-0e00-9436-e33cd2527faf", + "analyzer_one", + "analyzer_one", + ), + ("96b7bfe4-269f-092c-e685-14dba6a27271", "docs", ["doc3", "doc4"]), + ( + "96b7bfe4-269f-092c-e685-14dba6a27271", + "retriever_two", + "retriever_two", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5058-6a46-8000-086bffc73797", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.717827+00:00", + "id": "1f00fe48-5058-6a46-8000-086bffc73797", + "channel_versions": { + "__start__": "00000000000000000000000000000002.0.23383372151016169", + "query": "00000000000000000000000000000002.0.706632616485588", + "start:rewrite_query": "00000000000000000000000000000002.0.3082066433110763", + }, + "versions_seen": { + "__input__": {}, + "__start__": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + }, + "channel_values": { + "query": "what is weather in sf", + "start:rewrite_query": "__start__", + }, + "pending_sends": [], + }, + metadata={ + "source": "loop", + "writes": None, + "thread_id": "1", + "step": 0, + "parents": {}, + }, + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5057-62a4-bfff-1883a92a3e41", + } + }, + pending_writes=[ + ( + "058cf6d6-a83c-6509-b398-5dde0b6c5773", + "query", + "query: what is weather in sf", + ), + ( + "058cf6d6-a83c-6509-b398-5dde0b6c5773", + "rewrite_query", + "rewrite_query", + ), + ( + "058cf6d6-a83c-6509-b398-5dde0b6c5773", + "branch:to:retriever_two", + "rewrite_query", + ), + ], + ), + CheckpointTuple( + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": "1f00fe48-5057-62a4-bfff-1883a92a3e41", + } + }, + checkpoint={ + "v": 2, + "ts": "2025-04-02T17:04:20.717221+00:00", + "id": "1f00fe48-5057-62a4-bfff-1883a92a3e41", + "channel_versions": { + "__start__": "00000000000000000000000000000001.0.9575279209966122" + }, + "versions_seen": {"__input__": {}}, + "channel_values": {"__start__": {"query": "what is weather in sf"}}, + "pending_sends": [], + }, + metadata={ + "source": "input", + "writes": {"__start__": {"query": "what is weather in sf"}}, + "thread_id": "1", + "step": -1, + "parents": {}, + }, + parent_config=None, + pending_writes=[ + ( + "891e8564-d78f-7fb2-f15d-bce2a0ddf1c6", + "query", + "what is weather in sf", + ), + ( + "891e8564-d78f-7fb2-f15d-bce2a0ddf1c6", + "start:rewrite_query", + "__start__", + ), + ], + ), + ], +} + + +def make_state_graph() -> StateGraph: + def sorted_add( + x: list[str], y: Union[list[str], list[tuple[str, str]]] + ) -> list[str]: + if isinstance(y[0], tuple): + for rem, _ in y: + x.remove(rem) + y = [t[1] for t in y] + return sorted(operator.add(x, y)) + + class State(TypedDict, total=False): + query: str + answer: str + docs: Annotated[list[str], sorted_add] + + def rewrite_query(data: State) -> State: + return {"query": f"query: {data['query']}"} + + def analyzer_one(data: State) -> State: + return {"query": f"analyzed: {data['query']}"} + + def retriever_one(data: State) -> State: + return {"docs": ["doc1", "doc2"]} + + def retriever_two(data: State) -> State: + time.sleep(0.1) + return {"docs": ["doc3", "doc4"]} + + def qa(data: State) -> State: + interrupt("") + return {"answer": ",".join(data["docs"])} + + def rewrite_query_then(data: State) -> Literal["retriever_two"]: + return "retriever_two" + + workflow = StateGraph(State) + + workflow.add_node("rewrite_query", rewrite_query) + workflow.add_node("analyzer_one", analyzer_one) + workflow.add_node("retriever_one", retriever_one) + workflow.add_node("retriever_two", retriever_two) + workflow.add_node("qa", qa) + + workflow.set_entry_point("rewrite_query") + workflow.add_edge("rewrite_query", "analyzer_one") + workflow.add_edge("analyzer_one", "retriever_one") + workflow.add_conditional_edges("rewrite_query", rewrite_query_then) + workflow.add_edge("retriever_one", "qa") + workflow.set_finish_point("qa") + return workflow + + +def test_migrate_checkpoints() -> None: + # Check that the migration function works as expected + builder = make_state_graph() + graph = builder.compile() + + source = list(reversed(SAVED_CHECKPOINTS["2-start:*"])) + target = list(reversed(SAVED_CHECKPOINTS["3"])) + assert len(source) == len(target) + for idx, (source_checkpoint, target_checkpoint) in enumerate(zip(source, target)): + # copy the checkpoint to avoid modifying the original + migrated = copy_checkpoint(source_checkpoint.checkpoint) + # migrate the checkpoint + graph._migrate_checkpoint(migrated) + # replace values that don't need to match exactly + migrated["id"] = AnyStr() + migrated["ts"] = AnyStr() + for k in migrated["channel_values"]: + migrated["channel_values"][k] = AnyObject() + for v in migrated["channel_versions"]: + migrated["channel_versions"][v] = AnyStr( + migrated["channel_versions"][v].split(".")[0] + ) + for c in migrated["versions_seen"]: + for v in migrated["versions_seen"][c]: + migrated["versions_seen"][c][v] = AnyStr( + migrated["versions_seen"][c][v].split(".")[0] + ) + # check that the migrated checkpoint matches the target checkpoint + assert ( + migrated == target_checkpoint.checkpoint + ), "Checkpoint mismatch at index {}".format(idx) + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_latest_checkpoint_state_graph( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + builder = make_state_graph() + app = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert [*app.stream({"query": "what is weather in sf"}, config)] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + { + "__interrupt__": ( + Interrupt( + value="", + resumable=True, + ns=[AnyStr("qa:")], + ), + ) + }, + ] + + assert [*app.stream(Command(resume=""), config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + # check history with current checkpoints matches expected history + history = [*app.get_state_history(config)] + expected_history = get_expected_history() + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_latest_checkpoint_state_graph_async(checkpointer_name: str) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + builder = make_state_graph() + app = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert [ + c async for c in app.astream({"query": "what is weather in sf"}, config) + ] == [ + {"rewrite_query": {"query": "query: what is weather in sf"}}, + {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}}, + {"retriever_two": {"docs": ["doc3", "doc4"]}}, + {"retriever_one": {"docs": ["doc1", "doc2"]}}, + { + "__interrupt__": ( + Interrupt( + value="", + resumable=True, + ns=[AnyStr("qa:")], + ), + ) + }, + ] + + assert [c async for c in app.astream(Command(resume=""), config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + + # check history with current checkpoints matches expected history + history = [c async for c in app.aget_state_history(config)] + expected_history = get_expected_history() + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + +@pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*"]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_SYNC) +def test_saved_checkpoint_state_graph( + request: pytest.FixtureRequest, + checkpointer_name: str, + checkpoint_version: str, +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + builder = make_state_graph() + app = builder.compile(checkpointer=checkpointer) + + thread1 = "1" + config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} + + # save checkpoints + parent_id: Optional[str] = None + for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]): + grouped_writes = defaultdict(list) + for write in checkpoint.pending_writes: + grouped_writes[write[0]].append(write[1:]) + for tid, group in grouped_writes.items(): + checkpointer.put_writes(checkpoint.config, group, tid) + checkpointer.put( + patch_configurable(config, {"checkpoint_id": parent_id}), + checkpoint.checkpoint, + checkpoint.metadata, + checkpoint.checkpoint["channel_versions"], + ) + parent_id = checkpoint.checkpoint["id"] + + # load history + history = [*app.get_state_history(config)] + # check history with saved checkpoints matches expected history + expected_history = get_expected_history(exc_task_results=checkpoint_version != "3") + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + # resume from 2nd to latest checkpoint + assert [*app.stream(Command(resume=""), history[1].config)] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + # new checkpoint should match the latest checkpoint in history + latest_state = app.get_state(config) + assert ( + StateSnapshot( + values=latest_state.values, + next=latest_state.next, + config=patch_configurable(latest_state.config, {"checkpoint_id": AnyStr()}), + metadata=AnyDict(latest_state.metadata), + created_at=AnyStr(), + parent_config=latest_state.parent_config, + tasks=latest_state.tasks, + ) + == history[0] + ) + + +@pytest.mark.parametrize("checkpoint_version", ["3", "2-start:*"]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_saved_checkpoint_state_graph_async( + checkpointer_name: str, + checkpoint_version: str, +) -> None: + async with awith_checkpointer(checkpointer_name) as checkpointer: + builder = make_state_graph() + app = builder.compile(checkpointer=checkpointer) + + thread1 = "1" + config = {"configurable": {"thread_id": thread1, "checkpoint_ns": ""}} + + # save checkpoints + parent_id: Optional[str] = None + for checkpoint in reversed(SAVED_CHECKPOINTS[checkpoint_version]): + grouped_writes = defaultdict(list) + for write in checkpoint.pending_writes: + grouped_writes[write[0]].append(write[1:]) + for tid, group in grouped_writes.items(): + await checkpointer.aput_writes(checkpoint.config, group, tid) + await checkpointer.aput( + patch_configurable(config, {"checkpoint_id": parent_id}), + checkpoint.checkpoint, + checkpoint.metadata, + checkpoint.checkpoint["channel_versions"], + ) + parent_id = checkpoint.checkpoint["id"] + + # load history + history = [c async for c in app.aget_state_history(config)] + # check history with saved checkpoints matches expected history + expected_history = get_expected_history( + exc_task_results=checkpoint_version != "3" + ) + assert len(history) == len(expected_history) + assert history[0] == expected_history[0] + assert history[1] == expected_history[1] + assert history[2] == expected_history[2] + assert history[3] == expected_history[3] + assert history[4] == expected_history[4] + assert history[5] == expected_history[5] + + # resume from 2nd to latest checkpoint + assert [ + c async for c in app.astream(Command(resume=""), history[1].config) + ] == [ + {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, + ] + # new checkpoint should match the latest checkpoint in history + latest_state = await app.aget_state(config) + assert ( + StateSnapshot( + values=latest_state.values, + next=latest_state.next, + config=patch_configurable( + latest_state.config, {"checkpoint_id": AnyStr()} + ), + metadata=AnyDict(latest_state.metadata), + created_at=AnyStr(), + parent_config=latest_state.parent_config, + tasks=latest_state.tasks, + ) + == history[0] + )