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

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -7,7 +7,7 @@ from starlette.requests import Request
 from strawberry import UNSET, Info
 
 from phoenix.db import models
-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 BadRequest, NotFound, Unauthorized
 from phoenix.server.api.helpers.annotations import get_user_identifier
@@ -21,7 +21,7 @@ from phoenix.server.api.queries import Query
 from phoenix.server.api.types.AnnotationSource import AnnotationSource
 from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
 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.api.types.SpanAnnotation import SpanAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import SpanAnnotationDeleteEvent, SpanAnnotationInsertEvent
 
@@ -34,7 +34,7 @@ class SpanAnnotationMutationPayload:
 
 @strawberry.type
 class SpanAnnotationMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def create_span_annotations(
         self, info: Info[Context, None], input: list[CreateSpanAnnotationInput]
     ) -> SpanAnnotationMutationPayload:
@@ -138,7 +138,7 @@ class SpanAnnotationMutationMixin:
 
             # Convert the fully loaded annotations to GQL types
             returned_annotations = [
-                to_gql_span_annotation(anno) for anno in ordered_final_annotations
+                SpanAnnotation(id=anno.id, db_record=anno) for anno in ordered_final_annotations
             ]
 
             await session.commit()
@@ -148,7 +148,7 @@ class SpanAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def create_span_note(
         self, info: Info[Context, None], annotation_input: CreateSpanNoteInput
     ) -> SpanAnnotationMutationPayload:
@@ -184,14 +184,16 @@ class SpanAnnotationMutationMixin:
             processed_annotation = result.one()
 
             info.context.event_queue.put(SpanAnnotationInsertEvent((processed_annotation.id,)))
