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

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -1,19 +1,23 @@
-from collections.abc import Sequence
+from typing import Optional
 
 import strawberry
-from sqlalchemy import delete, insert, update
-from strawberry import UNSET
-from strawberry.types import Info
+from sqlalchemy import delete, insert, select
+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.context import Context
+from phoenix.server.api.exceptions import BadRequest, NotFound, Unauthorized
+from phoenix.server.api.helpers.annotations import get_user_identifier
 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.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.bearer_auth import PhoenixUser
 from phoenix.server.dml_event import TraceAnnotationDeleteEvent, TraceAnnotationInsertEvent
 
 
@@ -29,33 +33,91 @@ class TraceAnnotationMutationMixin:
     async def create_trace_annotations(
         self, info: Info[Context, None], input: list[CreateTraceAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
-        inserted_annotations: Sequence[models.TraceAnnotation] = []
+        if not input:
+            raise BadRequest("No trace annotations provided.")
+
+        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)
+
+        processed_annotations_map: dict[int, models.TraceAnnotation] = {}
+
+        trace_rowids = []
+        for idx, annotation_input in enumerate(input):
+            try:
+                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}"
+                )
+            trace_rowids.append(trace_rowid)
+
         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 idx, (trace_rowid, annotation_input) in enumerate(zip(trace_rowids, input)):
+                resolved_identifier = ""
+                if isinstance(annotation_input.identifier, str):
+                    resolved_identifier = annotation_input.identifier
+                elif annotation_input.source == AnnotationSource.APP and user_id is not None:
+                    resolved_identifier = get_user_identifier(user_id)
+                values = {
+                    "trace_rowid": trace_rowid,
+                    "name": annotation_input.name,
+                    "label": annotation_input.label,
+                    "score": annotation_input.score,
+                    "explanation": annotation_input.explanation,
+                    "annotator_kind": annotation_input.annotator_kind.value,
+                    "metadata_": annotation_input.metadata,
+                    "identifier": resolved_identifier,
+                    "source": annotation_input.source.value,
+                    "user_id": user_id,
+                }
+
+                processed_annotation: Optional[models.TraceAnnotation] = None
+
+                # Check if an annotation with this trace_rowid, name, and identifier already exists
+                q = select(models.TraceAnnotation).where(
+                    models.TraceAnnotation.trace_rowid == trace_rowid,
+                    models.TraceAnnotation.name == annotation_input.name,
+                    models.TraceAnnotation.identifier == resolved_identifier,
                 )
-                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))
-            )
+                existing_annotation = await session.scalar(q)
+
+                if existing_annotation:
+                    # Update existing annotation
+                    existing_annotation.name = values["name"]
+                    existing_annotation.label = values["label"]
+                    existing_annotation.score = values["score"]
+                    existing_annotation.explanation = values["explanation"]
+                    existing_annotation.metadata_ = values["metadata_"]
+                    existing_annotation.annotator_kind = values["annotator_kind"]
+                    existing_annotation.source = values["source"]
+                    existing_annotation.user_id = values["user_id"]
+                    session.add(existing_annotation)
+                    processed_annotation = existing_annotation
+
+                if processed_annotation is None:
+                    stmt = insert(models.TraceAnnotation).values(**values)
+                    stmt = stmt.returning(models.TraceAnnotation)
+                    result = await session.scalars(stmt)
+                    processed_annotation = result.one()
+
+                processed_annotations_map[idx] = processed_annotation
+
+            await session.commit()
+
+        inserted_annotation_ids = tuple(anno.id for anno in processed_annotations_map.values())
+        if inserted_annotation_ids:
+            info.context.event_queue.put(TraceAnnotationInsertEvent(inserted_annotation_ids))
+
+        returned_annotations = [
+            to_gql_trace_annotation(processed_annotations_map[i])
+            for i in sorted(processed_annotations_map.keys())
+        ]
+
         return TraceAnnotationMutationPayload(
-            trace_annotations=[
-                to_gql_trace_annotation(annotation) for annotation in inserted_annotations
-            ],
+            trace_annotations=returned_annotations,
             query=Query(),
         )
 
@@ -63,65 +125,132 @@ class TraceAnnotationMutationMixin:
     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:
