WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-03 04:58:43 +02:00
Code Issues Packages Projects Releases Wiki Activity

Lint

Browse Source
This commit is contained in:
Nuno Campos
2024-10-23 17:01:58 -07:00
parent dc8260bb72
commit 69227daff3
5 changed files with 42 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/graph.py
+8
View File
@@ -514,6 +514,14 @@ class CompiledGraph(Pregel):
self.nodes[end].triggers.append(channel_name)
cast(list[str], self.nodes[end].channels).append(channel_name)
async def aget_graph(
self,
config: Optional[RunnableConfig] = None,
*,
xray: Union[int, bool] = False,
) -> DrawableGraph:
return self.get_graph(config, xray=xray)
def get_graph(
self,
config: Optional[RunnableConfig] = None,
libs/langgraph/langgraph/pregel/__init__.py
+12 -4
View File
@@ -25,7 +25,6 @@ from uuid import UUID, uuid5
from langchain_core.globals import get_debug
from langchain_core.runnables import (
Runnable,
RunnableSequence,
)
from langchain_core.runnables.base import Input, Output
@@ -34,6 +33,7 @@ from langchain_core.runnables.config import (
get_async_callback_manager_for_config,
get_callback_manager_for_config,
)
from langchain_core.runnables.graph import Graph
from langchain_core.runnables.utils import (
ConfigurableFieldSpec,
get_unique_config_specs,
@@ -180,9 +180,7 @@ class Channel:
)
class Pregel(
Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], PregelProtocol
):
class Pregel(PregelProtocol):
nodes: dict[str, PregelNode]
channels: dict[str, Union[BaseChannel, ManagedValueSpec]]
@@ -262,6 +260,16 @@ class Pregel(
if auto_validate:
self.validate()
def get_graph(
self, config: RunnableConfig | None = None, *, xray: int | bool = False
) -> Graph:
raise NotImplementedError
async def aget_graph(
self, config: RunnableConfig | None = None, *, xray: int | bool = False
) -> Graph:
raise NotImplementedError
def copy(self, update: dict[str, Any] | None = None) -> Self:
attrs = {**self.__dict__, **(update or {})}
return self.__class__(**attrs)
libs/langgraph/langgraph/pregel/protocol.py
+18 -13
View File
@@ -1,27 +1,29 @@
from abc import ABC, abstractmethod
from typing import (
Any,
AsyncIterator,
Iterator,
Optional,
Protocol,
Sequence,
Union,
runtime_checkable,
)
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.graph import Graph as DrawableGraph
from typing_extensions import Self
from langgraph.pregel.types import All, StateSnapshot, StreamMode
@runtime_checkable
class PregelProtocol(Protocol):
class PregelProtocol(
Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]], ABC
):
@abstractmethod
def with_config(
self, config: Optional[RunnableConfig] = None, **kwargs: Any
) -> Self: ...
@abstractmethod
def get_graph(
self,
config: Optional[RunnableConfig] = None,
@@ -29,6 +31,7 @@ class PregelProtocol(Protocol):
xray: Union[int, bool] = False,
) -> DrawableGraph: ...
@abstractmethod
async def aget_graph(
self,
config: Optional[RunnableConfig] = None,
@@ -36,22 +39,17 @@ class PregelProtocol(Protocol):
xray: Union[int, bool] = False,
) -> DrawableGraph: ...
def get_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> Iterator[tuple[str, "PregelProtocol"]]: ...
def aget_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> AsyncIterator[tuple[str, "PregelProtocol"]]: ...
@abstractmethod
def get_state(
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot: ...
@abstractmethod
async def aget_state(
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot: ...
@abstractmethod
def get_state_history(
self,
config: RunnableConfig,
@@ -61,6 +59,7 @@ class PregelProtocol(Protocol):
limit: Optional[int] = None,
) -> Iterator[StateSnapshot]: ...
@abstractmethod
def aget_state_history(
self,
config: RunnableConfig,
@@ -70,6 +69,7 @@ class PregelProtocol(Protocol):
limit: Optional[int] = None,
) -> AsyncIterator[StateSnapshot]: ...
@abstractmethod
def update_state(
self,
config: RunnableConfig,
@@ -77,6 +77,7 @@ class PregelProtocol(Protocol):
as_node: Optional[str] = None,
) -> RunnableConfig: ...
@abstractmethod
async def aupdate_state(
self,
config: RunnableConfig,
@@ -84,6 +85,7 @@ class PregelProtocol(Protocol):
as_node: Optional[str] = None,
) -> RunnableConfig: ...
@abstractmethod
def stream(
self,
input: Union[dict[str, Any], Any],
@@ -95,6 +97,7 @@ class PregelProtocol(Protocol):
subgraphs: bool = False,
) -> Iterator[Union[dict[str, Any], Any]]: ...
@abstractmethod
def astream(
self,
input: Union[dict[str, Any], Any],
@@ -106,6 +109,7 @@ class PregelProtocol(Protocol):
subgraphs: bool = False,
) -> AsyncIterator[Union[dict[str, Any], Any]]: ...
@abstractmethod
def invoke(
self,
input: Union[dict[str, Any], Any],
@@ -115,6 +119,7 @@ class PregelProtocol(Protocol):
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
) -> Union[dict[str, Any], Any]: ...
@abstractmethod
async def ainvoke(
self,
input: Union[dict[str, Any], Any],
libs/langgraph/langgraph/pregel/remote.py
+4 -26
View File
@@ -9,7 +9,7 @@ from typing import (
)
import orjson
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables import RunnableConfig
from langchain_core.runnables.graph import (
Edge as DrawableEdge,
)
@@ -44,7 +44,9 @@ class RemoteException(Exception):
pass
class RemoteGraph(PregelProtocol, Runnable):
class RemoteGraph(PregelProtocol):
name: str
def __init__(
self,
name: str, # graph_id
@@ -124,30 +126,6 @@ class RemoteGraph(PregelProtocol, Runnable):
edges=[DrawableEdge(**edge) for edge in graph["edges"]],
)
def get_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> Iterator[tuple[str, "PregelProtocol"]]:
subgraphs = self.sync_client.assistants.get_subgraphs(
assistant_id=self.name,
namespace=namespace,
recurse=recurse,
)
for namespace, graph_schema in subgraphs.items():
remote_subgraph = self.copy({"name": graph_schema["graph_id"]})
yield (namespace, remote_subgraph)
async def aget_subgraphs(
self, namespace: Optional[str] = None, recurse: bool = False
) -> AsyncIterator[tuple[str, "PregelProtocol"]]:
subgraphs = await self.client.assistants.get_subgraphs(
assistant_id=self.name,
namespace=namespace,
recurse=recurse,
)
for namespace, graph_schema in subgraphs.items():
remote_subgraph = self.copy({"name": graph_schema["graph_id"]})
yield (namespace, remote_subgraph)
def _create_state_snapshot(self, state: ThreadState) -> StateSnapshot:
tasks = []
for task in state["tasks"]:
libs/langgraph/tests/test_remote_graph.py
-90
View File
@@ -135,90 +135,6 @@ async def test_aget_graph():
]
def test_get_subgraphs():
# set up test
mock_sync_client = MagicMock()
mock_sync_client.assistants.get_subgraphs.return_value = {
"namespace_1": {
"graph_id": "test_graph_id_2",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
"namespace_2": {
"graph_id": "test_graph_id_3",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
}
remote_pregel = RemoteGraph("test_graph_id_1", sync_client=mock_sync_client)
# call method / assertions
subgraphs = list(remote_pregel.get_subgraphs())
assert len(subgraphs) == 2
subgraph_1 = subgraphs[0]
ns_1 = subgraph_1[0]
remote_pregel_1: RemoteGraph = subgraph_1[1]
assert ns_1 == "namespace_1"
assert remote_pregel_1.graph_id == "test_graph_id_2"
subgraph_2 = subgraphs[1]
ns_2 = subgraph_2[0]
remote_pregel_2: RemoteGraph = subgraph_2[1]
assert ns_2 == "namespace_2"
assert remote_pregel_2.graph_id == "test_graph_id_3"
@pytest.mark.anyio
async def test_aget_subgraphs():
# set up test
mock_async_client = AsyncMock()
mock_async_client.assistants.get_subgraphs.return_value = {
"namespace_1": {
"graph_id": "test_graph_id_2",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
"namespace_2": {
"graph_id": "test_graph_id_3",
"input_schema": {},
"output_schema": {},
"state_schema": {},
"config_schema": {},
},
}
remote_pregel = RemoteGraph(
"test_graph_id_1",
client=mock_async_client,
)
# call method / assertions
subgraphs = []
async for subgraph in remote_pregel.aget_subgraphs():
subgraphs.append(subgraph)
assert len(subgraphs) == 2
subgraph_1 = subgraphs[0]
ns_1 = subgraph_1[0]
remote_pregel_1: RemoteGraph = subgraph_1[1]
assert ns_1 == "namespace_1"
assert remote_pregel_1.graph_id == "test_graph_id_2"
subgraph_2 = subgraphs[1]
ns_2 = subgraph_2[0]
remote_pregel_2: RemoteGraph = subgraph_2[1]
assert ns_2 == "namespace_2"
assert remote_pregel_2.graph_id == "test_graph_id_3"
def test_get_state():
# set up test
mock_sync_client = MagicMock()
@@ -866,9 +782,3 @@ async def test_langgraph_cloud_integration():
remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID
graph = await remote_pregel.aget_graph(xray=True)
print("graph:", graph)
# test get subgraphs
remote_pregel.graph_id = "fe096781-5601-53d2-b2f6-0d3403f7e9ca" # must be UUID
async for name, pregel in remote_pregel.aget_subgraphs():
print("name:", name)
print("pregel:", pregel)