-            returned_annotation = to_gql_span_annotation(processed_annotation)
+            returned_annotation = SpanAnnotation(
+                id=processed_annotation.id, db_record=processed_annotation
+            )
             await session.commit()
         return SpanAnnotationMutationPayload(
             span_annotations=[returned_annotation],
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_span_annotations(
         self, info: Info[Context, None], input: list[PatchAnnotationInput]
     ) -> SpanAnnotationMutationPayload:
@@ -256,7 +258,7 @@ class SpanAnnotationMutationMixin:
                 session.add(span_annotation)
 
             patched_annotations = [
-                to_gql_span_annotation(span_annotation)
+                SpanAnnotation(id=span_annotation.id, db_record=span_annotation)
                 for span_annotation in span_annotations_by_id.values()
             ]
 
@@ -268,7 +270,7 @@ class SpanAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
     async def delete_span_annotations(
         self, info: Info[Context, None], input: DeleteAnnotationsInput
     ) -> SpanAnnotationMutationPayload:
@@ -320,7 +322,10 @@ class SpanAnnotationMutationMixin:
                 )
 
         deleted_annotations_gql = [
-            to_gql_span_annotation(deleted_annotations_by_id[id]) for id in span_annotation_ids
+            SpanAnnotation(
+                id=deleted_annotations_by_id[id].id, db_record=deleted_annotations_by_id[id]
+            )
+            for id in span_annotation_ids
         ]
         info.context.event_queue.put(
             SpanAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -6,7 +6,7 @@ from starlette.requests import Request
 from strawberry import UNSET, Info
 
 from phoenix.db import models
-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 BadRequest, NotFound, Unauthorized
 from phoenix.server.api.helpers.annotations import get_user_identifier
@@ -16,7 +16,7 @@ from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationI
 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.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationDeleteEvent, TraceAnnotationInsertEvent
 
@@ -29,7 +29,7 @@ class TraceAnnotationMutationPayload:
 
 @strawberry.type
 class TraceAnnotationMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def create_trace_annotations(
         self, info: Info[Context, None], input: list[CreateTraceAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
@@ -49,8 +49,7 @@ class TraceAnnotationMutationMixin:
                 trace_rowid = from_global_id_with_expected_type(annotation_input.trace_id, "Trace")
             except ValueError:
                 raise BadRequest(
-                    f"Invalid trace ID for annotation at index {idx}: "
-                    f"{annotation_input.trace_id}"
+                    f"Invalid trace ID for annotation at index {idx}: {annotation_input.trace_id}"
                 )
             trace_rowids.append(trace_rowid)
 
@@ -112,7 +111,9 @@ class TraceAnnotationMutationMixin:
             info.context.event_queue.put(TraceAnnotationInsertEvent(inserted_annotation_ids))
 
         returned_annotations = [
-            to_gql_trace_annotation(processed_annotations_map[i])
+            TraceAnnotation(
+                id=processed_annotations_map[i].id, db_record=processed_annotations_map[i]
+            )
             for i in sorted(processed_annotations_map.keys())
         ]
 
@@ -121,7 +122,7 @@ class TraceAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def patch_trace_annotations(
         self, info: Info[Context, None], input: list[PatchAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
@@ -187,7 +188,7 @@ class TraceAnnotationMutationMixin:
             await session.commit()
 
         patched_annotations = [
-            to_gql_trace_annotation(trace_annotation)
+            TraceAnnotation(id=trace_annotation.id, db_record=trace_annotation)
             for trace_annotation in trace_annotations_by_id.values()
         ]
         info.context.event_queue.put(TraceAnnotationInsertEvent(tuple(patch_by_id.keys())))
@@ -196,7 +197,7 @@ class TraceAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
     async def delete_trace_annotations(
         self, info: Info[Context, None], input: DeleteAnnotationsInput
     ) -> TraceAnnotationMutationPayload:
@@ -246,7 +247,10 @@ class TraceAnnotationMutationMixin:
                 )
 
         deleted_gql_annotations = [
-            to_gql_trace_annotation(deleted_annotations_by_id[id]) for id in trace_annotation_ids
+            TraceAnnotation(
+                id=deleted_annotations_by_id[id].id, db_record=deleted_annotations_by_id[id]
+            )
+            for id in trace_annotation_ids
         ]
         info.context.event_queue.put(
             TraceAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))

phoenix/server/api/mutations/trace_mutations.py

@@ -1,12 +1,12 @@
 import strawberry
-from sqlalchemy import and_, delete, not_, select
+from sqlalchemy import and_, delete, not_, select, update
 from sqlalchemy.orm import load_only
 from sqlalchemy.sql import literal
 from strawberry.relay import GlobalID
 from strawberry.types import Info
 
 from phoenix.db import models
-from phoenix.server.api.auth import IsNotReadOnly
+from phoenix.server.api.auth import IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.queries import Query
@@ -16,7 +16,7 @@ from phoenix.server.dml_event import SpanDeleteEvent
 
 @strawberry.type
 class TraceMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
     async def delete_traces(
         self,
         info: Info[Context, None],
@@ -72,3 +72,47 @@ class TraceMutationMixin:
                 )
             info.context.event_queue.put(SpanDeleteEvent(project_ids))
         return Query()
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def transfer_traces_to_project(
+        self,
+        info: Info[Context, None],
+        trace_ids: list[GlobalID],
+        project_id: GlobalID,
+    ) -> Query:
+        if not trace_ids:
+            raise BadRequest("Must provide at least one trace ID to transfer")
+        trace_ids = list(set(trace_ids))
+        try:
+            trace_rowids = [
+                from_global_id_with_expected_type(global_id=id, expected_type_name="Trace")
+                for id in trace_ids
+            ]
+            dest_project_rowid = from_global_id_with_expected_type(
+                global_id=project_id, expected_type_name="Project"
+            )
+        except ValueError as error:
+            raise BadRequest(str(error))
+
+        async with info.context.db() as session:
+            dest_project = await session.get(models.Project, dest_project_rowid)
+            if dest_project is None:
+                raise BadRequest("Destination project does not exist")
+
+            traces = (
+                await session.scalars(select(models.Trace).where(models.Trace.id.in_(trace_rowids)))
+            ).all()
+            if len(traces) < len(trace_rowids):
+                raise BadRequest("Invalid trace IDs provided")
+
+            source_project_ids = set(trace.project_rowid for trace in traces)
+            if len(source_project_ids) > 1:
+                raise BadRequest("Cannot transfer traces from multiple projects")
+
+            await session.execute(
+                update(models.Trace)
+                .where(models.Trace.id.in_(trace_rowids))
+                .values(project_rowid=dest_project_rowid)
+            )
+
+        return Query()

phoenix/server/api/mutations/user_mutations.py

@@ -21,18 +21,19 @@ from phoenix.auth import (
     PASSWORD_REQUIREMENTS,
     PHOENIX_ACCESS_TOKEN_COOKIE_NAME,
     PHOENIX_REFRESH_TOKEN_COOKIE_NAME,
+    sanitize_email,
     validate_email_format,
     validate_password_format,
 )
 from phoenix.config import get_env_disable_basic_auth
-from phoenix.db import enums, models
-from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly
+from phoenix.db import models
+from phoenix.server.api.auth import IsAdmin, IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, Conflict, NotFound, Unauthorized
 from phoenix.server.api.input_types.UserRoleInput import UserRoleInput
 from phoenix.server.api.types.AuthMethod import AuthMethod
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.api.types.User import User, to_gql_user
+from phoenix.server.api.types.User import User
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.types import AccessTokenId, ApiKeyId, PasswordResetTokenId, RefreshTokenId
 
@@ -49,7 +50,10 @@ class CreateUserInput:
     auth_method: Optional[AuthMethod] = AuthMethod.LOCAL
 
     def __post_init__(self) -> None:
-        if self.auth_method is AuthMethod.OAUTH2:
+        if self.auth_method is AuthMethod.LDAP:
+            if self.password:
+                raise BadRequest("Password is not allowed for LDAP authentication")
+        elif self.auth_method is AuthMethod.OAUTH2:
             if self.password:
                 raise BadRequest("Password is not allowed for OAuth2 authentication")
         elif get_env_disable_basic_auth():
@@ -102,6 +106,11 @@ class DeleteUsersInput:
     user_ids: list[GlobalID]
 
 
+@strawberry.type
+class DeleteUsersPayload:
+    user_ids: list[GlobalID]
+
+
 @strawberry.type
 class UserMutationPayload:
     user: User
@@ -109,26 +118,34 @@ class UserMutationPayload:
 
 @strawberry.type
 class UserMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin, IsLocked])  # type: ignore
     async def create_user(
         self,
         info: Info[Context, None],
         input: CreateUserInput,
     ) -> UserMutationPayload:
+        # Sanitize email by trimming and lowercasing
+        email = sanitize_email(input.email)
+
         user: models.User
-        if input.auth_method is AuthMethod.OAUTH2:
+        if input.auth_method is AuthMethod.LDAP:
+            user = models.LDAPUser(
+                email=email,
+                username=input.username,
+            )
+        elif input.auth_method is AuthMethod.OAUTH2:
             user = models.OAuth2User(
-                email=input.email,
+                email=email,
                 username=input.username,
             )
         else:
             assert input.password
-            validate_email_format(input.email)
+            validate_email_format(email)
             validate_password_format(input.password)
             salt = secrets.token_bytes(DEFAULT_SECRET_LENGTH)
             password_hash = await info.context.hash_password(Secret(input.password), salt)
             user = models.LocalUser(
-                email=input.email,
+                email=email,
                 username=input.username,
                 password_hash=password_hash,
                 password_salt=salt,
@@ -151,9 +168,9 @@ class UserMutationMixin:
             except Exception as error:
                 # Log the error but do not raise it
                 logger.error(f"Failed to send welcome email: {error}")
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin])  # type: ignore
     async def patch_user(
         self,
         info: Info[Context, None],
@@ -170,6 +187,7 @@ class UserMutationMixin:
             if not (user := await session.scalar(_select_user_by_id(user_id))):
                 raise NotFound("User not found")
             stack.enter_context(session.no_autoflush)
+            should_log_out = False
             if input.new_role:
                 if user.email == DEFAULT_ADMIN_EMAIL:
                     raise Unauthorized("Cannot modify role for the default admin user")
@@ -179,6 +197,7 @@ class UserMutationMixin:
                 if user_role_id is None:
                     raise NotFound(f"Role {input.new_role.value} not found")
                 user.user_role_id = user_role_id
+                should_log_out = True
             if password := input.new_password:
                 if user.auth_method != "LOCAL":
                     raise Conflict("Cannot modify password for non-local user")
@@ -187,6 +206,7 @@ class UserMutationMixin:
                 user.password_salt = salt
                 user.password_hash = await info.context.hash_password(Secret(password), salt)
                 user.reset_password = True
+                should_log_out = True
             if username := input.new_username:
                 user.username = username
             assert user in session.dirty
@@ -195,11 +215,11 @@ class UserMutationMixin:
             except (PostgreSQLIntegrityError, SQLiteIntegrityError) as error:
                 raise Conflict(_user_operation_error_message(error, "modify"))
         assert user
-        if input.new_password:
+        if should_log_out:
             await info.context.log_out(user.id)
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_viewer(
         self,
         info: Info[Context, None],
@@ -239,33 +259,28 @@ class UserMutationMixin:
             response = info.context.get_response()
             response.delete_cookie(PHOENIX_REFRESH_TOKEN_COOKIE_NAME)
             response.delete_cookie(PHOENIX_ACCESS_TOKEN_COOKIE_NAME)
-        return UserMutationPayload(user=to_gql_user(user))
+        return UserMutationPayload(user=User(id=user.id, db_record=user))
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin, IsLocked])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer, IsAdmin, IsLocked])  # type: ignore
     async def delete_users(
         self,
         info: Info[Context, None],
         input: DeleteUsersInput,
-    ) -> None:
+    ) -> DeleteUsersPayload:
         assert (token_store := info.context.token_store) is not None
         if not input.user_ids:
