Ensure node tracing output is exactly what was returned from node

langgraph/graph/graph.py

@@ -40,12 +40,17 @@ class Branch(NamedTuple):
     condition: Runnable[Any, str]
     ends: Optional[dict[str, str]]
 
-    def run(self, writer: Callable[[str], Optional[Runnable]]) -> None:
+    def run(
+        self,
+        writer: Callable[[str], Optional[Runnable]],
+        reader: Optional[Callable[[RunnableConfig], Any]] = None,
+    ) -> None:
         return ChannelWrite.register_writer(
             RunnableCallable(
                 func=self._route,
                 afunc=self._aroute,
                 writer=writer,
+                reader=reader,
                 name=None,
                 trace=False,
             )
@@ -56,9 +61,10 @@ class Branch(NamedTuple):
         input: Any,
         config: RunnableConfig,
         *,
+        reader: Optional[Callable[[], Any]],
         writer: Callable[[str], Optional[Runnable]],
     ) -> Runnable:
-        result = self.condition.invoke(input, config)
+        result = self.condition.invoke(reader(config) if reader else input, config)
         if self.ends:
             destination = self.ends[result]
         else:
@@ -70,9 +76,12 @@ class Branch(NamedTuple):
         input: Any,
         config: RunnableConfig,
         *,
+        reader: Optional[Callable[[], Any]],
         writer: Callable[[str], Optional[Runnable]],
     ) -> Runnable:
-        result = await self.condition.ainvoke(input, config)
+        result = await self.condition.ainvoke(
+            reader(config) if reader else input, config
+        )
         if self.ends:
             destination = self.ends[result]
         else:

langgraph/graph/message.py

@@ -3,13 +3,14 @@ from typing import Annotated, Union
 
 from langchain_core.messages import (
     AnyMessage,
+    MessageLikeRepresentation,
     convert_to_messages,
     message_chunk_to_message,
 )
 
 from langgraph.graph.state import StateGraph
 
-Messages = Union[list[AnyMessage], AnyMessage]
+Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation]
 
 
 def add_messages(left: Messages, right: Messages) -> Messages:

langgraph/graph/state.py

@@ -142,18 +142,6 @@ class CompiledStateGraph(CompiledGraph):
             )
             for key in state_keys
         ]
-        # node that reads current state with (this node's) updates applied
-        state_reader = ChannelRead(
-            state_keys[0] if state_keys == ["__root__"] else state_keys,
-            tags=[TAG_HIDDEN],
-            fresh=True,
-            # coerce state dict to schema class (eg. pydantic model)
-            mapper=(
-                None
-                if state_keys == ["__root__"]
-                else partial(_coerce_state, self.graph.schema)
-            ),
-        )
 
         # add node and output channel
         if key == START:
@@ -163,8 +151,6 @@ class CompiledStateGraph(CompiledGraph):
                 channels=[START],
                 writers=[
                     ChannelWrite(state_write_entries, tags=[TAG_HIDDEN]),
-                    # read back state with updates applied
-                    state_reader,
                 ],
             )
         else:
@@ -178,15 +164,17 @@ class CompiledStateGraph(CompiledGraph):
                     else {chan: chan for chan in state_keys}
                 ),
                 # coerce state dict to schema class (eg. pydantic model)
-                mapper=state_reader.mapper,
+                mapper=(
+                    None
+                    if state_keys == ["__root__"]
+                    else partial(_coerce_state, self.graph.schema)
+                ),
                 writers=[
                     # publish to this channel and state keys
                     ChannelWrite(
                         [ChannelWriteEntry(key)] + state_write_entries,
                         tags=[TAG_HIDDEN],
                     ),
-                    # read back state with updates applied
-                    state_reader,
                 ],
             ).pipe(node)
 
@@ -226,7 +214,7 @@ class CompiledStateGraph(CompiledGraph):
                 )
 
         # attach branch publisher
-        self.nodes[start] |= branch.run(branch_writer)
+        self.nodes[start] |= branch.run(branch_writer, _get_state_reader(self.graph))
 
         # attach branch subscribers
         ends = branch.ends.values() if branch.ends else [node for node in self.nodes]
@@ -237,6 +225,19 @@ class CompiledStateGraph(CompiledGraph):
                 self.nodes[end].triggers.append(channel_name)
 
 
