diff --git a/py/src/braintrust/framework.py b/py/src/braintrust/framework.py index 3f0bd86e..341db8ba 100644 --- a/py/src/braintrust/framework.py +++ b/py/src/braintrust/framework.py @@ -2011,7 +2011,7 @@ async def ensure_spans_flushed(): set_current=False, ).as_dataset() - if inspect.isfunction(data_iterator) or inspect.isroutine(data_iterator): + if callable(data_iterator) and not isinstance(data_iterator, Dataset): data_iterator = data_iterator() if not inspect.isasyncgen(data_iterator): diff --git a/py/src/braintrust/test_framework.py b/py/src/braintrust/test_framework.py index afbfe026..48d5a792 100644 --- a/py/src/braintrust/test_framework.py +++ b/py/src/braintrust/test_framework.py @@ -5,6 +5,7 @@ import sys import threading from concurrent.futures import ThreadPoolExecutor +from functools import partial from unittest.mock import MagicMock, patch import pytest @@ -1833,6 +1834,49 @@ async def test_eval_enable_cache(): state.span_cache.stop.assert_called() +@pytest.mark.parametrize("eval_fn", [Eval, EvalAsync]) +@pytest.mark.parametrize("data_kind", ["partial", "callable_object", "function", "iterator", "list"]) +@pytest.mark.asyncio +async def test_eval_callable_data_factories(eval_fn, data_kind): + factory_calls = [] + + def make_data(offset=2): + factory_calls.append(offset) + return iter([EvalCase(input=x, expected=x * 2) for x in (offset + 1, offset + 2)]) + + class DataFactory: + def __call__(self): + return make_data() + + if data_kind == "partial": + data = partial(make_data, offset=2) + elif data_kind == "callable_object": + data = DataFactory() + elif data_kind == "function": + data = make_data + elif data_kind == "iterator": + data = make_data() + else: + data = list(make_data()) + + def exact_match(output, expected): + return output == expected + + result = await eval_fn( + "test-callable-data", + data=data, + task=lambda input: input * 2, + scores=[exact_match], + no_send_logs=True, + enable_cache=False, + ) + + assert factory_calls == [2] + assert [row.input for row in result.results] == [3, 4] + assert [row.output for row in result.results] == [6, 8] + assert all(row.error is None and row.scores == {"exact_match": 1.0} for row in result.results) + + @pytest.mark.asyncio async def test_run_evaluator_empty_dataset_warns(capsys): """Warn when run_evaluator receives an empty dataset."""