arize-phoenix 4.12.1rc1__py3-none-any.whl → 4.15.0__py3-none-any.whl

phoenix/server/api/dataloaders/experiment_error_rates.py

@@ -1,16 +1,14 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
 )
 
-from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy import case, func, select
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExperimentID: TypeAlias = int
 ErrorRate: TypeAlias = float
@@ -21,23 +19,43 @@ Result: TypeAlias = Optional[ErrorRate]
 class ExperimentErrorRatesDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
     ) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
     async def _load_fn(self, keys: List[Key]) -> List[Result]:
         experiment_ids = keys
+        resolved_experiment_ids = (
+            select(models.Experiment.id)
+            .where(models.Experiment.id.in_(set(experiment_ids)))
+            .subquery()
+        )
+        query = (
+            select(
+                resolved_experiment_ids.c.id,
+                case(
+                    (
+                        func.count(models.ExperimentRun.id) != 0,
+                        func.count(models.ExperimentRun.error)
+                        / func.count(models.ExperimentRun.id),
+                    ),
+                    else_=None,
+                ),
+            )
+            .outerjoin_from(
+                from_=resolved_experiment_ids,
+                target=models.ExperimentRun,
+                onclause=resolved_experiment_ids.c.id == models.ExperimentRun.experiment_id,
+            )
+            .group_by(resolved_experiment_ids.c.id)
+        )
         async with self._db() as session:
             error_rates = {
                 experiment_id: error_rate
-                async for experiment_id, error_rate in await session.stream(
-                    select(
-                        models.ExperimentRun.experiment_id,
-                        func.count(models.ExperimentRun.error) / func.count(),
-                    )
-                    .group_by(models.ExperimentRun.experiment_id)
-                    .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
-                )
+                async for experiment_id, error_rate in await session.stream(query)
             }
-        return [error_rates.get(experiment_id) for experiment_id in experiment_ids]
+        return [
+            error_rates.get(experiment_id, ValueError(f"Unknown experiment ID: {experiment_id}"))
+            for experiment_id in experiment_ids
+        ]

phoenix/server/api/dataloaders/experiment_run_counts.py

@@ -1,15 +1,13 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
 )
 
 from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExperimentID: TypeAlias = int
 RunCount: TypeAlias = int
@@ -20,21 +18,34 @@ Result: TypeAlias = RunCount
 class ExperimentRunCountsDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
     ) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 
     async def _load_fn(self, keys: List[Key]) -> List[Result]:
         experiment_ids = keys
+        resolved_experiment_ids = (
+            select(models.Experiment.id)
+            .where(models.Experiment.id.in_(set(experiment_ids)))
+            .subquery()
+        )
+        query = (
+            select(
+                resolved_experiment_ids.c.id,
+                func.count(models.ExperimentRun.experiment_id),
+            )
+            .outerjoin_from(
+                from_=resolved_experiment_ids,
+                target=models.ExperimentRun,
+                onclause=resolved_experiment_ids.c.id == models.ExperimentRun.experiment_id,
+            )
+            .group_by(resolved_experiment_ids.c.id)
+        )
         async with self._db() as session:
             run_counts = {
                 experiment_id: run_count
-                async for experiment_id, run_count in await session.stream(
-                    select(models.ExperimentRun.experiment_id, func.count())
-                    .where(models.ExperimentRun.experiment_id.in_(set(experiment_ids)))
-                    .group_by(models.ExperimentRun.experiment_id)
-                )
+                async for experiment_id, run_count in await session.stream(query)
             }
         return [
             run_counts.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))

phoenix/server/api/dataloaders/experiment_sequence_number.py

@@ -1,16 +1,14 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
 )
 
 from sqlalchemy import distinct, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ExperimentId: TypeAlias = int
 Key: TypeAlias = ExperimentId
@@ -18,7 +16,7 @@ Result: TypeAlias = Optional[int]
 
 
 class ExperimentSequenceNumberDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/latency_ms_quantile.py

@@ -2,9 +2,7 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
     AsyncIterator,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -36,6 +34,7 @@ from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["span", "trace"]
