From d635a22302df57ad8df7fca2aa61c20ace982ea3 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 22 Jul 2024 12:14:33 -0400 Subject: [PATCH] sort memory checkpoints monotonically decreasing --- libs/langgraph/langgraph/checkpoint/memory.py | 4 +- libs/langgraph/tests/test_pregel.py | 140 +++++++++--------- libs/langgraph/tests/test_pregel_async.py | 140 +++++++++--------- 3 files changed, 143 insertions(+), 141 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index bd5dc6fd1..eec3fcdf7 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -119,7 +119,9 @@ class MemorySaver(BaseCheckpointSaver): """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage for thread_id in thread_ids: - for ts, (checkpoint, metadata_b) in self.storage[thread_id].items(): + for ts, (checkpoint, metadata_b) in sorted( + self.storage[thread_id].items(), key=lambda x: x[0], reverse=True + ): # filter by thread_ts if before and ts >= before["configurable"]["thread_ts"]: continue diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 89db73b2a..f6f1257aa 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -608,33 +608,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: # list history assert [c for c in app.get_state_history({"configurable": {"thread_id": 1}})] == [ StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": -1, "writes": 2}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, + values={"inbox": 4, "output": 5, "input": 3}, next=(), config={ "configurable": { @@ -642,47 +616,8 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 6, "writes": 5}, created_at=AnyStr(), - metadata={"source": "loop", "step": 1, "writes": 4}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 2, "writes": 20}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 3, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 4, "writes": 3}, parent_config=None, ), StateSnapshot( @@ -694,12 +629,51 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: "thread_ts": AnyStr(), } }, - created_at=AnyStr(), metadata={"source": "loop", "step": 5, "writes": None}, + created_at=AnyStr(), parent_config=None, ), StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, next=(), config={ "configurable": { @@ -707,8 +681,34 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, created_at=AnyStr(), - metadata={"source": "loop", "step": 6, "writes": 5}, parent_config=None, ), ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index a43d22f8a..7e60c77fb 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -727,33 +727,7 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N c async for c in app.aget_state_history({"configurable": {"thread_id": 1}}) ] == [ StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": -1, "writes": 2}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 0, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, + values={"inbox": 4, "output": 5, "input": 3}, next=(), config={ "configurable": { @@ -761,47 +735,8 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 6, "writes": 5}, created_at=AnyStr(), - metadata={"source": "loop", "step": 1, "writes": 4}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 2, "writes": 20}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "loop", "step": 3, "writes": None}, - parent_config=None, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - created_at=AnyStr(), - metadata={"source": "input", "step": 4, "writes": 3}, parent_config=None, ), StateSnapshot( @@ -813,12 +748,51 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N "thread_ts": AnyStr(), } }, - created_at=AnyStr(), metadata={"source": "loop", "step": 5, "writes": None}, + created_at=AnyStr(), parent_config=None, ), StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, next=(), config={ "configurable": { @@ -826,8 +800,34 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N "thread_ts": AnyStr(), } }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=None, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": 1, + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, created_at=AnyStr(), - metadata={"source": "loop", "step": 6, "writes": 5}, parent_config=None, ), ]