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

Implement input_schema, output_schema, get_graph for Graph, StateGraph, MessageGraph

Browse Source
This commit is contained in:
Nuno Campos
2024-02-19 13:07:34 -08:00
parent 50b1a455bd
commit 27464c0f46
7 changed files with 316 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+34 -2
View File
@@ -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
langgraph/graph/state.py
+5 -4
View File
@@ -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",
langgraph/pregel/__init__.py
+1
View File
@@ -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]
poetry.lock
Generated
+33 -3
View File
@@ -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"
pyproject.toml
+1
View File
@@ -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"
tests/__snapshots__/test_pregel.ambr
+230
View File
@@ -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<agent_outcome>"
}
},
{
"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 |
+---------+ +-------+
'''
# ---
tests/test_pregel.py
+12 -3
View File
@@ -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": [