langgraph: add config metadata to pregel loop (#2323)

libs/langgraph/langgraph/pregel/__init__.py

@@ -814,7 +814,7 @@ class Pregel(PregelProtocol):
                 raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
 
         # get last checkpoint
-        config = merge_configs(self.config, config) if self.config else config
+        config = ensure_config(self.config, config)
         saved = checkpointer.get_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
         checkpoint_previous_versions = (
@@ -826,14 +826,17 @@ class Pregel(PregelProtocol):
             config,
             {CONFIG_KEY_CHECKPOINT_NS: config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")},
         )
+        checkpoint_metadata = config["metadata"]
         if saved:
             checkpoint_config = patch_configurable(config, saved.config[CONF])
+            checkpoint_metadata = {**saved.metadata, **checkpoint_metadata}
         # find last node that updated the state, if not provided
         if values is None and as_node is None:
             next_config = checkpointer.put(
                 checkpoint_config,
                 create_checkpoint(checkpoint, None, step),
                 {
+                    **checkpoint_metadata,
                     "source": "update",
                     "step": step + 1,
                     "writes": {},
@@ -922,6 +925,7 @@ class Pregel(PregelProtocol):
                 checkpoint_config,
                 checkpoint,
                 {
+                    **checkpoint_metadata,
                     "source": "update",
                     "step": step + 1,
                     "writes": {as_node: values},
@@ -966,7 +970,7 @@ class Pregel(PregelProtocol):
                 raise ValueError(f"Subgraph {recast_checkpoint_ns} not found")
 
         # get last checkpoint
-        config = merge_configs(self.config, config) if self.config else config
+        config = ensure_config(self.config, config)
         saved = await checkpointer.aget_tuple(config)
         checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint()
         checkpoint_previous_versions = (
@@ -978,14 +982,17 @@ class Pregel(PregelProtocol):
             config,
             {CONFIG_KEY_CHECKPOINT_NS: config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")},
         )
+        checkpoint_metadata = config["metadata"]
         if saved:
             checkpoint_config = patch_configurable(config, saved.config[CONF])
+            checkpoint_metadata = {**saved.metadata, **checkpoint_metadata}
         # find last node that updated the state, if not provided
         if values is None and as_node is None:
             next_config = await checkpointer.aput(
                 checkpoint_config,
                 create_checkpoint(checkpoint, None, step),
                 {
+                    **checkpoint_metadata,
                     "source": "update",
                     "step": step + 1,
                     "writes": {},
@@ -1072,6 +1079,7 @@ class Pregel(PregelProtocol):
                 checkpoint_config,
                 checkpoint,
                 {
+                    **checkpoint_metadata,
                     "source": "update",
                     "step": step + 1,
                     "writes": {as_node: values},

libs/langgraph/langgraph/pregel/loop.py

@@ -491,6 +491,8 @@ class PregelLoop(LoopProtocol):
             )
 
     def _put_checkpoint(self, metadata: CheckpointMetadata) -> None:
+        for k, v in self.config["metadata"].items():
+            metadata.setdefault(k, v)  # type: ignore
         # assign step and parents
         metadata["step"] = self.step
         metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {})

libs/langgraph/tests/test_prebuilt.py

@@ -158,6 +158,7 @@ def test_no_modifier(request: pytest.FixtureRequest, checkpointer_name: str) ->
             "source": "loop",
             "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
             "step": 1,
+            "thread_id": "123",
         }
         assert saved.pending_writes == []
 
@@ -189,6 +190,7 @@ async def test_no_modifier_async(checkpointer_name: str) -> None:
                 "source": "loop",
                 "writes": {"agent": {"messages": [AIMessage(content="hi?", id="0")]}},
                 "step": 1,
+                "thread_id": "123",
             }
             assert saved.pending_writes == []
 

libs/langgraph/tests/test_remote_graph.py

@@ -73,7 +73,9 @@ def test_get_graph():
         "__start__": DrawableNode(
             id="__start__", name="__start__", data="__start__", metadata=None
         ),
-        "__end__": DrawableNode(id="__end__", name="__end__", data="__end__", metadata=None),
+        "__end__": DrawableNode(
+            id="__end__", name="__end__", data="__end__", metadata=None
+        ),
         "agent": DrawableNode(
             id="agent",
             name="agent_1",
@@ -120,7 +122,9 @@ async def test_aget_graph():
         "__start__": DrawableNode(
             id="__start__", name="__start__", data="__start__", metadata=None
         ),
-        "__end__": DrawableNode(id="__end__", name="__end__", data="__end__", metadata=None),
+        "__end__": DrawableNode(
+            id="__end__", name="__end__", data="__end__", metadata=None
+        ),
         "agent": DrawableNode(
             id="agent",
             name="agent_1",