Merge pull request #147 from langchain-ai/nc/24feb/optimize-run-tree

Remove unnecessary runs from StateGraph/MessageGraph run tree

langgraph/graph/state.py

@@ -3,7 +3,7 @@ from functools import partial
 from inspect import signature
 from typing import Any, Optional, Sequence, Type
 
-from langchain_core.runnables import RunnableLambda, RunnablePassthrough
+from langchain_core.runnables import RunnableLambda
 from langchain_core.runnables.base import RunnableLike
 
 from langgraph.channels.any_value import AnyValue
@@ -14,8 +14,8 @@ from langgraph.channels.last_value import LastValue
 from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph.graph import END, START, CompiledGraph, Graph
 from langgraph.pregel import Channel
-from langgraph.pregel.read import ChannelRead
-from langgraph.pregel.write import SKIP_WRITE, ChannelWrite
+from langgraph.pregel.read import ChannelInvoke
+from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry
 
 
 class StateGraph(Graph):
@@ -46,18 +46,25 @@ class StateGraph(Graph):
 
         state_keys = list(self.channels)
         state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
+        state_channels = (
+            {chan: chan for chan in state_keys}
+            if isinstance(state_keys_read, list)
+            else {None: state_keys_read}
+        )
         update_channels = (
-            [("__root__", None, True)]
+            [ChannelWriteEntry("__root__", None, True)]
             if not isinstance(state_keys_read, list)
             else [
-                (key, RunnableLambda(partial(_dict_getter, state_keys, key)), False)
+                ChannelWriteEntry(
+                    key, RunnableLambda(partial(_dict_getter, state_keys, key)), False
+                )
                 for key in state_keys_read
             ]
         )
         coerce_state = (
             partial(_coerce_state, self.schema)
             if isinstance(state_keys_read, list)
-            else RunnablePassthrough()
+            else None
         )
 
         outgoing_edges = defaultdict(list)
@@ -66,10 +73,15 @@ class StateGraph(Graph):
 
         nodes = {
             key: (
-                Channel.subscribe_to(f"{key}:inbox")
-                | coerce_state  # coerce/validate using schema
+                ChannelInvoke(
+                    triggers=[f"{key}:inbox"],
+                    channels=state_channels,
+                    mapper=coerce_state,
+                )
                 | node
-                | ChannelWrite(channels=[(key, None, False)] + update_channels)
+                | ChannelWrite(
+                    channels=[ChannelWriteEntry(key, None, False)] + update_channels
+                )
             )
             for key, node in self.nodes.items()
         }
@@ -89,11 +101,16 @@ class StateGraph(Graph):
             outgoing = outgoing_edges[key]
             edges_key = f"{key}:edges"
             if outgoing or key in self.branches:
-                nodes[edges_key] = Channel.subscribe_to(
-                    key, tags=["langsmith:hidden"]
-                ) | ChannelRead(state_keys_read)
+                nodes[edges_key] = ChannelInvoke(
+                    triggers=[key], tags=["langsmith:hidden"], channels=state_channels
+                )
             if outgoing:
-                nodes[edges_key] |= Channel.write_to(*[dest for dest in outgoing])
+                nodes[edges_key] |= ChannelWrite(
+                    channels=[
+                        ChannelWriteEntry(dest, None if dest == END else key, True)
+                        for dest in outgoing
+                    ]
+                )
             if key in self.branches:
                 for branch in self.branches[key]:
                     nodes[edges_key] |= RunnableLambda(
@@ -102,10 +119,12 @@ class StateGraph(Graph):
 
         nodes[START] = Channel.subscribe_to(
             f"{START}:inbox", tags=["langsmith:hidden"]
-        ) | ChannelWrite(channels=[(START, None, False)] + update_channels)
-        nodes[f"{START}:edges"] = Channel.subscribe_to(
-            START, tags=["langsmith:hidden"]
-        ) | ChannelRead(state_keys_read)
+        ) | ChannelWrite(
+            channels=[ChannelWriteEntry(START, None, False)] + update_channels
+        )
+        nodes[f"{START}:edges"] = ChannelInvoke(
+            triggers=[START], tags=["langsmith:hidden"], channels=state_channels
+        )
         if self.entry_point:
             nodes[f"{START}:edges"] |= Channel.write_to(f"{self.entry_point}:inbox")
         elif self.entry_point_branch:

langgraph/pregel/__init__.py

@@ -63,7 +63,7 @@ from langgraph.pregel.log import logger
 from langgraph.pregel.read import ChannelBatch, ChannelInvoke
 from langgraph.pregel.reserved import ReservedChannels
 from langgraph.pregel.validate import validate_graph, validate_keys
