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

phoenix/server/api/mutations/api_key_mutations.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import Optional
+from typing import Literal, Optional
 
 import strawberry
 from sqlalchemy import select
@@ -9,7 +9,7 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.db.models import UserRoleName
-from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly
+from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import Unauthorized
 from phoenix.server.api.queries import Query
@@ -61,7 +61,7 @@ class DeleteApiKeyMutationPayload:
 
 @strawberry.type
 class ApiKeyMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin, IsLocked])  # type: ignore
     async def create_system_api_key(
         self, info: Info[Context, None], input: CreateApiKeyInput
     ) -> CreateSystemApiKeyMutationPayload:
@@ -92,13 +92,7 @@ class ApiKeyMutationMixin:
         token, token_id = await token_store.create_api_key(claims)
         return CreateSystemApiKeyMutationPayload(
             jwt=token,
-            api_key=SystemApiKey(
-                id_attr=int(token_id),
-                name=input.name,
-                description=input.description or None,
-                created_at=issued_at,
-                expires_at=input.expires_at or None,
-            ),
+            api_key=SystemApiKey(id=int(token_id)),
             query=Query(),
         )
 
@@ -113,12 +107,20 @@ class ApiKeyMutationMixin:
         except AttributeError:
             raise ValueError("User not found")
         issued_at = datetime.now(timezone.utc)
+        # Determine user role for API key
+        user_role: Literal["ADMIN", "MEMBER", "VIEWER"]
+        if user.is_admin:
+            user_role = "ADMIN"
+        elif user.is_viewer:
+            user_role = "VIEWER"
+        else:
+            user_role = "MEMBER"
         claims = ApiKeyClaims(
             subject=user.identity,
             issued_at=issued_at,
             expiration_time=input.expires_at or None,
             attributes=ApiKeyAttributes(
-                user_role="ADMIN" if user.is_admin else "MEMBER",
+                user_role=user_role,
                 name=input.name,
                 description=input.description,
             ),
@@ -126,18 +128,11 @@ class ApiKeyMutationMixin:
         token, token_id = await token_store.create_api_key(claims)
         return CreateUserApiKeyMutationPayload(
             jwt=token,
-            api_key=UserApiKey(
-                id_attr=int(token_id),
-                name=input.name,
-                description=input.description or None,
-                created_at=issued_at,
-                expires_at=input.expires_at or None,
-                user_id=int(user.identity),
-            ),
+            api_key=UserApiKey(id=int(token_id)),
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin])  # type: ignore
     async def delete_system_api_key(
         self, info: Info[Context, None], input: DeleteApiKeyInput
     ) -> DeleteApiKeyMutationPayload:

phoenix/server/api/mutations/chat_mutations.py

@@ -4,7 +4,7 @@ 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
@@ -26,8 +26,11 @@ 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
@@ -46,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,
@@ -80,7 +84,7 @@ 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
@@ -96,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
@@ -126,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,
@@ -181,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,
@@ -200,14 +215,24 @@ class ChatCompletionMutationMixin:
                 repetitions=1,
                 metadata_=input.experiment_metadata or dict(),
                 project_name=project_name,
+                user_id=user_id,
             )
-            session.add(experiment)
-            await session.flush()
+            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[ChatCompletionMutationPayload, BaseException]] = []
+        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(
@@ -224,10 +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,
             )
@@ -239,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,
@@ -261,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)
 
@@ -272,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
@@ -316,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(
@@ -324,9 +391,10 @@ class ChatCompletionMutationMixin:
         info: Info[Context, None],
         llm_client: PlaygroundStreamingClient,
         input: ChatCompletionInput,
+        repetition_number: int,
         project_name: str = PLAYGROUND_PROJECT_NAME,
         project_description: str = "Traces from prompt playground",
-    ) -> ChatCompletionMutationPayload:
+    ) -> tuple[ChatCompletionRepetition, models.Span]:
         attributes: dict[str, Any] = {}
         attributes.update(dict(prompt_metadata(input.prompt_name)))
 
@@ -473,26 +541,27 @@ class ChatCompletionMutationMixin:
                 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(

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(),
+        )