WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-27 12:04:58 +02:00
Code Issues Packages Projects Releases Wiki Activity

Update stream() and astream() to process 'updates' event types.

Browse Source
This commit is contained in:
Andrew Nguonly
2024-10-11 18:40:40 -07:00
parent 0557fb03a4
commit 2f819a6a9b
2 changed files with 127 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/remote.py
+52 -3
View File
@@ -26,10 +26,12 @@ from langgraph_sdk.client import (
get_client,
get_sync_client,
)
from langgraph_sdk.schema import Checkpoint, ThreadState
from langgraph_sdk.schema import Checkpoint, StreamPart, ThreadState
from typing_extensions import Self
from langgraph.checkpoint.base import CheckpointMetadata
from langgraph.constants import INTERRUPT
from langgraph.errors import GraphInterrupt
from langgraph.pregel.protocol import PregelProtocol
from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
from langgraph.types import Interrupt
@@ -345,6 +347,35 @@ class RemoteGraph(PregelProtocol, Runnable):
)
return self._get_config(response["checkpoint"])
def _get_stream_modes(
self,
stream_mode: Optional[Union[StreamMode, list[StreamMode]]],
) -> tuple[list[StreamMode], bool]:
"""Return a tuple of the final list of stream modes sent to the
remote graph and a boolean flag indicating if stream mode 'updates'
was present in the original list of stream modes.
'updates' mode is added to the list of stream modes so that interrupts
can be detected in the remote graph.
"""
updated_stream_modes = []
updates_mode = False
if stream_mode:
if isinstance(stream_mode, str):
updated_stream_modes.append(stream_mode)
else:
updated_stream_modes.extend(stream_mode)
if "updates" in updated_stream_modes:
updates_mode = True
else:
updated_stream_modes.append("updates")
else:
updated_stream_modes.extend(["values", "updates"])
return (updated_stream_modes, updates_mode)
def stream(
self,
input: Union[dict[str, Any], Any],
@@ -357,17 +388,26 @@ class RemoteGraph(PregelProtocol, Runnable):
) -> Iterator[Union[dict[str, Any], Any]]:
merged_config = merge_configs(self.config, config)
sanitized_config = self._sanitize_config(merged_config)
updated_stream_modes, include_updates = self._get_stream_modes(stream_mode)
for chunk in self.sync_client.runs.stream(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
input=input,
config=sanitized_config,
stream_mode=stream_mode, # type: ignore
stream_mode=updated_stream_modes, # type: ignore
interrupt_before=interrupt_before, # type: ignore
interrupt_after=interrupt_after, # type: ignore
stream_subgraphs=subgraphs,
):
if chunk.event == INTERRUPT:
raise GraphInterrupt()
# Don't emit 'updates' events if the original list of stream modes
# didn't include it.
if chunk.event == "updates" and not include_updates:
continue
yield chunk
async def astream(
@@ -382,17 +422,26 @@ class RemoteGraph(PregelProtocol, Runnable):
) -> AsyncIterator[Union[dict[str, Any], Any]]:
merged_config = merge_configs(self.config, config)
sanitized_config = self._sanitize_config(merged_config)
updated_stream_modes, include_updates = self._get_stream_modes(stream_mode)
async for chunk in self.client.runs.stream(
thread_id=sanitized_config["configurable"]["thread_id"],
assistant_id=self.graph_id,
input=input,
config=sanitized_config,
stream_mode=stream_mode if stream_mode else "values", # type: ignore
stream_mode=updated_stream_modes, # type: ignore
interrupt_before=interrupt_before, # type: ignore
interrupt_after=interrupt_after, # type: ignore
stream_subgraphs=subgraphs,
):
if chunk.event == INTERRUPT:
raise GraphInterrupt()
# Don't emit 'updates' events if the original list of stream modes
# didn't include it.
if chunk.event == "updates" and not include_updates:
continue
yield chunk
async def astream_events(
libs/langgraph/tests/test_remote_pregel.py → libs/langgraph/tests/test_remote_graph.py
+75 -16
View File
@@ -7,7 +7,9 @@ from langchain_core.runnables.graph import (
from langchain_core.runnables.graph import (
Node as DrawableNode,
)
from langgraph_sdk.schema import StreamPart
from langgraph.errors import GraphInterrupt
from langgraph.pregel.remote import RemoteGraph
from langgraph.pregel.types import StateSnapshot
@@ -473,17 +475,46 @@ def test_stream():
# set up test
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = [
{"chunk": "data1"},
{"chunk": "data2"},
{"chunk": "data3"},
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="__interrupt__", data={}),
]
# call method / assertions
remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id")
config = {"configurable": {"thread_id": "thread_1"}}
result = list(remote_pregel.stream({"input": "data"}, config))
assert result == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
# stream modes doesn't include 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"}, config={"configurable": {"thread_id": "thread_1"}}
):
stream_parts.append(stream_part)
assert stream_parts == [
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
]
# stream modes includes 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
):
stream_parts.append(stream_part)
assert stream_parts == [
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
]
@pytest.mark.anyio
@@ -492,20 +523,47 @@ async def test_astream():
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
{"chunk": "data1"},
{"chunk": "data2"},
{"chunk": "data3"},
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(event="__interrupt__", data={}),
]
mock_async_client.runs.stream.return_value = async_iter
# call method / assertions
remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id")
config = {"configurable": {"thread_id": "thread_1"}}
chunks = []
async for chunk in remote_pregel.astream({"input": "data"}, config):
chunks.append(chunk)
assert chunks == [{"chunk": "data1"}, {"chunk": "data2"}, {"chunk": "data3"}]
# stream modes doesn't include 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"}, config={"configurable": {"thread_id": "thread_1"}}
):
stream_parts.append(stream_part)
assert stream_parts == [
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
]
# stream modes includes 'updates'
stream_parts = []
with pytest.raises(GraphInterrupt):
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode=["updates"],
):
stream_parts.append(stream_part)
assert stream_parts == [
StreamPart(event="values", data={"chunk": "data1"}),
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
]
def test_invoke():
@@ -572,7 +630,7 @@ async def test_langgraph_cloud_integration():
"messages": [
{
"role": "human",
"content": "Hello world!",
"content": "What's the weather in SF?",
}
]
}
@@ -580,7 +638,8 @@ async def test_langgraph_cloud_integration():
# test invoke
response = app.invoke(
input,
config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}},
config={"configurable": {"thread_id": "39a6104a-34e7-4f83-929c-d9eb163003c9"}},
interrupt_before=["agent"],
)
print("response:", response["messages"][-1].content)