diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 173798b3d..89004d23c 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -1,5 +1,5 @@ import logging -from collections import defaultdict +from collections import Counter, defaultdict from typing import ( Any, Awaitable, @@ -402,6 +402,9 @@ class CompiledGraph(Pregel): end_nodes[key] = n for start, end in sorted(self.graph._all_edges): graph.add_edge(start_nodes[start], end_nodes[end]) + branches_by_name = Counter( + name for _, branches in self.graph.branches.items() for name in branches + ) for start, branches in self.graph.branches.items(): for name, branch in branches.items(): ends = branch.ends or { @@ -410,7 +413,10 @@ class CompiledGraph(Pregel): } if add_condition_nodes is True: - cond = graph.add_node(branch.condition, name) + cond = graph.add_node( + branch.condition, + f"{start}_{name}" if branches_by_name[name] > 1 else name, + ) graph.add_edge(start_nodes[start], cond) for label, end in ends.items(): graph.add_edge(cond, end_nodes[end], label, conditional=True) diff --git a/tests/__snapshots__/test_pregel.ambr b/tests/__snapshots__/test_pregel.ambr index eca8025a3..3605c0071 100644 --- a/tests/__snapshots__/test_pregel.ambr +++ b/tests/__snapshots__/test_pregel.ambr @@ -1657,7 +1657,7 @@ ''' # --- # 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"]}}}' + '{"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"}}, "is_last_step": {"title": "Is Last Step", "type": "boolean"}}, "required": ["messages", "is_last_step"]}}}' # --- # name: test_prebuilt_chat.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "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"]}}}' @@ -1771,7 +1771,7 @@ ''' # --- # name: test_prebuilt_tool_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"]}}}' + '{"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"}}, "is_last_step": {"title": "Is Last Step", "type": "boolean"}}, "required": ["messages", "is_last_step"]}}}' # --- # name: test_prebuilt_tool_chat.1 '{"title": "LangGraphOutput", "type": "object", "properties": {"messages": {"title": "Messages", "type": "array", "items": {"$ref": "#/definitions/BaseMessage"}}}, "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"]}}}' diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 25accdf50..1ca671249 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -3534,3 +3534,42 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert app.invoke({"my_key": "my value"}) == { "my_key": "my value there and back again" } + + +def test_repeat_condition(snapshot: SnapshotAssertion) -> None: + class AgentState(TypedDict): + hello: str + + def router(state: AgentState) -> str: + return "hmm" + + workflow = StateGraph(AgentState) + workflow.add_node("Researcher", lambda x: x) + workflow.add_node("Chart Generator", lambda x: x) + workflow.add_node("Call Tool", lambda x: x) + workflow.add_conditional_edges( + "Researcher", + router, + {"continue": "Chart Generator", "call_tool": "Call Tool", "end": END}, + ) + workflow.add_conditional_edges( + "Chart Generator", + router, + {"continue": "Researcher", "call_tool": "Call Tool", "end": END}, + ) + workflow.add_conditional_edges( + "Call Tool", + # Each agent node updates the 'sender' field + # the tool calling node does not, meaning + # this edge will route back to the original agent + # who invoked the tool + lambda x: x["sender"], + { + "Researcher": "Researcher", + "Chart Generator": "Chart Generator", + }, + ) + workflow.set_entry_point("Researcher") + + app = workflow.compile() + assert app.get_graph().draw_mermaid(with_styles=False) == snapshot