WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-24 16:42:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Simplify path for control branch attached to every node (#4247)

Browse Source 
- 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
This commit is contained in:
Nuno Campos
2025-04-11 09:44:18 -07:00
committed by GitHub
parent 0eb32a4251 5a7edead8c
commit d6f2f0c90d
3 changed files with 79 additions and 56 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/branch.py
+1
View File
@@ -138,6 +138,7 @@ class Branch(NamedTuple):
reader=reader,
name=None,
trace=False,
set_context=False,
func_accepts_config=True,
)
)
libs/langgraph/langgraph/graph/state.py
+56 -40
View File
@@ -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(
libs/langgraph/langgraph/utils/runnable.py
+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