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

phoenix/server/api/mutations/dataset_label_mutations.py

@@ -16,8 +16,8 @@ from phoenix.server.api.auth import IsLocked, 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.Dataset import Dataset, to_gql_dataset
-from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
+from phoenix.server.api.types.Dataset import Dataset
+from phoenix.server.api.types.DatasetLabel import DatasetLabel
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 
 
@@ -111,9 +111,12 @@ class DatasetLabelMutationMixin:
                 await session.commit()
 
         return CreateDatasetLabelMutationPayload(
-            dataset_label=to_gql_dataset_label(dataset_label_orm),
+            dataset_label=DatasetLabel(id=dataset_label_orm.id, db_record=dataset_label_orm),
             datasets=[
-                to_gql_dataset(datasets_by_id[dataset_rowid]) for dataset_rowid in dataset_rowids
+                Dataset(
+                    id=datasets_by_id[dataset_rowid].id, db_record=datasets_by_id[dataset_rowid]
+                )
+                for dataset_rowid in dataset_rowids
             ],
         )
 
@@ -145,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
             ]
         )
@@ -232,6 +238,6 @@ class DatasetLabelMutationMixin:
                 )
 
         return SetDatasetLabelsMutationPayload(
-            dataset=to_gql_dataset(dataset),
+            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(),
             )
 

phoenix/server/api/mutations/prompt_mutations.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union, cast
+from typing import Optional
 
 import strawberry
 from fastapi import Request
@@ -12,18 +12,11 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.db.types.identifier import Identifier as IdentifierModel
-from phoenix.db.types.model_provider import ModelProvider
 from phoenix.server.api.auth import IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
