mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-31 20:29:46 +02:00
Implement xray mode for graph draw
This commit is contained in:
Nuno Campos
+43
-15
@@ -10,7 +10,12 @@ from langchain_core.runnables.base import (
|
||||
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 (
|
||||
Graph as RunnableGraph,
|
||||
)
|
||||
from langchain_core.runnables.graph import (
|
||||
Node as RunnableGraphNode,
|
||||
)
|
||||
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.checkpoint import BaseCheckpointSaver
|
||||
@@ -238,48 +243,71 @@ class Graph:
|
||||
class CompiledGraph(Pregel):
|
||||
graph: Graph
|
||||
|
||||
def get_graph(self, config: Optional[RunnableConfig] = None) -> RunnableGraph:
|
||||
def get_graph(
|
||||
self, config: Optional[RunnableConfig] = None, *, xray: bool = False
|
||||
) -> RunnableGraph:
|
||||
graph = RunnableGraph()
|
||||
graph.add_node(self.get_input_schema(config), START)
|
||||
graph.add_node(self.get_output_schema(config), END)
|
||||
start_nodes: dict[str, RunnableGraphNode] = {
|
||||
START: graph.add_node(self.get_input_schema(config), START)
|
||||
}
|
||||
end_nodes: dict[str, RunnableGraphNode] = {
|
||||
END: graph.add_node(self.get_output_schema(config), END)
|
||||
}
|
||||
|
||||
for key, node in self.graph.nodes.items():
|
||||
graph.add_node(node, key)
|
||||
if xray:
|
||||
subgraph = (
|
||||
node.get_graph(config=config, xray=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()
|
||||
else:
|
||||
n = graph.add_node(node, key)
|
||||
start_nodes[key] = n
|
||||
end_nodes[key] = n
|
||||
else:
|
||||
n = graph.add_node(node, key)
|
||||
start_nodes[key] = n
|
||||
end_nodes[key] = n
|
||||
for start, end in self.graph.edges:
|
||||
graph.add_edge(graph.nodes[start], graph.nodes[end])
|
||||
graph.add_edge(start_nodes[start], end_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(
|
||||
cond = graph.add_node(
|
||||
RunnableLambda(branch.runnable, name=branch.condition.__name__),
|
||||
name,
|
||||
)
|
||||
graph.add_edge(graph.nodes[start], graph.nodes[name])
|
||||
graph.add_edge(start_nodes[start], cond)
|
||||
ends = branch.ends or {
|
||||
**{k: k for k in self.graph.nodes},
|
||||
END: END,
|
||||
}
|
||||
for label, end in ends.items():
|
||||
graph.add_edge(graph.nodes[name], graph.nodes[end], label)
|
||||
graph.add_edge(cond, end_nodes[end], label)
|
||||
if self.graph.entry_point_branch:
|
||||
graph.add_node(
|
||||
cond = graph.add_node(
|
||||
RunnableLambda(
|
||||
self.graph.entry_point_branch.runnable,
|
||||
name=self.graph.entry_point_branch.condition.__name__,
|
||||
),
|
||||
f"{START}_condition",
|
||||
)
|
||||
graph.add_edge(graph.nodes[START], graph.nodes[f"{START}_condition"])
|
||||
graph.add_edge(start_nodes[START], cond)
|
||||
ends = self.graph.entry_point_branch.ends or {
|
||||
k: k for k in self.graph.nodes
|
||||
}
|
||||
for label, end in ends.items():
|
||||
graph.add_edge(
|
||||
graph.nodes[f"{START}_condition"], graph.nodes[end], label
|
||||
)
|
||||
graph.add_edge(cond, end_nodes[end], label)
|
||||
elif self.graph.entry_point:
|
||||
graph.add_edge(graph.nodes[START], graph.nodes[self.graph.entry_point])
|
||||
graph.add_edge(start_nodes[START], end_nodes[self.graph.entry_point])
|
||||
|
||||
return graph
|
||||
|
||||
@@ -455,6 +455,301 @@
|
||||
+---------+ +-------+
|
||||
'''
|
||||
# ---
|
||||
# name: test_conditional_graph.2
|
||||
'''
|
||||
{
|
||||
"nodes": [
|
||||
{
|
||||
"id": "__start__",
|
||||
"type": "schema",
|
||||
"data": {
|
||||
"title": "LangGraphInput"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "__end__",
|
||||
"type": "schema",
|
||||
"data": {
|
||||
"title": "LangGraphOutput"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "schema",
|
||||
"data": {
|
||||
"title": "RunnableParallel<agent_outcome>Input",
|
||||
"type": "object",
|
||||
"properties": {}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "schema",
|
||||
"data": {
|
||||
"title": "RunnableParallel<agent_outcome>Output",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"agent_outcome": {
|
||||
"title": "Agent Outcome",
|
||||
"anyOf": [
|
||||
{
|
||||
"$ref": "#/definitions/AgentAction"
|
||||
},
|
||||
{
|
||||
"$ref": "#/definitions/AgentFinish"
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"definitions": {
|
||||
"AgentAction": {
|
||||
"title": "AgentAction",
|
||||
"description": "A full description of an action for an ActionAgent to execute.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"tool": {
|
||||
"title": "Tool",
|
||||
"type": "string"
|
||||
},
|
||||
"tool_input": {
|
||||
"title": "Tool Input",
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "object"
|
||||
}
|
||||
]
|
||||
},
|
||||
"log": {
|
||||
"title": "Log",
|
||||
"type": "string"
|
||||
},
|
||||
"type": {
|
||||
"title": "Type",
|
||||
"default": "AgentAction",
|
||||
"enum": [
|
||||
"AgentAction"
|
||||
],
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"tool",
|
||||
"tool_input",
|
||||
"log"
|
||||
]
|
||||
},
|
||||
"AgentFinish": {
|
||||
"title": "AgentFinish",
|
||||
"description": "The final return value of an ActionAgent.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"return_values": {
|
||||
"title": "Return Values",
|
||||
"type": "object"
|
||||
},
|
||||
"log": {
|
||||
"title": "Log",
|
||||
"type": "string"
|
||||
},
|
||||
"type": {
|
||||
"title": "Type",
|
||||
"default": "AgentFinish",
|
||||
"enum": [
|
||||
"AgentFinish"
|
||||
],
|
||||
"type": "string"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"return_values",
|
||||
"log"
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain",
|
||||
"prompts",
|
||||
"prompt",
|
||||
"PromptTemplate"
|
||||
],
|
||||
"name": "PromptTemplate"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain_community",
|
||||
"llms",
|
||||
"fake",
|
||||
"FakeStreamingListLLM"
|
||||
],
|
||||
"name": "FakeStreamingListLLM"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain_core",
|
||||
"runnables",
|
||||
"base",
|
||||
"RunnableLambda"
|
||||
],
|
||||
"name": "agent_parser"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"langchain",
|
||||
"schema",
|
||||
"runnable",
|
||||
"RunnablePassthrough"
|
||||
],
|
||||
"name": "RunnablePassthrough"
|
||||
}
|
||||
},
|
||||
{
|
||||
"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": 4,
|
||||
"target": 5
|
||||
},
|
||||
{
|
||||
"source": 5,
|
||||
"target": 6
|
||||
},
|
||||
{
|
||||
"source": 2,
|
||||
"target": 4
|
||||
},
|
||||
{
|
||||
"source": 6,
|
||||
"target": 3
|
||||
},
|
||||
{
|
||||
"source": 2,
|
||||
"target": 7
|
||||
},
|
||||
{
|
||||
"source": 7,
|
||||
"target": 3
|
||||
},
|
||||
{
|
||||
"source": "tools",
|
||||
"target": 2
|
||||
},
|
||||
{
|
||||
"source": 3,
|
||||
"target": "agent_should_continue"
|
||||
},
|
||||
{
|
||||
"source": "agent_should_continue",
|
||||
"target": "tools",
|
||||
"data": "continue"
|
||||
},
|
||||
{
|
||||
"source": "agent_should_continue",
|
||||
"target": "__end__",
|
||||
"data": "exit"
|
||||
},
|
||||
{
|
||||
"source": "__start__",
|
||||
"target": 2
|
||||
}
|
||||
]
|
||||
}
|
||||
'''
|
||||
# ---
|
||||
# name: test_conditional_graph.3
|
||||
'''
|
||||
+-----------+
|
||||
| __start__ |
|
||||
+-----------+
|
||||
*
|
||||
*
|
||||
*
|
||||
+------------------------------+
|
||||
| Parallel<agent_outcome>Input |
|
||||
+------------------------------+
|
||||
****** * *****
|
||||
****** * *****
|
||||
*** * ******
|
||||
+----------------+ * ***
|
||||
| PromptTemplate | * *
|
||||
+----------------+ * *
|
||||
* * *
|
||||
* * *
|
||||
* * *
|
||||
+----------------------+ * *
|
||||
| FakeStreamingListLLM | * *
|
||||
+----------------------+ * *
|
||||
* * *
|
||||
* * *
|
||||
* * *
|
||||
+----------------------+ +-------------+ *
|
||||
| Lambda(agent_parser) | | Passthrough | *
|
||||
+----------------------+ +-------------+ *
|
||||
** ** *
|
||||
*** *** *
|
||||
** ** *
|
||||
+-------------------------------+ *
|
||||
| Parallel<agent_outcome>Output | *
|
||||
+-------------------------------+ *
|
||||
* *
|
||||
* *
|
||||
* *
|
||||
+-----------------------+ *
|
||||
| agent_should_continue | *
|
||||
+-----------------------+******** *
|
||||
* ******** *
|
||||
* ******** *
|
||||
* ******** *
|
||||
+---------+ ***+-------+
|
||||
| __end__ | | tools |
|
||||
+---------+ +-------+
|
||||
'''
|
||||
# ---
|
||||
# name: test_conditional_graph_state
|
||||
'{"title": "LangGraphInput", "$ref": "#/definitions/AgentState", "definitions": {"AgentAction": {"title": "AgentAction", "description": "A full description of an action for an ActionAgent to execute.", "type": "object", "properties": {"tool": {"title": "Tool", "type": "string"}, "tool_input": {"title": "Tool Input", "anyOf": [{"type": "string"}, {"type": "object"}]}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentAction", "enum": ["AgentAction"], "type": "string"}}, "required": ["tool", "tool_input", "log"]}, "AgentFinish": {"title": "AgentFinish", "description": "The final return value of an ActionAgent.", "type": "object", "properties": {"return_values": {"title": "Return Values", "type": "object"}, "log": {"title": "Log", "type": "string"}, "type": {"title": "Type", "default": "AgentFinish", "enum": ["AgentFinish"], "type": "string"}}, "required": ["return_values", "log"]}, "AgentState": {"title": "AgentState", "type": "object", "properties": {"input": {"title": "Input", "type": "string"}, "agent_outcome": {"title": "Agent Outcome", "anyOf": [{"$ref": "#/definitions/AgentAction"}, {"$ref": "#/definitions/AgentFinish"}]}, "intermediate_steps": {"title": "Intermediate Steps", "type": "array", "items": {"type": "array", "minItems": 2, "maxItems": 2, "items": [{"$ref": "#/definitions/AgentAction"}, {"type": "string"}]}}}}}}'
|
||||
# ---
|
||||
@@ -1470,6 +1765,7 @@
|
||||
"type": "runnable",
|
||||
"data": {
|
||||
"id": [
|
||||
"tests",
|
||||
"test_pregel",
|
||||
"FakeFuntionChatModel"
|
||||
],
|
||||
|
||||
@@ -753,6 +753,8 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None:
|
||||
|
||||
assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot
|
||||
assert app.get_graph().draw_ascii() == snapshot
|
||||
assert json.dumps(app.get_graph(xray=True).to_json(), indent=2) == snapshot
|
||||
assert app.get_graph(xray=True).draw_ascii() == snapshot
|
||||
|
||||
assert app.invoke({"input": "what is weather in sf"}) == {
|
||||
"input": "what is weather in sf",
|
||||
|
||||
Reference in New Issue
Block a user