WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-26 01:22:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

In stream_mode=values yield all values, instead of only changed ones

Browse Source
This commit is contained in:
Nuno Campos
2024-04-11 11:12:10 -07:00
parent 9ffd6d9c78
commit a8af7467b8
4 changed files with 55 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+4 -6
View File
@@ -971,9 +971,8 @@ class Pregel(
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
output_keys = output_keys if output_keys is not None else self.output_channels
output_is_dict = not isinstance(output_keys, str)
if stream_mode == "values":
latest: Union[dict[str, Any], Any] = {} if output_is_dict else None
latest: Union[dict[str, Any], Any] = None
else:
chunks = []
for chunk in self.stream(
@@ -988,7 +987,7 @@ class Pregel(
**kwargs,
):
if stream_mode == "values":
latest = {**latest, **chunk} if output_is_dict else chunk
latest = chunk
else:
chunks.append(chunk)
if stream_mode == "values":
@@ -1010,9 +1009,8 @@ class Pregel(
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
output_keys = output_keys if output_keys is not None else self.output_channels
output_is_dict = not isinstance(output_keys, str)
if stream_mode == "values":
latest: Union[dict[str, Any], Any] = {} if output_is_dict else None
latest: Union[dict[str, Any], Any] = None
else:
chunks = []
async for chunk in self.astream(
@@ -1027,7 +1025,7 @@ class Pregel(
**kwargs,
):
if stream_mode == "values":
latest = {**latest, **chunk} if output_is_dict else chunk
latest = chunk
else:
chunks.append(chunk)
if stream_mode == "values":
langgraph/pregel/io.py
+2 -2
View File
@@ -71,8 +71,8 @@ def map_output_values(
if any(chan == output_channels for chan, _ in pending_writes):
yield read_channel(channels, output_channels)
else:
if updated := {c for c, _ in pending_writes if c in output_channels}:
yield read_channels(channels, updated)
if {c for c, _ in pending_writes if c in output_channels}:
yield read_channels(channels, output_channels)
def map_output_updates(
tests/test_pregel.py
+21 -8
View File
@@ -272,6 +272,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
nodes={"one": one, "two": two},
channels={"inbox": Topic(int)},
input_channels=["input", "inbox"],
stream_channels=["output", "inbox"],
)
# [12 + 1, 2 + 1 + 1]
@@ -294,7 +295,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
]
assert [*app.stream({"input": 2, "inbox": 12})] == [
{"inbox": [3], "output": 13},
{"output": 4},
{"inbox": [], "output": 4},
]
@@ -637,13 +638,16 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
)
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(nodes={"one": one, "two": two})
app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"])
assert [c for c in app.stream(2, stream_mode="updates")] == [
{"one": {"between": 3, "output": 3}},
{"two": {"output": 4}},
]
assert [c for c in app.stream(2)] == [{"between": 3, "output": 3}, {"output": 4}]
assert [c for c in app.stream(2)] == [
{"between": 3, "output": 3},
{"between": 3, "output": 4},
]
def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
@@ -695,6 +699,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
"ctx": Context(an_int, typ=int),
},
output_channels=["inbox", "output"],
stream_channels=["inbox", "output"],
)
assert setup.call_count == 0
@@ -705,7 +710,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
assert chunk == {"output": 4}
assert chunk == {"inbox": [], "output": 4}
else:
assert False, "Expected only two chunks"
assert cleanup.call_count == 1, "Expected cleanup to be called once"
@@ -1825,6 +1830,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None
assert app.invoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"output": "what is weather in sf->right",
"steps": [],
}
assert [*app.stream({"input": "what is weather in sf"})] == [
@@ -2887,10 +2893,17 @@ def test_in_one_fan_out_out_one_graph_state() -> None:
]
assert [*app.stream({"query": "what is weather in sf"}, stream_mode="values")] == [
{"query": "what is weather in sf"},
{"query": "query: what is weather in sf"},
{"docs": ["doc1", "doc2", "doc3", "doc4"]},
{"answer": "doc1,doc2,doc3,doc4"},
{"query": "what is weather in sf", "docs": []},
{"query": "query: what is weather in sf", "docs": []},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
},
]
tests/test_pregel_async.py
+28 -5
View File
@@ -164,7 +164,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output")
app = Pregel(nodes={"one": one, "two": two})
app = Pregel(nodes={"one": one, "two": two}, stream_channels=["inbox", "output"])
assert await app.ainvoke(2) == 4
@@ -182,6 +182,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
}
elif step == 2:
assert values == {
"inbox": 3,
"output": 4,
}
assert step == 2
@@ -280,6 +281,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
nodes={"one": one, "two": two},
channels={"inbox": Topic(int)},
input_channels=["input", "inbox"],
stream_channels=["inbox", "output"],
)
# [12 + 1, 2 + 1 + 1]
@@ -304,7 +306,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
]
assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [
{"inbox": [3], "output": 13},
{"output": 4},
{"inbox": [], "output": 4},
]
@@ -669,12 +671,12 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
)
two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(nodes={"one": one, "two": two})
app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"])
# Then invoke pubsub
assert [c async for c in app.astream(2)] == [
{"between": 3, "output": 3},
{"output": 4},
{"between": 3, "output": 4},
]
@@ -727,6 +729,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
"ctx": Context(an_int, an_int_async, typ=int),
},
output_channels=["inbox", "output"],
stream_channels=["inbox", "output"],
)
async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]:
@@ -747,7 +750,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
if i == 0:
assert chunk == {"inbox": [3]}
elif i == 1:
assert chunk == {"output": 4}
assert chunk == {"inbox": [], "output": 4}
else:
assert False, "Expected only two chunks"
assert setup_sync.call_count == 0
@@ -1956,6 +1959,7 @@ async def test_conditional_entrypoint_graph_state() -> None:
assert await app.ainvoke({"input": "what is weather in sf"}) == {
"input": "what is weather in sf",
"output": "what is weather in sf->right",
"steps": [],
}
assert [c async for c in app.astream({"input": "what is weather in sf"})] == [
@@ -2653,6 +2657,25 @@ async def test_in_one_fan_out_out_one_graph_state() -> None:
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
assert [
c
async for c in app.astream(
{"query": "what is weather in sf"}, stream_mode="values"
)
] == [
{"query": "what is weather in sf", "docs": []},
{"query": "query: what is weather in sf", "docs": []},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
},
{
"query": "query: what is weather in sf",
"docs": ["doc1", "doc2", "doc3", "doc4"],
"answer": "doc1,doc2,doc3,doc4",
},
]
@pytest.mark.parametrize(
"checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]