feat: trace annotation dataloader (#12699)

* refactor: use dataloader for trace_annotations resolver with filter support

* Fix trace annotation sort stability

* Fix formatting in Trace.py to pass CI lint check

* Fix mypy type error in test_TraceAnnotation.py

* Fix default sort to preserve original created_at DESC behavior

The dataloader refactor inadvertently changed the default sort from
created_at DESC to name ASC. Restore the original default so clients
that omit the sort argument still get most-recent-first ordering.

Author: ehutt

app/schema.graphql

@@ -3564,7 +3564,7 @@ type Trace implements Node {
   spans(first: Int!, last: Int, after: String, before: String, rootSpansOnly: Boolean, orphanSpanAsRootSpan: Boolean = true): SpanConnection!
 
   """Annotations associated with the trace."""
-  traceAnnotations(sort: TraceAnnotationSort = null): [TraceAnnotation!]!
+  traceAnnotations(sort: TraceAnnotationSort, filter: AnnotationFilter = null): [TraceAnnotation!]!
 
   """Summarizes each annotation (by name) associated with the trace"""
   traceAnnotationSummaries(filter: AnnotationFilter = null): [AnnotationSummary!]!

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

@@ -19,7 +19,10 @@
 from phoenix.server.api.context import Context
 from phoenix.server.api.extensions import RequireForwardPaginationExtension
 from phoenix.server.api.input_types.AnnotationFilter import AnnotationFilter, satisfies_filter
-from phoenix.server.api.input_types.TraceAnnotationSort import TraceAnnotationSort
+from phoenix.server.api.input_types.TraceAnnotationSort import (
+    TraceAnnotationColumn,
+    TraceAnnotationSort,
+)
 from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.pagination import (
@@ -283,19 +286,24 @@ async def spans(
     async def trace_annotations(
         self,
         info: Info[Context, None],
-        sort: Optional[TraceAnnotationSort] = None,
+        sort: Optional[TraceAnnotationSort] = UNSET,
+        filter: Optional[AnnotationFilter] = None,
     ) -> list[TraceAnnotation]:
-        async with info.context.db.read() as session:
-            stmt = select(models.TraceAnnotation).filter_by(trace_rowid=self.id)
-            if sort:
-                sort_col = getattr(models.TraceAnnotation, sort.col.value)
-                if sort.dir is SortDir.desc:
-                    stmt = stmt.order_by(sort_col.desc(), models.TraceAnnotation.id.desc())
-                else:
-                    stmt = stmt.order_by(sort_col.asc(), models.TraceAnnotation.id.asc())
-            else:
-                stmt = stmt.order_by(models.TraceAnnotation.created_at.desc())
-            annotations = await session.scalars(stmt)
+        annotations = list(await info.context.data_loaders.trace_annotations_by_trace.load(self.id))
+        sort_key = TraceAnnotationColumn.createdAt.value
+        sort_descending = True
+        if filter:
+            annotations = [
+                annotation for annotation in annotations if satisfies_filter(annotation, filter)
+            ]
+        if sort:
+            sort_key = sort.col.value
+            sort_descending = sort.dir is SortDir.desc
+        annotations = sorted(
+            annotations,
+            key=lambda annotation: (getattr(annotation, sort_key), annotation.id),
+            reverse=sort_descending,
+        )
         return [
             TraceAnnotation(id=annotation.id, db_record=annotation) for annotation in annotations
         ]

tests/unit/server/api/types/test_TraceAnnotation.py

@@ -194,3 +194,67 @@ async def test_annotating_a_trace(
             select(models.TraceAnnotation).where(models.TraceAnnotation.id == annotation_id)
         )
     assert not orm_annotation
+
+
+async def test_trace_annotations_sort_uses_id_as_tiebreaker(
+    gql_client: AsyncGraphQLClient,
+    db: DbSessionFactory,
+    project_with_a_single_trace_and_span: Any,
+) -> None:
+    async with db() as session:
+        trace_id = await session.scalar(select(models.Trace.id))
+        assert trace_id is not None
+        created_at = datetime.fromisoformat("2021-01-01T00:00:00.000+00:00")
+        session.add_all(
+            [
+                models.TraceAnnotation(
+                    trace_rowid=trace_id,
+                    name="same-name",
+                    label="older-id",
+                    score=0.1,
+                    explanation=None,
+                    metadata_={},
+                    annotator_kind="HUMAN",
+                    created_at=created_at,
+                    updated_at=created_at,
+                    identifier="a",
+                    source="API",
+                    user_id=None,
+                ),
+                models.TraceAnnotation(
+                    trace_rowid=trace_id,
+                    name="same-name",
+                    label="newer-id",
+                    score=0.2,
+                    explanation=None,
+                    metadata_={},
+                    annotator_kind="HUMAN",
+                    created_at=created_at,
+                    updated_at=created_at,
+                    identifier="b",
+                    source="API",
+                    user_id=None,
+                ),
+            ]
+        )
+
+    response = await gql_client.execute(
+        query="""
+            query GetTraceAnnotations($id: ID!) {
+                node(id: $id) {
+                    ... on Trace {
+                        traceAnnotations(sort: {col: createdAt, dir: desc}) {
+                            label
+                        }
+                    }
+                }
+            }
+        """,
+        variables={"id": str(GlobalID("Trace", str(trace_id)))},
+    )
+    assert not response.errors
+    assert response.data is not None
+    assert response.data["node"]["traceAnnotations"] == [
+        {"label": "newer-id"},
+        {"label": "older-id"},
+    ]