arize-phoenix 7.12.3__py3-none-any.whl → 8.0.0__py3-none-any.whl

phoenix/server/api/mutations/prompt_mutations.py

@@ -0,0 +1,312 @@
+from typing import Any, Optional, Union, cast
+
+import strawberry
+from fastapi import Request
+from pydantic import ValidationError
+from sqlalchemy import delete, select, update
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.orm import joinedload
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry.relay.types import GlobalID
+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.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,
+    upsert_prompt_version_tag,
+)
+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.bearer_auth import PhoenixUser
+
+
+@strawberry.input
+class CreateChatPromptInput:
+    name: Identifier
+    description: Optional[str] = None
+    prompt_version: ChatPromptVersionInput
+
+
+@strawberry.input
+class CreateChatPromptVersionInput:
+    prompt_id: GlobalID
+    prompt_version: ChatPromptVersionInput
+    tags: Optional[list[SetPromptVersionTagInput]] = None
+
+
+@strawberry.input
+class DeletePromptInput:
+    prompt_id: GlobalID
+
+
+@strawberry.input
+class ClonePromptInput:
+    name: Identifier
+    description: Optional[str] = None
+    prompt_id: GlobalID
+
+
+@strawberry.input
+class PatchPromptInput:
+    prompt_id: GlobalID
+    description: str
+
+
+@strawberry.type
+class DeletePromptMutationPayload:
+    query: Query
+
+
+@strawberry.type
+class PromptMutationMixin:
+    @strawberry.mutation
+    async def create_chat_prompt(
+        self, info: Info[Context, None], input: CreateChatPromptInput
+    ) -> Prompt:
+        user_id: Optional[int] = None
+        assert isinstance(request := info.context.request, Request)
+        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,
+            )
+        except ValidationError as error:
+            raise BadRequest(str(error))
+
+        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)
+
+    @strawberry.mutation
+    async def create_chat_prompt_version(
+        self,
+        info: Info[Context, None],
+        input: CreateChatPromptVersionInput,
+    ) -> Prompt:
+        user_id: Optional[int] = None
+        assert isinstance(request := info.context.request, Request)
+        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,
+            )
+        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__
+        )
+        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)
+
+    @strawberry.mutation
+    async def delete_prompt(
+        self, info: Info[Context, None], input: DeletePromptInput
+    ) -> DeletePromptMutationPayload:
+        prompt_id = from_global_id_with_expected_type(
+            global_id=input.prompt_id, expected_type_name=Prompt.__name__
+        )
+        async with info.context.db() as session:
+            stmt = delete(models.Prompt).where(models.Prompt.id == prompt_id)
+            result = await session.execute(stmt)
+
+            if result.rowcount == 0:
+                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
+
+            await session.commit()
+        return DeletePromptMutationPayload(query=Query())
+
+    @strawberry.mutation
+    async def clone_prompt(self, info: Info[Context, None], input: ClonePromptInput) -> Prompt:
+        prompt_id = from_global_id_with_expected_type(
+            global_id=input.prompt_id, expected_type_name=Prompt.__name__
+        )
+        async with info.context.db() as session:
+            # Load prompt with all versions
+            stmt = (
+                select(models.Prompt)
+                .options(joinedload(models.Prompt.prompt_versions))
+                .where(models.Prompt.id == prompt_id)
+            )
+            prompt = await session.scalar(stmt)
+
+            if not prompt:
+                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
+
+            # Create new prompt
+            name = IdentifierModel.model_validate(str(input.name))
+            new_prompt = models.Prompt(
+                name=name,
+                description=input.description,
+                source_prompt_id=prompt_id,
+            )
+
+            # Create copies of all versions
+            new_versions = [
+                models.PromptVersion(
+                    prompt_id=new_prompt.id,
+                    user_id=version.user_id,
+                    description=version.description,
+                    template_type=version.template_type,
+                    template_format=version.template_format,
+                    template=version.template,
+                    invocation_parameters=version.invocation_parameters,
+                    tools=version.tools,
+                    response_format=version.response_format,
+                    model_provider=version.model_provider,
+                    model_name=version.model_name,
+                )
+                for version in prompt.prompt_versions
+            ]
+            # Add all version copies to the new prompt
+            new_prompt.prompt_versions = new_versions
+
+            session.add(new_prompt)
+
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"A prompt named '{input.name}' already exists")
+        return to_gql_prompt_from_orm(new_prompt)
+
+    @strawberry.mutation
+    async def patch_prompt(self, info: Info[Context, None], input: PatchPromptInput) -> Prompt:
+        prompt_id = from_global_id_with_expected_type(
+            global_id=input.prompt_id, expected_type_name=Prompt.__name__
+        )
+
+        async with info.context.db() as session:
+            stmt = (
+                update(models.Prompt)
+                .where(models.Prompt.id == prompt_id)
+                .values(description=input.description)
+                .returning(models.Prompt)
+            )
+
+            result = await session.execute(stmt)
+            prompt = result.scalar_one_or_none()
+
+            if prompt is None:
+                raise NotFound(f"Prompt with ID '{input.prompt_id}' not found")
+
+        return to_gql_prompt_from_orm(prompt)

