arize-phoenix 8.32.0__py3-none-any.whl → 9.0.0__py3-none-any.whl

phoenix/server/api/mutations/annotation_config_mutations.py

@@ -0,0 +1,413 @@
+from typing import Optional
+
+import strawberry
+from sqlalchemy import delete, select, tuple_
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry.relay.types import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.db.types.annotation_configs import (
+    AnnotationConfigType,
+    AnnotationType,
+    CategoricalAnnotationValue,
+    OptimizationDirection,
+)
+from phoenix.db.types.annotation_configs import (
+    CategoricalAnnotationConfig as CategoricalAnnotationConfigModel,
+)
+from phoenix.db.types.annotation_configs import (
+    ContinuousAnnotationConfig as ContinuousAnnotationConfigModel,
+)
+from phoenix.db.types.annotation_configs import (
+    FreeformAnnotationConfig as FreeformAnnotationConfigModel,
+)
+from phoenix.server.api.auth import IsNotReadOnly
+from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.AnnotationConfig import (
+    AnnotationConfig,
+    CategoricalAnnotationConfig,
+    ContinuousAnnotationConfig,
+    FreeformAnnotationConfig,
+    to_gql_annotation_config,
+)
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.Project import Project
+
+ANNOTATION_TYPE_NAMES = (
+    CategoricalAnnotationConfig.__name__,
+    ContinuousAnnotationConfig.__name__,
+    FreeformAnnotationConfig.__name__,
+)
+
+
+@strawberry.input
+class CategoricalAnnotationConfigValueInput:
+    label: str
+    score: Optional[float] = None
+
+
+@strawberry.input
+class CategoricalAnnotationConfigInput:
+    name: str
+    description: Optional[str] = None
+    optimization_direction: OptimizationDirection
+    values: list[CategoricalAnnotationConfigValueInput]
+
+
+@strawberry.input
+class ContinuousAnnotationConfigInput:
+    name: str
+    description: Optional[str] = None
+    optimization_direction: OptimizationDirection
+    lower_bound: Optional[float] = None
+    upper_bound: Optional[float] = None
+
+
+@strawberry.input
+class FreeformAnnotationConfigInput:
+    name: str
+    description: Optional[str] = None
+
+
+@strawberry.input(one_of=True)
+class AnnotationConfigInput:
+    categorical: Optional[CategoricalAnnotationConfigInput] = strawberry.UNSET
+    continuous: Optional[ContinuousAnnotationConfigInput] = strawberry.UNSET
+    freeform: Optional[FreeformAnnotationConfigInput] = strawberry.UNSET
+
+    def __post_init__(self) -> None:
+        if (
+            sum(
+                [
+                    self.categorical is not strawberry.UNSET,
+                    self.continuous is not strawberry.UNSET,
+                    self.freeform is not strawberry.UNSET,
+                ]
+            )
+            != 1
+        ):
+            raise BadRequest("Exactly one of categorical, continuous, or freeform must be set")
+
+
+@strawberry.input
+class CreateAnnotationConfigInput:
+    annotation_config: AnnotationConfigInput
+
+
+@strawberry.type
+class CreateAnnotationConfigPayload:
+    query: Query
+    annotation_config: AnnotationConfig
+
+
+@strawberry.input
+class UpdateAnnotationConfigInput:
+    id: GlobalID
+    annotation_config: AnnotationConfigInput
+
+
+@strawberry.type
+class UpdateAnnotationConfigPayload:
+    query: Query
+    annotation_config: AnnotationConfig
+
+
+@strawberry.input
+class DeleteAnnotationConfigsInput:
+    ids: list[GlobalID]
+
+
+@strawberry.type
+class DeleteAnnotationConfigsPayload:
+    query: Query
+    annotation_configs: list[AnnotationConfig]
+
+
+@strawberry.input
+class AddAnnotationConfigToProjectInput:
+    project_id: GlobalID
+    annotation_config_id: GlobalID
+
+
+@strawberry.type
+class AddAnnotationConfigToProjectPayload:
+    query: Query
+    project: Project
+
+
+@strawberry.input
+class RemoveAnnotationConfigFromProjectInput:
+    project_id: GlobalID
+    annotation_config_id: GlobalID
+
+
+@strawberry.type
+class RemoveAnnotationConfigFromProjectPayload:
+    query: Query
+    project: Project
+
+
+def _to_pydantic_categorical_annotation_config(
+    input: CategoricalAnnotationConfigInput,
+) -> CategoricalAnnotationConfigModel:
+    try:
+        return CategoricalAnnotationConfigModel(
+            type=AnnotationType.CATEGORICAL.value,
+            description=input.description,
+            optimization_direction=input.optimization_direction,
+            values=[
+                CategoricalAnnotationValue(label=value.label, score=value.score)
+                for value in input.values
+            ],
+        )
+    except ValueError as error:
+        raise BadRequest(str(error))
+
+
+def _to_pydantic_continuous_annotation_config(
+    input: ContinuousAnnotationConfigInput,
+) -> ContinuousAnnotationConfigModel:
+    try:
+        return ContinuousAnnotationConfigModel(
+            type=AnnotationType.CONTINUOUS.value,
+            description=input.description,
+            optimization_direction=input.optimization_direction,
+            lower_bound=input.lower_bound,
+            upper_bound=input.upper_bound,
+        )
+    except ValueError as error:
+        raise BadRequest(str(error))
+
+
+def _to_pydantic_freeform_annotation_config(
+    input: FreeformAnnotationConfigInput,
+) -> FreeformAnnotationConfigModel:
+    try:
+        return FreeformAnnotationConfigModel(
+            type=AnnotationType.FREEFORM.value,
+            description=input.description,
+        )
+    except ValueError as error:
+        raise BadRequest(str(error))
+
+
+@strawberry.type
+class AnnotationConfigMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    async def create_annotation_config(
+        self,
+        info: Info[Context, None],
+        input: CreateAnnotationConfigInput,
+    ) -> CreateAnnotationConfigPayload:
+        input_annotation_config = input.annotation_config
+        config: AnnotationConfigType
+        name: str
+        if categorical_input := input_annotation_config.categorical:
+            name = categorical_input.name
+            config = _to_pydantic_categorical_annotation_config(categorical_input)
+        elif continuous_input := input_annotation_config.continuous:
+            name = input_annotation_config.continuous.name
+            config = _to_pydantic_continuous_annotation_config(continuous_input)
+        elif freeform_input := input_annotation_config.freeform:
+            name = freeform_input.name
+            config = _to_pydantic_freeform_annotation_config(freeform_input)
+        else:
+            raise BadRequest("No annotation config provided")
+
+        if name == "note":
+            raise BadRequest("The name 'note' is reserved for span notes")
+
+        async with info.context.db() as session:
+            annotation_config = models.AnnotationConfig(
+                name=name,
+                config=config,
+            )
+            session.add(annotation_config)
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Annotation configuration with name '{name}' already exists")
+        return CreateAnnotationConfigPayload(
+            query=Query(),
+            annotation_config=to_gql_annotation_config(annotation_config),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    async def update_annotation_config(
+        self,
+        info: Info[Context, None],
+        input: UpdateAnnotationConfigInput,
+    ) -> UpdateAnnotationConfigPayload:
+        try:
+            config_id = int(input.id.node_id)
+        except ValueError:
+            raise BadRequest("Invalid annotation config ID")
+
+        if input.id.type_name not in ANNOTATION_TYPE_NAMES:
+            raise BadRequest("Invalid annotation config ID")
+
+        input_annotation_config = input.annotation_config
+        config: AnnotationConfigType
+        name: str
+        if categorical_input := input_annotation_config.categorical:
+            name = categorical_input.name
+            config = _to_pydantic_categorical_annotation_config(categorical_input)
+        elif continuous_input := input_annotation_config.continuous:
+            name = input_annotation_config.continuous.name
+            config = _to_pydantic_continuous_annotation_config(continuous_input)
+        elif freeform_input := input_annotation_config.freeform:
+            name = freeform_input.name
+            config = _to_pydantic_freeform_annotation_config(freeform_input)
+        else:
+            raise BadRequest("No annotation config provided")
+
+        if name == "note":
+            raise BadRequest("The name 'note' is reserved for span notes")
+
+        async with info.context.db() as session:
+            annotation_config = await session.get(models.AnnotationConfig, config_id)
+            if not annotation_config:
+                raise NotFound("Annotation config not found")
+
+            annotation_config.name = name
+            annotation_config.config = config
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Annotation configuration with name '{name}' already exists")
+
+        return UpdateAnnotationConfigPayload(
+            query=Query(),
+            annotation_config=to_gql_annotation_config(annotation_config),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    async def delete_annotation_configs(
+        self,
+        info: Info[Context, None],
+        input: DeleteAnnotationConfigsInput,
+    ) -> DeleteAnnotationConfigsPayload:
+        config_ids = set()
+        for config_gid in input.ids:
+            if (type_name := config_gid.type_name) not in ANNOTATION_TYPE_NAMES:
+                raise BadRequest(f"Unexpected type name in Relay ID: {type_name}")
+            config_ids.add(int(config_gid.node_id))
+
+        async with info.context.db() as session:
+            result = await session.scalars(
+                delete(models.AnnotationConfig)
+                .where(models.AnnotationConfig.id.in_(config_ids))
+                .returning(models.AnnotationConfig)
+            )
+            deleted_annotation_configs = result.all()
+            if len(deleted_annotation_configs) < len(config_ids):
+                await session.rollback()
+                raise NotFound(
+                    "Could not find one or more annotation configs to delete, deletion aborted."
+                )
+        return DeleteAnnotationConfigsPayload(
+            query=Query(),
+            annotation_configs=[
+                to_gql_annotation_config(annotation_config)
+                for annotation_config in deleted_annotation_configs
+            ],
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    async def add_annotation_config_to_project(
+        self,
+        info: Info[Context, None],
+        input: list[AddAnnotationConfigToProjectInput],
+    ) -> AddAnnotationConfigToProjectPayload:
+        if not input:
+            raise BadRequest("No project annotation config associations provided")
+        project_annotation_config_ids: set[tuple[int, int]] = set()
+        for item in input:
+            project_id = from_global_id_with_expected_type(
+                global_id=item.project_id, expected_type_name="Project"
+            )
+            if (item.annotation_config_id.type_name) not in ANNOTATION_TYPE_NAMES:
+                raise BadRequest(
+                    f"Invalidation ID for annotation config: {str(item.annotation_config_id)}"
+                )
+            annotation_config_id = int(item.annotation_config_id.node_id)
+            project_annotation_config_ids.add((project_id, annotation_config_id))
+        project_ids = [project_id for project_id, _ in project_annotation_config_ids]
+        annotation_config_ids = [
+            annotation_config_id for _, annotation_config_id in project_annotation_config_ids
+        ]
+
+        async with info.context.db() as session:
+            result = await session.scalars(
+                select(models.Project.id).where(models.Project.id.in_(project_ids))
+            )
+            resolved_project_ids = result.all()
+            if set(project_ids) - set(resolved_project_ids):
+                raise NotFound("One or more projects were not found")
+
+            result = await session.scalars(
+                select(models.AnnotationConfig.id).where(
+                    models.AnnotationConfig.id.in_(annotation_config_ids)
+                )
+            )
+            resolved_annotation_config_ids = result.all()
+            if set(annotation_config_ids) - set(resolved_annotation_config_ids):
+                raise NotFound("One or more annotation configs were not found")
+
+            for project_id, annotation_config_id in project_annotation_config_ids:
+                project_annotation_config = models.ProjectAnnotationConfig(
+                    project_id=project_id,
+                    annotation_config_id=annotation_config_id,
+                )
+                session.add(project_annotation_config)
+
+            try:
+                await session.commit()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                await session.rollback()
+                raise Conflict(
+                    "One or more annotation configs have already been added to the project"
+                )
+            return AddAnnotationConfigToProjectPayload(
+                query=Query(),
+                project=Project(project_rowid=project_id),
+            )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore[misc]
+    async def remove_annotation_config_from_project(
+        self,
+        info: Info[Context, None],
+        input: list[RemoveAnnotationConfigFromProjectInput],
+    ) -> RemoveAnnotationConfigFromProjectPayload:
+        project_annotation_config_associations = set()
+        for item in input:
+            project_id = from_global_id_with_expected_type(
+                global_id=item.project_id, expected_type_name="Project"
+            )
+            if (type_name := item.annotation_config_id.type_name) not in ANNOTATION_TYPE_NAMES:
+                raise BadRequest(f"Unexpected type name in Relay ID: {type_name}")
+            annotation_config_id = int(item.annotation_config_id.node_id)
+            project_annotation_config_associations.add((project_id, annotation_config_id))
+        async with info.context.db() as session:
+            result = await session.scalars(
+                delete(models.ProjectAnnotationConfig)
+                .where(
+                    tuple_(
+                        models.ProjectAnnotationConfig.project_id,
+                        models.ProjectAnnotationConfig.annotation_config_id,
+                    ).in_(project_annotation_config_associations)
+                )
+                .returning(models.ProjectAnnotationConfig)
+            )
+            annotation_configs = result.all()
+            if len(annotation_configs) < len(project_annotation_config_associations):
+                await session.rollback()
+                raise NotFound("Could not find one or more input project annotation configs")
+        return RemoveAnnotationConfigFromProjectPayload(
+            query=Query(),
+            project=Project(project_rowid=project_id),
+        )

phoenix/server/api/mutations/dataset_mutations.py

@@ -11,7 +11,9 @@ from openinference.semconv.trace import (
     ToolCallAttributes,
 )
 from sqlalchemy import and_, delete, distinct, func, insert, select, update
+from sqlalchemy.orm import contains_eager
 from strawberry import UNSET
+from strawberry.relay.types import GlobalID
 from strawberry.types import Info
 
 from phoenix.db import models
@@ -130,43 +132,40 @@ class DatasetMutationMixin:
                 raise ValueError(
                     f"Unknown dataset: {dataset_id}"
                 )  # todo: implement error types https://github.com/Arize-ai/phoenix/issues/3221
-            dataset_version_rowid = await session.scalar(
-                insert(models.DatasetVersion)
-                .values(
-                    dataset_id=dataset_rowid,
-                    description=dataset_version_description,
-                    metadata_=dataset_version_metadata,
-                )
-                .returning(models.DatasetVersion.id)
+            dataset_version = models.DatasetVersion(
+                dataset_id=dataset_rowid,
+                description=dataset_version_description,
+                metadata_=dataset_version_metadata or {},
             )
+            session.add(dataset_version)
+            await session.flush()
             spans = (
-                await session.scalars(select(models.Span).where(models.Span.id.in_(span_rowids)))
-            ).all()
-            if missing_span_rowids := span_rowids - {span.id for span in spans}:
-                raise ValueError(
-                    f"Could not find spans with rowids: {', '.join(map(str, missing_span_rowids))}"
-                )  # todo: implement error handling types https://github.com/Arize-ai/phoenix/issues/3221
-
-            span_annotations = (
-                await session.scalars(
-                    select(models.SpanAnnotation).where(
-                        models.SpanAnnotation.span_rowid.in_(span_rowids)
+                (
+                    await session.scalars(
+                        select(models.Span)
+                        .outerjoin(
+                            models.SpanAnnotation,
+                            models.Span.id == models.SpanAnnotation.span_rowid,
+                        )
+                        .outerjoin(models.User, models.SpanAnnotation.user_id == models.User.id)
+                        .order_by(
+                            models.Span.id,
+                            models.SpanAnnotation.name,
+                            models.User.username,
+                        )
+                        .where(models.Span.id.in_(span_rowids))
+                        .options(
+                            contains_eager(models.Span.span_annotations).contains_eager(
+                                models.SpanAnnotation.user
+                            )
+                        )
                     )
                 )
-            ).all()
-
-            span_annotations_by_span: dict[int, dict[Any, Any]] = {span.id: {} for span in spans}
-            for annotation in span_annotations:
-                span_id = annotation.span_rowid
-                if span_id not in span_annotations_by_span:
-                    span_annotations_by_span[span_id] = dict()
-                span_annotations_by_span[span_id][annotation.name] = {
-                    "label": annotation.label,
-                    "score": annotation.score,
-                    "explanation": annotation.explanation,
-                    "metadata": annotation.metadata_,
-                    "annotator_kind": annotation.annotator_kind,
-                }
+                .unique()
+                .all()
+            )
+            if span_rowids - {span.id for span in spans}:
+                raise NotFound("Some spans could not be found")
 
             DatasetExample = models.DatasetExample
             dataset_example_rowids = (
@@ -201,7 +200,7 @@ class DatasetMutationMixin:
                 [
                     {
                         DatasetExampleRevision.dataset_example_id.key: dataset_example_rowid,
-                        DatasetExampleRevision.dataset_version_id.key: dataset_version_rowid,
+                        DatasetExampleRevision.dataset_version_id.key: dataset_version.id,
                         DatasetExampleRevision.input.key: get_dataset_example_input(span),
                         DatasetExampleRevision.output.key: get_dataset_example_output(span),
                         DatasetExampleRevision.metadata_.key: {
@@ -212,11 +211,7 @@ class DatasetMutationMixin:
                                 if k in nonprivate_span_attributes
                             },
                             "span_kind": span.span_kind,
-                            **(
-                                {"annotations": annotations}
-                                if (annotations := span_annotations_by_span[span.id])
-                                else {}
-                            ),
+                            "annotations": _gather_span_annotations_by_name(span.span_annotations),
                         },
                         DatasetExampleRevision.revision_kind.key: "CREATE",
                     }
@@ -602,6 +597,34 @@ def _to_orm_revision(
     }
 
 
+def _gather_span_annotations_by_name(
+    span_annotations: list[models.SpanAnnotation],
+) -> dict[str, list[dict[str, Any]]]:
+    span_annotations_by_name: dict[str, list[dict[str, Any]]] = {}
+    for span_annotation in span_annotations:
+        if span_annotation.name not in span_annotations_by_name:
+            span_annotations_by_name[span_annotation.name] = []
+        span_annotations_by_name[span_annotation.name].append(
+            _to_span_annotation_dict(span_annotation)
+        )
+    return span_annotations_by_name
+
+
+def _to_span_annotation_dict(span_annotation: models.SpanAnnotation) -> dict[str, Any]:
+    return {
+        "label": span_annotation.label,
+        "score": span_annotation.score,
+        "explanation": span_annotation.explanation,
+        "metadata": span_annotation.metadata_,
+        "annotator_kind": span_annotation.annotator_kind,
+        "user_id": str(GlobalID(models.User.__name__, str(user_id)))
+        if (user_id := span_annotation.user_id) is not None
+        else None,
+        "username": user.username if (user := span_annotation.user) is not None else None,
+        "email": user.email if user is not None else None,
+    }
+
+
 INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
 INPUT_VALUE = SpanAttributes.INPUT_VALUE
 OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE