From 0f667cd9ebf1c6b5a3b19e322bdafd7c460c5bce Mon Sep 17 00:00:00 2001 From: Angel Igareta Date: Mon, 15 Apr 2024 10:56:19 +0200 Subject: [PATCH 1/2] Add add_condition_nodes parameter to offer possibility of not including conditional nodes --- langgraph/graph/graph.py | 15 ++-- tests/__snapshots__/test_pregel.ambr | 128 ++++++++++++++++++++++++++- tests/test_pregel.py | 2 + 3 files changed, 138 insertions(+), 7 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 5830e3a9c..d3a11cedc 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -319,7 +319,7 @@ class CompiledGraph(Pregel): self.nodes[end].channels.append(channel_name) def get_graph( - self, config: Optional[RunnableConfig] = None, *, xray: bool = False + self, config: Optional[RunnableConfig] = None, *, xray: bool = False, add_condition_nodes: bool = True ) -> RunnableGraph: """Returns a drawable representation of the computation graph.""" graph = RunnableGraph() @@ -355,13 +355,18 @@ class CompiledGraph(Pregel): graph.add_edge(start_nodes[start], end_nodes[end]) for start, branches in self.graph.branches.items(): for name, branch in branches.items(): - cond = graph.add_node(branch.condition, 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(cond, end_nodes[end], label) + + if add_condition_nodes is True: + cond = graph.add_node(branch.condition, name) + graph.add_edge(start_nodes[start], cond) + for label, end in ends.items(): + graph.add_edge(cond, end_nodes[end], label) + else: + for label, end in ends.items(): + graph.add_edge(start_nodes[start], end_nodes[end], label) return graph diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index 129044fcb..c9b710050 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -410,6 +410,25 @@ ''' # --- # name: test_conditional_graph[end_of_run].2 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * * + ** ** + * * + +-------+ +---------+ + | tools | | __end__ | + +-------+ +---------+ + ''' +# --- +# name: test_conditional_graph[end_of_run].3 ''' { "nodes": [ @@ -563,7 +582,7 @@ } ''' # --- -# name: test_conditional_graph[end_of_run].3 +# name: test_conditional_graph[end_of_run].4 ''' +-----------+ | __start__ | @@ -612,6 +631,49 @@ +---------+ +-------+ ''' # --- +# name: test_conditional_graph[end_of_run].5 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +------------------------------+ + | ParallelInput | + +------------------------------+ + ***** * ***** + ****** * **** + *** * ***** + +----------------+ * *** + | PromptTemplate | * * + +----------------+ * * + * * * + * * * + * * * + +----------------------+ * * + | FakeStreamingListLLM | * * + +----------------------+ * * + * * * + * * * + * * * + +----------------------+ +-------------+ * + | Lambda(agent_parser) | | Passthrough | * + +----------------------+ +-------------+ * + *** *** * + ** ** * + ** ** * + +-------------------------------+ * + | ParallelOutput | * + +-------------------------------+***** * + * ********* * + * ******** * + * ***** * + +---------+ +-------+ + | __end__ | | tools | + +---------+ +-------+ + ''' +# --- # name: test_conditional_graph[end_of_step] ''' { @@ -719,6 +781,25 @@ ''' # --- # name: test_conditional_graph[end_of_step].2 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +-------+ + | agent | + +-------+ + * * + ** ** + * * + +-------+ +---------+ + | tools | | __end__ | + +-------+ +---------+ + ''' +# --- +# name: test_conditional_graph[end_of_step].3 ''' { "nodes": [ @@ -872,7 +953,7 @@ } ''' # --- -# name: test_conditional_graph[end_of_step].3 +# name: test_conditional_graph[end_of_step].4 ''' +-----------+ | __start__ | @@ -921,6 +1002,49 @@ +---------+ +-------+ ''' # --- +# name: test_conditional_graph[end_of_step].5 + ''' + +-----------+ + | __start__ | + +-----------+ + * + * + * + +------------------------------+ + | ParallelInput | + +------------------------------+ + ***** * ***** + ****** * **** + *** * ***** + +----------------+ * *** + | PromptTemplate | * * + +----------------+ * * + * * * + * * * + * * * + +----------------------+ * * + | FakeStreamingListLLM | * * + +----------------------+ * * + * * * + * * * + * * * + +----------------------+ +-------------+ * + | Lambda(agent_parser) | | Passthrough | * + +----------------------+ +-------------+ * + *** *** * + ** ** * + ** ** * + +-------------------------------+ * + | ParallelOutput | * + +-------------------------------+***** * + * ********* * + * ******** * + * ***** * + +---------+ +-------+ + | __end__ | | tools | + +---------+ +-------+ + ''' +# --- # name: test_conditional_graph_state[end_of_run] '{"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"}]}}}}}}' # --- diff --git a/tests/test_pregel.py b/tests/test_pregel.py index d234ca4a9..a681a1fbd 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -796,8 +796,10 @@ def test_conditional_graph( assert json.dumps(app.get_graph().to_json(), indent=2) == snapshot assert app.get_graph().draw_ascii() == snapshot + assert app.get_graph(add_condition_nodes=False).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.get_graph(xray=True, add_condition_nodes=False).draw_ascii() == snapshot assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", From 4e8e51126d035d009cd962e3639848a01075f0f4 Mon Sep 17 00:00:00 2001 From: Angel Igareta Date: Wed, 17 Apr 2024 18:03:22 +0200 Subject: [PATCH 2/2] Fix formatting in graph --- langgraph/graph/graph.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 0d4a8a093..0be33a53b 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -364,7 +364,11 @@ class CompiledGraph(Pregel): self.nodes[end].channels.append(channel_name) def get_graph( - self, config: Optional[RunnableConfig] = None, *, xray: bool = False, add_condition_nodes: bool = True + self, + config: Optional[RunnableConfig] = None, + *, + xray: bool = False, + add_condition_nodes: bool = True, ) -> RunnableGraph: """Returns a drawable representation of the computation graph.""" graph = RunnableGraph() @@ -405,7 +409,7 @@ class CompiledGraph(Pregel): END: END, } - if add_condition_nodes is True: + if add_condition_nodes is True: cond = graph.add_node(branch.condition, name) graph.add_edge(start_nodes[start], cond) for label, end in ends.items():