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 217f92fa3..ce2f898d0 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 ( @@ -671,9 +670,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) @@ -735,28 +734,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]]]: @@ -1085,9 +1062,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) -> Any: if isinstance(value, Send): - return [value] + ChannelWrite.do_write(config, (value,)) + return value commands: list[Command] = [] if isinstance(value, Command): commands.append(value) @@ -1095,22 +1073,30 @@ 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) -> Any: if isinstance(value, Send): - return [value] + ChannelWrite.do_write(config, (value,)) + return value commands: list[Command] = [] if isinstance(value, Command): commands.append(value) @@ -1118,17 +1104,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( @@ -1137,9 +1130,32 @@ CONTROL_BRANCH_PATH = RunnableCallable( tags=[TAG_HIDDEN], trace=False, recurse=False, - func_accepts_config=False, + set_context=False, + func_accepts_config=True, ) -CONTROL_BRANCH = Branch(CONTROL_BRANCH_PATH, None) + + +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( 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