arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/api/mutations/chat_mutations.py

@@ -1,9 +1,10 @@
 import asyncio
+import logging
 from dataclasses import asdict, field
 from datetime import datetime, timezone
 from itertools import chain, islice
 from traceback import format_exc
-from typing import Any, Iterable, Iterator, List, Optional, TypeVar, Union
+from typing import Any, Iterable, Iterator, Optional, TypeVar, Union
 
 import strawberry
 from openinference.instrumentation import safe_json_dumps
@@ -22,14 +23,19 @@ from strawberry.relay import GlobalID
 from strawberry.types import Info
 from typing_extensions import assert_never
 
+from phoenix.config import PLAYGROUND_PROJECT_NAME
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
-from phoenix.db.helpers import get_dataset_example_revisions
-from phoenix.server.api.auth import IsLocked, IsNotReadOnly
+from phoenix.db.helpers import (
+    get_dataset_example_revisions,
+    insert_experiment_with_examples_snapshot,
+)
+from phoenix.server.api.auth import IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
 from phoenix.server.api.helpers.dataset_helpers import get_dataset_example_output
 from phoenix.server.api.helpers.playground_clients import (
+    PlaygroundClientCredential,
     PlaygroundStreamingClient,
     initialize_playground_clients,
 )
@@ -43,6 +49,7 @@ from phoenix.server.api.helpers.playground_spans import (
     llm_tools,
     prompt_metadata,
 )
