arize-phoenix 11.38.0__py3-none-any.whl → 12.2.0__py3-none-any.whl

phoenix/server/api/mutations/dataset_split_mutations.py

@@ -0,0 +1,423 @@
+from typing import Optional
+
+import strawberry
+from sqlalchemy import delete, func, insert, 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 import GlobalID
+from strawberry.scalars import JSON
+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.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
+
+
+@strawberry.input
+class CreateDatasetSplitInput:
+    name: str
+    description: Optional[str] = UNSET
+    color: str
+    metadata: Optional[JSON] = UNSET
+
+
+@strawberry.input
+class PatchDatasetSplitInput:
+    dataset_split_id: GlobalID
+    name: Optional[str] = UNSET
+    description: Optional[str] = UNSET
+    color: Optional[str] = UNSET
+    metadata: Optional[JSON] = UNSET
+
+
+@strawberry.input
+class DeleteDatasetSplitInput:
+    dataset_split_ids: list[GlobalID]
+
+
+@strawberry.input
+class AddDatasetExamplesToDatasetSplitsInput:
+    dataset_split_ids: list[GlobalID]
+    example_ids: list[GlobalID]
+
+
+@strawberry.input
+class RemoveDatasetExamplesFromDatasetSplitsInput:
+    dataset_split_ids: list[GlobalID]
+    example_ids: list[GlobalID]
+
+
+@strawberry.input
+class CreateDatasetSplitWithExamplesInput:
+    name: str
+    description: Optional[str] = UNSET
+    color: str
+    metadata: Optional[JSON] = UNSET
+    example_ids: list[GlobalID]
+
+
+@strawberry.type
+class DatasetSplitMutationPayload:
+    dataset_split: DatasetSplit
+    query: "Query"
+
+
+@strawberry.type
+class DatasetSplitMutationPayloadWithExamples:
+    dataset_split: DatasetSplit
+    query: "Query"
+    examples: list[DatasetExample]
+
+
+@strawberry.type
+class DeleteDatasetSplitsMutationPayload:
+    dataset_splits: list[DatasetSplit]
+    query: "Query"
+
+
+@strawberry.type
+class AddDatasetExamplesToDatasetSplitsMutationPayload:
+    query: "Query"
+    examples: list[DatasetExample]
+
+
+@strawberry.type
+class RemoveDatasetExamplesFromDatasetSplitsMutationPayload:
+    query: "Query"
+    examples: list[DatasetExample]
+
+
+@strawberry.type
+class DatasetSplitMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_dataset_split(
+        self, info: Info[Context, None], input: CreateDatasetSplitInput
+    ) -> DatasetSplitMutationPayload:
+        user_id = get_user(info)
+        validated_name = _validated_name(input.name)
+        async with info.context.db() as session:
+            dataset_split_orm = models.DatasetSplit(
+                name=validated_name,
+                description=input.description,
+                color=input.color,
+                metadata_=input.metadata or {},
+                user_id=user_id,
+            )
+            session.add(dataset_split_orm)
+            try:
+                await session.commit()
+            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()
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def patch_dataset_split(
+        self, info: Info[Context, None], input: PatchDatasetSplitInput
+    ) -> DatasetSplitMutationPayload:
+        validated_name = _validated_name(input.name) if input.name else None
+        async with info.context.db() as session:
+            dataset_split_id = from_global_id_with_expected_type(
+                input.dataset_split_id, DatasetSplit.__name__
+            )
+            dataset_split_orm = await session.get(models.DatasetSplit, dataset_split_id)
+            if not dataset_split_orm:
+                raise NotFound(f"Dataset split with ID {input.dataset_split_id} not found")
+
+            if validated_name:
+                dataset_split_orm.name = validated_name
+            if input.description:
+                dataset_split_orm.description = input.description
+            if input.color:
+                dataset_split_orm.color = input.color
+            if isinstance(input.metadata, dict):
+                dataset_split_orm.metadata_ = input.metadata
+
+            gql_dataset_split = to_gql_dataset_split(dataset_split_orm)
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict("A dataset split with this name already exists")
+
+        return DatasetSplitMutationPayload(
+            dataset_split=gql_dataset_split,
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_dataset_splits(
+        self, info: Info[Context, None], input: DeleteDatasetSplitInput
+    ) -> DeleteDatasetSplitsMutationPayload:
+        unique_dataset_split_rowids: dict[int, None] = {}  # use a dict to preserve ordering
+        for dataset_split_gid in input.dataset_split_ids:
+            try:
+                dataset_split_rowid = from_global_id_with_expected_type(
+                    dataset_split_gid, DatasetSplit.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid dataset split ID: {dataset_split_gid}")
+            unique_dataset_split_rowids[dataset_split_rowid] = None
+        dataset_split_rowids = list(unique_dataset_split_rowids.keys())
+
+        async with info.context.db() as session:
+            deleted_splits_by_id = {
+                split.id: split
+                for split in (
+                    await session.scalars(
+                        delete(models.DatasetSplit)
+                        .where(models.DatasetSplit.id.in_(dataset_split_rowids))
+                        .returning(models.DatasetSplit)
+                    )
+                ).all()
+            }
+            if len(deleted_splits_by_id) < len(dataset_split_rowids):
+                await session.rollback()
+                raise NotFound("One or more dataset splits not found")
+            await session.commit()
+
+        return DeleteDatasetSplitsMutationPayload(
+            dataset_splits=[
+                to_gql_dataset_split(deleted_splits_by_id[dataset_split_rowid])
+                for dataset_split_rowid in dataset_split_rowids
+            ],
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def add_dataset_examples_to_dataset_splits(
+        self, info: Info[Context, None], input: AddDatasetExamplesToDatasetSplitsInput
+    ) -> AddDatasetExamplesToDatasetSplitsMutationPayload:
+        if not input.example_ids:
+            raise BadRequest("No examples provided.")
+        if not input.dataset_split_ids:
+            raise BadRequest("No dataset splits provided.")
+
+        unique_dataset_split_rowids: set[int] = set()
+        for dataset_split_gid in input.dataset_split_ids:
+            try:
+                dataset_split_rowid = from_global_id_with_expected_type(
+                    dataset_split_gid, DatasetSplit.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid dataset split ID: {dataset_split_gid}")
+            unique_dataset_split_rowids.add(dataset_split_rowid)
+        dataset_split_rowids = list(unique_dataset_split_rowids)
+
+        unique_example_rowids: set[int] = set()
+        for example_gid in input.example_ids:
+            try:
+                example_rowid = from_global_id_with_expected_type(
+                    example_gid, models.DatasetExample.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid example ID: {example_gid}")
+            unique_example_rowids.add(example_rowid)
+        example_rowids = list(unique_example_rowids)
+
+        async with info.context.db() as session:
+            existing_dataset_split_ids = (
+                await session.scalars(
+                    select(models.DatasetSplit.id).where(
+                        models.DatasetSplit.id.in_(dataset_split_rowids)
+                    )
+                )
+            ).all()
+            if len(existing_dataset_split_ids) != len(dataset_split_rowids):
+                raise NotFound("One or more dataset splits not found")
+
+            # Find existing (dataset_split_id, dataset_example_id) keys to avoid duplicates
+            # Users can submit multiple examples at once which can have
+            # indeterminate participation in multiple splits
+            existing_dataset_example_split_keys = await session.execute(
+                select(
+                    models.DatasetSplitDatasetExample.dataset_split_id,
+                    models.DatasetSplitDatasetExample.dataset_example_id,
+                ).where(
+                    models.DatasetSplitDatasetExample.dataset_split_id.in_(dataset_split_rowids)
+                    & models.DatasetSplitDatasetExample.dataset_example_id.in_(example_rowids)
+                )
+            )
+            unique_dataset_example_split_keys = set(existing_dataset_example_split_keys.all())
+
+            # Compute all desired pairs and insert only missing
+            values = []
+            for dataset_split_rowid in dataset_split_rowids:
+                for example_rowid in example_rowids:
+                    # if the keys already exists, skip
+                    if (dataset_split_rowid, example_rowid) in unique_dataset_example_split_keys:
+                        continue
+                    dataset_split_id_key = models.DatasetSplitDatasetExample.dataset_split_id.key
+                    dataset_example_id_key = (
+                        models.DatasetSplitDatasetExample.dataset_example_id.key
+                    )
+                    values.append(
+                        {
+                            dataset_split_id_key: dataset_split_rowid,
+                            dataset_example_id_key: example_rowid,
+                        }
+                    )
+
+            if values:
+                try:
+                    await session.execute(insert(models.DatasetSplitDatasetExample), values)
+                    await session.flush()
+                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
+    async def remove_dataset_examples_from_dataset_splits(
+        self, info: Info[Context, None], input: RemoveDatasetExamplesFromDatasetSplitsInput
+    ) -> RemoveDatasetExamplesFromDatasetSplitsMutationPayload:
+        if not input.dataset_split_ids:
+            raise BadRequest("No dataset splits provided.")
+        if not input.example_ids:
+            raise BadRequest("No examples provided.")
+
+        unique_dataset_split_rowids: set[int] = set()
+        for dataset_split_gid in input.dataset_split_ids:
+            try:
+                dataset_split_rowid = from_global_id_with_expected_type(
+                    dataset_split_gid, DatasetSplit.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid dataset split ID: {dataset_split_gid}")
+            unique_dataset_split_rowids.add(dataset_split_rowid)
+        dataset_split_rowids = list(unique_dataset_split_rowids)
+
+        unique_example_rowids: set[int] = set()
+        for example_gid in input.example_ids:
+            try:
+                example_rowid = from_global_id_with_expected_type(
+                    example_gid, models.DatasetExample.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid example ID: {example_gid}")
+            unique_example_rowids.add(example_rowid)
+        example_rowids = list(unique_example_rowids)
+
+        stmt = delete(models.DatasetSplitDatasetExample).where(
+            models.DatasetSplitDatasetExample.dataset_split_id.in_(dataset_split_rowids)
+            & models.DatasetSplitDatasetExample.dataset_example_id.in_(example_rowids)
+        )
+        async with info.context.db() as session:
+            existing_dataset_split_ids = (
+                await session.scalars(
+                    select(models.DatasetSplit.id).where(
+                        models.DatasetSplit.id.in_(dataset_split_rowids)
+                    )
+                )
+            ).all()
+            if len(existing_dataset_split_ids) != len(dataset_split_rowids):
+                raise NotFound("One or more dataset splits not found")
+
+            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
+    ) -> DatasetSplitMutationPayloadWithExamples:
+        user_id = get_user(info)
+        validated_name = _validated_name(input.name)
+        unique_example_rowids: set[int] = set()
+        for example_gid in input.example_ids:
+            try:
+                example_rowid = from_global_id_with_expected_type(
+                    example_gid, models.DatasetExample.__name__
+                )
+                unique_example_rowids.add(example_rowid)
+            except ValueError:
+                raise BadRequest(f"Invalid example ID: {example_gid}")
+        example_rowids = list(unique_example_rowids)
+        async with info.context.db() as session:
+            if example_rowids:
+                found_count = await session.scalar(
+                    select(func.count(models.DatasetExample.id)).where(
+                        models.DatasetExample.id.in_(example_rowids)
+                    )
+                )
+                if found_count is None or found_count < len(example_rowids):
+                    raise NotFound("One or more dataset examples were not found.")
+
+            dataset_split_orm = models.DatasetSplit(
+                name=validated_name,
+                description=input.description or None,
+                color=input.color,
+                metadata_=input.metadata or {},
+                user_id=user_id,
+            )
+            session.add(dataset_split_orm)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"A dataset split named '{validated_name}' already exists.")
+
+            if example_rowids:
+                values = [
+                    {
+                        models.DatasetSplitDatasetExample.dataset_split_id.key: dataset_split_orm.id,  # noqa: E501
+                        models.DatasetSplitDatasetExample.dataset_example_id.key: example_id,
+                    }
+                    for example_id in example_rowids
+                ]
+                try:
+                    await session.execute(insert(models.DatasetSplitDatasetExample), values)
+                except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
+                    # Roll back the transaction on association failure
+                    await session.rollback()
+                    raise Conflict(
+                        "Failed to associate examples with the new dataset split."
+                    ) from e
+
+            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],
+        )
+
+
+def _validated_name(name: str) -> str:
+    validated_name = name.strip()
+    if not validated_name:
+        raise BadRequest("Name cannot be empty")
+    return validated_name

phoenix/server/api/mutations/project_session_annotations_mutations.py

@@ -0,0 +1,161 @@
+from typing import Optional
+
+import strawberry
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from starlette.requests import Request
+from strawberry import Info
+from strawberry.relay import GlobalID
+
+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, Unauthorized
+from phoenix.server.api.helpers.annotations import get_user_identifier
+from phoenix.server.api.input_types.CreateProjectSessionAnnotationInput import (
+    CreateProjectSessionAnnotationInput,
+)
+from phoenix.server.api.input_types.UpdateAnnotationInput import UpdateAnnotationInput
+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.bearer_auth import PhoenixUser
+from phoenix.server.dml_event import (
+    ProjectSessionAnnotationDeleteEvent,
+    ProjectSessionAnnotationInsertEvent,
+)
+
+
+@strawberry.type
+class ProjectSessionAnnotationMutationPayload:
+    project_session_annotation: ProjectSessionAnnotation
+    query: Query
+
+
+@strawberry.type
+class ProjectSessionAnnotationMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_project_session_annotations(
+        self, info: Info[Context, None], input: CreateProjectSessionAnnotationInput
+    ) -> ProjectSessionAnnotationMutationPayload:
+        assert isinstance(request := info.context.request, Request)
+        user_id: Optional[int] = None
+        if "user" in request.scope and isinstance((user := info.context.user), PhoenixUser):
+            user_id = int(user.identity)
+
+        try:
+            project_session_id = from_global_id_with_expected_type(
+                input.project_session_id, "ProjectSession"
+            )
+        except ValueError:
+            raise BadRequest(f"Invalid session ID: {input.project_session_id}")
+
+        identifier = ""
+        if isinstance(input.identifier, str):
+            identifier = input.identifier  # Already trimmed in __post_init__
+        elif input.source == AnnotationSource.APP and user_id is not None:
+            identifier = get_user_identifier(user_id)
+
+        try:
+            async with info.context.db() as session:
+                anno = models.ProjectSessionAnnotation(
+                    project_session_id=project_session_id,
+                    name=input.name,
+                    label=input.label,
+                    score=input.score,
+                    explanation=input.explanation,
+                    annotator_kind=input.annotator_kind.value,
+                    metadata_=input.metadata,
+                    identifier=identifier,
+                    source=input.source.value,
+                    user_id=user_id,
+                )
+                session.add(anno)
+        except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
+            raise Conflict(f"Error creating annotation: {e}")
+
+        info.context.event_queue.put(ProjectSessionAnnotationInsertEvent((anno.id,)))
+
+        return ProjectSessionAnnotationMutationPayload(
+            project_session_annotation=to_gql_project_session_annotation(anno),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def update_project_session_annotations(
+        self, info: Info[Context, None], input: UpdateAnnotationInput
+    ) -> ProjectSessionAnnotationMutationPayload:
+        assert isinstance(request := info.context.request, Request)
+        user_id: Optional[int] = None
+        if "user" in request.scope and isinstance((user := info.context.user), PhoenixUser):
+            user_id = int(user.identity)
+
+        try:
+            id_ = from_global_id_with_expected_type(input.id, "ProjectSessionAnnotation")
+        except ValueError:
+            raise BadRequest(f"Invalid session annotation ID: {input.id}")
+
+        async with info.context.db() as session:
+            if not (anno := await session.get(models.ProjectSessionAnnotation, id_)):
+                raise NotFound(f"Could not find session annotation with ID: {input.id}")
+            if anno.user_id != user_id:
+                raise Unauthorized("Session annotation is not associated with the current user.")
+
+            # Update the annotation fields
+            anno.name = input.name
+            anno.label = input.label
+            anno.score = input.score
+            anno.explanation = input.explanation
+            anno.annotator_kind = input.annotator_kind.value
+            anno.metadata_ = input.metadata
+            anno.source = input.source.value
+
+            session.add(anno)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError) as e:
+                raise Conflict(f"Error updating annotation: {e}")
+
+        info.context.event_queue.put(ProjectSessionAnnotationInsertEvent((anno.id,)))
+        return ProjectSessionAnnotationMutationPayload(
+            project_session_annotation=to_gql_project_session_annotation(anno),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_project_session_annotation(
+        self, info: Info[Context, None], id: GlobalID
+    ) -> ProjectSessionAnnotationMutationPayload:
+        try:
+            id_ = from_global_id_with_expected_type(id, "ProjectSessionAnnotation")
+        except ValueError:
+            raise BadRequest(f"Invalid session annotation ID: {id}")
+
+        assert isinstance(request := info.context.request, Request)
+        user_id: Optional[int] = None
+        user_is_admin = False
+        if "user" in request.scope and isinstance((user := info.context.user), PhoenixUser):
+            user_id = int(user.identity)
+            user_is_admin = user.is_admin
+
+        async with info.context.db() as session:
+            if not (anno := await session.get(models.ProjectSessionAnnotation, id_)):
+                raise NotFound(f"Could not find session annotation with ID: {id}")
+
+            if not user_is_admin and anno.user_id != user_id:
+                raise Unauthorized(
+                    "Session annotation is not associated with the current user and "
+                    "the current user is not an admin."
+                )
+
+            await session.delete(anno)
+
+        deleted_gql_annotation = to_gql_project_session_annotation(anno)
+        info.context.event_queue.put(ProjectSessionAnnotationDeleteEvent((id_,)))
+        return ProjectSessionAnnotationMutationPayload(
+            project_session_annotation=deleted_gql_annotation, query=Query()
+        )

phoenix/server/api/queries.py

@@ -48,6 +48,8 @@ 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 (
     DEFAULT_CLUSTER_SELECTION_EPSILON,
@@ -959,6 +961,14 @@ class Query:
                 id_attr=example.id,
                 created_at=example.created_at,
             )
+        elif type_name == DatasetSplit.__name__:
+            async with info.context.db() as session:
+                dataset_split = await session.scalar(
+                    select(models.DatasetSplit).where(models.DatasetSplit.id == node_id)
+                )
+            if not dataset_split:
+                raise NotFound(f"Unknown dataset split: {id}")
+            return to_gql_dataset_split(dataset_split)
         elif type_name == Experiment.__name__:
             async with info.context.db() as session:
                 experiment = await session.scalar(
@@ -1140,6 +1150,49 @@ 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,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[DatasetSplit]:
+        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:
+            splits = await session.stream_scalars(select(models.DatasetSplit))
+            data = [to_gql_dataset_split(split) async for split in splits]
+            return connection_from_list(
+                data=data,
+                args=args,
+            )
+
     @strawberry.field
     async def annotation_configs(
         self,