diff --git a/pydantic_evals/pydantic_evals/evaluators/context.py b/pydantic_evals/pydantic_evals/evaluators/context.py index 4808361ca6..7d7ac20bd6 100644 --- a/pydantic_evals/pydantic_evals/evaluators/context.py +++ b/pydantic_evals/pydantic_evals/evaluators/context.py @@ -9,6 +9,7 @@ from dataclasses import dataclass, field from typing import Any, Generic +from pydantic import TypeAdapter from typing_extensions import TypeVar from ..otel._errors import SpanTreeRecordingError @@ -100,3 +101,7 @@ def span_tree(self) -> SpanTree: # In this case, there was a reason we couldn't record the SpanTree. We raise that now raise self._span_tree return self._span_tree + + +EVALUATOR_CONTEXT_ADAPTER = TypeAdapter(EvaluatorContext) +"""This adapter can be used to serialize and deserialize `EvaluatorContext` objects to and from JSON.""" diff --git a/pydantic_evals/pydantic_evals/otel/_errors.py b/pydantic_evals/pydantic_evals/otel/_errors.py index faca85dbe2..181e9285eb 100644 --- a/pydantic_evals/pydantic_evals/otel/_errors.py +++ b/pydantic_evals/pydantic_evals/otel/_errors.py @@ -1,7 +1,34 @@ +from typing import Any + +from pydantic_core import core_schema + + class SpanTreeRecordingError(Exception): """An exception that is used to provide the reason why a SpanTree was not recorded by `context_subtree`. This will either be due to missing dependencies or because a tracer provider had not been set. """ - pass + message: str + + def __init__(self, message: str): + self.message = message + super().__init__(message) + + @classmethod + def __get_pydantic_core_schema__(cls, _: Any, __: Any) -> core_schema.CoreSchema: + """Pydantic core schema to allow `SpanTreeRecordingError` to be (de)serialized.""" + schema = core_schema.typed_dict_schema( + { + 'message': core_schema.typed_dict_field(core_schema.str_schema()), + 'kind': core_schema.typed_dict_field(core_schema.literal_schema(['span-tree-recording-error'])), + } + ) + return core_schema.no_info_after_validator_function( + lambda dct: SpanTreeRecordingError(dct['message']), + schema, + serialization=core_schema.plain_serializer_function_ser_schema( + lambda x: {'message': x.message, 'kind': 'span-tree-recording-error'}, + return_schema=schema, + ), + )