arize-phoenix 5.2.1__py3-none-any.whl → 5.3.0__py3-none-any.whl

phoenix/server/api/subscriptions.py

@@ -0,0 +1,430 @@
+import json
+from collections import defaultdict
+from dataclasses import fields
+from datetime import datetime
+from enum import Enum
+from itertools import chain
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    AsyncIterator,
+    DefaultDict,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import strawberry
+from openinference.instrumentation import safe_json_dumps
+from openinference.semconv.trace import (
+    MessageAttributes,
+    OpenInferenceMimeTypeValues,
+    OpenInferenceSpanKindValues,
+    SpanAttributes,
+    ToolAttributes,
+    ToolCallAttributes,
+)
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import StatusCode
+from sqlalchemy import insert, select
+from strawberry import UNSET
+from strawberry.scalars import JSON as JSONScalarType
+from strawberry.types import Info
+from typing_extensions import TypeAlias, assert_never
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.ChatCompletionMessageInput import ChatCompletionMessageInput
+from phoenix.server.api.input_types.InvocationParameters import InvocationParameters
+from phoenix.server.api.types.ChatCompletionMessageRole import ChatCompletionMessageRole
+from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
+from phoenix.server.dml_event import SpanInsertEvent
+from phoenix.trace.attributes import unflatten
+from phoenix.utilities.json import jsonify
+from phoenix.utilities.template_formatters import (
+    FStringTemplateFormatter,
+    MustacheTemplateFormatter,
+    TemplateFormatter,
+)
+
+if TYPE_CHECKING:
+    from openai.types.chat import (
+        ChatCompletionMessageParam,
+    )
+
+PLAYGROUND_PROJECT_NAME = "playground"
+
+ToolCallIndex: TypeAlias = int
+
+
+@strawberry.enum
+class TemplateLanguage(Enum):
+    MUSTACHE = "MUSTACHE"
+    F_STRING = "F_STRING"
+
+
+@strawberry.input
+class TemplateOptions:
+    variables: JSONScalarType
+    language: TemplateLanguage
+
+
+@strawberry.type
+class TextChunk:
+    content: str
+
+
+@strawberry.type
+class FunctionCallChunk:
+    name: str
+    arguments: str
+
+
+@strawberry.type
+class ToolCallChunk:
+    id: str
+    function: FunctionCallChunk
+
+
+ChatCompletionChunk: TypeAlias = Annotated[
+    Union[TextChunk, ToolCallChunk], strawberry.union("ChatCompletionChunk")
+]
+
+
+@strawberry.input
+class GenerativeModelInput:
+    provider_key: GenerativeProviderKey
+    name: str
+    """ The name of the model. Or the Deployment Name for Azure OpenAI models. """
+    endpoint: Optional[str] = UNSET
+    """ The endpoint to use for the model. Only required for Azure OpenAI models. """
+    api_version: Optional[str] = UNSET
+    """ The API version to use for the model. """
+
+
+@strawberry.input
+class ChatCompletionInput:
+    messages: List[ChatCompletionMessageInput]
+    model: GenerativeModelInput
+    invocation_parameters: InvocationParameters
+    tools: Optional[List[JSONScalarType]] = UNSET
+    template: Optional[TemplateOptions] = UNSET
+    api_key: Optional[str] = strawberry.field(default=None)
+
+
+def to_openai_chat_completion_param(
+    role: ChatCompletionMessageRole, content: JSONScalarType
+) -> "ChatCompletionMessageParam":
+    from openai.types.chat import (
+        ChatCompletionAssistantMessageParam,
+        ChatCompletionSystemMessageParam,
+        ChatCompletionUserMessageParam,
+    )
+
+    if role is ChatCompletionMessageRole.USER:
+        return ChatCompletionUserMessageParam(
+            {
+                "content": content,
+                "role": "user",
+            }
+        )
+    if role is ChatCompletionMessageRole.SYSTEM:
+        return ChatCompletionSystemMessageParam(
+            {
+                "content": content,
+                "role": "system",
+            }
+        )
+    if role is ChatCompletionMessageRole.AI:
+        return ChatCompletionAssistantMessageParam(
+            {
+                "content": content,
+                "role": "assistant",
+            }
+        )
+    if role is ChatCompletionMessageRole.TOOL:
+        raise NotImplementedError
+    assert_never(role)
+
+
+@strawberry.type
+class Subscription:
+    @strawberry.subscription
+    async def chat_completion(
+        self, info: Info[Context, None], input: ChatCompletionInput
+    ) -> AsyncIterator[ChatCompletionChunk]:
+        from openai import NOT_GIVEN, AsyncAzureOpenAI, AsyncOpenAI
+
+        client: Union[AsyncAzureOpenAI, AsyncOpenAI]
+
+        if input.model.provider_key == GenerativeProviderKey.AZURE_OPENAI:
+            if input.model.endpoint is None or input.model.api_version is None:
+                raise ValueError("endpoint and api_version are required for Azure OpenAI models")
+            client = AsyncAzureOpenAI(
+                api_key=input.api_key,
+                azure_endpoint=input.model.endpoint,
+                api_version=input.model.api_version,
+            )
+        else:
+            client = AsyncOpenAI(api_key=input.api_key)
+
+        invocation_parameters = jsonify(input.invocation_parameters)
+
+        messages: List[Tuple[ChatCompletionMessageRole, str]] = [
+            (message.role, message.content) for message in input.messages
+        ]
+        if template_options := input.template:
+            messages = list(_formatted_messages(messages, template_options))
+        openai_messages = [to_openai_chat_completion_param(*message) for message in messages]
+
+        in_memory_span_exporter = InMemorySpanExporter()
+        tracer_provider = TracerProvider()
+        tracer_provider.add_span_processor(
+            span_processor=SimpleSpanProcessor(span_exporter=in_memory_span_exporter)
+        )
+        tracer = tracer_provider.get_tracer(__name__)
+        span_name = "ChatCompletion"
+        with tracer.start_span(
+            span_name,
+            attributes=dict(
+                chain(
+                    _llm_span_kind(),
+                    _llm_model_name(input.model.name),
+                    _llm_tools(input.tools or []),
+                    _llm_input_messages(messages),
+                    _llm_invocation_parameters(invocation_parameters),
+                    _input_value_and_mime_type(input),
+                )
+            ),
+        ) as span:
+            response_chunks = []
+            text_chunks: List[TextChunk] = []
+            tool_call_chunks: DefaultDict[ToolCallIndex, List[ToolCallChunk]] = defaultdict(list)
+            role: Optional[str] = None
+            async for chunk in await client.chat.completions.create(
+                messages=openai_messages,
+                model=input.model.name,
+                stream=True,
+                tools=input.tools or NOT_GIVEN,
+                **invocation_parameters,
+            ):
+                response_chunks.append(chunk)
+                choice = chunk.choices[0]
+                delta = choice.delta
+                if role is None:
+                    role = delta.role
+                if choice.finish_reason is None:
+                    if isinstance(chunk_content := delta.content, str):
+                        text_chunk = TextChunk(content=chunk_content)
+                        yield text_chunk
+                        text_chunks.append(text_chunk)
+                    if (tool_calls := delta.tool_calls) is not None:
+                        for tool_call_index, tool_call in enumerate(tool_calls):
+                            if (function := tool_call.function) is not None:
+                                if (tool_call_id := tool_call.id) is None:
+                                    first_tool_call_chunk = tool_call_chunks[tool_call_index][0]
+                                    tool_call_id = first_tool_call_chunk.id
+                                tool_call_chunk = ToolCallChunk(
+                                    id=tool_call_id,
+                                    function=FunctionCallChunk(
+                                        name=function.name or "",
+                                        arguments=function.arguments or "",
+                                    ),
+                                )
+                                yield tool_call_chunk
+                                tool_call_chunks[tool_call_index].append(tool_call_chunk)
+            span.set_status(StatusCode.OK)
+            assert role is not None
+            span.set_attributes(
+                dict(
+                    chain(
+                        _output_value_and_mime_type(response_chunks),
+                        _llm_output_messages(text_chunks, tool_call_chunks),
+                    )
+                )
+            )
+        assert len(spans := in_memory_span_exporter.get_finished_spans()) == 1
+        finished_span = spans[0]
+        assert finished_span.start_time is not None
+        assert finished_span.end_time is not None
+        assert (attributes := finished_span.attributes) is not None
+        start_time = _datetime(epoch_nanoseconds=finished_span.start_time)
+        end_time = _datetime(epoch_nanoseconds=finished_span.end_time)
+        trace_id = _hex(finished_span.context.trace_id)
+        span_id = _hex(finished_span.context.span_id)
+        status = finished_span.status
+        async with info.context.db() as session:
+            if (
+                playground_project_id := await session.scalar(
+                    select(models.Project.id).where(models.Project.name == PLAYGROUND_PROJECT_NAME)
+                )
+            ) is None:
+                playground_project_id = await session.scalar(
+                    insert(models.Project)
+                    .returning(models.Project.id)
+                    .values(
+                        name=PLAYGROUND_PROJECT_NAME,
+                        description="Traces from prompt playground",
+                    )
+                )
+            trace_rowid = await session.scalar(
+                insert(models.Trace)
+                .returning(models.Trace.id)
+                .values(
+                    project_rowid=playground_project_id,
+                    trace_id=trace_id,
+                    start_time=start_time,
+                    end_time=end_time,
+                )
+            )
+            await session.execute(
+                insert(models.Span).values(
+                    trace_rowid=trace_rowid,
+                    span_id=span_id,
+                    parent_id=None,
+                    name=span_name,
+                    span_kind=LLM,
+                    start_time=start_time,
+                    end_time=end_time,
+                    attributes=unflatten(attributes.items()),
+                    events=finished_span.events,
+                    status_code=status.status_code.name,
+                    status_message=status.description or "",
+                    cumulative_error_count=int(not status.is_ok),
+                    cumulative_llm_token_count_prompt=0,
+                    cumulative_llm_token_count_completion=0,
+                    llm_token_count_prompt=0,
+                    llm_token_count_completion=0,
+                )
+            )
+        info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
+
+
+def _llm_span_kind() -> Iterator[Tuple[str, Any]]:
+    yield OPENINFERENCE_SPAN_KIND, LLM
+
+
+def _llm_model_name(model_name: str) -> Iterator[Tuple[str, Any]]:
+    yield LLM_MODEL_NAME, model_name
+
+
+def _llm_invocation_parameters(invocation_parameters: Dict[str, Any]) -> Iterator[Tuple[str, Any]]:
+    yield LLM_INVOCATION_PARAMETERS, safe_json_dumps(invocation_parameters)
+
+
+def _llm_tools(tools: List[JSONScalarType]) -> Iterator[Tuple[str, Any]]:
+    for tool_index, tool in enumerate(tools):
+        yield f"{LLM_TOOLS}.{tool_index}.{TOOL_JSON_SCHEMA}", json.dumps(tool)
+
+
+def _input_value_and_mime_type(input: ChatCompletionInput) -> Iterator[Tuple[str, Any]]:
+    assert any(field.name == (api_key := "api_key") for field in fields(ChatCompletionInput))
+    yield INPUT_MIME_TYPE, JSON
+    yield INPUT_VALUE, safe_json_dumps({k: v for k, v in jsonify(input).items() if k != api_key})
+
+
+def _output_value_and_mime_type(output: Any) -> Iterator[Tuple[str, Any]]:
+    yield OUTPUT_MIME_TYPE, JSON
+    yield OUTPUT_VALUE, safe_json_dumps(jsonify(output))
+
+
+def _llm_input_messages(
+    messages: Iterable[Tuple[ChatCompletionMessageRole, str]],
+) -> Iterator[Tuple[str, Any]]:
+    for i, (role, content) in enumerate(messages):
+        yield f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_ROLE}", role.value.lower()
+        yield f"{LLM_INPUT_MESSAGES}.{i}.{MESSAGE_CONTENT}", content
+
+
+def _llm_output_messages(
+    text_chunks: List[TextChunk],
+    tool_call_chunks: DefaultDict[ToolCallIndex, List[ToolCallChunk]],
+) -> Iterator[Tuple[str, Any]]:
+    yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_ROLE}", "assistant"
+    if content := "".join(chunk.content for chunk in text_chunks):
+        yield f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_CONTENT}", content
+    for tool_call_index, tool_call_chunks_ in tool_call_chunks.items():
+        if tool_call_chunks_ and (name := tool_call_chunks_[0].function.name):
+            yield (
+                f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_NAME}",
+                name,
+            )
+        if arguments := "".join(chunk.function.arguments for chunk in tool_call_chunks_):
+            yield (
+                f"{LLM_OUTPUT_MESSAGES}.0.{MESSAGE_TOOL_CALLS}.{tool_call_index}.{TOOL_CALL_FUNCTION_ARGUMENTS_JSON}",
+                arguments,
+            )
+
+
+def _hex(number: int) -> str:
+    """
+    Converts an integer to a hexadecimal string.
+    """
+    return hex(number)[2:]
+
+
+def _datetime(*, epoch_nanoseconds: float) -> datetime:
+    """
+    Converts a Unix epoch timestamp in nanoseconds to a datetime.
+    """
+    epoch_seconds = epoch_nanoseconds / 1e9
+    return datetime.fromtimestamp(epoch_seconds)
+
+
+def _formatted_messages(
+    messages: Iterable[Tuple[ChatCompletionMessageRole, str]], template_options: TemplateOptions
+) -> Iterator[Tuple[ChatCompletionMessageRole, str]]:
+    """
+    Formats the messages using the given template options.
+    """
+    template_formatter = _template_formatter(template_language=template_options.language)
+    roles, templates = zip(*messages)
+    formatted_templates = map(
+        lambda template: template_formatter.format(template, **template_options.variables),
+        templates,
+    )
+    formatted_messages = zip(roles, formatted_templates)
+    return formatted_messages
+
+
+def _template_formatter(template_language: TemplateLanguage) -> TemplateFormatter:
+    """
+    Instantiates the appropriate template formatter for the template language.
+    """
+    if template_language is TemplateLanguage.MUSTACHE:
+        return MustacheTemplateFormatter()
+    if template_language is TemplateLanguage.F_STRING:
+        return FStringTemplateFormatter()
+    assert_never(template_language)
+
+
+JSON = OpenInferenceMimeTypeValues.JSON.value
+
+LLM = OpenInferenceSpanKindValues.LLM.value
+
+OPENINFERENCE_SPAN_KIND = SpanAttributes.OPENINFERENCE_SPAN_KIND
+INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE
+INPUT_VALUE = SpanAttributes.INPUT_VALUE
+OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE
+LLM_INPUT_MESSAGES = SpanAttributes.LLM_INPUT_MESSAGES
+LLM_OUTPUT_MESSAGES = SpanAttributes.LLM_OUTPUT_MESSAGES
+LLM_MODEL_NAME = SpanAttributes.LLM_MODEL_NAME
+LLM_INVOCATION_PARAMETERS = SpanAttributes.LLM_INVOCATION_PARAMETERS
+LLM_TOOLS = SpanAttributes.LLM_TOOLS
+
+MESSAGE_CONTENT = MessageAttributes.MESSAGE_CONTENT
+MESSAGE_ROLE = MessageAttributes.MESSAGE_ROLE
+MESSAGE_TOOL_CALLS = MessageAttributes.MESSAGE_TOOL_CALLS
+
+TOOL_CALL_FUNCTION_NAME = ToolCallAttributes.TOOL_CALL_FUNCTION_NAME
+TOOL_CALL_FUNCTION_ARGUMENTS_JSON = ToolCallAttributes.TOOL_CALL_FUNCTION_ARGUMENTS_JSON
+
+TOOL_JSON_SCHEMA = ToolAttributes.TOOL_JSON_SCHEMA