+def _get_state_reader(graph: StateGraph) -> ChannelRead:
+    state_keys = list(graph.channels)
+    return partial(
+        ChannelRead.do_read,
+        channel=state_keys[0] if state_keys == ["__root__"] else state_keys,
+        fresh=True,
+        # coerce state dict to schema class (eg. pydantic model)
+        mapper=(
+            None if state_keys == ["__root__"] else partial(_coerce_state, graph.schema)
+        ),
+    )
+
+
 def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
     return schema(**input)
 

langgraph/pregel/read.py

@@ -65,30 +65,34 @@ class ChannelRead(RunnableCallable):
         return super().get_name(suffix, name=name)
 
     def _read(self, _: Any, config: RunnableConfig) -> Any:
-        try:
-            read: READ_TYPE = config["configurable"][CONFIG_KEY_READ]
-        except KeyError:
-            raise RuntimeError(
-                f"Runnable {self} is not configured with a read function"
-                "Make sure to call in the context of a Pregel process"
-            )
-        if self.mapper:
-            return self.mapper(read(self.channel, self.fresh))
-        else:
-            return read(self.channel, self.fresh)
+        return self.do_read(
+            config, channel=self.channel, fresh=self.fresh, mapper=self.mapper
+        )
 
     async def _aread(self, _: Any, config: RunnableConfig) -> Any:
+        return self.do_read(
+            config, channel=self.channel, fresh=self.fresh, mapper=self.mapper
+        )
+
+    @staticmethod
+    def do_read(
+        config: RunnableConfig,
+        *,
+        channel: Union[str, list[str]],
+        fresh: bool = False,
+        mapper: Optional[Callable[[Any], Any]] = None,
+    ) -> Any:
         try:
             read: READ_TYPE = config["configurable"][CONFIG_KEY_READ]
         except KeyError:
             raise RuntimeError(
-                f"Runnable {self} is not configured with a read function"
+                "Not configured with a read function"
                 "Make sure to call in the context of a Pregel process"
             )
-        if self.mapper:
-            return self.mapper(read(self.channel, self.fresh))
+        if mapper:
+            return mapper(read(channel, fresh))
         else:
-            return read(self.channel, self.fresh)
+            return read(channel, fresh)
 
 
 DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough()
@@ -110,20 +114,6 @@ class PregelNode(RunnableBindingBase):
     def get_writers(self) -> list[Runnable]:
         """Get writers with optimizations applied."""
         writers = self.writers.copy()
-        while writers and isinstance(writers[-1], ChannelRead):
-            # we can avoid reads if no writers would be called after them
-            writers.pop()
-        while (
-            len(writers) > 1
-            and isinstance(writers[-1], ChannelWrite)
-            and all(
-                write.value is not None and not isinstance(write.value, Runnable)
-                for write in writers[-1].writes
-            )
-            and isinstance(writers[-2], ChannelRead)
-        ):
-            # we can avoid reads if all subsequent write values don't use the input
-            writers.pop(-2)
         while (
             len(writers) > 1
             and isinstance(writers[-1], ChannelWrite)

tests/test_pregel.py

@@ -2307,7 +2307,7 @@ def test_message_graph(
         FunctionMessage(
             content="result for query",
             name="search_api",
-            id="00000000-0000-4000-8000-000000000013",
+            id="00000000-0000-4000-8000-000000000012",
         ),
         AIMessage(
             content="",
@@ -2319,7 +2319,7 @@ def test_message_graph(
         FunctionMessage(
             content="result for another",
             name="search_api",
-            id="00000000-0000-4000-8000-000000000024",
+            id="00000000-0000-4000-8000-000000000022",
         ),
         AIMessage(content="answer", id="ai3"),
     ]
@@ -2338,7 +2338,7 @@ def test_message_graph(
             "action": FunctionMessage(
                 content="result for query",
                 name="search_api",
-                id="00000000-0000-4000-8000-000000000043",
+                id="00000000-0000-4000-8000-000000000039",
             )
         },
         {
@@ -2354,7 +2354,7 @@ def test_message_graph(
             "action": FunctionMessage(
                 content="result for another",
                 name="search_api",
-                id="00000000-0000-4000-8000-000000000054",
+                id="00000000-0000-4000-8000-000000000049",
             )
         },
         {"agent": AIMessage(content="answer", id="ai3")},