From f2803372ec83790bc7d7e613aa2eb2eda7c81b06 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 24 Apr 2024 11:19:40 -0700 Subject: [PATCH] Improve graph validation - Move all validation to compile() This allows adding edges before nodes - Detect more cases of missing edges with shorthand branches --- langgraph/graph/graph.py | 67 +++++++++++++++---------------- tests/test_pregel.py | 87 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 120 insertions(+), 34 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 89004d23c..9415ef597 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -128,11 +128,6 @@ class Graph: raise ValueError("END cannot be a start node") if end_key == START: raise ValueError("START cannot be an end node") - if start_key not in self.nodes and start_key != START: - raise ValueError(f"Need to add_node `{start_key}` first") - if end_key not in self.nodes and end_key != END: - raise ValueError(f"Need to add_node `{end_key}` first") - if not self.support_multiple_edges and start_key in set( start for start, _ in self.edges ): @@ -180,17 +175,6 @@ class Graph: condition = coerce_to_runnable(condition) name = condition.name or "condition" # validate the condition - if start_key not in self.nodes and start_key != START: - raise ValueError(f"Need to add_node `{start_key}` first") - if conditional_edge_mapping and set( - conditional_edge_mapping.values() - ).difference([END]).difference(self.nodes): - raise ValueError( - f"Missing nodes which are in conditional edge mapping. Mapping " - f"contains possible destinations: " - f"{list(conditional_edge_mapping.values())}. Possible nodes are " - f"{list(self.nodes.keys())}." - ) if name in self.branches[start_key]: raise ValueError( f"Branch with name `{condition.name}` already exists for node " @@ -242,31 +226,46 @@ class Graph: return self.add_edge(key, END) def validate(self, interrupt: Optional[Sequence[str]] = None) -> None: - all_starts = {src for src, _ in self._all_edges} | { + # assemble sources + all_sources = {src for src, _ in self._all_edges} | { src for src in self.branches } + # validate sources for node in self.nodes: - if node not in all_starts: - raise ValueError(f"Node `{node}` is a dead-end") - - all_branches = [ - branch - for branches in self.branches.values() - for branch in branches.values() - ] - if all(branch.ends is not None for branch in all_branches): - all_ends = {end for _, end in self._all_edges} | { - end for branch in all_branches for end in branch.ends.values() - } - - for node in self.nodes: - if node not in all_ends: - raise ValueError(f"Node `{node}` is not reachable") + if node not in all_sources: + raise ValueError(f"Node '{node}' is a dead-end") + for source in all_sources: + if node not in self.nodes and node != START: + raise ValueError(f"Found edge starting at unkown node '{source}'") + # assemble targets + all_targets = {end for _, end in self._all_edges} + for start, branches in self.branches.items(): + for cond, branch in branches.items(): + if branch.ends is not None: + for end in branch.ends.values(): + if end not in self.nodes and end != END: + raise ValueError( + f"At '{start}' node, '{cond}' branch found unknown target '{end}'" + ) + all_targets.add(end) + else: + all_targets.add(END) + for node in self.nodes: + if node != start: + all_targets.add(node) + # validate targets + for node in self.nodes: + if node not in all_targets: + raise ValueError(f"Node `{node}` is not reachable") + for target in all_targets: + if target not in self.nodes and target != END: + raise ValueError(f"Found edge ending at unknown node `{target}`") + # validate interrupts if interrupt: for node in interrupt: if node not in self.nodes: - raise ValueError(f"Node `{node}` is not present") + raise ValueError(f"Interrupt node `{node}` not found") self.compiled = True diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 1ca671249..4aeb543d4 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -31,6 +31,93 @@ from tests.any_str import AnyStr from tests.memory_assert import MemorySaverAssertImmutable +def test_graph_validation() -> None: + def logic(inp: str) -> str: + return "" + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.set_entry_point("agent") + workflow.set_finish_point("agent") + assert workflow.compile(), "valid graph" + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.set_entry_point("agent") + with pytest.raises(ValueError, match="dead-end"): + workflow.compile() + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.set_finish_point("agent") + with pytest.raises(ValueError, match="not reachable"): + workflow.compile() + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + workflow.set_entry_point("agent") + workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) + workflow.add_edge("tools", "agent") + assert workflow.compile(), "valid graph" + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + workflow.set_entry_point("tools") + workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) + workflow.add_edge("tools", "agent") + assert workflow.compile(), "valid graph" + + workflow = Graph() + workflow.set_entry_point("tools") + workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) + workflow.add_edge("tools", "agent") + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + assert workflow.compile(), "valid graph" + + workflow = Graph() + workflow.set_entry_point("tools") + workflow.add_conditional_edges( + "agent", logic, {"continue": "tools", "exit": END, "hmm": "extra"} + ) + workflow.add_edge("tools", "agent") + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + with pytest.raises(ValueError, match="unknown"): # extra is not defined + workflow.compile() + + workflow = Graph() + workflow.set_entry_point("agent") + workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) + workflow.add_edge("tools", "extra") + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + with pytest.raises(ValueError, match="unknown"): # extra is not defined + workflow.compile() + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + workflow.add_node("extra", logic) + workflow.set_entry_point("agent") + workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END}) + workflow.add_edge("tools", "agent") + with pytest.raises(ValueError): # extra is dead-end / not reachable + workflow.compile() + + workflow = Graph() + workflow.add_node("agent", logic) + workflow.add_node("tools", logic) + workflow.add_node("extra", logic) + workflow.set_entry_point("agent") + workflow.add_conditional_edges("agent", logic) + workflow.add_edge("tools", "agent") + with pytest.raises(ValueError): # extra is dead-end + workflow.compile() + + def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")