Serialization / Deserialization for `EvaluatorContext` and separate task execution and evaluator execution

[hon-gyu]

Description

Motivation:

• Some evalution tasks can be time-consuming / expensive to run. It would be nice to be able to explore what metrics can be derived from evaluation results (EvaluatorContext) without re-running the task again
• This also enable data-science-like workflow with evaluation results, as DataFrame can be constructed from evaluation results.

Proposed implementation:

• Add serde schema for SpanTreeRecordingError so that TypeAdapter can be used on EvaluatorContext.
• Split pydantic_evals/pydantic_evals/dataset.py::_run_task_and_evaluators() function into two functions, one for running task and one for running evaluators, and make them public.
• (Optional?) A ContextCaptureEvaluator that saves EvaluatorContext to disk.

References

No response

[dmontagu]

• I'm totally fine introducing some way of serializing the SpanTreeRecordingError, but note that it will probably be hard for it to have the same data you'd have if it was recorded the real way (e.g., __cause__/__context__, the full traceback, etc.). I think that's not necessary, but my point is it will probably be hard/impossible/not-a-reasonable-goal to get idempotant serialization->deserialization. I could imagine introducing a new BaseModel type for deserializing the exception or similar, and adding it to the union of values allowed on the EvaluatorContext during serialization/deserialization; I'm flexible.

• Split pydantic_evals/pydantic_evals/dataset.py::_run_task_and_evaluators() function into two functions, one for running task and one for running evaluators, and make them public.

  I'm okay with this in principle, but I'll note the current design purposely does not run all tasks and then all evaluations on those tasks, so that if tasks have large outputs you don't run into OOM issues or need to spill to disk or whatever; functionally, we are streaming the execution of tasks + their evaluations. I think a change here still needs to retain the ability to conveniently run the tasks with a maximum concurrency and without being forced to dump results to disk prior to running evaluations. But I'm okay (and think it would be a good improvement to the library) if we introduce backwards compatible APIs that make it possible to run evaluators after the fact, etc.

More concrete proposals and/or PRs doing this would be very welcome!

[hon-gyu]

@dmontagu Thanks for your reply. I'll try to draft a PR.

[hon-gyu]

Hi @dmontagu, I have a PR adding core schema for SpanTreeRecordingError, enabling the use of TypeAdapter on EvaluatorContext. You mentioned idempotant serialization / deserialization could be unrealisitc, but I think it's actually idempotant as SpanTreeRecordingError only captures either missing dependencies or tracer provider not being set. There's no extra information except for a message string.

Here's an example usage of capturing the task output on disk for additional evaluation experiment.

from pathlib import Path
from uuid import uuid4

import logfire

from pydantic_evals import Case, Dataset
from pydantic_evals.evaluators import (
    Evaluator,
    EvaluatorContext,
    EvaluatorOutput,
)
from pydantic_evals.evaluators.context import EVALUATOR_CONTEXT_ADAPTER


@logfire.instrument()
def to_title_case(text: str) -> str:
    """Convert text to title case."""
    return text.title()


class ContextCaptureEvaluator(Evaluator):
    def __init__(self, save_dir: Path):
        self.save_dir = save_dir
        if not self.save_dir.exists():
            self.save_dir.mkdir(parents=True)
        super().__init__()

    def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
        name = ctx.name or uuid4().hex[:8]
        with (self.save_dir / f'evaluator_context_{name}.json').open('wb') as f:
            f.write(EVALUATOR_CONTEXT_ADAPTER.dump_json(ctx, indent=2))
        return {'captured': True}


dataset = Dataset(
    name='title_case_validation',
    cases=[
        Case(
            name='basic_lowercase',
            inputs='hello world',
            expected_output='Hello World',
        ),
        Case(
            name='basic_uppercase',
            inputs='HELLO WORLD',
            expected_output='Hello World',
        ),
        Case(
            name='mixed_case',
            inputs='HeLLo WoRLd',
            expected_output='Hello World',
        ),
    ],
    evaluators=[ContextCaptureEvaluator(Path('captured_contexts'))],
)

