diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 7c0923751..2071caedc 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -170,7 +170,7 @@ class Graph: def add_node( self, node: Union[str, RunnableLike], - action: Optional[RunnableLike] = None, + action: Optional[Union[RunnableLike]] = None, *, metadata: Optional[dict[str, Any]] = None, ) -> Self: diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index b535a06e1..f894c6420 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -86,6 +86,7 @@ from langgraph.pregel.io import read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager from langgraph.pregel.messages import StreamMessagesHandler +from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner @@ -179,7 +180,9 @@ class Channel: ) -class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): +class Pregel( + Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], PregelProtocol +): nodes: dict[str, PregelNode] channels: dict[str, Union[BaseChannel, ManagedValueSpec]] diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index f0001659b..dfd8ccd11 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -47,7 +47,9 @@ class RemoteException(Exception): class RemoteGraph(PregelProtocol, Runnable): def __init__( self, - graph_id: str, + name: str, # graph_id + /, + *, url: Optional[str] = None, api_key: Optional[str] = None, headers: Optional[dict[str, str]] = None, @@ -60,7 +62,7 @@ class RemoteGraph(PregelProtocol, Runnable): If `client` or `sync_client` are provided, they will be used instead of the default clients. See `LangGraphClient` and `SyncLangGraphClient` for details on the default clients. """ - self.graph_id = graph_id + self.name = name self.config = config self.client = client or get_client(url=url, api_key=api_key, headers=headers) self.sync_client = sync_client or get_sync_client( @@ -69,7 +71,7 @@ class RemoteGraph(PregelProtocol, Runnable): def copy(self, update: dict[str, Any]) -> Self: attrs = {**self.__dict__, **update} - return self.__class__(**attrs) + return self.__class__(attrs.pop("name"), **attrs) def with_config( self, config: Optional[RunnableConfig] = None, **kwargs: Any @@ -99,7 +101,7 @@ class RemoteGraph(PregelProtocol, Runnable): xray: Union[int, bool] = False, ) -> DrawableGraph: graph = self.sync_client.assistants.get_graph( - assistant_id=self.graph_id, + assistant_id=self.name, xray=xray, ) return DrawableGraph( @@ -114,7 +116,7 @@ class RemoteGraph(PregelProtocol, Runnable): xray: Union[int, bool] = False, ) -> DrawableGraph: graph = await self.client.assistants.get_graph( - assistant_id=self.graph_id, + assistant_id=self.name, xray=xray, ) return DrawableGraph( @@ -126,24 +128,24 @@ class RemoteGraph(PregelProtocol, Runnable): self, namespace: Optional[str] = None, recurse: bool = False ) -> Iterator[tuple[str, "PregelProtocol"]]: subgraphs = self.sync_client.assistants.get_subgraphs( - assistant_id=self.graph_id, + assistant_id=self.name, namespace=namespace, recurse=recurse, ) for namespace, graph_schema in subgraphs.items(): - remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]}) + remote_subgraph = self.copy({"name": graph_schema["graph_id"]}) yield (namespace, remote_subgraph) async def aget_subgraphs( self, namespace: Optional[str] = None, recurse: bool = False ) -> AsyncIterator[tuple[str, "PregelProtocol"]]: subgraphs = await self.client.assistants.get_subgraphs( - assistant_id=self.graph_id, + assistant_id=self.name, namespace=namespace, recurse=recurse, ) for namespace, graph_schema in subgraphs.items(): - remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]}) + remote_subgraph = self.copy({"name": graph_schema["graph_id"]}) yield (namespace, remote_subgraph) def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot: @@ -403,7 +405,7 @@ class RemoteGraph(PregelProtocol, Runnable): for chunk in self.sync_client.runs.stream( thread_id=cast(str, sanitized_config["configurable"]["thread_id"]), - assistant_id=self.graph_id, + assistant_id=self.name, input=input, config=sanitized_config, stream_mode=stream_modes, @@ -450,7 +452,7 @@ class RemoteGraph(PregelProtocol, Runnable): async for chunk in self.client.runs.stream( thread_id=sanitized_config["configurable"]["thread_id"], - assistant_id=self.graph_id, + assistant_id=self.name, input=input, config=sanitized_config, stream_mode=stream_modes, @@ -481,6 +483,22 @@ class RemoteGraph(PregelProtocol, Runnable): else: yield chunk + async def astream_events( + self, + input: Any, + config: RunnableConfig | None = None, + *, + version: All | All, + include_names: Sequence[All] | None = None, + include_types: Sequence[All] | None = None, + include_tags: Sequence[All] | None = None, + exclude_names: Sequence[All] | None = None, + exclude_types: Sequence[All] | None = None, + exclude_tags: Sequence[All] | None = None, + **kwargs: Any, + ) -> AsyncIterator[dict[str, Any]]: + raise NotImplementedError + def invoke( self, input: Union[dict[str, Any], Any], @@ -494,7 +512,7 @@ class RemoteGraph(PregelProtocol, Runnable): return self.sync_client.runs.wait( thread_id=sanitized_config["configurable"]["thread_id"], - assistant_id=self.graph_id, + assistant_id=self.name, input=input, config=sanitized_config, interrupt_before=interrupt_before, @@ -515,7 +533,7 @@ class RemoteGraph(PregelProtocol, Runnable): return await self.client.runs.wait( thread_id=sanitized_config["configurable"]["thread_id"], - assistant_id=self.graph_id, + assistant_id=self.name, input=input, config=sanitized_config, interrupt_before=interrupt_before, diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 2b09f8f75..cc7221ea4 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -4,6 +4,7 @@ from langchain_core.runnables import RunnableLambda, RunnableSequence from langchain_core.runnables.utils import get_function_nonlocals from langgraph.checkpoint.base import ChannelVersions +from langgraph.pregel.protocol import PregelProtocol from langgraph.utils.runnable import Runnable, RunnableCallable, RunnableSeq @@ -32,9 +33,9 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]: for c in candidates: if ( - isinstance(c, Pregel) + isinstance(c, PregelProtocol) # subgraphs that disabled checkpointing are not considered - and c.checkpointer is not False + and (not isinstance(c, Pregel) or c.checkpointer is not False) ): return c elif isinstance(c, RunnableSequence) or isinstance(c, RunnableSeq): diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 46b3ab9a4..dfd4ad57c 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -17,7 +17,7 @@ from langgraph.pregel.types import StateSnapshot def test_with_config(): # set up test remote_pregel = RemoteGraph( - graph_id="test_graph_id", + "test_graph_id", config={ "configurable": { "foo": "bar", @@ -64,7 +64,7 @@ def test_get_graph(): ], } - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client) # call method / assertions drawable_graph = remote_pregel.get_graph() @@ -111,7 +111,7 @@ async def test_aget_graph(): ], } - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client) # call method / assertions drawable_graph = await remote_pregel.aget_graph() @@ -155,9 +155,7 @@ def test_get_subgraphs(): }, } - remote_pregel = RemoteGraph( - sync_client=mock_sync_client, graph_id="test_graph_id_1" - ) + remote_pregel = RemoteGraph("test_graph_id_1", sync_client=mock_sync_client) # call method / assertions subgraphs = list(remote_pregel.get_subgraphs()) @@ -198,8 +196,8 @@ async def test_aget_subgraphs(): } remote_pregel = RemoteGraph( + "test_graph_id_1", client=mock_async_client, - graph_id="test_graph_id_1", ) # call method / assertions @@ -240,7 +238,10 @@ def test_get_state(): } # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_snapshot = remote_pregel.get_state(config) @@ -287,7 +288,10 @@ async def test_aget_state(): } # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_snapshot = await remote_pregel.aget_state(config) @@ -338,7 +342,10 @@ def test_get_state_history(): ] # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_history_snapshot = list( @@ -386,7 +393,10 @@ async def test_aget_state_history(): ] # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread1"}} state_history_snapshot = [] @@ -427,7 +437,10 @@ def test_update_state(): } # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread1"}} response = remote_pregel.update_state(config, {"key": "value"}) @@ -456,7 +469,10 @@ async def test_aupdate_state(): } # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread1"}} response = await remote_pregel.aupdate_state(config, {"key": "value"}) @@ -483,7 +499,10 @@ def test_stream(): ] # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) # stream modes doesn't include 'updates' stream_parts = [] @@ -583,7 +602,10 @@ async def test_astream(): mock_async_client.runs.stream.return_value = async_iter # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) # stream modes doesn't include 'updates' stream_parts = [] @@ -717,7 +739,10 @@ def test_invoke(): } # call method / assertions - remote_pregel = RemoteGraph(sync_client=mock_sync_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + sync_client=mock_sync_client, + ) config = {"configurable": {"thread_id": "thread_1"}} result = remote_pregel.invoke( @@ -736,7 +761,10 @@ async def test_ainvoke(): } # call method / assertions - remote_pregel = RemoteGraph(client=mock_async_client, graph_id="test_graph_id") + remote_pregel = RemoteGraph( + "test_graph_id", + client=mock_async_client, + ) config = {"configurable": {"thread_id": "thread_1"}} result = await remote_pregel.ainvoke( @@ -758,7 +786,9 @@ async def test_langgraph_cloud_integration(): client = get_client() sync_client = get_sync_client() remote_pregel = RemoteGraph( - client=client, sync_client=sync_client, graph_id="agent" + "agent", + client=client, + sync_client=sync_client, ) # define graph