From 3d12a2df59277fa8c14f3addd1980b0e2a93b494 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 11 Apr 2025 09:00:16 -0700 Subject: [PATCH 1/3] Simplify path for control branch attached to every node - attached to every node to handle command/send return values - used to be a full blown conditional edge, can be simpler by doing all of it in a single function --- libs/langgraph/langgraph/graph/state.py | 51 ++++++++++++++++--------- 1 file changed, 34 insertions(+), 17 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index fde945709..417e52a55 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -43,7 +43,6 @@ from langgraph.constants import ( MISSING, NS_END, NS_SEP, - SELF, TAG_HIDDEN, ) from langgraph.errors import ( @@ -670,9 +669,9 @@ class StateGraph(Graph): for key, node in self.nodes.items(): compiled.attach_node(key, node) - compiled.attach_branch(START, SELF, CONTROL_BRANCH, with_reader=False) - for key, node in self.nodes.items(): - compiled.attach_branch(key, SELF, CONTROL_BRANCH, with_reader=False) + compiled.nodes[START].writers.append(CONTROL_BRANCH_PATH) + for key in self.nodes: + compiled.nodes[key].writers.append(CONTROL_BRANCH_PATH) for start, end in self.edges: compiled.attach_edge(start, end) @@ -1084,9 +1083,10 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) -def _control_branch(value: Any) -> Sequence[Union[str, Send]]: +def _control_branch(value: Any, config: RunnableConfig) -> Sequence[Union[str, Send]]: if isinstance(value, Send): - return [value] + ChannelWrite.do_write(config, (value,)) + return value commands: list[Command] = [] if isinstance(value, Command): commands.append(value) @@ -1094,22 +1094,32 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]: for cmd in value: if isinstance(cmd, Command): commands.append(cmd) - rtn: list[Union[str, Send]] = [] + rtn: list[Union[ChannelWriteEntry, Send]] = [] for command in commands: if command.graph == Command.PARENT: raise ParentCommand(command) if isinstance(command.goto, Send): rtn.append(command.goto) elif isinstance(command.goto, str): - rtn.append(command.goto) + rtn.append(ChannelWriteEntry(CHANNEL_BRANCH_TO.format(command.goto), None)) else: - rtn.extend(command.goto) - return rtn + rtn.extend( + go + if isinstance(go, Send) + else ChannelWriteEntry(CHANNEL_BRANCH_TO.format(go), None) + for go in command.goto + ) + if rtn: + ChannelWrite.do_write(config, rtn) + return value -async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: +async def _acontrol_branch( + value: Any, config: RunnableConfig +) -> Sequence[Union[str, Send]]: if isinstance(value, Send): - return [value] + ChannelWrite.do_write(config, (value,)) + return value commands: list[Command] = [] if isinstance(value, Command): commands.append(value) @@ -1117,17 +1127,24 @@ async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]: for cmd in value: if isinstance(cmd, Command): commands.append(cmd) - rtn: list[Union[str, Send]] = [] + rtn: list[Union[ChannelWriteEntry, Send]] = [] for command in commands: if command.graph == Command.PARENT: raise ParentCommand(command) if isinstance(command.goto, Send): rtn.append(command.goto) elif isinstance(command.goto, str): - rtn.append(command.goto) + rtn.append(ChannelWriteEntry(CHANNEL_BRANCH_TO.format(command.goto), None)) else: - rtn.extend(command.goto) - return rtn + rtn.extend( + go + if isinstance(go, Send) + else ChannelWriteEntry(CHANNEL_BRANCH_TO.format(go), None) + for go in command.goto + ) + if rtn: + ChannelWrite.do_write(config, rtn) + return value CONTROL_BRANCH_PATH = RunnableCallable( @@ -1136,7 +1153,7 @@ CONTROL_BRANCH_PATH = RunnableCallable( tags=[TAG_HIDDEN], trace=False, recurse=False, - func_accepts_config=False, + func_accepts_config=True, ) CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) From 8ff5c43cf054fcb6a58fc66e5165c06b365a6cc2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 11 Apr 2025 09:10:58 -0700 Subject: [PATCH 2/3] 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 From 5a7edead8c6a1a2cde2fdde5fd41ebc93cfe8fbe Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 11 Apr 2025 09:20:10 -0700 Subject: [PATCH 3/3] Lint --- libs/langgraph/langgraph/graph/state.py | 51 ++++++++++++------------- 1 file changed, 25 insertions(+), 26 deletions(-) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index e31557191..73d242dfd 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -733,28 +733,6 @@ class CompiledStateGraph(CompiledGraph): if is_writable_managed_value(v) ] - def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]: - if isinstance(input, Command): - if input.graph == Command.PARENT: - return () - return input._update_as_tuples() - elif ( - isinstance(input, (list, tuple)) - and input - and any(isinstance(i, Command) for i in input) - ): - updates: list[tuple[str, Any]] = [] - for i in input: - if isinstance(i, Command): - if i.graph == Command.PARENT: - continue - updates.extend(i._update_as_tuples()) - else: - updates.append(("__root__", i)) - return updates - elif input is not None: - return [("__root__", input)] - def _get_updates( input: Union[None, dict, Any], ) -> Optional[Sequence[tuple[str, Any]]]: @@ -1083,7 +1061,7 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]: return schema(**input) -def _control_branch(value: Any, config: RunnableConfig) -> Sequence[Union[str, Send]]: +def _control_branch(value: Any, config: RunnableConfig) -> Any: if isinstance(value, Send): ChannelWrite.do_write(config, (value,)) return value @@ -1114,9 +1092,7 @@ def _control_branch(value: Any, config: RunnableConfig) -> Sequence[Union[str, S return value -async def _acontrol_branch( - value: Any, config: RunnableConfig -) -> Sequence[Union[str, Send]]: +async def _acontrol_branch(value: Any, config: RunnableConfig) -> Any: if isinstance(value, Send): ChannelWrite.do_write(config, (value,)) return value @@ -1158,6 +1134,29 @@ CONTROL_BRANCH_PATH = RunnableCallable( ) +def _get_root(input: Any) -> Optional[Sequence[tuple[str, Any]]]: + if isinstance(input, Command): + if input.graph == Command.PARENT: + return () + return input._update_as_tuples() + elif ( + isinstance(input, (list, tuple)) + and input + and any(isinstance(i, Command) for i in input) + ): + updates: list[tuple[str, Any]] = [] + for i in input: + if isinstance(i, Command): + if i.graph == Command.PARENT: + continue + updates.extend(i._update_as_tuples()) + else: + updates.append(("__root__", i)) + return updates + elif input is not None: + return [("__root__", input)] + + def _get_channels( schema: Type[dict], ) -> tuple[dict[str, BaseChannel], dict[str, ManagedValueSpec], dict[str, Any]]: