arize-phoenix 4.35.2__py3-none-any.whl → 5.0.0__py3-none-any.whl

phoenix/server/main.py

@@ -1,6 +1,5 @@
 import atexit
 import codecs
-import logging
 import os
 import sys
 from argparse import ArgumentParser
@@ -11,27 +10,42 @@ from time import sleep, time
 from typing import List, Optional
 from urllib.parse import urljoin
 
+from fastapi_mail import ConnectionConfig
 from jinja2 import BaseLoader, Environment
 from uvicorn import Config, Server
 
 import phoenix.trace.v1 as pb
 from phoenix.config import (
     EXPORT_DIR,
-    get_auth_settings,
+    get_env_access_token_expiry,
+    get_env_auth_settings,
     get_env_database_connection_str,
     get_env_database_schema,
+    get_env_db_logging_level,
     get_env_enable_prometheus,
     get_env_grpc_port,
     get_env_host,
     get_env_host_root_path,
+    get_env_log_migrations,
+    get_env_logging_level,
+    get_env_logging_mode,
+    get_env_oauth2_settings,
+    get_env_password_reset_token_expiry,
     get_env_port,
+    get_env_refresh_token_expiry,
+    get_env_smtp_hostname,
+    get_env_smtp_mail_from,
+    get_env_smtp_password,
+    get_env_smtp_port,
+    get_env_smtp_username,
+    get_env_smtp_validate_certs,
     get_pids_path,
-    get_working_dir,
 )
 from phoenix.core.model_schema_adapter import create_model_from_inferences
 from phoenix.db import get_printable_db_url
 from phoenix.inferences.fixtures import FIXTURES, get_inferences
 from phoenix.inferences.inferences import EMPTY_INFERENCES, Inferences
