Add original_ctx to NodeOutput for context tracking.

src/modelplane/evaluator/context.py

@@ -9,6 +9,7 @@
 @dataclass
 class NodeOutput:
     value: Any
+    original_ctx: EvalContext
     realized_cost: RealizedCost = field(default_factory=RealizedCost)
     updated_ctx: Optional[EvalContext] = None
 
@@ -17,6 +18,7 @@ def to_dict(self) -> dict:
             "value": str(self.value),
             "realized_cost": self.realized_cost.to_dict(),
             "updated_ctx": self.updated_ctx.to_dict() if self.updated_ctx else None,
+            "original_ctx": self.original_ctx.to_dict(),
         }
 
 

src/modelplane/evaluator/nodes.py

@@ -63,6 +63,7 @@ def build_output(
         return NodeOutput(
             value=value,
             realized_cost=self.realized_cost(ctx),
+            original_ctx=ctx,
             updated_ctx=updated_ctx,
         )
 

src/modelplane/evaluator/safety.py

@@ -71,8 +71,8 @@ def run(self, ctx: EvalContext) -> NodeOutput:
         if key in self._cache:
             val = self._cache[key]
             assert isinstance(val, Safety)
-            return NodeOutput(value=val)
+            return NodeOutput(value=val, original_ctx=ctx)
         else:
             val = self._run(ctx)
             self._cache[key] = val
-            return NodeOutput(value=val)
+            return NodeOutput(value=val, original_ctx=ctx)

tests/unit/evaluator/test_context.py

@@ -45,10 +45,12 @@ def test_with_different_parent_outputs(sample_ctx):
     parent_outputs = {
         "parent1": NodeOutput(
             value="output1",
+            original_ctx=sample_ctx,
             updated_ctx=sample_ctx.with_response("Updated response for parent1"),
         ),
         "parent2": NodeOutput(
             value="output2",
+            original_ctx=sample_ctx,
             updated_ctx=sample_ctx.with_prompt("Updated prompt for parent2"),
         ),
     }

tests/unit/evaluator/test_dag.py

@@ -93,7 +93,7 @@ def test_dag_run_with_dataframe(simple_dag, tmp_path):
     df = pd.DataFrame(
         {
             "prompt": ["a", "ab", "abc", "abcd"],  # odd, even, odd, even
-            "response": ["hello world", "helloworld", "hello world", "helloworld"],
+            "response": ["Hello world", "Helloworld", "Hello world", "Helloworld"],
         }
     )
     result_df = simple_dag.run_dataframe(df)
@@ -117,6 +117,12 @@ def test_dag_run_with_dataframe(simple_dag, tmp_path):
     for dag_run_json in result_df[simple_dag.df_dag_run_col]:
         dag_run = json.loads(dag_run_json)
         assert "always_true" in dag_run
+        if "lower_caser" in dag_run:
+            # confirm original ctx is present and different from updated ctx
+            assert (
+                dag_run["lower_caser"]["original_ctx"]["response"].lower()
+                == dag_run["lower_caser"]["updated_ctx"]["response"]
+            )
 
 
 def test_dag_run_with_dataframe_parallel(simple_dag):

tests/unit/evaluator/test_nodes.py

@@ -36,8 +36,8 @@ def test_fixed_scorer(sample_ctx, score_1):
 
 def test_consistent_arbiters(sample_ctx, score_1, score_2, always_unsafe, always_safe):
     parent_outputs = {
-        score_1.name: NodeOutput(value=SCORE1),
-        score_2.name: NodeOutput(value=SCORE2),
+        score_1.name: NodeOutput(value=SCORE1, original_ctx=sample_ctx),
+        score_2.name: NodeOutput(value=SCORE2, original_ctx=sample_ctx),
     }
     run_ctx = sample_ctx.with_parent_outputs(parent_outputs)
     output = always_unsafe.run(run_ctx)
@@ -48,15 +48,21 @@ def test_consistent_arbiters(sample_ctx, score_1, score_2, always_unsafe, always
 
 def test_threshold_arbiter_true(sample_ctx, threshold_arbiter):
     run_ctx = sample_ctx.with_parent_outputs(
-        {"parent0": NodeOutput(value=SCORE2), "parent1": NodeOutput(value=SCORE2)}
+        {
+            "parent0": NodeOutput(value=SCORE2, original_ctx=sample_ctx),
+            "parent1": NodeOutput(value=SCORE2, original_ctx=sample_ctx),
+        }
     )
     output = threshold_arbiter.run(run_ctx)
     assert output.value.name == "UNSAFE"
 
 
 def test_threshold_arbiter_false(sample_ctx, threshold_arbiter):
     run_ctx = sample_ctx.with_parent_outputs(
-        {"parent0": NodeOutput(value=SCORE1), "parent1": NodeOutput(value=SCORE1)}
+        {
+            "parent0": NodeOutput(value=SCORE1, original_ctx=sample_ctx),
+            "parent1": NodeOutput(value=SCORE1, original_ctx=sample_ctx),
+        }
     )
     output = threshold_arbiter.run(run_ctx)
     assert output.value.name == "SAFE"