arize-phoenix 5.5.2__py3-none-any.whl → 5.7.0__py3-none-any.whl

phoenix/server/api/types/Span.py

@@ -1,8 +1,9 @@
 import json
+from collections.abc import Mapping, Sized
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum
-from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Sized, cast
+from typing import TYPE_CHECKING, Any, Optional, cast
 
 import numpy as np
 import strawberry
@@ -19,10 +20,12 @@ from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
 )
+from phoenix.server.api.input_types.InvocationParameters import InvocationParameter
 from phoenix.server.api.input_types.SpanAnnotationSort import (
     SpanAnnotationColumn,
     SpanAnnotationSort,
 )
+from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.SpanAnnotation import to_gql_span_annotation
 from phoenix.trace.attributes import get_attribute_value
@@ -152,7 +155,7 @@ class Span(Node):
     token_count_completion: Optional[int]
     input: Optional[SpanIOValue]
     output: Optional[SpanIOValue]
-    events: List[SpanEvent]
+    events: list[SpanEvent]
     cumulative_token_count_total: Optional[int] = strawberry.field(
         description="Cumulative (prompt plus completion) token count from "
         "self and all descendant spans (children, grandchildren, etc.)",
@@ -180,7 +183,7 @@ class Span(Node):
         self,
         info: Info[Context, None],
         sort: Optional[SpanAnnotationSort] = UNSET,
-    ) -> List[SpanAnnotation]:
+    ) -> list[SpanAnnotation]:
         span_id = self.id_attr
         annotations = await info.context.data_loaders.span_annotations.load(span_id)
         sort_key = SpanAnnotationColumn.name.value
@@ -201,7 +204,7 @@ class Span(Node):
         "a list, and each evaluation is identified by its document's (zero-based) "
         "index in that list."
     )  # type: ignore
-    async def document_evaluations(self, info: Info[Context, None]) -> List[DocumentEvaluation]:
+    async def document_evaluations(self, info: Info[Context, None]) -> list[DocumentEvaluation]:
         return await info.context.data_loaders.document_evaluations.load(self.id_attr)
 
     @strawberry.field(
@@ -211,7 +214,7 @@ class Span(Node):
         self,
         info: Info[Context, None],
         evaluation_name: Optional[str] = UNSET,
-    ) -> List[DocumentRetrievalMetrics]:
+    ) -> list[DocumentRetrievalMetrics]:
         if not self.num_documents:
             return []
         return await info.context.data_loaders.document_retrieval_metrics.load(
@@ -224,7 +227,7 @@ class Span(Node):
     async def descendants(
         self,
         info: Info[Context, None],
-    ) -> List["Span"]:
+    ) -> list["Span"]:
         span_id = str(self.context.span_id)
         spans = await info.context.data_loaders.span_descendants.load(span_id)
         return [to_gql_span(span) for span in spans]
@@ -290,9 +293,44 @@ class Span(Node):
         examples = await info.context.data_loaders.span_dataset_examples.load(self.id_attr)
         return bool(examples)
 
+    @strawberry.field(description="Invocation parameters for the span")  # type: ignore
+    async def invocation_parameters(self, info: Info[Context, None]) -> list[InvocationParameter]:
+        from phoenix.server.api.helpers.playground_clients import OpenAIStreamingClient
+        from phoenix.server.api.helpers.playground_registry import PLAYGROUND_CLIENT_REGISTRY
+
+        db_span = self.db_span
+        attributes = db_span.attributes
+        llm_provider: GenerativeProviderKey = (
+            get_attribute_value(attributes, SpanAttributes.LLM_PROVIDER)
+            or GenerativeProviderKey.OPENAI
+        )
+        llm_model = get_attribute_value(attributes, SpanAttributes.LLM_MODEL_NAME)
+        invocation_parameters = get_attribute_value(
+            attributes, SpanAttributes.LLM_INVOCATION_PARAMETERS
+        )
+        if invocation_parameters is None:
+            return []
+        invocation_parameters = json.loads(invocation_parameters)
+        # find the client class for the provider, if there is no client class or provider,
+        # return openai as default
+        client_class = PLAYGROUND_CLIENT_REGISTRY.get_client(llm_provider, llm_model)
+        if not client_class:
+            client_class = OpenAIStreamingClient
+        supported_invocation_parameters = client_class.supported_invocation_parameters()
+        # filter supported invocation parameters down to those whose canonical_name is in the
+        # invocation_parameters keys
+        return [
+            ip
+            for ip in supported_invocation_parameters
+            if (
+                ip.canonical_name in invocation_parameters
+                or ip.invocation_name in invocation_parameters
+            )
+        ]
+
 
 def to_gql_span(span: models.Span) -> Span:
