arize-phoenix 4.23.0__py3-none-any.whl → 4.25.0__py3-none-any.whl

phoenix/server/api/mutations/__init__.py

@@ -1,5 +1,6 @@
 import strawberry
 
+from phoenix.server.api.mutations.api_key_mutations import ApiKeyMutationMixin
 from phoenix.server.api.mutations.dataset_mutations import DatasetMutationMixin
 from phoenix.server.api.mutations.experiment_mutations import ExperimentMutationMixin
 from phoenix.server.api.mutations.export_events_mutations import ExportEventsMutationMixin
@@ -16,5 +17,6 @@ class Mutation(
     ExportEventsMutationMixin,
     SpanAnnotationMutationMixin,
     TraceAnnotationMutationMixin,
+    ApiKeyMutationMixin,
 ):
     pass

phoenix/server/api/mutations/api_key_mutations.py

@@ -0,0 +1,119 @@
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import jwt
+import strawberry
+from sqlalchemy import insert, select
+from strawberry import UNSET
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.mutations.auth import HasSecret, IsAuthenticated
+from phoenix.server.api.queries import Query
+from phoenix.server.api.types.SystemApiKey import SystemApiKey
+
+
+@strawberry.type
+class CreateSystemApiKeyMutationPayload:
+    jwt: str
+    api_key: SystemApiKey
+    query: Query
+
+
+@strawberry.input
+class CreateApiKeyInput:
+    name: str
+    description: Optional[str] = UNSET
+    expires_at: Optional[datetime] = UNSET
+
+
+@strawberry.type
+class ApiKeyMutationMixin:
+    @strawberry.mutation(permission_classes=[HasSecret, IsAuthenticated])  # 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
+        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
+                .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,
+        )
+        return CreateSystemApiKeyMutationPayload(
+            jwt=encoded_jwt,
+            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,
+            ),
+            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

phoenix/server/api/mutations/auth.py

@@ -9,3 +9,10 @@ class IsAuthenticated(BasePermission):
 
     async def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
         return not info.context.read_only
+
+
+class HasSecret(BasePermission):
+    message = "Application secret is not set"
+
+    async def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        return info.context.secret is not None

phoenix/server/api/queries.py

@@ -66,8 +66,11 @@ from phoenix.server.api.types.pagination import (
 from phoenix.server.api.types.Project import Project
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span, to_gql_span
+from phoenix.server.api.types.SystemApiKey import SystemApiKey
 from phoenix.server.api.types.Trace import Trace
 from phoenix.server.api.types.User import User
+from phoenix.server.api.types.UserApiKey import UserApiKey
+from phoenix.server.api.types.UserRole import UserRole
 
 
 @strawberry.type
@@ -87,7 +90,13 @@ class Query:
             last=last,
             before=before if isinstance(before, CursorString) else None,
         )
-        stmt = select(models.User).order_by(models.User.email)
+        stmt = (
+            select(models.User)
+            .join(models.UserRole)
+            .where(models.UserRole.name != "SYSTEM")
+            .order_by(models.User.email)
+            .options(joinedload(models.User.role))
+        )
         async with info.context.db() as session:
             users = await session.stream_scalars(stmt)
             data = [
@@ -96,11 +105,77 @@ class Query:
                     email=user.email,
                     username=user.username,
                     created_at=user.created_at,
+                    role=UserRole(
+                        id_attr=user.role.id,
+                        name=user.role.name,
+                    ),
                 )
                 async for user in users
             ]
         return connection_from_list(data=data, args=args)
 
