diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 8bcbb2b94..4c26324b7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -69,6 +69,8 @@ CONFIG_KEY_ENSURE_LATEST = sys.intern("__pregel_ensure_latest") # (for distributed mode) CONFIG_KEY_DELEGATE = sys.intern("__pregel_delegate") # holds a boolean indicating whether to delegate subgraphs (for distributed mode) +CONFIG_KEY_THREAD_ID = sys.intern("thread_id") +# holds the thread ID for the current invocation CONFIG_KEY_CHECKPOINT_MAP = sys.intern("checkpoint_map") # holds a mapping of checkpoint_ns -> checkpoint_id for parent graphs CONFIG_KEY_CHECKPOINT_ID = sys.intern("checkpoint_id") diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index e1de05602..ad75921d3 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -52,6 +52,7 @@ from langgraph.constants import ( CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_STREAM, CONFIG_KEY_TASK_ID, + CONFIG_KEY_THREAD_ID, EMPTY_SEQ, ERROR, INPUT, @@ -285,6 +286,12 @@ class PregelLoop(LoopProtocol): ) else: self.checkpoint_config = self.config + if thread_id := config[CONF].get(CONFIG_KEY_THREAD_ID): + if not isinstance(thread_id, str): + self.checkpoint_config = patch_configurable( + self.checkpoint_config, + {CONFIG_KEY_THREAD_ID: str(thread_id)}, + ) 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/test_pregel.py b/libs/langgraph/tests/test_pregel.py index cf1c70dab..b79d569c5 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6156,7 +6156,7 @@ def test_falsy_return_from_task( falsy_task().result() interrupt("test") - configurable = {"configurable": {"thread_id": str(uuid.uuid4())}} + configurable = {"configurable": {"thread_id": uuid.uuid4()}} assert [ chunk for chunk in graph.stream({"a": 5}, configurable, stream_mode="debug") ] == [