Merge configurable fields with previous checkpoint config before each run (#510)

* Merge configurable fields with previous checkpoint config.

* Update update_state() and aupdate_state() to merge configurable fields with previous checkpoint config.

* Update tests to verify that all checkpoint metadata contain the expected configurable field keys. This assertion is needed because a run can have an arbitrary number of steps based on the construction of the graph.

langgraph/pregel/__init__.py

@@ -550,8 +550,19 @@ class Pregel(
             # apply to checkpoint and save
             _apply_writes(checkpoint, channels, task.writes)
             step = saved.metadata.get("step", -2) + 1 if saved else -1
+
+            # merge configurable fields with previous checkpoint config
+            checkpoint_config = config
+            if saved:
+                checkpoint_config = {
+                    "configurable": {
+                        **config.get("configurable", {}),
+                        **saved.config["configurable"],
+                    }
+                }
+
             return self.checkpointer.put(
-                saved.config if saved else config,
+                checkpoint_config,
                 create_checkpoint(checkpoint, channels, step),
                 {
                     "source": "update",
@@ -625,8 +636,19 @@ class Pregel(
             # apply to checkpoint and save
             _apply_writes(checkpoint, channels, task.writes)
             step = saved.metadata.get("step", -2) + 1 if saved else -1
+
+            # merge configurable fields with previous checkpoint config
+            checkpoint_config = config
+            if saved:
+                checkpoint_config = {
+                    "configurable": {
+                        **config.get("configurable", {}),
+                        **saved.config["configurable"],
+                    }
+                }
+
             return await self.checkpointer.aput(
-                saved.config if saved else config,
+                checkpoint_config,
                 create_checkpoint(checkpoint, channels, step),
                 {
                     "source": "update",
@@ -728,7 +750,17 @@ class Pregel(
             # get checkpoint from saver, or create an empty one
             saved = self.checkpointer.get_tuple(config) if self.checkpointer else None
             checkpoint = saved.checkpoint if saved else empty_checkpoint()
-            checkpoint_config = saved.config if saved else config
+
+            # merge configurable fields with previous checkpoint config
+            checkpoint_config = config
+            if saved:
+                checkpoint_config = {
+                    "configurable": {
+                        **config.get("configurable", {}),
+                        **saved.config["configurable"],
+                    }
+                }
+
             start = saved.metadata.get("step", -2) + 1 if saved else -1
             # create channels from checkpoint
             with ChannelsManager(
@@ -1002,7 +1034,17 @@ class Pregel(
                 else None
             )
             checkpoint = saved.checkpoint if saved else empty_checkpoint()
-            checkpoint_config = saved.config if saved else config
+
+            # merge configurable fields with previous checkpoint config
+            checkpoint_config = config
+            if saved:
+                checkpoint_config = {
+                    "configurable": {
+                        **config.get("configurable", {}),
+                        **saved.config["configurable"],
+                    }
+                }
+
             start = saved.metadata.get("step", -2) + 1 if saved else -1
             # create channels from checkpoint
             async with AsyncChannelsManager(

tests/memory_assert.py

@@ -1,6 +1,9 @@
+import asyncio
 from collections import defaultdict
 from typing import Any, Optional
 
+from langchain_core.runnables import RunnableConfig
+
 from langgraph.checkpoint.base import (
     Checkpoint,
     CheckpointMetadata,
@@ -46,3 +49,63 @@ class MemorySaverAssertImmutable(MemorySaver):
         )
         # call super to write checkpoint
         return super().put(config, checkpoint, metadata)
+
+
+class MemorySaverAssertCheckpointMetadata(MemorySaver):
+    """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.
+    """
+
+    serde = NoopSerializer()
+
+    def __init__(
+        self,
+        *,
+        serde: Optional[SerializerProtocol] = None,
+    ) -> None:
+        super().__init__(serde=serde)
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: Optional[CheckpointMetadata] = None,
+    ) -> 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 thread_ts to make testing simpler
+        configurable.pop("thread_ts", None)
+
+        self.storage[config["configurable"]["thread_id"]].update(
+            {
+                checkpoint["id"]: (
+                    self.serde.dumps(checkpoint),
+                    # merge configurable fields and metadata
+                    self.serde.dumps({**configurable, **metadata}),
+                )
+            }
+        )
+        return {
+            "configurable": {
+                "thread_id": config["configurable"]["thread_id"],
+                "thread_ts": checkpoint["id"],
+            }
+        }
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+    ) -> RunnableConfig:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.put, config, checkpoint, metadata
+        )

tests/test_pregel.py

@@ -41,7 +41,10 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from tests.any_str import AnyStr
-from tests.memory_assert import MemorySaverAssertImmutable
+from tests.memory_assert import (
+    MemorySaverAssertCheckpointMetadata,
+    MemorySaverAssertImmutable,
+)
 
 
 def test_graph_validation() -> None:
@@ -5492,3 +5495,161 @@ def test_repeat_condition(snapshot: SnapshotAssertion) -> None:
 
     app = workflow.compile()
     assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
+
+
+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.
+    """
+    # set up test
+    from langchain_core.language_models.fake_chat_models import (
+        FakeMessagesListChatModel,
+    )
+    from langchain_core.messages import AIMessage, AnyMessage
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_core.tools import tool
+
+    # graph state
+    class BaseState(TypedDict):
+        messages: Annotated[list[AnyMessage], add_messages]
+
+    # initialize graph nodes
+    @tool()
+    def search_api(query: str) -> str:
+        """Searches the API for the query."""
+        return f"result for {query}"
+
+    tools = [search_api]
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a nice assistant."),
+            ("placeholder", "{messages}"),
+        ]
+    )
+
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="",
+                tool_calls=[
+                    {
+                        "id": "tool_call123",
+                        "name": "search_api",
+                        "args": {"query": "query"},
+                    },
+                ],
+            ),
+            AIMessage(content="answer"),
+        ]
+    )
+
+    def agent(state: BaseState, config: RunnableConfig) -> BaseState:
+        formatted = prompt.invoke(state)
+        response = model.invoke(formatted)
+        return {"messages": response}
+
+    def should_continue(data: BaseState) -> str:
+        # Logic to decide whether to continue in the loop or exit
+        if not data["messages"][-1].tool_calls:
+            return "exit"
+        else:
+            return "continue"
+
+    # define graphs w/ and w/o interrupt
+    workflow = StateGraph(BaseState)
+    workflow.add_node("agent", agent)
+    workflow.add_node("tools", ToolNode(tools))
+    workflow.set_entry_point("agent")
+    workflow.add_conditional_edges(
+        "agent", should_continue, {"continue": "tools", "exit": END}
+    )
+    workflow.add_edge("tools", "agent")
+
+    # graph w/o interrupt
+    checkpointer_1 = MemorySaverAssertCheckpointMetadata()
+    app = workflow.compile(checkpointer=checkpointer_1)
+
+    # graph w/ interrupt
+    checkpointer_2 = MemorySaverAssertCheckpointMetadata()
+    app_w_interrupt = workflow.compile(
+        checkpointer=checkpointer_2, interrupt_before=["tools"]
+    )
+
+    # assertions
+
+    # invoke graph w/o interrupt
+    app.invoke(
+        {"messages": ["what is weather in sf"]},
+        {
+            "configurable": {
+                "thread_id": "1",
+                "test_config_1": "foo",
+                "test_config_2": "bar",
+            },
+        },
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    # assert that checkpoint metadata contains the run's configurable fields
+    chkpnt_metadata_1 = checkpointer_1.get_tuple(config).metadata
+    assert chkpnt_metadata_1["thread_id"] == "1"
+    assert chkpnt_metadata_1["test_config_1"] == "foo"
+    assert chkpnt_metadata_1["test_config_2"] == "bar"
+
+    # Verify that all checkpoint metadata have the expected keys. This check
+    # is needed because a run may have an arbitrary number of steps depending
+    # on how the graph is constructed.
+    chkpnt_tuples_1 = checkpointer_1.list(config)
+    for chkpnt_tuple in chkpnt_tuples_1:
+        assert chkpnt_tuple.metadata["thread_id"] == "1"
+        assert chkpnt_tuple.metadata["test_config_1"] == "foo"
+        assert chkpnt_tuple.metadata["test_config_2"] == "bar"
+
+    # invoke graph, but interrupt before tool call
+    app_w_interrupt.invoke(
+        {"messages": ["what is weather in sf"]},
+        {
+            "configurable": {
+                "thread_id": "2",
+                "test_config_3": "foo",
+                "test_config_4": "bar",
+            },
+        },
+    )
+
+    config = {"configurable": {"thread_id": "2"}}
+
+    # assert that checkpoint metadata contains the run's configurable fields
+    chkpnt_metadata_2 = checkpointer_2.get_tuple(config).metadata
+    assert chkpnt_metadata_2["thread_id"] == "2"
+    assert chkpnt_metadata_2["test_config_3"] == "foo"
+    assert chkpnt_metadata_2["test_config_4"] == "bar"
+
+    # resume graph execution
+    app_w_interrupt.invoke(
+        input=None,
+        config={
+            "configurable": {
+                "thread_id": "2",
+                "test_config_3": "foo",
+                "test_config_4": "bar",
+            }
+        },
+    )
+
+    # assert that checkpoint metadata contains the run's configurable fields
+    chkpnt_metadata_3 = checkpointer_2.get_tuple(config).metadata
+    assert chkpnt_metadata_3["thread_id"] == "2"
+    assert chkpnt_metadata_3["test_config_3"] == "foo"
+    assert chkpnt_metadata_3["test_config_4"] == "bar"
+
+    # Verify that all checkpoint metadata have the expected keys. This check
+    # is needed because a run may have an arbitrary number of steps depending
+    # on how the graph is constructed.
+    chkpnt_tuples_2 = checkpointer_2.list(config)
+    for chkpnt_tuple in chkpnt_tuples_2:
+        assert chkpnt_tuple.metadata["thread_id"] == "2"
+        assert chkpnt_tuple.metadata["test_config_3"] == "foo"
+        assert chkpnt_tuple.metadata["test_config_4"] == "bar"

tests/test_pregel_async.py

@@ -41,7 +41,10 @@ from langgraph.prebuilt.tool_node import ToolNode
 from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot
 from langgraph.pregel.retry import RetryPolicy
 from tests.any_str import AnyStr
-from tests.memory_assert import MemorySaverAssertImmutable
+from tests.memory_assert import (
+    MemorySaverAssertCheckpointMetadata,
+    MemorySaverAssertImmutable,
+)
 
 
 async def test_node_cancellation_on_external_cancel() -> None:
@@ -4779,3 +4782,161 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
                 ]
             }
     assert times_called == 1
+
+
+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.
+    """
+    # set up test
+    from langchain_core.language_models.fake_chat_models import (
+        FakeMessagesListChatModel,
+    )
+    from langchain_core.messages import AIMessage, AnyMessage
+    from langchain_core.prompts import ChatPromptTemplate
+    from langchain_core.tools import tool
+
+    # graph state
+    class BaseState(TypedDict):
+        messages: Annotated[list[AnyMessage], add_messages]
+
+    # initialize graph nodes
+    @tool()
+    def search_api(query: str) -> str:
+        """Searches the API for the query."""
+        return f"result for {query}"
+
+    tools = [search_api]
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", "You are a nice assistant."),
+            ("placeholder", "{messages}"),
+        ]
+    )
+
+    model = FakeMessagesListChatModel(
+        responses=[
+            AIMessage(
+                content="",
+                tool_calls=[
+                    {
+                        "id": "tool_call123",
+                        "name": "search_api",
+                        "args": {"query": "query"},
+                    },
+                ],
+            ),
+            AIMessage(content="answer"),
+        ]
+    )
+
+    def agent(state: BaseState, config: RunnableConfig) -> BaseState:
+        formatted = prompt.invoke(state)
+        response = model.invoke(formatted)
+        return {"messages": response}
+
+    def should_continue(data: BaseState) -> str:
+        # Logic to decide whether to continue in the loop or exit
+        if not data["messages"][-1].tool_calls:
+            return "exit"
+        else:
+            return "continue"
+
+    # define graphs w/ and w/o interrupt
+    workflow = StateGraph(BaseState)
+    workflow.add_node("agent", agent)
+    workflow.add_node("tools", ToolNode(tools))
+    workflow.set_entry_point("agent")
+    workflow.add_conditional_edges(
+        "agent", should_continue, {"continue": "tools", "exit": END}
+    )
+    workflow.add_edge("tools", "agent")
+
+    # graph w/o interrupt
+    checkpointer_1 = MemorySaverAssertCheckpointMetadata()
+    app = workflow.compile(checkpointer=checkpointer_1)
+
+    # graph w/ interrupt
+    checkpointer_2 = MemorySaverAssertCheckpointMetadata()
+    app_w_interrupt = workflow.compile(
+        checkpointer=checkpointer_2, interrupt_before=["tools"]
+    )
+
+    # assertions
+
+    # invoke graph w/o interrupt
+    await app.ainvoke(
+        {"messages": ["what is weather in sf"]},
+        {
+            "configurable": {
+                "thread_id": "1",
+                "test_config_1": "foo",
+                "test_config_2": "bar",
+            },
+        },
+    )
+
+    config = {"configurable": {"thread_id": "1"}}
+
+    # assert that checkpoint metadata contains the run's configurable fields
+    chkpnt_metadata_1 = (await checkpointer_1.aget_tuple(config)).metadata
+    assert chkpnt_metadata_1["thread_id"] == "1"
+    assert chkpnt_metadata_1["test_config_1"] == "foo"
+    assert chkpnt_metadata_1["test_config_2"] == "bar"
+
+    # Verify that all checkpoint metadata have the expected keys. This check
+    # is needed because a run may have an arbitrary number of steps depending
+    # on how the graph is constructed.
+    chkpnt_tuples_1 = checkpointer_1.alist(config)
+    async for chkpnt_tuple in chkpnt_tuples_1:
+        assert chkpnt_tuple.metadata["thread_id"] == "1"
+        assert chkpnt_tuple.metadata["test_config_1"] == "foo"
+        assert chkpnt_tuple.metadata["test_config_2"] == "bar"
+
+    # invoke graph, but interrupt before tool call
+    await app_w_interrupt.ainvoke(
+        {"messages": ["what is weather in sf"]},
+        {
+            "configurable": {
+                "thread_id": "2",
+                "test_config_3": "foo",
+                "test_config_4": "bar",
+            },
+        },
+    )
+
+    config = {"configurable": {"thread_id": "2"}}
+
+    # assert that checkpoint metadata contains the run's configurable fields
+    chkpnt_metadata_2 = (await checkpointer_2.aget_tuple(config)).metadata
+    assert chkpnt_metadata_2["thread_id"] == "2"
+    assert chkpnt_metadata_2["test_config_3"] == "foo"
+    assert chkpnt_metadata_2["test_config_4"] == "bar"
+
+    # resume graph execution
+    await app_w_interrupt.ainvoke(
+        input=None,
+        config={
+            "configurable": {
+                "thread_id": "2",
+                "test_config_3": "foo",
+                "test_config_4": "bar",
+            }
+        },
+    )
+
+    # assert that checkpoint metadata contains the run's configurable fields
+    chkpnt_metadata_3 = (await checkpointer_2.aget_tuple(config)).metadata
+    assert chkpnt_metadata_3["thread_id"] == "2"
+    assert chkpnt_metadata_3["test_config_3"] == "foo"
+    assert chkpnt_metadata_3["test_config_4"] == "bar"
+
+    # Verify that all checkpoint metadata have the expected keys. This check
+    # is needed because a run may have an arbitrary number of steps depending
+    # on how the graph is constructed.
+    chkpnt_tuples_2 = checkpointer_2.alist(config)
+    async for chkpnt_tuple in chkpnt_tuples_2:
+        assert chkpnt_tuple.metadata["thread_id"] == "2"
+        assert chkpnt_tuple.metadata["test_config_3"] == "foo"
+        assert chkpnt_tuple.metadata["test_config_4"] == "bar"