From 52f87321f9cffeee6bf0752179e42cfea10752de Mon Sep 17 00:00:00 2001
From: Vishal Doshi <vishal@mlcommons.org>
Date: Wed, 10 Jun 2026 09:17:17 -0400
Subject: [PATCH] Add method for metadata updates in EvalContext.

---
 src/modelplane/evaluator/context.py  | 17 +++++++++++++++++
 tests/unit/evaluator/test_context.py | 18 +++++++++++++++++-
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/src/modelplane/evaluator/context.py b/src/modelplane/evaluator/context.py
index 04d575d..1c2546b 100644
--- a/src/modelplane/evaluator/context.py
+++ b/src/modelplane/evaluator/context.py
@@ -87,6 +87,23 @@ def with_response(self, new_response: str) -> EvalContext:
         )
 
     def with_metadata(self, new_metadata: dict[str, Any]) -> EvalContext:
+        """
+        Return a new EvalContext with the provided metadata replacing the
+        original metadata.
+        """
+        return EvalContext(
+            prompt=self.prompt,
+            response=self.response,
+            metadata=new_metadata,
+        )
+
+    def with_metadata_updates(self, updates: dict[str, Any]) -> EvalContext:
+        """
+        Return a new EvalContext with the original metadata updated with the
+        provided updates.
+        """
+        new_metadata = self.metadata.copy()
+        new_metadata.update(updates)
         return EvalContext(
             prompt=self.prompt,
             response=self.response,
diff --git a/tests/unit/evaluator/test_context.py b/tests/unit/evaluator/test_context.py
index 31ec8f9..7c6fdb0 100644
--- a/tests/unit/evaluator/test_context.py
+++ b/tests/unit/evaluator/test_context.py
@@ -27,6 +27,18 @@ def test_metadata_replacement(sample_ctx):
     assert new_ctx.metadata == new_metadata
 
 
+def test_metadata_updates():
+    original_metadata = {"key1": "value1"}
+    sample_ctx = EvalContext(
+        prompt="Prompt", response="Response", metadata=original_metadata
+    )
+    additional_metadata = {"key2": "value2"}
+    new_ctx = sample_ctx.with_metadata_updates(additional_metadata)
+    assert new_ctx.prompt == sample_ctx.prompt
+    assert new_ctx.response == sample_ctx.response
+    assert new_ctx.metadata == {**sample_ctx.metadata, **additional_metadata}
+
+
 def test_with_updates(sample_ctx):
     new_prompt = "Updated prompt"
     new_response = "Updated response"
@@ -69,5 +81,9 @@ def test_hash(sample_ctx):
     other_ctx = EvalContext(prompt=sample_ctx.prompt, response=sample_ctx.response)
     assert sample_ctx.hash() == other_ctx.hash()
 
-    other_ctx = EvalContext(prompt=sample_ctx.prompt, response=sample_ctx.response, metadata={"key": "value"})
+    other_ctx = EvalContext(
+        prompt=sample_ctx.prompt,
+        response=sample_ctx.response,
+        metadata={"key": "value"},
+    )
     assert sample_ctx.hash() != other_ctx.hash()