gather span annotations by name

Author: axiomofjoy

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

@@ -213,10 +213,7 @@ async def add_spans_to_dataset(
                                 if k in nonprivate_span_attributes
                             },
                             "span_kind": span.span_kind,
-                            "annotations": [
-                                _serialize_span_annotation(span_annotation)
-                                for span_annotation in span.span_annotations
-                            ],
+                            "annotations": _gather_span_annotations_by_name(span.span_annotations),
                         },
                         DatasetExampleRevision.revision_kind.key: "CREATE",
                     }
@@ -602,6 +599,19 @@ def _to_orm_revision(
     }
 
 
+def _gather_span_annotations_by_name(
+    span_annotations: list[models.SpanAnnotation],
+) -> dict[str, list[dict[str, Any]]]:
+    span_annotations_by_name: dict[str, list[dict[str, Any]]] = {}
+    for span_annotation in span_annotations:
+        if span_annotation.name not in span_annotations_by_name:
+            span_annotations_by_name[span_annotation.name] = []
+        span_annotations_by_name[span_annotation.name].append(
+            _serialize_span_annotation(span_annotation)
+        )
+    return span_annotations_by_name
+
+
 def _serialize_span_annotation(span_annotation: models.SpanAnnotation) -> dict[str, Any]:
     return {
         "label": span_annotation.label,