From 8ff5c43cf054fcb6a58fc66e5165c06b365a6cc2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 11 Apr 2025 09:10:58 -0700 Subject: [PATCH] Avoid creating contexts for control branches --- libs/langgraph/langgraph/graph/branch.py | 1 + libs/langgraph/langgraph/graph/state.py | 2 +- libs/langgraph/langgraph/utils/runnable.py | 38 +++++++++++++--------- 3 files changed, 24 insertions(+), 17 deletions(-) diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py index a1ae358d6..fd2039f20 100644 --- a/libs/langgraph/langgraph/graph/branch.py +++ b/libs/langgraph/langgraph/graph/branch.py @@ -138,6 +138,7 @@ class Branch(NamedTuple): reader=reader, name=None, trace=False, + set_context=False, func_accepts_config=True, ) ) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 417e52a55..e31557191 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1153,9 +1153,9 @@ CONTROL_BRANCH_PATH = RunnableCallable( tags=[TAG_HIDDEN], trace=False, recurse=False, + set_context=False, func_accepts_config=True, ) -CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) def _get_channels( diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 4dab44023..2736cc7d5 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -254,6 +254,7 @@ class RunnableCallable(Runnable): tags: Optional[Sequence[str]] = None, trace: bool = True, recurse: bool = True, + set_context: bool = True, explode_args: bool = False, func_accepts_config: Optional[bool] = None, **kwargs: Any, @@ -277,6 +278,7 @@ class RunnableCallable(Runnable): self.kwargs = kwargs self.trace = trace self.recurse = recurse + self.set_context = set_context self.explode_args = explode_args # check signature if func is None and afunc is None: @@ -363,17 +365,22 @@ class RunnableCallable(Runnable): ) try: child_config = patch_config(config, callbacks=run_manager.get_child()) - with set_config_context(child_config) as context: - ret = context.run(self.func, *args, **kwargs) + if self.set_context: + with set_config_context(child_config) as context: + ret = context.run(self.func, *args, **kwargs) + else: + ret = self.func(*args, **kwargs) except BaseException as e: run_manager.on_chain_error(e) raise else: run_manager.on_chain_end(ret) - else: + elif self.set_context: with set_config_context(config) as context: ret = context.run(self.func, *args, **kwargs) - if isinstance(ret, Runnable) and self.recurse: + else: + ret = self.func(*args, **kwargs) + if self.recurse and isinstance(ret, Runnable): return ret.invoke(input, config) return ret @@ -417,25 +424,24 @@ class RunnableCallable(Runnable): ) try: child_config = patch_config(config, callbacks=run_manager.get_child()) - with set_config_context(child_config) as context: - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - if ASYNCIO_ACCEPTS_CONTEXT: + coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) + if ASYNCIO_ACCEPTS_CONTEXT and self.set_context: + with set_config_context(child_config) as context: ret = await asyncio.create_task(coro, context=context) - else: - ret = await coro + else: + ret = await coro except BaseException as e: await run_manager.on_chain_error(e) raise else: await run_manager.on_chain_end(ret) - else: + elif ASYNCIO_ACCEPTS_CONTEXT and self.set_context: with set_config_context(config) as context: - if ASYNCIO_ACCEPTS_CONTEXT: - coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) - ret = await asyncio.create_task(coro, context=context) - else: - ret = await self.afunc(*args, **kwargs) - if isinstance(ret, Runnable) and self.recurse: + coro = cast(Coroutine[None, None, Any], self.afunc(*args, **kwargs)) + ret = await asyncio.create_task(coro, context=context) + else: + ret = await self.afunc(*args, **kwargs) + if self.recurse and isinstance(ret, Runnable): return await ret.ainvoke(input, config) return ret