mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-06 09:47:51 +02:00
Merge pull request #339 from langchain-ai/nc/24apr/draw-xray
Improvements to get_graph(xray=)
This commit is contained in:
Nuno Campos
+11
-11
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
Generated
+4
-4
@@ -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"
|
||||
|
||||
+1
-1
@@ -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]
|
||||
|
||||
@@ -1901,6 +1901,218 @@
|
||||
|
||||
'''
|
||||
# ---
|
||||
# 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
|
||||
'''
|
||||
%%{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"]}}}'
|
||||
# ---
|
||||
|
||||
@@ -3848,6 +3848,35 @@ 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_mermaid() == snapshot
|
||||
|
||||
|
||||
def test_nested_graph(snapshot: SnapshotAssertion) -> None:
|
||||
def never_called_fn(state: Any):
|
||||
assert 0, "This function should never be called"
|
||||
|
||||
Reference in New Issue
Block a user