Clean up PregelNode attributes (#5049)

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(
                             [

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

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(

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}"

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__(