phoenix/server/api/mutations/prompt_version_tag_mutations.py

@@ -0,0 +1,148 @@
+from typing import Optional
+
+import strawberry
+from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry.relay import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.db.types.identifier import Identifier as IdentifierModel
+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.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.PromptVersion import PromptVersion
+from phoenix.server.api.types.PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+
+
+@strawberry.input
+class DeletePromptVersionTagInput:
+    prompt_version_tag_id: GlobalID
+
+
+@strawberry.input
+class SetPromptVersionTagInput:
+    prompt_version_id: GlobalID
+    name: Identifier
+    description: Optional[str] = None
+
+
+@strawberry.type
+class PromptVersionTagMutationPayload:
+    prompt_version_tag: Optional[PromptVersionTag]
+    prompt: Prompt
+    query: Query
+
+
+@strawberry.type
+class PromptVersionTagMutationMixin:
+    @strawberry.mutation
+    async def delete_prompt_version_tag(
+        self, info: Info[Context, None], input: DeletePromptVersionTagInput
+    ) -> PromptVersionTagMutationPayload:
+        async with info.context.db() as session:
+            prompt_version_tag_id = from_global_id_with_expected_type(
+                input.prompt_version_tag_id, PromptVersionTag.__name__
+            )
+            stmt = (
+                select(models.PromptVersionTag, models.Prompt)
+                .join(
+                    models.PromptVersion,
+                    models.PromptVersion.id == models.PromptVersionTag.prompt_version_id,
+                )
+                .join(models.Prompt, models.Prompt.id == models.PromptVersion.prompt_id)
+                .where(models.PromptVersionTag.id == prompt_version_tag_id)
+            )
+            result = await session.execute(stmt)
+            if results := result.one_or_none():
+                prompt_version_tag, prompt = results
+
+            if not prompt_version_tag:
+                raise NotFound(f"PromptVersionTag with ID {input.prompt_version_tag_id} not found")
+
+            if not prompt:
+                raise BadRequest(
+                    f"PromptVersionTag with ID {input.prompt_version_tag_id} "
+                    "does not belong to a prompt"
+                )
+
+            await session.delete(prompt_version_tag)
+            await session.commit()
+            return PromptVersionTagMutationPayload(
+                prompt_version_tag=None, query=Query(), prompt=to_gql_prompt_from_orm(prompt)
+            )
+
+    @strawberry.mutation
+    async def set_prompt_version_tag(
+        self, info: Info[Context, None], input: SetPromptVersionTagInput
+    ) -> PromptVersionTagMutationPayload:
+        async with info.context.db() as session:
+            prompt_version_id = from_global_id_with_expected_type(
+                input.prompt_version_id, PromptVersion.__name__
+            )
+            prompt_version = await session.scalar(
+                select(models.PromptVersion).where(models.PromptVersion.id == prompt_version_id)
+            )
+            if not prompt_version:
+                raise BadRequest(f"PromptVersion with ID {input.prompt_version_id} not found.")
+
+            prompt_id = prompt_version.prompt_id
+            prompt = await session.scalar(
+                select(models.Prompt).where(models.Prompt.id == prompt_id)
+            )
+            if not prompt:
+                raise BadRequest("All prompt version tags must belong to a prompt")
+
+            updated_tag = await upsert_prompt_version_tag(
+                session, prompt_id, prompt_version_id, input.name, input.description
+            )
+
+            if not updated_tag:
+                raise BadRequest("Failed to create or update PromptVersionTag.")
+
+            try:
+                await session.commit()
+            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()
+            )
+
+
+async def upsert_prompt_version_tag(
+    session: AsyncSession,
+    prompt_id: int,
+    prompt_version_id: int,
+    name_str: str,
+    description: Optional[str] = None,
+) -> models.PromptVersionTag:
+    name = IdentifierModel.model_validate(name_str)
+
+    existing_tag = await session.scalar(
+        select(models.PromptVersionTag).where(
+            models.PromptVersionTag.prompt_id == prompt_id,
+            models.PromptVersionTag.name == name,
+        )
+    )
+
+    if existing_tag:
+        existing_tag.prompt_version_id = prompt_version_id
+        if description is not None:
+            existing_tag.description = description
+        return existing_tag
+    else:
+        new_tag = models.PromptVersionTag(
+            name=name,
+            description=description,
+            prompt_id=prompt_id,
+            prompt_version_id=prompt_version_id,
+        )
+        session.add(new_tag)
+        return new_tag

