From 3daa4ae46654a2fcaf96bfe1b2dfc83e5cbdd581 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Wed, 5 Feb 2025 22:04:25 -0800 Subject: [PATCH] Update name to InMemorySaver (#2044) (Keep MemorySaver around for backwards compatibility) "MemorySaver" is ambiguous: is it saving memories? Where is it saving memories to? InMemorySaver aligns naming InMemoryStore as well as similar LangChain objects (InMemoryVectorStore, etc.) --- docs/docs/concepts/persistence.md | 2 +- .../langgraph/checkpoint/memory/__init__.py | 69 +++++++++++-------- libs/checkpoint/pyproject.toml | 2 +- libs/checkpoint/tests/test_memory.py | 10 ++- libs/langgraph/tests/memory_assert.py | 8 +-- libs/langgraph/tests/test_pregel.py | 18 ++--- libs/langgraph/tests/test_pregel_async.py | 16 ++--- 7 files changed, 70 insertions(+), 55 deletions(-) diff --git a/docs/docs/concepts/persistence.md b/docs/docs/concepts/persistence.md index 629a85897..7f4519e3f 100644 --- a/docs/docs/concepts/persistence.md +++ b/docs/docs/concepts/persistence.md @@ -433,7 +433,7 @@ See the [deployment guide](../cloud/deployment/semantic_search.md) for more deta Under the hood, checkpointing is powered by checkpointer objects that conform to [BaseCheckpointSaver][langgraph.checkpoint.base.BaseCheckpointSaver] interface. LangGraph provides several checkpointer implementations, all implemented via standalone, installable libraries: -* `langgraph-checkpoint`: The base interface for checkpointer savers ([BaseCheckpointSaver][langgraph.checkpoint.base.BaseCheckpointSaver]) and serialization/deserialization interface ([SerializerProtocol][langgraph.checkpoint.serde.base.SerializerProtocol]). Includes in-memory checkpointer implementation ([MemorySaver][langgraph.checkpoint.memory.MemorySaver]) for experimentation. LangGraph comes with `langgraph-checkpoint` included. +* `langgraph-checkpoint`: The base interface for checkpointer savers ([BaseCheckpointSaver][langgraph.checkpoint.base.BaseCheckpointSaver]) and serialization/deserialization interface ([SerializerProtocol][langgraph.checkpoint.serde.base.SerializerProtocol]). Includes in-memory checkpointer implementation ([InMemorySaver][langgraph.checkpoint.memory.InMemorySaver]) for experimentation. LangGraph comes with `langgraph-checkpoint` included. * `langgraph-checkpoint-sqlite`: An implementation of LangGraph checkpointer that uses SQLite database ([SqliteSaver][langgraph.checkpoint.sqlite.SqliteSaver] / [AsyncSqliteSaver][langgraph.checkpoint.sqlite.aio.AsyncSqliteSaver]). Ideal for experimentation and local workflows. Needs to be installed separately. * `langgraph-checkpoint-postgres`: An advanced checkpointer that uses Postgres database ([PostgresSaver][langgraph.checkpoint.postgres.PostgresSaver] / [AsyncPostgresSaver][langgraph.checkpoint.postgres.aio.AsyncPostgresSaver]), used in LangGraph Cloud. Ideal for using in production. Needs to be installed separately. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 56a55a5c6..e16754d6e 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -26,7 +26,7 @@ from langgraph.checkpoint.serde.types import TASKS, ChannelProtocol logger = logging.getLogger(__name__) -class MemorySaver( +class InMemorySaver( BaseCheckpointSaver[str], AbstractContextManager, AbstractAsyncContextManager ): """An in-memory checkpoint saver. @@ -34,7 +34,7 @@ class MemorySaver( This checkpoint saver stores checkpoints in memory using a defaultdict. Note: - Only use `MemorySaver` for debugging or testing purposes. + Only use `InMemorySaver` for debugging or testing purposes. For production use cases we recommend installing [langgraph-checkpoint-postgres](https://pypi.org/project/langgraph-checkpoint-postgres/) and using `PostgresSaver` / `AsyncPostgresSaver`. Args: @@ -44,7 +44,7 @@ class MemorySaver( import asyncio - from langgraph.checkpoint.memory import MemorySaver + from langgraph.checkpoint.memory import InMemorySaver from langgraph.graph import StateGraph builder = StateGraph(int) @@ -52,7 +52,7 @@ class MemorySaver( builder.set_entry_point("add_one") builder.set_finish_point("add_one") - memory = MemorySaver() + memory = InMemorySaver() graph = builder.compile(checkpointer=memory) coro = graph.ainvoke(1, {"configurable": {"thread_id": "thread-1"}}) asyncio.run(coro) # Output: 2 @@ -84,7 +84,7 @@ class MemorySaver( self.stack.enter_context(self.storage) # type: ignore[arg-type] self.stack.enter_context(self.writes) # type: ignore[arg-type] - def __enter__(self) -> "MemorySaver": + def __enter__(self) -> "InMemorySaver": return self.stack.__enter__() def __exit__( @@ -95,7 +95,7 @@ class MemorySaver( ) -> Optional[bool]: return self.stack.__exit__(exc_type, exc_value, traceback) - async def __aenter__(self) -> "MemorySaver": + async def __aenter__(self) -> "InMemorySaver": return self.stack.__enter__() async def __aexit__( @@ -149,15 +149,17 @@ class MemorySaver( pending_writes=[ (id, c, self.serde.loads_typed(v)) for id, c, v, _ in writes ], - parent_config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": parent_checkpoint_id, + parent_config=( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, + } } - } - if parent_checkpoint_id - else None, + if parent_checkpoint_id + else None + ), ) else: if checkpoints := self.storage[thread_id][checkpoint_ns]: @@ -193,15 +195,17 @@ class MemorySaver( pending_writes=[ (id, c, self.serde.loads_typed(v)) for id, c, v, _ in writes ], - parent_config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": parent_checkpoint_id, + parent_config=( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, + } } - } - if parent_checkpoint_id - else None, + if parent_checkpoint_id + else None + ), ) def list( @@ -307,15 +311,17 @@ class MemorySaver( ], }, metadata=metadata, - parent_config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": parent_checkpoint_id, + parent_config=( + { + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, + } } - } - if parent_checkpoint_id - else None, + if parent_checkpoint_id + else None + ), pending_writes=[ (id, c, self.serde.loads_typed(v)) for id, c, v, _ in writes ], @@ -492,6 +498,9 @@ class MemorySaver( return f"{next_v:032}.{next_h:016}" +MemorySaver = InMemorySaver # Kept for backwards compatibility + + class PersistentDict(defaultdict): """Persistent dictionary with an API compatible with shelve and anydbm. diff --git a/libs/checkpoint/pyproject.toml b/libs/checkpoint/pyproject.toml index f59f45ff4..7e9c41d3b 100644 --- a/libs/checkpoint/pyproject.toml +++ b/libs/checkpoint/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langgraph-checkpoint" -version = "2.0.10" +version = "2.0.11" description = "Library with base interfaces for LangGraph checkpoint savers." authors = [] license = "MIT" diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 578437233..847cf931f 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -9,13 +9,13 @@ from langgraph.checkpoint.base import ( create_checkpoint, empty_checkpoint, ) -from langgraph.checkpoint.memory import MemorySaver +from langgraph.checkpoint.memory import InMemorySaver class TestMemorySaver: @pytest.fixture(autouse=True) def setup(self) -> None: - self.memory_saver = MemorySaver() + self.memory_saver = InMemorySaver() # objects for test setup self.config_1: RunnableConfig = { @@ -138,3 +138,9 @@ class TestMemorySaver: c async for c in self.memory_saver.alist(None, filter=query_4) ] assert len(search_results_4) == 0 + + +def test_memory_saver() -> None: + from langgraph.checkpoint.memory import MemorySaver + + assert isinstance(MemorySaver(), InMemorySaver) diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index 0a9f13a47..f88f0358f 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -15,7 +15,7 @@ from langgraph.checkpoint.base import ( SerializerProtocol, copy_checkpoint, ) -from langgraph.checkpoint.memory import MemorySaver, PersistentDict +from langgraph.checkpoint.memory import InMemorySaver, PersistentDict class NoopSerializer(SerializerProtocol): @@ -26,7 +26,7 @@ class NoopSerializer(SerializerProtocol): return "type", obj -class MemorySaverAssertImmutable(MemorySaver): +class MemorySaverAssertImmutable(InMemorySaver): storage_for_copies: defaultdict[str, dict[str, dict[str, Checkpoint]]] def __init__( @@ -71,7 +71,7 @@ class MemorySaverAssertImmutable(MemorySaver): return super().put(config, checkpoint, metadata, new_versions) -class MemorySaverAssertCheckpointMetadata(MemorySaver): +class MemorySaverAssertCheckpointMetadata(InMemorySaver): """This custom checkpointer is for verifying that a run's configurable fields are merged with the previous checkpoint config for each step in the run. This is the desired behavior. Because the checkpointer's (a)put() @@ -126,7 +126,7 @@ class MemorySaverAssertCheckpointMetadata(MemorySaver): ) -class MemorySaverNoPending(MemorySaver): +class MemorySaverNoPending(InMemorySaver): def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: result = super().get_tuple(config) if result: diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a54e77889..248bb7be2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -52,7 +52,7 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, ) -from langgraph.checkpoint.memory import MemorySaver +from langgraph.checkpoint.memory import InMemorySaver, MemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START from langgraph.errors import InvalidUpdateError from langgraph.func import entrypoint, task @@ -222,7 +222,7 @@ def test_graph_validation_with_command() -> None: def test_checkpoint_errors() -> None: - class FaultyGetCheckpointer(MemorySaver): + class FaultyGetCheckpointer(InMemorySaver): def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: raise ValueError("Faulty get_tuple") @@ -236,13 +236,13 @@ def test_checkpoint_errors() -> None: ) -> RunnableConfig: raise ValueError("Faulty put") - class FaultyPutWritesCheckpointer(MemorySaver): + class FaultyPutWritesCheckpointer(InMemorySaver): def put_writes( self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str ) -> RunnableConfig: raise ValueError("Faulty put_writes") - class FaultyVersionCheckpointer(MemorySaver): + class FaultyVersionCheckpointer(InMemorySaver): def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: raise ValueError("Faulty get_next_version") @@ -4060,7 +4060,7 @@ def test_xray_lance(snapshot: SnapshotAssertion): interview_builder.add_conditional_edges("answer_question", route_messages) # Set up memory - memory = MemorySaver() + memory = InMemorySaver() # Interview interview_graph = interview_builder.compile(checkpointer=memory).with_config( @@ -4288,7 +4288,7 @@ def test_subgraph_retries(): parent.add_edge("parent_node", "child_graph") parent.set_entry_point("parent_node") - checkpointer = MemorySaver() + checkpointer = InMemorySaver() app = parent.compile(checkpointer=checkpointer) with pytest.raises(RandomError): app.invoke({"count": 0}, {"configurable": {"thread_id": "foo"}}) @@ -4411,7 +4411,7 @@ def test_debug_retry(): builder.add_edge("one", "two") builder.add_edge("two", END) - saver = MemorySaver() + saver = InMemorySaver() graph = builder.compile(checkpointer=saver) @@ -4479,7 +4479,7 @@ def test_debug_subgraphs(): parent.add_edge("p_one", "p_two") parent.add_edge("p_two", END) - graph = parent.compile(checkpointer=MemorySaver()) + graph = parent.compile(checkpointer=InMemorySaver()) config = {"configurable": {"thread_id": "1"}} events = [ @@ -4555,7 +4555,7 @@ def test_debug_nested_subgraphs(): grand_parent.add_edge("gp_one", "gp_two") grand_parent.add_edge("gp_two", END) - graph = grand_parent.compile(checkpointer=MemorySaver()) + graph = grand_parent.compile(checkpointer=InMemorySaver()) config = {"configurable": {"thread_id": "1"}} events = [ diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 5db7f6e57..f369a7aa7 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -48,7 +48,7 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, ) -from langgraph.checkpoint.memory import MemorySaver +from langgraph.checkpoint.memory import InMemorySaver, MemorySaver from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START from langgraph.errors import InvalidUpdateError, NodeInterrupt from langgraph.func import entrypoint, task @@ -99,11 +99,11 @@ NEEDS_CONTEXTVARS = pytest.mark.skipif( async def test_checkpoint_errors() -> None: - class FaultyGetCheckpointer(MemorySaver): + class FaultyGetCheckpointer(InMemorySaver): async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: raise ValueError("Faulty get_tuple") - class FaultyPutCheckpointer(MemorySaver): + class FaultyPutCheckpointer(InMemorySaver): async def aput( self, config: RunnableConfig, @@ -113,13 +113,13 @@ async def test_checkpoint_errors() -> None: ) -> RunnableConfig: raise ValueError("Faulty put") - class FaultyPutWritesCheckpointer(MemorySaver): + class FaultyPutWritesCheckpointer(InMemorySaver): async def aput_writes( self, config: RunnableConfig, writes: List[Tuple[str, Any]], task_id: str ) -> RunnableConfig: raise ValueError("Faulty put_writes") - class FaultyVersionCheckpointer(MemorySaver): + class FaultyVersionCheckpointer(InMemorySaver): def get_next_version(self, current: Optional[int], channel: BaseChannel) -> int: raise ValueError("Faulty get_next_version") @@ -5866,7 +5866,7 @@ async def test_debug_retry(): builder.add_edge("one", "two") builder.add_edge("two", END) - saver = MemorySaver() + saver = InMemorySaver() graph = builder.compile(checkpointer=saver) @@ -5939,7 +5939,7 @@ async def test_debug_subgraphs(): parent.add_edge("p_one", "p_two") parent.add_edge("p_two", END) - graph = parent.compile(checkpointer=MemorySaver()) + graph = parent.compile(checkpointer=InMemorySaver()) config = {"configurable": {"thread_id": "1"}} events = [ @@ -6014,7 +6014,7 @@ async def test_debug_nested_subgraphs(): grand_parent.add_edge("gp_one", "gp_two") grand_parent.add_edge("gp_two", END) - graph = grand_parent.compile(checkpointer=MemorySaver()) + graph = grand_parent.compile(checkpointer=InMemorySaver()) config = {"configurable": {"thread_id": "1"}} events = [