diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index aea9069b9..25575e481 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -364,8 +364,12 @@ class MemorySaver( checkpoint_ns = config["configurable"]["checkpoint_ns"] checkpoint_id = config["configurable"]["checkpoint_id"] outer_key = (thread_id, checkpoint_ns, checkpoint_id) + outer_writes_ = self.writes.get(outer_key) for idx, (c, v) in enumerate(writes): inner_key = (task_id, WRITES_IDX_MAP.get(c, idx)) + if inner_key[1] >= 0 and outer_writes_ and inner_key in outer_writes_: + continue + self.writes[outer_key][inner_key] = (task_id, c, self.serde.dumps_typed(v)) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index 491624228..09fba4e25 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint" -version = "2.0.2" +version = "2.0.3" description = "Library with base interfaces for LangGraph checkpoint savers." authors = [] license = "MIT" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index 05c4d1836..00c42fce5 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph" -version = "0.2.45" +version = "0.2.46" description = "Building stateful, multi-actor applications with LLMs" authors = [] license = "MIT" diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 20b12beb2..4026cecd6 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1115,6 +1115,86 @@ def test_fork_always_re_runs_nodes( ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_run_from_checkpoint_id_retains_previous_writes( + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture +) -> None: + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + + class MyState(TypedDict): + myval: Annotated[int, operator.add] + otherval: bool + + class Anode: + def __init__(self): + self.switch = False + + def __call__(self, state: MyState): + self.switch = not self.switch + return {"myval": 2 if self.switch else 1, "otherval": self.switch} + + builder = StateGraph(MyState) + thenode = Anode() # Fun. + builder.add_node("node_one", thenode) + builder.add_node("node_two", thenode) + builder.add_edge(START, "node_one") + + def _getedge(src: str): + swap = "node_one" if src == "node_two" else "node_two" + + def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + if st["myval"] > 3: + return END + if st["otherval"]: + return swap + return src + + return _edge + + builder.add_conditional_edges("node_one", _getedge("node_one")) + builder.add_conditional_edges("node_two", _getedge("node_two")) + graph = builder.compile(checkpointer=checkpointer) + + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} + + result = graph.invoke({"myval": 1}, thread1) + assert result["myval"] == 4 + history = [c for c in graph.get_state_history(thread1)] + + assert len(history) == 4 + assert history[-1].values == {"myval": 0} + assert history[0].values == {"myval": 4, "otherval": False} + + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = graph.invoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} + + new_history = [ + c + for c in graph.get_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] + + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] + + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] + + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + + def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 66d3c2b8f..e20a76cbe 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1960,6 +1960,85 @@ async def test_pending_writes_resume( ) +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_run_from_checkpoint_id_retains_previous_writes( + request: pytest.FixtureRequest, checkpointer_name: str, mocker: MockerFixture +) -> None: + class MyState(TypedDict): + myval: Annotated[int, operator.add] + otherval: bool + + class Anode: + def __init__(self): + self.switch = False + + async def __call__(self, state: MyState): + self.switch = not self.switch + return {"myval": 2 if self.switch else 1, "otherval": self.switch} + + builder = StateGraph(MyState) + thenode = Anode() # Fun. + builder.add_node("node_one", thenode) + builder.add_node("node_two", thenode) + builder.add_edge(START, "node_one") + + def _getedge(src: str): + swap = "node_one" if src == "node_two" else "node_two" + + def _edge(st: MyState) -> Literal["__end__", "node_one", "node_two"]: + if st["myval"] > 3: + return END + if st["otherval"]: + return swap + return src + + return _edge + + builder.add_conditional_edges("node_one", _getedge("node_one")) + builder.add_conditional_edges("node_two", _getedge("node_two")) + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + + thread_id = uuid.uuid4() + thread1 = {"configurable": {"thread_id": str(thread_id)}} + + result = await graph.ainvoke({"myval": 1}, thread1) + assert result["myval"] == 4 + history = [c async for c in graph.aget_state_history(thread1)] + + assert len(history) == 4 + assert history[-1].values == {"myval": 0} + assert history[0].values == {"myval": 4, "otherval": False} + + second_run_config = { + **thread1, + "configurable": { + **thread1["configurable"], + "checkpoint_id": history[1].config["configurable"]["checkpoint_id"], + }, + } + second_result = await graph.ainvoke(None, second_run_config) + assert second_result == {"myval": 5, "otherval": True} + + new_history = [ + c + async for c in graph.aget_state_history( + {"configurable": {"thread_id": str(thread_id), "checkpoint_ns": ""}} + ) + ] + + assert len(new_history) == len(history) + 1 + for original, new in zip(history, new_history[1:]): + assert original.values == new.values + assert original.next == new.next + assert original.metadata["step"] == new.metadata["step"] + + def _get_tasks(hist: list, start: int): + return [h.tasks for h in hist[start:]] + + assert _get_tasks(new_history, 1) == _get_tasks(history, 0) + + async def test_cond_edge_after_send() -> None: class Node: def __init__(self, name: str):