From 69ad42cac583769a456892d38f3d6facd3bc7f96 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 17 Feb 2025 17:33:38 -0800 Subject: [PATCH] Add CONFIG_KEY_RUNNER_SUBMIT --- libs/langgraph/langgraph/constants.py | 2 ++ libs/langgraph/langgraph/pregel/__init__.py | 5 +++-- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index c37a48fa0..4fde5e7dd 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -81,6 +81,8 @@ CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad") # holds a mutable dict for temporary storage scoped to the current task CONFIG_KEY_PREVIOUS = sys.intern("__pregel_previous") # holds the previous return value from a stateful Pregel graph. +CONFIG_KEY_RUNNER_SUBMIT = sys.intern("__pregel_runner_submit") +# holds a function that receives tasks from runner, executes them and returns results # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 173763a9c..5e1301b03 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -59,6 +59,7 @@ from langgraph.constants import ( CONFIG_KEY_NODE_FINISHED, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, + CONFIG_KEY_RUNNER_SUBMIT, CONFIG_KEY_SEND, CONFIG_KEY_STORE, CONFIG_KEY_STREAM, @@ -1755,7 +1756,7 @@ class Pregel(PregelProtocol): ) as loop: # create runner runner = PregelRunner( - submit=loop.submit, + submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), put_writes=loop.put_writes, schedule_task=loop.accept_push, node_finished=config[CONF].get(CONFIG_KEY_NODE_FINISHED), @@ -2047,7 +2048,7 @@ class Pregel(PregelProtocol): ) as loop: # create runner runner = PregelRunner( - submit=loop.submit, + submit=config[CONF].get(CONFIG_KEY_RUNNER_SUBMIT, loop.submit), put_writes=loop.put_writes, schedule_task=loop.accept_push, use_astream=do_stream is not None,