diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 14860dd9d..72ead903b 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -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 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index ff507145a..fb5f2b850 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -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": "RunnableParallelInput", + "type": "object", + "properties": {} + } + }, + { + "id": 3, + "type": "schema", + "data": { + "title": "RunnableParallelOutput", + "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__ | + +-----------+ + * + * + * + +------------------------------+ + | ParallelInput | + +------------------------------+ + ****** * ***** + ****** * ***** + *** * ****** + +----------------+ * *** + | PromptTemplate | * * + +----------------+ * * + * * * + * * * + * * * + +----------------------+ * * + | FakeStreamingListLLM | * * + +----------------------+ * * + * * * + * * * + * * * + +----------------------+ +-------------+ * + | Lambda(agent_parser) | | Passthrough | * + +----------------------+ +-------------+ * + ** ** * + *** *** * + ** ** * + +-------------------------------+ * + | ParallelOutput | * + +-------------------------------+ * + * * + * * + * * + +-----------------------+ * + | 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" ], diff --git a/tests/test_pregel.py b/tests/test_pregel.py index b943c4506..c57dabff0 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -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",