arize-phoenix 3.16.1__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/{types/ExportEventsMutation.py → mutations/export_events_mutations.py}

@@ -1,23 +1,25 @@
 import asyncio
 from collections import defaultdict
 from datetime import datetime
-from typing import Dict, List, Optional, Tuple
+from typing import Optional
 
 import strawberry
 from strawberry import ID, UNSET
 from strawberry.types import Info
 
 import phoenix.core.model_schema as ms
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.ClusterInput import ClusterInput
-from phoenix.server.api.types.DatasetRole import AncillaryDatasetRole, DatasetRole
-from phoenix.server.api.types.Event import parse_event_ids_by_dataset_role, unpack_event_id
+from phoenix.server.api.types.Event import parse_event_ids_by_inferences_role, unpack_event_id
 from phoenix.server.api.types.ExportedFile import ExportedFile
+from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
 
 
 @strawberry.type
-class ExportEventsMutation:
+class ExportEventsMutationMixin:
     @strawberry.mutation(
+        permission_classes=[IsNotReadOnly],
         description=(
             "Given a list of event ids, export the corresponding data subset in Parquet format."
             " File name is optional, but if specified, should be without file extension. By default"
@@ -27,16 +29,16 @@ class ExportEventsMutation:
     async def export_events(
         self,
         info: Info[Context, None],
-        event_ids: List[ID],
+        event_ids: list[ID],
         file_name: Optional[str] = UNSET,
     ) -> ExportedFile:
         if not isinstance(file_name, str):
             file_name = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-        row_ids = parse_event_ids_by_dataset_role(event_ids)
+        row_ids = parse_event_ids_by_inferences_role(event_ids)
         exclude_corpus_row_ids = {}
-        for dataset_role in list(row_ids.keys()):
-            if isinstance(dataset_role, DatasetRole):
-                exclude_corpus_row_ids[dataset_role.value] = row_ids[dataset_role]
+        for inferences_role in list(row_ids.keys()):
+            if isinstance(inferences_role, InferencesRole):
+                exclude_corpus_row_ids[inferences_role.value] = row_ids[inferences_role]
         path = info.context.export_path
         with open(path / (file_name + ".parquet"), "wb") as fd:
             loop = asyncio.get_running_loop()
@@ -49,6 +51,7 @@ class ExportEventsMutation:
         return ExportedFile(file_name=file_name)
 
     @strawberry.mutation(
+        permission_classes=[IsNotReadOnly],
         description=(
             "Given a list of clusters, export the corresponding data subset in Parquet format."
             " File name is optional, but if specified, should be without file extension. By default"
@@ -58,7 +61,7 @@ class ExportEventsMutation:
     async def export_clusters(
         self,
         info: Info[Context, None],
-        clusters: List[ClusterInput],
+        clusters: list[ClusterInput],
         file_name: Optional[str] = UNSET,
     ) -> ExportedFile:
         if not isinstance(file_name, str):
@@ -78,14 +81,14 @@ class ExportEventsMutation:
 
 
 def _unpack_clusters(
-    clusters: List[ClusterInput],
-) -> Tuple[Dict[ms.DatasetRole, List[int]], Dict[ms.DatasetRole, Dict[int, str]]]:
-    row_numbers: Dict[ms.DatasetRole, List[int]] = defaultdict(list)
-    cluster_ids: Dict[ms.DatasetRole, Dict[int, str]] = defaultdict(dict)
+    clusters: list[ClusterInput],
+) -> tuple[dict[ms.InferencesRole, list[int]], dict[ms.InferencesRole, dict[int, str]]]:
+    row_numbers: dict[ms.InferencesRole, list[int]] = defaultdict(list)
+    cluster_ids: dict[ms.InferencesRole, dict[int, str]] = defaultdict(dict)
     for i, cluster in enumerate(clusters):
-        for row_number, dataset_role in map(unpack_event_id, cluster.event_ids):
-            if isinstance(dataset_role, AncillaryDatasetRole):
+        for row_number, inferences_role in map(unpack_event_id, cluster.event_ids):
+            if isinstance(inferences_role, AncillaryInferencesRole):
                 continue
-            row_numbers[dataset_role.value].append(row_number)
-            cluster_ids[dataset_role.value][row_number] = cluster.id or str(i)
+            row_numbers[inferences_role.value].append(row_number)
+            cluster_ids[inferences_role.value][row_number] = cluster.id or str(i)
     return row_numbers, cluster_ids

phoenix/server/api/mutations/project_mutations.py

@@ -0,0 +1,57 @@
+import strawberry
+from sqlalchemy import delete, select
+from sqlalchemy.orm import load_only
+from strawberry.relay import GlobalID
+from strawberry.types import Info
+
+from phoenix.config import DEFAULT_PROJECT_NAME
+from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
+from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.ClearProjectInput import ClearProjectInput
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.dml_event import ProjectDeleteEvent, SpanDeleteEvent
+
+
+@strawberry.type
+class ProjectMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
+        project_id = from_global_id_with_expected_type(global_id=id, expected_type_name="Project")
+        async with info.context.db() as session:
+            project = await session.scalar(
+                select(models.Project)
+                .where(models.Project.id == project_id)
+                .options(load_only(models.Project.name))
+            )
+            if project is None:
+                raise ValueError(f"Unknown project: {id}")
+            if project.name == DEFAULT_PROJECT_NAME:
+                raise ValueError(f"Cannot delete the {DEFAULT_PROJECT_NAME} project")
+            await session.delete(project)
+        info.context.event_queue.put(ProjectDeleteEvent((project_id,)))
+        return Query()
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def clear_project(self, info: Info[Context, None], input: ClearProjectInput) -> Query:
+        project_id = from_global_id_with_expected_type(
+            global_id=input.id, expected_type_name="Project"
+        )
+        delete_statement = (
+            delete(models.Trace)
+            .where(models.Trace.project_rowid == project_id)
+            .returning(models.Trace.project_session_rowid)
+        )
+        if input.end_time:
+            delete_statement = delete_statement.where(models.Trace.start_time < input.end_time)
+        async with info.context.db() as session:
+            deleted_trace_project_session_ids = await session.scalars(delete_statement)
+            if deleted_trace_project_session_ids:
+                await session.execute(
+                    delete(models.ProjectSession).where(
+                        models.ProjectSession.id.in_(set(deleted_trace_project_session_ids))
+                    )
+                )
+        info.context.event_queue.put(SpanDeleteEvent((project_id,)))
+        return Query()

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -0,0 +1,128 @@
+from collections.abc import Sequence
+
+import strawberry
+from sqlalchemy import delete, insert, update
+from strawberry import UNSET
+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.input_types.CreateSpanAnnotationInput import CreateSpanAnnotationInput
+from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
+from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.dml_event import SpanAnnotationDeleteEvent, SpanAnnotationInsertEvent
+
+
+@strawberry.type
+class SpanAnnotationMutationPayload:
+    span_annotations: list[SpanAnnotation]
+    query: Query
+
+
+@strawberry.type
+class SpanAnnotationMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_span_annotations(
+        self, info: Info[Context, None], input: list[CreateSpanAnnotationInput]
+    ) -> SpanAnnotationMutationPayload:
+        inserted_annotations: Sequence[models.SpanAnnotation] = []
+        async with info.context.db() as session:
+            values_list = [
+                dict(
+                    span_rowid=from_global_id_with_expected_type(annotation.span_id, "Span"),
+                    name=annotation.name,
+                    label=annotation.label,
+                    score=annotation.score,
+                    explanation=annotation.explanation,
+                    annotator_kind=annotation.annotator_kind.value,
+                    metadata_=annotation.metadata,
+                )
+                for annotation in input
+            ]
+            stmt = (
+                insert(models.SpanAnnotation).values(values_list).returning(models.SpanAnnotation)
+            )
+            result = await session.scalars(stmt)
+            inserted_annotations = result.all()
+        if inserted_annotations:
+            info.context.event_queue.put(
+                SpanAnnotationInsertEvent(tuple(anno.id for anno in inserted_annotations))
+            )
+        return SpanAnnotationMutationPayload(
+            span_annotations=[
+                to_gql_span_annotation(annotation) for annotation in inserted_annotations
+            ],
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def patch_span_annotations(
+        self, info: Info[Context, None], input: list[PatchAnnotationInput]
+    ) -> SpanAnnotationMutationPayload:
+        patched_annotations = []
+        async with info.context.db() as session:
+            for annotation in input:
+                span_annotation_id = from_global_id_with_expected_type(
+                    annotation.annotation_id, "SpanAnnotation"
+                )
+                patch = {
+                    column.key: patch_value
+                    for column, patch_value, column_is_nullable in (
+                        (models.SpanAnnotation.name, annotation.name, False),
+                        (
+                            models.SpanAnnotation.annotator_kind,
+                            annotation.annotator_kind.value
+                            if annotation.annotator_kind is not None
+                            and annotation.annotator_kind is not UNSET
+                            else None,
+                            False,
+                        ),
+                        (models.SpanAnnotation.label, annotation.label, True),
+                        (models.SpanAnnotation.score, annotation.score, True),
+                        (models.SpanAnnotation.explanation, annotation.explanation, True),
+                        (models.SpanAnnotation.metadata_, annotation.metadata, False),
+                    )
+                    if patch_value is not UNSET and (patch_value is not None or column_is_nullable)
+                }
+                span_annotation = await session.scalar(
+                    update(models.SpanAnnotation)
+                    .where(models.SpanAnnotation.id == span_annotation_id)
+                    .values(**patch)
+                    .returning(models.SpanAnnotation)
+                )
+                if span_annotation is not None:
+                    patched_annotations.append(to_gql_span_annotation(span_annotation))
+                    info.context.event_queue.put(SpanAnnotationInsertEvent((span_annotation.id,)))
+        return SpanAnnotationMutationPayload(span_annotations=patched_annotations, query=Query())
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_span_annotations(
+        self, info: Info[Context, None], input: DeleteAnnotationsInput
+    ) -> SpanAnnotationMutationPayload:
+        span_annotation_ids = [
+            from_global_id_with_expected_type(global_id, "SpanAnnotation")
+            for global_id in input.annotation_ids
+        ]
+        async with info.context.db() as session:
+            stmt = (
+                delete(models.SpanAnnotation)
+                .where(models.SpanAnnotation.id.in_(span_annotation_ids))
+                .returning(models.SpanAnnotation)
+            )
+            result = await session.scalars(stmt)
+            deleted_annotations = result.all()
+
+            deleted_annotations_gql = [
+                to_gql_span_annotation(annotation) for annotation in deleted_annotations
+            ]
+        if deleted_annotations:
+            info.context.event_queue.put(
+                SpanAnnotationDeleteEvent(tuple(anno.id for anno in deleted_annotations))
+            )
+        return SpanAnnotationMutationPayload(
+            span_annotations=deleted_annotations_gql, query=Query()
+        )

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -0,0 +1,127 @@
+from collections.abc import Sequence
+
+import strawberry
+from sqlalchemy import delete, insert, update
+from strawberry import UNSET
+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.input_types.CreateTraceAnnotationInput import CreateTraceAnnotationInput
+from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
+from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
+from phoenix.server.dml_event import TraceAnnotationDeleteEvent, TraceAnnotationInsertEvent
+
+
+@strawberry.type
+class TraceAnnotationMutationPayload:
+    trace_annotations: list[TraceAnnotation]
+    query: Query
+
+
+@strawberry.type
+class TraceAnnotationMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def create_trace_annotations(
+        self, info: Info[Context, None], input: list[CreateTraceAnnotationInput]
+    ) -> TraceAnnotationMutationPayload:
+        inserted_annotations: Sequence[models.TraceAnnotation] = []
+        async with info.context.db() as session:
+            values_list = [
+                dict(
+                    trace_rowid=from_global_id_with_expected_type(annotation.trace_id, "Trace"),
+                    name=annotation.name,
+                    label=annotation.label,
+                    score=annotation.score,
+                    explanation=annotation.explanation,
+                    annotator_kind=annotation.annotator_kind.value,
+                    metadata_=annotation.metadata,
+                )
+                for annotation in input
+            ]
+            stmt = (
+                insert(models.TraceAnnotation).values(values_list).returning(models.TraceAnnotation)
+            )
+            result = await session.scalars(stmt)
+            inserted_annotations = result.all()
+        if inserted_annotations:
+            info.context.event_queue.put(
+                TraceAnnotationInsertEvent(tuple(anno.id for anno in inserted_annotations))
+            )
+        return TraceAnnotationMutationPayload(
+            trace_annotations=[
+                to_gql_trace_annotation(annotation) for annotation in inserted_annotations
+            ],
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    async def patch_trace_annotations(
+        self, info: Info[Context, None], input: list[PatchAnnotationInput]
+    ) -> TraceAnnotationMutationPayload:
+        patched_annotations = []
+        async with info.context.db() as session:
+            for annotation in input:
+                trace_annotation_id = from_global_id_with_expected_type(
+                    annotation.annotation_id, "TraceAnnotation"
+                )
+                patch = {
+                    column.key: patch_value
+                    for column, patch_value, column_is_nullable in (
+                        (models.TraceAnnotation.name, annotation.name, False),
+                        (
+                            models.TraceAnnotation.annotator_kind,
+                            annotation.annotator_kind.value
+                            if annotation.annotator_kind is not None
+                            else None,
+                            False,
+                        ),
+                        (models.TraceAnnotation.label, annotation.label, True),
+                        (models.TraceAnnotation.score, annotation.score, True),
+                        (models.TraceAnnotation.explanation, annotation.explanation, True),
+                        (models.TraceAnnotation.metadata_, annotation.metadata, False),
+                    )
+                    if patch_value is not UNSET and (patch_value is not None or column_is_nullable)
+                }
+                trace_annotation = await session.scalar(
+                    update(models.TraceAnnotation)
+                    .where(models.TraceAnnotation.id == trace_annotation_id)
+                    .values(**patch)
+                    .returning(models.TraceAnnotation)
+                )
+                if trace_annotation:
+                    patched_annotations.append(to_gql_trace_annotation(trace_annotation))
+                    info.context.event_queue.put(TraceAnnotationInsertEvent((trace_annotation.id,)))
+        return TraceAnnotationMutationPayload(trace_annotations=patched_annotations, query=Query())
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_trace_annotations(
+        self, info: Info[Context, None], input: DeleteAnnotationsInput
+    ) -> TraceAnnotationMutationPayload:
+        trace_annotation_ids = [
+            from_global_id_with_expected_type(global_id, "TraceAnnotation")
+            for global_id in input.annotation_ids
+        ]
+        async with info.context.db() as session:
+            stmt = (
+                delete(models.TraceAnnotation)
+                .where(models.TraceAnnotation.id.in_(trace_annotation_ids))
+                .returning(models.TraceAnnotation)
+            )
+            result = await session.scalars(stmt)
+            deleted_annotations = result.all()
+
+            deleted_annotations_gql = [
+                to_gql_trace_annotation(annotation) for annotation in deleted_annotations
+            ]
+        if deleted_annotations:
+            info.context.event_queue.put(
+                TraceAnnotationDeleteEvent(tuple(anno.id for anno in deleted_annotations))
+            )
+        return TraceAnnotationMutationPayload(
+            trace_annotations=deleted_annotations_gql, query=Query()
+        )