diff --git a/py/src/braintrust/framework.py b/py/src/braintrust/framework.py index 3b4c2e52..4099a552 100644 --- a/py/src/braintrust/framework.py +++ b/py/src/braintrust/framework.py @@ -1236,10 +1236,14 @@ def __init__( self.update({"expected": expected}) self.update({"trial_index": trial_index}) self._span = None - if tags is not None: - self.update({"tags": tags}) - else: + if tags is None: self.update({"tags": []}) + elif isinstance(tags, (list, tuple, set)): + # Copy so hook mutations don't leak into the caller's tags. + self.update({"tags": list(tags)}) + else: + # Leave other values as-is instead of splitting a string or dict into tags. + self.update({"tags": tags}) self._report_progress = report_progress self._parameters = parameters @@ -1756,6 +1760,7 @@ async def run_evaluator_task(datum, trial_index=0): # In most cases this will be a no-op span, but if the parent is set, it will use that ctx. root_span = start_span(state=state, **base_event) + hooks = None with root_span: try: @@ -1950,7 +1955,10 @@ async def ensure_spans_flushed(): root_span.log(metadata=metadata, tags=tags) except Exception as e: exc_type, exc_value, tb = sys.exc_info() - root_span.log(error=stringify_exception(exc_type, exc_value, tb)) + if hooks is not None: + # Keep tags the task added before it failed. + tags = hooks.tags if hooks.tags else None + root_span.log(error=stringify_exception(exc_type, exc_value, tb), tags=tags) error = e # Python3.10 has a different set of arguments to format_exception than earlier versions, diff --git a/py/src/braintrust/test_framework.py b/py/src/braintrust/test_framework.py index c3284e53..e01076cb 100644 --- a/py/src/braintrust/test_framework.py +++ b/py/src/braintrust/test_framework.py @@ -1,4 +1,5 @@ import asyncio +import copy import importlib.util import re import sys @@ -1444,6 +1445,91 @@ def task_with_hooks(input, hooks): assert root_span[0].get("tags") == expected_tags +@pytest.mark.parametrize("use_trials", [False, True], ids=["shared-row-tags", "trials"]) +@pytest.mark.parametrize("initial_tags", [None, ["seed"], ("seed",), {"seed"}], ids=["none", "list", "tuple", "set"]) +def test_hooks_tags_mutation_is_isolated_per_evaluation(with_memory_logger, use_trials, initial_tags): + original_tags = copy.copy(initial_tags) + base_tags = list(initial_tags or []) + inputs = ["row"] if use_trials else ["first", "second"] + + def task(input, hooks): + own_tag = f"{input}:{hooks.trial_index}" + hooks.tags.append(own_tag) + return own_tag + + def tag_isolation(output, tags): + return float(tags == [*base_tags, output]) + + result = Eval( + "isolated-hook-tags", + data=[EvalCase(input=value, tags=initial_tags) for value in inputs], + task=task, + scores=[tag_isolation], + trial_count=2 if use_trials else 1, + max_concurrency=1, + no_send_logs=True, + enable_cache=False, + ) + + assert [row.error for row in result.results] == [None, None] + assert [row.tags for row in result.results] == [[*base_tags, row.output] for row in result.results] + assert [row.scores["tag_isolation"] for row in result.results] == [1.0, 1.0] + assert initial_tags == original_tags + + +@pytest.mark.asyncio +async def test_hooks_tags_are_kept_when_task_raises(with_memory_logger, with_simulate_login): + def task(input, hooks): + hooks.tags.append("failed-step") + raise ValueError("task failed") + + evaluator = Evaluator( + project_name="test-project", + eval_name="test-evaluator", + data=[EvalCase(input="row", tags=["seed"])], + task=task, + scores=[], + experiment_name=None, + metadata=None, + summarize_scores=False, + ) + + result = await run_evaluator( + experiment=init_test_exp("test-evaluator", "test-project"), + evaluator=evaluator, + position=None, + filters=[], + ) + + assert isinstance(result.results[0].error, ValueError) + assert result.results[0].tags == ["seed", "failed-step"] + root_spans = [log for log in with_memory_logger.pop() if not log["span_parents"]] + assert [span.get("tags") for span in root_spans] == [["seed", "failed-step"]] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("tags", ["prod", {"prod": True}], ids=["str", "dict"]) +async def test_non_sequence_eval_case_tags_are_rejected(with_memory_logger, with_simulate_login, tags): + evaluator = Evaluator( + project_name="test-project", + eval_name="test-evaluator", + data=[EvalCase(input="row", tags=tags)], + task=lambda input: input, + scores=[], + experiment_name=None, + metadata=None, + summarize_scores=False, + ) + + with pytest.raises(ValueError, match="tags must be a list, set, or tuple of strings"): + await run_evaluator( + experiment=init_test_exp("test-evaluator", "test-project"), + evaluator=evaluator, + position=None, + filters=[], + ) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("tags", "expected_tags"),