+from phoenix.server.api.helpers.playground_users import get_user
 from phoenix.server.api.helpers.prompts.models import PromptTemplateFormat
 from phoenix.server.api.input_types.ChatCompletionInput import (
     ChatCompletionInput,
@@ -62,6 +69,7 @@ from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Span import Span
 from phoenix.server.dml_event import SpanInsertEvent
+from phoenix.server.experiments.utils import generate_experiment_project_name
 from phoenix.trace.attributes import unflatten
 from phoenix.trace.schemas import SpanException
 from phoenix.utilities.json import jsonify
@@ -72,9 +80,11 @@ from phoenix.utilities.template_formatters import (
     TemplateFormatter,
 )
 
+logger = logging.getLogger(__name__)
+
 initialize_playground_clients()
 
-ChatCompletionMessage = tuple[ChatCompletionMessageRole, str, Optional[str], Optional[List[Any]]]
+ChatCompletionMessage = tuple[ChatCompletionMessageRole, str, Optional[str], Optional[list[Any]]]
 
 
 @strawberry.type
@@ -90,24 +100,25 @@ class ChatCompletionToolCall:
 
 
 @strawberry.type
-class ChatCompletionMutationPayload:
-    db_span: strawberry.Private[models.Span]
+class ChatCompletionRepetition:
+    repetition_number: int
     content: Optional[str]
-    tool_calls: List[ChatCompletionToolCall]
-    span: Span
+    tool_calls: list[ChatCompletionToolCall]
+    span: Optional[Span]
     error_message: Optional[str]
 
 
 @strawberry.type
-class ChatCompletionMutationError:
-    message: str
+class ChatCompletionMutationPayload:
+    repetitions: list[ChatCompletionRepetition]
 
 
 @strawberry.type
 class ChatCompletionOverDatasetMutationExamplePayload:
     dataset_example_id: GlobalID
+    repetition_number: int
     experiment_run_id: GlobalID
-    result: Union[ChatCompletionMutationPayload, ChatCompletionMutationError]
+    repetition: ChatCompletionRepetition
 
 
 @strawberry.type
@@ -120,7 +131,7 @@ class ChatCompletionOverDatasetMutationPayload:
 
 @strawberry.type
 class ChatCompletionMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     @classmethod
     async def chat_completion_over_dataset(
         cls,
@@ -132,9 +143,17 @@ class ChatCompletionMutationMixin:
         if llm_client_class is None:
             raise BadRequest(f"Unknown LLM provider: '{provider_key.value}'")
         try:
+            # Convert GraphQL credentials to PlaygroundCredential objects
+            credentials = None
+            if input.credentials:
+                credentials = [
+                    PlaygroundClientCredential(env_var_name=cred.env_var_name, value=cred.value)
+                    for cred in input.credentials
+                ]
+
             llm_client = llm_client_class(
                 model=input.model,
-                api_key=input.api_key,
+                credentials=credentials,
             )
         except CustomGraphQLError:
             raise
@@ -151,6 +170,7 @@ class ChatCompletionMutationMixin:
             if input.dataset_version_id
             else None
         )
+        project_name = generate_experiment_project_name()
         async with info.context.db() as session:
             dataset = await session.scalar(select(models.Dataset).filter_by(id=dataset_id))
             if dataset is None:
@@ -166,16 +186,26 @@ class ChatCompletionMutationMixin:
                     raise NotFound("No versions found for the given dataset")
             else:
                 resolved_version_id = dataset_version_id
+            # Parse split IDs if provided
+            resolved_split_ids: Optional[list[int]] = None
+            if input.split_ids is not None and len(input.split_ids) > 0:
+                resolved_split_ids = [
+                    from_global_id_with_expected_type(split_id, models.DatasetSplit.__name__)
+                    for split_id in input.split_ids
+                ]
+
             revisions = [
                 revision
                 async for revision in await session.stream_scalars(
-                    get_dataset_example_revisions(resolved_version_id).order_by(
-                        models.DatasetExampleRevision.id
-                    )
+                    get_dataset_example_revisions(
+                        resolved_version_id,
+                        split_ids=resolved_split_ids,
+                    ).order_by(models.DatasetExampleRevision.id)
                 )
             ]
             if not revisions:
                 raise NotFound("No examples found for the given dataset and version")
+            user_id = get_user(info)
             experiment = models.Experiment(
                 dataset_id=from_global_id_with_expected_type(input.dataset_id, Dataset.__name__),
                 dataset_version_id=resolved_version_id,
@@ -184,15 +214,25 @@ class ChatCompletionMutationMixin:
                 description=input.experiment_description,
                 repetitions=1,
                 metadata_=input.experiment_metadata or dict(),
-                project_name=PLAYGROUND_PROJECT_NAME,
+                project_name=project_name,
+                user_id=user_id,
             )
-            session.add(experiment)
-            await session.flush()
-
-        results: list[Union[ChatCompletionMutationPayload, BaseException]] = []
+            if resolved_split_ids:
+                experiment.experiment_dataset_splits = [
+                    models.ExperimentDatasetSplit(dataset_split_id=split_id)
+                    for split_id in resolved_split_ids
+                ]
+            await insert_experiment_with_examples_snapshot(session, experiment)
+
+        results: list[Union[tuple[ChatCompletionRepetition, models.Span], BaseException]] = []
         batch_size = 3
         start_time = datetime.now(timezone.utc)
-        for batch in _get_batches(revisions, batch_size):
+        unbatched_items = [
+            (revision, repetition_number)
+            for revision in revisions
+            for repetition_number in range(1, input.repetitions + 1)
+        ]
+        for batch in _get_batches(unbatched_items, batch_size):
             batch_results = await asyncio.gather(
                 *(
                     cls._chat_completion(
@@ -200,7 +240,7 @@ class ChatCompletionMutationMixin:
                         llm_client,
                         ChatCompletionInput(
                             model=input.model,
-                            api_key=input.api_key,
+                            credentials=input.credentials,
                             messages=input.messages,
                             tools=input.tools,
                             invocation_parameters=input.invocation_parameters,
@@ -209,9 +249,12 @@ class ChatCompletionMutationMixin:
                                 variables=revision.input,
                             ),
                             prompt_name=input.prompt_name,
+                            repetitions=repetition_number,
                         ),
+                        repetition_number=repetition_number,
+                        project_name=project_name,
                     )
-                    for revision in batch
+                    for revision, repetition_number in batch
                 ),
                 return_exceptions=True,
             )
@@ -223,19 +266,19 @@ class ChatCompletionMutationMixin:
             experiment_id=GlobalID(models.Experiment.__name__, str(experiment.id)),
         )
         experiment_runs = []
-        for revision, result in zip(revisions, results):
+        for (revision, repetition_number), result in zip(unbatched_items, results):
             if isinstance(result, BaseException):
                 experiment_run = models.ExperimentRun(
                     experiment_id=experiment.id,
                     dataset_example_id=revision.dataset_example_id,
                     output={},
-                    repetition_number=1,
+                    repetition_number=repetition_number,
                     start_time=start_time,
                     end_time=start_time,
                     error=str(result),
                 )
             else:
-                db_span: models.Span = result.db_span
+                repetition, db_span = result
                 experiment_run = models.ExperimentRun(
                     experiment_id=experiment.id,
                     dataset_example_id=revision.dataset_example_id,
@@ -245,10 +288,10 @@ class ChatCompletionMutationMixin:
                     ),
                     prompt_token_count=db_span.cumulative_llm_token_count_prompt,
                     completion_token_count=db_span.cumulative_llm_token_count_completion,