-            return
-        user_ids = tuple(
-            map(
-                lambda gid: from_global_id_with_expected_type(gid, User.__name__),
-                set(input.user_ids),
-            )
-        )
-        system_user_role_id = (
-            select(models.UserRole.id)
-            .where(models.UserRole.name == enums.UserRole.SYSTEM.value)
-            .scalar_subquery()
-        )
-        admin_user_role_id = (
-            select(models.UserRole.id)
-            .where(models.UserRole.name == enums.UserRole.ADMIN.value)
-            .scalar_subquery()
-        )
+            raise BadRequest("At least one user ID is required")
+        user_rowid_to_gid: dict[int, GlobalID] = {}
+        for user_gid in input.user_ids:
+            try:
+                user_rowid = from_global_id_with_expected_type(user_gid, User.__name__)
+            except ValueError:
+                raise BadRequest(f"Invalid user ID: '{user_gid}'")
+            user_rowid_to_gid[user_rowid] = user_gid
+
+        user_rowids = list(user_rowid_to_gid.keys())
+        system_user_role_id = select(models.UserRole.id).filter_by(name="SYSTEM").scalar_subquery()
+        admin_user_role_id = select(models.UserRole.id).filter_by(name="ADMIN").scalar_subquery()
         default_admin_user_id = (
             select(models.User.id)
             .where(
@@ -302,7 +317,7 @@ class UserMutationMixin:
                     .select_from(models.User)
                     .where(
                         and_(
-                            models.User.id.in_(user_ids),
+                            models.User.id.in_(user_rowids),
                             models.User.user_role_id != system_user_role_id,
                         )
                     )
@@ -310,41 +325,51 @@ class UserMutationMixin:
             ).all()
             if deletes_default_admin:
                 raise Conflict("Cannot delete the default admin user")
-            if num_resolved_user_ids < len(user_ids):
+            if num_resolved_user_ids < len(user_rowids):
                 raise NotFound("Some user IDs could not be found")
             password_reset_token_ids = [
                 PasswordResetTokenId(id_)
                 async for id_ in await session.stream_scalars(
                     select(models.PasswordResetToken.id).where(
-                        models.PasswordResetToken.user_id.in_(user_ids)
+                        models.PasswordResetToken.user_id.in_(user_rowids)
                     )
                 )
             ]
             access_token_ids = [
                 AccessTokenId(id_)
                 async for id_ in await session.stream_scalars(
-                    select(models.AccessToken.id).where(models.AccessToken.user_id.in_(user_ids))
+                    select(models.AccessToken.id).where(models.AccessToken.user_id.in_(user_rowids))
                 )
             ]
             refresh_token_ids = [
                 RefreshTokenId(id_)
                 async for id_ in await session.stream_scalars(
-                    select(models.RefreshToken.id).where(models.RefreshToken.user_id.in_(user_ids))
+                    select(models.RefreshToken.id).where(
+                        models.RefreshToken.user_id.in_(user_rowids)
+                    )
                 )
             ]
             api_key_ids = [
                 ApiKeyId(id_)
                 async for id_ in await session.stream_scalars(
-                    select(models.ApiKey.id).where(models.ApiKey.user_id.in_(user_ids))
+                    select(models.ApiKey.id).where(models.ApiKey.user_id.in_(user_rowids))
                 )
             ]
-            await session.execute(delete(models.User).where(models.User.id.in_(user_ids)))
+            deleted_user_ids = await session.scalars(
+                delete(models.User).where(models.User.id.in_(user_rowids)).returning(models.User.id)
+            )
         await token_store.revoke(
             *password_reset_token_ids,
             *access_token_ids,
             *refresh_token_ids,
             *api_key_ids,
         )
+        unique_deleted_user_ids = set(deleted_user_ids)
+        deleted_user_gids: list[GlobalID] = []
+        for user_rowid, user_gid in user_rowid_to_gid.items():
+            if user_rowid in unique_deleted_user_ids:
+                deleted_user_gids.append(user_gid)
+        return DeleteUsersPayload(user_ids=deleted_user_gids)
 
 
 def _select_role_id_by_name(role_name: str) -> Select[tuple[int]]: