mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 15:42:25 +02:00
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).
This commit is contained in:
Andrew Nguonly
@@ -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]: ...
|
||||
@@ -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
|
||||
)
|
||||
Generated
+28
-1
@@ -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"
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user