From e72c25873f00a55f46b6379bc45d0f0356361937 Mon Sep 17 00:00:00 2001 From: Andrew Nguonly Date: Fri, 11 Oct 2024 09:23:12 -0700 Subject: [PATCH] Implement `PregelProtocol` and `RemotePregel` class (attempt 2) (#2078) ### Summary Redo of [this PR](https://github.com/langchain-ai/langgraph/pull/2034) (branched from clean branch). --- libs/langgraph/langgraph/pregel/protocol.py | 125 ++++ libs/langgraph/langgraph/pregel/remote.py | 465 ++++++++++++++ libs/langgraph/poetry.lock | 29 +- libs/langgraph/pyproject.toml | 1 + libs/langgraph/tests/test_remote_pregel.py | 666 ++++++++++++++++++++ 5 files changed, 1285 insertions(+), 1 deletion(-) create mode 100644 libs/langgraph/langgraph/pregel/protocol.py create mode 100644 libs/langgraph/langgraph/pregel/remote.py create mode 100644 libs/langgraph/tests/test_remote_pregel.py diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py new file mode 100644 index 000000000..34789284b --- /dev/null +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -0,0 +1,125 @@ +from typing import ( + Any, + AsyncIterator, + Iterator, + Optional, + Protocol, + Sequence, + Union, + runtime_checkable, +) + +from langchain_core.runnables import 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): + def with_config( + self, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Self: ... + + def get_graph( + self, + config: Optional[RunnableConfig] = None, + *, + xray: Union[int, bool] = False, + ) -> DrawableGraph: ... + + async def aget_graph( + self, + config: Optional[RunnableConfig] = None, + *, + 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"]]: ... + + def get_state( + self, config: RunnableConfig, *, subgraphs: bool = False + ) -> StateSnapshot: ... + + async def aget_state( + self, config: RunnableConfig, *, subgraphs: bool = False + ) -> StateSnapshot: ... + + def get_state_history( + self, + config: RunnableConfig, + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[StateSnapshot]: ... + + def aget_state_history( + self, + config: RunnableConfig, + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> AsyncIterator[StateSnapshot]: ... + + def update_state( + self, + config: RunnableConfig, + values: Optional[Union[dict[str, Any], Any]], + as_node: Optional[str] = None, + ) -> RunnableConfig: ... + + async def aupdate_state( + self, + config: RunnableConfig, + values: Optional[Union[dict[str, Any], Any]], + as_node: Optional[str] = None, + ) -> RunnableConfig: ... + + def stream( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + subgraphs: bool = False, + ) -> Iterator[Union[dict[str, Any], Any]]: ... + + def astream( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + subgraphs: bool = False, + ) -> AsyncIterator[Union[dict[str, Any], Any]]: ... + + def invoke( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + ) -> Union[dict[str, Any], Any]: ... + + async def ainvoke( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + ) -> Union[dict[str, Any], Any]: ... diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py new file mode 100644 index 000000000..768727996 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -0,0 +1,465 @@ +from typing import ( + Any, + AsyncIterator, + Iterator, + Optional, + Sequence, + Union, + cast, +) + +import orjson +from langchain_core.runnables import Runnable, RunnableConfig +from langchain_core.runnables.graph import ( + Edge as DrawableEdge, +) +from langchain_core.runnables.graph import ( + Graph as DrawableGraph, +) +from langchain_core.runnables.graph import ( + Node as DrawableNode, +) +from langchain_core.runnables.schema import StandardStreamEvent, StreamEvent +from langgraph_sdk.client import ( + LangGraphClient, + SyncLangGraphClient, + get_client, + get_sync_client, +) +from langgraph_sdk.schema import Checkpoint, ThreadState +from typing_extensions import Self + +from langgraph.checkpoint.base import CheckpointMetadata +from langgraph.pregel.protocol import PregelProtocol +from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode +from langgraph.types import Interrupt +from langgraph.utils.config import merge_configs + + +class RemotePregel(PregelProtocol, Runnable): + def __init__( + self, + graph_id: str, + config: Optional[RunnableConfig] = None, + url: Optional[str] = None, + api_key: Optional[str] = None, + headers: Optional[dict[str, str]] = None, + client: Optional[LangGraphClient] = None, + sync_client: Optional[SyncLangGraphClient] = None, + ): + """Specify `url`, `api_key`, and/or `headers` to create default sync and async clients. + + 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.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( + url=url, api_key=api_key, headers=headers + ) + + def copy(self, update: dict[str, Any]) -> Self: + attrs = {**self.__dict__, **update} + return self.__class__(**attrs) + + def with_config( + self, config: Optional[RunnableConfig] = None, **kwargs: Any + ) -> Self: + return self.copy( + {"config": merge_configs(self.config, config, cast(RunnableConfig, kwargs))} + ) + + def _get_drawable_nodes( + self, graph: dict[str, list[dict[str, Any]]] + ) -> dict[str, DrawableNode]: + nodes = {} + for node in graph["nodes"]: + node_id = str(node["id"]) + nodes[node_id] = DrawableNode( + id=node_id, + name=node.get("name", ""), + data=node.get("data", {}), + metadata=node.get("metadata"), + ) + return nodes + + def get_graph( + self, + config: Optional[RunnableConfig] = None, + *, + xray: Union[int, bool] = False, + ) -> DrawableGraph: + graph = self.sync_client.assistants.get_graph( + assistant_id=self.graph_id, + xray=xray, + ) + return DrawableGraph( + nodes=self._get_drawable_nodes(graph), + edges=[DrawableEdge(**edge) for edge in graph["edges"]], + ) + + async def aget_graph( + self, + config: Optional[RunnableConfig] = None, + *, + xray: Union[int, bool] = False, + ) -> DrawableGraph: + graph = await self.client.assistants.get_graph( + assistant_id=self.graph_id, + xray=xray, + ) + return DrawableGraph( + nodes=self._get_drawable_nodes(graph), + 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.graph_id, + namespace=namespace, + recurse=recurse, + ) + for namespace, graph_schema in subgraphs.items(): + remote_subgraph = self.copy({"graph_id": 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, + namespace=namespace, + recurse=recurse, + ) + for namespace, graph_schema in subgraphs.items(): + remote_subgraph = self.copy({"graph_id": graph_schema["graph_id"]}) + yield (namespace, remote_subgraph) + + def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot: + tasks = [] + for task in state["tasks"]: + interrupts = [] + for interrupt in task["interrupts"]: + interrupts.append(Interrupt(**interrupt)) + + tasks.append( + PregelTask( + id=task["id"], + name=task["name"], + path=tuple(), + error=Exception(task["error"]) if task["error"] else None, + interrupts=tuple(interrupts), + state=self._create_state_snapshot(task["state"]) + if task["state"] + else None, + ) + ) + + return StateSnapshot( + values=state["values"], + next=tuple(state["next"]) if state["next"] else tuple(), + config={ + "configurable": { + "thread_id": state["checkpoint"]["thread_id"], + "checkpoint_ns": state["checkpoint"]["checkpoint_ns"], + "checkpoint_id": state["checkpoint"]["checkpoint_id"], + "checkpoint_map": state["checkpoint"].get("checkpoint_map", {}), + } + }, + metadata=CheckpointMetadata(**state["metadata"]), + created_at=state["created_at"], + parent_config={ + "configurable": { + "thread_id": state["parent_checkpoint"]["thread_id"], + "checkpoint_ns": state["parent_checkpoint"]["checkpoint_ns"], + "checkpoint_id": state["parent_checkpoint"]["checkpoint_id"], + "checkpoint_map": state["parent_checkpoint"].get( + "checkpoint_map", {} + ), + } + } + if state["parent_checkpoint"] + else None, + tasks=tuple(tasks), + ) + + def _get_checkpoint(self, config: Optional[RunnableConfig]) -> Optional[Checkpoint]: + if config is None: + return None + + checkpoint = {} + + if "thread_id" in config["configurable"]: + checkpoint["thread_id"] = config["configurable"]["thread_id"] + if "checkpoint_ns" in config["configurable"]: + checkpoint["checkpoint_ns"] = config["configurable"]["checkpoint_ns"] + if "checkpoint_id" in config["configurable"]: + checkpoint["checkpoint_id"] = config["configurable"]["checkpoint_id"] + if "checkpoint_map" in config["configurable"]: + checkpoint["checkpoint_map"] = config["configurable"]["checkpoint_map"] + + return checkpoint if checkpoint else None + + def _get_config(self, checkpoint: Checkpoint) -> RunnableConfig: + return { + "configurable": { + "thread_id": checkpoint["thread_id"], + "checkpoint_ns": checkpoint["checkpoint_ns"], + "checkpoint_id": checkpoint["checkpoint_id"], + "checkpoint_map": checkpoint.get("checkpoint_map", {}), + } + } + + def _sanitize_config(self, config: RunnableConfig) -> RunnableConfig: + reserved_configurable_keys = frozenset( + [ + "callbacks", + "checkpoint_map", + "checkpoint_id", + "checkpoint_ns", + ] + ) + + def _sanitize_obj(obj: Any) -> Any: + """Remove non-JSON serializable fields from the given object.""" + if isinstance(obj, dict): + return {k: _sanitize_obj(v) for k, v in obj.items()} + elif isinstance(obj, list): + return [_sanitize_obj(v) for v in obj] + else: + try: + orjson.dumps(obj) + return obj + except orjson.JSONEncodeError: + return None + + # Remove non-JSON serializable fields from the config. + config = _sanitize_obj(config) + + # Only include configurable keys that are not reserved and + # not starting with "__pregel_" prefix. + new_configurable = { + k: v + for k, v in config["configurable"].items() + if k not in reserved_configurable_keys and not k.startswith("__pregel_") + } + + return {"configurable": new_configurable} + + def get_state( + self, config: RunnableConfig, *, subgraphs: bool = False + ) -> StateSnapshot: + merged_config = merge_configs(self.config, config) + + state = self.sync_client.threads.get_state( + thread_id=merged_config["configurable"]["thread_id"], + checkpoint=self._get_checkpoint(merged_config), + subgraphs=subgraphs, + ) + return self._create_state_snapshot(state) + + async def aget_state( + self, config: RunnableConfig, *, subgraphs: bool = False + ) -> StateSnapshot: + merged_config = merge_configs(self.config, config) + + state = await self.client.threads.get_state( + thread_id=merged_config["configurable"]["thread_id"], + checkpoint=self._get_checkpoint(merged_config), + subgraphs=subgraphs, + ) + return self._create_state_snapshot(state) + + def get_state_history( + self, + config: RunnableConfig, + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[StateSnapshot]: + merged_config = merge_configs(self.config, config) + + states = self.sync_client.threads.get_history( + thread_id=merged_config["configurable"]["thread_id"], + limit=limit if limit else 10, + before=self._get_checkpoint(before), + metadata=filter, + checkpoint=self._get_checkpoint(merged_config), + ) + for state in states: + yield self._create_state_snapshot(state) + + async def aget_state_history( + self, + config: RunnableConfig, + *, + filter: Optional[dict[str, Any]] = None, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> AsyncIterator[StateSnapshot]: + merged_config = merge_configs(self.config, config) + + states = await self.client.threads.get_history( + thread_id=merged_config["configurable"]["thread_id"], + limit=limit if limit else 10, + before=self._get_checkpoint(before), + metadata=filter, + checkpoint=self._get_checkpoint(merged_config), + ) + for state in states: + yield self._create_state_snapshot(state) + + def update_state( + self, + config: RunnableConfig, + values: Optional[Union[dict[str, Any], Any]], + as_node: Optional[str] = None, + ) -> RunnableConfig: + merged_config = merge_configs(self.config, config) + + response: dict = self.sync_client.threads.update_state( # type: ignore + thread_id=merged_config["configurable"]["thread_id"], + values=values, # type: ignore + as_node=as_node, + checkpoint=self._get_checkpoint(merged_config), + ) + return self._get_config(response["checkpoint"]) + + async def aupdate_state( + self, + config: RunnableConfig, + values: Optional[Union[dict[str, Any], Any]], + as_node: Optional[str] = None, + ) -> RunnableConfig: + merged_config = merge_configs(self.config, config) + + response: dict = await self.client.threads.update_state( # type: ignore + thread_id=merged_config["configurable"]["thread_id"], + values=values, # type: ignore + as_node=as_node, + checkpoint=self._get_checkpoint(merged_config), + ) + return self._get_config(response["checkpoint"]) + + def stream( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + subgraphs: bool = False, + ) -> Iterator[Union[dict[str, Any], Any]]: + merged_config = merge_configs(self.config, config) + sanitized_config = self._sanitize_config(merged_config) + + 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 + interrupt_before=interrupt_before, # type: ignore + interrupt_after=interrupt_after, # type: ignore + stream_subgraphs=subgraphs, + ): + yield chunk + + async def astream( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + stream_mode: Optional[Union[StreamMode, list[StreamMode]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + subgraphs: bool = False, + ) -> AsyncIterator[Union[dict[str, Any], Any]]: + merged_config = merge_configs(self.config, config) + sanitized_config = self._sanitize_config(merged_config) + + 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 + interrupt_before=interrupt_before, # type: ignore + interrupt_after=interrupt_after, # type: ignore + stream_subgraphs=subgraphs, + ): + yield chunk + + async def astream_events( + self, + input: Any, + config: Optional[RunnableConfig] = None, + **kwargs: Any, + ) -> AsyncIterator[StreamEvent]: + merged_config = merge_configs(self.config, config) + sanitized_config = self._sanitize_config(merged_config) + + # manually add 'events' to stream modes list + stream_mode: list[str] = kwargs.get("stream_mode", []) + if "events" not in stream_mode: + stream_mode.append("events") + + 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, # type: ignore + interrupt_before=kwargs.get("interrupt_before"), + interrupt_after=kwargs.get("interrupt_after"), + stream_subgraphs=kwargs.get("subgraphs", False), + ): + yield StandardStreamEvent( + event=chunk.event, + data=chunk.data, + ) + + def invoke( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + ) -> Union[dict[str, Any], Any]: + merged_config = merge_configs(self.config, config) + sanitized_config = self._sanitize_config(merged_config) + + return self.sync_client.runs.wait( + thread_id=sanitized_config["configurable"]["thread_id"], + assistant_id=self.graph_id, + input=input, + config=sanitized_config, + interrupt_before=interrupt_before, # type: ignore + interrupt_after=interrupt_after, # type: ignore + ) + + async def ainvoke( + self, + input: Union[dict[str, Any], Any], + config: Optional[RunnableConfig] = None, + *, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + ) -> Union[dict[str, Any], Any]: + merged_config = merge_configs(self.config, config) + sanitized_config = self._sanitize_config(merged_config) + + return await self.client.runs.wait( + thread_id=sanitized_config["configurable"]["thread_id"], + assistant_id=self.graph_id, + input=input, + config=sanitized_config, + interrupt_before=interrupt_before, # type: ignore + interrupt_after=interrupt_after, # type: ignore + ) diff --git a/libs/langgraph/poetry.lock b/libs/langgraph/poetry.lock index 6e3cd6ffe..58b231fd5 100644 --- a/libs/langgraph/poetry.lock +++ b/libs/langgraph/poetry.lock @@ -687,6 +687,17 @@ cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"] http2 = ["h2 (>=3,<5)"] socks = ["socksio (==1.*)"] +[[package]] +name = "httpx-sse" +version = "0.4.0" +description = "Consume Server-Sent Event (SSE) messages with HTTPX." +optional = false +python-versions = ">=3.8" +files = [ + {file = "httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721"}, + {file = "httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f"}, +] + [[package]] name = "idna" version = "3.7" @@ -1289,6 +1300,22 @@ langgraph-checkpoint = "^2.0.0" type = "directory" url = "../checkpoint-sqlite" +[[package]] +name = "langgraph-sdk" +version = "0.1.32" +description = "SDK for interacting with LangGraph API" +optional = false +python-versions = "<4.0.0,>=3.9.0" +files = [ + {file = "langgraph_sdk-0.1.32-py3-none-any.whl", hash = "sha256:b77770f0641dc7b04f196d313233548d9d65888818aa58c1fdbcdbbf9b85d740"}, + {file = "langgraph_sdk-0.1.32.tar.gz", hash = "sha256:d0bd7bbdc44d6a3afa79a010f2eb4ea0aa2e50485e21b88004d05dff69d24a2e"}, +] + +[package.dependencies] +httpx = ">=0.25.2" +httpx-sse = ">=0.4.0" +orjson = ">=3.10.1" + [[package]] name = "langsmith" version = "0.1.129" @@ -3252,4 +3279,4 @@ test = ["big-O", "importlib-resources", "jaraco.functools", "jaraco.itertools", [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "6abc4c3073b1ce305ce15f891821f34fc5a8fd6b3bab98ad51bf899d7dac37b4" +content-hash = "e4e6d19d835c0c142af7d937afd2ec82293072bb7f0bdb41204da36187cf9774" diff --git a/libs/langgraph/pyproject.toml b/libs/langgraph/pyproject.toml index cffeea877..7b87e3378 100644 --- a/libs/langgraph/pyproject.toml +++ b/libs/langgraph/pyproject.toml @@ -11,6 +11,7 @@ repository = "https://www.github.com/langchain-ai/langgraph" python = ">=3.9.0,<4.0" langchain-core = ">=0.2.39,<0.4" langgraph-checkpoint = "^2.0.0" +langgraph-sdk = "^0.1.32" [tool.poetry.group.dev.dependencies] pytest = "^8.3.2" diff --git a/libs/langgraph/tests/test_remote_pregel.py b/libs/langgraph/tests/test_remote_pregel.py new file mode 100644 index 000000000..35aeaaa4a --- /dev/null +++ b/libs/langgraph/tests/test_remote_pregel.py @@ -0,0 +1,666 @@ +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.runnables.graph import ( + Edge as DrawableEdge, +) +from langchain_core.runnables.graph import ( + Node as DrawableNode, +) + +from langgraph.pregel.remote import RemotePregel +from langgraph.pregel.types import StateSnapshot + + +def test_with_config(): + # set up test + remote_pregel = RemotePregel( + graph_id="test_graph_id", + config={ + "configurable": { + "foo": "bar", + "thread_id": "thread_id_1", + } + }, + ) + + # call method / assertions + config = {"configurable": {"hello": "world"}} + remote_pregel_copy = remote_pregel.with_config(config) + + # assert that a copy was returned + assert remote_pregel_copy != remote_pregel + # assert that configs were merged + assert remote_pregel_copy.config == { + "configurable": { + "foo": "bar", + "thread_id": "thread_id_1", + "hello": "world", + } + } + + +def test_get_graph(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.assistants.get_graph.return_value = { + "nodes": [ + {"id": "__start__", "type": "schema", "data": "__start__"}, + {"id": "__end__", "type": "schema", "data": "__end__"}, + { + "id": "agent", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "RunnableCallable"], + "name": "agent", + }, + }, + ], + "edges": [ + {"source": "__start__", "target": "agent"}, + {"source": "agent", "target": "__end__"}, + ], + } + + remote_pregel = RemotePregel( + sync_client=mock_sync_client, graph_id="test_graph_id" + ) + + # call method / assertions + drawable_graph = remote_pregel.get_graph() + + assert drawable_graph.nodes == { + "__start__": DrawableNode( + id="__start__", name="", data="__start__", metadata=None + ), + "__end__": DrawableNode(id="__end__", name="", data="__end__", metadata=None), + "agent": DrawableNode( + id="agent", + name="", + data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent"}, + metadata=None, + ), + } + + assert drawable_graph.edges == [ + DrawableEdge(source="__start__", target="agent"), + DrawableEdge(source="agent", target="__end__"), + ] + + +@pytest.mark.anyio +async def test_aget_graph(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.assistants.get_graph.return_value = { + "nodes": [ + {"id": "__start__", "type": "schema", "data": "__start__"}, + {"id": "__end__", "type": "schema", "data": "__end__"}, + { + "id": "agent", + "type": "runnable", + "data": { + "id": ["langgraph", "utils", "RunnableCallable"], + "name": "agent", + }, + }, + ], + "edges": [ + {"source": "__start__", "target": "agent"}, + {"source": "agent", "target": "__end__"}, + ], + } + + remote_pregel = RemotePregel( + client=mock_async_client, graph_id="test_graph_id" + ) + + # call method / assertions + drawable_graph = await remote_pregel.aget_graph() + + assert drawable_graph.nodes == { + "__start__": DrawableNode( + id="__start__", name="", data="__start__", metadata=None + ), + "__end__": DrawableNode(id="__end__", name="", data="__end__", metadata=None), + "agent": DrawableNode( + id="agent", + name="", + data={"id": ["langgraph", "utils", "RunnableCallable"], "name": "agent"}, + metadata=None, + ), + } + + assert drawable_graph.edges == [ + DrawableEdge(source="__start__", target="agent"), + DrawableEdge(source="agent", target="__end__"), + ] + + +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 = RemotePregel( + sync_client=mock_sync_client, graph_id="test_graph_id_1" + ) + + # 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: RemotePregel = 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: RemotePregel = 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 = RemotePregel( + client=mock_async_client, + graph_id="test_graph_id_1", + ) + + # 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: RemotePregel = 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: RemotePregel = 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() + mock_sync_client.threads.get_state.return_value = { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": None, + "tasks": [], + } + + # call method / assertions + remote_pregel = RemotePregel( + sync_client=mock_sync_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_snapshot = remote_pregel.get_state(config) + + assert state_snapshot == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config=None, + tasks=(), + ) + + +@pytest.mark.anyio +async def test_aget_state(): + mock_async_client = AsyncMock() + mock_async_client.threads.get_state.return_value = { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_2", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "tasks": [], + } + + # call method / assertions + remote_pregel = RemotePregel( + client=mock_async_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_snapshot = await remote_pregel.aget_state(config) + + assert state_snapshot == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_2", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + tasks=(), + ) + + +def test_get_state_history(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.threads.get_history.return_value = [ + { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": None, + "tasks": [], + } + ] + + # call method / assertions + remote_pregel = RemotePregel( + sync_client=mock_sync_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_history_snapshot = list( + remote_pregel.get_state_history(config, filter=None, before=None, limit=None) + ) + + assert len(state_history_snapshot) == 1 + assert state_history_snapshot[0] == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config=None, + tasks=(), + ) + + +@pytest.mark.anyio +async def test_aget_state_history(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.threads.get_history.return_value = [ + { + "values": {"messages": [{"type": "human", "content": "hello"}]}, + "next": None, + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + }, + "metadata": {}, + "created_at": "timestamp", + "parent_checkpoint": None, + "tasks": [], + } + ] + + # call method / assertions + remote_pregel = RemotePregel( + client=mock_async_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread1"}} + state_history_snapshot = [] + async for state_snapshot in remote_pregel.aget_state_history( + config, filter=None, before=None, limit=None + ): + state_history_snapshot.append(state_snapshot) + + assert len(state_history_snapshot) == 1 + assert state_history_snapshot[0] == StateSnapshot( + values={"messages": [{"type": "human", "content": "hello"}]}, + next=(), + config={ + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + }, + metadata={}, + created_at="timestamp", + parent_config=None, + tasks=(), + ) + + +def test_update_state(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.threads.update_state.return_value = { + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + # call method / assertions + remote_pregel = RemotePregel( + sync_client=mock_sync_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread1"}} + response = remote_pregel.update_state(config, {"key": "value"}) + + assert response == { + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + +@pytest.mark.anyio +async def test_aupdate_state(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.threads.update_state.return_value = { + "checkpoint": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + # call method / assertions + remote_pregel = RemotePregel( + client=mock_async_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread1"}} + response = await remote_pregel.aupdate_state(config, {"key": "value"}) + + assert response == { + "configurable": { + "thread_id": "thread_1", + "checkpoint_ns": "ns", + "checkpoint_id": "checkpoint_1", + "checkpoint_map": {}, + } + } + + +def test_stream(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.runs.stream.return_value = [ + {"chunk": "data1"}, + {"chunk": "data2"}, + {"chunk": "data3"}, + ] + + # call method / assertions + remote_pregel = RemotePregel( + 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"}] + + +@pytest.mark.anyio +async def test_astream(): + # set up test + mock_async_client = MagicMock() + async_iter = MagicMock() + async_iter.__aiter__.return_value = [ + {"chunk": "data1"}, + {"chunk": "data2"}, + {"chunk": "data3"}, + ] + mock_async_client.runs.stream.return_value = async_iter + + # call method / assertions + remote_pregel = RemotePregel( + 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"}] + + +def test_invoke(): + # set up test + mock_sync_client = MagicMock() + mock_sync_client.runs.wait.return_value = { + "values": {"messages": [{"type": "human", "content": "world"}]} + } + + # call method / assertions + remote_pregel = RemotePregel( + sync_client=mock_sync_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread_1"}} + result = remote_pregel.invoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config + ) + + assert result == {"values": {"messages": [{"type": "human", "content": "world"}]}} + + +@pytest.mark.anyio +async def test_ainvoke(): + # set up test + mock_async_client = AsyncMock() + mock_async_client.runs.wait.return_value = { + "values": {"messages": [{"type": "human", "content": "world"}]} + } + + # call method / assertions + remote_pregel = RemotePregel( + client=mock_async_client, graph_id="test_graph_id" + ) + + config = {"configurable": {"thread_id": "thread_1"}} + result = await remote_pregel.ainvoke( + {"input": {"messages": [{"type": "human", "content": "hello"}]}}, config + ) + + assert result == {"values": {"messages": [{"type": "human", "content": "world"}]}} + + +@pytest.mark.skip("Unskip this test to manually test the LangGraph Cloud integration") +@pytest.mark.anyio +async def test_langgraph_cloud_integration(): + from langgraph_sdk.client import get_client, get_sync_client + + from langgraph.checkpoint.memory import MemorySaver + from langgraph.graph import END, START, MessagesState, StateGraph + + # create RemotePregel instance + client = get_client() + sync_client = get_sync_client() + remote_pregel = RemotePregel( + client=client, sync_client=sync_client, graph_id="agent" + ) + + # define graph + workflow = StateGraph(MessagesState) + workflow.add_node("agent", remote_pregel) + workflow.add_edge(START, "agent") + workflow.add_edge("agent", END) + app = workflow.compile(checkpointer=MemorySaver()) + + # test invocation + input = { + "messages": [ + { + "role": "human", + "content": "Hello world!", + } + ] + } + + # test invoke + response = app.invoke( + input, + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + ) + print("response:", response["messages"][-1].content) + + # test stream + async for chunk in app.astream( + input, + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + subgraphs=True, + stream_mode=["debug", "messages"], + ): + print("chunk:", chunk) + + # test stream events + async for chunk in remote_pregel.astream_events( + input, + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + version="v2", + subgraphs=True, + stream_mode=[], + ): + print("chunk:", chunk) + + # test get state + state_snapshot = await remote_pregel.aget_state( + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + subgraphs=True, + ) + print("state snapshot:", state_snapshot) + + # test update state + response = await remote_pregel.aupdate_state( + config={"configurable": {"thread_id": "6645e002-ed50-4022-92a3-d0d186fdf812"}}, + values={ + "messages": [ + { + "role": "ai", + "content": "Hello world again!", + } + ] + }, + ) + print("response:", response) + + # test get history + async for state in remote_pregel.aget_state_history( + config={"configurable": {"thread_id": "2dc3e3e7-39ac-4597-aa57-4404b944e82a"}}, + ): + print("state snapshot:", state) + + # test get graph + 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)