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

phoenix/server/api/mutations/project_session_annotations_mutations.py

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

phoenix/server/api/mutations/project_trace_retention_policy_mutations.py

@@ -16,7 +16,7 @@ from phoenix.db.types.trace_retention import (
     TraceRetentionCronExpression,
     TraceRetentionRule,
 )
-from phoenix.server.api.auth import IsAdminIfAuthEnabled, IsLocked, IsNotReadOnly
+from phoenix.server.api.auth import IsAdminIfAuthEnabled, IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, NotFound
 from phoenix.server.api.queries import Query
@@ -113,7 +113,9 @@ class ProjectTraceRetentionPolicyMutationPayload:
 
 @strawberry.type
 class ProjectTraceRetentionPolicyMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdminIfAuthEnabled, IsLocked])  # type: ignore
+    @strawberry.mutation(
+        permission_classes=[IsNotReadOnly, IsNotViewer, IsAdminIfAuthEnabled, IsLocked]
+    )  # type: ignore
     async def create_project_trace_retention_policy(
         self,
         info: Info[Context, None],
@@ -146,7 +148,9 @@ class ProjectTraceRetentionPolicyMutationMixin:
             node=ProjectTraceRetentionPolicy(id=policy.id, db_policy=policy),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdminIfAuthEnabled, IsLocked])  # type: ignore
+    @strawberry.mutation(
+        permission_classes=[IsNotReadOnly, IsNotViewer, IsAdminIfAuthEnabled, IsLocked]
+    )  # type: ignore
     async def patch_project_trace_retention_policy(
         self,
         info: Info[Context, None],
@@ -204,7 +208,7 @@ class ProjectTraceRetentionPolicyMutationMixin:
             node=ProjectTraceRetentionPolicy(id=policy.id, db_policy=policy),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdminIfAuthEnabled])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdminIfAuthEnabled])  # type: ignore
     async def delete_project_trace_retention_policy(
         self,
         info: Info[Context, None],

phoenix/server/api/mutations/prompt_label_mutations.py

@@ -10,79 +10,89 @@ from strawberry.relay import GlobalID
 from strawberry.types import Info
 
 from phoenix.db import models
-from phoenix.db.types.identifier import Identifier as IdentifierModel
-from phoenix.server.api.auth import IsLocked, IsNotReadOnly
+from phoenix.server.api.auth import IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import Conflict, NotFound
 from phoenix.server.api.queries import Query
-from phoenix.server.api.types.Identifier import Identifier
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Prompt import Prompt
-from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
+from phoenix.server.api.types.PromptLabel import PromptLabel
 
 
 @strawberry.input
 class CreatePromptLabelInput:
-    name: Identifier
+    name: str
     description: Optional[str] = None
+    color: str
 
 
 @strawberry.input
 class PatchPromptLabelInput:
     prompt_label_id: GlobalID
-    name: Optional[Identifier] = None
+    name: Optional[str] = None
     description: Optional[str] = None
 
 
 @strawberry.input
-class DeletePromptLabelInput:
-    prompt_label_id: GlobalID
+class DeletePromptLabelsInput:
+    prompt_label_ids: list[GlobalID]
 
 
 @strawberry.input
-class SetPromptLabelInput:
+class SetPromptLabelsInput:
     prompt_id: GlobalID
-    prompt_label_id: GlobalID
+    prompt_label_ids: list[GlobalID]
 
 
 @strawberry.input
-class UnsetPromptLabelInput:
+class UnsetPromptLabelsInput:
     prompt_id: GlobalID
-    prompt_label_id: GlobalID
+    prompt_label_ids: list[GlobalID]
 
 
 @strawberry.type
 class PromptLabelMutationPayload:
-    prompt_label: Optional["PromptLabel"]
+    prompt_labels: list["PromptLabel"]
+    query: "Query"
+
+
+@strawberry.type
+class PromptLabelDeleteMutationPayload:
+    deleted_prompt_label_ids: list["GlobalID"]
+    query: "Query"
+
+
+@strawberry.type
+class PromptLabelAssociationMutationPayload:
     query: "Query"
 
 
 @strawberry.type
 class PromptLabelMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def create_prompt_label(
         self, info: Info[Context, None], input: CreatePromptLabelInput
     ) -> PromptLabelMutationPayload:
         async with info.context.db() as session:
-            name = IdentifierModel.model_validate(str(input.name))
-            label_orm = models.PromptLabel(name=name, description=input.description)
+            label_orm = models.PromptLabel(
+                name=input.name, description=input.description, color=input.color
+            )
             session.add(label_orm)
 
             try:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
-                raise Conflict(f"A prompt label named '{name}' already exists.")
+                raise Conflict(f"A prompt label named '{input.name}' already exists.")
 
             return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm),
+                prompt_labels=[PromptLabel(id=label_orm.id, db_record=label_orm)],
                 query=Query(),
             )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_prompt_label(
         self, info: Info[Context, None], input: PatchPromptLabelInput
     ) -> PromptLabelMutationPayload:
