compute span annotation summaries

Author: axiomofjoy

src/phoenix/server/api/helpers/annotations.py

@@ -0,0 +1,35 @@
+from collections import defaultdict
+from dataclasses import asdict, dataclass
+from typing import Optional
+
+import pandas as pd
+
+from phoenix.db import models
+from phoenix.server.api.types.Span import AnnotationSummary
+
+
+def compute_span_annotation_summaries(
+    annotations: list[models.SpanAnnotation],
+) -> list[AnnotationSummary]:
+    @dataclass
+    class Metrics:
+        record_count: int = 0
+        label_count: int = 0
+        score_sum: float = 0
+        score_count: int = 0
+
+    summaries: defaultdict[str, defaultdict[Optional[str], Metrics]] = defaultdict(
+        lambda: defaultdict(Metrics)
+    )
+    for annotation in annotations:
+        metrics = summaries[annotation.name][annotation.label]
+        metrics.record_count += 1
+        metrics.label_count += int(annotation.label is not None)
+        metrics.score_sum += annotation.score or 0
+        metrics.score_count += int(annotation.score is not None)
+
+    result: list[AnnotationSummary] = []
+    for name, label_metrics in summaries.items():
+        rows = [{"label": label, **asdict(metrics)} for label, metrics in label_metrics.items()]
+        result.append(AnnotationSummary(name=name, df=pd.DataFrame(rows), simple_avg=True))
+    return result

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

@@ -1,14 +1,11 @@
 import json
 from asyncio import gather
-from collections import defaultdict
 from collections.abc import Mapping
-from dataclasses import asdict, dataclass
 from datetime import datetime
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Iterable, Optional, cast
 
 import numpy as np
-import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
 from strawberry import ID, UNSET
@@ -19,6 +16,7 @@
 import phoenix.trace.schemas as trace_schema
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.helpers.annotations import compute_span_annotation_summaries
 from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
@@ -558,28 +556,7 @@ async def span_annotation_summaries(
                 annotation for annotation in annotations if satisfies_filter(annotation, filter)
             ]
 
-        @dataclass
-        class Metrics:
-            record_count: int = 0
-            label_count: int = 0
-            score_sum: float = 0
-            score_count: int = 0
-
-        summaries: defaultdict[str, defaultdict[Optional[str], Metrics]] = defaultdict(
-            lambda: defaultdict(Metrics)
-        )
-        for annotation in annotations:
-            metrics = summaries[annotation.name][annotation.label]
-            metrics.record_count += 1
-            metrics.label_count += int(annotation.label is not None)
-            metrics.score_sum += annotation.score or 0
-            metrics.score_count += int(annotation.score is not None)
-
-        result: list[AnnotationSummary] = []
-        for name, label_metrics in summaries.items():
-            rows = [{"label": label, **asdict(metrics)} for label, metrics in label_metrics.items()]
-            result.append(AnnotationSummary(name=name, df=pd.DataFrame(rows), simple_avg=True))
-        return result
+        return compute_span_annotation_summaries(annotations)
 
     @strawberry.field(
         description="Evaluations of the documents associated with the span, e.g. "