From 516175780d92b87fee5e56deabb8fb3ce8e9eb9d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 10 Jun 2025 16:14:15 -0700 Subject: [PATCH] Clean up things for Matt! --- libs/langgraph/langgraph/func/__init__.py | 2 +- libs/langgraph/langgraph/graph/state.py | 10 +++---- libs/langgraph/langgraph/pregel/__init__.py | 21 ++++++-------- libs/langgraph/langgraph/pregel/algo.py | 31 ++++++++++----------- libs/langgraph/langgraph/pregel/read.py | 31 +++++---------------- 5 files changed, 36 insertions(+), 59 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 4c7595c7e..2b3095b23 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -499,7 +499,7 @@ class entrypoint: func.__name__: PregelNode( bound=bound, triggers=[START], - channels=[START], + channels=START, writers=[ ChannelWrite( [ diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b4210a351..9fa1fa082 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -996,18 +996,18 @@ class CompiledStateGraph( self.nodes[key] = PregelNode( tags=[TAG_HIDDEN], triggers=[START], - channels=[START], + channels=START, writers=[ChannelWrite(write_entries)], ) elif node is not None: input_schema = node.input if node else self.builder._state_schema - input_values = {k: k for k in self.builder.schemas[input_schema]} - is_single_input = len(input_values) == 1 and "__root__" in input_values + input_channels = list(self.builder.schemas[input_schema]) + is_single_input = len(input_channels) == 1 and "__root__" in input_channels if input_schema in self.schema_to_mapper: mapper = self.schema_to_mapper[input_schema] else: mapper = _pick_mapper( - list(input_values), + input_channels, input_schema, ) self.schema_to_mapper[input_schema] = mapper @@ -1021,7 +1021,7 @@ class CompiledStateGraph( self.nodes[key] = PregelNode( triggers=[branch_channel], # read state keys and managed values - channels=(list(input_values) if is_single_input else input_values), + channels=("__root__" if is_single_input else input_channels), # coerce state dict to schema class (eg. pydantic model) mapper=mapper, # publish to state keys diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 5519ea237..bd77d1487 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -145,7 +145,7 @@ class NodeBuilder: "_cache_policy", ) - _channels: list[str] | dict[str, str] + _channels: str | list[str] _triggers: list[str] _tags: list[str] _metadata: dict[str, Any] @@ -157,7 +157,7 @@ class NodeBuilder: def __init__( self, ) -> None: - self._channels = {} + self._channels = [] self._triggers = [] self._tags = [] self._metadata = {} @@ -171,10 +171,8 @@ class NodeBuilder: channel: str, ) -> Self: """Subscribe to a single channel.""" - if isinstance(self._channels, list): - self._channels.append(channel) - elif not self._channels: - self._channels = [channel] + if not self._channels: + self._channels = channel else: raise ValueError( "Cannot subscribe to single channels when other channels are already subscribed to" @@ -200,15 +198,15 @@ class NodeBuilder: Returns: Self for chaining """ - if isinstance(self._channels, list): + if isinstance(self._channels, str): raise ValueError( "Cannot subscribe to channels when subscribed to a single channel" ) if read: if not self._channels: - self._channels = {chan: chan for chan in channels} + self._channels = list(channels) else: - self._channels.update({chan: chan for chan in channels}) + self._channels.extend(channels) if isinstance(channels, str): self._triggers.append(channels) @@ -222,11 +220,10 @@ class NodeBuilder: *channels: str, ) -> Self: """Adds the specified channels to read from, without subscribing to them.""" - assert self._channels, "Channels must be specified first" - assert isinstance(self._channels, dict), ( + assert isinstance(self._channels, list), ( "Cannot read additional channels when subscribed to single channels" ) - self._channels.update({c: c for c in channels}) + self._channels.extend(channels) return self def do( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index eb83de225..4fd9cab60 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -922,18 +922,18 @@ def _triggers( seen: ChannelVersions | None, null_version: V, proc: PregelNode, -) -> Sequence[str]: +) -> bool: if seen is None: for chan in proc.triggers: if channels[chan].is_available(): - return (chan,) + return True else: for chan in proc.triggers: if channels[chan].is_available() and versions.get( # type: ignore[operator] chan, null_version ) > seen.get(chan, null_version): - return (chan,) - return EMPTY_SEQ + return True + return False def _scratchpad( @@ -1019,25 +1019,22 @@ def _proc_input( return copy(input_cache[proc.input_cache_key]) # If all trigger channels subscribed by this process are not empty # then invoke the process with the values of all non-empty channels - if isinstance(proc.channels, dict): + if isinstance(proc.channels, list): val: dict[str, Any] = {} - for k, chan in proc.channels.items(): - if chan in channels: - if channels[chan].is_available(): - val[k] = channels[chan].get() - else: - val[k] = managed[k].get(scratchpad) - elif isinstance(proc.channels, list): for chan in proc.channels: if chan in channels: if channels[chan].is_available(): - val = channels[chan].get() - break + val[chan] = channels[chan].get() else: - val = managed[chan].get(scratchpad) - break + val[chan] = managed[chan].get(scratchpad) + elif isinstance(proc.channels, str): + if proc.channels in channels: + if channels[proc.channels].is_available(): + val = channels[proc.channels].get() + else: + return MISSING else: - return MISSING + val = managed[proc.channels].get(scratchpad) else: raise RuntimeError( f"Invalid channels type, expected list or dict, got {proc.channels}" diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 91228e992..bccda5ead 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -101,11 +101,10 @@ class PregelNode(Runnable): itself, but instead acts as a container for the components necessary to make a PregelExecutableTask for a node.""" - channels: list[str] | Mapping[str, str] + channels: str | list[str] """The channels that will be passed as input to `bound`. - If a list, the node will be invoked with the first of that isn't empty. - If a dict, the keys are the names of the channels, and the values are the keys - to use in the input to `bound`.""" + If a str, the node will be invoked with its value if it isn't empty. + If a list, the node will be invoked with a dict of those channels' values.""" triggers: list[str] """If any of these channels is written to, this node will be triggered in @@ -140,7 +139,7 @@ class PregelNode(Runnable): def __init__( self, *, - channels: list[str] | Mapping[str, str], + channels: str | list[str], triggers: Sequence[str], mapper: Callable[[Any], Any] | None = None, writers: list[Runnable] | None = None, @@ -223,25 +222,9 @@ class PregelNode(Runnable): This is used to avoid calculating the same input multiple times.""" return ( self.mapper, - tuple(f"{key}:{value}" for key, value in self.channels.items()) - if isinstance(self.channels, dict) - else tuple(self.channels), - ) - - def join(self, channels: Sequence[str]) -> PregelNode: - assert isinstance(channels, list) or isinstance(channels, tuple), ( - "channels must be a list or tuple" - ) - assert isinstance(self.channels, dict), ( - "all channels must be named when using .join()" - ) - return self.copy( - update=dict( - channels={ - **self.channels, - **{chan: chan for chan in channels}, - } - ), + tuple(self.channels) + if isinstance(self.channels, list) + else (self.channels,), ) def __or__(