phoenix/server/api/types/ChatCompletionMessageRole.py

@@ -0,0 +1,11 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class ChatCompletionMessageRole(Enum):
+    USER = "USER"
+    SYSTEM = "SYSTEM"
+    TOOL = "TOOL"
+    AI = "AI"  # E.g. the assistant. Normalize to AI for consistency.

phoenix/server/api/types/GenerativeModel.py

@@ -0,0 +1,9 @@
+import strawberry
+
+from phoenix.server.api.types.GenerativeProvider import GenerativeProviderKey
+
+
+@strawberry.type
+class GenerativeModel:
+    name: str
+    provider_key: GenerativeProviderKey

phoenix/server/api/types/GenerativeProvider.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class GenerativeProviderKey(Enum):
+    OPENAI = "OPENAI"
+    ANTHROPIC = "ANTHROPIC"
+    AZURE_OPENAI = "AZURE_OPENAI"
+
+
+@strawberry.type
+class GenerativeProvider:
+    name: str
+    key: GenerativeProviderKey

phoenix/server/app.py

@@ -27,6 +27,7 @@ from typing import (
     Union,
     cast,
 )
+from urllib.parse import urlparse
 
 import strawberry
 from fastapi import APIRouter, Depends, FastAPI
@@ -35,26 +36,30 @@ 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.exceptions import HTTPException, WebSocketException
 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
