diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 22e075575..22eba86ad 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -9,6 +9,8 @@ from langchain_core.runnables.base import ( RunnableLike, coerce_to_runnable, ) +from langchain_core.runnables.config import RunnableConfig +from langchain_core.runnables.graph import Graph as RunnableGraph from langgraph.checkpoint import BaseCheckpointSaver from langgraph.pregel import Channel, Pregel @@ -150,7 +152,7 @@ class Graph: checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, - ) -> Pregel: + ) -> "CompiledGraph": interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] self.validate(interrupt=interrupt_before + interrupt_after) @@ -177,7 +179,8 @@ class Graph: branch.runnable, name=f"{key}_condition" ) - return Pregel( + return CompiledGraph( + graph=self, nodes=nodes, input=f"{self.entry_point}:inbox", output=END, @@ -188,3 +191,32 @@ class Graph: + [node for node in interrupt_after] ), ) + + +class CompiledGraph(Pregel): + graph: Graph + + def get_graph(self, config: RunnableConfig | None = None) -> RunnableGraph: + graph = RunnableGraph() + graph.add_node(self.get_input_schema(), "__start__") + graph.add_node(self.get_output_schema(), END) + + for key, node in self.graph.nodes.items(): + graph.add_node(node, key) + for start, end in self.graph.edges: + graph.add_edge(graph.nodes[start], graph.nodes[end]) + for start, branches in self.graph.branches.items(): + for i, branch in enumerate(branches): + name = f"{start}_{branch.condition.__name__}" + if i > 0: + name += f"_{i}" + graph.add_node( + RunnableLambda(branch.runnable, name=branch.condition.__name__), + name, + ) + graph.add_edge(graph.nodes[start], graph.nodes[name]) + for label, end in branch.ends.items(): + graph.add_edge(graph.nodes[name], graph.nodes[end], label) + graph.add_edge(graph.nodes["__start__"], graph.nodes[self.graph.entry_point]) + + return graph diff --git a/langgraph/graph/state.py b/langgraph/graph/state.py index 552b954a1..4a73f1a08 100644 --- a/langgraph/graph/state.py +++ b/langgraph/graph/state.py @@ -12,8 +12,8 @@ from langgraph.channels.binop import BinaryOperatorAggregate from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint import BaseCheckpointSaver -from langgraph.graph.graph import END, Graph -from langgraph.pregel import Channel, Pregel +from langgraph.graph.graph import END, CompiledGraph, Graph +from langgraph.pregel import Channel from langgraph.pregel.read import ChannelRead from langgraph.pregel.write import SKIP_WRITE, ChannelWrite @@ -41,7 +41,7 @@ class StateGraph(Graph): checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Sequence[str]] = None, interrupt_after: Optional[Sequence[str]] = None, - ) -> Pregel: + ) -> CompiledGraph: interrupt_before = interrupt_before or [] interrupt_after = interrupt_after or [] self.validate(interrupt=interrupt_before + interrupt_after) @@ -111,7 +111,8 @@ class StateGraph(Graph): | Channel.write_to(f"{self.entry_point}:inbox") ) - return Pregel( + return CompiledGraph( + graph=self, nodes=nodes, channels={**self.channels, **node_inboxes, **node_outboxes}, input=f"{START}:inbox", diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index d7fadd9e5..2413a78e9 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -237,6 +237,7 @@ class Pregel( self, config: Optional[RunnableConfig] = None ) -> Type[BaseModel]: if isinstance(self.output, str): + print(self.output, self.channels[self.output].UpdateType) return super().get_output_schema(config) else: return create_model( # type: ignore[call-overload] diff --git a/poetry.lock b/poetry.lock index 462f4e42e..03cbc892e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand. [[package]] name = "aiohttp" @@ -843,6 +843,23 @@ files = [ {file = "frozenlist-1.4.1.tar.gz", hash = "sha256:c037a86e8513059a2613aaba4d817bb90b9d9b6b69aace3ce9c877e8c8ed402b"}, ] +[[package]] +name = "grandalf" +version = "0.8" +description = "Graph and drawing algorithms framework" +optional = false +python-versions = "*" +files = [ + {file = "grandalf-0.8-py3-none-any.whl", hash = "sha256:793ca254442f4a79252ea9ff1ab998e852c1e071b863593e5383afee906b4185"}, + {file = "grandalf-0.8.tar.gz", hash = "sha256:2813f7aab87f0d20f334a3162ccfbcbf085977134a17a5b516940a93a77ea974"}, +] + +[package.dependencies] +pyparsing = "*" + +[package.extras] +full = ["numpy", "ply"] + [[package]] name = "greenlet" version = "3.0.3" @@ -2397,6 +2414,20 @@ files = [ plugins = ["importlib-metadata"] windows-terminal = ["colorama (>=0.4.6)"] +[[package]] +name = "pyparsing" +version = "3.1.1" +description = "pyparsing module - Classes and methods to define and execute parsing grammars" +optional = false +python-versions = ">=3.6.8" +files = [ + {file = "pyparsing-3.1.1-py3-none-any.whl", hash = "sha256:32c7c0b711493c72ff18a981d24f28aaf9c1fb7ed5e9667c9e84e3db623bdbfb"}, + {file = "pyparsing-3.1.1.tar.gz", hash = "sha256:ede28a1a32462f5a9705e07aea48001a08f7cf81a021585011deba701581a0db"}, +] + +[package.extras] +diagrams = ["jinja2", "railroad-diagrams"] + [[package]] name = "pytest" version = "7.4.4" @@ -2604,7 +2635,6 @@ files = [ {file = "PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28"}, {file = "PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9"}, - {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef"}, {file = "PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0"}, {file = "PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4"}, {file = "PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54"}, @@ -3729,4 +3759,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9.0,<4.0" -content-hash = "f489f2e9159e8db255a43027617c41a389afb5fe84ff94bf8521c0f98de8cd38" +content-hash = "883b0bc577972335d162a7fa7a3d28ed285808c6932ab0d56ed2cf10919f1d6c" diff --git a/pyproject.toml b/pyproject.toml index 8df024914..c0c95767a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,7 @@ httpx = "^0.26.0" pytest-watcher = "^0.3.4" langchain = "^0.1.0" aiosqlite = "^0.19.0" +grandalf = "^0.8" [tool.poetry.group.lint.dependencies] ruff = "^0.1.4" diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr new file mode 100644 index 000000000..249ae2c92 --- /dev/null +++ b/tests/__snapshots__/test_pregel.ambr @@ -0,0 +1,230 @@ +# serializer version: 1 +# name: test_conditional_graph + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput" + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput" + } + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableAssign" + ], + "name": "RunnableAssign" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "execute_tools" + } + }, + { + "id": "agent_should_continue", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_continue" + } + } + ], + "edges": [ + { + "source": "tools", + "target": "agent", + "data": null + }, + { + "source": "agent", + "target": "agent_should_continue", + "data": null + }, + { + "source": "agent_should_continue", + "target": "tools", + "data": "continue" + }, + { + "source": "agent_should_continue", + "target": "__end__", + "data": "exit" + }, + { + "source": "__start__", + "target": "agent", + "data": null + } + ] + } + ''' +# --- +# name: test_conditional_graph.1 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + *+-------+* + ** *** + ** ** + ** ** + +-----------------------+ ** + | agent_should_continue | * + +-----------------------+ * + * **** * + * ***** * + * *** * + +---------+ +-------+ + | __end__ | | tools | + +---------+ +-------+ + ''' +# --- +# name: test_conditional_graph_state + '{"title": "LangGraphOutput"}' +# --- +# name: test_conditional_graph_state.1 + ''' + { + "nodes": [ + { + "id": "__start__", + "type": "schema", + "data": { + "title": "LangGraphInput" + } + }, + { + "id": "__end__", + "type": "schema", + "data": { + "title": "LangGraphOutput" + } + }, + { + "id": "agent", + "type": "runnable", + "data": { + "id": [ + "langchain", + "schema", + "runnable", + "RunnableSequence" + ], + "name": "RunnableSequence" + } + }, + { + "id": "tools", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "execute_tools" + } + }, + { + "id": "agent_should_continue", + "type": "runnable", + "data": { + "id": [ + "langchain_core", + "runnables", + "base", + "RunnableLambda" + ], + "name": "should_continue" + } + } + ], + "edges": [ + { + "source": "tools", + "target": "agent", + "data": null + }, + { + "source": "agent", + "target": "agent_should_continue", + "data": null + }, + { + "source": "agent_should_continue", + "target": "tools", + "data": "continue" + }, + { + "source": "agent_should_continue", + "target": "__end__", + "data": "exit" + }, + { + "source": "__start__", + "target": "agent", + "data": null + } + ] + } + ''' +# --- +# name: test_conditional_graph_state.2 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + *+-------+* + ** *** + ** ** + ** ** + +-----------------------+ ** + | agent_should_continue | * + +-----------------------+ * + * **** * + * ***** * + * *** * + +---------+ +-------+ + | __end__ | | tools | + +---------+ +-------+ + ''' +# --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index cce76cf85..c79d2cf3d 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -9,6 +9,7 @@ from typing import Annotated, Generator, Optional, TypedDict, Union import pytest from langchain_core.runnables import RunnablePassthrough from pytest_mock import MockerFixture +from syrupy import SnapshotAssertion from langgraph.channels.base import InvalidUpdateError from langgraph.channels.binop import BinaryOperatorAggregate @@ -659,7 +660,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: assert cleanup.call_count == 1, "Expected cleanup to be called once" -def test_conditional_graph() -> None: +def test_conditional_graph(snapshot: SnapshotAssertion) -> None: from copy import deepcopy from langchain.llms.fake import FakeStreamingListLLM @@ -732,6 +733,9 @@ def test_conditional_graph() -> None: app = workflow.compile() + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", "intermediate_steps": [ @@ -879,13 +883,13 @@ def test_conditional_graph() -> None: ] -def test_conditional_graph_state() -> None: +def test_conditional_graph_state(snapshot: SnapshotAssertion) -> None: from langchain.llms.fake import FakeStreamingListLLM from langchain_community.tools import tool from langchain_core.agents import AgentAction, AgentFinish from langchain_core.prompts import PromptTemplate - class AgentState(TypedDict): + class AgentState(TypedDict, total=False): input: str agent_outcome: Optional[Union[AgentAction, AgentFinish]] intermediate_steps: Annotated[list[tuple[AgentAction, str]], operator.add] @@ -959,6 +963,11 @@ def test_conditional_graph_state() -> None: app = workflow.compile() + assert app.get_input_schema().schema_json() == snapshot + assert app.get_output_schema().schema_json() == snapshot + assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot + assert app.get_graph().draw_ascii() == snapshot + assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", "intermediate_steps": [