+        if not input:
+            raise BadRequest("No trace annotations provided.")
+
+        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)
+
+        patch_by_id = {}
+        for patch in input:
+            try:
                 trace_annotation_id = from_global_id_with_expected_type(
-                    annotation.annotation_id, "TraceAnnotation"
+                    patch.annotation_id, "TraceAnnotation"
+                )
+            except ValueError:
+                raise BadRequest(f"Invalid trace annotation ID: {patch.annotation_id}")
+            if trace_annotation_id in patch_by_id:
+                raise BadRequest(f"Duplicate patch for trace annotation ID: {trace_annotation_id}")
+            patch_by_id[trace_annotation_id] = patch
+
+        async with info.context.db() as session:
+            trace_annotations_by_id = {}
+            for trace_annotation in await session.scalars(
+                select(models.TraceAnnotation).where(
+                    models.TraceAnnotation.id.in_(patch_by_id.keys())
                 )
-                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 trace_annotation.user_id != user_id:
+                    raise Unauthorized(
+                        "At least one trace annotation is not associated with the current user."
                     )
-                    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)
+                trace_annotations_by_id[trace_annotation.id] = trace_annotation
+
+            missing_trace_annotation_ids = set(patch_by_id.keys()) - set(
+                trace_annotations_by_id.keys()
+            )
+            if missing_trace_annotation_ids:
+                raise NotFound(
+                    f"Could not find trace annotations with IDs: {missing_trace_annotation_ids}"
                 )
-                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())
+
+            for trace_annotation_id, patch in patch_by_id.items():
+                trace_annotation = trace_annotations_by_id[trace_annotation_id]
+                if patch.name:
+                    trace_annotation.name = patch.name
+                if patch.annotator_kind:
+                    trace_annotation.annotator_kind = patch.annotator_kind.value
+                if patch.label is not UNSET:
+                    trace_annotation.label = patch.label
+                if patch.score is not UNSET:
+                    trace_annotation.score = patch.score
+                if patch.explanation is not UNSET:
+                    trace_annotation.explanation = patch.explanation
+                if patch.metadata is not UNSET:
+                    assert isinstance(patch.metadata, dict)
+                    trace_annotation.metadata_ = patch.metadata
+                if patch.identifier is not UNSET:
+                    trace_annotation.identifier = patch.identifier or ""
+                session.add(trace_annotation)
+            await session.commit()
+
+        patched_annotations = [
+            to_gql_trace_annotation(trace_annotation)
+            for trace_annotation in trace_annotations_by_id.values()
+        ]
+        info.context.event_queue.put(TraceAnnotationInsertEvent(tuple(patch_by_id.keys())))
+        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
-        ]
+        if not input.annotation_ids:
+            raise BadRequest("No trace annotation IDs provided.")
+
+        trace_annotation_ids: dict[int, None] = {}  # use dict to preserve order
+        for annotation_gid in input.annotation_ids:
+            try:
+                annotation_id = from_global_id_with_expected_type(annotation_gid, "TraceAnnotation")
+            except ValueError:
+                raise BadRequest(f"Invalid trace annotation ID: {annotation_gid}")
+            if annotation_id in trace_annotation_ids:
+                raise BadRequest(f"Duplicate trace annotation ID: {annotation_id}")
+            trace_annotation_ids[annotation_id] = None
+
+        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:
-            stmt = (
+            result = await session.scalars(
                 delete(models.TraceAnnotation)
-                .where(models.TraceAnnotation.id.in_(trace_annotation_ids))
+                .where(models.TraceAnnotation.id.in_(trace_annotation_ids.keys()))
                 .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))
+            deleted_annotations_by_id = {annotation.id: annotation for annotation in result.all()}
+
+            if not user_is_admin and any(
+                annotation.user_id != user_id for annotation in deleted_annotations_by_id.values()
+            ):
+                await session.rollback()
+                raise Unauthorized(
+                    "At least one trace annotation is not associated with the current user "
+                    "and the current user is not an admin."
+                )
+
+            missing_trace_annotation_ids = set(trace_annotation_ids.keys()) - set(
+                deleted_annotations_by_id.keys()
             )
+            if missing_trace_annotation_ids:
+                raise NotFound(
+                    f"Could not find trace annotations with IDs: {missing_trace_annotation_ids}"
+                )
+
+        deleted_gql_annotations = [
+            to_gql_trace_annotation(deleted_annotations_by_id[id]) for id in trace_annotation_ids
+        ]
+        info.context.event_queue.put(
+            TraceAnnotationDeleteEvent(tuple(deleted_annotations_by_id.keys()))
+        )
         return TraceAnnotationMutationPayload(
-            trace_annotations=deleted_annotations_gql, query=Query()
+            trace_annotations=deleted_gql_annotations, query=Query()
         )

phoenix/server/api/queries.py

@@ -19,6 +19,7 @@ from phoenix.config import (
     getenv,
 )
 from phoenix.db import enums, models
+from phoenix.db.constants import DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
 from phoenix.db.helpers import SupportedSQLDialect, exclude_experiment_projects
 from phoenix.db.models import DatasetExample as OrmExample
 from phoenix.db.models import DatasetExampleRevision as OrmRevision
@@ -42,6 +43,7 @@ from phoenix.server.api.input_types.ClusterInput import ClusterInput
 from phoenix.server.api.input_types.Coordinates import InputCoordinate2D, InputCoordinate3D
 from phoenix.server.api.input_types.DatasetSort import DatasetSort
 from phoenix.server.api.input_types.InvocationParameters import InvocationParameter