+from starlette.responses import JSONResponse, PlainTextResponse, Response
 from starlette.staticfiles import StaticFiles
+from starlette.status import HTTP_401_UNAUTHORIZED
 from starlette.templating import Jinja2Templates
 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
 
-import phoenix
 import phoenix.trace.v1 as pb
 from phoenix.config import (
     DEFAULT_PROJECT_NAME,
+    ENV_PHOENIX_CSRF_TRUSTED_ORIGINS,
     SERVER_DIR,
     OAuth2ClientConfig,
+    get_env_csrf_trusted_origins,
     get_env_host,
     get_env_port,
     server_instrumentation_is_enabled,
@@ -131,6 +136,7 @@ from phoenix.trace.fixtures import (
 from phoenix.trace.otel import decode_otlp_span, encode_span_to_otlp
 from phoenix.trace.schemas import Span
 from phoenix.utilities.client import PHOENIX_SERVER_VERSION_HEADER
+from phoenix.version import __version__ as phoenix_version
 
 if TYPE_CHECKING:
     from opentelemetry.trace import TracerProvider
@@ -213,7 +219,7 @@ class Static(StaticFiles):
                     "n_neighbors": self._app_config.n_neighbors,
                     "n_samples": self._app_config.n_samples,
                     "basename": self._sanitize_basename(request.scope.get("root_path", "")),
-                    "platform_version": phoenix.__version__,
+                    "platform_version": phoenix_version,
                     "request": request,
                     "is_development": self._app_config.is_development,
                     "manifest": self._web_manifest,
@@ -226,13 +232,32 @@ class Static(StaticFiles):
         return response
 
 
+class RequestOriginHostnameValidator(BaseHTTPMiddleware):
+    def __init__(self, trusted_hostnames: List[str], *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._trusted_hostnames = trusted_hostnames
+
+    async def dispatch(
+        self,
+        request: Request,
+        call_next: RequestResponseEndpoint,
+    ) -> Response:
+        headers = request.headers
+        for key in "origin", "referer":
+            if not (url := headers.get(key)):
+                continue
+            if urlparse(url).hostname not in self._trusted_hostnames:
+                return Response(f"untrusted {key}", status_code=HTTP_401_UNAUTHORIZED)
+        return await call_next(request)
+
+
 class HeadersMiddleware(BaseHTTPMiddleware):
     async def dispatch(
         self,
         request: Request,
         call_next: RequestResponseEndpoint,
     ) -> Response:
-        from phoenix import __version__ as phoenix_version
+        from phoenix.version import __version__ as phoenix_version
 
         response = await call_next(request)
         response.headers["x-colab-notebook-cache-control"] = "no-cache"
@@ -245,7 +270,7 @@ ProjectRowId: TypeAlias = int
 
 @router.get("/arize_phoenix_version")
 async def version() -> PlainTextResponse:
-    return PlainTextResponse(f"{phoenix.__version__}")
+    return PlainTextResponse(f"{phoenix_version}")
 
 
 DB_MUTEX: Optional[asyncio.Lock] = None
@@ -557,6 +582,7 @@ def create_graphql_router(
         include_in_schema=False,
         prefix="/graphql",
         dependencies=(Depends(is_authenticated),) if authentication_enabled else (),
+        subscription_protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL],
     )
 
 
@@ -607,6 +633,29 @@ async def plain_text_http_exception_handler(request: Request, exc: HTTPException
     return PlainTextResponse(str(exc.detail), status_code=exc.status_code, headers=headers)
 
 
+async def websocket_denial_response_handler(websocket: WebSocket, exc: WebSocketException) -> None:
+    """
+    Overrides the default exception handler for WebSocketException to ensure
+    that the HTTP response returned when a WebSocket connection is denied has
+    the same status code as the raised exception. This is in keeping with the
+    WebSocket Denial Response Extension of the ASGI specificiation described
+    below.
+
+    "Websocket connections start with the client sending a HTTP request
+    containing the appropriate upgrade headers. On receipt of this request a
+    server can choose to either upgrade the connection or respond with an HTTP
+    response (denying the upgrade). The core ASGI specification does not allow
+    for any control over the denial response, instead specifying that the HTTP
+    status code 403 should be returned, whereas this extension allows an ASGI
+    framework to control the denial response."
+
+    For details, see:
+    - https://asgi.readthedocs.io/en/latest/extensions.html#websocket-denial-response
+    """
+    assert isinstance(exc, WebSocketException)
+    await websocket.send_denial_response(JSONResponse(status_code=exc.code, content=exc.reason))
+
+
 def create_app(
     db: DbSessionFactory,
     export_path: Path,
@@ -660,6 +709,16 @@ def create_app(
     )
     last_updated_at = LastUpdatedAt()
     middlewares: List[Middleware] = [Middleware(HeadersMiddleware)]
+    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))
+    elif email_sender or oauth2_client_configs:
+        logger.warning(
+            "CSRF protection can be enabled by listing trusted origins via "
+            f"the `{ENV_PHOENIX_CSRF_TRUSTED_ORIGINS}` environment variable. "
+            "This is recommended when setting up OAuth2 clients or sending "
+            "password reset emails."
+        )
     if authentication_enabled and secret:
         token_store = JwtStore(db, secret)
         middlewares.append(
@@ -743,7 +802,10 @@ def create_app(
             scaffolder_config=scaffolder_config,
         ),
         middleware=middlewares,
-        exception_handlers={HTTPException: plain_text_http_exception_handler},
+        exception_handlers={
+            HTTPException: plain_text_http_exception_handler,
+            WebSocketException: websocket_denial_response_handler,  # type: ignore[dict-item]
+        },
         debug=debug,
         swagger_ui_parameters={
             "defaultModelsExpandDepth": -1,  # hides the schema section in the Swagger UI

phoenix/server/bearer_auth.py

@@ -7,10 +7,11 @@ from typing import (
     Callable,
     Optional,
     Tuple,
+    cast,
 )
 
 import grpc
-from fastapi import HTTPException, Request
+from fastapi import HTTPException, Request, WebSocket, WebSocketException
 from grpc_interceptor import AsyncServerInterceptor
 from grpc_interceptor.exceptions import Unauthenticated
 from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
@@ -116,12 +117,19 @@ class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
         raise Unauthenticated()
 
 
-async def is_authenticated(request: Request) -> None:
+async def is_authenticated(
+    # fastapi dependencies require non-optional types
+    request: Request = cast(Request, None),
+    websocket: WebSocket = cast(WebSocket, None),
+) -> None:
     """
-    Raises a 401 if the request is not authenticated.
+    Raises a 401 if the request or websocket connection is not authenticated.
     """
-    if not isinstance((user := request.user), PhoenixUser):
+    assert request or websocket
+    if request and not isinstance((user := request.user), PhoenixUser):
         raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+    if websocket and not isinstance((user := websocket.user), PhoenixUser):
+        raise WebSocketException(code=HTTP_401_UNAUTHORIZED, reason="Invalid token")
     claims = user.claims
     if claims.status is ClaimSetStatus.EXPIRED:
         raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Expired token")

phoenix/server/main.py

@@ -3,7 +3,6 @@ import codecs
 import os
 import sys
 from argparse import SUPPRESS, ArgumentParser
-from importlib.metadata import version
 from pathlib import Path
 from threading import Thread
 from time import sleep, time
@@ -72,6 +71,7 @@ from phoenix.trace.fixtures import (
 )
 from phoenix.trace.otel import decode_otlp_span, encode_span_to_otlp
 from phoenix.trace.schemas import Span
+from phoenix.version import __version__ as phoenix_version
 
 _WELCOME_MESSAGE = Environment(loader=BaseLoader()).from_string("""
 
@@ -351,7 +351,7 @@ def main() -> None:
     # Print information about the server
     root_path = urljoin(f"http://{host}:{port}", host_root_path)
     msg = _WELCOME_MESSAGE.render(
-        version=version("arize-phoenix"),
+        version=phoenix_version,
         ui_path=root_path,
         grpc_path=f"http://{host}:{get_env_grpc_port()}",
         http_path=urljoin(root_path, "v1/traces"),