Repository: pydantic/pydantic-ai
File: tests/evals/test_dataset.py (Python)

Code before:
        return TaskOutput(answer='Unknown')  # pragma: no cover

    await example_dataset.evaluate(mock_async_task)

    spans = capfire.exporter.exported_spans_as_dict(parse_json_attributes=True)
    spans.sort(key=lambda s: s['start_time'])

    for span in spans:
        # These may or may not be present and may have weird values due to things running in async
        span['attributes'].pop('code.filepath', None)
        span['attributes'].pop('code.function', None)
        span['attributes'].pop('code.lineno', None)

    assert [(span['name'], span['attributes']) for span in spans] == snapshot(
        [
            (
                'evaluate {name}',
                {
                    'name': 'mock_async_task',
                    'n_cases': 2,
                    'assertion_pass_rate': 1.0,
                    'logfire.msg_template': 'evaluate {name}',
                    'logfire.msg': 'evaluate mock_async_task',
                    'experiment.metadata': {
                        'n_cases': 2,
                        'averages': {
                            'name': 'Averages',
                            'scores': {'confidence': 1.0},
                            'labels': {},
                            'metrics': {},
                            'assertions': 1.0,
                            'task_duration': 1.0,
                            'total_duration': 9.0,
                        },
                    },
                    'logfire.span_type': 'span',
                    'logfire.json_schema': {
                        'type': 'object',
                        'properties': {
                            'name': {},
                            'n_cases': {},
                            'experiment.metadata': {
                                'type': 'object',
                                'properties': {
                                    'averages': {
                                        'type': 'object',
                                        'title': 'ReportCaseAggregate',
                                        'x-python-datatype': 'PydanticModel',
                                    }
                                },
                            },

Review comment (suggestion):
If we're going to be using this on the frontend, I don't know if it's wise to pass along the entire `ReportCaseAggregate` object or just the subset we need

Code after:
        output_type: type[str] | type[int]

    my_dataset = Dataset[MyInputs, Any, Any](
        cases=[
            Case(
                name='str',
                inputs=MyInputs(output_type=str),
                expected_output='abc',
            ),
            Case(
                name='int',
                inputs=MyInputs(output_type=int),
                expected_output=123,
            ),
        ],
    )

    async def my_task(my_inputs: MyInputs) -> int | str:
        if issubclass(my_inputs.output_type, str):
            return my_inputs.output_type('abc')
        else:
            return my_inputs.output_type(123)

    report = my_dataset.evaluate_sync(task=my_task)
    assert [c.inputs for c in report.cases] == snapshot([MyInputs(output_type=str), MyInputs(output_type=int)])

    table = report.console_table(include_input=True)
    assert render_table(table) == snapshot("""\
                                        Evaluation Summary: my_task
┏━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Case ID  ┃ Inputs                                                                             ┃ Duration ┃
┡━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ str      │ test_evaluate_non_serializable_inputs.<locals>.MyInputs(output_type=<class 'str'>) │     1.0s │
├──────────┼────────────────────────────────────────────────────────────────────────────────────┼──────────┤
│ int      │ test_evaluate_non_serializable_inputs.<locals>.MyInputs(output_type=<class 'int'>) │     1.0s │
├──────────┼────────────────────────────────────────────────────────────────────────────────────┼──────────┤
│ Averages │                                                                                    │     1.0s │
└──────────┴────────────────────────────────────────────────────────────────────────────────────┴──────────┘
""")


async def test_evaluate_async_logfire(
    example_dataset: Dataset[TaskInput, TaskOutput, TaskMetadata],
    simple_evaluator: type[Evaluator[TaskInput, TaskOutput, TaskMetadata]],
    capfire: CaptureLogfire,
):  # pragma: lax no cover
    """Test evaluating a dataset."""
    example_dataset.add_evaluator(simple_evaluator())

    async def mock_async_task(inputs: TaskInput) -> TaskOutput:
        if inputs.query == 'What is 2+2?':