From a876b2b6cda45453f063f268cdcaefe3bed86591 Mon Sep 17 00:00:00 2001 From: Curtis Galione Date: Wed, 5 Aug 2026 10:59:17 -0700 Subject: [PATCH 1/2] fix(framework): preserve remote eval row origins --- py/src/braintrust/framework.py | 81 +++++-- py/src/braintrust/test_framework.py | 205 +++++++++++++++++- .../type_tests/test_eval_generics.py | 23 ++ py/src/braintrust/types/_eval.py | 6 + 4 files changed, 299 insertions(+), 16 deletions(-) diff --git a/py/src/braintrust/framework.py b/py/src/braintrust/framework.py index 091f65e54..75ba13859 100644 --- a/py/src/braintrust/framework.py +++ b/py/src/braintrust/framework.py @@ -20,10 +20,12 @@ Protocol, TypedDict, TypeVar, + cast, ) from tqdm.asyncio import tqdm as async_tqdm from tqdm.auto import tqdm as std_tqdm +from typing_extensions import NotRequired from .generated_types import FunctionFormat, FunctionOutputType, ObjectReference from .git_fields import GitMetadataSettings, RepoInfo @@ -90,6 +92,7 @@ class EvalCase(SerializableDataClass, Generic[Input, Expected]): id: str | None = None _xact_id: str | None = None created: str | None = None + origin: ObjectReference | None = None # Inheritance doesn't quite work for dataclasses, so we redefine the fields @@ -107,6 +110,7 @@ class EvalResult(SerializableDataClass, Generic[Input, Output, Expected]): tags: list[str] | None = None error: Exception | None = None exc_info: str | None = None + origin: ObjectReference | None = None class TaskProgressEvent(TypedDict): @@ -135,7 +139,7 @@ class SSEProgressEvent(TaskProgressEvent): id: str object_type: str - origin: ObjectReference + origin: NotRequired[ObjectReference] name: str @@ -1372,6 +1376,40 @@ def _get_persisted_base_experiment_id(experiment: Experiment) -> str | None: return base_experiment_id if isinstance(base_experiment_id, str) and base_experiment_id else None +_OBJECT_REFERENCE_TYPES = frozenset({"project_logs", "experiment", "dataset", "prompt", "function", "prompt_session"}) + + +def _validated_object_reference(value: Any) -> ObjectReference | None: + if not isinstance(value, Mapping): + return None + + object_type = value.get("object_type") + object_id = value.get("object_id") + row_id = value.get("id") + if ( + not isinstance(object_type, str) + or object_type not in _OBJECT_REFERENCE_TYPES + or not isinstance(object_id, str) + or not isinstance(row_id, str) + ): + return None + + validated: dict[str, Any] = { + "object_type": object_type, + "object_id": object_id, + "id": row_id, + } + for optional_field in ("_xact_id", "created"): + if optional_field not in value: + continue + optional_value = value[optional_field] + if optional_value is not None and not isinstance(optional_value, str): + return None + validated[optional_field] = optional_value + + return cast(ObjectReference, validated) + + async def run_evaluator( experiment: Experiment | None, evaluator: Evaluator[Input, Output, Expected], @@ -1561,24 +1599,32 @@ async def run_evaluator_task(datum, trial_index=0): experiment.dataset if experiment else evaluator.data if isinstance(evaluator.data, Dataset) else None ) - origin = ( - { - "object_type": "dataset", - "object_id": event_dataset.id, - "id": datum.id, - "created": datum.created, - "_xact_id": datum._xact_id, - } - if event_dataset and datum.id and datum._xact_id + dataset_origin = ( + cast( + ObjectReference, + { + "object_type": "dataset", + "object_id": event_dataset.id, + "id": datum.id, + "_xact_id": datum._xact_id, + **({"created": datum.created} if isinstance(datum.created, str) else {}), + }, + ) + if event_dataset + and isinstance(datum.id, str) + and datum.id + and isinstance(datum._xact_id, str) + and datum._xact_id else None ) + origin = dataset_origin or _validated_object_reference(datum.origin) base_event = dict( name="eval", span_attributes={"type": SpanTypeAttribute.EVAL}, input=datum.input, expected=datum.expected, tags=tags, - origin=origin, + **({"origin": origin} if origin is not None else {}), ) if experiment: @@ -1593,11 +1639,15 @@ async def run_evaluator_task(datum, trial_index=0): def report_progress(event: TaskProgressEvent): if not stream: return - stream( - SSEProgressEvent( - id=root_span.id, origin=origin, name=evaluator.eval_name, object_type="task", **event - ) + progress = SSEProgressEvent( + id=root_span.id, + name=evaluator.eval_name, + object_type="task", + **event, ) + if origin is not None: + progress["origin"] = origin + stream(progress) hooks = DictEvalHooks( metadata, @@ -1786,6 +1836,7 @@ async def ensure_spans_flushed(): classifications=classifications or None, error=error, exc_info=exc_info, + origin=origin, ) data_iterator = evaluator.data diff --git a/py/src/braintrust/test_framework.py b/py/src/braintrust/test_framework.py index 5c000bdff..d566365c6 100644 --- a/py/src/braintrust/test_framework.py +++ b/py/src/braintrust/test_framework.py @@ -4,7 +4,8 @@ from unittest.mock import MagicMock, patch import pytest -from braintrust.logger import BraintrustState +from braintrust.logger import BraintrustState, Dataset, ObjectMetadata, ProjectDatasetMetadata +from braintrust.util import LazyValue from .framework import ( Eval, @@ -25,6 +26,208 @@ HAS_PYDANTIC = importlib.util.find_spec("pydantic") is not None +def make_dataset(dataset_id, row): + project_metadata = ObjectMetadata(id="test-project", name="test-project", full_info={}) + dataset_metadata = ObjectMetadata(id=dataset_id, name="test-dataset", full_info={}) + dataset = Dataset( + lazy_metadata=LazyValue( + lambda: ProjectDatasetMetadata(project=project_metadata, dataset=dataset_metadata), + use_mutex=False, + ), + state=BraintrustState(), + ) + return dataset, patch.object(dataset, "_refetch", return_value=[row]) + + +def progress_event(): + return { + "format": "code", + "output_type": "completion", + "event": "progress", + "data": "0.5", + } + + +def test_eval_case_from_dict_preserves_valid_origin(): + origin = { + "object_type": "project_logs", + "object_id": "source-project", + "id": "source-row", + "_xact_id": "source-xact", + "created": "2026-06-01T00:00:00.000Z", + } + + assert EvalCase.from_dict({"input": 1, "origin": origin}).origin == origin + + +@pytest.mark.asyncio +async def test_run_evaluator_preserves_inline_origin_for_span_progress_and_result( + with_memory_logger, with_simulate_login +): + origin = { + "object_type": "project_logs", + "object_id": "source-project", + "id": "source-row", + "_xact_id": "source-xact", + "created": "2026-06-01T00:00:00.000Z", + } + streamed_events = [] + + def task(input_value, hooks): + hooks.report_progress(progress_event()) + return input_value * 2 + + evaluator = Evaluator( + project_name="test-project", + eval_name="test-evaluator", + data=[{"input": 1, "origin": origin}], + task=task, + scores=[], + experiment_name=None, + metadata=None, + summarize_scores=False, + ) + exp = init_test_exp("test-evaluator", "test-project") + + result = await run_evaluator( + experiment=exp, + evaluator=evaluator, + position=None, + filters=[], + stream=streamed_events.append, + ) + + assert result.results[0].origin == origin + assert streamed_events[0]["origin"] == origin + root_spans = [log for log in with_memory_logger.pop() if not log["span_parents"]] + assert root_spans[0]["origin"] == origin + + +@pytest.mark.asyncio +async def test_dataset_row_origin_precedes_preserved_source_origin(): + source_origin = { + "object_type": "project_logs", + "object_id": "source-project", + "id": "source-row", + "_xact_id": "source-xact", + "created": "2026-06-01T00:00:00.000Z", + } + row = { + "input": 1, + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + "origin": source_origin, + } + dataset, patched_refetch = make_dataset("active-dataset", row) + streamed_events = [] + + def task(input_value, hooks): + hooks.report_progress(progress_event()) + return input_value * 2 + + evaluator = Evaluator( + project_name="test-project", + eval_name="test-evaluator", + data=dataset, + task=task, + scores=[], + experiment_name=None, + metadata=None, + summarize_scores=False, + ) + + with patched_refetch: + result = await run_evaluator( + experiment=None, + evaluator=evaluator, + position=None, + filters=[], + stream=streamed_events.append, + ) + + expected_origin = { + "object_type": "dataset", + "object_id": "active-dataset", + "id": "dataset-row", + "_xact_id": "dataset-xact", + "created": "2026-06-02T00:00:00.000Z", + } + assert result.results[0].origin == expected_origin + assert streamed_events[0]["origin"] == expected_origin + + +@pytest.mark.asyncio +async def test_incomplete_dataset_row_origin_falls_back_to_preserved_source_origin(): + source_origin = { + "object_type": "project_logs", + "object_id": "source-project", + "id": "source-row", + "_xact_id": "source-xact", + "created": "2026-06-01T00:00:00.000Z", + } + row = { + "input": 1, + "id": "dataset-row", + "created": "2026-06-02T00:00:00.000Z", + "origin": source_origin, + } + dataset, patched_refetch = make_dataset("active-dataset", row) + evaluator = Evaluator( + project_name="test-project", + eval_name="test-evaluator", + data=dataset, + task=lambda input_value: input_value * 2, + scores=[], + experiment_name=None, + metadata=None, + ) + + with patched_refetch: + result = await run_evaluator(experiment=None, evaluator=evaluator, position=None, filters=[]) + + assert result.results[0].origin == source_origin + + +@pytest.mark.asyncio +async def test_invalid_inline_origin_is_ignored_for_span_progress_and_result(with_memory_logger, with_simulate_login): + invalid_origin = { + "object_type": "dataset", + "object_id": "source-dataset", + "id": 123, + } + streamed_events = [] + + def task(input_value, hooks): + hooks.report_progress(progress_event()) + return input_value * 2 + + evaluator = Evaluator( + project_name="test-project", + eval_name="test-evaluator", + data=[{"input": 1, "origin": invalid_origin}], + task=task, + scores=[], + experiment_name=None, + metadata=None, + summarize_scores=False, + ) + exp = init_test_exp("test-evaluator", "test-project") + + result = await run_evaluator( + experiment=exp, + evaluator=evaluator, + position=None, + filters=[], + stream=streamed_events.append, + ) + + assert result.results[0].origin is None + assert "origin" not in streamed_events[0] + root_spans = [log for log in with_memory_logger.pop() if not log["span_parents"]] + assert "origin" not in root_spans[0] + + @pytest.mark.asyncio async def test_run_evaluator_basic(): """Test that run_evaluator correctly processes a simple evaluation.""" diff --git a/py/src/braintrust/type_tests/test_eval_generics.py b/py/src/braintrust/type_tests/test_eval_generics.py index 8ef793ac7..053355475 100644 --- a/py/src/braintrust/type_tests/test_eval_generics.py +++ b/py/src/braintrust/type_tests/test_eval_generics.py @@ -14,6 +14,7 @@ import pytest from braintrust.framework import EvalAsync, EvalCase, EvalResultWithSummary +from braintrust.generated_types import ObjectReference from braintrust.score import Score @@ -118,3 +119,25 @@ async def test_eval_divergent_output_and_expected(): assert result.results[0].output == ModelOutput(answer="4", confidence=0.99) assert isinstance(result.results[0].expected, frozenset) assert result.results[0].scores.get("match") == 1.0 + + +@pytest.mark.asyncio +async def test_eval_origin_types(): + """EvalCase, dictionary inputs, and EvalResult expose row origins.""" + origin: ObjectReference = { + "object_type": "project_logs", + "object_id": "source-project", + "id": "source-row", + } + + eval_case = EvalCase(input="case", origin=origin) + assert eval_case.origin == origin + + result = await EvalAsync( + "test-origin-types", + data=[{"input": "dictionary", "origin": origin}], + task=lambda input_value: input_value, + scores=[], + no_send_logs=True, + ) + assert result.results[0].origin == origin diff --git a/py/src/braintrust/types/_eval.py b/py/src/braintrust/types/_eval.py index 528df11e8..caf0ae165 100644 --- a/py/src/braintrust/types/_eval.py +++ b/py/src/braintrust/types/_eval.py @@ -10,6 +10,8 @@ from typing_extensions import NotRequired, TypedDict +from ..generated_types import ObjectReference + Input = TypeVar("Input") Expected = TypeVar("Expected") @@ -30,6 +32,8 @@ class EvalCaseDictNoOutput(Generic[Input], TypedDict): id: NotRequired[str | None] _xact_id: NotRequired[str | None] + created: NotRequired[str | None] + origin: NotRequired[ObjectReference | None] class EvalCaseDict(Generic[Input, Expected], EvalCaseDictNoOutput[Input]): @@ -51,3 +55,5 @@ class ExperimentDatasetEvent(TypedDict): input: Any | None expected: Any | None tags: Sequence[str] | None + created: NotRequired[str | None] + origin: NotRequired[ObjectReference | None] From ced60efb110d002eb177b1e247370fe7df8a0725 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Wed, 5 Aug 2026 15:11:55 -0400 Subject: [PATCH 2/2] simplify --- py/src/braintrust/framework.py | 20 ++-- py/src/braintrust/test_framework.py | 141 +++++++++------------------- 2 files changed, 53 insertions(+), 108 deletions(-) diff --git a/py/src/braintrust/framework.py b/py/src/braintrust/framework.py index 75ba13859..6bc0e2555 100644 --- a/py/src/braintrust/framework.py +++ b/py/src/braintrust/framework.py @@ -1599,8 +1599,14 @@ async def run_evaluator_task(datum, trial_index=0): experiment.dataset if experiment else evaluator.data if isinstance(evaluator.data, Dataset) else None ) - dataset_origin = ( - cast( + if ( + event_dataset + and isinstance(datum.id, str) + and datum.id + and isinstance(datum._xact_id, str) + and datum._xact_id + ): + origin = cast( ObjectReference, { "object_type": "dataset", @@ -1610,14 +1616,8 @@ async def run_evaluator_task(datum, trial_index=0): **({"created": datum.created} if isinstance(datum.created, str) else {}), }, ) - if event_dataset - and isinstance(datum.id, str) - and datum.id - and isinstance(datum._xact_id, str) - and datum._xact_id - else None - ) - origin = dataset_origin or _validated_object_reference(datum.origin) + else: + origin = _validated_object_reference(datum.origin) base_event = dict( name="eval", span_attributes={"type": SpanTypeAttribute.EVAL}, diff --git a/py/src/braintrust/test_framework.py b/py/src/braintrust/test_framework.py index d566365c6..9f1b390b3 100644 --- a/py/src/braintrust/test_framework.py +++ b/py/src/braintrust/test_framework.py @@ -39,98 +39,89 @@ def make_dataset(dataset_id, row): return dataset, patch.object(dataset, "_refetch", return_value=[row]) -def progress_event(): - return { - "format": "code", - "output_type": "completion", - "event": "progress", - "data": "0.5", - } +SOURCE_ORIGIN = { + "object_type": "project_logs", + "object_id": "source-project", + "id": "source-row", + "_xact_id": "source-xact", + "created": "2026-06-01T00:00:00.000Z", +} +INVALID_ORIGIN = { + "object_type": "dataset", + "object_id": "source-dataset", + "id": 123, +} + + +def reporting_task(input_value, hooks): + hooks.report_progress( + { + "format": "code", + "output_type": "completion", + "event": "progress", + "data": "0.5", + } + ) + return input_value * 2 def test_eval_case_from_dict_preserves_valid_origin(): - origin = { - "object_type": "project_logs", - "object_id": "source-project", - "id": "source-row", - "_xact_id": "source-xact", - "created": "2026-06-01T00:00:00.000Z", - } - - assert EvalCase.from_dict({"input": 1, "origin": origin}).origin == origin + assert EvalCase.from_dict({"input": 1, "origin": SOURCE_ORIGIN}).origin == SOURCE_ORIGIN +@pytest.mark.parametrize( + ("inline_origin", "expected_origin"), + [(SOURCE_ORIGIN, SOURCE_ORIGIN), (INVALID_ORIGIN, None)], + ids=["valid", "invalid"], +) @pytest.mark.asyncio -async def test_run_evaluator_preserves_inline_origin_for_span_progress_and_result( - with_memory_logger, with_simulate_login +async def test_run_evaluator_propagates_only_valid_inline_origins( + inline_origin, expected_origin, with_memory_logger, with_simulate_login ): - origin = { - "object_type": "project_logs", - "object_id": "source-project", - "id": "source-row", - "_xact_id": "source-xact", - "created": "2026-06-01T00:00:00.000Z", - } streamed_events = [] - - def task(input_value, hooks): - hooks.report_progress(progress_event()) - return input_value * 2 - evaluator = Evaluator( project_name="test-project", eval_name="test-evaluator", - data=[{"input": 1, "origin": origin}], - task=task, + data=[{"input": 1, "origin": inline_origin}], + task=reporting_task, scores=[], experiment_name=None, metadata=None, summarize_scores=False, ) - exp = init_test_exp("test-evaluator", "test-project") result = await run_evaluator( - experiment=exp, + experiment=init_test_exp("test-evaluator", "test-project"), evaluator=evaluator, position=None, filters=[], stream=streamed_events.append, ) - assert result.results[0].origin == origin - assert streamed_events[0]["origin"] == origin + assert result.results[0].origin == expected_origin root_spans = [log for log in with_memory_logger.pop() if not log["span_parents"]] - assert root_spans[0]["origin"] == origin + for event in (streamed_events[0], root_spans[0]): + assert event.get("origin") == expected_origin + assert ("origin" in event) is (expected_origin is not None) @pytest.mark.asyncio async def test_dataset_row_origin_precedes_preserved_source_origin(): - source_origin = { - "object_type": "project_logs", - "object_id": "source-project", - "id": "source-row", - "_xact_id": "source-xact", - "created": "2026-06-01T00:00:00.000Z", - } row = { "input": 1, "id": "dataset-row", "_xact_id": "dataset-xact", "created": "2026-06-02T00:00:00.000Z", - "origin": source_origin, + "origin": SOURCE_ORIGIN, } dataset, patched_refetch = make_dataset("active-dataset", row) streamed_events = [] - def task(input_value, hooks): - hooks.report_progress(progress_event()) - return input_value * 2 - evaluator = Evaluator( project_name="test-project", eval_name="test-evaluator", data=dataset, - task=task, + task=reporting_task, scores=[], experiment_name=None, metadata=None, @@ -159,18 +150,11 @@ def task(input_value, hooks): @pytest.mark.asyncio async def test_incomplete_dataset_row_origin_falls_back_to_preserved_source_origin(): - source_origin = { - "object_type": "project_logs", - "object_id": "source-project", - "id": "source-row", - "_xact_id": "source-xact", - "created": "2026-06-01T00:00:00.000Z", - } row = { "input": 1, "id": "dataset-row", "created": "2026-06-02T00:00:00.000Z", - "origin": source_origin, + "origin": SOURCE_ORIGIN, } dataset, patched_refetch = make_dataset("active-dataset", row) evaluator = Evaluator( @@ -186,46 +170,7 @@ async def test_incomplete_dataset_row_origin_falls_back_to_preserved_source_orig with patched_refetch: result = await run_evaluator(experiment=None, evaluator=evaluator, position=None, filters=[]) - assert result.results[0].origin == source_origin - - -@pytest.mark.asyncio -async def test_invalid_inline_origin_is_ignored_for_span_progress_and_result(with_memory_logger, with_simulate_login): - invalid_origin = { - "object_type": "dataset", - "object_id": "source-dataset", - "id": 123, - } - streamed_events = [] - - def task(input_value, hooks): - hooks.report_progress(progress_event()) - return input_value * 2 - - evaluator = Evaluator( - project_name="test-project", - eval_name="test-evaluator", - data=[{"input": 1, "origin": invalid_origin}], - task=task, - scores=[], - experiment_name=None, - metadata=None, - summarize_scores=False, - ) - exp = init_test_exp("test-evaluator", "test-project") - - result = await run_evaluator( - experiment=exp, - evaluator=evaluator, - position=None, - filters=[], - stream=streamed_events.append, - ) - - assert result.results[0].origin is None - assert "origin" not in streamed_events[0] - root_spans = [log for log in with_memory_logger.pop() if not log["span_parents"]] - assert "origin" not in root_spans[0] + assert result.results[0].origin == SOURCE_ORIGIN @pytest.mark.asyncio