-        validated_name = IdentifierModel.model_validate(str(input.name)) if input.name else None
         async with info.context.db() as session:
             label_id = from_global_id_with_expected_type(
                 input.prompt_label_id, PromptLabel.__name__
@@ -92,8 +102,8 @@ class PromptLabelMutationMixin:
             if not label_orm:
                 raise NotFound(f"PromptLabel with ID {input.prompt_label_id} not found")
 
-            if validated_name is not None:
-                label_orm.name = validated_name.root
+            if input.name is not None:
+                label_orm.name = input.name
             if input.description is not None:
                 label_orm.description = input.description
 
@@ -103,46 +113,48 @@ class PromptLabelMutationMixin:
                 raise Conflict("Error patching PromptLabel. Possibly a name conflict?")
 
             return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm),
+                prompt_labels=[PromptLabel(id=label_orm.id, db_record=label_orm)],
                 query=Query(),
             )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
-    async def delete_prompt_label(
-        self, info: Info[Context, None], input: DeletePromptLabelInput
-    ) -> PromptLabelMutationPayload:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def delete_prompt_labels(
+        self, info: Info[Context, None], input: DeletePromptLabelsInput
+    ) -> PromptLabelDeleteMutationPayload:
         """
         Deletes a PromptLabel (and any crosswalk references).
         """
         async with info.context.db() as session:
-            label_id = from_global_id_with_expected_type(
-                input.prompt_label_id, PromptLabel.__name__
-            )
-            stmt = delete(models.PromptLabel).where(models.PromptLabel.id == label_id)
-            result = await session.execute(stmt)
-
-            if result.rowcount == 0:
-                raise NotFound(f"PromptLabel with ID {input.prompt_label_id} not found")
+            label_ids = [
+                from_global_id_with_expected_type(prompt_label_id, PromptLabel.__name__)
+                for prompt_label_id in input.prompt_label_ids
+            ]
+            stmt = delete(models.PromptLabel).where(models.PromptLabel.id.in_(label_ids))
+            await session.execute(stmt)
 
             await session.commit()
 
-            return PromptLabelMutationPayload(
-                prompt_label=None,
+            return PromptLabelDeleteMutationPayload(
+                deleted_prompt_label_ids=input.prompt_label_ids,
                 query=Query(),
             )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
-    async def set_prompt_label(
-        self, info: Info[Context, None], input: SetPromptLabelInput
-    ) -> PromptLabelMutationPayload:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
+    async def set_prompt_labels(
+        self, info: Info[Context, None], input: SetPromptLabelsInput
+    ) -> PromptLabelAssociationMutationPayload:
         async with info.context.db() as session:
             prompt_id = from_global_id_with_expected_type(input.prompt_id, Prompt.__name__)
-            label_id = from_global_id_with_expected_type(
-                input.prompt_label_id, PromptLabel.__name__
-            )
+            label_ids = [
+                from_global_id_with_expected_type(prompt_label_id, PromptLabel.__name__)
+                for prompt_label_id in input.prompt_label_ids
+            ]
 
-            crosswalk = models.PromptPromptLabel(prompt_id=prompt_id, prompt_label_id=label_id)
-            session.add(crosswalk)
+            crosswalk_items = [
+                models.PromptPromptLabel(prompt_id=prompt_id, prompt_label_id=label_id)
+                for label_id in label_ids
+            ]
+            session.add_all(crosswalk_items)
 
             try:
                 await session.commit()
@@ -152,41 +164,38 @@ class PromptLabelMutationMixin:
                 # - Foreign key violation => prompt_id or label_id doesn't exist
                 raise Conflict("Failed to associate PromptLabel with Prompt.") from e
 
-            label_orm = await session.get(models.PromptLabel, label_id)
-            if not label_orm:
-                raise NotFound(f"PromptLabel with ID {input.prompt_label_id} not found")
-
-            return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm),
+            return PromptLabelAssociationMutationPayload(
                 query=Query(),
             )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
-    async def unset_prompt_label(
-        self, info: Info[Context, None], input: UnsetPromptLabelInput
-    ) -> PromptLabelMutationPayload:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def unset_prompt_labels(
+        self, info: Info[Context, None], input: UnsetPromptLabelsInput
+    ) -> PromptLabelAssociationMutationPayload:
         """
         Unsets a PromptLabel from a Prompt by removing the row in the crosswalk.
         """
         async with info.context.db() as session:
             prompt_id = from_global_id_with_expected_type(input.prompt_id, Prompt.__name__)
-            label_id = from_global_id_with_expected_type(
-                input.prompt_label_id, PromptLabel.__name__
-            )
+            label_ids = [
+                from_global_id_with_expected_type(prompt_label_id, PromptLabel.__name__)
+                for prompt_label_id in input.prompt_label_ids
+            ]
 
             stmt = delete(models.PromptPromptLabel).where(
                 (models.PromptPromptLabel.prompt_id == prompt_id)
-                & (models.PromptPromptLabel.prompt_label_id == label_id)
+                & (models.PromptPromptLabel.prompt_label_id.in_(label_ids))
             )
             result = await session.execute(stmt)
 
-            if result.rowcount == 0:
-                raise NotFound(f"No association between prompt={prompt_id} and label={label_id}.")
+            if result.rowcount != len(label_ids):  # type: ignore[attr-defined]
+                label_ids_str = ", ".join(str(i) for i in label_ids)
+                raise NotFound(
+                    f"No association between prompt={prompt_id} and labels={label_ids_str}."
+                )
 
             await session.commit()
 
-            label_orm = await session.get(models.PromptLabel, label_id)
-            return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm) if label_orm else None,
+            return PromptLabelAssociationMutationPayload(
                 query=Query(),
             )