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

langchain-ai · Nov 4, 2024 · 2e656d9 · 2e656d9
1 parent 1fb8e01
commit 2e656d9
Show file tree

Hide file tree

Showing 6 changed files with 695 additions and 59 deletions.
diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py
@@ -814,7 +814,7 @@ def update_state(
                 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 @@ def update_state(
             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 @@ def update_state(
                 checkpoint_config,
                 checkpoint,
                 {
+                    **checkpoint_metadata,
                     "source": "update",
                     "step": step + 1,
                     "writes": {as_node: values},
@@ -966,7 +970,7 @@ async def aupdate_state(
                 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 @@ async def aupdate_state(
             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 @@ async def aupdate_state(
                 checkpoint_config,
                 checkpoint,
                 {
+                    **checkpoint_metadata,
                     "source": "update",
                     "step": step + 1,
                     "writes": {as_node: values},

diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py
@@ -491,6 +491,8 @@ def _first(self, *, input_keys: Union[str, Sequence[str]]) -> None:
             )
 
     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, {})

diff --git a/libs/langgraph/tests/test_prebuilt.py b/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 == []