arize-phoenix 4.35.2__py3-none-any.whl → 5.0.0__py3-none-any.whl

phoenix/server/api/mutations/api_key_mutations.py

@@ -1,20 +1,22 @@
-from datetime import datetime
-from typing import Any, Dict, Optional
+from datetime import datetime, timezone
+from typing import Optional
 
-import jwt
 import strawberry
-from sqlalchemy import insert, select
+from sqlalchemy import select
 from strawberry import UNSET
 from strawberry.relay import GlobalID
 from strawberry.types import Info
 
-from phoenix.db import models
+from phoenix.db import enums, models
+from phoenix.server.api.auth import IsAdmin, IsNotReadOnly
 from phoenix.server.api.context import Context
-from phoenix.server.api.exceptions import NotFound
-from phoenix.server.api.mutations.auth import HasSecret, IsAuthenticated
+from phoenix.server.api.exceptions import Unauthorized
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
+from phoenix.server.api.types.UserApiKey import UserApiKey
+from phoenix.server.bearer_auth import PhoenixUser
+from phoenix.server.types import ApiKeyAttributes, ApiKeyClaims, ApiKeyId, UserId
 
 
 @strawberry.type
@@ -31,119 +33,135 @@ class CreateApiKeyInput:
     expires_at: Optional[datetime] = UNSET
 
 
+@strawberry.type
+class CreateUserApiKeyMutationPayload:
+    jwt: str
+    api_key: UserApiKey
+    query: Query
+
+
+@strawberry.input
+class CreateUserApiKeyInput:
+    name: str
+    description: Optional[str] = UNSET
+    expires_at: Optional[datetime] = UNSET
+
+
 @strawberry.input
 class DeleteApiKeyInput:
     id: GlobalID
 
 
 @strawberry.type
-class DeleteSystemApiKeyMutationPayload:
-    id: GlobalID
+class DeleteApiKeyMutationPayload:
+    apiKeyId: GlobalID
     query: Query
 
 
 @strawberry.type
 class ApiKeyMutationMixin:
