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

phoenix/server/dml_event_handler.py

@@ -2,24 +2,10 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from asyncio import gather
+from collections.abc import Callable, Iterable, Iterator, Mapping
 from inspect import getmro
 from itertools import chain
-from typing import (
-    Any,
-    Callable,
-    Generic,
-    Iterable,
-    Iterator,
-    Mapping,
-    Optional,
-    Set,
-    Tuple,
-    Type,
-    TypedDict,
-    TypeVar,
-    Union,
-    cast,
-)
+from typing import Any, Generic, Optional, TypedDict, TypeVar, Union, cast
 
 from sqlalchemy import Select, select
 from typing_extensions import TypeAlias, Unpack
@@ -54,7 +40,7 @@ _DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)
 class _DmlEventQueue(Generic[_DmlEventT]):
     def __init__(self, **kwargs: Any) -> None:
         super().__init__(**kwargs)
-        self._events: Set[_DmlEventT] = set()
+        self._events: set[_DmlEventT] = set()
 
     @property
     def empty(self) -> bool:
@@ -120,7 +106,7 @@ class _GenericDmlEventHandler(_DmlEventHandler[DmlEvent]):
             for id_ in e.ids:
                 self._update(e.table, id_)
 
-    def _update(self, table: Type[Base], id_: int) -> None:
+    def _update(self, table: type[Base], id_: int) -> None:
         self._last_updated_at.set(table, id_)
 
 
@@ -146,9 +132,9 @@ class _SpanDeleteEventHandler(_SpanDmlEventHandler):
 
 
 _AnnotationTable: TypeAlias = Union[
-    Type[SpanAnnotation],
-    Type[TraceAnnotation],
-    Type[DocumentAnnotation],
+    type[SpanAnnotation],
+    type[TraceAnnotation],
+    type[DocumentAnnotation],
 ]
 
 _AnnotationDmlEventT = TypeVar(
@@ -165,7 +151,7 @@ class _AnnotationDmlEventHandler(
     ABC,
 ):
     _table: _AnnotationTable
-    _base_stmt: Union[Select[Tuple[int, str]], Select[Tuple[int]]] = (
+    _base_stmt: Union[Select[tuple[int, str]], Select[tuple[int]]] = (
         select(Project.id).join_from(Project, Trace).distinct()
     )
 
@@ -175,7 +161,7 @@ class _AnnotationDmlEventHandler(
         if self._cache_for_dataloaders:
             self._stmt = self._stmt.add_columns(self._table.name)
 
-    def _get_stmt(self) -> Union[Select[Tuple[int, str]], Select[Tuple[int]]]:
+    def _get_stmt(self) -> Union[Select[tuple[int, str]], Select[tuple[int]]]:
         ids = set(chain.from_iterable(e.ids for e in self._batch))
         return self._stmt.where(self._table.id.in_(ids))
 
@@ -242,7 +228,7 @@ class DmlEventHandler:
             cache_for_dataloaders=cache_for_dataloaders,
             sleep_seconds=sleep_seconds,
         )
-        self._handlers: Mapping[Type[DmlEvent], Iterable[_DmlEventHandler[Any]]] = {
+        self._handlers: Mapping[type[DmlEvent], Iterable[_DmlEventHandler[Any]]] = {
             DmlEvent: [_GenericDmlEventHandler(**kwargs)],
             SpanDmlEvent: [_SpanDmlEventHandler(**kwargs)],
             SpanDeleteEvent: [_SpanDeleteEventHandler(**kwargs)],

phoenix/server/grpc_server.py

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, List, Optional
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Any, Optional
 
 import grpc
 from grpc.aio import RpcContext, Server, ServerInterceptor
@@ -66,7 +67,7 @@ class GrpcServer:
     async def __aenter__(self) -> None:
         if self._disabled:
             return
-        interceptors: List[ServerInterceptor] = []
+        interceptors: list[ServerInterceptor] = []
         if self._token_store:
             interceptors.append(ApiKeyInterceptor(self._token_store))
         if self._enable_prometheus:

phoenix/server/jwt_store.py

@@ -1,11 +1,12 @@
 import logging
 from abc import ABC, abstractmethod
 from asyncio import create_task, gather, sleep
+from collections.abc import Callable, Coroutine
 from copy import deepcopy
 from dataclasses import replace
 from datetime import datetime, timezone
 from functools import cached_property, singledispatchmethod
-from typing import Any, Callable, Coroutine, Dict, Generic, List, Optional, Tuple, Type, TypeVar
+from typing import Any, Generic, Optional, TypeVar
 
 from authlib.jose import jwt
 from authlib.jose.errors import JoseError
@@ -65,7 +66,7 @@ class JwtStore:
         self._api_key_store = _ApiKeyStore(*args, **kwargs)
 
     @cached_property
-    def _stores(self) -> Tuple[DaemonTask, ...]:
+    def _stores(self) -> tuple[DaemonTask, ...]:
         return tuple(dt for dt in self.__dict__.values() if isinstance(dt, _Store))
 
     async def __aenter__(self) -> None:
@@ -131,34 +132,34 @@ class JwtStore:
     async def create_password_reset_token(
         self,
         claim: PasswordResetTokenClaims,
-    ) -> Tuple[PasswordResetToken, PasswordResetTokenId]:
+    ) -> tuple[PasswordResetToken, PasswordResetTokenId]:
         return await self._password_reset_token_store.create(claim)
 
     async def create_access_token(
         self,
         claim: AccessTokenClaims,
-    ) -> Tuple[AccessToken, AccessTokenId]:
+    ) -> tuple[AccessToken, AccessTokenId]:
         return await self._access_token_store.create(claim)
 
     async def create_refresh_token(
         self,
         claim: RefreshTokenClaims,
-    ) -> Tuple[RefreshToken, RefreshTokenId]:
+    ) -> tuple[RefreshToken, RefreshTokenId]:
         return await self._refresh_token_store.create(claim)
 
     async def create_api_key(
         self,
         claim: ApiKeyClaims,
-    ) -> Tuple[ApiKey, ApiKeyId]:
+    ) -> tuple[ApiKey, ApiKeyId]:
         return await self._api_key_store.create(claim)
 
     async def revoke(self, *token_ids: TokenId) -> None:
         if not token_ids:
             return
