Avoid creating contexts for control branches

This commit is contained in:
Nuno Campos
2025-04-11 09:10:58 -07:00
parent 3d12a2df59
commit 8ff5c43cf0
3 changed files with 24 additions and 17 deletions
+1
View File
@@ -138,6 +138,7 @@ class Branch(NamedTuple):
reader=reader,
name=None,
trace=False,
set_context=False,
func_accepts_config=True,
)
)
+1 -1
View File
@@ -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(
+22 -16
View File
@@ -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