-                    repetition_number=1,
+                    repetition_number=repetition_number,
                     start_time=db_span.start_time,
                     end_time=db_span.end_time,
-                    error=str(result.error_message) if result.error_message else None,
+                    error=str(repetition.error_message) if repetition.error_message else None,
                 )
             experiment_runs.append(experiment_run)
 
@@ -256,22 +299,31 @@ class ChatCompletionMutationMixin:
             session.add_all(experiment_runs)
             await session.flush()
 
-        for revision, experiment_run, result in zip(revisions, experiment_runs, results):
+        for (revision, repetition_number), experiment_run, result in zip(
+            unbatched_items, experiment_runs, results
+        ):
             dataset_example_id = GlobalID(
                 models.DatasetExample.__name__, str(revision.dataset_example_id)
             )
             experiment_run_id = GlobalID(models.ExperimentRun.__name__, str(experiment_run.id))
             example_payload = ChatCompletionOverDatasetMutationExamplePayload(
                 dataset_example_id=dataset_example_id,
+                repetition_number=repetition_number,
                 experiment_run_id=experiment_run_id,
-                result=result
-                if isinstance(result, ChatCompletionMutationPayload)
-                else ChatCompletionMutationError(message=str(result)),
+                repetition=ChatCompletionRepetition(
+                    repetition_number=repetition_number,
+                    content=None,
+                    tool_calls=[],
+                    span=None,
+                    error_message=str(result),
+                )
+                if isinstance(result, BaseException)
+                else result[0],
             )
             payload.examples.append(example_payload)
         return payload
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     @classmethod
     async def chat_completion(
         cls, info: Info[Context, None], input: ChatCompletionInput
@@ -281,9 +333,17 @@ class ChatCompletionMutationMixin:
         if llm_client_class is None:
             raise BadRequest(f"Unknown LLM provider: '{provider_key.value}'")
         try:
+            # Convert GraphQL credentials to PlaygroundCredential objects
+            credentials = None
+            if input.credentials:
+                credentials = [
+                    PlaygroundClientCredential(env_var_name=cred.env_var_name, value=cred.value)
+                    for cred in input.credentials
+                ]
+
             llm_client = llm_client_class(
                 model=input.model,
-                api_key=input.api_key,
+                credentials=credentials,
             )
         except CustomGraphQLError:
             raise
@@ -292,7 +352,38 @@ class ChatCompletionMutationMixin:
                 f"Failed to connect to LLM API for {provider_key.value} {input.model.name}: "
                 f"{str(error)}"
             )
-        return await cls._chat_completion(info, llm_client, input)
+
+        results: list[Union[tuple[ChatCompletionRepetition, models.Span], BaseException]] = []
+        batch_size = 3
+        for batch in _get_batches(range(1, input.repetitions + 1), batch_size):
+            batch_results = await asyncio.gather(
+                *(
+                    cls._chat_completion(
+                        info, llm_client, input, repetition_number=repetition_number
+                    )
+                    for repetition_number in batch
+                ),
+                return_exceptions=True,
+            )
+            results.extend(batch_results)
+
+        repetitions: list[ChatCompletionRepetition] = []
+        for repetition_number, result in enumerate(results, start=1):
+            if isinstance(result, BaseException):
+                repetitions.append(
+                    ChatCompletionRepetition(
+                        repetition_number=repetition_number,
+                        content=None,
+                        tool_calls=[],
+                        span=None,
+                        error_message=str(result),
+                    )
+                )
+            else:
+                repetition, _ = result
+                repetitions.append(repetition)
+
+        return ChatCompletionMutationPayload(repetitions=repetitions)
 
     @classmethod
     async def _chat_completion(
@@ -300,7 +391,10 @@ class ChatCompletionMutationMixin:
         info: Info[Context, None],
         llm_client: PlaygroundStreamingClient,
         input: ChatCompletionInput,
-    ) -> ChatCompletionMutationPayload:
+        repetition_number: int,
+        project_name: str = PLAYGROUND_PROJECT_NAME,
+        project_description: str = "Traces from prompt playground",
+    ) -> tuple[ChatCompletionRepetition, models.Span]:
         attributes: dict[str, Any] = {}
         attributes.update(dict(prompt_metadata(input.prompt_name)))
 