if __name__ == '__main__':
    logfire.configure()
    dataset.evaluate_sync(to_title_case)

---

Regarding the second part about separating task execution and evaluation, I found that with the (de)serialization of evaluator contexts available, it's quite easy to adapt pydantic_evals.evaluators._run_evaluator::run_evaluator and pydantic_evals.dataset::_run_task for achieving the goal. So maybe it's not necessary to make them a public method of DataSet class. Here's my usage of constructing DataFrame from evaluation results for each eval cases.

Details

from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from uuid import uuid4

from pydantic_ai.retries import RetryConfig
from pydantic_evals import Case, Dataset
from pydantic_evals._utils import task_group_gather
from pydantic_evals.dataset import InputsT, MetadataT, OutputT, _group_evaluator_outputs_by_type, _run_task
from pydantic_evals.evaluators import (
    Contains,
    EqualsExpected,
    EvaluationResult,
    Evaluator,
    EvaluatorContext,
    EvaluatorFailure,
    IsInstance,
    MaxDuration,
)
from pydantic_evals.evaluators._run_evaluator import run_evaluator

try:
    import polars as pl
except ImportError:
    POLARS_AVAILABLE = False
else:
    POLARS_AVAILABLE = True


async def run_task(
    dataset: Dataset,
    task: Callable[[InputsT], Awaitable[OutputT] | OutputT],
    retry: RetryConfig | None = None,
) -> list[EvaluatorContext[InputsT, OutputT, MetadataT]]:
    """Run a task on all cases in the dataset, returning evaluation contexts."""
    return await task_group_gather([lambda case=case: _run_task(task, case, retry) for case in dataset.cases])


async def run_evaluator_for_contexts(
    scoring_contexts: list[EvaluatorContext[InputsT, OutputT, MetadataT]],
    evaluators: list[Evaluator[InputsT, OutputT, MetadataT]],
    retry: RetryConfig | None = None,
) -> list[list[EvaluationResult | EvaluatorFailure]]:
    """Run evaluators on multiple pre-computed contexts."""

    async def _run_evaluators_for_context(
        ctx: EvaluatorContext,
    ) -> list[EvaluationResult | EvaluatorFailure]:
        """Run all evaluators on a single context, returning flattened results."""
        evaluator_outputs = await task_group_gather([lambda ev=ev: run_evaluator(ev, ctx, retry) for ev in evaluators])

        # Flatten results
        flattened: list[EvaluationResult | EvaluatorFailure] = []
        for output in evaluator_outputs:
            if isinstance(output, EvaluatorFailure):
                flattened.append(output)
            else:
                # output is list[EvaluationResult], extend to flatten
                flattened.extend(output)

        return flattened

    # Run evaluators on all contexts
    return await task_group_gather([lambda ctx=ctx: _run_evaluators_for_context(ctx) for ctx in scoring_contexts])


def to_title_case(text: str) -> str:
    """Convert text to title case with some quirks."""
    # Sometimes return uppercase instead of title case
    if 'HELLO' in text.upper() and len(text) > 15:
        return text.upper()

    # Sometimes skip short words
    words = text.split()
    result_words = []
    for word in words:
        if len(word) <= 2 and word.lower() in ['a', 'an', 'or', 'is']:
            result_words.append(word.lower())  # Keep articles lowercase
        else:
            result_words.append(word.title())

    return ' '.join(result_words)


@dataclass
class CapitalizationScore(Evaluator):
    """Calculate the percentage of words that start with a capital letter."""

    def evaluate(self, ctx: EvaluatorContext) -> float:
        if not isinstance(ctx.output, str):
            return 0.0

        words = ctx.output.split()
        if not words:
            return 1.0  # Empty string is technically 100% correct

        capitalized_count = sum(1 for word in words if word and word[0].isupper())
        return capitalized_count / len(words)


@dataclass
class CapitalLetterRatio(Evaluator):
    """Calculate the ratio of capital letters to total letters."""

    def evaluate(self, ctx: EvaluatorContext) -> float:
        if not isinstance(ctx.output, str):
            return 0.0

        letters = [c for c in ctx.output if c.isalpha()]
        if not letters:
            return 0.0

        capital_count = sum(1 for c in letters if c.isupper())
        return capital_count / len(letters)


