diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 71ad967e5..53e77557e 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/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") diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index ff4223158..f1ef888ec 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -1392,6 +1392,7 @@ down --> __end__; side --> down; up --> down; + up --> other; up --> side; ''' diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 684c2df92..e025efac0 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d7303eb28..989d2a1e5 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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"}}, ]