@@ -394,15 +488,15 @@ class ChatCompletionMutationMixin:
             # Get or create the project ID
             if (
                 project_id := await session.scalar(
-                    select(models.Project.id).where(models.Project.name == PLAYGROUND_PROJECT_NAME)
+                    select(models.Project.id).where(models.Project.name == project_name)
                 )
             ) is None:
                 project_id = await session.scalar(
                     insert(models.Project)
                     .returning(models.Project.id)
                     .values(
-                        name=PLAYGROUND_PROJECT_NAME,
-                        description="Traces from prompt playground",
+                        name=project_name,
+                        description=project_description,
                     )
                 )
             trace = models.Trace(
@@ -433,27 +527,41 @@ class ChatCompletionMutationMixin:
             session.add(trace)
             session.add(span)
             await session.flush()
+            try:
+                span_cost = info.context.span_cost_calculator.calculate_cost(
+                    start_time=span.start_time,
+                    attributes=span.attributes,
+                )
+            except Exception as e:
+                logger.exception(f"Failed to calculate cost for span {span.id}: {e}")
+                span_cost = None
+            if span_cost:
+                span_cost.span_rowid = span.id
+                span_cost.trace_rowid = trace.id
+                session.add(span_cost)
+                await session.flush()
 
-        gql_span = Span(span_rowid=span.id, db_span=span)
+        gql_span = Span(id=span.id, db_record=span)
 
         info.context.event_queue.put(SpanInsertEvent(ids=(project_id,)))
 
         if status_code is StatusCode.ERROR:
-            return ChatCompletionMutationPayload(
-                db_span=span,
+            repetition = ChatCompletionRepetition(
+                repetition_number=repetition_number,
                 content=None,
                 tool_calls=[],
                 span=gql_span,
                 error_message=status_message,
             )
         else:
-            return ChatCompletionMutationPayload(
-                db_span=span,
+            repetition = ChatCompletionRepetition(
+                repetition_number=repetition_number,
                 content=text_content if text_content else None,
                 tool_calls=list(tool_calls.values()),
                 span=gql_span,
                 error_message=None,
             )
+        return repetition, span
 
 
 def _formatted_messages(
@@ -588,5 +696,4 @@ TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUME
 TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA
 PROMPT_TEMPLATE_VARIABLES = SpanAttributes.LLM_PROMPT_TEMPLATE_VARIABLES
 
-
-PLAYGROUND_PROJECT_NAME = "playground"
+LLM_PROVIDER = SpanAttributes.LLM_PROVIDER

phoenix/server/api/mutations/dataset_label_mutations.py

@@ -0,0 +1,243 @@
+from typing import Optional
+
+import sqlalchemy
+import strawberry
+from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.orm import joinedload
+from sqlalchemy.sql import tuple_
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry import UNSET
+from strawberry.relay.types import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+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
+from phoenix.server.api.types.DatasetLabel import DatasetLabel
+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
+    dataset_ids: Optional[list[GlobalID]] = UNSET
+
+
+@strawberry.type
+class CreateDatasetLabelMutationPayload:
+    dataset_label: DatasetLabel
+    datasets: list[Dataset]
+
+
+@strawberry.input
+class DeleteDatasetLabelsInput:
+    dataset_label_ids: list[GlobalID]
+
+
+@strawberry.type
+class DeleteDatasetLabelsMutationPayload:
+    dataset_labels: list[DatasetLabel]
+
+
+@strawberry.input
+class SetDatasetLabelsInput:
+    dataset_id: GlobalID
+    dataset_label_ids: list[GlobalID]
+
+
+@strawberry.type
+class SetDatasetLabelsMutationPayload:
+    query: Query
+    dataset: Dataset
+
+
+@strawberry.type
+class DatasetLabelMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, 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
+        dataset_rowids: dict[
+            int, None
+        ] = {}  # use dictionary to de-duplicate while preserving order
+        if input.dataset_ids:
+            for dataset_id in input.dataset_ids:
+                try:
+                    dataset_rowid = from_global_id_with_expected_type(dataset_id, Dataset.__name__)
+                except ValueError:
+                    raise BadRequest(f"Invalid dataset ID: {dataset_id}")
+                dataset_rowids[dataset_rowid] = None
+
+        async with info.context.db() as session:
+            dataset_label_orm = models.DatasetLabel(name=name, description=description, color=color)
+            session.add(dataset_label_orm)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"A dataset label named '{name}' already exists")
+            except sqlalchemy.exc.StatementError as error:
+                raise BadRequest(str(error.orig))
+
+            datasets_by_id: dict[int, models.Dataset] = {}
+            if dataset_rowids:
+                datasets_by_id = {
+                    dataset.id: dataset
+                    for dataset in await session.scalars(
+                        select(models.Dataset).where(models.Dataset.id.in_(dataset_rowids.keys()))
+                    )
+                }
+                if len(datasets_by_id) < len(dataset_rowids):
+                    raise NotFound("One or more datasets not found")
+                session.add_all(
+                    [
+                        models.DatasetsDatasetLabel(
+                            dataset_id=dataset_rowid,
+                            dataset_label_id=dataset_label_orm.id,
+                        )
+                        for dataset_rowid in dataset_rowids
+                    ]
+                )
+                await session.commit()
+
+        return CreateDatasetLabelMutationPayload(
+            dataset_label=DatasetLabel(id=dataset_label_orm.id, db_record=dataset_label_orm),
+            datasets=[
+                Dataset(
+                    id=datasets_by_id[dataset_rowid].id, db_record=datasets_by_id[dataset_rowid]
+                )
+                for dataset_rowid in dataset_rowids
+            ],
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
+    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=[
+                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
+            ]
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
+    async def set_dataset_labels(
+        self, info: Info[Context, None], input: SetDatasetLabelsInput
+    ) -> SetDatasetLabelsMutationPayload:
+        try:
+            dataset_id = from_global_id_with_expected_type(input.dataset_id, Dataset.__name__)
+        except ValueError:
+            raise BadRequest(f"Invalid dataset ID: {input.dataset_id}")
+
+        dataset_label_ids: dict[
+            int, None
+        ] = {}  # use dictionary to de-duplicate while preserving order
+        for dataset_label_gid in input.dataset_label_ids:
+            try:
+                dataset_label_id = from_global_id_with_expected_type(
+                    dataset_label_gid, DatasetLabel.__name__
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid dataset label ID: {dataset_label_gid}")
+            dataset_label_ids[dataset_label_id] = None
+
+        async with info.context.db() as session:
+            dataset = await session.scalar(
+                select(models.Dataset)
+                .where(models.Dataset.id == dataset_id)
+                .options(joinedload(models.Dataset.datasets_dataset_labels))
+            )
+
+            if not dataset:
+                raise NotFound(f"Dataset with ID {input.dataset_id} not found")
+
+            existing_label_ids = (
+                await session.scalars(
+                    select(models.DatasetLabel.id).where(
+                        models.DatasetLabel.id.in_(dataset_label_ids.keys())
+                    )
+                )
+            ).all()
+            if len(existing_label_ids) != len(dataset_label_ids):
+                raise NotFound("One or more dataset labels not found")
+
+            previously_applied_dataset_label_ids = {
+                dataset_dataset_label.dataset_label_id
+                for dataset_dataset_label in dataset.datasets_dataset_labels
+            }
+
+            datasets_dataset_labels_to_add = [
+                models.DatasetsDatasetLabel(
+                    dataset_id=dataset_id,
+                    dataset_label_id=dataset_label_id,
+                )
+                for dataset_label_id in dataset_label_ids
+                if dataset_label_id not in previously_applied_dataset_label_ids
+            ]
+            if datasets_dataset_labels_to_add:
+                session.add_all(datasets_dataset_labels_to_add)
+                await session.flush()
+
+            datasets_dataset_labels_to_delete = [
+                dataset_dataset_label
+                for dataset_dataset_label in dataset.datasets_dataset_labels
+                if dataset_dataset_label.dataset_label_id not in dataset_label_ids
+            ]
+            if datasets_dataset_labels_to_delete:
+                await session.execute(
+                    delete(models.DatasetsDatasetLabel).where(
+                        tuple_(
+                            models.DatasetsDatasetLabel.dataset_id,
+                            models.DatasetsDatasetLabel.dataset_label_id,
+                        ).in_(
+                            [
+                                (
+                                    datasets_dataset_labels.dataset_id,
+                                    datasets_dataset_labels.dataset_label_id,
+                                )
+                                for datasets_dataset_labels in datasets_dataset_labels_to_delete
+                            ]
+                        )
+                    )
+                )
+
+        return SetDatasetLabelsMutationPayload(
+            dataset=Dataset(id=dataset.id, db_record=dataset),
+            query=Query(),
+        )