-from phoenix.server.api.helpers.prompts.models import (
-    normalize_response_format,
-    normalize_tools,
-    validate_invocation_parameters,
-)
 from phoenix.server.api.input_types.PromptVersionInput import (
     ChatPromptVersionInput,
-    to_pydantic_prompt_chat_template_v1,
 )
 from phoenix.server.api.mutations.prompt_version_tag_mutations import (
     SetPromptVersionTagInput,
@@ -32,7 +25,7 @@ from phoenix.server.api.mutations.prompt_version_tag_mutations import (
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.Identifier import Identifier
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
+from phoenix.server.api.types.Prompt import Prompt
 from phoenix.server.bearer_auth import PhoenixUser
 
 
@@ -84,63 +77,23 @@ class PromptMutationMixin:
         if "user" in request.scope:
             assert isinstance(user := request.user, PhoenixUser)
             user_id = int(user.identity)
-
-        input_prompt_version = input.prompt_version
-        tool_definitions = [tool.definition for tool in input_prompt_version.tools]
-        tool_choice = cast(
-            Optional[Union[str, dict[str, Any]]],
-            cast(dict[str, Any], input.prompt_version.invocation_parameters).pop(
-                "tool_choice", None
-            ),
-        )
-        model_provider = ModelProvider(input_prompt_version.model_provider)
         try:
-            tools = (
-                normalize_tools(tool_definitions, model_provider, tool_choice)
-                if tool_definitions
-                else None
-            )
-            template = to_pydantic_prompt_chat_template_v1(input_prompt_version.template)
-            response_format = (
-                normalize_response_format(
-                    input_prompt_version.response_format.definition,
-                    model_provider,
-                )
-                if input_prompt_version.response_format
-                else None
-            )
-            invocation_parameters = validate_invocation_parameters(
-                input_prompt_version.invocation_parameters,
-                model_provider,
-            )
+            prompt_version = input.prompt_version.to_orm_prompt_version(user_id)
         except ValidationError as error:
             raise BadRequest(str(error))
-
+        name = IdentifierModel.model_validate(str(input.name))
+        prompt = models.Prompt(
+            name=name,
+            description=input.description,
+            prompt_versions=[prompt_version],
+        )
         async with info.context.db() as session:
-            prompt_version = models.PromptVersion(
-                description=input_prompt_version.description,
-                user_id=user_id,
-                template_type="CHAT",
-                template_format=input_prompt_version.template_format,
-                template=template,
-                invocation_parameters=invocation_parameters,
-                tools=tools,
-                response_format=response_format,
-                model_provider=input_prompt_version.model_provider,
-                model_name=input_prompt_version.model_name,
-            )
-            name = IdentifierModel.model_validate(str(input.name))
-            prompt = models.Prompt(
-                name=name,
-                description=input.description,
-                prompt_versions=[prompt_version],
-            )
             session.add(prompt)
             try:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A prompt named '{input.name}' already exists")
-        return to_gql_prompt_from_orm(prompt)
+        return Prompt(id=prompt.id, db_record=prompt)
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def create_chat_prompt_version(
@@ -153,72 +106,26 @@ class PromptMutationMixin:
         if "user" in request.scope:
             assert isinstance(user := request.user, PhoenixUser)
             user_id = int(user.identity)
-
-        input_prompt_version = input.prompt_version
-        tool_definitions = [tool.definition for tool in input.prompt_version.tools]
-        tool_choice = cast(
-            Optional[Union[str, dict[str, Any]]],
-            cast(dict[str, Any], input.prompt_version.invocation_parameters).pop(
-                "tool_choice", None
-            ),
-        )
-        model_provider = ModelProvider(input_prompt_version.model_provider)
         try:
-            tools = (
-                normalize_tools(tool_definitions, model_provider, tool_choice)
-                if tool_definitions
-                else None
-            )
-            template = to_pydantic_prompt_chat_template_v1(input_prompt_version.template)
-            response_format = (
-                normalize_response_format(
-                    input_prompt_version.response_format.definition,
-                    model_provider,
-                )
-                if input_prompt_version.response_format
-                else None
-            )
-            invocation_parameters = validate_invocation_parameters(
-                input_prompt_version.invocation_parameters,
-                model_provider,
-            )
+            prompt_version = input.prompt_version.to_orm_prompt_version(user_id)
         except ValidationError as error:
             raise BadRequest(str(error))
-
         prompt_id = from_global_id_with_expected_type(
             global_id=input.prompt_id, expected_type_name=Prompt.__name__
         )
+        prompt_version.prompt_id = prompt_id
         async with info.context.db() as session:
-            prompt = await session.get(models.Prompt, prompt_id)
-            if not prompt:
-                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
-
-            prompt_version = models.PromptVersion(
-                prompt_id=prompt_id,
-                description=input.prompt_version.description,
-                user_id=user_id,
-                template_type="CHAT",
-                template_format=input.prompt_version.template_format,
-                template=template,
-                invocation_parameters=invocation_parameters,
-                tools=tools,
-                response_format=response_format,
-                model_provider=input.prompt_version.model_provider,
-                model_name=input.prompt_version.model_name,
-            )
             session.add(prompt_version)
-
-        # ensure prompt_version is flushed to the database before creating tags against the
-        # prompt_version id
-        await session.flush()
-
-        if input.tags:
-            for tag in input.tags:
-                await upsert_prompt_version_tag(
-                    session, prompt_id, prompt_version.id, tag.name, tag.description
-                )
-
-        return to_gql_prompt_from_orm(prompt)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
+            if input.tags:
+                for tag in input.tags:
+                    await upsert_prompt_version_tag(
+                        session, prompt_id, prompt_version.id, tag.name, tag.description
+                    )
+        return Prompt(id=prompt_id)
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
     async def delete_prompt(
@@ -288,7 +195,7 @@ class PromptMutationMixin:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict(f"A prompt named '{input.name}' already exists")
-        return to_gql_prompt_from_orm(new_prompt)
+        return Prompt(id=new_prompt.id, db_record=new_prompt)
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_prompt(self, info: Info[Context, None], input: PatchPromptInput) -> Prompt:
@@ -310,4 +217,4 @@ class PromptMutationMixin:
             if prompt is None:
                 raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
 
-        return to_gql_prompt_from_orm(prompt)
+        return Prompt(id=prompt.id, db_record=prompt)

phoenix/server/api/mutations/prompt_version_tag_mutations.py

@@ -16,9 +16,9 @@ from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.Identifier import Identifier
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
+from phoenix.server.api.types.Prompt import Prompt
 from phoenix.server.api.types.PromptVersion import PromptVersion
-from phoenix.server.api.types.PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+from phoenix.server.api.types.PromptVersionTag import PromptVersionTag
 
 
 @strawberry.input
@@ -75,7 +75,9 @@ class PromptVersionTagMutationMixin:
             await session.delete(prompt_version_tag)
             await session.commit()
             return PromptVersionTagMutationPayload(
-                prompt_version_tag=None, query=Query(), prompt=to_gql_prompt_from_orm(prompt)
+                prompt_version_tag=None,
+                query=Query(),
+                prompt=Prompt(id=prompt.id, db_record=prompt),
             )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
@@ -111,9 +113,10 @@ class PromptVersionTagMutationMixin:
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
                 raise Conflict("Failed to update PromptVersionTag.")
 
-            version_tag = to_gql_prompt_version_tag(updated_tag)
             return PromptVersionTagMutationPayload(
-                prompt_version_tag=version_tag, prompt=to_gql_prompt_from_orm(prompt), query=Query()
+                prompt_version_tag=PromptVersionTag(id=updated_tag.id, db_record=updated_tag),
+                prompt=Prompt(id=prompt.id, db_record=prompt),
+                query=Query(),
             )
 
 

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -21,7 +21,7 @@ from phoenix.server.api.queries import Query
 from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 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
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanAnnotationDeleteEvent, SpanAnnotationInsertEvent
 
@@ -138,7 +138,7 @@ class SpanAnnotationMutationMixin:
 
             # Convert the fully loaded annotations to GQL types
             returned_annotations = [
-                to_gql_span_annotation(anno) for anno in ordered_final_annotations
+                SpanAnnotation(id=anno.id, db_record=anno) for anno in ordered_final_annotations
             ]
 
             await session.commit()
@@ -184,7 +184,9 @@ class SpanAnnotationMutationMixin:
             processed_annotation = result.one()
 
             info.context.event_queue.put(SpanAnnotationInsertEvent((processed_annotation.id,)))
-            returned_annotation = to_gql_span_annotation(processed_annotation)
+            returned_annotation = SpanAnnotation(
+                id=processed_annotation.id, db_record=processed_annotation
+            )
             await session.commit()
         return SpanAnnotationMutationPayload(
             span_annotations=[returned_annotation],
@@ -256,7 +258,7 @@ class SpanAnnotationMutationMixin:
                 session.add(span_annotation)
 
             patched_annotations = [
-                to_gql_span_annotation(span_annotation)
+                SpanAnnotation(id=span_annotation.id, db_record=span_annotation)
                 for span_annotation in span_annotations_by_id.values()
             ]
 
@@ -320,7 +322,10 @@ class SpanAnnotationMutationMixin:
                 )
 
         deleted_annotations_gql = [
-            to_gql_span_annotation(deleted_annotations_by_id[id]) for id in span_annotation_ids
+            SpanAnnotation(
+                id=deleted_annotations_by_id[id].id, db_record=deleted_annotations_by_id[id]
+            )
+            for id in span_annotation_ids
         ]
         info.context.event_queue.put(
             SpanAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -16,7 +16,7 @@ from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationI
 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.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationDeleteEvent, TraceAnnotationInsertEvent
 
@@ -111,7 +111,9 @@ class TraceAnnotationMutationMixin:
             info.context.event_queue.put(TraceAnnotationInsertEvent(inserted_annotation_ids))
 
         returned_annotations = [
-            to_gql_trace_annotation(processed_annotations_map[i])
+            TraceAnnotation(
+                id=processed_annotations_map[i].id, db_record=processed_annotations_map[i]
+            )
             for i in sorted(processed_annotations_map.keys())
         ]
 
@@ -186,7 +188,7 @@ class TraceAnnotationMutationMixin:
             await session.commit()
 
         patched_annotations = [
-            to_gql_trace_annotation(trace_annotation)
+            TraceAnnotation(id=trace_annotation.id, db_record=trace_annotation)
             for trace_annotation in trace_annotations_by_id.values()
         ]
         info.context.event_queue.put(TraceAnnotationInsertEvent(tuple(patch_by_id.keys())))
@@ -245,7 +247,10 @@ class TraceAnnotationMutationMixin:
                 )
 
         deleted_gql_annotations = [
-            to_gql_trace_annotation(deleted_annotations_by_id[id]) for id in trace_annotation_ids
+            TraceAnnotation(
+                id=deleted_annotations_by_id[id].id, db_record=deleted_annotations_by_id[id]
+            )
+            for id in trace_annotation_ids
         ]
         info.context.event_queue.put(
             TraceAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))

phoenix/server/api/mutations/user_mutations.py

@@ -33,7 +33,7 @@ from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound, Unauth
 from phoenix.server.api.input_types.UserRoleInput import UserRoleInput
 from phoenix.server.api.types.AuthMethod import AuthMethod
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.User import User, to_gql_user
+from phoenix.server.api.types.User import User
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.types import AccessTokenId, ApiKeyId, PasswordResetTokenId, RefreshTokenId
 
@@ -155,7 +155,7 @@ class UserMutationMixin:
             except Exception as error:
                 # Log the error but do not raise it
                 logger.error(f"Failed to send welcome email: {error}")
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin])  # type: ignore
     async def patch_user(
@@ -204,7 +204,7 @@ class UserMutationMixin:
         assert user
         if should_log_out:
             await info.context.log_out(user.id)
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_viewer(
@@ -246,7 +246,7 @@ class UserMutationMixin:
             response = info.context.get_response()
             response.delete_cookie(PHOENIX_REFRESH_TOKEN_COOKIE_NAME)
             response.delete_cookie(PHOENIX_ACCESS_TOKEN_COOKIE_NAME)
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin, IsLocked])  # type: ignore
     async def delete_users(