From c03f1c29003348dd5a6e964eed1123ed450337b6 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 4 Nov 2024 09:43:58 -0500 Subject: [PATCH] lint --- libs/langgraph/langgraph/constants.py | 2 ++ libs/langgraph/langgraph/pregel/loop.py | 3 ++- libs/langgraph/tests/test_remote_graph.py | 8 ++++++-- 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 36e375751..d38650214 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -67,6 +67,8 @@ # holds the current checkpoint_ns, "" for root graph CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished") # callback to be called when a node is finished +CONFIG_KEY_METADATA = sys.intern("metadata") +# holds metadata for the graph invocation # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 8cb66b3f3..dbb16258c 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -44,6 +44,7 @@ CONFIG_KEY_DEDUPE_TASKS, CONFIG_KEY_DELEGATE, CONFIG_KEY_ENSURE_LATEST, + CONFIG_KEY_METADATA, CONFIG_KEY_RESUMING, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, @@ -494,7 +495,7 @@ def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: # assign step and parents metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) - metadata.update(self.config["metadata"]) + metadata.update(self.config[CONFIG_KEY_METADATA]) # debug flag if self.debug: print_step_checkpoint( diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 21fb8278f..70857ed61 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/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",