fix: Propagate span annotation metadata to examples on all mutations

src/phoenix/server/api/mutations/dataset_mutations.py

@@ -230,7 +230,7 @@ async def add_examples_to_dataset(
     ) -> DatasetMutationPayload:
         dataset_id = input.dataset_id
         # Extract the span rowids from the input examples if they exist
-        span_ids = span_ids = [example.span_id for example in input.examples if example.span_id]
+        span_ids = [example.span_id for example in input.examples if example.span_id]
         span_rowids = {
             from_global_id_with_expected_type(global_id=span_id, expected_type_name=Span.__name__)
             for span_id in set(span_ids)
@@ -260,16 +260,45 @@ async def add_examples_to_dataset(
                 )
                 .returning(models.DatasetVersion.id)
             )
+
+            # Fetch spans and span annotations
             spans = (
                 await session.execute(
                     select(models.Span.id)
                     .select_from(models.Span)
                     .where(models.Span.id.in_(span_rowids))
                 )
             ).all()
-            # Just validate that the number of spans matches the number of span_ids
-            # to ensure that the span_ids are valid
-            assert len(spans) == len(span_rowids)
+
+            span_annotations = (
+                await session.execute(
+                    select(
+                        models.SpanAnnotation.span_rowid,
+                        models.SpanAnnotation.name,
+                        models.SpanAnnotation.label,
+                        models.SpanAnnotation.score,
+                        models.SpanAnnotation.explanation,
+                        models.SpanAnnotation.metadata_,
+                        models.SpanAnnotation.annotator_kind,
+                    )
+                    .select_from(models.SpanAnnotation)
+                    .where(models.SpanAnnotation.span_rowid.in_(span_rowids))
+                )
+            ).all()
+
+            span_annotations_by_span: Dict[int, Dict[Any, Any]] = {span.id: {} for span in spans}
+            for annotation in span_annotations:
+                span_id = annotation.span_rowid
+                if span_id not in span_annotations_by_span:
+                    span_annotations_by_span[span_id] = dict()
+                span_annotations_by_span[span_id][annotation.name] = {
+                    "label": annotation.label,
+                    "score": annotation.score,
+                    "explanation": annotation.explanation,
+                    "metadata": annotation.metadata_,
+                    "annotator_kind": annotation.annotator_kind,
+                }
+
             DatasetExample = models.DatasetExample
             dataset_example_rowids = (
                 await session.scalars(
@@ -291,21 +320,32 @@ async def add_examples_to_dataset(
             assert len(dataset_example_rowids) == len(input.examples)
             assert all(map(lambda id: isinstance(id, int), dataset_example_rowids))
             DatasetExampleRevision = models.DatasetExampleRevision
-            await session.execute(
-                insert(DatasetExampleRevision),
-                [
+
+            dataset_example_revisions = []
+            for dataset_example_rowid, example in zip(dataset_example_rowids, input.examples):
+                span_annotation = {}
+                if example.span_id:
+                    span_id = from_global_id_with_expected_type(
+                        global_id=example.span_id,
+                        expected_type_name=Span.__name__,
+                    )
+                    span_annotation = span_annotations_by_span.get(span_id, {})
+                dataset_example_revisions.append(
                     {
                         DatasetExampleRevision.dataset_example_id.key: dataset_example_rowid,
                         DatasetExampleRevision.dataset_version_id.key: dataset_version_rowid,
                         DatasetExampleRevision.input.key: example.input,
                         DatasetExampleRevision.output.key: example.output,
-                        DatasetExampleRevision.metadata_.key: example.metadata,
+                        DatasetExampleRevision.metadata_.key: {
+                            **(example.metadata or {}),
+                            "annotations": span_annotation,
+                        },
                         DatasetExampleRevision.revision_kind.key: "CREATE",
                     }
-                    for dataset_example_rowid, example in zip(
-                        dataset_example_rowids, input.examples
-                    )
-                ],
+                )
+            await session.execute(
+                insert(DatasetExampleRevision),
+                dataset_example_revisions,
             )
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))

src/phoenix/server/api/types/Span.py

@@ -240,7 +240,7 @@ async def descendants(
     @strawberry.field(
         description="The span's attributes translated into an example revision for a dataset",
     )  # type: ignore
-    def as_example_revision(self) -> SpanAsExampleRevision:
+    async def as_example_revision(self, info: Info[Context, None]) -> SpanAsExampleRevision:
         db_span = self.db_span
         attributes = db_span.attributes
         span_io = _SpanIO(
@@ -256,10 +256,28 @@ def as_example_revision(self) -> SpanAsExampleRevision:
             llm_output_messages=get_attribute_value(attributes, LLM_OUTPUT_MESSAGES),
             retrieval_documents=get_attribute_value(attributes, RETRIEVAL_DOCUMENTS),
         )
+
+        # Fetch annotations associated with this span
+        span_annotations = await self.span_annotations(info)
+        annotations = dict()
+        for annotation in span_annotations:
+            annotations[annotation.name] = {
+                "label": annotation.label,
+                "score": annotation.score,
+                "explanation": annotation.explanation,
+                "metadata": annotation.metadata,
+                "annotator_kind": annotation.annotator_kind.value,
+            }
+        # Merge annotations into the metadata
+        metadata = {
+            **attributes,
+            "annotations": annotations,
+        }
+
         return SpanAsExampleRevision(
             input=get_dataset_example_input(span_io),
             output=get_dataset_example_output(span_io),
-            metadata=attributes,
+            metadata=metadata,
         )
 
     @strawberry.field(description="The project that this span belongs to.")  # type: ignore