phoenix/server/api/mutations/user_mutations.py

@@ -1,12 +1,13 @@
 import secrets
 from contextlib import AsyncExitStack
 from datetime import datetime, timezone
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 
 import strawberry
 from sqlalchemy import Boolean, Select, and_, case, cast, delete, distinct, func, select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
 from sqlalchemy.orm import joinedload
-from sqlean.dbapi2 import IntegrityError  # type: ignore[import-untyped]
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 from strawberry import UNSET
 from strawberry.relay import GlobalID
 from strawberry.types import Info
@@ -108,7 +109,7 @@ class UserMutationMixin:
             session.add(user)
             try:
                 await session.flush()
-            except IntegrityError as error:
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError) as error:
                 raise Conflict(_user_operation_error_message(error))
         return UserMutationPayload(user=to_gql_user(user))
 
@@ -148,7 +149,7 @@ class UserMutationMixin:
             assert user in session.dirty
             try:
                 await session.flush()
-            except IntegrityError as error:
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError) as error:
                 raise Conflict(_user_operation_error_message(error, "modify"))
         assert user
         if input.new_password:
@@ -186,7 +187,7 @@ class UserMutationMixin:
             user.updated_at = datetime.now(timezone.utc)
             try:
                 await session.flush()
-            except IntegrityError as error:
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError) as error:
                 raise Conflict(_user_operation_error_message(error, "modify"))
         assert user
         if input.new_password:
