Merge pull request #294 from langchain-ai/nc/10apr/fix-null-output

Fix null output for node with conditional edge returning END

langgraph/graph/graph.py

@@ -27,7 +27,7 @@ from langgraph.checkpoint import BaseCheckpointSaver
 from langgraph.constants import TAG_HIDDEN
 from langgraph.pregel import Channel, Pregel
 from langgraph.pregel.read import PregelNode
-from langgraph.pregel.write import ChannelWrite
+from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
 from langgraph.utils import RunnableCallable
 
 logger = logging.getLogger(__name__)
@@ -71,7 +71,7 @@ class Branch(NamedTuple):
             destinations = [self.ends[r] for r in result]
         else:
             destinations = result
-        return writer(destinations)
+        return writer(destinations) or input
 
     async def _aroute(
         self,
@@ -90,7 +90,7 @@ class Branch(NamedTuple):
             destinations = [self.ends[r] for r in result]
         else:
             destinations = result
-        return writer(destinations)
+        return writer(destinations) or input
 
 
 class Graph:
@@ -278,14 +278,16 @@ class CompiledGraph(Pregel):
         self.nodes[key] = (
             PregelNode(channels=[], triggers=[])
             | node
-            | Channel.write_to(key, tags=[TAG_HIDDEN])
+            | ChannelWrite([ChannelWriteEntry(key)], tags=[TAG_HIDDEN])
         )
         cast(list[str], self.stream_channels).append(key)
 
     def attach_edge(self, start: str, end: str) -> None:
         if end == END:
             # publish to end channel
-            self.nodes[start].writers.append(Channel.write_to(END, tags=[TAG_HIDDEN]))
+            self.nodes[start].writers.append(
+                ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])
+            )
         else:
             # subscribe to start channel
             self.nodes[end].triggers.append(start)
@@ -296,7 +298,9 @@ class CompiledGraph(Pregel):
             channels = [
                 f"branch:{start}:{name}:{end}" if end != END else END for end in ends
             ]
-            return Channel.write_to(*channels, tags=[TAG_HIDDEN])
+            return ChannelWrite(
+                [ChannelWriteEntry(ch) for ch in channels], tags=[TAG_HIDDEN]
+            )
 
         # add hidden start node
         if start == START and start not in self.nodes:

langgraph/pregel/__init__.py

@@ -1106,7 +1106,7 @@ def _apply_writes(
             checkpoint["channel_versions"][chan] = max_version + 1
             updated_channels.add(chan)
         else:
-            logger.warning(f"Skipping write for channel {chan} which has no readers")
+            logger.warning(f"Skipping write for channel '{chan}' which has no readers")
     # Channels that weren't updated in this step are notified of a new step
     for chan in channels:
         if chan not in updated_channels:

tests/test_pregel_async.py

@@ -947,6 +947,69 @@ async def test_conditional_graph(checkpoint_at: CheckpointAt) -> None:
 
     # Check that agent (one of the nodes) has its output streamed to the logs
     assert "/logs/agent/streamed_output/-" in patch_paths
+    assert "/logs/agent:2/streamed_output/-" in patch_paths
+    assert "/logs/agent:3/streamed_output/-" in patch_paths
+    # Check that agent (one of the nodes) has its final output set in the logs
+    assert "/logs/agent/final_output" in patch_paths
+    assert "/logs/agent:2/final_output" in patch_paths
+    assert "/logs/agent:3/final_output" in patch_paths
+    assert [
+        p["value"]
+        for log in patches
+        for p in log.ops
+        if p["path"] == "/logs/agent/final_output"
+        or p["path"] == "/logs/agent:2/final_output"
+        or p["path"] == "/logs/agent:3/final_output"
+    ] == [
+        {
+            "input": "what is weather in sf",
+            "agent_outcome": AgentAction(
+                tool="search_api", tool_input="query", log="tool:search_api:query"
+            ),
+        },
+        {
+            "input": "what is weather in sf",
+            "intermediate_steps": [
+                (
+                    AgentAction(
+                        tool="search_api",
+                        tool_input="query",
+                        log="tool:search_api:query",
+                    ),
+                    "result for query",
+                )
+            ],
+            "agent_outcome": AgentAction(
+                tool="search_api",
+                tool_input="another",
+                log="tool:search_api:another",
+            ),
+        },
+        {
+            "input": "what is weather in sf",
+            "intermediate_steps": [
+                (
+                    AgentAction(
+                        tool="search_api",
+                        tool_input="query",
+                        log="tool:search_api:query",
+                    ),
+                    "result for query",
+                ),
+                (
+                    AgentAction(
+                        tool="search_api",
+                        tool_input="another",
+                        log="tool:search_api:another",
+                    ),
+                    "result for another",
+                ),
+            ],
+            "agent_outcome": AgentFinish(
+                return_values={"answer": "answer"}, log="finish:answer"
+            ),
+        },
+    ]
 
     # test state get/update methods with interrupt_after
 
@@ -1534,6 +1597,38 @@ async def test_conditional_graph_state(checkpoint_at: CheckpointAt) -> None:
         },
     ]
 
+    patches = [c async for c in app.astream_log({"input": "what is weather in sf"})]
+    patch_paths = {op["path"] for log in patches for op in log.ops}
+
+    # Check that agent (one of the nodes) has its output streamed to the logs
+    assert "/logs/agent/streamed_output/-" in patch_paths
+    # Check that agent (one of the ndoes) has its final output set in the logs
+    assert "/logs/agent/final_output" in patch_paths
+    assert [
+        p["value"]
+        for log in patches
+        for p in log.ops
+        if p["path"] == "/logs/agent/final_output"
+        or p["path"] == "/logs/agent:2/final_output"
+        or p["path"] == "/logs/agent:3/final_output"
+    ] == [
+        {
+            "agent_outcome": AgentAction(
+                tool="search_api", tool_input="query", log="tool:search_api:query"
+            )
+        },
+        {
+            "agent_outcome": AgentAction(
+                tool="search_api", tool_input="another", log="tool:search_api:another"
+            )
+        },
+        {
+            "agent_outcome": AgentFinish(
+                return_values={"answer": "answer"}, log="finish:answer"
+            ),
+        },
+    ]
+
     # test state get/update methods with interrupt_after
 
     app_w_interrupt = workflow.compile(