Merge pull request #147 from mlcommons/source-node-in-exception

node info in exception

Author: bkorycki

src/modelplane/evaluator/dag.py

@@ -45,6 +45,15 @@ class FailedDAGOutput(_DAGOutput):
     error: Exception
 
 
+class NodeExecutionError(Exception):
+    def __init__(self, node_name: str, original_error: Exception):
+        self.node_name = node_name
+        self.original_error = original_error
+        super().__init__(
+            f"Error while executing node '{node_name}': {original_error}"
+        )
+
+
 class Composer:
     """DAG of ComposerNodes.
 
@@ -225,9 +234,10 @@ def _run_traced(
             try:
                 output = self._run_node(node, ctx)
             except Exception as e:
+                wrapped_error = NodeExecutionError(node.name, e)
                 return (
                     FailedDAGOutput(
-                        node_outputs=node_outputs, total_cost=total_cost, error=e
+                        node_outputs=node_outputs, total_cost=total_cost, error=wrapped_error
                     ),
                     traversed_edges,
                 )

tests/unit/evaluator/test_safety.py

@@ -4,7 +4,7 @@
 from modelgauge.prompt import TextPrompt
 from modelgauge.sut import SUTResponse
 
-from modelplane.evaluator.dag import Composer, FailedDAGOutput
+from modelplane.evaluator.dag import Composer, FailedDAGOutput, NodeExecutionError
 from modelplane.evaluator.safety import AnnotatorArbiter, Safety, SafetyDAGAnnotator
 from modelplane.evaluator.verdict import Verdict
 
@@ -50,13 +50,14 @@ def test_safety_dag_with_bad_node(sample_ctx, threshold_arbiter):
     )
     dag_output = dag.run(sample_ctx)
     assert isinstance(dag_output, FailedDAGOutput)
-    assert str(dag_output.error) == "I'm afraid I can't do that, Dave."
+    assert str(dag_output.error.original_error) == "I'm afraid I can't do that, Dave."
 
     dag_annotator = SafetyDAGAnnotator("safety_annotator", dag)
     with pytest.raises(
-        type(dag_output.error), match="I'm afraid I can't do that, Dave."
-    ):
+        NodeExecutionError, match="Error while executing node 'failing_node': I'm afraid I can't do that, Dave."
+    ) as e:
         dag_annotator.process(
             prompt=TextPrompt(text=sample_ctx.prompt),
             response=SUTResponse(text=sample_ctx.response),
         )
+        assert type(e.value.original_error) == type(dag_output.error)