arize-phoenix 4.36.0__py3-none-any.whl → 5.1.0__py3-none-any.whl

phoenix/server/app.py

@@ -20,21 +20,23 @@ from typing import (
     List,
     NamedTuple,
     Optional,
+    Sequence,
     Tuple,
+    TypedDict,
     Union,
     cast,
 )
 
 import strawberry
-from fastapi import APIRouter, FastAPI
+from fastapi import APIRouter, Depends, FastAPI
 from fastapi.middleware.gzip import GZipMiddleware
-from fastapi.responses import FileResponse
 from fastapi.utils import is_body_allowed_for_status_code
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
 from starlette.datastructures import State as StarletteState
 from starlette.exceptions import HTTPException
 from starlette.middleware import Middleware
+from starlette.middleware.authentication import AuthenticationMiddleware
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
 from starlette.responses import PlainTextResponse, Response
@@ -50,6 +52,7 @@ import phoenix.trace.v1 as pb
 from phoenix.config import (
     DEFAULT_PROJECT_NAME,
     SERVER_DIR,
+    OAuth2ClientConfig,
     get_env_host,
     get_env_port,
     server_instrumentation_is_enabled,
@@ -58,6 +61,7 @@ from phoenix.core.model_schema import Model
 from phoenix.db import models
 from phoenix.db.bulk_inserter import BulkInserter
 from phoenix.db.engines import create_engine
+from phoenix.db.facilitator import Facilitator
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.exceptions import PhoenixMigrationError
 from phoenix.pointcloud.umap_parameters import UMAPParameters
@@ -86,13 +90,24 @@ from phoenix.server.api.dataloaders import (
     SpanProjectsDataLoader,
     TokenCountDataLoader,
     TraceRowIdsDataLoader,
+    UserRolesDataLoader,
+    UsersDataLoader,
+)
+from phoenix.server.api.routers import (
+    auth_router,
+    create_embeddings_router,
+    create_v1_router,
+    oauth2_router,
 )
 from phoenix.server.api.routers.v1 import REST_API_VERSION
-from phoenix.server.api.routers.v1 import router as v1_router
 from phoenix.server.api.schema import schema
+from phoenix.server.bearer_auth import BearerTokenAuthBackend, is_authenticated
 from phoenix.server.dml_event import DmlEvent
 from phoenix.server.dml_event_handler import DmlEventHandler
+from phoenix.server.email.types import EmailSender
 from phoenix.server.grpc_server import GrpcServer
+from phoenix.server.jwt_store import JwtStore
+from phoenix.server.oauth2 import OAuth2Clients
 from phoenix.server.telemetry import initialize_opentelemetry_tracer_provider
 from phoenix.server.types import (
     CanGetLastUpdatedAt,
@@ -100,6 +115,7 @@ from phoenix.server.types import (
     DaemonTask,
     DbSessionFactory,
     LastUpdatedAt,
+    TokenStore,
 )
 from phoenix.trace.fixtures import (
     TracesFixture,
@@ -135,6 +151,11 @@ ProjectName: TypeAlias = str
 _Callback: TypeAlias = Callable[[], Union[None, Awaitable[None]]]
 
 
+class OAuth2Idp(TypedDict):
+    name: str
+    displayName: str
+
+
 class AppConfig(NamedTuple):
     has_inferences: bool
     """ Whether the model has inferences (e.g. a primary dataset) """
@@ -146,6 +167,7 @@ class AppConfig(NamedTuple):
     web_manifest_path: Path
     authentication_enabled: bool
     """ Whether authentication is enabled """
+    oauth2_idps: Sequence[OAuth2Idp]
 
 
 class Static(StaticFiles):
@@ -194,6 +216,7 @@ class Static(StaticFiles):
                     "is_development": self._app_config.is_development,
                     "manifest": self._web_manifest,
                     "authentication_enabled": self._app_config.authentication_enabled,
+                    "oauth2_idps": self._app_config.oauth2_idps,
                 },
             )
         except Exception as e:
@@ -218,18 +241,6 @@ class HeadersMiddleware(BaseHTTPMiddleware):
 ProjectRowId: TypeAlias = int
 
 
-@router.get("/exports")
-async def download_exported_file(request: Request, filename: str) -> FileResponse:
-    file = request.app.state.export_path / (filename + ".parquet")
-    if not file.is_file():
-        raise HTTPException(status_code=404)
-    return FileResponse(
-        path=file,
-        filename=file.name,
-        media_type="application/x-octet-stream",
-    )
-
-
 @router.get("/arize_phoenix_version")
 async def version() -> PlainTextResponse:
     return PlainTextResponse(f"{phoenix.__version__}")
@@ -238,13 +249,15 @@ async def version() -> PlainTextResponse:
 DB_MUTEX: Optional[asyncio.Lock] = None
 
 
-def _db(engine: AsyncEngine) -> Callable[[], AsyncContextManager[AsyncSession]]:
+def _db(
+    engine: AsyncEngine, bypass_lock: bool = False
+) -> Callable[[], AsyncContextManager[AsyncSession]]:
     Session = async_sessionmaker(engine, expire_on_commit=False)
 
     @contextlib.asynccontextmanager
     async def factory() -> AsyncIterator[AsyncSession]:
         async with contextlib.AsyncExitStack() as stack:
-            if DB_MUTEX:
+            if not bypass_lock and DB_MUTEX:
                 await stack.enter_async_context(DB_MUTEX)
             yield await stack.enter_async_context(Session.begin())
 
@@ -283,9 +296,6 @@ class Scaffolder(DaemonTask):
             return
         await self.start()
 
-    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
-        await self.stop()
-
     async def _run(self) -> None:
         """
         Main entry point for Scaffolder.
@@ -374,6 +384,7 @@ def _lifespan(
     db: DbSessionFactory,
     bulk_inserter: BulkInserter,
     dml_event_handler: DmlEventHandler,
+    token_store: Optional[TokenStore] = None,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
     startup_callbacks: Iterable[_Callback] = (),
@@ -400,6 +411,7 @@ def _lifespan(
                 disabled=read_only,
                 tracer_provider=tracer_provider,
                 enable_prometheus=enable_prometheus,
+                token_store=token_store,
             )
             await stack.enter_async_context(grpc_server)
             await stack.enter_async_context(dml_event_handler)
@@ -410,6 +422,8 @@ def _lifespan(
                     queue_evaluation=queue_evaluation,
                 )
                 await stack.enter_async_context(scaffolder)
+            if isinstance(token_store, AsyncContextManager):
+                await stack.enter_async_context(token_store)
             yield {
                 "event_queue": dml_event_handler,
                 "enqueue": enqueue,
@@ -436,11 +450,13 @@ def create_graphql_router(
     model: Model,
     export_path: Path,
     last_updated_at: CanGetLastUpdatedAt,
+    authentication_enabled: bool,
     corpus: Optional[Model] = None,
     cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
     event_queue: CanPutItem[DmlEvent],
     read_only: bool = False,
     secret: Optional[str] = None,
+    token_store: Optional[TokenStore] = None,
 ) -> GraphQLRouter:  # type: ignore[type-arg]
     """Creates the GraphQL router.
 
@@ -450,6 +466,7 @@ def create_graphql_router(
         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.
+        authentication_enabled (bool): Whether authentication is enabled.
         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.
@@ -521,18 +538,23 @@ def create_graphql_router(
                 ),
                 trace_row_ids=TraceRowIdsDataLoader(db),
                 project_by_name=ProjectByNameDataLoader(db),
+                users=UsersDataLoader(db),
+                user_roles=UserRolesDataLoader(db),
             ),
             cache_for_dataloaders=cache_for_dataloaders,
             read_only=read_only,
+            auth_enabled=authentication_enabled,
             secret=secret,
+            token_store=token_store,
         )
 
     return GraphQLRouter(
         schema,
-        graphiql=True,
+        graphql_ide="graphiql",
         context_getter=get_context,
         include_in_schema=False,
         prefix="/graphql",
+        dependencies=(Depends(is_authenticated),) if authentication_enabled else (),
     )
 
 
