diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 2071caedc..7947544cb 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -514,6 +514,14 @@ class CompiledGraph(Pregel): self.nodes[end].triggers.append(channel_name) cast(list[str], self.nodes[end].channels).append(channel_name) + async def aget_graph( + self, + config: Optional[RunnableConfig] = None, + *, + xray: Union[int, bool] = False, + ) -> DrawableGraph: + return self.get_graph(config, xray=xray) + def get_graph( self, config: Optional[RunnableConfig] = None, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f894c6420..c6018cbc7 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -25,7 +25,6 @@ from uuid import UUID, uuid5 from langchain_core.globals import get_debug from langchain_core.runnables import ( - Runnable, RunnableSequence, ) from langchain_core.runnables.base import Input, Output @@ -34,6 +33,7 @@ from langchain_core.runnables.config import ( get_async_callback_manager_for_config, get_callback_manager_for_config, ) +from langchain_core.runnables.graph import Graph from langchain_core.runnables.utils import ( ConfigurableFieldSpec, get_unique_config_specs, @@ -180,9 +180,7 @@ class Channel: ) -class Pregel( - Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], PregelProtocol -): +class Pregel(PregelProtocol): nodes: dict[str, PregelNode] channels: dict[str, Union[BaseChannel, ManagedValueSpec]] @@ -262,6 +260,16 @@ class Pregel( if auto_validate: self.validate() + def get_graph( + self, config: RunnableConfig | None = None, *, xray: int | bool = False + ) -> Graph: + raise NotImplementedError + + async def aget_graph( + self, config: RunnableConfig | None = None, *, xray: int | bool = False + ) -> Graph: + raise NotImplementedError + def copy(self, update: dict[str, Any] | None = None) -> Self: attrs = {**self.__dict__, **(update or {})} return self.__class__(**attrs) diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index 34789284b..ac046e949 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -1,27 +1,29 @@ +from abc import ABC, abstractmethod from typing import ( Any, AsyncIterator, Iterator, Optional, - Protocol, Sequence, Union, - runtime_checkable, ) -from langchain_core.runnables import RunnableConfig +from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.graph import Graph as DrawableGraph from typing_extensions import Self from langgraph.pregel.types import All, StateSnapshot, StreamMode -@runtime_checkable -class PregelProtocol(Protocol): +class PregelProtocol( + Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], ABC +): + @abstractmethod def with_config( self, config: Optional[RunnableConfig] = None, **kwargs: Any ) -> Self: ... + @abstractmethod def get_graph( self, config: Optional[RunnableConfig] = None, @@ -29,6 +31,7 @@ class PregelProtocol(Protocol): xray: Union[int, bool] = False, ) -> DrawableGraph: ... + @abstractmethod async def aget_graph( self, config: Optional[RunnableConfig] = None, @@ -36,22 +39,17 @@ class PregelProtocol(Protocol): xray: Union[int, bool] = False, ) -> DrawableGraph: ... - def get_subgraphs( - self, namespace: Optional[str] = None, recurse: bool = False - ) -> Iterator[tuple[str, "PregelProtocol"]]: ... - - def aget_subgraphs( - self, namespace: Optional[str] = None, recurse: bool = False - ) -> AsyncIterator[tuple[str, "PregelProtocol"]]: ... - + @abstractmethod def get_state( self, config: RunnableConfig, *, subgraphs: bool = False ) -> StateSnapshot: ... + @abstractmethod async def aget_state( self, config: RunnableConfig, *, subgraphs: bool = False ) -> StateSnapshot: ... + @abstractmethod def get_state_history( self, config: RunnableConfig, @@ -61,6 +59,7 @@ class PregelProtocol(Protocol): limit: Optional[int] = None, ) -> Iterator[StateSnapshot]: ... + @abstractmethod def aget_state_history( self, config: RunnableConfig, @@ -70,6 +69,7 @@ class PregelProtocol(Protocol): limit: Optional[int] = None, ) -> AsyncIterator[StateSnapshot]: ... + @abstractmethod def update_state( self, config: RunnableConfig, @@ -77,6 +77,7 @@ class PregelProtocol(Protocol): as_node: Optional[str] = None, ) -> RunnableConfig: ... + @abstractmethod async def aupdate_state( self, config: RunnableConfig, @@ -84,6 +85,7 @@ class PregelProtocol(Protocol): as_node: Optional[str] = None, ) -> RunnableConfig: ... + @abstractmethod def stream( self, input: Union[dict[str, Any], Any], @@ -95,6 +97,7 @@ class PregelProtocol(Protocol): subgraphs: bool = False, ) -> Iterator[Union[dict[str, Any], Any]]: ... + @abstractmethod def astream( self, input: Union[dict[str, Any], Any], @@ -106,6 +109,7 @@ class PregelProtocol(Protocol): subgraphs: bool = False, ) -> AsyncIterator[Union[dict[str, Any], Any]]: ... + @abstractmethod def invoke( self, input: Union[dict[str, Any], Any], @@ -115,6 +119,7 @@ class PregelProtocol(Protocol): interrupt_after: Optional[Union[All, Sequence[str]]] = None, ) -> Union[dict[str, Any], Any]: ... + @abstractmethod async def ainvoke( self, input: Union[dict[str, Any], Any], diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index dfd8ccd11..80295dae5 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -9,7 +9,7 @@ from typing import ( ) import orjson -from langchain_core.runnables import Runnable, RunnableConfig +from langchain_core.runnables import RunnableConfig from langchain_core.runnables.graph import ( Edge as DrawableEdge, ) @@ -44,7 +44,9 @@ class RemoteException(Exception): pass -class RemoteGraph(PregelProtocol, Runnable): +class RemoteGraph(PregelProtocol): + name: str + def __init__( self, name: str, # graph_id @@ -124,30 +126,6 @@ class RemoteGraph(PregelProtocol, Runnable): edges=[DrawableEdge(**edge) for edge in graph["edges"]], ) - def get_subgraphs( - self, namespace: Optional[str] = None, recurse: bool = False - ) -> Iterator[tuple[str, "PregelProtocol"]]: - subgraphs = self.sync_client.assistants.get_subgraphs( - assistant_id=self.name, - namespace=namespace, - recurse=recurse, - ) - for namespace, graph_schema in subgraphs.items(): - 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.name, - namespace=namespace, - recurse=recurse, - ) - for namespace, graph_schema in subgraphs.items(): - remote_subgraph = self.copy({"name": graph_schema["graph_id"]}) - yield (namespace, remote_subgraph) - def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot: tasks = [] for task in state["tasks"]: diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index dfd4ad57c..8ee0a3f9a 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/libs/langgraph/tests/test_remote_graph.py @@ -135,90 +135,6 @@ async def test_aget_graph(): ] -def test_get_subgraphs(): - # set up test - mock_sync_client = MagicMock() - mock_sync_client.assistants.get_subgraphs.return_value = { - "namespace_1": { - "graph_id": "test_graph_id_2", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - "namespace_2": { - "graph_id": "test_graph_id_3", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - } - - remote_pregel = RemoteGraph("test_graph_id_1", sync_client=mock_sync_client) - - # call method / assertions - subgraphs = list(remote_pregel.get_subgraphs()) - assert len(subgraphs) == 2 - - subgraph_1 = subgraphs[0] - ns_1 = subgraph_1[0] - remote_pregel_1: RemoteGraph = subgraph_1[1] - assert ns_1 == "namespace_1" - assert remote_pregel_1.graph_id == "test_graph_id_2" - - subgraph_2 = subgraphs[1] - ns_2 = subgraph_2[0] - remote_pregel_2: RemoteGraph = subgraph_2[1] - assert ns_2 == "namespace_2" - assert remote_pregel_2.graph_id == "test_graph_id_3" - - -@pytest.mark.anyio -async def test_aget_subgraphs(): - # set up test - mock_async_client = AsyncMock() - mock_async_client.assistants.get_subgraphs.return_value = { - "namespace_1": { - "graph_id": "test_graph_id_2", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - "namespace_2": { - "graph_id": "test_graph_id_3", - "input_schema": {}, - "output_schema": {}, - "state_schema": {}, - "config_schema": {}, - }, - } - - remote_pregel = RemoteGraph( - "test_graph_id_1", - client=mock_async_client, - ) - - # call method / assertions - subgraphs = [] - async for subgraph in remote_pregel.aget_subgraphs(): - subgraphs.append(subgraph) - assert len(subgraphs) == 2 - - subgraph_1 = subgraphs[0] - ns_1 = subgraph_1[0] - remote_pregel_1: RemoteGraph = subgraph_1[1] - assert ns_1 == "namespace_1" - assert remote_pregel_1.graph_id == "test_graph_id_2" - - subgraph_2 = subgraphs[1] - ns_2 = subgraph_2[0] - remote_pregel_2: RemoteGraph = subgraph_2[1] - assert ns_2 == "namespace_2" - assert remote_pregel_2.graph_id == "test_graph_id_3" - - def test_get_state(): # set up test mock_sync_client = MagicMock() @@ -866,9 +782,3 @@ async def test_langgraph_cloud_integration(): remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID graph = await remote_pregel.aget_graph(xray=True) print("graph:", graph) - - # test get subgraphs - remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID - async for name, pregel in remote_pregel.aget_subgraphs(): - print("name:", name) - print("pregel:", pregel)