diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 6eeb8867b..d2202693f 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -70,6 +70,12 @@ class InMemorySaver( tuple[str, str, str], dict[tuple[str, int], tuple[str, str, tuple[str, bytes], str]], ] + blobs: dict[ + tuple[ + str, str, str, str | int | float + ], # thread id, checkpoint ns, channel, version + tuple[str, bytes], + ] def __init__( self, @@ -80,6 +86,7 @@ class InMemorySaver( super().__init__(serde=serde) self.storage = factory(lambda: defaultdict(dict)) self.writes = factory(dict) + self.blobs = factory() self.stack = ExitStack() if factory is not defaultdict: self.stack.enter_context(self.storage) # type: ignore[arg-type] @@ -107,6 +114,18 @@ class InMemorySaver( ) -> Optional[bool]: return self.stack.__exit__(__exc_type, __exc_value, __traceback) + def _load_blobs( + self, thread_id: str, checkpoint_ns: str, versions: ChannelVersions + ) -> dict[str, Any]: + channel_values: dict[str, Any] = {} + for k, v in versions.items(): + kk = (thread_id, checkpoint_ns, k, v) + if kk in self.blobs: + vv = self.blobs[kk] + if vv[0] != "empty": + channel_values[k] = self.serde.loads_typed(vv) + return channel_values + def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: """Get a checkpoint tuple from the in-memory storage. @@ -121,8 +140,8 @@ class InMemorySaver( Returns: Optional[CheckpointTuple]: The retrieved checkpoint tuple, or None if no matching checkpoint was found. """ - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + thread_id: str = config["configurable"]["thread_id"] + checkpoint_ns: str = config["configurable"].get("checkpoint_ns", "") if checkpoint_id := get_checkpoint_id(config): if saved := self.storage[thread_id][checkpoint_ns].get(checkpoint_id): checkpoint, metadata, parent_checkpoint_id = saved @@ -140,10 +159,14 @@ class InMemorySaver( ) else: sends = [] + checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) return CheckpointTuple( config=config, checkpoint={ - **self.serde.loads_typed(checkpoint), + **checkpoint_, + "channel_values": self._load_blobs( + thread_id, checkpoint_ns, checkpoint_["channel_versions"] + ), "pending_sends": [self.serde.loads_typed(s[2]) for s in sends], }, metadata=self.serde.loads_typed(metadata), @@ -180,6 +203,9 @@ class InMemorySaver( ) else: sends = [] + + checkpoint_ = self.serde.loads_typed(checkpoint) + return CheckpointTuple( config={ "configurable": { @@ -189,7 +215,10 @@ class InMemorySaver( } }, checkpoint={ - **self.serde.loads_typed(checkpoint), + **checkpoint_, + "channel_values": self._load_blobs( + thread_id, checkpoint_ns, checkpoint_["channel_versions"] + ), "pending_sends": [self.serde.loads_typed(s[2]) for s in sends], }, metadata=self.serde.loads_typed(metadata), @@ -297,6 +326,8 @@ class InMemorySaver( else: sends = [] + checkpoint_: Checkpoint = self.serde.loads_typed(checkpoint) + yield CheckpointTuple( config={ "configurable": { @@ -306,7 +337,12 @@ class InMemorySaver( } }, checkpoint={ - **self.serde.loads_typed(checkpoint), + **checkpoint_, + "channel_values": self._load_blobs( + thread_id, + checkpoint_ns, + checkpoint_["channel_versions"], + ), "pending_sends": [ self.serde.loads_typed(s[2]) for s in sends ], @@ -353,6 +389,11 @@ class InMemorySaver( c.pop("pending_sends") # type: ignore[misc] thread_id = config["configurable"]["thread_id"] checkpoint_ns = config["configurable"]["checkpoint_ns"] + values: dict[str, Any] = c.pop("channel_values") # type: ignore[misc] + for k, v in new_versions.items(): + self.blobs[(thread_id, checkpoint_ns, k, v)] = ( + self.serde.dumps_typed(values[k]) if k in values else ("empty", b"") + ) self.storage[thread_id][checkpoint_ns].update( { checkpoint["id"]: ( diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index 3c49219f9..ad2dbdb1e 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -68,7 +68,9 @@ class TestMemorySaver: }, "metadata": {"run_id": "my_run_id"}, } - self.memory_saver.put(config, self.chkpnt_2, self.metadata_2, {}) + self.memory_saver.put( + config, self.chkpnt_2, self.metadata_2, self.chkpnt_2["channel_versions"] + ) checkpoint = self.memory_saver.get_tuple(config) assert checkpoint is not None assert checkpoint.metadata == { @@ -80,9 +82,24 @@ class TestMemorySaver: async def test_search(self) -> None: # set up test # save checkpoints - self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {}) - self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {}) - self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {}) + self.memory_saver.put( + self.config_1, + self.chkpnt_1, + self.metadata_1, + self.chkpnt_1["channel_versions"], + ) + self.memory_saver.put( + self.config_2, + self.chkpnt_2, + self.metadata_2, + self.chkpnt_2["channel_versions"], + ) + self.memory_saver.put( + self.config_3, + self.chkpnt_3, + self.metadata_3, + self.chkpnt_3["channel_versions"], + ) # call method / assertions query_1 = {"source": "input"} # search by 1 key @@ -129,9 +146,24 @@ class TestMemorySaver: async def test_asearch(self) -> None: # set up test # save checkpoints - self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {}) - self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {}) - self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {}) + self.memory_saver.put( + self.config_1, + self.chkpnt_1, + self.metadata_1, + self.chkpnt_1["channel_versions"], + ) + self.memory_saver.put( + self.config_2, + self.chkpnt_2, + self.metadata_2, + self.chkpnt_2["channel_versions"], + ) + self.memory_saver.put( + self.config_3, + self.chkpnt_3, + self.metadata_3, + self.chkpnt_3["channel_versions"], + ) # call method / assertions query_1 = {"source": "input"} # search by 1 key diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 94faff87f..3449ccde3 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -269,13 +269,13 @@ class PregelLoop(LoopProtocol): self.checkpoint_config = patch_configurable( self.config, { - CONFIG_KEY_CHECKPOINT_ID: config[CONF][CONFIG_KEY_CHECKPOINT_MAP][ - self.config[CONF][CONFIG_KEY_CHECKPOINT_NS] - ] + CONFIG_KEY_CHECKPOINT_ID: self.config[CONF][ + CONFIG_KEY_CHECKPOINT_MAP + ][self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]] }, ) else: - self.checkpoint_config = config + self.checkpoint_config = self.config self.checkpoint_ns = ( tuple(cast(str, self.config[CONF][CONFIG_KEY_CHECKPOINT_NS]).split(NS_SEP)) if self.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS) diff --git a/libs/langgraph/tests/memory_assert.py b/libs/langgraph/tests/memory_assert.py index f88f0358f..3a1ef4536 100644 --- a/libs/langgraph/tests/memory_assert.py +++ b/libs/langgraph/tests/memory_assert.py @@ -1,4 +1,3 @@ -import asyncio import os import tempfile from collections import defaultdict @@ -13,7 +12,6 @@ from langgraph.checkpoint.base import ( CheckpointMetadata, CheckpointTuple, SerializerProtocol, - copy_checkpoint, ) from langgraph.checkpoint.memory import InMemorySaver, PersistentDict @@ -63,69 +61,14 @@ class MemorySaverAssertImmutable(InMemorySaver): self.storage_for_copies[thread_id][checkpoint_ns][saved["id"]] ) == saved - ) + ), config["configurable"]["checkpoint_ns"] self.storage_for_copies[thread_id][checkpoint_ns][checkpoint["id"]] = ( - self.serde.dumps_typed(copy_checkpoint(checkpoint)) + self.serde.dumps_typed(checkpoint) ) # call super to write checkpoint return super().put(config, checkpoint, metadata, new_versions) -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() - method is called for each step, the implementation of this checkpointer - should produce a side effect that can be asserted. - """ - - def put( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> None: - """The implementation of put() merges config["configurable"] (a run's - configurable fields) with the metadata field. The state of the - checkpoint metadata can be asserted to confirm that the run's - configurable fields were merged with the previous checkpoint config. - """ - configurable = config["configurable"].copy() - - # remove checkpoint_id to make testing simpler - checkpoint_id = configurable.pop("checkpoint_id", None) - thread_id = config["configurable"]["thread_id"] - checkpoint_ns = config["configurable"]["checkpoint_ns"] - self.storage[thread_id][checkpoint_ns].update( - { - checkpoint["id"]: ( - self.serde.dumps_typed(checkpoint), - # merge configurable fields and metadata - self.serde.dumps_typed({**configurable, **metadata}), - checkpoint_id, - ) - } - ) - return { - "configurable": { - "thread_id": config["configurable"]["thread_id"], - "checkpoint_id": checkpoint["id"], - } - } - - async def aput( - self, - config: RunnableConfig, - checkpoint: Checkpoint, - metadata: CheckpointMetadata, - new_versions: ChannelVersions, - ) -> RunnableConfig: - return await asyncio.get_running_loop().run_in_executor( - None, self.put, config, checkpoint, metadata, new_versions - ) - - class MemorySaverNoPending(InMemorySaver): def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: result = super().get_tuple(config) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 40fc5f1b9..9f547e940 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -84,7 +84,6 @@ from tests.conftest import ( REGULAR_CHECKPOINTERS_SYNC, SHOULD_CHECK_SNAPSHOTS, ) -from tests.memory_assert import MemorySaverAssertCheckpointMetadata from tests.messages import ( _AnyIdAIMessage, _AnyIdAIMessageChunk, @@ -4213,11 +4212,11 @@ def test_checkpoint_metadata() -> None: workflow.add_edge("tools", "agent") # graph w/o interrupt - checkpointer_1 = MemorySaverAssertCheckpointMetadata() + checkpointer_1 = InMemorySaver() app = workflow.compile(checkpointer=checkpointer_1) # graph w/ interrupt - checkpointer_2 = MemorySaverAssertCheckpointMetadata() + checkpointer_2 = InMemorySaver() app_w_interrupt = workflow.compile( checkpointer=checkpointer_2, interrupt_before=["tools"] ) @@ -4635,59 +4634,6 @@ def test_multiple_sinks_subgraphs(snapshot: SnapshotAssertion) -> None: assert app.get_graph(xray=True).draw_mermaid() == snapshot -def test_subgraph_retries(): - class State(TypedDict): - count: int - - class ChildState(State): - some_list: Annotated[list, operator.add] - - called_times = 0 - - class RandomError(ValueError): - """This will be retried on.""" - - def parent_node(state: State): - return {"count": state["count"] + 1} - - def child_node_a(state: ChildState): - nonlocal called_times - # We want it to retry only on node_b - # NOT re-compute the whole graph. - assert not called_times - called_times += 1 - return {"some_list": ["val"]} - - def child_node_b(state: ChildState): - raise RandomError("First attempt fails") - - child = StateGraph(ChildState) - child.add_node(child_node_a) - child.add_node(child_node_b) - child.add_edge("__start__", "child_node_a") - child.add_edge("child_node_a", "child_node_b") - - parent = StateGraph(State) - parent.add_node("parent_node", parent_node) - parent.add_node( - "child_graph", - child.compile(), - retry=RetryPolicy( - max_attempts=3, - retry_on=(RandomError,), - backoff_factor=0.0001, - initial_interval=0.0001, - ), - ) - parent.add_edge("parent_node", "child_graph") - parent.set_entry_point("parent_node") - - checkpointer = InMemorySaver() - app = parent.compile(checkpointer=checkpointer) - with pytest.raises(RandomError): - app.invoke({"count": 0}, {"configurable": {"thread_id": "foo"}}) - - @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @pytest.mark.parametrize("store_name", ALL_STORES_SYNC) def test_store_injected( @@ -6294,6 +6240,7 @@ def test_double_interrupt_subgraph( def invoke_sub_agent(state: AgentState): return subgraph.invoke(state) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} parent_agent = ( StateGraph(AgentState) .add_node("invoke_sub_agent", invoke_sub_agent) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 58bb15cdc..437004ca3 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -77,10 +77,7 @@ from tests.conftest import ( awith_store, ) from tests.fake_tracer import FakeTracer -from tests.memory_assert import ( - MemorySaverAssertCheckpointMetadata, - MemorySaverNoPending, -) +from tests.memory_assert import MemorySaverNoPending from tests.messages import ( _AnyIdAIMessage, _AnyIdAIMessageChunk, @@ -5762,11 +5759,11 @@ async def test_checkpoint_metadata() -> None: workflow.add_edge("tools", "agent") # graph w/o interrupt - checkpointer_1 = MemorySaverAssertCheckpointMetadata() + checkpointer_1 = InMemorySaver() app = workflow.compile(checkpointer=checkpointer_1) # graph w/ interrupt - checkpointer_2 = MemorySaverAssertCheckpointMetadata() + checkpointer_2 = InMemorySaver() app_w_interrupt = workflow.compile( checkpointer=checkpointer_2, interrupt_before=["tools"] ) @@ -7012,6 +7009,8 @@ async def test_double_interrupt_subgraph(checkpointer_name: str) -> None: def invoke_sub_agent(state: AgentState): return subgraph.invoke(state) + thread = {"configurable": {"thread_id": str(uuid.uuid4())}} + parent_agent = ( StateGraph(AgentState) .add_node("invoke_sub_agent", invoke_sub_agent)