arize-phoenix 4.4.4rc4__py3-none-any.whl → 4.4.4rc5__py3-none-any.whl

phoenix/datasets/experiments.py

@@ -8,11 +8,11 @@ from itertools import product
 from typing import (
     Any,
     Awaitable,
-    Callable,
-    Coroutine,
+    Dict,
     Iterable,
     Mapping,
     Optional,
+    Sequence,
     Tuple,
     Type,
     Union,
@@ -42,20 +42,23 @@ from phoenix.config import (
     get_env_host,
     get_env_port,
 )
+from phoenix.datasets.evaluators.utils import (
+    Evaluator,
+    EvaluatorName,
+    ExperimentEvaluator,
+    create_evaluator,
+)
 from phoenix.datasets.tracing import capture_spans
 from phoenix.datasets.types import (
-    CanAsyncEvaluate,
-    CanEvaluate,
     Dataset,
     EvaluationResult,
     Example,
     Experiment,
     ExperimentEvaluationRun,
-    ExperimentEvaluator,
     ExperimentResult,
     ExperimentRun,
     ExperimentRunId,
-    JSONSerializable,
+    ExperimentTask,
     TestCase,
 )
 from phoenix.evals.executors import get_executor_on_sync_context
@@ -65,11 +68,6 @@ from phoenix.session.session import active_session
 from phoenix.trace.attributes import flatten
 from phoenix.utilities.json import jsonify
 
-ExperimentTask: TypeAlias = Union[
-    Callable[[Example], JSONSerializable],
-    Callable[[Example], Coroutine[None, None, JSONSerializable]],
-]
-
 
 def _get_base_url() -> str:
     host = get_env_host()
@@ -98,10 +96,22 @@ def _get_dataset_experiments_url(*, dataset_id: str) -> str:
     return f"{_get_web_base_url()}datasets/{dataset_id}/experiments"
 
 
-def _phoenix_client() -> httpx.Client:
+def _phoenix_clients() -> Tuple[httpx.Client, httpx.AsyncClient]:
     headers = get_env_client_headers()
-    client = httpx.Client(base_url=_get_base_url(), headers=headers)
-    return client
+    return httpx.Client(
+        base_url=_get_base_url(),
+        headers=headers,
+    ), httpx.AsyncClient(
+        base_url=_get_base_url(),
+        headers=headers,
+    )
+
+
+Evaluators: TypeAlias = Union[
+    ExperimentEvaluator,
+    Sequence[ExperimentEvaluator],
+    Mapping[EvaluatorName, ExperimentEvaluator],
+]
 
 
 def run_experiment(
@@ -111,16 +121,17 @@ def run_experiment(
     experiment_name: Optional[str] = None,
     experiment_description: Optional[str] = None,
     experiment_metadata: Optional[Mapping[str, Any]] = None,
-    evaluators: Optional[Union[ExperimentEvaluator, Iterable[ExperimentEvaluator]]] = None,
+    evaluators: Optional[Evaluators] = None,
     rate_limit_errors: Optional[Union[Type[BaseException], Tuple[Type[BaseException], ...]]] = None,
 ) -> Experiment:
     # Add this to the params once supported in the UI
     repetitions = 1
     assert repetitions > 0, "Must run the experiment at least once."
+    evaluators_by_name = _evaluators_by_name(evaluators)
 
-    client = _phoenix_client()
+    sync_client, async_client = _phoenix_clients()
 
-    experiment_response = client.post(
+    experiment_response = sync_client.post(
         f"/v1/datasets/{dataset.id}/experiments",
         json={
             "version-id": dataset.version_id,
@@ -141,8 +152,8 @@ def run_experiment(
         SimpleSpanProcessor(OTLPSpanExporter(urljoin(f"{_get_base_url()}", "v1/traces")))
     )
     tracer = tracer_provider.get_tracer(__name__)
-    root_span_name = f"Task: {task.__qualname__}"
-    root_span_kind = CHAIN.value
+    root_span_name = f"Task: {_get_task_name(task)}"
+    root_span_kind = CHAIN
 
     dataset_experiments_url = _get_dataset_experiments_url(dataset_id=dataset.id)
     experiment_compare_url = _get_experiment_url(dataset_id=dataset.id, experiment_id=experiment_id)
@@ -207,6 +218,10 @@ def run_experiment(
             error=repr(error) if error else None,
             trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore[no-untyped-call]
         )
+        resp = sync_client.post(
+            f"/v1/experiments/{experiment_id}/runs", json=jsonify(experiment_run)
+        )
+        resp.raise_for_status()
         return experiment_run
 
     async def async_run_experiment(test_case: TestCase) -> ExperimentRun:
@@ -257,6 +272,10 @@ def run_experiment(
             error=repr(error) if error else None,
             trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore[no-untyped-call]
         )
+        resp = await async_client.post(
+            f"/v1/experiments/{experiment_id}/runs", json=jsonify(experiment_run)
+        )
+        resp.raise_for_status()
         return experiment_run
 
     rate_limited_sync_run_experiment = functools.reduce(
@@ -279,12 +298,7 @@ def run_experiment(
         TestCase(example=ex, repetition_number=rep)
         for ex, rep in product(dataset.examples, range(1, repetitions + 1))
     ]
-    experiment_payloads, _execution_details = executor.run(test_cases)
-    for payload in experiment_payloads:
-        if payload is not None:
-            resp = client.post(f"/v1/experiments/{experiment_id}/runs", json=jsonify(payload))
-            resp.raise_for_status()
-
+    _, _execution_details = executor.run(test_cases)
     experiment = Experiment(
         id=experiment_id,
         dataset_id=dataset.id,
@@ -293,26 +307,34 @@ def run_experiment(
     )
 
     print("✅ Task runs completed.")
-    print("🧠 Evaluation started.")
 
-    if evaluators is not None:
-        _evaluate_experiment(experiment, evaluators, dataset.examples, client)
+    if evaluators_by_name:
+        _evaluate_experiment(
+            experiment,
+            evaluators=evaluators_by_name,
+            dataset_examples=dataset.examples,
+            clients=(sync_client, async_client),
+        )
 
     return experiment
 
 
 def evaluate_experiment(
     experiment: Experiment,
-    evaluators: Union[ExperimentEvaluator, Iterable[ExperimentEvaluator]],
+    evaluators: Union[
+        ExperimentEvaluator,
+        Sequence[ExperimentEvaluator],
+        Mapping[EvaluatorName, ExperimentEvaluator],
+    ],
 ) -> None:
-    client = _phoenix_client()
+    sync_client, async_client = _phoenix_clients()
     dataset_id = experiment.dataset_id
     dataset_version_id = experiment.dataset_version_id
 
     dataset_examples = [
         Example.from_dict(ex)
         for ex in (
-            client.get(
+            sync_client.get(
                 f"/v1/datasets/{dataset_id}/examples",
                 params={"version-id": str(dataset_version_id)},
             )
@@ -321,26 +343,29 @@ def evaluate_experiment(
             .get("examples", [])
         )
     ]
-    _evaluate_experiment(experiment, evaluators, dataset_examples, client)
-
-
-ExperimentEvaluatorName: TypeAlias = str
+    _evaluate_experiment(
+        experiment,
+        evaluators=evaluators,
+        dataset_examples=dataset_examples,
+        clients=(sync_client, async_client),
+    )
 
 
 def _evaluate_experiment(
     experiment: Experiment,
-    evaluators: Union[ExperimentEvaluator, Iterable[ExperimentEvaluator]],
+    *,
+    evaluators: Evaluators,
     dataset_examples: Iterable[Example],
-    client: httpx.Client,
+    clients: Tuple[httpx.Client, httpx.AsyncClient],
 ) -> None:
-    if isinstance(evaluators, (CanEvaluate, CanAsyncEvaluate)):
-        evaluators = [evaluators]
-
+    evaluators_by_name = _evaluators_by_name(evaluators)
+    if not evaluators_by_name:
+        raise ValueError("Must specify at least one Evaluator")
     experiment_id = experiment.id
-
+    sync_client, async_client = clients
     experiment_runs = [
         ExperimentRun.from_dict(exp_run)
-        for exp_run in client.get(f"/v1/experiments/{experiment_id}/runs").json()
+        for exp_run in sync_client.get(f"/v1/experiments/{experiment_id}/runs").json()
     ]
 
     # not all dataset examples have associated experiment runs, so we need to pair them up
@@ -350,9 +375,9 @@ def _evaluate_experiment(
         example = examples_by_id.get(exp_run.dataset_example_id)
         if example:
             example_run_pairs.append((deepcopy(example), exp_run))
-    evaluation_inputs = [
-        (example, run, evaluator.name, evaluator)
-        for (example, run), evaluator in product(example_run_pairs, evaluators)
+    evaluation_input = [
+        (example, run, evaluator)
+        for (example, run), evaluator in product(example_run_pairs, evaluators_by_name.values())
     ]
 
     project_name = "evaluators"
@@ -362,36 +387,34 @@ def _evaluate_experiment(
         SimpleSpanProcessor(OTLPSpanExporter(urljoin(f"{_get_base_url()}", "v1/traces")))
     )
     tracer = tracer_provider.get_tracer(__name__)
-    root_span_kind = "EVALUATOR"
+    root_span_kind = EVALUATOR
 
     def sync_evaluate_run(
-        obj: Tuple[Example, ExperimentRun, ExperimentEvaluatorName, ExperimentEvaluator],
+        obj: Tuple[Example, ExperimentRun, Evaluator],
     ) -> ExperimentEvaluationRun:
-        example, experiment_run, name, evaluator = obj
+        example, experiment_run, evaluator = obj
         result: Optional[EvaluationResult] = None
         error: Optional[BaseException] = None
         status = Status(StatusCode.OK)
-        root_span_name = f"Evaluation: {name}"
+        root_span_name = f"Evaluation: {evaluator.name}"
         with ExitStack() as stack:
             span: Span = stack.enter_context(
                 tracer.start_as_current_span(root_span_name, context=Context())
             )
             stack.enter_context(capture_spans(resource))
             try:
-                # Do not use keyword arguments, which can fail at runtime
-                # even when function obeys protocol, because keyword arguments
-                # are implementation details.
-                if not isinstance(evaluator, CanEvaluate):
-                    raise RuntimeError("Task is async but running in sync context")
-                _output = evaluator.evaluate(example, experiment_run)
-                if isinstance(_output, Awaitable):
-                    raise RuntimeError("Task is async but running in sync context")
-                result = _output
+                result = evaluator.evaluate(
+                    output=None if experiment_run.output is None else experiment_run.output.result,
+                    expected=example.output,
+                    input=example.input,
+                    metadata=example.metadata,
+                )
             except BaseException as exc:
                 span.record_exception(exc)
                 status = Status(StatusCode.ERROR, f"{type(exc).__name__}: {exc}")
                 error = exc
-            span.set_attributes(dict(flatten(jsonify(result), recurse_on_sequence=True)))
+            if result:
+                span.set_attributes(dict(flatten(jsonify(result), recurse_on_sequence=True)))
             span.set_attribute(OPENINFERENCE_SPAN_KIND, root_span_kind)
             span.set_status(status)
 
@@ -400,43 +423,41 @@ def _evaluate_experiment(
             start_time=_decode_unix_nano(cast(int, span.start_time)),
             end_time=_decode_unix_nano(cast(int, span.end_time)),
             name=evaluator.name,
-            annotator_kind=evaluator.annotator_kind,
+            annotator_kind=evaluator.kind,
             error=repr(error) if error else None,
             result=result,
             trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore[no-untyped-call]
         )
+        resp = sync_client.post("/v1/experiment_evaluations", json=jsonify(evaluator_payload))
+        resp.raise_for_status()
         return evaluator_payload
 
     async def async_evaluate_run(
-        obj: Tuple[Example, ExperimentRun, ExperimentEvaluatorName, ExperimentEvaluator],
+        obj: Tuple[Example, ExperimentRun, Evaluator],
     ) -> ExperimentEvaluationRun:
-        example, experiment_run, name, evaluator = obj
+        example, experiment_run, evaluator = obj
         result: Optional[EvaluationResult] = None
         error: Optional[BaseException] = None
         status = Status(StatusCode.OK)
-        root_span_name = f"Evaluation: {name}"
+        root_span_name = f"Evaluation: {evaluator.name}"
         with ExitStack() as stack:
             span: Span = stack.enter_context(
                 tracer.start_as_current_span(root_span_name, context=Context())
             )
             stack.enter_context(capture_spans(resource))
             try:
-                # Do not use keyword arguments, which can fail at runtime
-                # even when function obeys protocol, because keyword arguments
-                # are implementation details.
-                if isinstance(evaluator, CanAsyncEvaluate):
-                    result = await evaluator.async_evaluate(example, experiment_run)
-                else:
-                    _output = evaluator.evaluate(example, experiment_run)
-                    if isinstance(_output, Awaitable):
-                        result = await _output
-                    else:
-                        result = _output
+                result = await evaluator.async_evaluate(
+                    output=None if experiment_run.output is None else experiment_run.output.result,
+                    expected=example.output,
+                    input=example.input,
+                    metadata=example.metadata,
+                )
             except BaseException as exc:
                 span.record_exception(exc)
                 status = Status(StatusCode.ERROR, f"{type(exc).__name__}: {exc}")
                 error = exc
-            span.set_attributes(dict(flatten(jsonify(result), recurse_on_sequence=True)))
+            if result:
+                span.set_attributes(dict(flatten(jsonify(result), recurse_on_sequence=True)))
             span.set_attribute(OPENINFERENCE_SPAN_KIND, root_span_kind)
             span.set_status(status)
 
@@ -445,11 +466,15 @@ def _evaluate_experiment(
             start_time=_decode_unix_nano(cast(int, span.start_time)),
             end_time=_decode_unix_nano(cast(int, span.end_time)),
             name=evaluator.name,
-            annotator_kind=evaluator.annotator_kind,
+            annotator_kind=evaluator.kind,
             error=repr(error) if error else None,
             result=result,
             trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore[no-untyped-call]
         )
+        resp = await async_client.post(
+            "/v1/experiment_evaluations", json=jsonify(evaluator_payload)
+        )
+        resp.raise_for_status()
         return evaluator_payload
 
     executor = get_executor_on_sync_context(
@@ -460,11 +485,38 @@ def _evaluate_experiment(
         fallback_return_value=None,
         tqdm_bar_format=get_tqdm_progress_bar_formatter("running experiment evaluations"),
     )
-    evaluation_payloads, _execution_details = executor.run(evaluation_inputs)
-    for payload in evaluation_payloads:
-        if payload is not None:
-            resp = client.post("/v1/experiment_evaluations", json=jsonify(payload))
-            resp.raise_for_status()
+    print("🧠 Evaluation started.")
+    _, _execution_details = executor.run(evaluation_input)
+
+
+def _evaluators_by_name(obj: Optional[Evaluators]) -> Mapping[EvaluatorName, Evaluator]:
+    evaluators_by_name: Dict[EvaluatorName, Evaluator] = {}
+    if obj is None:
+        return evaluators_by_name
+    if isinstance(mapping := obj, Mapping):
+        for name, value in mapping.items():
+            evaluator = (
+                create_evaluator(name=name)(value) if not isinstance(value, Evaluator) else value
+            )
+            name = evaluator.name
+            if name in evaluators_by_name:
+                raise ValueError(f"Two evaluators have the same name: {name}")
+            evaluators_by_name[name] = evaluator
+    elif isinstance(seq := obj, Sequence):
+        for value in seq:
+            evaluator = create_evaluator()(value) if not isinstance(value, Evaluator) else value
+            name = evaluator.name
+            if name in evaluators_by_name:
+                raise ValueError(f"Two evaluators have the same name: {name}")
+            evaluators_by_name[name] = evaluator
+    else:
+        assert not isinstance(obj, Mapping) and not isinstance(obj, Sequence)
+        evaluator = create_evaluator()(obj) if not isinstance(obj, Evaluator) else obj
+        name = evaluator.name
+        if name in evaluators_by_name:
+            raise ValueError(f"Two evaluators have the same name: {name}")
+        evaluators_by_name[name] = evaluator
+    return evaluators_by_name
 
 
 def _str_trace_id(id_: int) -> str:
@@ -475,11 +527,24 @@ def _decode_unix_nano(time_unix_nano: int) -> datetime:
     return datetime.fromtimestamp(time_unix_nano / 1e9, tz=timezone.utc)
 
 
+def _get_task_name(task: ExperimentTask) -> str:
+    """
+    Makes a best-effort attempt to get the name of the task.
+    """
+
+    if isinstance(task, functools.partial):
+        return task.func.__qualname__
+    if hasattr(task, "__qualname__"):
+        return task.__qualname__
+    return str(task)
+
+
 INPUT_VALUE = SpanAttributes.INPUT_VALUE
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
 INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
 OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
 OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
 
-CHAIN = OpenInferenceSpanKindValues.CHAIN
+CHAIN = OpenInferenceSpanKindValues.CHAIN.value
+EVALUATOR = OpenInferenceSpanKindValues.EVALUATOR.value
 JSON = OpenInferenceMimeTypeValues.JSON

phoenix/datasets/tracing.py

@@ -12,12 +12,22 @@ from wrapt import apply_patch, resolve_path, wrap_function_wrapper
 
 
 class SpanModifier:
+    """
+    A class that modifies spans with the specified resource attributes.
+    """
+
     __slots__ = ("_resource",)
 
     def __init__(self, resource: Resource) -> None:
         self._resource = resource
 
     def modify_resource(self, span: ReadableSpan) -> None:
+        """
+        Takes a span and merges in the resource attributes specified in the constructor.
+
+        Args:
+          span: ReadableSpan: the span to modify
+        """
         if (ctx := span._context) is None or ctx.span_id == INVALID_TRACE_ID:
             return
         span._resource = span._resource.merge(self._resource)
@@ -59,6 +69,15 @@ def _monkey_patch_span_init() -> Iterator[None]:
 
 @contextmanager
 def capture_spans(resource: Resource) -> Iterator[SpanModifier]:
+    """
+    A context manager that captures spans and modifies them with the specified resources.
+
+    Args:
+      resource: Resource: The resource to merge into the spans created within the context.
+
+    Returns:
+        modifier: Iterator[SpanModifier]: The span modifier that is active within the context.
+    """
     modifier = SpanModifier(resource)
     with _monkey_patch_span_init():
         token = _ACTIVE_MODIFIER.set(modifier)

phoenix/datasets/types.py

@@ -2,22 +2,27 @@ from __future__ import annotations
 
 from dataclasses import dataclass, field
 from datetime import datetime
-from types import MappingProxyType
+from enum import Enum
 from typing import (
-    TYPE_CHECKING,
     Any,
+    Awaitable,
+    Callable,
     Dict,
     List,
     Mapping,
     Optional,
-    Protocol,
     Sequence,
     Union,
-    runtime_checkable,
 )
 
 from typing_extensions import TypeAlias
 
+
+class AnnotatorKind(Enum):
+    CODE = "CODE"
+    LLM = "LLM"
+
+
 JSONSerializable: TypeAlias = Optional[Union[Dict[str, Any], List[Any], str, int, float, bool]]
 
 ExperimentId: TypeAlias = str
@@ -28,6 +33,8 @@ RepetitionNumber: TypeAlias = int
 ExperimentRunId: TypeAlias = str
 TraceId: TypeAlias = str
 
+TaskOutput: TypeAlias = JSONSerializable
+
 
 @dataclass(frozen=True)
 class Example:
@@ -35,7 +42,7 @@ class Example:
     updated_at: datetime
     input: Mapping[str, JSONSerializable]
     output: Mapping[str, JSONSerializable]
-    metadata: Mapping[str, JSONSerializable] = field(default_factory=lambda: MappingProxyType({}))
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)
 
     @classmethod
     def from_dict(cls, obj: Mapping[str, Any]) -> Example:
@@ -71,7 +78,7 @@ class Experiment:
 
 @dataclass(frozen=True)
 class ExperimentResult:
-    result: JSONSerializable
+    result: TaskOutput
 
     @classmethod
     def from_dict(cls, obj: Optional[Mapping[str, Any]]) -> Optional[ExperimentResult]:
@@ -116,7 +123,7 @@ class EvaluationResult:
     score: Optional[float] = None
     label: Optional[str] = None
     explanation: Optional[str] = None
-    metadata: Mapping[str, JSONSerializable] = field(default_factory=lambda: MappingProxyType({}))
+    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)
 
     @classmethod
     def from_dict(cls, obj: Optional[Mapping[str, Any]]) -> Optional[EvaluationResult]:
@@ -165,48 +172,7 @@ class ExperimentEvaluationRun:
             ValueError("Must specify either result or error")
 
 
-class _HasName(Protocol):
-    name: str
-
-
-class _HasKind(Protocol):
-    @property
-    def annotator_kind(self) -> str: ...
-
-
-@runtime_checkable
-class CanEvaluate(_HasName, _HasKind, Protocol):
-    def evaluate(
-        self,
-        example: Example,
-        experiment_run: ExperimentRun,
-    ) -> EvaluationResult: ...
-
-
-@runtime_checkable
-class CanAsyncEvaluate(_HasName, _HasKind, Protocol):
-    async def async_evaluate(
-        self,
-        example: Example,
-        experiment_run: ExperimentRun,
-    ) -> EvaluationResult: ...
-
-
-ExperimentEvaluator: TypeAlias = Union[CanEvaluate, CanAsyncEvaluate]
-
-
-# Someday we'll do type checking in unit tests.
-if TYPE_CHECKING:
-
-    class _EvaluatorDummy:
-        annotator_kind: str
-        name: str
-
-        def evaluate(self, _: Example, __: ExperimentRun) -> EvaluationResult:
-            raise NotImplementedError
-
-        async def async_evaluate(self, _: Example, __: ExperimentRun) -> EvaluationResult:
-            raise NotImplementedError
-
-    _: ExperimentEvaluator
-    _ = _EvaluatorDummy()
+ExperimentTask: TypeAlias = Union[
+    Callable[[Example], TaskOutput],
+    Callable[[Example], Awaitable[TaskOutput]],
+]

phoenix/db/insertion/dataset.py

@@ -1,17 +1,17 @@
 import logging
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from enum import Enum
 from itertools import chain
 from typing import (
     Any,
     Awaitable,
+    Dict,
     FrozenSet,
     Iterable,
     Iterator,
     Mapping,
     Optional,
-    Sequence,
     Union,
     cast,
 )
@@ -30,7 +30,16 @@ DatasetVersionId: TypeAlias = int
 DatasetExampleId: TypeAlias = int
 DatasetExampleRevisionId: TypeAlias = int
 SpanRowId: TypeAlias = int
-Examples: TypeAlias = Iterable[Mapping[str, Any]]
+
+
+@dataclass(frozen=True)
+class ExampleContent:
+    input: Dict[str, Any] = field(default_factory=dict)
+    output: Dict[str, Any] = field(default_factory=dict)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+Examples: TypeAlias = Iterable[ExampleContent]
 
 
 @dataclass(frozen=True)
@@ -149,14 +158,10 @@ async def add_dataset_examples(
     session: AsyncSession,
     name: str,
     examples: Union[Examples, Awaitable[Examples]],
-    input_keys: Sequence[str],
-    output_keys: Sequence[str],
-    metadata_keys: Sequence[str] = (),
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     action: DatasetAction = DatasetAction.CREATE,
 ) -> Optional[DatasetExampleAdditionEvent]:
-    keys = DatasetKeys(frozenset(input_keys), frozenset(output_keys), frozenset(metadata_keys))
     created_at = datetime.now(timezone.utc)
     dataset_id: Optional[DatasetId] = None
     if action is DatasetAction.APPEND and name:
@@ -173,9 +178,7 @@ async def add_dataset_examples(
                 created_at=created_at,
             )
         except Exception:
-            logger.exception(
-                f"Fail to insert dataset: {input_keys=}, {output_keys=}, {metadata_keys=}"
-            )
+            logger.exception(f"Failed to insert dataset: {name=}")
             raise
     try:
         dataset_version_id = await insert_dataset_version(
@@ -184,7 +187,7 @@ async def add_dataset_examples(
             created_at=created_at,
         )
     except Exception:
-        logger.exception(f"Fail to insert dataset version for {dataset_id=}")
+        logger.exception(f"Failed to insert dataset version for {dataset_id=}")
         raise
     for example in (await examples) if isinstance(examples, Awaitable) else examples:
         try:
@@ -194,21 +197,21 @@ async def add_dataset_examples(
                 created_at=created_at,
             )
         except Exception:
-            logger.exception(f"Fail to insert dataset example for {dataset_id=}")
+            logger.exception(f"Failed to insert dataset example for {dataset_id=}")
             raise
         try:
             await insert_dataset_example_revision(
                 session=session,
                 dataset_version_id=dataset_version_id,
                 dataset_example_id=dataset_example_id,
-                input={key: example.get(key) for key in keys.input},
-                output={key: example.get(key) for key in keys.output},
-                metadata={key: example.get(key) for key in keys.metadata},
+                input=example.input,
+                output=example.output,
+                metadata=example.metadata,
                 created_at=created_at,
             )
         except Exception:
             logger.exception(
-                f"Fail to insert dataset example revision for {dataset_version_id=}, "
+                f"Failed to insert dataset example revision for {dataset_version_id=}, "
                 f"{dataset_example_id=}"
             )
             raise