@@ -88,7 +87,7 @@ class LatencyMsQuantileCache(
 class LatencyMsQuantileDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

phoenix/server/api/dataloaders/min_start_or_max_end_times.py

@@ -1,8 +1,6 @@
 from collections import defaultdict
 from datetime import datetime
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -12,12 +10,12 @@ from typing import (
 
 from cachetools import LFUCache
 from sqlalchemy import func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias, assert_never
 
 from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
+from phoenix.server.types import DbSessionFactory
 
 Kind: TypeAlias = Literal["start", "end"]
 ProjectRowId: TypeAlias = int
@@ -50,7 +48,7 @@ class MinStartOrMaxEndTimeCache(
 class MinStartOrMaxEndTimeDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

phoenix/server/api/dataloaders/project_by_name.py

@@ -1,12 +1,12 @@
 from collections import defaultdict
-from typing import AsyncContextManager, Callable, DefaultDict, List, Optional
+from typing import DefaultDict, List, Optional
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 ProjectName: TypeAlias = str
 Key: TypeAlias = ProjectName
@@ -14,7 +14,7 @@ Result: TypeAlias = Optional[models.Project]
 
 
 class ProjectByNameDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/record_counts.py

@@ -2,8 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -13,13 +11,13 @@ from typing import (
 
 from cachetools import LFUCache, TTLCache
 from sqlalchemy import Select, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias, assert_never
 
 from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["span", "trace"]
@@ -69,7 +67,7 @@ class RecordCountCache(
 class RecordCountDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

phoenix/server/api/dataloaders/span_annotations.py

@@ -1,25 +1,23 @@
 from collections import defaultdict
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
 )
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[SpanAnnotation]
 
 
 class SpanAnnotationsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/span_dataset_examples.py

@@ -0,0 +1,36 @@
+from typing import (
+    Dict,
+    List,
+)
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanID: TypeAlias = int
+Key: TypeAlias = SpanID
+Result: TypeAlias = List[models.DatasetExample]
+
+
+class SpanDatasetExamplesDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        span_rowids = keys
+        async with self._db() as session:
+            dataset_examples: Dict[Key, List[models.DatasetExample]] = {
+                span_rowid: [] for span_rowid in span_rowids
+            }
+            async for span_rowid, dataset_example in await session.stream(
+                select(models.Span.id, models.DatasetExample)
+                .select_from(models.Span)
+                .join(models.DatasetExample, models.DatasetExample.span_rowid == models.Span.id)
+                .where(models.Span.id.in_(span_rowids))
+            ):
+                dataset_examples[span_rowid].append(dataset_example)
+        return [dataset_examples.get(span_rowid, []) for span_rowid in span_rowids]

phoenix/server/api/dataloaders/span_descendants.py

@@ -1,19 +1,17 @@
 from random import randint
 from typing import (
-    AsyncContextManager,
-    Callable,
     Dict,
     List,
 )
 
 from aioitertools.itertools import groupby
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 SpanId: TypeAlias = str
 
@@ -22,7 +20,7 @@ Result: TypeAlias = List[models.Span]
 
 
 class SpanDescendantsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/span_evaluations.py

@@ -1,25 +1,23 @@
 from collections import defaultdict
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
 )
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.server.api.types.Evaluation import SpanEvaluation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[SpanEvaluation]
 
 
 class SpanEvaluationsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/span_projects.py

@@ -1,11 +1,11 @@
-from typing import AsyncContextManager, Callable, List, Union
+from typing import List, Union
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 SpanID: TypeAlias = int
 Key: TypeAlias = SpanID
@@ -13,7 +13,7 @@ Result: TypeAlias = models.Project
 
 
 class SpanProjectsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/token_counts.py

@@ -2,8 +2,6 @@ from collections import defaultdict
 from datetime import datetime
 from typing import (
     Any,
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
     Literal,
@@ -14,7 +12,6 @@ from typing import (
 from cachetools import LFUCache, TTLCache
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import Select, func, select
-from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql.functions import coalesce
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias
@@ -22,6 +19,7 @@ from typing_extensions import TypeAlias
 from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
 Kind: TypeAlias = Literal["prompt", "completion", "total"]
@@ -71,7 +69,7 @@ class TokenCountCache(
 class TokenCountDataLoader(DataLoader[Key, Result]):
     def __init__(
         self,
-        db: Callable[[], AsyncContextManager[AsyncSession]],
+        db: DbSessionFactory,
         cache_map: Optional[AbstractCache[Key, Result]] = None,
     ) -> None:
         super().__init__(

phoenix/server/api/dataloaders/trace_evaluations.py

@@ -1,25 +1,23 @@
 from collections import defaultdict
 from typing import (
-    AsyncContextManager,
-    Callable,
     DefaultDict,
     List,
 )
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.server.api.types.Evaluation import TraceEvaluation
+from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
 Result: TypeAlias = List[TraceEvaluation]
 
 
 class TraceEvaluationsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/dataloaders/trace_row_ids.py

@@ -1,17 +1,15 @@
 from typing import (
-    AsyncContextManager,
-    Callable,
     List,
     Optional,
     Tuple,
 )
 
 from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
 
 TraceId: TypeAlias = str
 Key: TypeAlias = TraceId
@@ -21,7 +19,7 @@ Result: TypeAlias = Optional[Tuple[TraceRowId, ProjectRowId]]
 
 
 class TraceRowIdsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+    def __init__(self, db: DbSessionFactory) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
 

phoenix/server/api/input_types/{CreateSpanAnnotationsInput.py → CreateSpanAnnotationInput.py}

@@ -4,12 +4,14 @@ import strawberry
 from strawberry.relay import GlobalID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
 
 @strawberry.input
-class CreateSpanAnnotationsInput:
+class CreateSpanAnnotationInput:
     span_id: GlobalID
     name: str
-    annotator_kind: str
+    annotator_kind: AnnotatorKind
     label: Optional[str] = None
     score: Optional[float] = None
     explanation: Optional[str] = None

phoenix/server/api/input_types/{CreateTraceAnnotationsInput.py → CreateTraceAnnotationInput.py}

@@ -4,12 +4,14 @@ import strawberry
 from strawberry.relay import GlobalID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
 
 @strawberry.input
-class CreateTraceAnnotationsInput:
+class CreateTraceAnnotationInput:
     trace_id: GlobalID
     name: str
-    annotator_kind: str
+    annotator_kind: AnnotatorKind
     label: Optional[str] = None
     score: Optional[float] = None
     explanation: Optional[str] = None

phoenix/server/api/input_types/{PatchAnnotationsInput.py → PatchAnnotationInput.py}

@@ -5,12 +5,14 @@ from strawberry import UNSET
 from strawberry.relay import GlobalID
 from strawberry.scalars import JSON
 
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
 
 @strawberry.input
-class PatchAnnotationsInput:
+class PatchAnnotationInput:
     annotation_id: GlobalID
     name: Optional[str] = UNSET
-    annotator_kind: Optional[str] = UNSET
+    annotator_kind: Optional[AnnotatorKind] = UNSET
     label: Optional[str] = UNSET
     score: Optional[float] = UNSET
     explanation: Optional[str] = UNSET

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -7,10 +7,11 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
-from phoenix.server.api.input_types.CreateSpanAnnotationsInput import CreateSpanAnnotationsInput
+from phoenix.server.api.input_types.CreateSpanAnnotationInput import CreateSpanAnnotationInput
 from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
-from phoenix.server.api.input_types.PatchAnnotationsInput import PatchAnnotationsInput
+from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
 from phoenix.server.api.mutations.auth import IsAuthenticated
+from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
 
@@ -18,13 +19,14 @@ from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_
 @strawberry.type
 class SpanAnnotationMutationPayload:
     span_annotations: List[SpanAnnotation]
+    query: Query
 
 
 @strawberry.type
 class SpanAnnotationMutationMixin:
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
     async def create_span_annotations(
-        self, info: Info[Context, None], input: List[CreateSpanAnnotationsInput]
+        self, info: Info[Context, None], input: List[CreateSpanAnnotationInput]
     ) -> SpanAnnotationMutationPayload:
         inserted_annotations: Sequence[models.SpanAnnotation] = []
         async with info.context.db() as session:
@@ -35,7 +37,7 @@ class SpanAnnotationMutationMixin:
                     label=annotation.label,
                     score=annotation.score,
                     explanation=annotation.explanation,
-                    annotator_kind=annotation.annotator_kind,
+                    annotator_kind=annotation.annotator_kind.value,
                     metadata_=annotation.metadata,
                 )
                 for annotation in input
@@ -49,12 +51,13 @@ class SpanAnnotationMutationMixin:
         return SpanAnnotationMutationPayload(
             span_annotations=[
                 to_gql_span_annotation(annotation) for annotation in inserted_annotations
-            ]
+            ],
+            query=Query(),
         )
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
     async def patch_span_annotations(
-        self, info: Info[Context, None], input: List[PatchAnnotationsInput]
+        self, info: Info[Context, None], input: List[PatchAnnotationInput]
     ) -> SpanAnnotationMutationPayload:
         patched_annotations = []
         async with info.context.db() as session:
@@ -66,7 +69,13 @@ class SpanAnnotationMutationMixin:
                     column.key: patch_value
                     for column, patch_value, column_is_nullable in (
                         (models.SpanAnnotation.name, annotation.name, False),
-                        (models.SpanAnnotation.annotator_kind, annotation.annotator_kind, False),
+                        (
+                            models.SpanAnnotation.annotator_kind,
+                            annotation.annotator_kind.value
+                            if annotation.annotator_kind is not None
+                            else None,
+                            False,
+                        ),
                         (models.SpanAnnotation.label, annotation.label, True),
                         (models.SpanAnnotation.score, annotation.score, True),
                         (models.SpanAnnotation.explanation, annotation.explanation, True),
@@ -83,7 +92,7 @@ class SpanAnnotationMutationMixin:
                 if span_annotation is not None:
                     patched_annotations.append(to_gql_span_annotation(span_annotation))
 
-        return SpanAnnotationMutationPayload(span_annotations=patched_annotations)
+        return SpanAnnotationMutationPayload(span_annotations=patched_annotations, query=Query())
 
     @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
     async def delete_span_annotations(
@@ -105,4 +114,6 @@ class SpanAnnotationMutationMixin:
             deleted_annotations_gql = [
                 to_gql_span_annotation(annotation) for annotation in deleted_annotations
             ]
-        return SpanAnnotationMutationPayload(span_annotations=deleted_annotations_gql)
+        return SpanAnnotationMutationPayload(
+            span_annotations=deleted_annotations_gql, query=Query()
+        )