Merge pull request #1265 from langchain-ai/nc/7aug/stream-nodes-wout-writes

Include in stream output all nodes ran, even if they returned no writes to state keys

libs/langgraph/langgraph/pregel/io.py

@@ -1,4 +1,3 @@
-from collections import defaultdict
 from typing import Any, Iterator, Mapping, Optional, Sequence, TypeVar, Union
 
 from langchain_core.runnables.utils import AddableDict
@@ -102,36 +101,33 @@ def map_output_updates(
     output_tasks = [
         t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags")
     ]
+    if not output_tasks:
+        return
     if isinstance(output_channels, str):
-        if updated := [
+        updated = [
             (task.name, value)
             for task in output_tasks
             for chan, value in task.writes
             if chan == output_channels
-        ]:
-            grouped = defaultdict(list)
-            for node, value in updated:
-                grouped[node].append(value)
-            for node, value in grouped.items():
-                if len(value) == 1:
-                    grouped[node] = value[0]
-            yield AddableUpdatesDict(grouped)
+        ]
     else:
-        if updated := [
+        updated = [
             (
                 task.name,
                 {chan: value for chan, value in task.writes if chan in output_channels},
             )
             for task in output_tasks
             if any(chan in output_channels for chan, _ in task.writes)
-        ]:
-            grouped = defaultdict(list)
-            for node, value in updated:
-                grouped[node].append(value)
-            for node, value in grouped.items():
-                if len(value) == 1:
-                    grouped[node] = value[0]
-            yield AddableUpdatesDict(grouped)
+        ]
+    grouped = {t.name: [] for t in output_tasks}
+    for node, value in updated:
+        grouped[node].append(value)
+    for node, value in grouped.items():
+        if len(value) == 0:
+            grouped[node] = None
+        if len(value) == 1:
+            grouped[node] = value[0]
+    yield AddableUpdatesDict(grouped)
 
 
 T = TypeVar("T")

libs/langgraph/tests/__snapshots__/test_pregel.ambr

@@ -1392,6 +1392,7 @@
   	down --> __end__;
   	side --> down;
   	up --> down;
+  	up --> other;
   	up --> side;
   
   '''

libs/langgraph/tests/test_pregel.py

@@ -330,7 +330,9 @@ def test_node_schemas_custom_output() -> None:
             }
         )
     ] == [
+        {"a": None},
         {"b": {"hello": "again", "now": 123}},
+        {"c": None},
     ]
 
 
@@ -959,6 +961,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
             {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates"
         )
     ] == [
+        {"one": None},
         {"two": 13},
         {"two": 4},
     ]
@@ -7265,6 +7268,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
         {
             "analyzer_one": {
@@ -7273,6 +7277,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         },
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
 
@@ -7400,6 +7405,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
         {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
         {
             "analyzer_one": {
@@ -7408,6 +7414,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> None:
         },
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
 
@@ -7422,6 +7429,9 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
     def side(state: State):
         pass
 
+    def other(state: State):
+        return {"my_key": "_more"}
+
     def down(state: State):
         pass
 
@@ -7429,17 +7439,25 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
 
     graph.add_node("up", up)
     graph.add_node("side", side)
+    graph.add_node("other", other)
     graph.add_node("down", down)
 
     graph.set_entry_point("up")
     graph.add_edge("up", "side")
+    graph.add_edge("up", "other")
     graph.add_edge(["up", "side"], "down")
     graph.set_finish_point("down")
 
     app = graph.compile()
 
     assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
-    assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
+    assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value_more"}
+    assert [*app.stream({"my_key": "my_value"})] == [
+        {"up": None},
+        {"side": None},
+        {"other": {"my_key": "_more"}},
+        {"down": None},
+    ]
 
 
 def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:

libs/langgraph/tests/test_pregel_async.py

@@ -491,7 +491,9 @@ async def test_node_schemas_custom_output() -> None:
             }
         )
     ] == [
+        {"a": None},
         {"b": {"hello": "again", "now": 123}},
+        {"c": None},
     ]
 
 
@@ -1078,6 +1080,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
             {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates"
         )
     ] == [
+        {"one": None},
         {"two": 13},
         {"two": 4},
     ]
@@ -5878,6 +5881,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
         {
             "analyzer_one": {
@@ -5886,6 +5890,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple() -> None:
         },
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]
 
@@ -5962,6 +5967,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N
         {"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"rewrite_query": {"query": "query: analyzed: query: what is weather in sf"}},
         {
             "analyzer_one": {
@@ -5970,6 +5976,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple_cond_edge() -> N
         },
         {"retriever_two": {"docs": ["doc3", "doc4"]}},
         {"retriever_one": {"docs": ["doc1", "doc2"]}},
+        {"decider": None},
         {"qa": {"answer": "doc1,doc1,doc2,doc2,doc3,doc3,doc4,doc4"}},
     ]