diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 2b7bce5d4..08feeabfa 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -66,7 +66,7 @@ class JsonPlusSerializer(SerializerProtocol): return self._encode_constructor_args(obj.__class__, args=[obj.value]) elif isinstance(obj, SendProtocol): return self._encode_constructor_args( - obj.__class__, kwargs={"node": obj.node, "arg": obj.arg} + obj.__class__, kwargs={"node": obj.node, "arg": obj.arg, "id": obj.id} ) elif isinstance(obj, (bytes, bytearray)): return self._encode_constructor_args( diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index de61ef78b..cc5c1fa8b 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -55,6 +55,7 @@ class SendProtocol(Protocol): # Mirrors langgraph.constants.Send node: str arg: Any + id: str def __hash__(self) -> int: ... diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index ae447d62e..89a2effa7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,4 +1,5 @@ -from typing import Any +from typing import Any, Optional +from uuid import uuid4 INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" @@ -22,6 +23,7 @@ START = "__start__" END = "__end__" CHECKPOINT_NAMESPACE_SEPARATOR = "|" +SEND_CHECKPOINT_NAMESPACE_SEPARATOR = ":" class Send: @@ -40,6 +42,7 @@ class Send: Attributes: node (str): The name of the target node to send the message to. arg (Any): The state or message to send to the target node. + id (str): ID associated with the Send. Examples: >>> from typing import Annotated @@ -67,23 +70,26 @@ class Send: node: str arg: Any + id: Optional[str] - def __init__(self, /, node: str, arg: Any) -> None: + def __init__(self, /, node: str, arg: Any, id: Optional[str] = None) -> None: """ Initialize a new instance of the Send class. Args: node (str): The name of the target node to send the message to. arg (Any): The state or message to send to the target node. + id (str): ID associated with the Send. """ self.node = node self.arg = arg + self.id = id or str(uuid4()) def __hash__(self) -> int: return hash((self.node, self.arg)) def __repr__(self) -> str: - return f"Send(node={self.node!r}, arg={self.arg!r})" + return f"Send(node={self.node!r}, arg={self.arg!r}, id={self.id!r})" def __eq__(self, value: object) -> bool: return ( diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index c17373820..6fc4eb8cd 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -27,6 +27,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.constants import ( CHECKPOINT_NAMESPACE_SEPARATOR, END, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, START, TAG_HIDDEN, Send, @@ -159,10 +160,15 @@ class Graph: *, metadata: Optional[dict[str, Any]] = None, ) -> None: - if isinstance(node, str) and CHECKPOINT_NAMESPACE_SEPARATOR in node: - raise ValueError( - f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names." - ) + if isinstance(node, str): + for character in ( + CHECKPOINT_NAMESPACE_SEPARATOR, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, + ): + if character in node: + raise ValueError( + f"'{character}' is a reserved character and is not allowed in the node names." + ) if self.compiled: logger.warning( diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index e9971a5f1..8e7298616 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -29,7 +29,11 @@ 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 CHECKPOINT_NAMESPACE_SEPARATOR, TAG_HIDDEN +from langgraph.constants import ( + CHECKPOINT_NAMESPACE_SEPARATOR, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, + TAG_HIDDEN, +) from langgraph.errors import InvalidUpdateError from langgraph.graph.graph import ( END, @@ -312,10 +316,14 @@ class StateGraph(Graph): if node == END or node == START: raise ValueError(f"Node `{node}` is reserved.") - if CHECKPOINT_NAMESPACE_SEPARATOR in node: - raise ValueError( - f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names." - ) + for character in ( + CHECKPOINT_NAMESPACE_SEPARATOR, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, + ): + if character in node: + raise ValueError( + f"'{character}' is a reserved character and is not allowed in the node names." + ) try: if isfunction(action) and ( diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 58e34dd59..96f720be5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -72,6 +72,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -370,6 +371,15 @@ class Pregel( nodes = self.nodes channels = self.channels for subgraph_node_name in path: + # if we have this separator it means we have a node that was triggered by Send + if SEND_CHECKPOINT_NAMESPACE_SEPARATOR in subgraph_node_name: + name_parts = subgraph_node_name.split( + SEND_CHECKPOINT_NAMESPACE_SEPARATOR + ) + if len(name_parts) != 2: + raise ValueError(f"Malformed node name '{subgraph_node_name}'") + + subgraph_node_name = name_parts[0] if subgraph_node_name not in nodes: raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") @@ -1190,8 +1200,7 @@ class Pregel( ) if not done: break # timed out - for fut in done: - task = futures.pop(fut) + for fut, task in zip(done, [futures.pop(fut) for fut in done]): if fut.exception() is not None: # we got an exception, break out of while loop # exception will be handled in panic_or_proceed @@ -1435,8 +1444,7 @@ class Pregel( ) if not done: break # timed out - for fut in done: - task = futures.pop(fut) + for fut, task in zip(done, [futures.pop(fut) for fut in done]): if fut.exception() is not None: # we got an exception, break out of while loop # exception will be handled in panic_or_proceed diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 2ddbcd873..ec8c18650 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -280,9 +280,9 @@ def prepare_next_tasks( "langgraph_task_idx": len(tasks), } checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}" + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}" if parent_ns - else packet.node + else f"{packet.node}:{packet.id}" ) task_id = str( uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) @@ -318,6 +318,10 @@ def prepare_next_tasks( PregelTaskWrites(packet.node, writes, triggers), config, ), + CONFIG_KEY_CHECKPOINTER: checkpointer, + CONFIG_KEY_RESUMING: is_resuming, + "checkpoint_id": checkpoint["id"], + "checkpoint_ns": checkpoint_ns, # in Send we can't checkpoint nested graphs # as they could be running in parallel }, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8780bc042..881487c93 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9807,6 +9807,366 @@ def test_doubly_nested_graph_state( ) +@pytest.mark.repeat(10) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_send_to_nested_graphs( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class OverallState(TypedDict): + subjects: list[str] + jokes: Annotated[list[str], operator.add] + + def continue_to_jokes(state: OverallState): + return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] + + class JokeState(TypedDict): + subject: str + + def edit(state: JokeState): + subject = state["subject"] + return {"subject": f"{subject} - hohoho"} + + # subgraph + subgraph = StateGraph(input=JokeState, output=OverallState) + subgraph.add_node("edit", edit) + subgraph.add_node( + "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} + ) + subgraph.set_entry_point("edit") + subgraph.add_edge("edit", "generate") + subgraph.set_finish_point("generate") + + # parent graph + builder = StateGraph(OverallState) + builder.add_node( + "generate_joke", + subgraph.compile( + checkpointer=INHERIT_CHECKPOINTER, interrupt_before=["generate"] + ), + ) + builder.add_conditional_edges(START, continue_to_jokes) + builder.add_edge("generate_joke", END) + + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # invoke and pause at nested interrupt + assert graph.invoke({"subjects": ["cats", "dogs"]}, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + actual_snapshot = graph.get_state(config, include_subgraph_state=True) + subgraph_nodes = list(actual_snapshot.subgraph_state_snapshots.keys()) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + + expected_snapshot = StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + assert actual_snapshot == expected_snapshot + + # continue past interrupt + assert graph.invoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + } + + actual_snapshot = graph.get_state(config, include_subgraph_state=True) + subgraph_nodes, _ = zip( + *( + sorted( + actual_snapshot.subgraph_state_snapshots.items(), + key=lambda x: x[1].values["jokes"][0], + ) + ) + ) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = list(graph.get_state_history(config, include_subgraph_state=True)) + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ), + StateSnapshot( + values={"jokes": []}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"subjects": ["cats", "dogs"]}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + assert actual_history == expected_history + + def test_repeat_condition(snapshot: SnapshotAssertion) -> None: class AgentState(TypedDict): hello: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index af624f639..43eab4ea6 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8318,6 +8318,368 @@ async def test_doubly_nested_graph_state( ) +@pytest.mark.repeat(10) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_send_to_nested_graphs( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class OverallState(TypedDict): + subjects: list[str] + jokes: Annotated[list[str], operator.add] + + async def continue_to_jokes(state: OverallState): + return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] + + class JokeState(TypedDict): + subject: str + + async def edit(state: JokeState): + subject = state["subject"] + return {"subject": f"{subject} - hohoho"} + + # subgraph + subgraph = StateGraph(input=JokeState, output=OverallState) + subgraph.add_node("edit", edit) + subgraph.add_node( + "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} + ) + subgraph.set_entry_point("edit") + subgraph.add_edge("edit", "generate") + subgraph.set_finish_point("generate") + + # parent graph + builder = StateGraph(OverallState) + builder.add_node( + "generate_joke", + subgraph.compile( + checkpointer=INHERIT_CHECKPOINTER, interrupt_before=["generate"] + ), + ) + builder.add_conditional_edges(START, continue_to_jokes) + builder.add_edge("generate_joke", END) + + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # invoke and pause at nested interrupt + assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + actual_snapshot = await graph.aget_state(config, include_subgraph_state=True) + subgraph_nodes = list(actual_snapshot.subgraph_state_snapshots.keys()) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + + expected_snapshot = StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + assert actual_snapshot == expected_snapshot + + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + } + + actual_snapshot = await graph.aget_state(config, include_subgraph_state=True) + subgraph_nodes, _ = zip( + *( + sorted( + actual_snapshot.subgraph_state_snapshots.items(), + key=lambda x: x[1].values["jokes"][0], + ) + ) + ) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = [ + c async for c in graph.aget_state_history(config, include_subgraph_state=True) + ] + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ), + StateSnapshot( + values={"jokes": []}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"subjects": ["cats", "dogs"]}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + assert actual_history == expected_history + + async def test_checkpoint_metadata() -> None: """This test verifies that a run's configurable fields are merged with the previous checkpoint config for each step in the run.