Refactor validation a bit and add test.

Author: superdosh

src/modelplane/evaluator/dag.py

@@ -91,8 +91,14 @@ def _validate_and_build(self) -> None:
         if self._validated:
             return
 
-        # check that all route targets reference registered nodes or instances of the output type
+        # check that all route targets reference registered nodes or instances
+        # of the output type, and that all Arbiters have compatible output types
         for node_name, node in self._nodes.items():
+            if isinstance(node, Arbiter):
+                if not issubclass(node.output_type, self.output_type):
+                    raise ValueError(
+                        f"Node {node_name} is an Arbiter with output_type {node.output_type.__name__}, which is not compatible with the DAG's output_type {self.output_type.__name__}."
+                    )
             for target in node.all_routes():
                 if target not in self._nodes and not isinstance(
                     target, self.output_type
@@ -127,16 +133,6 @@ def _validate_and_build(self) -> None:
             nodes_in_cycle = set(self._nodes) - set(ordered)
             raise ValueError(f"DAG contains a cycle. Nodes in cycle: {nodes_in_cycle}")
 
-        # check all terminal Arbiter nodes have correct output types
-        terminal_nodes = [n for n in self._nodes if not all_routes.get(n)]
-        for terminal in terminal_nodes:
-            node = self._nodes[terminal]
-            if isinstance(node, Arbiter):
-                if not issubclass(node.output_type, self.output_type):
-                    raise ValueError(
-                        f"Terminal node {terminal} has output_type {node.output_type.__name__}, which is not compatible with the DAG's output_type {self.output_type.__name__}."
-                    )
-
         # build predecessors
         for name, node in self._nodes.items():
             for target in node.all_routes():

tests/unit/evaluator/conftest.py

@@ -22,6 +22,7 @@
     PromptLengthGate,
     ThresholdArbiter,
     UnexpectedArbiter,
+    UnexpectedOutput,
     UpperCaser,
     UpperCaseScorer,
 )
@@ -183,3 +184,18 @@ def bad_dag_with_bad_arbiter():
     dag = EvaluatorDAG("test", output_type=Safety)
     dag.add_node(BadArbiter(name="bad_arbiter"))
     return dag
+
+
+@pytest.fixture
+def bad_one_step_dag():
+    return (
+        EvaluatorDAG("one_step", output_type=Safety)
+        .add_node(
+            AlwaysFalse(
+                name="gate",
+                routes_true=[UnexpectedOutput()],
+                routes_false=["always_unsafe"],
+            )
+        )
+        .add_node(AlwaysUnsafe(name="always_unsafe"))
+    )

tests/unit/evaluator/test_dag.py

@@ -55,6 +55,14 @@ def test_dag_with_bad_arbiter(bad_dag_with_bad_arbiter, sample_ctx):
         bad_dag_with_bad_arbiter.run(sample_ctx)
 
 
+def test_dag_with_bad_output_route(bad_one_step_dag, sample_ctx):
+    with pytest.raises(
+        ValueError,
+        match=r"incompatible output",
+    ):
+        bad_one_step_dag.run(sample_ctx)
+
+
 def test_dag_run(simple_dag, sample_ctx):
     result = simple_dag.run(sample_ctx)
     assert result.name == "UNSAFE"