+from phoenix.server.api.types.AnnotationConfig import AnnotationConfig, to_gql_annotation_config
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
 from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
@@ -65,14 +67,17 @@ from phoenix.server.api.types.node import from_global_id, from_global_id_with_ex
 from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
 from phoenix.server.api.types.Project import Project
 from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
+from phoenix.server.api.types.ProjectTraceRetentionPolicy import ProjectTraceRetentionPolicy
 from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
 from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
 from phoenix.server.api.types.PromptVersion import PromptVersion, to_gql_prompt_version
 from phoenix.server.api.types.PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
 from phoenix.server.api.types.Trace import Trace
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
 from phoenix.server.api.types.User import User, to_gql_user
 from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
 from phoenix.server.api.types.UserRole import UserRole
@@ -588,6 +593,26 @@ class Query:
                 if not (prompt_version_tag := await session.get(models.PromptVersionTag, node_id)):
                     raise NotFound(f"Unknown prompt version tag: {id}")
             return to_gql_prompt_version_tag(prompt_version_tag)
+        elif type_name == ProjectTraceRetentionPolicy.__name__:
+            async with info.context.db() as session:
+                db_policy = await session.scalar(
+                    select(models.ProjectTraceRetentionPolicy).filter_by(id=node_id)
+                )
+                if not db_policy:
+                    raise NotFound(f"Unknown project trace retention policy: {id}")
+            return ProjectTraceRetentionPolicy(id=db_policy.id, db_policy=db_policy)
+        elif type_name == SpanAnnotation.__name__:
+            async with info.context.db() as session:
+                span_annotation = await session.get(models.SpanAnnotation, node_id)
+                if not span_annotation:
+                    raise NotFound(f"Unknown span annotation: {id}")
+            return to_gql_span_annotation(span_annotation)
+        elif type_name == TraceAnnotation.__name__:
+            async with info.context.db() as session:
+                trace_annotation = await session.get(models.TraceAnnotation, node_id)
+                if not trace_annotation:
+                    raise NotFound(f"Unknown trace annotation: {id}")
+            return to_gql_trace_annotation(trace_annotation)
         raise NotFound(f"Unknown node type: {type_name}")
 
     @strawberry.field
@@ -657,6 +682,28 @@ class Query:
                 args=args,
             )
 
+    @strawberry.field
+    async def annotation_configs(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = None,
+        after: Optional[str] = None,
+        before: Optional[str] = None,
+    ) -> Connection[AnnotationConfig]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        async with info.context.db() as session:
+            configs = await session.stream_scalars(
+                select(models.AnnotationConfig).order_by(models.AnnotationConfig.name)
+            )
+            data = [to_gql_annotation_config(config) async for config in configs]
+            return connection_from_list(data=data, args=args)
+
     @strawberry.field
     def clusters(
         self,
@@ -785,6 +832,45 @@ class Query:
             clustered_events=clustered_events,
         )
 
+    @strawberry.field
+    async def default_project_trace_retention_policy(
+        self,
+        info: Info[Context, None],
+    ) -> ProjectTraceRetentionPolicy:
+        stmt = select(models.ProjectTraceRetentionPolicy).filter_by(
+            id=DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
+        )
+        async with info.context.db() as session:
+            db_policy = await session.scalar(stmt)
+        assert db_policy
+        return ProjectTraceRetentionPolicy(id=db_policy.id, db_policy=db_policy)
+
+    @strawberry.field
+    async def project_trace_retention_policies(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 100,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[ProjectTraceRetentionPolicy]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        stmt = select(models.ProjectTraceRetentionPolicy).order_by(
+            models.ProjectTraceRetentionPolicy.id
+        )
+        async with info.context.db() as session:
+            result = await session.stream_scalars(stmt)
+            data = [
+                ProjectTraceRetentionPolicy(id=db_policy.id, db_policy=db_policy)
+                async for db_policy in result
+            ]
+        return connection_from_list(data=data, args=args)
+
     @strawberry.field(
         description="The allocated storage capacity of the database in bytes. "
         "Return None if this information is unavailable.",

phoenix/server/api/routers/v1/__init__.py

@@ -4,6 +4,8 @@ from starlette.status import HTTP_403_FORBIDDEN
 
 from phoenix.server.bearer_auth import is_authenticated
 
+from .annotation_configs import router as annotation_configs_router
+from .annotations import router as annotations_router
 from .datasets import router as datasets_router
 from .evaluations import router as evaluations_router
 from .experiment_evaluations import router as experiment_evaluations_router
@@ -56,6 +58,8 @@ def create_v1_router(authentication_enabled: bool) -> APIRouter:
             ]
         ),
     )
+    router.include_router(annotation_configs_router)
+    router.include_router(annotations_router)
     router.include_router(datasets_router)
     router.include_router(experiments_router)
     router.include_router(experiment_runs_router)