-    @strawberry.mutation(permission_classes=[HasSecret, IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
     async def create_system_api_key(
         self, info: Info[Context, None], input: CreateApiKeyInput
     ) -> CreateSystemApiKeyMutationPayload:
-        # TODO(auth): safe guard against auth being disabled and secret not being set
+        assert (token_store := info.context.token_store) is not None
+        user_role = enums.UserRole.SYSTEM
         async with info.context.db() as session:
             # Get the system user - note this could be pushed into a dataloader
             system_user = await session.scalar(
                 select(models.User)
                 .join(models.UserRole)  # Join User with UserRole
-                .where(models.UserRole.name == "SYSTEM")  # Filter where role is SYSTEM
+                .where(models.UserRole.name == user_role.value)  # Filter where role is SYSTEM
+                .order_by(models.User.id)
                 .limit(1)
             )
             if system_user is None:
                 raise ValueError("System user not found")
-
-            insert_stmt = (
-                insert(models.APIKey)
-                .values(
-                    user_id=system_user.id,
-                    name=input.name,
-                    description=input.description or None,
-                    expires_at=input.expires_at or None,
-                )
-                .returning(models.APIKey)
-            )
-            api_key = await session.scalar(insert_stmt)
-            assert api_key is not None
-
-        encoded_jwt = create_jwt(
-            secret=info.context.get_secret(),
-            name=api_key.name,
-            id=api_key.id,
-            description=api_key.description,
-            iat=api_key.created_at,
-            exp=api_key.expires_at,
+        issued_at = datetime.now(timezone.utc)
+        claims = ApiKeyClaims(
+            subject=UserId(system_user.id),
+            issued_at=issued_at,
+            expiration_time=input.expires_at or None,
+            attributes=ApiKeyAttributes(
+                user_role=user_role,
+                name=input.name,
+                description=input.description,
+            ),
         )
+        token, token_id = await token_store.create_api_key(claims)
         return CreateSystemApiKeyMutationPayload(
-            jwt=encoded_jwt,
+            jwt=token,
             api_key=SystemApiKey(
-                id_attr=api_key.id,
-                name=api_key.name,
-                description=api_key.description,
-                created_at=api_key.created_at,
-                expires_at=api_key.expires_at,
+                id_attr=int(token_id),
+                name=input.name,
+                description=input.description or None,
+                created_at=issued_at,
+                expires_at=input.expires_at or None,
             ),
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[HasSecret, IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def create_user_api_key(
+        self, info: Info[Context, None], input: CreateUserApiKeyInput
+    ) -> CreateUserApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
+        try:
+            user = info.context.request.user  # type: ignore
+            assert isinstance(user, PhoenixUser)
+        except AttributeError:
+            raise ValueError("User not found")
+        issued_at = datetime.now(timezone.utc)
+        claims = ApiKeyClaims(
+            subject=user.identity,
+            issued_at=issued_at,
+            expiration_time=input.expires_at or None,
+            attributes=ApiKeyAttributes(
+                user_role=enums.UserRole.MEMBER,
+                name=input.name,
+                description=input.description,
+            ),
+        )
+        token, token_id = await token_store.create_api_key(claims)
+        return CreateUserApiKeyMutationPayload(
+            jwt=token,
+            api_key=UserApiKey(
+                id_attr=int(token_id),
+                name=input.name,
+                description=input.description or None,
+                created_at=issued_at,
+                expires_at=input.expires_at or None,
+                user_id=int(user.identity),
+            ),
+            query=Query(),
+        )
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly, IsAdmin])  # type: ignore
     async def delete_system_api_key(
         self, info: Info[Context, None], input: DeleteApiKeyInput
-    ) -> DeleteSystemApiKeyMutationPayload:
+    ) -> DeleteApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
         api_key_id = from_global_id_with_expected_type(
             input.id, expected_type_name=SystemApiKey.__name__
         )
+        await token_store.revoke(ApiKeyId(api_key_id))
+        return DeleteApiKeyMutationPayload(apiKeyId=input.id, query=Query())
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def delete_user_api_key(
+        self, info: Info[Context, None], input: DeleteApiKeyInput
+    ) -> DeleteApiKeyMutationPayload:
+        assert (token_store := info.context.token_store) is not None
+        api_key_id = from_global_id_with_expected_type(
+            input.id, expected_type_name=UserApiKey.__name__
+        )
         async with info.context.db() as session:
-            api_key = await session.get(models.APIKey, api_key_id)
+            api_key = await session.scalar(
+                select(models.ApiKey).where(models.ApiKey.id == api_key_id)
+            )
             if api_key is None:
-                raise NotFound(f"Unknown System API Key: {input.id}")
-
-            await session.delete(api_key)
-
-        return DeleteSystemApiKeyMutationPayload(id=input.id, query=Query())
-
-
-def create_jwt(
-    *,
-    secret: str,
-    algorithm: str = "HS256",
-    name: str,
-    description: Optional[str],
-    iat: datetime,
-    exp: Optional[datetime],
-    id: int,
-) -> str:
-    """Create a signed JSON Web Token for authentication
-
-    Args:
-        secret (str): the secret to sign with
-        name (str): name of the key / token
-        description (Optional[str]): description of the token
-        iat (datetime): the issued at time
-        exp (Optional[datetime]): the expiry, if set
-        id (int): the id of the key
-        algorithm (str, optional): the algorithm to use. Defaults to "HS256".
-
-    Returns:
-        str: The encoded JWT
-    """
-    payload: Dict[str, Any] = {
-        "name": name,
-        "description": description,
-        "iat": iat.utcnow(),
-        "id": id,
-    }
-    if exp is not None:
-        payload["exp"] = exp.utcnow()
-
-    # Encode the payload to create the JWT
-    token = jwt.encode(payload, secret, algorithm=algorithm)
-
-    return token
+                raise ValueError(f"API key with id {input.id} not found")
+        if int((user := info.context.user).identity) != api_key.user_id and not user.is_admin:
+            raise Unauthorized("User not authorized to delete")
+        await token_store.revoke(ApiKeyId(api_key_id))
+        return DeleteApiKeyMutationPayload(apiKeyId=input.id, query=Query())

phoenix/server/api/mutations/dataset_mutations.py

@@ -12,6 +12,7 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.db.helpers import get_eval_trace_ids_for_datasets, get_project_names_for_datasets
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, NotFound
 from phoenix.server.api.helpers.dataset_helpers import (
@@ -28,7 +29,6 @@ from phoenix.server.api.input_types.PatchDatasetExamplesInput import (
     PatchDatasetExamplesInput,
 )
 from phoenix.server.api.input_types.PatchDatasetInput import PatchDatasetInput
-from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
 from phoenix.server.api.types.node import from_global_id_with_expected_type
@@ -44,7 +44,7 @@ class DatasetMutationPayload:
 
 @strawberry.type
 class DatasetMutationMixin:
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def create_dataset(
         self,
         info: Info[Context, None],
@@ -67,7 +67,7 @@ class DatasetMutationMixin:
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_dataset(
         self,
         info: Info[Context, None],
@@ -96,7 +96,7 @@ class DatasetMutationMixin:
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def add_spans_to_dataset(
         self,
         info: Info[Context, None],
@@ -225,7 +225,7 @@ class DatasetMutationMixin:
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def add_examples_to_dataset(
         self, info: Info[Context, None], input: AddExamplesToDatasetInput
     ) -> DatasetMutationPayload:
@@ -351,7 +351,7 @@ class DatasetMutationMixin:
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_dataset(
         self,
         info: Info[Context, None],
@@ -382,7 +382,7 @@ class DatasetMutationMixin:
         info.context.event_queue.put(DatasetDeleteEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_dataset_examples(
         self,
         info: Info[Context, None],
@@ -474,7 +474,7 @@ class DatasetMutationMixin:
         info.context.event_queue.put(DatasetInsertEvent((dataset.id,)))
         return DatasetMutationPayload(dataset=to_gql_dataset(dataset))
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_dataset_examples(
         self, info: Info[Context, None], input: DeleteDatasetExamplesInput
     ) -> DatasetMutationPayload:

phoenix/server/api/mutations/experiment_mutations.py

@@ -8,10 +8,10 @@ from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.db.helpers import get_eval_trace_ids_for_experiments, get_project_names_for_experiments
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import CustomGraphQLError
 from phoenix.server.api.input_types.DeleteExperimentsInput import DeleteExperimentsInput
-from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.utils import delete_projects, delete_traces
@@ -25,7 +25,7 @@ class ExperimentMutationPayload:
 
 @strawberry.type
 class ExperimentMutationMixin:
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_experiments(
         self,
         info: Info[Context, None],

phoenix/server/api/mutations/export_events_mutations.py

@@ -8,9 +8,9 @@ from strawberry import ID, UNSET
 from strawberry.types import Info
 
 import phoenix.core.model_schema as ms
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.ClusterInput import ClusterInput
-from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.types.Event import parse_event_ids_by_inferences_role, unpack_event_id
 from phoenix.server.api.types.ExportedFile import ExportedFile
 from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
@@ -19,7 +19,7 @@ from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, Inf
 @strawberry.type
 class ExportEventsMutationMixin:
     @strawberry.mutation(
-        permission_classes=[IsAuthenticated],
+        permission_classes=[IsNotReadOnly],
         description=(
             "Given a list of event ids, export the corresponding data subset in Parquet format."
             " File name is optional, but if specified, should be without file extension. By default"
@@ -51,7 +51,7 @@ class ExportEventsMutationMixin:
         return ExportedFile(file_name=file_name)
 
     @strawberry.mutation(
-        permission_classes=[IsAuthenticated],
+        permission_classes=[IsNotReadOnly],
         description=(
             "Given a list of clusters, export the corresponding data subset in Parquet format."
             " File name is optional, but if specified, should be without file extension. By default"

phoenix/server/api/mutations/project_mutations.py

@@ -6,9 +6,9 @@ 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.context import Context
 from phoenix.server.api.input_types.ClearProjectInput import ClearProjectInput
-from phoenix.server.api.mutations.auth import IsAuthenticated
 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
@@ -16,7 +16,7 @@ from phoenix.server.dml_event import ProjectDeleteEvent, SpanDeleteEvent
 
 @strawberry.type
 class ProjectMutationMixin:
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # 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 +33,7 @@ class ProjectMutationMixin:
         info.context.event_queue.put(ProjectDeleteEvent((project_id,)))
         return Query()
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # 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"

phoenix/server/api/mutations/span_annotations_mutations.py

@@ -6,11 +6,11 @@ from strawberry import UNSET
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.CreateSpanAnnotationInput import CreateSpanAnnotationInput
 from phoenix.server.api.input_types.DeleteAnnotationsInput import DeleteAnnotationsInput
 from phoenix.server.api.input_types.PatchAnnotationInput import PatchAnnotationInput
-from phoenix.server.api.mutations.auth import IsAuthenticated
 from phoenix.server.api.queries import Query
 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
@@ -25,7 +25,7 @@ class SpanAnnotationMutationPayload:
 
 @strawberry.type
 class SpanAnnotationMutationMixin:
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def create_span_annotations(
         self, info: Info[Context, None], input: List[CreateSpanAnnotationInput]
     ) -> SpanAnnotationMutationPayload:
@@ -59,7 +59,7 @@ class SpanAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_span_annotations(
         self, info: Info[Context, None], input: List[PatchAnnotationInput]
     ) -> SpanAnnotationMutationPayload:
@@ -99,7 +99,7 @@ class SpanAnnotationMutationMixin:
                     info.context.event_queue.put(SpanAnnotationInsertEvent((span_annotation.id,)))
         return SpanAnnotationMutationPayload(span_annotations=patched_annotations, query=Query())
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_span_annotations(
         self, info: Info[Context, None], input: DeleteAnnotationsInput
     ) -> SpanAnnotationMutationPayload:

phoenix/server/api/mutations/trace_annotations_mutations.py

@@ -6,11 +6,11 @@ from strawberry import UNSET
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
 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.mutations.auth import IsAuthenticated
 from phoenix.server.api.queries import Query
 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
@@ -25,7 +25,7 @@ class TraceAnnotationMutationPayload:
 
 @strawberry.type
 class TraceAnnotationMutationMixin:
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def create_trace_annotations(
         self, info: Info[Context, None], input: List[CreateTraceAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
@@ -59,7 +59,7 @@ class TraceAnnotationMutationMixin:
             query=Query(),
         )
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def patch_trace_annotations(
         self, info: Info[Context, None], input: List[PatchAnnotationInput]
     ) -> TraceAnnotationMutationPayload:
@@ -98,7 +98,7 @@ class TraceAnnotationMutationMixin:
                     info.context.event_queue.put(TraceAnnotationInsertEvent((trace_annotation.id,)))
         return TraceAnnotationMutationPayload(trace_annotations=patched_annotations, query=Query())
 
-    @strawberry.mutation(permission_classes=[IsAuthenticated])  # type: ignore
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_trace_annotations(
         self, info: Info[Context, None], input: DeleteAnnotationsInput
     ) -> TraceAnnotationMutationPayload: