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

@@ -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
 

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")