+from phoenix.logging import setup_logging
 from phoenix.pointcloud.umap_parameters import (
     DEFAULT_MIN_DIST,
     DEFAULT_N_NEIGHBORS,
@@ -45,6 +59,7 @@ from phoenix.server.app import (
     create_engine_and_run_migrations,
     instrument_engine_if_enabled,
 )
+from phoenix.server.email.sender import EMAIL_TEMPLATE_FOLDER, FastMailSender
 from phoenix.server.types import DbSessionFactory
 from phoenix.settings import Settings
 from phoenix.trace.fixtures import (
@@ -59,9 +74,6 @@ from phoenix.trace.fixtures import (
 from phoenix.trace.otel import decode_otlp_span, encode_span_to_otlp
 from phoenix.trace.schemas import Span
 
-logger = logging.getLogger(__name__)
-logger.addHandler(logging.NullHandler())
-
 _WELCOME_MESSAGE = Environment(loader=BaseLoader()).from_string("""
 
 ██████╗ ██╗  ██╗ ██████╗ ███████╗███╗   ██╗██╗██╗  ██╗
@@ -83,6 +95,7 @@ _WELCOME_MESSAGE = Environment(loader=BaseLoader()).from_string("""
 |
 |  🚀 Phoenix Server 🚀
 |  Phoenix UI: {{ ui_path }}
+|  Authentication: {{ auth_enabled }}
 |  Log traces:
 |    - gRPC: {{ grpc_path }}
 |    - HTTP: {{ http_path }}
@@ -92,11 +105,6 @@ _WELCOME_MESSAGE = Environment(loader=BaseLoader()).from_string("""
 {% endif -%}
 """)
 
-_EXPERIMENTAL_WARNING = """
-🚨 WARNING: Phoenix is running in experimental mode. 🚨
-|  Authentication enabled: {auth_enabled}
-"""
-
 
 def _write_pid_file_when_ready(
     server: Server,
@@ -121,19 +129,16 @@ def _get_pid_file() -> Path:
 
 DEFAULT_UMAP_PARAMS_STR = f"{DEFAULT_MIN_DIST},{DEFAULT_N_NEIGHBORS},{DEFAULT_N_SAMPLES}"
 
-if __name__ == "__main__":
+
+def main() -> None:
     primary_inferences_name: str
     reference_inferences_name: Optional[str]
     trace_dataset_name: Optional[str] = None
-    simulate_streaming: Optional[bool] = None
 
     primary_inferences: Inferences = EMPTY_INFERENCES
     reference_inferences: Optional[Inferences] = None
     corpus_inferences: Optional[Inferences] = None
 
-    # Initialize the settings for the Server
-    Settings.log_migrations = True
-
     # automatically remove the pid file when the process is being gracefully terminated
     atexit.register(_remove_pid_file)
 
@@ -230,7 +235,7 @@ if __name__ == "__main__":
     db_connection_str = (
         args.database_url if args.database_url else get_env_database_connection_str()
     )
-    export_path = Path(args.export_path) if args.export_path else EXPORT_DIR
+    export_path = Path(args.export_path) if args.export_path else Path(EXPORT_DIR)
 
     force_fixture_ingestion = False
     scaffold_datasets = False
@@ -260,7 +265,6 @@ if __name__ == "__main__":
             reference_inferences = None
     elif args.command == "trace-fixture":
         trace_dataset_name = args.fixture
-        simulate_streaming = args.simulate_streaming
     elif args.command == "demo":
         fixture_name = args.fixture
         primary_inferences, reference_inferences, corpus_inferences = get_inferences(
@@ -268,7 +272,6 @@ if __name__ == "__main__":
             args.no_internet,
         )
         trace_dataset_name = args.trace_fixture
-        simulate_streaming = args.simulate_streaming
     elif args.command == "serve":
         # We use sets to avoid duplicates
         if args.with_fixture:
@@ -290,12 +293,6 @@ if __name__ == "__main__":
         force_fixture_ingestion = args.force_fixture_ingestion
         scaffold_datasets = args.scaffold_datasets
     host: Optional[str] = args.host or get_env_host()
-    display_host = host or "localhost"
-    # If the host is "::", the convention is to bind to all interfaces. However, uvicorn
-    # does not support this directly unless the host is set to None.
-    if host and ":" in host:
-        # format IPv6 hosts in brackets
-        display_host = f"[{host}]"
     if host == "::":
         # TODO(dustin): why is this necessary? it's not type compliant
         host = None
@@ -309,7 +306,7 @@ if __name__ == "__main__":
         reference_inferences,
     )
 
-    authentication_enabled, secret = get_auth_settings()
+    authentication_enabled, secret = get_env_auth_settings()
 
     fixture_spans: List[Span] = []
     fixture_evals: List[pb.Evaluation] = []
@@ -336,13 +333,11 @@ if __name__ == "__main__":
         n_samples=int(umap_params_list[2]),
     )
 
-    logger.info(f"Server umap params: {umap_params}")
     if enable_prometheus := get_env_enable_prometheus():
         from phoenix.server.prometheus import start_prometheus
 
         start_prometheus()
 
-    working_dir = get_working_dir().resolve()
     engine = create_engine_and_run_migrations(db_connection_str)
     instrumentation_cleanups = instrument_engine_if_enabled(engine)
     factory = DbSessionFactory(db=_db(engine), dialect=engine.dialect.name)
@@ -358,9 +353,8 @@ if __name__ == "__main__":
         http_path=urljoin(root_path, "v1/traces"),
         storage=get_printable_db_url(db_connection_str),
         schema=get_env_database_schema(),
+        auth_enabled=authentication_enabled,
     )
-    if authentication_enabled:
-        msg += _EXPERIMENTAL_WARNING.format(auth_enabled=True)
     if sys.platform.startswith("win"):
         msg = codecs.encode(msg, "ascii", errors="ignore").decode("ascii").strip()
     scaffolder_config = ScaffolderConfig(
@@ -370,6 +364,25 @@ if __name__ == "__main__":
         scaffold_datasets=scaffold_datasets,
         phoenix_url=root_path,
     )
+    email_sender = None
+    if mail_sever := get_env_smtp_hostname():
+        assert (mail_username := get_env_smtp_username()), "SMTP username is required"
+        assert (mail_password := get_env_smtp_password()), "SMTP password is required"
+        assert (mail_from := get_env_smtp_mail_from()), "SMTP mail_from is required"
+        email_sender = FastMailSender(
+            ConnectionConfig(
+                MAIL_USERNAME=mail_username,
+                MAIL_PASSWORD=mail_password,
+                MAIL_FROM=mail_from,
+                MAIL_SERVER=mail_sever,
+                MAIL_PORT=get_env_smtp_port(),
+                VALIDATE_CERTS=get_env_smtp_validate_certs(),
+                USE_CREDENTIALS=True,
+                MAIL_STARTTLS=True,
+                MAIL_SSL_TLS=False,
+                TEMPLATE_FOLDER=EMAIL_TEMPLATE_FOLDER,
+            )
+        )
     app = create_app(
         db=factory,
         export_path=export_path,
@@ -387,10 +400,28 @@ if __name__ == "__main__":
         startup_callbacks=[lambda: print(msg)],
         shutdown_callbacks=instrumentation_cleanups,
         secret=secret,
+        password_reset_token_expiry=get_env_password_reset_token_expiry(),
+        access_token_expiry=get_env_access_token_expiry(),
+        refresh_token_expiry=get_env_refresh_token_expiry(),
         scaffolder_config=scaffolder_config,
+        email_sender=email_sender,
+        oauth2_client_configs=get_env_oauth2_settings(),
     )
     server = Server(config=Config(app, host=host, port=port, root_path=host_root_path))  # type: ignore
     Thread(target=_write_pid_file_when_ready, args=(server,), daemon=True).start()
 
     # Start the server
     server.run()
+
+
+def initialize_settings() -> None:
+    Settings.logging_mode = get_env_logging_mode()
+    Settings.logging_level = get_env_logging_level()
+    Settings.db_logging_level = get_env_db_logging_level()
+    Settings.log_migrations = get_env_log_migrations()
+
+
+if __name__ == "__main__":
+    initialize_settings()
+    setup_logging()
+    main()

phoenix/server/oauth2.py

@@ -0,0 +1,51 @@
+from typing import Any, Dict, Iterable
+
+from authlib.integrations.base_client import BaseApp
+from authlib.integrations.base_client.async_app import AsyncOAuth2Mixin
+from authlib.integrations.base_client.async_openid import AsyncOpenIDMixin
+from authlib.integrations.httpx_client import AsyncOAuth2Client as AsyncHttpxOAuth2Client
+
+from phoenix.config import OAuth2ClientConfig
+
+
+class OAuth2Client(AsyncOAuth2Mixin, AsyncOpenIDMixin, BaseApp):  # type:ignore[misc]
+    """
+    An OAuth2 client class that supports OpenID Connect. Adapted from authlib's
+    `StarletteOAuth2App` to be useable without integration with Starlette.
+
+    https://github.com/lepture/authlib/blob/904d66bebd79bf39fb8814353a22bab7d3e092c4/authlib/integrations/starlette_client/apps.py#L58
+    """
+
+    client_cls = AsyncHttpxOAuth2Client
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(framework=None, *args, **kwargs)
+
+
+class OAuth2Clients:
+    def __init__(self) -> None:
+        self._clients: Dict[str, OAuth2Client] = {}
+
+    def add_client(self, config: OAuth2ClientConfig) -> None:
+        if (idp_name := config.idp_name) in self._clients:
+            raise ValueError(f"oauth client already registered: {idp_name}")
+        client = OAuth2Client(
+            client_id=config.client_id,
+            client_secret=config.client_secret,
+            server_metadata_url=config.oidc_config_url,
+            client_kwargs={"scope": "openid email profile"},
+        )
+        assert isinstance(client, OAuth2Client)
+        self._clients[config.idp_name] = client
+
+    def get_client(self, idp_name: str) -> OAuth2Client:
+        if (client := self._clients.get(idp_name)) is None:
+            raise ValueError(f"unknown or unregistered OAuth2 client: {idp_name}")
+        return client
+
+    @classmethod
+    def from_configs(cls, configs: Iterable[OAuth2ClientConfig]) -> "OAuth2Clients":
+        oauth2_clients = cls()
+        for config in configs:
+            oauth2_clients.add_client(config)
+        return oauth2_clients

phoenix/server/prometheus.py

@@ -50,6 +50,26 @@ BULK_LOADER_EXCEPTIONS = Counter(
     documentation="Total count of bulk loader exceptions",
 )
 
+RATE_LIMITER_CACHE_SIZE = Gauge(
+    name="rate_limiter_cache_size",
+    documentation="Current size of the rate limiter cache",
+)
+
+RATE_LIMITER_THROTTLES = Counter(
+    name="rate_limiter_throttles_total",
+    documentation="Total count of rate limiter throttles",
+)
+
+JWT_STORE_TOKENS_ACTIVE = Gauge(
+    name="jwt_store_tokens_active",
+    documentation="Current number of refresh tokens in the JWT store",
+)
+
+JWT_STORE_API_KEYS_ACTIVE = Gauge(
+    name="jwt_store_api_keys_active",
+    documentation="Current number of API keys in the JWT store",
+)
+
 
 class PrometheusMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:

phoenix/server/rate_limiters.py

@@ -0,0 +1,191 @@
+import re
+import time
+from collections import defaultdict
+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,
+)
+
+from fastapi import HTTPException, Request
+
+from phoenix.config import get_env_enable_prometheus
+from phoenix.exceptions import PhoenixException
+
+
+class UnavailableTokensError(PhoenixException):
+    pass
+
+
+class TokenBucket:
+    """
+    An implementation of the token-bucket algorithm for use as a rate limiter.
+
+    Args:
+        per_second_request_rate (float): The allowed request rate.
+        enforcement_window_minutes (float): The time window over which the rate limit is enforced.
+    """
+
+    def __init__(
+        self,
+        per_second_request_rate: float,
+        enforcement_window_seconds: float = 1,
+    ):
+        self.enforcement_window = enforcement_window_seconds
+        self.rate = per_second_request_rate
+
+        now = time.time()
+        self.last_checked = now
+        self.tokens = self.max_tokens()
+
+    def max_tokens(self) -> float:
+        return self.rate * self.enforcement_window
+
+    def available_tokens(self) -> float:
+        now = time.time()
+        time_since_last_checked = now - self.last_checked
+        self.tokens = min(self.max_tokens(), self.rate * time_since_last_checked + self.tokens)
+        self.last_checked = now
+        return self.tokens
+
+    def make_request_if_ready(self) -> None:
+        if self.available_tokens() < 1:
+            if get_env_enable_prometheus():
+                from phoenix.server.prometheus import RATE_LIMITER_THROTTLES
+
+                RATE_LIMITER_THROTTLES.inc()
+            raise UnavailableTokensError
+        self.tokens -= 1
+
+
+class ServerRateLimiter:
+    """
+    This rate limiter holds a cache of token buckets that enforce rate limits.
+
+    The cache is kept in partitions that rotate every `partition_seconds`. Each user's rate limiter
+    can be accessed from all active partitions, the number of active partitions is set with
+    `active_partitions`. This guarantees that a user's rate limiter will sit in the cache for at
+    least:
+
+        minimum_cache_lifetime = (active_partitions - 1) * partition_seconds
+
+    Every time the cache is accessed, inactive partitions are purged. If enough time has passed,
+    the entire cache is purged.
+    """
+
+    def __init__(
+        self,
+        per_second_rate_limit: float = 0.5,
+        enforcement_window_seconds: float = 5,
+        partition_seconds: float = 60,
+        active_partitions: int = 2,
+    ):
+        self.bucket_factory = partial(
+            TokenBucket,
+            per_second_request_rate=per_second_rate_limit,
+            enforcement_window_seconds=enforcement_window_seconds,
+        )
+        self.partition_seconds = partition_seconds
+        self.active_partitions = active_partitions
+        self.num_partitions = active_partitions + 2  # two overflow partitions to avoid edge cases
+        self._reset_rate_limiters()
+        self._last_cleanup_time = time.time()
+
+    def _reset_rate_limiters(self) -> None:
+        self.cache_partitions: List[DefaultDict[Any, TokenBucket]] = [
+            defaultdict(self.bucket_factory) for _ in range(self.num_partitions)
+        ]
+
+    def _current_partition_index(self, timestamp: float) -> int:
+        return (
+            int(timestamp // self.partition_seconds) % self.num_partitions
+        )  # a cyclic bucket index
+
+    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]:
+        active_indices = set(self._active_partition_indices(current_index))
+        all_indices = set(range(self.num_partitions))
+        return list(all_indices - active_indices)
+
+    def _cleanup_expired_limiters(self, request_time: float) -> None:
+        time_since_last_cleanup = request_time - self._last_cleanup_time
+        if time_since_last_cleanup >= ((self.num_partitions - 1) * self.partition_seconds):
+            # Reset the cache to avoid "looping" back to the same partitions
+            self._reset_rate_limiters()
+            self._last_cleanup_time = request_time
+            return
+
+        current_partition_index = self._current_partition_index(request_time)
+        inactive_indices = self._inactive_partition_indices(current_partition_index)
+        for ii in inactive_indices:
+            self.cache_partitions[ii] = defaultdict(self.bucket_factory)
+        self._last_cleanup_time = request_time
+
+    def _fetch_token_bucket(self, key: str, request_time: float) -> TokenBucket:
+        current_partition_index = self._current_partition_index(request_time)
+        active_indices = self._active_partition_indices(current_partition_index)
+        bucket: Optional[TokenBucket] = None
+        for ii in active_indices:
+            partition = self.cache_partitions[ii]
+            if key in partition:
+                bucket = partition.pop(key)
+                break
+
+        current_partition = self.cache_partitions[current_partition_index]
+        if key not in current_partition and bucket is not None:
+            current_partition[key] = bucket
+        return current_partition[key]
+
+    def make_request(self, key: str) -> None:
+        request_time = time.time()
+        self._cleanup_expired_limiters(request_time)
+        rate_limiter = self._fetch_token_bucket(key, request_time)
+        rate_limiter.make_request_if_ready()
+        if get_env_enable_prometheus():
+            from phoenix.server.prometheus import RATE_LIMITER_CACHE_SIZE
+
+            RATE_LIMITER_CACHE_SIZE.set(sum(len(partition) for partition in self.cache_partitions))
+
+
+def fastapi_ip_rate_limiter(
+    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):
+            client = request.client
+            if client:  # bypasses rate limiter if no client
+                client_ip = client.host
+                try:
+                    rate_limiter.make_request(client_ip)
+                except UnavailableTokensError:
+                    raise HTTPException(status_code=429, detail="Too Many Requests")
+        return request
+
+    return dependency
+
+
+def fastapi_route_rate_limiter(
+    rate_limiter: ServerRateLimiter,
+) -> Callable[[Request], Coroutine[Any, Any, Request]]:
+    async def dependency(request: Request) -> Request:
+        try:
+            rate_limiter.make_request(request.url.path)
+        except UnavailableTokensError:
+            raise HTTPException(status_code=429, detail="Too Many Requests")
+        return request
+
+    return dependency
+
+
+def path_match(path: str, match_pattern: Union[str, Pattern[str]]) -> bool:
+    if isinstance(match_pattern, re.Pattern):
+        return bool(match_pattern.match(path))
+    return path == match_pattern

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

@@ -1,32 +1,32 @@
 {
-  "_components-Dte7_KRd.js": {
-    "file": "assets/components-Dte7_KRd.js",
+  "_components-REunxTt6.js": {
+    "file": "assets/components-REunxTt6.js",
     "name": "components",
     "imports": [
-      "_vendor-BC3OPQuM.js",
-      "_vendor-arizeai-NjB3cZzD.js",
-      "_pages-CnTvEGEN.js",
+      "_vendor-B5IC0ivG.js",
+      "_vendor-arizeai-aFbT4kl1.js",
+      "_pages-1VrMk2pW.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-codemirror-gE_JCOgX.js"
+      "_vendor-codemirror-BEGorXSV.js"
     ]
   },
-  "_pages-CnTvEGEN.js": {
-    "file": "assets/pages-CnTvEGEN.js",
+  "_pages-1VrMk2pW.js": {
+    "file": "assets/pages-1VrMk2pW.js",
     "name": "pages",
     "imports": [
-      "_vendor-BC3OPQuM.js",
-      "_components-Dte7_KRd.js",
-      "_vendor-arizeai-NjB3cZzD.js",
-      "_vendor-recharts-BXLYwcXF.js",
-      "_vendor-codemirror-gE_JCOgX.js"
+      "_vendor-B5IC0ivG.js",
+      "_vendor-arizeai-aFbT4kl1.js",
+      "_components-REunxTt6.js",
+      "_vendor-recharts-6nUU7gU_.js",
+      "_vendor-codemirror-BEGorXSV.js"
     ]
   },
   "_vendor-!~{003}~.js": {
     "file": "assets/vendor-DxkFTwjz.css",
     "src": "_vendor-!~{003}~.js"
   },
-  "_vendor-BC3OPQuM.js": {
-    "file": "assets/vendor-BC3OPQuM.js",
+  "_vendor-B5IC0ivG.js": {
+    "file": "assets/vendor-B5IC0ivG.js",
     "name": "vendor",
     "imports": [
       "_vendor-three-DwGkEfCM.js"
@@ -35,25 +35,25 @@
       "assets/vendor-DxkFTwjz.css"
     ]
   },
-  "_vendor-arizeai-NjB3cZzD.js": {
-    "file": "assets/vendor-arizeai-NjB3cZzD.js",
+  "_vendor-arizeai-aFbT4kl1.js": {
+    "file": "assets/vendor-arizeai-aFbT4kl1.js",
     "name": "vendor-arizeai",
     "imports": [
-      "_vendor-BC3OPQuM.js"
+      "_vendor-B5IC0ivG.js"
     ]
   },
-  "_vendor-codemirror-gE_JCOgX.js": {
-    "file": "assets/vendor-codemirror-gE_JCOgX.js",
+  "_vendor-codemirror-BEGorXSV.js": {
+    "file": "assets/vendor-codemirror-BEGorXSV.js",
     "name": "vendor-codemirror",
     "imports": [
-      "_vendor-BC3OPQuM.js"
+      "_vendor-B5IC0ivG.js"
     ]
   },
-  "_vendor-recharts-BXLYwcXF.js": {
-    "file": "assets/vendor-recharts-BXLYwcXF.js",
+  "_vendor-recharts-6nUU7gU_.js": {
+    "file": "assets/vendor-recharts-6nUU7gU_.js",
     "name": "vendor-recharts",
     "imports": [
-      "_vendor-BC3OPQuM.js"
+      "_vendor-B5IC0ivG.js"
     ]
   },
   "_vendor-three-DwGkEfCM.js": {
@@ -61,18 +61,18 @@
     "name": "vendor-three"
   },
   "index.tsx": {
-    "file": "assets/index-fq1-hCK4.js",
+    "file": "assets/index-DAPJxlCw.js",
     "name": "index",
     "src": "index.tsx",
     "isEntry": true,
     "imports": [
-      "_vendor-BC3OPQuM.js",
-      "_vendor-arizeai-NjB3cZzD.js",
-      "_pages-CnTvEGEN.js",
-      "_components-Dte7_KRd.js",
+      "_vendor-B5IC0ivG.js",
+      "_vendor-arizeai-aFbT4kl1.js",
+      "_pages-1VrMk2pW.js",
+      "_components-REunxTt6.js",
       "_vendor-three-DwGkEfCM.js",
-      "_vendor-recharts-BXLYwcXF.js",
-      "_vendor-codemirror-gE_JCOgX.js"
+      "_vendor-recharts-6nUU7gU_.js",
+      "_vendor-codemirror-BEGorXSV.js"
     ]
   }
 }