Merge pull request #75 from langchain-ai/nc/31jan/rm-double-channelwrite

Remove double channelwrite

langgraph/graph/state.py

@@ -3,7 +3,7 @@ from functools import partial
 from inspect import signature
 from typing import Any, Optional, Type
 
-from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnablePassthrough
+from langchain_core.runnables import RunnableLambda, RunnablePassthrough
 from langchain_core.runnables.base import RunnableLike
 
 from langgraph.channels.base import BaseChannel, InvalidUpdateError
@@ -13,7 +13,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.graph.graph import END, Graph
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import ChannelRead
-from langgraph.pregel.write import ChannelWrite
+from langgraph.pregel.write import SKIP_WRITE, ChannelWrite
 
 START = "__start__"
 
@@ -39,10 +39,13 @@ class StateGraph(Graph):
 
         state_keys = list(self.channels)
         state_keys_read = state_keys[0] if state_keys == ["__root__"] else state_keys
-        update_state = (
-            partial(_update_state_dict, state_keys)
-            if isinstance(state_keys_read, list)
-            else _update_state_root
+        update_channels = (
+            [("__root__", None, True)]
+            if not isinstance(state_keys_read, list)
+            else [
+                (key, RunnableLambda(partial(_dict_getter, state_keys, key)), False)
+                for key in state_keys_read
+            ]
         )
         coerce_state = (
             partial(_coerce_state, self.schema)
@@ -59,8 +62,7 @@ class StateGraph(Graph):
                 Channel.subscribe_to(f"{key}:inbox")
                 | coerce_state  # coerce/validate using schema
                 | node
-                | partial(update_state, key)
-                | Channel.write_to(key)
+                | ChannelWrite(channels=[(key, None, False)] + update_channels)
             )
             for key, node in self.nodes.items()
         }
@@ -80,11 +82,9 @@ class StateGraph(Graph):
                         branch.runnable, name=f"{key}_condition"
                     )
 
-        nodes[START] = (
-            Channel.subscribe_to(f"{START}:inbox", tags=["langsmith:hidden"])
-            | partial(update_state, START)
-            | Channel.write_to(START)
-        )
+        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)
@@ -105,25 +105,16 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 
 
-def _update_state_dict(
-    state_keys: list[str], node_name: str, input: dict[str, Any], config: RunnableConfig
-) -> dict[str, Any]:
+def _dict_getter(allowed_keys: str, key: str, input: dict) -> Any:
     if input is not None:
-        if not isinstance(input, dict) or any(key not in state_keys for key in input):
+        if not isinstance(input, dict) or any(key not in allowed_keys for key in input):
             raise InvalidUpdateError(
-                f"Invalid state update from node {node_name},"
-                f" expected dict with one or more of {state_keys}, got {input}"
+                f"Invalid state update,"
+                f" expected dict with one or more of {allowed_keys}, got {input}"
             )
-        ChannelWrite.do_write(config, **input)
-    return input
-
-
-def _update_state_root(
-    node_name: str, input: Any, config: RunnableConfig
-) -> dict[str, Any]:
-    if input is not None:
-        ChannelWrite.do_write(config, __root__=input)
-    return input
+        return input.get(key, SKIP_WRITE)
+    else:
+        return SKIP_WRITE
 
 
 def _get_channels(schema: Type[dict]) -> dict[str, BaseChannel]:

langgraph/pregel/__init__.py

@@ -149,8 +149,8 @@ class Channel:
         """Writes to channels the result of the lambda, or None to skip writing."""
         return ChannelWrite(
             channels=(
-                [(c, None) for c in channels]
-                + [(k, _coerce_write_value(v)) for k, v in kwargs.items()]
+                [(c, None, False) for c in channels]
+                + [(k, _coerce_write_value(v), True) for k, v in kwargs.items()]
             )
         )
 

langgraph/pregel/write.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 from typing import Any, Callable, Optional, Sequence
 
 from langchain_core.runnables import (
@@ -14,8 +15,11 @@ from langgraph.constants import CONFIG_KEY_SEND
 TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
 
+SKIP_WRITE = object()
+
+
 class ChannelWrite(RunnablePassthrough):
-    channels: Sequence[tuple[str, Optional[Runnable]]]
+    channels: Sequence[tuple[str, Optional[Runnable], bool]]
     """
     Mapping of write channels to Runnables that return the value to be written,
     or None to skip writing.
@@ -27,10 +31,10 @@ class ChannelWrite(RunnablePassthrough):
     def __init__(
         self,
         *,
-        channels: Sequence[tuple[str, Optional[Runnable]]],
+        channels: Sequence[tuple[str, Optional[Runnable], bool]],
     ):
         super().__init__(func=self._write, afunc=self._awrite, channels=channels)
-        self.name = f"ChannelWrite<{','.join(chan for chan, _ in self.channels)}>"
+        self.name = f"ChannelWrite<{','.join(chan for chan, _, _ in self.channels)}>"
 
     def __repr_args__(self) -> Any:
         return [("channels", self.channels)]
@@ -49,25 +53,28 @@ class ChannelWrite(RunnablePassthrough):
 
     def _write(self, input: Any, config: RunnableConfig) -> None:
         values = [
-            (chan, r.invoke(input, config) if r else input) for chan, r in self.channels
+            (chan, r.invoke(input, config) if r else input)
+            for chan, r, _ in self.channels
         ]
         values = [
             write
-            for write, chan in zip(values, self.channels)
-            if chan[1] is None or write[1] is not None
+            for write, (_, _, skip_none) in zip(values, self.channels)
+            if not skip_none or write[1] is not None
         ]
 
         self.do_write(config, **dict(values))
 
     async def _awrite(self, input: Any, config: RunnableConfig) -> None:
+        values = await asyncio.gather(
+            *(
+                r.ainvoke(input, config) if r else _mk_future(input)
+                for _, r, _ in self.channels
+            )
+        )
         values = [
-            (chan, await r.ainvoke(input, config) if r else input)
-            for chan, r in self.channels
-        ]
-        values = [
-            write
-            for write, chan in zip(values, self.channels)
-            if chan[1] is None or write[1] is not None
+            (chan, val)
+            for val, (chan, _, skip_none) in zip(values, self.channels)
+            if not skip_none or val is not None
         ]
 
         self.do_write(config, **dict(values))
@@ -75,4 +82,10 @@ class ChannelWrite(RunnablePassthrough):
     @staticmethod
     def do_write(config: RunnableConfig, **values: Any) -> None:
         write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND]
-        write([(chan, val) for chan, val in values.items()])
+        write([(chan, val) for chan, val in values.items() if val is not SKIP_WRITE])
+
+
+def _mk_future(val: Any) -> asyncio.Future:
+    fut = asyncio.Future()
+    fut.set_result(val)
+    return fut