# Create evaluation dataset
evaluators = [
    # Boolean assertions
    IsInstance(type_name='str'),
    EqualsExpected(),
    Contains(value='H', evaluation_name='has_capitals'),
    MaxDuration(seconds=0.001),
    # Numeric scores for more interesting analysis
    CapitalizationScore(),
    CapitalLetterRatio(),
]
dataset = Dataset(
    name='title_case_validation',
    cases=[
        # Basic functionality
        Case(
            name='basic_lowercase',
            inputs='hello world',
            expected_output='Hello World',
        ),
        Case(
            name='basic_uppercase',
            inputs='HELLO WORLD',
            expected_output='Hello World',
        ),
        Case(
            name='mixed_case',
            inputs='HeLLo WoRLd',
            expected_output='Hello World',
        ),
        # Edge cases
        Case(
            name='empty_string',
            inputs='',
            expected_output='',
        ),
        Case(
            name='single_word',
            inputs='hello',
            expected_output='Hello',
        ),
        Case(
            name='with_punctuation',
            inputs='hello, world!',
            expected_output='Hello, World!',
        ),
        Case(
            name='with_numbers',
            inputs='hello 123 world',
            expected_output='Hello 123 World',
        ),
        Case(
            name='apostrophes',
            inputs="don't stop believin'",
            expected_output="Don'T Stop Believin'",
        ),
        # Additional cases for more interesting numeric scores
        Case(
            name='long_sentence',
            inputs='the quick brown fox jumps over the lazy dog',
            expected_output='The Quick Brown Fox Jumps Over The Lazy Dog',
        ),
        Case(
            name='with_articles',
            inputs='a tale of two cities',
            expected_output='A Tale Of Two Cities',
        ),
    ],
    evaluators=evaluators,
)


async def main():
    """Main async function to run the evaluation workflow."""
    # Run tasks
    contexts = await run_task(dataset, to_title_case)

    # Run evaluators on contexts
    results: list[list[EvaluationResult | EvaluatorFailure]] = await run_evaluator_for_contexts(contexts, evaluators)

    print_dataframe_analysis(results)


def print_dataframe_analysis(results: list[list[EvaluationResult | EvaluatorFailure]]):
    if POLARS_AVAILABLE:
        # Convert to DataFrame for analysis
        def results_to_dataframe(
            case_names: list[str],
            results: list[list[EvaluationResult | EvaluatorFailure]],
        ) -> pl.DataFrame:
            """Convert evaluation results to a Polars DataFrame in wide format.

            Args:
                case_names: List of case names, used as row index
                results: List of evaluation results (one list per case)

            Returns:
                DataFrame where each row is a case and each column is an evaluator's metric.
                Failed evaluations have None values.
            """
            rows: list[dict] = []

            for case_name, context_results in zip(case_names, results, strict=True):
                # Separate successful results from failures
                successful_results = [r for r in context_results if isinstance(r, EvaluationResult)]
                failures = [r for r in context_results if isinstance(r, EvaluatorFailure)]

                # Group results by type (bool, numeric, string)
                assertions, scores, labels = _group_evaluator_outputs_by_type(successful_results)

                # Build row with case name and all evaluator results as columns
                row = {'case_name': case_name}

                # Add assertions as columns
                for evaluator_name, result in assertions.items():
                    row[evaluator_name] = result.value

                # Add scores as columns
                for evaluator_name, result in scores.items():
                    row[evaluator_name] = result.value

                # Add labels as columns
                for evaluator_name, result in labels.items():
                    row[evaluator_name] = result.value

                # Add failures as None columns
                for failure in failures:
                    row[failure.name] = None

                rows.append(row)

            return pl.DataFrame(rows)

        df = results_to_dataframe([c.name or uuid4().hex[:8] for c in dataset.cases], results)
        print(df)
        print(df.select(pl.exclude('case_name')).describe())


if __name__ == '__main__':
    import asyncio

    asyncio.run(main())

Result: