diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 2f1c0c41e..4b7e87421 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -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( diff --git a/libs/langgraph/tests/test_remote_pregel.py b/libs/langgraph/tests/test_remote_graph.py similarity index 85% rename from libs/langgraph/tests/test_remote_pregel.py rename to libs/langgraph/tests/test_remote_graph.py index f2ca53583..dd07d27bb 100644 --- a/libs/langgraph/tests/test_remote_pregel.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -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)