-        password_reset_token_ids: List[PasswordResetTokenId] = []
-        access_token_ids: List[AccessTokenId] = []
-        refresh_token_ids: List[RefreshTokenId] = []
-        api_key_ids: List[ApiKeyId] = []
+        password_reset_token_ids: list[PasswordResetTokenId] = []
+        access_token_ids: list[AccessTokenId] = []
+        refresh_token_ids: list[RefreshTokenId] = []
+        api_key_ids: list[ApiKeyId] = []
         for token_id in token_ids:
             if isinstance(token_id, PasswordResetTokenId):
                 password_reset_token_ids.append(token_id)
@@ -168,7 +169,7 @@ class JwtStore:
                 refresh_token_ids.append(token_id)
             elif isinstance(token_id, ApiKeyId):
                 api_key_ids.append(token_id)
-        coroutines: List[Coroutine[None, None, None]] = []
+        coroutines: list[Coroutine[None, None, None]] = []
         if password_reset_token_ids:
             coroutines.append(self._password_reset_token_store.revoke(*password_reset_token_ids))
         if access_token_ids:
@@ -202,7 +203,7 @@ _RecordT = TypeVar(
 
 class _Claims(Generic[_TokenIdT, _ClaimSetT]):
     def __init__(self) -> None:
-        self._cache: Dict[_TokenIdT, _ClaimSetT] = {}
+        self._cache: dict[_TokenIdT, _ClaimSetT] = {}
 
     def __getitem__(self, token_id: _TokenIdT) -> Optional[_ClaimSetT]:
         claim = self._cache.get(token_id)
@@ -223,7 +224,7 @@ class _Claims(Generic[_TokenIdT, _ClaimSetT]):
 
 
 class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC):
-    _table: Type[_RecordT]
+    _table: type[_RecordT]
     _token_id: Callable[[int], _TokenIdT]
     _token: Callable[[str], _TokenT]
 
@@ -244,7 +245,7 @@ class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC)
         self._algorithm = algorithm
 
     def _encode(self, claim: ClaimSet) -> str:
-        payload: Dict[str, Any] = dict(jti=claim.token_id)
+        payload: dict[str, Any] = dict(jti=claim.token_id)
         header = {"alg": self._algorithm}
         jwt_bytes: bytes = jwt.encode(header=header, payload=payload, key=self._secret)
         return jwt_bytes.decode()