-    events: List[SpanEvent] = list(map(SpanEvent.from_dict, span.events))
+    events: list[SpanEvent] = list(map(SpanEvent.from_dict, span.events))
     input_value = cast(Optional[str], get_attribute_value(span.attributes, INPUT_VALUE))
     output_value = cast(Optional[str], get_attribute_value(span.attributes, OUTPUT_VALUE))
     retrieval_documents = get_attribute_value(span.attributes, RETRIEVAL_DOCUMENTS)

phoenix/server/api/types/TemplateLanguage.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class TemplateLanguage(Enum):
+    MUSTACHE = "MUSTACHE"
+    F_STRING = "F_STRING"

phoenix/server/api/types/TimeSeries.py

@@ -1,7 +1,8 @@
+from collections.abc import Iterable
 from dataclasses import replace
 from datetime import datetime, timedelta
 from functools import total_ordering
-from typing import Iterable, List, Optional, Tuple, Union, cast
+from typing import Optional, Union, cast
 
 import pandas as pd
 import strawberry
@@ -39,7 +40,7 @@ def to_gql_datapoints(
     df: pd.DataFrame,
     metric: Metric,
     timestamps: Iterable[datetime],
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     data = []
     for timestamp in timestamps:
         try:
@@ -59,7 +60,7 @@ def to_gql_datapoints(
 class TimeSeries:
     """A collection of data points over time"""
 
-    data: List[TimeSeriesDataPoint]
+    data: list[TimeSeriesDataPoint]
 
 
 def get_timeseries_data(
@@ -67,7 +68,7 @@ def get_timeseries_data(
     metric: Metric,
     time_range: TimeRange,
     granularity: Granularity,
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     return df.pipe(
         timeseries(
             start_time=time_range.start,
@@ -98,7 +99,7 @@ def get_data_quality_timeseries_data(
     time_range: TimeRange,
     granularity: Granularity,
     inferences_role: InferencesRole,
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     metric_instance = metric.value()
     if isinstance(metric_instance, UnaryOperator):
         metric_instance = replace(
@@ -128,7 +129,7 @@ def get_drift_timeseries_data(
     time_range: TimeRange,
     granularity: Granularity,
     reference_data: pd.DataFrame,
-) -> List[TimeSeriesDataPoint]:
+) -> list[TimeSeriesDataPoint]:
     metric_instance = metric.value()
     metric_instance = replace(
         metric_instance,
@@ -163,7 +164,7 @@ def ensure_timeseries_parameters(
     inferences: Inferences,
     time_range: Optional[TimeRange] = UNSET,
     granularity: Optional[Granularity] = UNSET,
-) -> Tuple[TimeRange, Granularity]:
+) -> tuple[TimeRange, Granularity]:
     if not isinstance(time_range, TimeRange):
         start, stop = inferences.time_range
         time_range = TimeRange(start=start, end=stop)

phoenix/server/api/types/Trace.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import List, Optional
+from typing import Optional
 
 import strawberry
 from sqlalchemy import desc, select
@@ -69,7 +69,7 @@ class Trace(Node):
         self,
         info: Info[Context, None],
         sort: Optional[TraceAnnotationSort] = None,
-    ) -> List[TraceAnnotation]:
+    ) -> list[TraceAnnotation]:
         async with info.context.db() as session:
             stmt = select(models.TraceAnnotation).filter_by(span_rowid=self.id_attr)
             if sort:

phoenix/server/api/types/UMAPPoints.py

@@ -1,4 +1,4 @@
-from typing import List, Union
+from typing import Union
 
 import numpy as np
 import numpy.typing as npt
@@ -57,8 +57,8 @@ class UMAPPoint:
 
 @strawberry.type
 class UMAPPoints:
-    data: List[UMAPPoint]
-    reference_data: List[UMAPPoint]
-    clusters: List[Cluster]
-    corpus_data: List[UMAPPoint] = strawberry.field(default_factory=list)
-    context_retrievals: List[Retrieval] = strawberry.field(default_factory=list)
+    data: list[UMAPPoint]
+    reference_data: list[UMAPPoint]
+    clusters: list[Cluster]
+    corpus_data: list[UMAPPoint] = strawberry.field(default_factory=list)
+    context_retrievals: list[Retrieval] = strawberry.field(default_factory=list)

phoenix/server/api/types/User.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import Optional
 
 import strawberry
 from sqlalchemy import select
@@ -35,7 +35,7 @@ class User(Node):
         return to_gql_user_role(role)
 
     @strawberry.field
-    async def api_keys(self, info: Info[Context, None]) -> List[UserApiKey]:
+    async def api_keys(self, info: Info[Context, None]) -> list[UserApiKey]:
         async with info.context.db() as session:
             api_keys = await session.scalars(
                 select(models.ApiKey).where(models.ApiKey.user_id == self.id_attr)
@@ -43,7 +43,7 @@ class User(Node):
         return [to_gql_api_key(api_key) for api_key in api_keys]
 
 
-def to_gql_user(user: models.User, api_keys: Optional[List[models.ApiKey]] = None) -> User:
+def to_gql_user(user: models.User, api_keys: Optional[list[models.ApiKey]] = None) -> User:
     """
     Converts an ORM user to a GraphQL user.
     """

phoenix/server/api/types/node.py

@@ -1,9 +1,7 @@
-from typing import Tuple
-
 from strawberry.relay import GlobalID
 
 
-def from_global_id(global_id: GlobalID) -> Tuple[str, int]:
+def from_global_id(global_id: GlobalID) -> tuple[str, int]:
     """
     Decode the given global id into a type and id.
 

phoenix/server/api/types/pagination.py

@@ -2,7 +2,7 @@ import base64
 from dataclasses import dataclass
 from datetime import datetime
 from enum import Enum, auto
-from typing import Any, ClassVar, List, Optional, Tuple, Union
+from typing import Any, ClassVar, Optional, Union
 
 from strawberry import UNSET
 from strawberry.relay.types import Connection, Edge, NodeType, PageInfo
@@ -176,7 +176,7 @@ class ConnectionArgs:
 
 
 def connection_from_list(
-    data: List[NodeType],
+    data: list[NodeType],
     args: ConnectionArgs,
 ) -> Connection[NodeType]:
     """
@@ -188,7 +188,7 @@ def connection_from_list(
 
 
 def connection_from_list_slice(
-    list_slice: List[NodeType],
+    list_slice: list[NodeType],
     args: ConnectionArgs,
     slice_start: int,
     list_length: int,
@@ -254,7 +254,7 @@ def connection_from_list_slice(
 
 
 def connection_from_cursors_and_nodes(
-    cursors_and_nodes: List[Tuple[Any, NodeType]],
+    cursors_and_nodes: list[tuple[Any, NodeType]],
     has_previous_page: bool,
     has_next_page: bool,
 ) -> Connection[NodeType]:

phoenix/server/api/utils.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from sqlalchemy import delete
 
 from phoenix.db import models
@@ -9,7 +7,7 @@ from phoenix.server.types import DbSessionFactory
 async def delete_projects(
     db: DbSessionFactory,
     *project_names: str,
-) -> List[int]:
+) -> list[int]:
     if not project_names:
         return []
     stmt = (
@@ -24,7 +22,7 @@ async def delete_projects(
 async def delete_traces(
     db: DbSessionFactory,
     *trace_ids: str,
-) -> List[int]:
+) -> list[int]:
     if not trace_ids:
         return []
     stmt = (

phoenix/server/app.py

@@ -1,8 +1,11 @@
 import asyncio
 import contextlib
+import importlib
 import json
 import logging
-from contextlib import AsyncExitStack
+import os
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta, timezone
 from functools import cached_property
@@ -11,18 +14,8 @@ from types import MethodType
 from typing import (
     TYPE_CHECKING,
     Any,
-    AsyncContextManager,
-    AsyncIterator,
-    Awaitable,
-    Callable,
-    Dict,
-    Iterable,
-    List,
     NamedTuple,
     Optional,
-    Sequence,
-    Tuple,
-    Type,
     TypedDict,
     Union,
     cast,
@@ -49,7 +42,6 @@ from starlette.types import Scope, StatefulLifespan
 from starlette.websockets import WebSocket
 from strawberry.extensions import SchemaExtension
 from strawberry.fastapi import GraphQLRouter
-from strawberry.schema import BaseSchema
 from strawberry.subscriptions import GRAPHQL_TRANSPORT_WS_PROTOCOL
 from typing_extensions import TypeAlias
 
@@ -60,6 +52,8 @@ from phoenix.config import (
     SERVER_DIR,
     OAuth2ClientConfig,
     get_env_csrf_trusted_origins,
+    get_env_fastapi_middleware_paths,
+    get_env_gql_extension_paths,
     get_env_host,
     get_env_port,
     server_instrumentation_is_enabled,
@@ -107,7 +101,7 @@ from phoenix.server.api.routers import (
     oauth2_router,
 )
 from phoenix.server.api.routers.v1 import REST_API_VERSION
-from phoenix.server.api.schema import schema
+from phoenix.server.api.schema import build_graphql_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
@@ -159,6 +153,28 @@ ProjectName: TypeAlias = str
 _Callback: TypeAlias = Callable[[], Union[None, Awaitable[None]]]
 
 
+def import_object_from_file(file_path: str, object_name: str) -> Any:
+    """Import an object (class or function) from a Python file."""
+    try:
+        if not os.path.isfile(file_path):
+            raise FileNotFoundError(f"File '{file_path}' does not exist.")
+        module_name = f"custom_module_{hash(file_path)}"
+        spec = importlib.util.spec_from_file_location(module_name, file_path)
+        if spec is None:
+            raise ImportError(f"Could not load spec for '{file_path}'")
+        module = importlib.util.module_from_spec(spec)
+        loader = spec.loader
+        if loader is None:
+            raise ImportError(f"No loader found for '{file_path}'")
+        loader.exec_module(module)
+        try:
+            return getattr(module, object_name)
+        except AttributeError:
+            raise ImportError(f"Module '{file_path}' does not have an object '{object_name}'.")
+    except Exception as e:
+        raise ImportError(f"Could not import '{object_name}' from '{file_path}': {e}")
+
+
 class OAuth2Idp(TypedDict):
     name: str
     displayName: str
@@ -175,6 +191,7 @@ class AppConfig(NamedTuple):
     web_manifest_path: Path
     authentication_enabled: bool
     """ Whether authentication is enabled """
+    websockets_enabled: bool
     oauth2_idps: Sequence[OAuth2Idp]
 
 
@@ -188,10 +205,10 @@ class Static(StaticFiles):
         super().__init__(**kwargs)
 
     @cached_property
-    def _web_manifest(self) -> Dict[str, Any]:
+    def _web_manifest(self) -> dict[str, Any]:
         try:
             with open(self._app_config.web_manifest_path, "r") as f:
-                return cast(Dict[str, Any], json.load(f))
+                return cast(dict[str, Any], json.load(f))
         except FileNotFoundError as e:
             if self._app_config.is_development:
                 return {}
@@ -225,6 +242,7 @@ class Static(StaticFiles):
                     "manifest": self._web_manifest,
                     "authentication_enabled": self._app_config.authentication_enabled,
                     "oauth2_idps": self._app_config.oauth2_idps,
+                    "websockets_enabled": self._app_config.websockets_enabled,
                 },
             )
         except Exception as e:
@@ -233,7 +251,7 @@ class Static(StaticFiles):
 
 
 class RequestOriginHostnameValidator(BaseHTTPMiddleware):
-    def __init__(self, trusted_hostnames: List[str], *args: Any, **kwargs: Any) -> None:
+    def __init__(self, trusted_hostnames: list[str], *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
         self._trusted_hostnames = trusted_hostnames
 
@@ -265,6 +283,28 @@ class HeadersMiddleware(BaseHTTPMiddleware):
         return response
 
 
+def user_fastapi_middlewares() -> list[Middleware]:
+    paths = get_env_fastapi_middleware_paths()
+    middlewares = []
+    for file_path, object_name in paths:
+        middleware_class = import_object_from_file(file_path, object_name)
+        if not issubclass(middleware_class, BaseHTTPMiddleware):
+            raise TypeError(f"{middleware_class} is not a subclass of BaseHTTPMiddleware")
+        middlewares.append(Middleware(middleware_class))
+    return middlewares
+
+
+def user_gql_extensions() -> list[Union[type[SchemaExtension], SchemaExtension]]:
+    paths = get_env_gql_extension_paths()
+    extensions = []
+    for file_path, object_name in paths:
+        extension_class = import_object_from_file(file_path, object_name)
+        if not issubclass(extension_class, SchemaExtension):
+            raise TypeError(f"{extension_class} is not a subclass of SchemaExtension")
+        extensions.append(extension_class)
+    return extensions
+
+
 ProjectRowId: TypeAlias = int
 
 
@@ -278,7 +318,7 @@ DB_MUTEX: Optional[asyncio.Lock] = None
 
 def _db(
     engine: AsyncEngine, bypass_lock: bool = False
-) -> Callable[[], AsyncContextManager[AsyncSession]]:
+) -> Callable[[], AbstractAsyncContextManager[AsyncSession]]:
     Session = async_sessionmaker(engine, expire_on_commit=False)
 
     @contextlib.asynccontextmanager
@@ -420,7 +460,7 @@ def _lifespan(
     scaffolder_config: Optional[ScaffolderConfig] = None,
 ) -> StatefulLifespan[FastAPI]:
     @contextlib.asynccontextmanager
-    async def lifespan(_: FastAPI) -> AsyncIterator[Dict[str, Any]]:
+    async def lifespan(_: FastAPI) -> AsyncIterator[dict[str, Any]]:
         for callback in startup_callbacks:
             if isinstance((res := callback()), Awaitable):
                 await res
@@ -449,7 +489,7 @@ def _lifespan(
                     queue_evaluation=queue_evaluation,
                 )
                 await stack.enter_async_context(scaffolder)
-            if isinstance(token_store, AsyncContextManager):
+            if isinstance(token_store, AbstractAsyncContextManager):
                 await stack.enter_async_context(token_store)
             yield {
                 "event_queue": dml_event_handler,
@@ -472,7 +512,7 @@ async def check_healthz(_: Request) -> PlainTextResponse:
 
 def create_graphql_router(
     *,
-    schema: BaseSchema,
+    graphql_schema: strawberry.Schema,
     db: DbSessionFactory,
     model: Model,
     export_path: Path,
@@ -576,7 +616,7 @@ def create_graphql_router(
         )
 
     return GraphQLRouter(
-        schema,
+        graphql_schema,
         graphql_ide="graphiql",
         context_getter=get_context,
         include_in_schema=False,
@@ -607,7 +647,7 @@ def create_engine_and_run_migrations(
         raise PhoenixMigrationError(msg) from e
 
 
-def instrument_engine_if_enabled(engine: AsyncEngine) -> List[Callable[[], None]]:
+def instrument_engine_if_enabled(engine: AsyncEngine) -> list[Callable[[], None]]:
     instrumentation_cleanups = []
     if server_instrumentation_is_enabled():
         from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
@@ -662,12 +702,13 @@ def create_app(
     model: Model,
     authentication_enabled: bool,
     umap_params: UMAPParameters,
+    enable_websockets: bool,
     corpus: Optional[Model] = None,
     debug: bool = False,
     dev: bool = False,
     read_only: bool = False,
     enable_prometheus: bool = False,
-    initial_spans: Optional[Iterable[Union[Span, Tuple[Span, str]]]] = None,
+    initial_spans: Optional[Iterable[Union[Span, tuple[Span, str]]]] = None,
     initial_evaluations: Optional[Iterable[pb.Evaluation]] = None,
     serve_ui: bool = True,
     startup_callbacks: Iterable[_Callback] = (),
@@ -678,7 +719,7 @@ def create_app(
     refresh_token_expiry: Optional[timedelta] = None,
     scaffolder_config: Optional[ScaffolderConfig] = None,
     email_sender: Optional[EmailSender] = None,
-    oauth2_client_configs: Optional[List[OAuth2ClientConfig]] = None,
+    oauth2_client_configs: Optional[list[OAuth2ClientConfig]] = None,
     bulk_inserter_factory: Optional[Callable[..., BulkInserter]] = None,
 ) -> FastAPI:
     if model.embedding_dimensions:
@@ -692,10 +733,10 @@ def create_app(
             ) from exc
     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: 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]] = (
+    initial_batch_of_spans: Iterable[tuple[Span, str]] = (
         ()
         if initial_spans is None
         else (
@@ -708,7 +749,8 @@ def create_app(
         CacheForDataLoaders() if db.dialect is SupportedSQLDialect.SQLITE else None
     )
     last_updated_at = LastUpdatedAt()
-    middlewares: List[Middleware] = [Middleware(HeadersMiddleware)]
+    middlewares: list[Middleware] = [Middleware(HeadersMiddleware)]
+    middlewares.extend(user_fastapi_middlewares())
     if origins := get_env_csrf_trusted_origins():
         trusted_hostnames = [h for o in origins if o and (h := urlparse(o).hostname)]
         middlewares.append(Middleware(RequestOriginHostnameValidator, trusted_hostnames))
@@ -742,8 +784,9 @@ def create_app(
         initial_batch_of_evaluations=initial_batch_of_evaluations,
     )
     tracer_provider = None
-    strawberry_extensions: List[Union[Type[SchemaExtension], SchemaExtension]] = []
-    strawberry_extensions.extend(schema.get_extensions())
+    graphql_schema_extensions: list[Union[type[SchemaExtension], SchemaExtension]] = []
+    graphql_schema_extensions.extend(user_gql_extensions())
+
     if server_instrumentation_is_enabled():
         tracer_provider = initialize_opentelemetry_tracer_provider()
         from opentelemetry.trace import TracerProvider
@@ -761,16 +804,11 @@ def create_app(
                 # used by OpenInference.
                 self._tracer = cast(TracerProvider, tracer_provider).get_tracer("strawberry")
 
-        strawberry_extensions.append(_OpenTelemetryExtension)
+        graphql_schema_extensions.append(_OpenTelemetryExtension)
 
     graphql_router = create_graphql_router(
         db=db,
-        schema=strawberry.Schema(
-            query=schema.query,
-            mutation=schema.mutation,
-            subscription=schema.subscription,
-            extensions=strawberry_extensions,
-        ),
+        graphql_schema=build_graphql_schema(graphql_schema_extensions),
         model=model,
         corpus=corpus,
         authentication_enabled=authentication_enabled,
@@ -839,6 +877,7 @@ def create_app(
                     authentication_enabled=authentication_enabled,
                     web_manifest_path=web_manifest_path,
                     oauth2_idps=oauth2_idps,
+                    websockets_enabled=enable_websockets,
                 ),
             ),
             name="static",

phoenix/server/bearer_auth.py

@@ -1,14 +1,8 @@
 from abc import ABC
+from collections.abc import Awaitable, Callable
 from datetime import datetime, timedelta, timezone
 from functools import cached_property
-from typing import (
-    Any,
-    Awaitable,
-    Callable,
-    Optional,
-    Tuple,
-    cast,
-)
+from typing import Any, Optional, cast
 
 import grpc
 from fastapi import HTTPException, Request, WebSocket, WebSocketException
@@ -51,7 +45,7 @@ class BearerTokenAuthBackend(HasTokenStore, AuthenticationBackend):
     async def authenticate(
         self,
         conn: HTTPConnection,
-    ) -> Optional[Tuple[AuthCredentials, BaseUser]]:
+    ) -> Optional[tuple[AuthCredentials, BaseUser]]:
         if header := conn.headers.get("Authorization"):
             scheme, _, token = header.partition(" ")
             if scheme.lower() != "bearer" or not token:
@@ -143,7 +137,7 @@ async def create_access_and_refresh_tokens(
     user: OrmUser,
     access_token_expiry: timedelta,
     refresh_token_expiry: timedelta,
-) -> Tuple[AccessToken, RefreshToken]:
+) -> tuple[AccessToken, RefreshToken]:
     issued_at = datetime.now(timezone.utc)
     user_id = UserId(user.id)
     user_role = UserRole(user.role.name)

phoenix/server/dml_event.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from abc import ABC
 from dataclasses import dataclass, field
-from typing import ClassVar, Tuple, Type
+from typing import ClassVar
 
 from phoenix.db import models
 
@@ -14,8 +14,8 @@ class DmlEvent(ABC):
     operation, e.g. insertion, update, or deletion.
     """
 
-    table: ClassVar[Type[models.Base]]
-    ids: Tuple[int, ...] = field(default_factory=tuple)
+    table: ClassVar[type[models.Base]]
+    ids: tuple[int, ...] = field(default_factory=tuple)
 
     def __bool__(self) -> bool:
         return bool(self.ids)