arize-phoenix 12.7.1__py3-none-any.whl → 12.9.0__py3-none-any.whl

phoenix/server/api/mutations/dataset_label_mutations.py

@@ -4,6 +4,8 @@ import sqlalchemy
 import strawberry
 from sqlalchemy import delete, select
 from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql import tuple_
 from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 from strawberry import UNSET
 from strawberry.relay.types import GlobalID
@@ -15,7 +17,7 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.Dataset import Dataset
-from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
+from phoenix.server.api.types.DatasetLabel import DatasetLabel
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 
 
@@ -24,11 +26,13 @@ class CreateDatasetLabelInput:
     name: str
     description: Optional[str] = UNSET
     color: str
+    dataset_ids: Optional[list[GlobalID]] = UNSET
 
 
 @strawberry.type
 class CreateDatasetLabelMutationPayload:
     dataset_label: DatasetLabel
+    datasets: list[Dataset]
 
 
 @strawberry.input
@@ -41,39 +45,16 @@ class DeleteDatasetLabelsMutationPayload:
     dataset_labels: list[DatasetLabel]
 
 
-@strawberry.input
-class UpdateDatasetLabelInput:
-    dataset_label_id: GlobalID
-    name: str
-    description: Optional[str] = None
-    color: str
-
-
-@strawberry.type
-class UpdateDatasetLabelMutationPayload:
-    dataset_label: DatasetLabel
-
-
 @strawberry.input
 class SetDatasetLabelsInput:
+    dataset_id: GlobalID
     dataset_label_ids: list[GlobalID]
-    dataset_ids: list[GlobalID]
 
 
 @strawberry.type
 class SetDatasetLabelsMutationPayload:
-    query: "Query"
-
-
-@strawberry.input
-class UnsetDatasetLabelsInput:
-    dataset_label_ids: list[GlobalID]
-    dataset_ids: list[GlobalID]
-
-
-@strawberry.type
-class UnsetDatasetLabelsMutationPayload:
-    query: "Query"
+    query: Query
+    dataset: Dataset
 
 
 @strawberry.type
@@ -87,50 +68,56 @@ class DatasetLabelMutationMixin:
         name = input.name
         description = input.description
         color = input.color
+        dataset_rowids: dict[
+            int, None
+        ] = {}  # use dictionary to de-duplicate while preserving order
+        if input.dataset_ids:
+            for dataset_id in input.dataset_ids:
+                try:
+                    dataset_rowid = from_global_id_with_expected_type(dataset_id, Dataset.__name__)
+                except ValueError:
+                    raise BadRequest(f"Invalid dataset ID: {dataset_id}")
+                dataset_rowids[dataset_rowid] = None
+
         async with info.context.db() as session:
             dataset_label_orm = models.DatasetLabel(name=name, description=description, color=color)
             session.add(dataset_label_orm)
             try:
-                await session.commit()
+                await session.flush()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A dataset label named '{name}' already exists")
             except sqlalchemy.exc.StatementError as error:
                 raise BadRequest(str(error.orig))
-        return CreateDatasetLabelMutationPayload(
-            dataset_label=to_gql_dataset_label(dataset_label_orm)
-        )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
-    async def update_dataset_label(
-        self, info: Info[Context, None], input: UpdateDatasetLabelInput
-    ) -> UpdateDatasetLabelMutationPayload:
-        if not input.name or not input.name.strip():
-            raise BadRequest("Dataset label name cannot be empty")
-
-        try:
-            dataset_label_id = from_global_id_with_expected_type(
-                input.dataset_label_id, DatasetLabel.__name__
-            )
-        except ValueError:
-            raise BadRequest(f"Invalid dataset label ID: {input.dataset_label_id}")
-
-        async with info.context.db() as session:
-            dataset_label_orm = await session.get(models.DatasetLabel, dataset_label_id)
-            if not dataset_label_orm:
-                raise NotFound(f"DatasetLabel with ID {input.dataset_label_id} not found")
-
-            dataset_label_orm.name = input.name.strip()
-            dataset_label_orm.description = input.description
-            dataset_label_orm.color = input.color.strip()
-
-            try:
+            datasets_by_id: dict[int, models.Dataset] = {}
+            if dataset_rowids:
+                datasets_by_id = {
+                    dataset.id: dataset
+                    for dataset in await session.scalars(
+                        select(models.Dataset).where(models.Dataset.id.in_(dataset_rowids.keys()))
+                    )
+                }
+                if len(datasets_by_id) < len(dataset_rowids):
+                    raise NotFound("One or more datasets not found")
+                session.add_all(
+                    [
+                        models.DatasetsDatasetLabel(
+                            dataset_id=dataset_rowid,
+                            dataset_label_id=dataset_label_orm.id,
+                        )
+                        for dataset_rowid in dataset_rowids
+                    ]
+                )
                 await session.commit()