@@ -600,11 +622,19 @@ def create_app(
     startup_callbacks: Iterable[_Callback] = (),
     shutdown_callbacks: Iterable[_Callback] = (),
     secret: Optional[str] = None,
+    password_reset_token_expiry: Optional[timedelta] = None,
+    access_token_expiry: Optional[timedelta] = None,
+    refresh_token_expiry: Optional[timedelta] = None,
     scaffolder_config: Optional[ScaffolderConfig] = None,
+    email_sender: Optional[EmailSender] = None,
+    oauth2_client_configs: Optional[List[OAuth2ClientConfig]] = None,
+    bulk_inserter_factory: Optional[Callable[..., BulkInserter]] = None,
 ) -> FastAPI:
     logger.info(f"Server umap params: {umap_params}")
+    bulk_inserter_factory = bulk_inserter_factory or BulkInserter
     startup_callbacks_list: List[_Callback] = list(startup_callbacks)
     shutdown_callbacks_list: List[_Callback] = list(shutdown_callbacks)
+    startup_callbacks_list.append(Facilitator(db=db))
     initial_batch_of_spans: Iterable[Tuple[Span, str]] = (
         ()
         if initial_spans is None
@@ -619,12 +649,22 @@ def create_app(
     )
     last_updated_at = LastUpdatedAt()
     middlewares: List[Middleware] = [Middleware(HeadersMiddleware)]
+    if authentication_enabled and secret:
+        token_store = JwtStore(db, secret)
+        middlewares.append(
+            Middleware(
+                AuthenticationMiddleware,
+                backend=BearerTokenAuthBackend(token_store),
+            )
+        )
+    else:
+        token_store = None
     dml_event_handler = DmlEventHandler(
         db=db,
         cache_for_dataloaders=cache_for_dataloaders,
         last_updated_at=last_updated_at,
     )
-    bulk_inserter = BulkInserter(
+    bulk_inserter = bulk_inserter_factory(
         db,
         enable_prometheus=enable_prometheus,
         event_queue=dml_event_handler,
@@ -662,12 +702,14 @@ def create_app(
         ),
         model=model,
         corpus=corpus,
+        authentication_enabled=authentication_enabled,
         export_path=export_path,
         last_updated_at=last_updated_at,
         event_queue=dml_event_handler,
         cache_for_dataloaders=cache_for_dataloaders,
         read_only=read_only,
         secret=secret,
+        token_store=token_store,
     )
     if enable_prometheus:
         from phoenix.server.prometheus import PrometheusMiddleware
@@ -681,6 +723,7 @@ def create_app(
             read_only=read_only,
             bulk_inserter=bulk_inserter,
             dml_event_handler=dml_event_handler,
+            token_store=token_store,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
             shutdown_callbacks=shutdown_callbacks_list,
@@ -694,14 +737,20 @@ def create_app(
             "defaultModelsExpandDepth": -1,  # hides the schema section in the Swagger UI
         },
     )
-    app.state.read_only = read_only
-    app.state.export_path = export_path
-    app.include_router(v1_router)
+    app.include_router(create_v1_router(authentication_enabled))
+    app.include_router(create_embeddings_router(authentication_enabled))
     app.include_router(router)
     app.include_router(graphql_router)
+    if authentication_enabled:
+        app.include_router(auth_router)
+        app.include_router(oauth2_router)
     app.add_middleware(GZipMiddleware)
     web_manifest_path = SERVER_DIR / "static" / ".vite" / "manifest.json"
     if serve_ui and web_manifest_path.is_file():
+        oauth2_idps = [
+            OAuth2Idp(name=config.idp_name, displayName=config.idp_display_name)
+            for config in oauth2_client_configs or []
+        ]
         app.mount(
             "/",
             app=Static(
@@ -715,11 +764,21 @@ def create_app(
                     is_development=dev,
                     authentication_enabled=authentication_enabled,
                     web_manifest_path=web_manifest_path,
+                    oauth2_idps=oauth2_idps,
                 ),
             ),
             name="static",
         )
-    app = _update_app_state(app, db=db, secret=secret)
+    app.state.read_only = read_only
+    app.state.export_path = export_path
+    app.state.password_reset_token_expiry = password_reset_token_expiry
+    app.state.access_token_expiry = access_token_expiry
+    app.state.refresh_token_expiry = refresh_token_expiry
+    app.state.oauth2_clients = OAuth2Clients.from_configs(oauth2_client_configs or [])
+    app.state.db = db
+    app.state.email_sender = email_sender
+    app = _add_get_secret_method(app=app, secret=secret)
+    app = _add_get_token_store_method(app=app, token_store=token_store)
     if tracer_provider:
         from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 
@@ -729,13 +788,10 @@ def create_app(
     return app
 
 
-def _update_app_state(app: FastAPI, /, *, db: DbSessionFactory, secret: Optional[str]) -> FastAPI:
+def _add_get_secret_method(*, app: FastAPI, secret: Optional[str]) -> FastAPI:
     """
-    Dynamically updates the app's `state` to include useful fields and methods
-    (at the time of this writing, FastAPI does not support setting this state
-    during the creation of the app).
+    Dynamically adds a `get_secret` method to the app's `state`.
     """
-    app.state.db = db
     app.state._secret = secret
 
     def get_secret(self: StarletteState) -> str:
@@ -746,3 +802,19 @@ def _update_app_state(app: FastAPI, /, *, db: DbSessionFactory, secret: Optional
 
     app.state.get_secret = MethodType(get_secret, app.state)
     return app
+
+
+def _add_get_token_store_method(*, app: FastAPI, token_store: Optional[JwtStore]) -> FastAPI:
+    """
+    Dynamically adds a `get_token_store` method to the app's `state`.
+    """
+    app.state._token_store = token_store
+
+    def get_token_store(self: StarletteState) -> JwtStore:
+        if (token_store := self._token_store) is None:
+            raise ValueError("token store is not set on the app")
+        assert isinstance(token_store, JwtStore)
+        return token_store
+
+    app.state.get_token_store = MethodType(get_token_store, app.state)
+    return app

phoenix/server/bearer_auth.py

@@ -0,0 +1,161 @@
+from abc import ABC
+from datetime import datetime, timedelta, timezone
+from functools import cached_property
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Optional,
+    Tuple,
+)
+
+import grpc
+from fastapi import HTTPException, Request
+from grpc_interceptor import AsyncServerInterceptor
+from grpc_interceptor.exceptions import Unauthenticated
+from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
+from starlette.requests import HTTPConnection
+from starlette.status import HTTP_401_UNAUTHORIZED
+
+from phoenix.auth import (
+    PHOENIX_ACCESS_TOKEN_COOKIE_NAME,
+    CanReadToken,
+    ClaimSetStatus,
+    Token,
+)
+from phoenix.db import enums
+from phoenix.db.enums import UserRole
+from phoenix.db.models import User as OrmUser
+from phoenix.server.types import (
+    AccessToken,
+    AccessTokenAttributes,
+    AccessTokenClaims,
+    ApiKeyClaims,
+    RefreshToken,
+    RefreshTokenAttributes,
+    RefreshTokenClaims,
+    TokenStore,
+    UserClaimSet,
+    UserId,
+)
+
+
+class HasTokenStore(ABC):
+    def __init__(self, token_store: CanReadToken) -> None:
+        super().__init__()
+        self._token_store = token_store
+
+
+class BearerTokenAuthBackend(HasTokenStore, AuthenticationBackend):
+    async def authenticate(
+        self,
+        conn: HTTPConnection,
+    ) -> Optional[Tuple[AuthCredentials, BaseUser]]:
+        if header := conn.headers.get("Authorization"):
+            scheme, _, token = header.partition(" ")
+            if scheme.lower() != "bearer" or not token:
+                return None
+        elif access_token := conn.cookies.get(PHOENIX_ACCESS_TOKEN_COOKIE_NAME):
+            token = access_token
+        else:
+            return None
+        claims = await self._token_store.read(Token(token))
+        if not (isinstance(claims, UserClaimSet) and isinstance(claims.subject, UserId)):
+            return None
+        if not isinstance(claims, (ApiKeyClaims, AccessTokenClaims)):
+            return None
+        return AuthCredentials(), PhoenixUser(claims.subject, claims)
+
+
+class PhoenixUser(BaseUser):
+    def __init__(self, user_id: UserId, claims: UserClaimSet) -> None:
+        self._user_id = user_id
+        self.claims = claims
+        assert claims.attributes
+        self._is_admin = (
+            claims.status is ClaimSetStatus.VALID
+            and claims.attributes.user_role == enums.UserRole.ADMIN
+        )
+
+    @cached_property
+    def is_admin(self) -> bool:
+        return self._is_admin
+
+    @cached_property
+    def identity(self) -> UserId:
+        return self._user_id
+
+    @cached_property
+    def is_authenticated(self) -> bool:
+        return True
+
+
+class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
+    async def intercept(
+        self,
+        method: Callable[[Any, grpc.ServicerContext], Awaitable[Any]],
+        request_or_iterator: Any,
+        context: grpc.ServicerContext,
+        method_name: str,
+    ) -> Any:
+        for datum in context.invocation_metadata():
+            if datum.key.lower() == "authorization":
+                scheme, _, token = datum.value.partition(" ")
+                if scheme.lower() != "bearer" or not token:
+                    break
+                claims = await self._token_store.read(Token(token))
+                if not (isinstance(claims, UserClaimSet) and isinstance(claims.subject, UserId)):
+                    break
+                if not isinstance(claims, (ApiKeyClaims, AccessTokenClaims)):
+                    raise Unauthenticated(details="Invalid token")
+                if claims.status is ClaimSetStatus.EXPIRED:
+                    raise Unauthenticated(details="Expired token")
+                if claims.status is ClaimSetStatus.VALID:
+                    return await method(request_or_iterator, context)
+                raise Unauthenticated()
+        raise Unauthenticated()
+
+
+async def is_authenticated(request: Request) -> None:
+    """
+    Raises a 401 if the request is not authenticated.
+    """
+    if not isinstance((user := request.user), PhoenixUser):
+        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+    claims = user.claims
+    if claims.status is ClaimSetStatus.EXPIRED:
+        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Expired token")
+    if claims.status is not ClaimSetStatus.VALID:
+        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+
+
+async def create_access_and_refresh_tokens(
+    *,
+    token_store: TokenStore,
+    user: OrmUser,
+    access_token_expiry: timedelta,
+    refresh_token_expiry: timedelta,
+) -> Tuple[AccessToken, RefreshToken]:
+    issued_at = datetime.now(timezone.utc)
+    user_id = UserId(user.id)
+    user_role = UserRole(user.role.name)
+    refresh_token_claims = RefreshTokenClaims(
+        subject=user_id,
+        issued_at=issued_at,
+        expiration_time=issued_at + refresh_token_expiry,
+        attributes=RefreshTokenAttributes(
+            user_role=user_role,
+        ),
+    )
+    refresh_token, refresh_token_id = await token_store.create_refresh_token(refresh_token_claims)
+    access_token_claims = AccessTokenClaims(
+        subject=user_id,
+        issued_at=issued_at,
+        expiration_time=issued_at + access_token_expiry,
+        attributes=AccessTokenAttributes(
+            user_role=user_role,
+            refresh_token_id=refresh_token_id,
+        ),
+    )
+    access_token, _ = await token_store.create_access_token(access_token_claims)
+    return access_token, refresh_token

phoenix/server/email/sender.py

@@ -0,0 +1,26 @@
+from pathlib import Path
+
+from fastapi_mail import ConnectionConfig, FastMail, MessageSchema
+
+EMAIL_TEMPLATE_FOLDER = Path(__file__).parent / "templates"
+
+
+class FastMailSender:
+    def __init__(self, conf: ConnectionConfig) -> None:
+        self._fm = FastMail(conf)
+
+    async def send_password_reset_email(
+        self,
+        email: str,
+        reset_url: str,
+    ) -> None:
+        message = MessageSchema(
+            subject="[Phoenix] Password Reset Request",
+            recipients=[email],
+            template_body=dict(reset_url=reset_url),
+            subtype="html",
+        )
+        await self._fm.send_message(
+            message,
+            template_name="password_reset.html",
+        )

phoenix/server/email/templates/password_reset.html

@@ -0,0 +1,19 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="UTF-8" />
+    <title>Password Reset</title>
+  </head>
+  <body>
+    <p>Hello.</p>
+    <p>
+      You have requested a password reset. Please click on the link below to
+      reset your password:
+    </p>
+    <p>
+      <a id="reset-url" href="{{ reset_url }}">Reset Password</a
+      >
+    </p>
+    <p>If you did not make this request, please contact your administrator.</p>
+  </body>
+</html>

phoenix/server/email/types.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class EmailSender(Protocol):
+    async def send_password_reset_email(
+        self,
+        email: str,
+        reset_url: str,
+    ) -> None: ...

phoenix/server/grpc_server.py

@@ -12,7 +12,9 @@ from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import (
 )
 from typing_extensions import TypeAlias
 
+from phoenix.auth import CanReadToken
 from phoenix.config import get_env_grpc_port
+from phoenix.server.bearer_auth import ApiKeyInterceptor
 from phoenix.trace.otel import decode_otlp_span
 from phoenix.trace.schemas import Span
 from phoenix.utilities.project import get_project_name
@@ -52,17 +54,21 @@ class GrpcServer:
         tracer_provider: Optional["TracerProvider"] = None,
         enable_prometheus: bool = False,
         disabled: bool = False,
+        token_store: Optional[CanReadToken] = None,
     ) -> None:
         self._callback = callback
         self._server: Optional[Server] = None
         self._tracer_provider = tracer_provider
         self._enable_prometheus = enable_prometheus
         self._disabled = disabled
+        self._token_store = token_store
 
     async def __aenter__(self) -> None:
         if self._disabled:
             return
         interceptors: List[ServerInterceptor] = []
+        if self._token_store:
+            interceptors.append(ApiKeyInterceptor(self._token_store))
         if self._enable_prometheus:
             ...
             # TODO: convert to async interceptor