@@ -313,7 +314,7 @@ def _select_user_by_id(user_id: int) -> Select[tuple[models.User]]:
 
 
 def _user_operation_error_message(
-    error: IntegrityError,
+    error: Union[PostgreSQLIntegrityError, SQLiteIntegrityError],
     operation: Literal["create", "modify"] = "create",
 ) -> str:
     """

phoenix/server/api/openapi/schema.py

@@ -13,4 +13,5 @@ def get_openapi_schema() -> dict[str, Any]:
         openapi_version="3.1.0",
         description="Schema for Arize-Phoenix REST API",
         routes=v1_router.routes,
+        separate_input_output_schemas=False,
     )

phoenix/server/api/queries.py

@@ -14,22 +14,12 @@ from strawberry.types import Info
 from typing_extensions import Annotated, TypeAlias
 
 from phoenix.db import enums, models
-from phoenix.db.models import (
-    DatasetExample as OrmExample,
-)
-from phoenix.db.models import (
-    DatasetExampleRevision as OrmRevision,
-)
-from phoenix.db.models import (
-    DatasetVersion as OrmVersion,
-)
-from phoenix.db.models import (
-    Experiment as OrmExperiment,
-)
+from phoenix.db.models import DatasetExample as OrmExample
+from phoenix.db.models import DatasetExampleRevision as OrmRevision
+from phoenix.db.models import DatasetVersion as OrmVersion
+from phoenix.db.models import Experiment as OrmExperiment
 from phoenix.db.models import ExperimentRun as OrmExperimentRun
-from phoenix.db.models import (
-    Trace as OrmTrace,
-)
+from phoenix.db.models import Trace as OrmTrace
 from phoenix.pointcloud.clustering import Hdbscan
 from phoenix.server.api.auth import MSG_ADMIN_ONLY, IsAdmin
 from phoenix.server.api.context import Context
@@ -43,14 +33,9 @@ from phoenix.server.api.helpers.experiment_run_filters import (
 from phoenix.server.api.helpers.playground_clients import initialize_playground_clients
 from phoenix.server.api.helpers.playground_registry import PLAYGROUND_CLIENT_REGISTRY
 from phoenix.server.api.input_types.ClusterInput import ClusterInput
-from phoenix.server.api.input_types.Coordinates import (
-    InputCoordinate2D,
-    InputCoordinate3D,
-)
+from phoenix.server.api.input_types.Coordinates import InputCoordinate2D, InputCoordinate3D
 from phoenix.server.api.input_types.DatasetSort import DatasetSort
-from phoenix.server.api.input_types.InvocationParameters import (
-    InvocationParameter,
-)
+from phoenix.server.api.input_types.InvocationParameters import InvocationParameter
 from phoenix.server.api.subscriptions import PLAYGROUND_PROJECT_NAME
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
 from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
@@ -68,20 +53,16 @@ from phoenix.server.api.types.ExperimentComparison import ExperimentComparison,
 from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
 from phoenix.server.api.types.Functionality import Functionality
 from phoenix.server.api.types.GenerativeModel import GenerativeModel
-from phoenix.server.api.types.GenerativeProvider import (
-    GenerativeProvider,
-    GenerativeProviderKey,
-)
+from phoenix.server.api.types.GenerativeProvider import GenerativeProvider, GenerativeProviderKey
 from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
 from phoenix.server.api.types.Model import Model
 from phoenix.server.api.types.node import from_global_id, from_global_id_with_expected_type
-from phoenix.server.api.types.pagination import (
-    ConnectionArgs,
-    CursorString,
-    connection_from_list,
-)
+from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
 from phoenix.server.api.types.Project import Project
 from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
+from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
+from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
+from phoenix.server.api.types.PromptVersion import PromptVersion, to_gql_prompt_version
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
@@ -587,6 +568,31 @@ class Query:
                 ):
                     raise NotFound(f"Unknown user: {id}")
             return to_gql_project_session(project_session)
+        elif type_name == Prompt.__name__:
+            async with info.context.db() as session:
+                if orm_prompt := await session.scalar(
+                    select(models.Prompt).where(models.Prompt.id == node_id)
+                ):
+                    return to_gql_prompt_from_orm(orm_prompt)
+                else:
+                    raise NotFound(f"Unknown prompt: {id}")
+        elif type_name == PromptVersion.__name__:
+            async with info.context.db() as session:
+                if orm_prompt_version := await session.scalar(
+                    select(models.PromptVersion).where(models.PromptVersion.id == node_id)
+                ):
+                    return to_gql_prompt_version(orm_prompt_version)
+                else:
+                    raise NotFound(f"Unknown prompt version: {id}")
+        elif type_name == PromptLabel.__name__:
+            async with info.context.db() as session:
+                if not (
+                    prompt_label := await session.scalar(
+                        select(models.PromptLabel).where(models.PromptLabel.id == node_id)
+                    )
+                ):
+                    raise NotFound(f"Unknown prompt label: {id}")
+            return to_gql_prompt_label(prompt_label)
         raise NotFound(f"Unknown node type: {type_name}")
 
     @strawberry.field
@@ -609,6 +615,53 @@ class Query:
                 return None
         return to_gql_user(user)
 
+    @strawberry.field
+    async def prompts(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[Prompt]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        stmt = select(models.Prompt)
+        async with info.context.db() as session:
+            orm_prompts = await session.stream_scalars(stmt)
+            data = [to_gql_prompt_from_orm(orm_prompt) async for orm_prompt in orm_prompts]
+            return connection_from_list(
+                data=data,
+                args=args,
+            )
+
+    @strawberry.field
+    async def prompt_labels(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[PromptLabel]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        async with info.context.db() as session:
+            prompt_labels = await session.stream_scalars(select(models.PromptLabel))
+            data = [to_gql_prompt_label(prompt_label) async for prompt_label in prompt_labels]
+            return connection_from_list(
+                data=data,
+                args=args,
+            )
+
     @strawberry.field
     def clusters(
         self,

phoenix/server/api/routers/oauth2.py

@@ -11,9 +11,10 @@ from authlib.jose import jwt
 from authlib.jose.errors import JoseError
 from fastapi import APIRouter, Cookie, Depends, Path, Query, Request
 from sqlalchemy import Boolean, and_, case, cast, func, insert, or_, select, update
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
-from sqlean.dbapi2 import IntegrityError  # type: ignore[import-untyped]
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 from starlette.datastructures import URL, URLPath
 from starlette.responses import RedirectResponse
 from starlette.routing import Router
@@ -323,7 +324,7 @@ async def _update_user_email(session: AsyncSession, /, *, user_id: int, email: s
             .values(email=email)
             .options(joinedload(models.User.role))
         )
-    except IntegrityError:
+    except (PostgreSQLIntegrityError, SQLiteIntegrityError):
         raise EmailAlreadyInUse(f"An account for {email} is already in use.")
     user = await session.scalar(
         select(models.User).where(models.User.id == user_id).options(joinedload(models.User.role))