From a0425dd19fc1064254c5de81352eaa36a44d3ee6 Mon Sep 17 00:00:00 2001 From: Barbara Korycki Date: Tue, 26 May 2026 14:40:01 -0700 Subject: [PATCH] more info in excpetion --- src/modelplane/evaluator/dag.py | 12 +++++++++++- tests/unit/evaluator/test_safety.py | 9 +++++---- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/src/modelplane/evaluator/dag.py b/src/modelplane/evaluator/dag.py index 891c984..455a7e6 100644 --- a/src/modelplane/evaluator/dag.py +++ b/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, ) diff --git a/tests/unit/evaluator/test_safety.py b/tests/unit/evaluator/test_safety.py index 56a8a38..e714723 100644 --- a/tests/unit/evaluator/test_safety.py +++ b/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)