@@ -275,12 +276,12 @@ class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC)
             await session.execute(stmt)
 
     @abstractmethod
-    def _from_db(self, record: _RecordT, role: UserRole) -> Tuple[_TokenIdT, _ClaimSetT]: ...
+    def _from_db(self, record: _RecordT, role: UserRole) -> tuple[_TokenIdT, _ClaimSetT]: ...
 
     @abstractmethod
     def _to_db(self, claims: _ClaimSetT) -> _RecordT: ...
 
-    async def create(self, claim: _ClaimSetT) -> Tuple[_TokenT, _TokenIdT]:
+    async def create(self, claim: _ClaimSetT) -> tuple[_TokenT, _TokenIdT]:
         record = self._to_db(claim)
         async with self._db() as session:
             session.add(record)
@@ -303,7 +304,7 @@ class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC)
         self._claims = claims
 
     @cached_property
-    def _update_stmt(self) -> Select[Tuple[_RecordT, str]]:
+    def _update_stmt(self) -> Select[tuple[_RecordT, str]]:
         return (
             select(self._table, models.UserRole.name)
             .join_from(self._table, models.User)
@@ -340,7 +341,7 @@ class _PasswordResetTokenStore(
         self,
         record: models.PasswordResetToken,
         user_role: UserRole,
-    ) -> Tuple[PasswordResetTokenId, PasswordResetTokenClaims]:
+    ) -> tuple[PasswordResetTokenId, PasswordResetTokenClaims]:
         token_id = PasswordResetTokenId(record.id)
         return token_id, PasswordResetTokenClaims(
             token_id=token_id,
@@ -379,7 +380,7 @@ class _AccessTokenStore(
         self,
         record: models.AccessToken,
         user_role: UserRole,
-    ) -> Tuple[AccessTokenId, AccessTokenClaims]:
+    ) -> tuple[AccessTokenId, AccessTokenClaims]:
         token_id = AccessTokenId(record.id)
         refresh_token_id = RefreshTokenId(record.refresh_token_id)
         return token_id, AccessTokenClaims(
@@ -423,7 +424,7 @@ class _RefreshTokenStore(
         self,
         record: models.RefreshToken,
         user_role: UserRole,
-    ) -> Tuple[RefreshTokenId, RefreshTokenClaims]:
+    ) -> tuple[RefreshTokenId, RefreshTokenClaims]:
         token_id = RefreshTokenId(record.id)
         return token_id, RefreshTokenClaims(
             token_id=token_id,
@@ -469,7 +470,7 @@ class _ApiKeyStore(
         self,
         record: models.ApiKey,
         user_role: UserRole,
-    ) -> Tuple[ApiKeyId, ApiKeyClaims]:
+    ) -> tuple[ApiKeyId, ApiKeyClaims]:
         token_id = ApiKeyId(record.id)
         return token_id, ApiKeyClaims(
             token_id=token_id,

phoenix/server/main.py

@@ -6,7 +6,7 @@ from argparse import SUPPRESS, ArgumentParser
 from pathlib import Path
 from threading import Thread
 from time import sleep, time
-from typing import List, Optional
+from typing import Optional
 from urllib.parse import urljoin
 
 from jinja2 import BaseLoader, Environment
@@ -21,6 +21,7 @@ from phoenix.config import (
     get_env_database_schema,
     get_env_db_logging_level,
     get_env_enable_prometheus,
+    get_env_enable_websockets,
     get_env_grpc_port,
     get_env_host,
     get_env_host_root_path,
@@ -95,6 +96,7 @@ _WELCOME_MESSAGE = Environment(loader=BaseLoader()).from_string("""
 |  🚀 Phoenix Server 🚀
 |  Phoenix UI: {{ ui_path }}
 |  Authentication: {{ auth_enabled }}
+|  Websockets: {{ websockets_enabled }}
 |  Log traces:
 |    - gRPC: {{ grpc_path }}
 |    - HTTP: {{ http_path }}
@@ -162,7 +164,7 @@ def main() -> None:
     parser.add_argument("--debug", action="store_true", help=SUPPRESS)
     parser.add_argument("--dev", action="store_true", help=SUPPRESS)
     parser.add_argument("--no-ui", action="store_true", help=SUPPRESS)
-
+    parser.add_argument("--enable-websockets", type=str, help=SUPPRESS)
     subparsers = parser.add_subparsers(dest="command", required=True, help=SUPPRESS)
 
     serve_parser = subparsers.add_parser("serve")
@@ -312,8 +314,8 @@ def main() -> None:
 
     authentication_enabled, secret = get_env_auth_settings()
 
-    fixture_spans: List[Span] = []
-    fixture_evals: List[pb.Evaluation] = []
+    fixture_spans: list[Span] = []
+    fixture_evals: list[pb.Evaluation] = []
     if trace_dataset_name is not None:
         fixture_spans, fixture_evals = reset_fixture_span_ids_and_timestamps(
             (
@@ -348,6 +350,14 @@ def main() -> None:
     corpus_model = (
         None if corpus_inferences is None else create_model_from_inferences(corpus_inferences)
     )
+
+    # Get enable_websockets from environment variable or command line argument
+    enable_websockets = get_env_enable_websockets()
+    if args.enable_websockets is not None:
+        enable_websockets = args.enable_websockets.lower() == "true"
+    if enable_websockets is None:
+        enable_websockets = True
+
     # Print information about the server
     root_path = urljoin(f"http://{host}:{port}", host_root_path)
     msg = _WELCOME_MESSAGE.render(
@@ -358,6 +368,7 @@ def main() -> None:
         storage=get_printable_db_url(db_connection_str),
         schema=get_env_database_schema(),
         auth_enabled=authentication_enabled,
+        websockets_enabled=enable_websockets,
     )
     if sys.platform.startswith("win"):
         msg = codecs.encode(msg, "ascii", errors="ignore").decode("ascii").strip()
@@ -382,10 +393,12 @@ def main() -> None:
             connection_method="STARTTLS",
             validate_certs=get_env_smtp_validate_certs(),
         )
+
     app = create_app(
         db=factory,
         export_path=export_path,
         model=model,
+        enable_websockets=enable_websockets,
         authentication_enabled=authentication_enabled,
         umap_params=umap_params,
         corpus=corpus_model,

phoenix/server/oauth2.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, Iterable
+from collections.abc import Iterable
+from typing import Any
 
 from authlib.integrations.base_client import BaseApp
 from authlib.integrations.base_client.async_app import AsyncOAuth2Mixin
@@ -24,7 +25,7 @@ class OAuth2Client(AsyncOAuth2Mixin, AsyncOpenIDMixin, BaseApp):  # type:ignore[
 
 class OAuth2Clients:
     def __init__(self) -> None:
-        self._clients: Dict[str, OAuth2Client] = {}
+        self._clients: dict[str, OAuth2Client] = {}
 
     def add_client(self, config: OAuth2ClientConfig) -> None:
         if (idp_name := config.idp_name) in self._clients:

phoenix/server/rate_limiters.py

@@ -1,13 +1,10 @@
 import re
 import time
 from collections import defaultdict
+from collections.abc import Callable, Coroutine
 from functools import partial
 from typing import (
     Any,
-    Callable,
-    Coroutine,
-    DefaultDict,
-    List,
     Optional,
     Pattern,  # import from re module when we drop support for 3.8
     Union,
@@ -98,7 +95,7 @@ class ServerRateLimiter:
         self._last_cleanup_time = time.time()
 
     def _reset_rate_limiters(self) -> None:
-        self.cache_partitions: List[DefaultDict[Any, TokenBucket]] = [
+        self.cache_partitions: list[defaultdict[Any, TokenBucket]] = [
             defaultdict(self.bucket_factory) for _ in range(self.num_partitions)
         ]
 
@@ -107,10 +104,10 @@ class ServerRateLimiter:
             int(timestamp // self.partition_seconds) % self.num_partitions
         )  # a cyclic bucket index
 
-    def _active_partition_indices(self, current_index: int) -> List[int]:
+    def _active_partition_indices(self, current_index: int) -> list[int]:
         return [(current_index - ii) % self.num_partitions for ii in range(self.active_partitions)]
 
-    def _inactive_partition_indices(self, current_index: int) -> List[int]:
+    def _inactive_partition_indices(self, current_index: int) -> list[int]:
         active_indices = set(self._active_partition_indices(current_index))
         all_indices = set(range(self.num_partitions))
         return list(all_indices - active_indices)
@@ -156,7 +153,7 @@ class ServerRateLimiter:
 
 
 def fastapi_ip_rate_limiter(
-    rate_limiter: ServerRateLimiter, paths: Optional[List[Union[str, Pattern[str]]]] = None
+    rate_limiter: ServerRateLimiter, paths: Optional[list[Union[str, Pattern[str]]]] = None
 ) -> Callable[[Request], Coroutine[Any, Any, Request]]:
     async def dependency(request: Request) -> Request:
         if paths is None or any(path_match(request.url.path, path) for path in paths):

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

@@ -1,32 +1,32 @@
 {
-  "_components-hX0LgYz3.js": {
-    "file": "assets/components-hX0LgYz3.js",
+  "_components-Csu8UKOs.js": {
+    "file": "assets/components-Csu8UKOs.js",
     "name": "components",
     "imports": [
-      "_vendor-6IcPAw_j.js",
-      "_vendor-arizeai-DRZuoyuF.js",
-      "_pages-CAL1FDMt.js",
+      "_vendor-CtqfhlbC.js",
+      "_vendor-arizeai-C_3SBz56.js",
+      "_pages-UeWaKXNs.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-codemirror-DVE2_WBr.js"
+      "_vendor-codemirror-wfdk9cjp.js"
     ]
   },
-  "_pages-CAL1FDMt.js": {
-    "file": "assets/pages-CAL1FDMt.js",
+  "_pages-UeWaKXNs.js": {
+    "file": "assets/pages-UeWaKXNs.js",
     "name": "pages",
     "imports": [
-      "_vendor-6IcPAw_j.js",
-      "_vendor-arizeai-DRZuoyuF.js",
-      "_components-hX0LgYz3.js",
-      "_vendor-recharts-DwrexFA4.js",
-      "_vendor-codemirror-DVE2_WBr.js"
+      "_vendor-CtqfhlbC.js",
+      "_vendor-arizeai-C_3SBz56.js",
+      "_components-Csu8UKOs.js",
+      "_vendor-recharts-BiVnSv90.js",
+      "_vendor-codemirror-wfdk9cjp.js"
     ]
   },
   "_vendor-!~{003}~.js": {
     "file": "assets/vendor-DxkFTwjz.css",
     "src": "_vendor-!~{003}~.js"
   },
-  "_vendor-6IcPAw_j.js": {
-    "file": "assets/vendor-6IcPAw_j.js",
+  "_vendor-CtqfhlbC.js": {
+    "file": "assets/vendor-CtqfhlbC.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-DwGkEfCM.js"
@@ -35,25 +35,25 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-DRZuoyuF.js": {
-    "file": "assets/vendor-arizeai-DRZuoyuF.js",
+  "_vendor-arizeai-C_3SBz56.js": {
+    "file": "assets/vendor-arizeai-C_3SBz56.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-6IcPAw_j.js"
+      "_vendor-CtqfhlbC.js"
     ]
   },
-  "_vendor-codemirror-DVE2_WBr.js": {
-    "file": "assets/vendor-codemirror-DVE2_WBr.js",
+  "_vendor-codemirror-wfdk9cjp.js": {
+    "file": "assets/vendor-codemirror-wfdk9cjp.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-6IcPAw_j.js"
+      "_vendor-CtqfhlbC.js"
     ]
   },
-  "_vendor-recharts-DwrexFA4.js": {
-    "file": "assets/vendor-recharts-DwrexFA4.js",
+  "_vendor-recharts-BiVnSv90.js": {
+    "file": "assets/vendor-recharts-BiVnSv90.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-6IcPAw_j.js"
+      "_vendor-CtqfhlbC.js"
     ]
   },
   "_vendor-three-DwGkEfCM.js": {
@@ -61,18 +61,18 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-DCzakdJq.js",
+    "file": "assets/index-Bk5C9EA7.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-6IcPAw_j.js",
-      "_vendor-arizeai-DRZuoyuF.js",
-      "_pages-CAL1FDMt.js",
-      "_components-hX0LgYz3.js",
+      "_vendor-CtqfhlbC.js",
+      "_vendor-arizeai-C_3SBz56.js",
+      "_pages-UeWaKXNs.js",
+      "_components-Csu8UKOs.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-recharts-DwrexFA4.js",
-      "_vendor-codemirror-DVE2_WBr.js"
+      "_vendor-recharts-BiVnSv90.js",
+      "_vendor-codemirror-wfdk9cjp.js"
     ]
   }
 }