WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-21 23:22:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Interop of RemoteGraph w core lib

Browse Source
This commit is contained in:
Nuno Campos
2024-10-23 15:37:17 -07:00
parent 62a5ec509d
commit dc8260bb72
5 changed files with 87 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/graph.py
+1 -1
View File
@@ -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:
libs/langgraph/langgraph/pregel/__init__.py
+4 -1
View File
@@ -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]]
libs/langgraph/langgraph/pregel/remote.py
+31 -13
View File
@@ -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,
libs/langgraph/langgraph/pregel/utils.py
+3 -2
View File
@@ -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):
libs/langgraph/tests/test_remote_graph.py
+48 -18
View File
@@ -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