-            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
-                raise Conflict(f"A dataset label named '{input.name}' already exists")
-            except sqlalchemy.exc.StatementError as error:
-                raise BadRequest(str(error.orig))
-        return UpdateDatasetLabelMutationPayload(
-            dataset_label=to_gql_dataset_label(dataset_label_orm)
+
+        return CreateDatasetLabelMutationPayload(
+            dataset_label=DatasetLabel(id=dataset_label_orm.id, db_record=dataset_label_orm),
+            datasets=[
+                Dataset(
+                    id=datasets_by_id[dataset_rowid].id, db_record=datasets_by_id[dataset_rowid]
+                )
+                for dataset_rowid in dataset_rowids
+            ],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
@@ -161,7 +148,10 @@ class DatasetLabelMutationMixin:
         }
         return DeleteDatasetLabelsMutationPayload(
             dataset_labels=[
-                to_gql_dataset_label(deleted_dataset_labels_by_id[dataset_label_row_id])
+                DatasetLabel(
+                    id=deleted_dataset_labels_by_id[dataset_label_row_id].id,
+                    db_record=deleted_dataset_labels_by_id[dataset_label_row_id],
+                )
                 for dataset_label_row_id in dataset_label_row_ids
             ]
         )
@@ -170,122 +160,84 @@ class DatasetLabelMutationMixin:
     async def set_dataset_labels(
         self, info: Info[Context, None], input: SetDatasetLabelsInput
     ) -> SetDatasetLabelsMutationPayload:
-        if not input.dataset_ids:
-            raise BadRequest("No datasets provided.")
-        if not input.dataset_label_ids:
-            raise BadRequest("No dataset labels provided.")
-
-        unique_dataset_rowids: set[int] = set()
-        for dataset_gid in input.dataset_ids:
-            try:
-                dataset_rowid = from_global_id_with_expected_type(dataset_gid, Dataset.__name__)
-            except ValueError:
-                raise BadRequest(f"Invalid dataset ID: {dataset_gid}")
-            unique_dataset_rowids.add(dataset_rowid)
-        dataset_rowids = list(unique_dataset_rowids)
+        try:
+            dataset_id = from_global_id_with_expected_type(input.dataset_id, Dataset.__name__)
+        except ValueError:
+            raise BadRequest(f"Invalid dataset ID: {input.dataset_id}")
 
-        unique_dataset_label_rowids: set[int] = set()
+        dataset_label_ids: dict[
+            int, None
+        ] = {}  # use dictionary to de-duplicate while preserving order
         for dataset_label_gid in input.dataset_label_ids:
             try:
-                dataset_label_rowid = from_global_id_with_expected_type(
+                dataset_label_id = from_global_id_with_expected_type(
                     dataset_label_gid, DatasetLabel.__name__
                 )
             except ValueError:
                 raise BadRequest(f"Invalid dataset label ID: {dataset_label_gid}")
-            unique_dataset_label_rowids.add(dataset_label_rowid)
-        dataset_label_rowids = list(unique_dataset_label_rowids)
+            dataset_label_ids[dataset_label_id] = None
 
         async with info.context.db() as session:
-            existing_dataset_ids = (
-                await session.scalars(
-                    select(models.Dataset.id).where(models.Dataset.id.in_(dataset_rowids))
-                )
-            ).all()
-            if len(existing_dataset_ids) != len(dataset_rowids):
-                raise NotFound("One or more datasets not found")
+            dataset = await session.scalar(
+                select(models.Dataset)
+                .where(models.Dataset.id == dataset_id)
+                .options(joinedload(models.Dataset.datasets_dataset_labels))
+            )
 
-            existing_dataset_label_ids = (
+            if not dataset:
+                raise NotFound(f"Dataset with ID {input.dataset_id} not found")
+
+            existing_label_ids = (
                 await session.scalars(
                     select(models.DatasetLabel.id).where(
-                        models.DatasetLabel.id.in_(dataset_label_rowids)
+                        models.DatasetLabel.id.in_(dataset_label_ids.keys())
                     )
                 )
             ).all()
-            if len(existing_dataset_label_ids) != len(dataset_label_rowids):
+            if len(existing_label_ids) != len(dataset_label_ids):
                 raise NotFound("One or more dataset labels not found")
 
-            existing_dataset_label_keys = await session.execute(
-                select(
-                    models.DatasetsDatasetLabel.dataset_id,
-                    models.DatasetsDatasetLabel.dataset_label_id,
-                ).where(
-                    models.DatasetsDatasetLabel.dataset_id.in_(dataset_rowids)
-                    & models.DatasetsDatasetLabel.dataset_label_id.in_(dataset_label_rowids)
-                )
-            )
-            unique_dataset_label_keys = set(existing_dataset_label_keys.all())
+            previously_applied_dataset_label_ids = {
+                dataset_dataset_label.dataset_label_id
+                for dataset_dataset_label in dataset.datasets_dataset_labels
+            }
 
-            datasets_dataset_labels = []
-            for dataset_rowid in dataset_rowids:
-                for dataset_label_rowid in dataset_label_rowids:
-                    if (dataset_rowid, dataset_label_rowid) in unique_dataset_label_keys:
-                        continue
-                    datasets_dataset_labels.append(
-                        models.DatasetsDatasetLabel(
-                            dataset_id=dataset_rowid,
-                            dataset_label_id=dataset_label_rowid,
+            datasets_dataset_labels_to_add = [
+                models.DatasetsDatasetLabel(
+                    dataset_id=dataset_id,
+                    dataset_label_id=dataset_label_id,
+                )
+                for dataset_label_id in dataset_label_ids
+                if dataset_label_id not in previously_applied_dataset_label_ids
+            ]
+            if datasets_dataset_labels_to_add:
+                session.add_all(datasets_dataset_labels_to_add)
+                await session.flush()
+
+            datasets_dataset_labels_to_delete = [
+                dataset_dataset_label
+                for dataset_dataset_label in dataset.datasets_dataset_labels
+                if dataset_dataset_label.dataset_label_id not in dataset_label_ids
+            ]
+            if datasets_dataset_labels_to_delete:
+                await session.execute(
+                    delete(models.DatasetsDatasetLabel).where(
+                        tuple_(
+                            models.DatasetsDatasetLabel.dataset_id,
+                            models.DatasetsDatasetLabel.dataset_label_id,
+                        ).in_(
+                            [
+                                (
+                                    datasets_dataset_labels.dataset_id,
+                                    datasets_dataset_labels.dataset_label_id,
+                                )
+                                for datasets_dataset_labels in datasets_dataset_labels_to_delete
+                            ]
                         )
                     )
-            session.add_all(datasets_dataset_labels)
-
-            if datasets_dataset_labels:
-                try:
-                    await session.commit()
-                except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
-                    raise Conflict("Failed to add dataset labels to datasets.") from e
-
-        return SetDatasetLabelsMutationPayload(
-            query=Query(),
-        )
-
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
-    async def unset_dataset_labels(
-        self, info: Info[Context, None], input: UnsetDatasetLabelsInput
-    ) -> UnsetDatasetLabelsMutationPayload:
-        if not input.dataset_ids:
-            raise BadRequest("No datasets provided.")
-        if not input.dataset_label_ids:
-            raise BadRequest("No dataset labels provided.")
-
-        unique_dataset_rowids: set[int] = set()
-        for dataset_gid in input.dataset_ids:
-            try:
-                dataset_rowid = from_global_id_with_expected_type(dataset_gid, Dataset.__name__)
-            except ValueError:
-                raise BadRequest(f"Invalid dataset ID: {dataset_gid}")
-            unique_dataset_rowids.add(dataset_rowid)
-        dataset_rowids = list(unique_dataset_rowids)
-
-        unique_dataset_label_rowids: set[int] = set()
-        for dataset_label_gid in input.dataset_label_ids:
-            try:
-                dataset_label_rowid = from_global_id_with_expected_type(
-                    dataset_label_gid, DatasetLabel.__name__
-                )
-            except ValueError:
-                raise BadRequest(f"Invalid dataset label ID: {dataset_label_gid}")
-            unique_dataset_label_rowids.add(dataset_label_rowid)
-        dataset_label_rowids = list(unique_dataset_label_rowids)
-
-        async with info.context.db() as session:
-            await session.execute(
-                delete(models.DatasetsDatasetLabel).where(
-                    models.DatasetsDatasetLabel.dataset_id.in_(dataset_rowids)
-                    & models.DatasetsDatasetLabel.dataset_label_id.in_(dataset_label_rowids)
                 )
-            )
-            await session.commit()
 
-        return UnsetDatasetLabelsMutationPayload(
+        return SetDatasetLabelsMutationPayload(
+            dataset=Dataset(id=dataset.id, db_record=dataset),
             query=Query(),
         )

phoenix/server/api/mutations/dataset_mutations.py

@@ -35,7 +35,7 @@ from phoenix.server.api.input_types.PatchDatasetExamplesInput import (
     PatchDatasetExamplesInput,
 )
 from phoenix.server.api.input_types.PatchDatasetInput import PatchDatasetInput
-from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
+from phoenix.server.api.types.Dataset import Dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span
@@ -72,7 +72,7 @@ class DatasetMutationMixin:
             )
             assert dataset is not None
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_dataset(
@@ -101,7 +101,7 @@ class DatasetMutationMixin:
             )
             assert dataset is not None
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def add_spans_to_dataset(
@@ -221,7 +221,7 @@ class DatasetMutationMixin:
                 ],
             )
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def add_examples_to_dataset(
@@ -348,7 +348,7 @@ class DatasetMutationMixin:
                 dataset_example_revisions,
             )
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
     async def delete_dataset(
@@ -379,7 +379,7 @@ class DatasetMutationMixin:
             return_exceptions=True,
         )
         info.context.event_queue.put(DatasetDeleteEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_dataset_examples(
@@ -472,7 +472,7 @@ class DatasetMutationMixin:
                 ],
             )
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def delete_dataset_examples(
@@ -556,7 +556,7 @@ class DatasetMutationMixin:
                 ],
             )
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
-        return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
+        return DatasetMutationPayload(dataset=Dataset(id=dataset.id, db_record=dataset))
 
 
 def _span_attribute(semconv: str) -> Any:

phoenix/server/api/mutations/dataset_split_mutations.py

@@ -15,8 +15,8 @@ from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.helpers.playground_users import get_user
 from phoenix.server.api.queries import Query
-from phoenix.server.api.types.DatasetExample import DatasetExample, to_gql_dataset_example
-from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
+from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetSplit import DatasetSplit
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 
 
@@ -116,7 +116,8 @@ class DatasetSplitMutationMixin:
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A dataset split named '{input.name}' already exists.")
         return DatasetSplitMutationPayload(
-            dataset_split=to_gql_dataset_split(dataset_split_orm), query=Query()
+            dataset_split=DatasetSplit(id=dataset_split_orm.id, db_record=dataset_split_orm),
+            query=Query(),
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
@@ -141,7 +142,7 @@ class DatasetSplitMutationMixin:
             if isinstance(input.metadata, dict):
                 dataset_split_orm.metadata_ = input.metadata
 
-            gql_dataset_split = to_gql_dataset_split(dataset_split_orm)
+            gql_dataset_split = DatasetSplit(id=dataset_split_orm.id, db_record=dataset_split_orm)
             try:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
@@ -185,7 +186,10 @@ class DatasetSplitMutationMixin:
 
         return DeleteDatasetSplitsMutationPayload(
             dataset_splits=[
-                to_gql_dataset_split(deleted_splits_by_id[dataset_split_rowid])
+                DatasetSplit(
+                    id=deleted_splits_by_id[dataset_split_rowid].id,
+                    db_record=deleted_splits_by_id[dataset_split_rowid],
+                )
                 for dataset_split_rowid in dataset_split_rowids
             ],
             query=Query(),
@@ -281,7 +285,7 @@ class DatasetSplitMutationMixin:
             ).all()
         return AddDatasetExamplesToDatasetSplitsMutationPayload(
             query=Query(),
-            examples=[to_gql_dataset_example(example) for example in examples],
+            examples=[DatasetExample(id=example.id, db_record=example) for example in examples],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
@@ -342,7 +346,7 @@ class DatasetSplitMutationMixin:
 
         return RemoveDatasetExamplesFromDatasetSplitsMutationPayload(
             query=Query(),
-            examples=[to_gql_dataset_example(example) for example in examples],
+            examples=[DatasetExample(id=example.id, db_record=example) for example in examples],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
@@ -410,9 +414,9 @@ class DatasetSplitMutationMixin:
             ).all()
 
         return DatasetSplitMutationPayloadWithExamples(
-            dataset_split=to_gql_dataset_split(dataset_split_orm),
+            dataset_split=DatasetSplit(id=dataset_split_orm.id, db_record=dataset_split_orm),
             query=Query(),
-            examples=[to_gql_dataset_example(example) for example in examples],
+            examples=[DatasetExample(id=example.id, db_record=example) for example in examples],
         )
 
 

phoenix/server/api/mutations/model_mutations.py

@@ -16,7 +16,7 @@ from phoenix.server.api.auth import IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.queries import Query
-from phoenix.server.api.types.GenerativeModel import GenerativeModel, to_gql_generative_model
+from phoenix.server.api.types.GenerativeModel import GenerativeModel
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.TokenPrice import TokenKind
 
@@ -110,7 +110,7 @@ class ModelMutationMixin:
                 raise Conflict(f"Model with name '{input.name}' already exists")
 
         return CreateModelMutationPayload(
-            model=to_gql_generative_model(model),
+            model=GenerativeModel(id=model.id, db_record=model),
             query=Query(),
         )
 
@@ -163,7 +163,7 @@ class ModelMutationMixin:
             await session.refresh(model)
 
         return UpdateModelMutationPayload(
-            model=to_gql_generative_model(model),
+            model=GenerativeModel(id=model.id, db_record=model),
             query=Query(),
         )
 
@@ -192,7 +192,7 @@ class ModelMutationMixin:
                 await session.rollback()
                 raise BadRequest("Cannot delete built-in model")
         return DeleteModelMutationPayload(
-            model=to_gql_generative_model(model),
+            model=GenerativeModel(id=model.id, db_record=model),
             query=Query(),
         )
 

phoenix/server/api/mutations/project_session_annotations_mutations.py

@@ -19,10 +19,7 @@ from phoenix.server.api.input_types.UpdateAnnotationInput import UpdateAnnotatio
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.ProjectSessionAnnotation import (
-    ProjectSessionAnnotation,
-    to_gql_project_session_annotation,
-)
+from phoenix.server.api.types.ProjectSessionAnnotation import ProjectSessionAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import (
     ProjectSessionAnnotationDeleteEvent,
@@ -81,7 +78,7 @@ class ProjectSessionAnnotationMutationMixin:
         info.context.event_queue.put(ProjectSessionAnnotationInsertEvent((anno.id,)))
 
         return ProjectSessionAnnotationMutationPayload(
-            project_session_annotation=to_gql_project_session_annotation(anno),
+            project_session_annotation=ProjectSessionAnnotation(id=anno.id, db_record=anno),
             query=Query(),
         )
 
@@ -122,7 +119,7 @@ class ProjectSessionAnnotationMutationMixin:
 
         info.context.event_queue.put(ProjectSessionAnnotationInsertEvent((anno.id,)))
         return ProjectSessionAnnotationMutationPayload(
-            project_session_annotation=to_gql_project_session_annotation(anno),
+            project_session_annotation=ProjectSessionAnnotation(id=anno.id, db_record=anno),
             query=Query(),
         )
 
@@ -154,7 +151,7 @@ class ProjectSessionAnnotationMutationMixin:
 
             await session.delete(anno)
 
-        deleted_gql_annotation = to_gql_project_session_annotation(anno)
+        deleted_gql_annotation = ProjectSessionAnnotation(id=anno.id, db_record=anno)
         info.context.event_queue.put(ProjectSessionAnnotationDeleteEvent((id_,)))
         return ProjectSessionAnnotationMutationPayload(
             project_session_annotation=deleted_gql_annotation, query=Query()

phoenix/server/api/mutations/prompt_label_mutations.py

@@ -16,7 +16,7 @@ from phoenix.server.api.exceptions import Conflict, NotFound
 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.Prompt import Prompt
-from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
+from phoenix.server.api.types.PromptLabel import PromptLabel
 
 
 @strawberry.input
@@ -85,7 +85,7 @@ class PromptLabelMutationMixin:
                 raise Conflict(f"A prompt label named '{input.name}' already exists.")
 
             return PromptLabelMutationPayload(
-                prompt_labels=[to_gql_prompt_label(label_orm)],
+                prompt_labels=[PromptLabel(id=label_orm.id, db_record=label_orm)],
                 query=Query(),
             )
 
@@ -113,7 +113,7 @@ class PromptLabelMutationMixin:
                 raise Conflict("Error patching PromptLabel. Possibly a name conflict?")
 
             return PromptLabelMutationPayload(
-                prompt_labels=[to_gql_prompt_label(label_orm)],
+                prompt_labels=[PromptLabel(id=label_orm.id, db_record=label_orm)],
                 query=Query(),
             )