arize-phoenix 12.0.0__tar.gz → 12.2.0__tar.gz

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arize-phoenix
-Version: 12.0.0
+Version: 12.2.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://arize.com/docs/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/db/insertion/document_annotation.py

@@ -181,7 +181,7 @@ def _key(p: Received[Precursors.DocumentAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.DocumentAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid, p.item.document_position, p.item.identifier
+    return p.item.obj.name, p.item.span_rowid, p.item.document_position, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/db/insertion/session_annotation.py

@@ -169,7 +169,7 @@ def _key(p: Received[Precursors.SessionAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.SessionAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.project_session_rowid, p.item.identifier
+    return p.item.obj.name, p.item.project_session_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/db/insertion/span_annotation.py

@@ -167,7 +167,7 @@ def _key(p: Received[Precursors.SpanAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.SpanAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid, p.item.identifier
+    return p.item.obj.name, p.item.span_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/db/insertion/trace_annotation.py

@@ -166,7 +166,7 @@ def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.trace_rowid, p.item.identifier
+    return p.item.obj.name, p.item.trace_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/db/insertion/types.py

@@ -248,7 +248,6 @@ class Insertables(ABC):
     class SpanAnnotation(Precursors.SpanAnnotation):
         updated_at: datetime
         span_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.SpanAnnotation:
@@ -261,7 +260,6 @@ class Insertables(ABC):
     class TraceAnnotation(Precursors.TraceAnnotation):
         updated_at: datetime
         trace_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.TraceAnnotation:
@@ -274,7 +272,6 @@ class Insertables(ABC):
     class DocumentAnnotation(Precursors.DocumentAnnotation):
         updated_at: datetime
         span_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.DocumentAnnotation:
@@ -287,7 +284,6 @@ class Insertables(ABC):
     class SessionAnnotation(Precursors.SessionAnnotation):
         updated_at: datetime
         project_session_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.ProjectSessionAnnotation:

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/db/models.py

@@ -443,6 +443,27 @@ class _RegexStr(TypeDecorator[re.Pattern[str]]):
         return re.compile(value)
 
 
+_HEX_COLOR_PATTERN = re.compile(r"^#([0-9a-f]{6})$")
+
+
+class _HexColor(TypeDecorator[str]):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    cache_ok = True
+    impl = String
+
+    def process_bind_param(self, value: Optional[str], _: Dialect) -> Optional[str]:
+        if value is None:
+            return None
+        if not _HEX_COLOR_PATTERN.match(value):
+            raise ValueError(f"Expected a hex color, got {value}")
+        return value
+
+    def process_result_value(self, value: Optional[str], _: Dialect) -> Optional[str]:
+        if value is None:
+            return None
+        return value
+
+
 class ExperimentRunOutput(TypedDict, total=False):
     task_output: Any
 
@@ -1125,7 +1146,7 @@ class DatasetLabel(HasId):
     __tablename__ = "dataset_labels"
     name: Mapped[str] = mapped_column(unique=True)
     description: Mapped[Optional[str]]
-    color: Mapped[str] = mapped_column(String, nullable=False)
+    color: Mapped[str] = mapped_column(_HexColor, nullable=False)
     datasets_dataset_labels: Mapped[list["DatasetsDatasetLabel"]] = relationship(
         "DatasetsDatasetLabel", back_populates="dataset_label"
     )

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/server/api/context.py

@@ -76,6 +76,7 @@ from phoenix.server.api.dataloaders import (
     UserRolesDataLoader,
     UsersDataLoader,
 )
+from phoenix.server.api.dataloaders.dataset_labels import DatasetLabelsDataLoader
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import DmlEvent
@@ -99,6 +100,7 @@ class DataLoaders:
     average_experiment_run_latency: AverageExperimentRunLatencyDataLoader
     dataset_example_revisions: DatasetExampleRevisionsDataLoader
     dataset_example_spans: DatasetExampleSpansDataLoader
+    dataset_labels: DatasetLabelsDataLoader
     dataset_examples_and_versions_by_experiment_run: (
         DatasetExamplesAndVersionsByExperimentRunDataLoader
     )

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/server/api/dataloaders/__init__.py

@@ -16,6 +16,7 @@ from .dataset_example_splits import DatasetExampleSplitsDataLoader
 from .dataset_examples_and_versions_by_experiment_run import (
     DatasetExamplesAndVersionsByExperimentRunDataLoader,
 )
+from .dataset_labels import DatasetLabelsDataLoader
 from .document_evaluation_summaries import (
     DocumentEvaluationSummaryCache,
     DocumentEvaluationSummaryDataLoader,
@@ -87,6 +88,7 @@ __all__ = [
     "DatasetExampleSpansDataLoader",
     "DatasetExamplesAndVersionsByExperimentRunDataLoader",
     "DatasetExampleSplitsDataLoader",
+    "DatasetLabelsDataLoader",
     "DocumentEvaluationSummaryDataLoader",
     "DocumentEvaluationsDataLoader",
     "DocumentRetrievalMetricsDataLoader",

arize_phoenix-12.2.0/src/phoenix/server/api/dataloaders/dataset_labels.py

@@ -0,0 +1,36 @@
+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
+
+DatasetID: TypeAlias = int
+Key: TypeAlias = DatasetID
+Result: TypeAlias = list[models.DatasetLabel]
+
+
+class DatasetLabelsDataLoader(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]:
+        dataset_ids = keys
+        async with self._db() as session:
+            labels: dict[Key, Result] = {}
+            for dataset_id, label in await session.execute(
+                select(models.DatasetsDatasetLabel.dataset_id, models.DatasetLabel)
+                .select_from(models.DatasetLabel)
+                .join(
+                    models.DatasetsDatasetLabel,
+                    models.DatasetLabel.id == models.DatasetsDatasetLabel.dataset_label_id,
+                )
+                .where(models.DatasetsDatasetLabel.dataset_id.in_(dataset_ids))
+            ):
+                if dataset_id not in labels:
+                    labels[dataset_id] = []
+                labels[dataset_id].append(label)
+        return [
+            sorted(labels.get(dataset_id, []), key=lambda label: label.name) for dataset_id in keys
+        ]

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/server/api/helpers/playground_clients.py

@@ -1677,6 +1677,7 @@ class AnthropicStreamingClient(PlaygroundStreamingClient):
 @register_llm_client(
     provider_key=GenerativeProviderKey.ANTHROPIC,
     model_names=[
+        "claude-sonnet-4-5",
         "claude-sonnet-4-0",
         "claude-sonnet-4-20250514",
         "claude-opus-4-1",

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/server/api/mutations/__init__.py

@@ -5,6 +5,7 @@ from phoenix.server.api.mutations.api_key_mutations import ApiKeyMutationMixin
 from phoenix.server.api.mutations.chat_mutations import (
     ChatCompletionMutationMixin,
 )
+from phoenix.server.api.mutations.dataset_label_mutations import DatasetLabelMutationMixin
 from phoenix.server.api.mutations.dataset_mutations import DatasetMutationMixin
 from phoenix.server.api.mutations.dataset_split_mutations import DatasetSplitMutationMixin
 from phoenix.server.api.mutations.experiment_mutations import ExperimentMutationMixin
@@ -31,6 +32,7 @@ class Mutation(
     AnnotationConfigMutationMixin,
     ApiKeyMutationMixin,
     ChatCompletionMutationMixin,
+    DatasetLabelMutationMixin,
     DatasetMutationMixin,
     DatasetSplitMutationMixin,
     ExperimentMutationMixin,

arize_phoenix-12.2.0/src/phoenix/server/api/mutations/dataset_label_mutations.py

@@ -0,0 +1,291 @@
+from typing import Optional
+
+import sqlalchemy
+import strawberry
+from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry import UNSET
+from strawberry.relay.types import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.auth import IsLocked, IsNotReadOnly
+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.node import from_global_id_with_expected_type
+
+
+@strawberry.input
+class CreateDatasetLabelInput:
+    name: str
+    description: Optional[str] = UNSET
+    color: str
+
+
+@strawberry.type
+class CreateDatasetLabelMutationPayload:
+    dataset_label: DatasetLabel
+
+
+@strawberry.input
+class DeleteDatasetLabelsInput:
+    dataset_label_ids: list[GlobalID]
+
+
+@strawberry.type
+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_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"
+
+
+@strawberry.type
+class DatasetLabelMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_dataset_label(
+        self,
+        info: Info[Context, None],
+        input: CreateDatasetLabelInput,
+    ) -> CreateDatasetLabelMutationPayload:
+        name = input.name
+        description = input.description
+        color = input.color
+        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()
+            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, 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:
+                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)
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def delete_dataset_labels(
+        self, info: Info[Context, None], input: DeleteDatasetLabelsInput
+    ) -> DeleteDatasetLabelsMutationPayload:
+        dataset_label_row_ids: dict[int, None] = {}
+        for dataset_label_node_id in input.dataset_label_ids:
+            try:
+                dataset_label_row_id = from_global_id_with_expected_type(
+                    dataset_label_node_id, DatasetLabel.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Unknown dataset label: {dataset_label_node_id}")
+            dataset_label_row_ids[dataset_label_row_id] = None
+        async with info.context.db() as session:
+            stmt = (
+                delete(models.DatasetLabel)
+                .where(models.DatasetLabel.id.in_(dataset_label_row_ids.keys()))
+                .returning(models.DatasetLabel)
+            )
+            deleted_dataset_labels = (await session.scalars(stmt)).all()
+            if len(deleted_dataset_labels) < len(dataset_label_row_ids):
+                await session.rollback()
+                raise NotFound("Could not find one or more dataset labels with given IDs")
+        deleted_dataset_labels_by_id = {
+            dataset_label.id: dataset_label for dataset_label in deleted_dataset_labels
+        }
+        return DeleteDatasetLabelsMutationPayload(
+            dataset_labels=[
+                to_gql_dataset_label(deleted_dataset_labels_by_id[dataset_label_row_id])
+                for dataset_label_row_id in dataset_label_row_ids
+            ]
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    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)
+
+        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:
+            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")
+
+            existing_dataset_label_ids = (
+                await session.scalars(
+                    select(models.DatasetLabel.id).where(
+                        models.DatasetLabel.id.in_(dataset_label_rowids)
+                    )
+                )
+            ).all()
+            if len(existing_dataset_label_ids) != len(dataset_label_rowids):
+                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())
+
+            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,
+                        )
+                    )
+            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, 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(
+            query=Query(),
+        )

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/server/api/mutations/dataset_split_mutations.py

@@ -15,6 +15,7 @@ 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.node import from_global_id_with_expected_type
 
@@ -68,6 +69,13 @@ class DatasetSplitMutationPayload:
     query: "Query"
 
 
+@strawberry.type
+class DatasetSplitMutationPayloadWithExamples:
+    dataset_split: DatasetSplit
+    query: "Query"
+    examples: list[DatasetExample]
+
+
 @strawberry.type
 class DeleteDatasetSplitsMutationPayload:
     dataset_splits: list[DatasetSplit]
@@ -77,11 +85,13 @@ class DeleteDatasetSplitsMutationPayload:
 @strawberry.type
 class AddDatasetExamplesToDatasetSplitsMutationPayload:
     query: "Query"
+    examples: list[DatasetExample]
 
 
 @strawberry.type
 class RemoveDatasetExamplesFromDatasetSplitsMutationPayload:
     query: "Query"
+    examples: list[DatasetExample]
 
 
 @strawberry.type
@@ -262,8 +272,16 @@ class DatasetSplitMutationMixin:
                 except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
                     raise Conflict("Failed to add examples to dataset splits.") from e
 
+            examples = (
+                await session.scalars(
+                    select(models.DatasetExample).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+            ).all()
         return AddDatasetExamplesToDatasetSplitsMutationPayload(
             query=Query(),
+            examples=[to_gql_dataset_example(example) for example in examples],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
@@ -314,14 +332,23 @@ class DatasetSplitMutationMixin:
 
             await session.execute(stmt)
 
+            examples = (
+                await session.scalars(
+                    select(models.DatasetExample).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+            ).all()
+
         return RemoveDatasetExamplesFromDatasetSplitsMutationPayload(
             query=Query(),
+            examples=[to_gql_dataset_example(example) for example in examples],
         )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
     async def create_dataset_split_with_examples(
         self, info: Info[Context, None], input: CreateDatasetSplitWithExamplesInput
-    ) -> DatasetSplitMutationPayload:
+    ) -> DatasetSplitMutationPayloadWithExamples:
         user_id = get_user(info)
         validated_name = _validated_name(input.name)
         unique_example_rowids: set[int] = set()
@@ -374,9 +401,18 @@ class DatasetSplitMutationMixin:
                         "Failed to associate examples with the new dataset split."
                     ) from e
 
-        return DatasetSplitMutationPayload(
+            examples = (
+                await session.scalars(
+                    select(models.DatasetExample).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+            ).all()
+
+        return DatasetSplitMutationPayloadWithExamples(
             dataset_split=to_gql_dataset_split(dataset_split_orm),
             query=Query(),
+            examples=[to_gql_dataset_example(example) for example in examples],
         )
 
 

{arize_phoenix-12.0.0 → arize_phoenix-12.2.0}/src/phoenix/server/api/queries.py

@@ -48,6 +48,7 @@ from phoenix.server.api.types.AnnotationConfig import AnnotationConfig, to_gql_a
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
 from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
 from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
 from phoenix.server.api.types.Dimension import to_gql_dimension
 from phoenix.server.api.types.EmbeddingDimension import (
@@ -1149,6 +1150,26 @@ class Query:
                 args=args,
             )
 
+    @strawberry.field
+    async def dataset_labels(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[DatasetLabel]:
+        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:
+            dataset_labels = await session.scalars(select(models.DatasetLabel))
+        data = [to_gql_dataset_label(dataset_label) for dataset_label in dataset_labels]
+        return connection_from_list(data=data, args=args)
+
     @strawberry.field
     async def dataset_splits(
         self,