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

phoenix/server/api/mutations/model_mutations.py

@@ -0,0 +1,210 @@
+import re
+from datetime import datetime, timezone
+from typing import Optional
+
+import sqlalchemy as sa
+import strawberry
+from sqlalchemy import delete
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
+from sqlalchemy.orm import joinedload
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
+from strawberry.relay import GlobalID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly, IsNotViewer
+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.GenerativeModel import GenerativeModel
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.TokenPrice import TokenKind
+
+
+@strawberry.input
+class TokenPriceInput:
+    token_type: str
+    cost_per_million_tokens: float
+    kind: TokenKind
+
+    @property
+    def token_prices(self) -> models.TokenPrice:
+        """Generate TokenPrice instances based on the input."""
+        return models.TokenPrice(
+            token_type=self.token_type,
+            is_prompt=self.kind == TokenKind.PROMPT,
+            base_rate=self.cost_per_million_tokens / 1_000_000,
+        )
+
+
+@strawberry.input
+class CreateModelMutationInput:
+    name: str
+    provider: Optional[str] = None
+    name_pattern: str
+    costs: list[TokenPriceInput]
+    start_time: Optional[datetime] = None
+
+
+@strawberry.type
+class CreateModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.input
+class UpdateModelMutationInput:
+    id: GlobalID
+    name: str
+    provider: Optional[str]
+    name_pattern: str
+    costs: list[TokenPriceInput]
+    start_time: Optional[datetime] = None
+
+
+@strawberry.type
+class UpdateModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.input
+class DeleteModelMutationInput:
+    id: GlobalID
+
+
+@strawberry.type
+class DeleteModelMutationPayload:
+    model: GenerativeModel
+    query: Query
+
+
+@strawberry.type
+class ModelMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def create_model(
+        self,
+        info: Info[Context, None],
+        input: CreateModelMutationInput,
+    ) -> CreateModelMutationPayload:
+        cost_types = set(cost.token_type for cost in input.costs)
+        if "input" not in cost_types:
+            raise BadRequest("input cost is required")
+        if "output" not in cost_types:
+            raise BadRequest("output cost is required")
+        name_pattern = _compile_regular_expression(input.name_pattern)
+        token_prices = [cost.token_prices for cost in input.costs]
+        model = models.GenerativeModel(
+            name=input.name,
+            provider=input.provider,
+            name_pattern=name_pattern,
+            is_built_in=False,
+            token_prices=token_prices,
+            start_time=input.start_time,
+        )
+        async with info.context.db() as session:
+            session.add(model)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Model with name '{input.name}' already exists")
+
+        return CreateModelMutationPayload(
+            model=GenerativeModel(id=model.id, db_record=model),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def update_model(
+        self,
+        info: Info[Context, None],
+        input: UpdateModelMutationInput,
+    ) -> UpdateModelMutationPayload:
+        try:
+            model_id = from_global_id_with_expected_type(input.id, GenerativeModel.__name__)
+        except ValueError:
+            raise BadRequest(f'Invalid model id: "{input.id}"')
+
+        cost_types = set(cost.token_type for cost in input.costs)
+        if "input" not in cost_types:
+            raise BadRequest("input cost is required")
+        if "output" not in cost_types:
+            raise BadRequest("output cost is required")
+        name_pattern = _compile_regular_expression(input.name_pattern)
+        token_prices = [cost.token_prices for cost in input.costs]
+        async with info.context.db() as session:
+            model = await session.scalar(
+                sa.select(models.GenerativeModel)
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .where(models.GenerativeModel.id == model_id)
+                .options(joinedload(models.GenerativeModel.token_prices))
+            )
+            if model is None:
+                raise NotFound(f'Model "{input.id}" not found')
+            if model.is_built_in:
+                raise BadRequest("Cannot update built-in model")
+
+            await session.execute(
+                delete(models.TokenPrice).where(models.TokenPrice.model_id == model.id)
+            )
+
+            await session.refresh(model)
+
+            model.name = input.name
+            model.provider = input.provider or ""
+            model.name_pattern = name_pattern
+            model.token_prices = token_prices
+            model.start_time = input.start_time
+            # Explicitly set updated_at so the GenerativeModelStore daemon picks up this
+            # change (SQLAlchemy's onupdate may not trigger for relationship-only changes).
+            model.updated_at = datetime.now(timezone.utc)
+            session.add(model)
+            try:
+                await session.flush()
+            except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+                raise Conflict(f"Model with name '{input.name}' already exists")
+
+        return UpdateModelMutationPayload(
+            model=GenerativeModel(id=model.id, db_record=model),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def delete_model(
+        self,
+        info: Info[Context, None],
+        input: DeleteModelMutationInput,
+    ) -> DeleteModelMutationPayload:
+        try:
+            model_id = from_global_id_with_expected_type(input.id, GenerativeModel.__name__)
+        except ValueError:
+            raise BadRequest(f'Invalid model id: "{input.id}"')
+
+        async with info.context.db() as session:
+            model = await session.scalar(
+                sa.update(models.GenerativeModel)
+                .values(deleted_at=datetime.now(timezone.utc))
+                .where(models.GenerativeModel.deleted_at.is_(None))
+                .where(models.GenerativeModel.id == model_id)
+                .returning(models.GenerativeModel)
+            )
+            if model is None:
+                raise NotFound(f'Model "{input.id}" not found')
+            if model.is_built_in:
+                await session.rollback()
+                raise BadRequest("Cannot delete built-in model")
+        return DeleteModelMutationPayload(
+            model=GenerativeModel(id=model.id, db_record=model),
+            query=Query(),
+        )
+
+
+def _compile_regular_expression(maybe_regex: str) -> re.Pattern[str]:
+    """
+    Compile the given string as a regular expression.
+    Raises a BadRequest error if the given string is not a valid regex.
+    """
+    try:
+        return re.compile(maybe_regex)
+    except re.error as error:
+        raise BadRequest(f"Invalid regex: {str(error)}")

phoenix/server/api/mutations/project_mutations.py

@@ -1,22 +1,58 @@
 import strawberry
 from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
 from sqlalchemy.orm import load_only
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 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.auth import IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest, Conflict
 from phoenix.server.api.input_types.ClearProjectInput import ClearProjectInput
+from phoenix.server.api.input_types.CreateProjectInput import CreateProjectInput
 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
+from phoenix.server.api.types.Project import Project, to_gql_project
+from phoenix.server.dml_event import ProjectDeleteEvent, ProjectInsertEvent, SpanDeleteEvent
+
+
+@strawberry.type
+class ProjectMutationPayload:
+    project: Project
+    query: Query
 
 
 @strawberry.type
 class ProjectMutationMixin:
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # type: ignore
+    async def create_project(
+        self,
+        info: Info[Context, None],
+        input: CreateProjectInput,
+    ) -> ProjectMutationPayload:
+        if not (name := input.name.strip()):
+            raise BadRequest("Name cannot be empty")
+        description = (input.description or "").strip() or None
+        gradient_start_color = (input.gradient_start_color or "").strip() or None
+        gradient_end_color = (input.gradient_end_color or "").strip() or None
+        project = models.Project(
+            name=name,
+            description=description,
+            gradient_start_color=gradient_start_color,
+            gradient_end_color=gradient_end_color,
+        )
+        try:
+            async with info.context.db() as session:
+                session.add(project)
+        except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+            raise Conflict(f"Project with name '{name}' already exists")
+        info.context.event_queue.put(ProjectInsertEvent((project.id,)))
+        return ProjectMutationPayload(project=to_gql_project(project), query=Query())
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # 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:
@@ -33,7 +69,7 @@ class ProjectMutationMixin:
         info.context.event_queue.put(ProjectDeleteEvent((project_id,)))
         return Query()
 
-    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsNotViewer])  # 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"
@@ -47,11 +83,14 @@ class ProjectMutationMixin:
             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))
-                    )
-                )
+            session_ids_to_delete = [
+                id_ for id_ in set(deleted_trace_project_session_ids) if id_ is not None
+            ]
+            # Process deletions in chunks of 10000 to avoid PostgreSQL argument limit
+            chunk_size = 10000
+            stmt = delete(models.ProjectSession)
+            for i in range(0, len(session_ids_to_delete), chunk_size):
+                chunk = session_ids_to_delete[i : i + chunk_size]
+                await session.execute(stmt.where(models.ProjectSession.id.in_(chunk)))
         info.context.event_queue.put(SpanDeleteEvent((project_id,)))
         return Query()

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],