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

Improvements to get_graph(xray=)

Browse Source 
- Support xray: int, to control depth, eg xray=1 only exposes one level deep
- In mermaid draw a containing box around subgraphs created from xray
This commit is contained in:
Nuno Campos
2024-04-25 11:27:13 -07:00
parent daf8bf2eef
commit 42b13ef576
6 changed files with 327 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/graph/graph.py
+11 -11
View File
@@ -15,9 +15,6 @@ from typing import (
from langchain_core.runnables import Runnable
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 langchain_core.runnables.graph import (
Node as RunnableGraphNode,
)
@@ -28,7 +25,7 @@ from langgraph.constants import TAG_HIDDEN
from langgraph.pregel import Channel, Pregel
from langgraph.pregel.read import PregelNode
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
from langgraph.utils import RunnableCallable
from langgraph.utils import DrawableGraph, RunnableCallable
logger = logging.getLogger(__name__)
@@ -380,11 +377,11 @@ class CompiledGraph(Pregel):
self,
config: Optional[RunnableConfig] = None,
*,
xray: bool = False,
xray: Union[int, bool] = False,
add_condition_nodes: bool = True,
) -> RunnableGraph:
) -> DrawableGraph:
"""Returns a drawable representation of the computation graph."""
graph = RunnableGraph()
graph = DrawableGraph()
start_nodes: dict[str, RunnableGraphNode] = {
START: graph.add_node(self.get_input_schema(config), START)
}
@@ -395,16 +392,19 @@ class CompiledGraph(Pregel):
for key, node in self.graph.nodes.items():
if xray:
subgraph = (
node.get_graph(config=config, xray=xray)
node.get_graph(
config=config,
xray=xray - 1 if isinstance(xray, int) and xray > 0 else xray,
)
if isinstance(node, CompiledGraph)
else node.get_graph(config=config)
)
subgraph.trim_first_node()
subgraph.trim_last_node()
if len(subgraph.nodes) > 1:
graph.extend(subgraph)
start_nodes[key] = subgraph.last_node()
end_nodes[key] = subgraph.first_node()
end_nodes[key], start_nodes[key] = graph.extend(
subgraph, prefix=key
)
else:
n = graph.add_node(node, key)
start_nodes[key] = n
langgraph/utils.py
+32
View File
@@ -3,6 +3,7 @@ from typing import Any, Awaitable, Callable, Optional
from langchain_core.runnables import Runnable, RunnableConfig
from langchain_core.runnables.config import merge_configs
from langchain_core.runnables.graph import Edge, Graph, Node, is_uuid
# Before Python 3.11 native StrEnum is not available
@@ -67,3 +68,34 @@ class RunnableCallable(Runnable):
if isinstance(ret, Runnable) and self.recurse:
return await ret.ainvoke(input, config)
return ret
class DrawableGraph(Graph):
def extend(
self, graph: Graph, prefix: str = ""
) -> tuple[Optional[Node], Optional[Node]]:
if all(is_uuid(node.id) for node in graph.nodes.values()):
super().extend(graph)
return graph.first_node(), graph.last_node()
new_nodes = {
f"{prefix}:{k}": Node(f"{prefix}:{k}", v.data)
for k, v in graph.nodes.items()
}
new_edges = [
Edge(
f"{prefix}:{edge.source}",
f"{prefix}:{edge.target}",
edge.data,
edge.conditional,
)
for edge in graph.edges
]
self.nodes.update(new_nodes)
self.edges.extend(new_edges)
first = graph.first_node()
last = graph.last_node()
return (
Node(f"{prefix}:{first.id}", first.data) if first else None,
Node(f"{prefix}:{last.id}", last.data) if last else None,
)
poetry.lock
Generated
+4 -4
View File
@@ -1710,13 +1710,13 @@ extended-testing = ["aiosqlite (>=0.19.0,<0.20.0)", "aleph-alpha-client (>=2.15.
[[package]]
name = "langchain-core"
version = "0.1.45"
version = "0.1.46"
description = "Building applications with LLMs through composability"
optional = false
python-versions = "<4.0,>=3.8.1"
files = [
{file = "langchain_core-0.1.45-py3-none-any.whl", hash = "sha256:91eff20de0bcf5f025e1d8c4582cb597a9c17527965eb03b314486e7c834e7df"},
{file = "langchain_core-0.1.45.tar.gz", hash = "sha256:526532c1af279a9e2debe7a4e143ba6e980cf90b5ab2e0991c2230ee04c693e2"},
{file = "langchain_core-0.1.46-py3-none-any.whl", hash = "sha256:1c0befcd2665dd4aa153318aa9bf729071644b4c179e491769b8e583b4bf7441"},
{file = "langchain_core-0.1.46.tar.gz", hash = "sha256:17c416349f5c7a9808e70e3725749a3a2df5088f1ecca045c883871aa95f9c9e"},
]
[package.dependencies]
@@ -4094,4 +4094,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 = "5d0e5b014e355f65a1731548531a53313cff61d729312d64f1f3b339673e2f5c"
content-hash = "4c5fea39af2c255ddbc99a3b0ccbdb94e53c3526e210f023f71c95b1d34c0a98"
pyproject.toml
+1 -1
View File
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
[tool.poetry.dependencies]
python = ">=3.9.0,<4.0"
langchain-core = "^0.1.42"
langchain-core = "^0.1.46"
[tool.poetry.group.test.dependencies]
tests/__snapshots__/test_pregel.ambr
+249
View File
@@ -1901,6 +1901,255 @@
'''
# ---
# name: test_nested_graph_xray
dict({
'edges': list([
dict({
'conditional': True,
'data': 'tool_two_slow',
'source': 'tool_two:condition',
'target': 'tool_two:tool_two_slow',
}),
dict({
'source': 'tool_two:tool_two_slow',
'target': 'tool_two:__end__',
}),
dict({
'conditional': True,
'data': 'tool_two_fast',
'source': 'tool_two:condition',
'target': 'tool_two:tool_two_fast',
}),
dict({
'source': 'tool_two:tool_two_fast',
'target': 'tool_two:__end__',
}),
dict({
'source': '__start__',
'target': 'condition',
}),
dict({
'conditional': True,
'data': 'tool_one',
'source': 'condition',
'target': 'tool_one',
}),
dict({
'source': 'tool_one',
'target': '__end__',
}),
dict({
'conditional': True,
'data': 'tool_two',
'source': 'condition',
'target': 'tool_two:condition',
}),
dict({
'source': 'tool_two:__end__',
'target': '__end__',
}),
dict({
'conditional': True,
'data': 'tool_three',
'source': 'condition',
'target': 'tool_three',
}),
dict({
'source': 'tool_three',
'target': '__end__',
}),
]),
'nodes': list([
dict({
'data': '__start__',
'id': '__start__',
'type': 'schema',
}),
dict({
'data': '__end__',
'id': '__end__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langchain_core',
'runnables',
'base',
'RunnableLambda',
]),
'name': 'logic',
}),
'id': 'tool_one',
'type': 'runnable',
}),
dict({
'data': 'tool_two:__end__',
'id': 'tool_two:__end__',
'type': 'schema',
}),
dict({
'data': dict({
'id': list([
'langchain_core',
'runnables',
'base',
'RunnableLambda',
]),
'name': 'logic',
}),
'id': 'tool_two:tool_two_slow',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langchain_core',
'runnables',
'base',
'RunnableLambda',
]),
'name': 'logic',
}),
'id': 'tool_two:tool_two_fast',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langchain_core',
'runnables',
'base',
'RunnableLambda',
]),
'name': 'RunnableLambda',
}),
'id': 'tool_two:condition',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langchain_core',
'runnables',
'base',
'RunnableLambda',
]),
'name': 'logic',
}),
'id': 'tool_three',
'type': 'runnable',
}),
dict({
'data': dict({
'id': list([
'langchain_core',
'runnables',
'base',
'RunnableLambda',
]),
'name': 'RunnableLambda',
}),
'id': 'condition',
'type': 'runnable',
}),
]),
})
# ---
# name: test_nested_graph_xray.1
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-----------+
| condition |
..+-----------+...
..... . .....
... . ...
... . ...
+----------+ +----------+ +------------+
| tool_one |* | tool_two | | tool_three |
+----------+ *** +----------+ **+------------+
***** * *****
*** * ***
*** * ***
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_nested_graph_xray.2
'''
+-----------+
| __start__ |
+-----------+
*
*
*
+-----------+
| condition |.....
.....+-----------+... .......
.... .... .........
..... ... .......
... .... .........
+--------------------+ .. ....
| tool_two:condition | . .
+--------------------+ . .
... ... . .
... ... . .
.. .. . .
+------------------------+ +------------------------+ . .
| tool_two:tool_two_slow | | tool_two:tool_two_fast | . .
+------------------------+ +------------------------+ . .
*** *** . .
*** *** . .
** ** . .
+------------------+ +----------+ +------------+
| tool_two:__end__ | | tool_one | ****| tool_three |
+------------------+** ***+----------+******** +------------+
**** *** ********
***** *************
*** ******
+---------+
| __end__ |
+---------+
'''
# ---
# name: test_nested_graph_xray.3
'''
%%{init: {'flowchart': {'curve': 'linear'}}}%%
graph TD;
__start__[__start__]:::startclass;
__end__[__end__]:::endclass;
tool_one([tool_one]):::otherclass;
tool_two___end__([__end__]):::otherclass;
tool_two_tool_two_slow([tool_two_slow]):::otherclass;
tool_two_tool_two_fast([tool_two_fast]):::otherclass;
tool_two_condition([condition]):::otherclass;
tool_three([tool_three]):::otherclass;
condition([condition]):::otherclass;
subgraph tool_two
tool_two_condition -. tool_two_slow .-> tool_two_tool_two_slow;
tool_two_tool_two_slow --> tool_two___end__;
tool_two_condition -. tool_two_fast .-> tool_two_tool_two_fast;
tool_two_tool_two_fast --> tool_two___end__;
end
__start__ --> condition;
condition -. tool_one .-> tool_one;
tool_one --> __end__;
condition -. tool_two .-> tool_two_condition;
tool_two___end__ --> __end__;
condition -. tool_three .-> tool_three;
tool_three --> __end__;
classDef startclass fill:#ffdfba;
classDef endclass fill:#baffc9;
classDef otherclass fill:#fad7de;
'''
# ---
# name: test_prebuilt_chat
'{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"BaseMessage": {"title": "BaseMessage", "description": "Base abstract Message class.\\n\\nMessages are the inputs and outputs of ChatModels.", "type": "object", "properties": {"content": {"title": "Content", "anyOf": [{"type": "string"}, {"type": "array", "items": {"anyOf": [{"type": "string"}, {"type": "object"}]}}]}, "additional_kwargs": {"title": "Additional Kwargs", "type": "object"}, "response_metadata": {"title": "Response Metadata", "type": "object"}, "type": {"title": "Type", "type": "string"}, "name": {"title": "Name", "type": "string"}, "id": {"title": "Id", "type": "string"}}, "required": ["content", "type"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "required": ["messages"]}}}'
# ---
tests/test_pregel.py
+30
View File
@@ -3848,6 +3848,36 @@ def test_simple_multi_edge(snapshot: SnapshotAssertion) -> None:
assert app.invoke({"my_key": "my_value"}) == {"my_key": "my_value"}
def test_nested_graph_xray(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
def logic(state: State):
pass
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two_slow", logic)
tool_two_graph.add_node("tool_two_fast", logic)
tool_two_graph.set_conditional_entry_point(
lambda s: "tool_two_slow" if s["market"] == "DE" else "tool_two_fast",
then=END,
)
tool_two = tool_two_graph.compile()
graph = StateGraph(State)
graph.add_node("tool_one", logic)
graph.add_node("tool_two", tool_two)
graph.add_node("tool_three", logic)
graph.set_conditional_entry_point(lambda s: "tool_one", then=END)
app = graph.compile()
assert app.get_graph(xray=True).to_json() == snapshot
assert app.get_graph().draw_ascii() == snapshot
assert app.get_graph(xray=True).draw_ascii() == snapshot
assert app.get_graph(xray=True).draw_mermaid() == snapshot
def test_nested_graph(snapshot: SnapshotAssertion) -> None:
def never_called_fn(state: Any):
assert 0, "This function should never be called"