-from langgraph.pregel.write import ChannelWrite
+from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 
 WriteValue = Union[
     Runnable[Input, Output],
@@ -149,8 +149,11 @@ class Channel:
         """Writes to channels the result of the lambda, or None to skip writing."""
         return ChannelWrite(
             channels=(
-                [(c, None, False) for c in channels]
-                + [(k, _coerce_write_value(v), True) for k, v in kwargs.items()]
+                [ChannelWriteEntry(c, None, False) for c in channels]
+                + [
+                    ChannelWriteEntry(k, _coerce_write_value(v), True)
+                    for k, v in kwargs.items()
+                ]
             )
         )
 
@@ -789,7 +792,8 @@ def _prepare_next_tasks(
                 checkpoint["channel_versions"][chan] > seen[chan]
                 for chan in proc.triggers
             ):
-                # If all channels subscribed by this process are not empty
+                # If all trigger channels subscribed by this process are not empty
+                # then invoke the process with the values of all non-empty channels
                 try:
                     val: Any = {
                         k: _read_channel(
@@ -800,6 +804,10 @@ def _prepare_next_tasks(
                 except EmptyChannelError:
                     continue
 
+                # If the process has a mapper, apply it to the value
+                if proc.mapper is not None:
+                    val = proc.mapper(val)
+
                 # Processes that subscribe to a single keyless channel get
                 # the value directly, instead of a dict
                 if list(proc.channels.keys()) == [None]:

langgraph/pregel/read.py

@@ -79,6 +79,8 @@ class ChannelInvoke(RunnableBindingBase):
 
     triggers: list[str] = Field(default_factory=list)
 
+    mapper: Optional[Callable[[Any], Any]] = None
+
     when: Optional[Callable[[Any], bool]] = None
 
     bound: Runnable[Any, Any] = Field(default=default_bound)
@@ -89,6 +91,7 @@ class ChannelInvoke(RunnableBindingBase):
         self,
         channels: Mapping[None, str] | Mapping[str, str],
         triggers: Sequence[str],
+        mapper: Optional[Callable[[Any], Any]] = None,
         when: Optional[Callable[[Any], bool]] = None,
         tags: Optional[list[str]] = None,
         *,
@@ -100,6 +103,7 @@ class ChannelInvoke(RunnableBindingBase):
         super().__init__(
             channels=channels,
             triggers=triggers,
+            mapper=mapper,
             when=when,
             bound=bound or default_bound,
             kwargs=kwargs or {},
@@ -120,6 +124,7 @@ class ChannelInvoke(RunnableBindingBase):
                 **{chan: chan for chan in channels},
             },
             triggers=self.triggers,
+            mapper=self.mapper,
             when=self.when,
             bound=self.bound,
             kwargs=self.kwargs,
@@ -138,6 +143,7 @@ class ChannelInvoke(RunnableBindingBase):
             return ChannelInvoke(
                 channels=self.channels,
                 triggers=self.triggers,
+                mapper=self.mapper,
                 when=self.when,
                 bound=coerce_to_runnable(other),
                 kwargs=self.kwargs,
@@ -147,6 +153,7 @@ class ChannelInvoke(RunnableBindingBase):
             return ChannelInvoke(
                 channels=self.channels,
                 triggers=self.triggers,
+                mapper=self.mapper,
                 when=self.when,
                 # delegate to __or__ in self.bound
                 bound=self.bound | other,

langgraph/pregel/write.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import asyncio
-from typing import Any, Callable, Optional, Sequence
+from typing import Any, Callable, NamedTuple, Optional, Sequence, Union
 
 from langchain_core.runnables import (
     Runnable,
@@ -18,21 +18,25 @@ TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 SKIP_WRITE = object()
 
 
+class ChannelWriteEntry(NamedTuple):
+    channel: str
+    value: Optional[Union[Any, Runnable]]
+    skip_none: bool
+
+
 class ChannelWrite(RunnablePassthrough):
-    channels: Sequence[tuple[str, Optional[Runnable], bool]]
+    channels: Sequence[ChannelWriteEntry]
     """
-    Mapping of write channels to Runnables that return the value to be written,
-    or None to skip writing.
+    Sequence of write entries, each of which is a tuple of:
+    - channel name
+    - runnable to map input, or None to use the input, or any other value to use instead
+    - whether to skip writing if the mapped value is None
     """
 
     class Config:
         arbitrary_types_allowed = True
 
-    def __init__(
-        self,
-        *,
-        channels: Sequence[tuple[str, Optional[Runnable], bool]],
-    ):
+    def __init__(self, *, channels: Sequence[ChannelWriteEntry]):
         super().__init__(func=self._write, afunc=self._awrite, channels=channels)
         self.name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.channels)}>"
 
@@ -53,7 +57,14 @@ class ChannelWrite(RunnablePassthrough):
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
         values = [
-            (chan, r.invoke(input, config) if r else input)
+            (
+                chan,
+                r.invoke(input, config)
+                if isinstance(r, Runnable)
+                else r
+                if r is not None
+                else input,
+            )
             for chan, r, _ in self.channels
         ]
         values = [
@@ -67,7 +78,11 @@ class ChannelWrite(RunnablePassthrough):
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
         values = await asyncio.gather(
             *(
-                r.ainvoke(input, config) if r else _mk_future(input)
+                r.ainvoke(input, config)
+                if isinstance(r, Runnable)
+                else _mk_future(r)
+                if r is not None
+                else _mk_future(input)
                 for _, r, _ in self.channels
             )
         )