+    @strawberry.field
+    async def user_roles(
+        self,
+        info: Info[Context, None],
+    ) -> List[UserRole]:
+        async with info.context.db() as session:
+            roles = await session.scalars(
+                select(models.UserRole).where(models.UserRole.name != "SYSTEM")
+            )
+        return [
+            UserRole(
+                id_attr=role.id,
+                name=role.name,
+            )
+            for role in roles
+        ]
+
+    @strawberry.field
+    async def user_api_keys(self, info: Info[Context, None]) -> List[UserApiKey]:
+        # TODO(auth): add access control
+        stmt = (
+            select(models.APIKey)
+            .join(models.User)
+            .join(models.UserRole)
+            .where(models.UserRole.name != "SYSTEM")
+        )
+        async with info.context.db() as session:
+            api_keys = await session.scalars(stmt)
+        return [
+            UserApiKey(
+                id_attr=api_key.id,
+                user_id=api_key.user_id,
+                name=api_key.name,
+                description=api_key.description,
+                created_at=api_key.created_at,
+                expires_at=api_key.expires_at,
+            )
+            for api_key in api_keys
+        ]
+
+    @strawberry.field
+    async def system_api_keys(self, info: Info[Context, None]) -> List[SystemApiKey]:
+        # TODO(auth): add access control
+        stmt = (
+            select(models.APIKey)
+            .join(models.User)
+            .join(models.UserRole)
+            .where(models.UserRole.name == "SYSTEM")
+        )
+        async with info.context.db() as session:
+            api_keys = await session.scalars(stmt)
+        return [
+            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,
+            )
+            for api_key in api_keys
+        ]
+
     @strawberry.field
     async def projects(
         self,

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

@@ -71,6 +71,7 @@ from .utils import (
 )
 
 logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
 
 DATASET_NODE_NAME = DatasetNodeType.__name__
 DATASET_VERSION_NODE_NAME = DatasetVersionNodeType.__name__

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

@@ -196,7 +196,7 @@ class AnnotateSpansResponseBody(ResponseBody[List[InsertedSpanAnnotation]]):
 async def annotate_spans(
     request: Request,
     request_body: AnnotateSpansRequestBody,
-    sync: bool = Query(default=True, description="If true, fulfill request synchronously."),
+    sync: bool = Query(default=False, description="If true, fulfill request synchronously."),
 ) -> AnnotateSpansResponseBody:
     if not request_body.data:
         return AnnotateSpansResponseBody(data=[])

phoenix/server/api/types/ApiKey.py

@@ -0,0 +1,16 @@
+from datetime import datetime
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.interface
+class ApiKey:
+    name: str = strawberry.field(description="Name of the API key.")
+    description: Optional[str] = strawberry.field(description="Description of the API key.")
+    created_at: datetime = strawberry.field(
+        description="The date and time the API key was created."
+    )
+    expires_at: Optional[datetime] = strawberry.field(
+        description="The date and time the API key will expire."
+    )

phoenix/server/api/types/SystemApiKey.py

@@ -0,0 +1,9 @@
+import strawberry
+from strawberry.relay import Node, NodeID
+
+from .ApiKey import ApiKey
+
+
+@strawberry.type
+class SystemApiKey(ApiKey, Node):
+    id_attr: NodeID[int]

phoenix/server/api/types/User.py

@@ -4,6 +4,8 @@ from typing import Optional
 import strawberry
 from strawberry.relay import Node, NodeID
 
+from .UserRole import UserRole
+
 
 @strawberry.type
 class User(Node):
@@ -11,3 +13,4 @@ class User(Node):
     email: str
     username: Optional[str]
     created_at: datetime
+    role: UserRole

phoenix/server/api/types/UserApiKey.py

@@ -0,0 +1,11 @@
+import strawberry
+from strawberry import Private
+from strawberry.relay.types import Node, NodeID
+
+from .ApiKey import ApiKey
+
+
+@strawberry.type
+class UserApiKey(ApiKey, Node):
+    id_attr: NodeID[int]
+    user_id: Private[int]

phoenix/server/api/types/UserRole.py

@@ -0,0 +1,8 @@
+import strawberry
+from strawberry.relay import Node, NodeID
+
+
+@strawberry.type
+class UserRole(Node):
+    id_attr: NodeID[int]
+    name: str

phoenix/server/app.py

@@ -100,6 +100,7 @@ if TYPE_CHECKING:
     from opentelemetry.trace import TracerProvider
 
 logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
 
 router = APIRouter(include_in_schema=False)
 
@@ -229,7 +230,8 @@ def _lifespan(
     dml_event_handler: DmlEventHandler,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
-    clean_ups: Iterable[Callable[[], None]] = (),
+    startup_callbacks: Iterable[Callable[[], None]] = (),
+    shutdown_callbacks: Iterable[Callable[[], None]] = (),
     read_only: bool = False,
 ) -> StatefulLifespan[FastAPI]:
     @contextlib.asynccontextmanager
@@ -247,6 +249,8 @@ def _lifespan(
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
         ), dml_event_handler:
+            for callback in startup_callbacks:
+                callback()
             yield {
                 "event_queue": dml_event_handler,
                 "enqueue": enqueue,
@@ -254,8 +258,8 @@ def _lifespan(
                 "queue_evaluation_for_bulk_insert": queue_evaluation,
                 "enqueue_operation": enqueue_operation,
             }
-        for clean_up in clean_ups:
-            clean_up()
+        for callback in shutdown_callbacks:
+            callback()
 
     return lifespan
 
@@ -276,7 +280,26 @@ def create_graphql_router(
     cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
     event_queue: CanPutItem[DmlEvent],
     read_only: bool = False,
+    secret: Optional[str] = None,
 ) -> GraphQLRouter:  # type: ignore[type-arg]
+    """Creates the GraphQL router.
+
+    Args:
+        schema (BaseSchema): The GraphQL schema.
+        db (DbSessionFactory): The database session factory pointing to a SQL database.
+        model (Model): The Model representing inferences (legacy)
+        export_path (Path): the file path to export data to for download (legacy)
+        last_updated_at (CanGetLastUpdatedAt): How to get the last updated timestamp for updates.
+        event_queue (CanPutItem[DmlEvent]): The event queue for DML events.
+        corpus (Optional[Model], optional): the corpus for UMAP projection. Defaults to None.
+        cache_for_dataloaders (Optional[CacheForDataLoaders], optional): GraphQL data loaders.
+        read_only (bool, optional): Marks the app as read-only. Defaults to False.
+        secret (Optional[str], optional): The application secret for auth. Defaults to None.
+
+    Returns:
+        GraphQLRouter: The router mounted at /graphql
+    """
+
     def get_context() -> Context:
         return Context(
             db=db,
@@ -336,6 +359,7 @@ def create_graphql_router(
             ),
             cache_for_dataloaders=cache_for_dataloaders,
             read_only=read_only,
+            secret=secret,
         )
 
     return GraphQLRouter(
@@ -408,9 +432,12 @@ def create_app(
     initial_spans: Optional[Iterable[Union[Span, Tuple[Span, str]]]] = None,
     initial_evaluations: Optional[Iterable[pb.Evaluation]] = None,
     serve_ui: bool = True,
-    clean_up_callbacks: List[Callable[[], None]] = [],
+    startup_callbacks: Iterable[Callable[[], None]] = (),
+    shutdown_callbacks: Iterable[Callable[[], None]] = (),
+    secret: Optional[str] = None,
 ) -> FastAPI:
-    clean_ups: List[Callable[[], None]] = clean_up_callbacks  # To be called at app shutdown.
+    startup_callbacks_list: List[Callable[[], None]] = list(startup_callbacks)
+    shutdown_callbacks_list: List[Callable[[], None]] = list(shutdown_callbacks)
     initial_batch_of_spans: Iterable[Tuple[Span, str]] = (
         ()
         if initial_spans is None
@@ -472,6 +499,7 @@ def create_app(
         event_queue=dml_event_handler,
         cache_for_dataloaders=cache_for_dataloaders,
         read_only=read_only,
+        secret=secret,
     )
     if enable_prometheus:
         from phoenix.server.prometheus import PrometheusMiddleware
@@ -489,7 +517,8 @@ def create_app(
             dml_event_handler=dml_event_handler,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
-            clean_ups=clean_ups,
+            shutdown_callbacks=shutdown_callbacks_list,
+            startup_callbacks=startup_callbacks_list,
         ),
         middleware=[
             Middleware(HeadersMiddleware),
@@ -532,5 +561,5 @@ def create_app(
 
         FastAPIInstrumentor().instrument(tracer_provider=tracer_provider)
         FastAPIInstrumentor.instrument_app(app, tracer_provider=tracer_provider)
-        clean_ups.append(FastAPIInstrumentor().uninstrument)
+        shutdown_callbacks_list.append(FastAPIInstrumentor().uninstrument)
     return app

phoenix/server/main.py

@@ -1,13 +1,16 @@
 import atexit
+import codecs
 import logging
 import os
+import sys
 from argparse import ArgumentParser
-from pathlib import Path, PosixPath
+from importlib.metadata import version
+from pathlib import Path
 from threading import Thread
 from time import sleep, time
 from typing import List, Optional
+from urllib.parse import urljoin
 
-import pkg_resources
 from uvicorn import Config, Server
 
 import phoenix.trace.v1 as pb
@@ -53,6 +56,7 @@ from phoenix.trace.otel import decode_otlp_span, encode_span_to_otlp
 from phoenix.trace.schemas import Span
 
 logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
 
 _WELCOME_MESSAGE = """
 
@@ -137,6 +141,7 @@ if __name__ == "__main__":
     parser.add_argument("--debug", action="store_true")
     # Whether the app is running in a development environment
     parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--no-ui", action="store_true")
     subparsers = parser.add_subparsers(dest="command", required=True)
     serve_parser = subparsers.add_parser("serve")
     datasets_parser = subparsers.add_parser("datasets")
@@ -218,7 +223,7 @@ if __name__ == "__main__":
         reference_inferences,
     )
 
-    authentication_enabled, auth_secret = get_auth_settings()
+    authentication_enabled, secret = get_auth_settings()
 
     fixture_spans: List[Span] = []
     fixture_evals: List[pb.Evaluation] = []
@@ -255,6 +260,18 @@ if __name__ == "__main__":
     engine = create_engine_and_run_migrations(db_connection_str)
     instrumentation_cleanups = instrument_engine_if_enabled(engine)
     factory = DbSessionFactory(db=_db(engine), dialect=engine.dialect.name)
+    # Print information about the server
+    msg = _WELCOME_MESSAGE.format(
+        version=version("arize-phoenix"),
+        ui_path=urljoin(f"http://{host}:{port}", host_root_path),
+        grpc_path=f"http://{host}:{get_env_grpc_port()}",
+        http_path=urljoin(urljoin(f"http://{host}:{port}", host_root_path), "v1/traces"),
+        storage=get_printable_db_url(db_connection_str),
+    )
+    if authentication_enabled:
+        msg += _EXPERIMENTAL_WARNING.format(auth_enabled=True)
+    if sys.platform.startswith("win"):
+        msg = codecs.encode(msg, "ascii", errors="ignore").decode("ascii").strip()
     app = create_app(
         db=factory,
         export_path=export_path,
@@ -266,29 +283,17 @@ if __name__ == "__main__":
         else create_model_from_inferences(corpus_inferences),
         debug=args.debug,
         dev=args.dev,
+        serve_ui=not args.no_ui,
         read_only=read_only,
         enable_prometheus=enable_prometheus,
         initial_spans=fixture_spans,
         initial_evaluations=fixture_evals,
-        clean_up_callbacks=instrumentation_cleanups,
+        startup_callbacks=[lambda: print(msg)],
+        shutdown_callbacks=instrumentation_cleanups,
+        secret=secret,
     )
     server = Server(config=Config(app, host=host, port=port, root_path=host_root_path))  # type: ignore
     Thread(target=_write_pid_file_when_ready, args=(server,), daemon=True).start()
 
-    # Print information about the server
-    phoenix_version = pkg_resources.get_distribution("arize-phoenix").version
-    print(
-        _WELCOME_MESSAGE.format(
-            version=phoenix_version,
-            ui_path=PosixPath(f"http://{host}:{port}", host_root_path),
-            grpc_path=f"http://{host}:{get_env_grpc_port()}",
-            http_path=PosixPath(f"http://{host}:{port}", host_root_path, "v1/traces"),
-            storage=get_printable_db_url(db_connection_str),
-        )
-    )
-
-    if authentication_enabled:
-        print(_EXPERIMENTAL_WARNING.format(auth_enabled=authentication_enabled))
-
     # Start the server
     server.run()

phoenix/server/static/.vite/manifest.json

@@ -1,32 +1,32 @@
 {
-  "_components-DBYPF96c.js": {
-    "file": "assets/components-DBYPF96c.js",
+  "_components-1Ahruijo.js": {
+    "file": "assets/components-1Ahruijo.js",
     "name": "components",
     "imports": [
-      "_vendor-CIqy43_9.js",
-      "_vendor-arizeai-B1YgcWL8.js",
-      "_pages-BhOnrUmC.js",
+      "_vendor-aSQri0vz.js",
+      "_vendor-arizeai-CsdcB1NH.js",
+      "_pages-CFS6mPnW.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-codemirror-_bcwCA1C.js"
+      "_vendor-codemirror-CYHkhs7D.js"
     ]
   },
-  "_pages-BhOnrUmC.js": {
-    "file": "assets/pages-BhOnrUmC.js",
+  "_pages-CFS6mPnW.js": {
+    "file": "assets/pages-CFS6mPnW.js",
     "name": "pages",
     "imports": [
-      "_vendor-CIqy43_9.js",
-      "_components-DBYPF96c.js",
-      "_vendor-arizeai-B1YgcWL8.js",
-      "_vendor-recharts-C3pM_Wlg.js",
-      "_vendor-codemirror-_bcwCA1C.js"
+      "_vendor-aSQri0vz.js",
+      "_components-1Ahruijo.js",
+      "_vendor-arizeai-CsdcB1NH.js",
+      "_vendor-recharts-B0sannek.js",
+      "_vendor-codemirror-CYHkhs7D.js"
     ]
   },
   "_vendor-!~{003}~.js": {
     "file": "assets/vendor-DxkFTwjz.css",
     "src": "_vendor-!~{003}~.js"
   },
-  "_vendor-CIqy43_9.js": {
-    "file": "assets/vendor-CIqy43_9.js",
+  "_vendor-aSQri0vz.js": {
+    "file": "assets/vendor-aSQri0vz.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-DwGkEfCM.js"
@@ -35,25 +35,25 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-B1YgcWL8.js": {
-    "file": "assets/vendor-arizeai-B1YgcWL8.js",
+  "_vendor-arizeai-CsdcB1NH.js": {
+    "file": "assets/vendor-arizeai-CsdcB1NH.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-CIqy43_9.js"
+      "_vendor-aSQri0vz.js"
     ]
   },
-  "_vendor-codemirror-_bcwCA1C.js": {
-    "file": "assets/vendor-codemirror-_bcwCA1C.js",
+  "_vendor-codemirror-CYHkhs7D.js": {
+    "file": "assets/vendor-codemirror-CYHkhs7D.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-CIqy43_9.js"
+      "_vendor-aSQri0vz.js"
     ]
   },
-  "_vendor-recharts-C3pM_Wlg.js": {
-    "file": "assets/vendor-recharts-C3pM_Wlg.js",
+  "_vendor-recharts-B0sannek.js": {
+    "file": "assets/vendor-recharts-B0sannek.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-CIqy43_9.js"
+      "_vendor-aSQri0vz.js"
     ]
   },
   "_vendor-three-DwGkEfCM.js": {
@@ -61,18 +61,18 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-DNxu4viw.js",
+    "file": "assets/index-BEE_RWJx.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-CIqy43_9.js",
-      "_vendor-arizeai-B1YgcWL8.js",
-      "_pages-BhOnrUmC.js",
-      "_components-DBYPF96c.js",
+      "_vendor-aSQri0vz.js",
+      "_vendor-arizeai-CsdcB1NH.js",
+      "_pages-CFS6mPnW.js",
+      "_components-1Ahruijo.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-recharts-C3pM_Wlg.js",
-      "_vendor-codemirror-_bcwCA1C.js"
+      "_vendor-recharts-B0sannek.js",
+      "_vendor-codemirror-CYHkhs7D.js"
     ]
   }
 }