arize-phoenix 3.16.1__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/main.py

@@ -1,54 +1,87 @@
 import atexit
-import logging
+import codecs
 import os
-from argparse import ArgumentParser
+import sys
+from argparse import SUPPRESS, ArgumentParser
 from pathlib import Path
-from random import random
 from threading import Thread
 from time import sleep, time
-from typing import Iterable, Optional, Protocol, TypeVar
+from typing import Optional
+from urllib.parse import urljoin
 
-import pkg_resources
+from jinja2 import BaseLoader, Environment
 from uvicorn import Config, Server
 
+import phoenix.trace.v1 as pb
 from phoenix.config import (
     EXPORT_DIR,
+    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_enable_websockets,
+    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,
 )
-from phoenix.core.model_schema_adapter import create_model_from_datasets
-from phoenix.core.traces import Traces
-from phoenix.datasets.dataset import EMPTY_DATASET, Dataset
-from phoenix.datasets.fixtures import FIXTURES, get_datasets
+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,
     DEFAULT_N_SAMPLES,
     UMAPParameters,
 )
-from phoenix.server.app import create_app
-from phoenix.storage.span_store import SpanStore
+from phoenix.server.app import (
+    ScaffolderConfig,
+    _db,
+    create_app,
+    create_engine_and_run_migrations,
+    instrument_engine_if_enabled,
+)
+from phoenix.server.email.sender import SimpleEmailSender
+from phoenix.server.types import DbSessionFactory
+from phoenix.settings import Settings
 from phoenix.trace.fixtures import (
     TRACES_FIXTURES,
-    _download_traces_fixture,
-    _get_trace_fixture_by_name,
+    get_dataset_fixtures,
     get_evals_from_fixture,
+    get_trace_fixtures_by_project_name,
+    load_example_traces,
+    reset_fixture_span_ids_and_timestamps,
+    send_dataset_fixtures,
 )
-from phoenix.trace.otel import decode, encode
-from phoenix.trace.span_json_decoder import json_string_to_span
-from phoenix.utilities.span_store import get_span_store, load_traces_data_from_store
-
-logger = logging.getLogger(__name__)
+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 = """
+_WELCOME_MESSAGE = Environment(loader=BaseLoader()).from_string("""
 
 ██████╗ ██╗  ██╗ ██████╗ ███████╗███╗   ██╗██╗██╗  ██╗
 ██╔══██╗██║  ██║██╔═══██╗██╔════╝████╗  ██║██║╚██╗██╔╝
 ██████╔╝███████║██║   ██║█████╗  ██╔██╗ ██║██║ ╚███╔╝
 ██╔═══╝ ██╔══██║██║   ██║██╔══╝  ██║╚██╗██║██║ ██╔██╗
 ██║     ██║  ██║╚██████╔╝███████╗██║ ╚████║██║██╔╝ ██╗
-╚═╝     ╚═╝  ╚═╝ ╚═════╝ ╚══════╝╚═╝  ╚═══╝╚═╝╚═╝  ╚═╝ v{0}
+╚═╝     ╚═╝  ╚═╝ ╚═════╝ ╚══════╝╚═╝  ╚═══╝╚═╝╚═╝  ╚═╝ v{{ version }}
 
 |
 |  🌎 Join our Community 🌎
@@ -61,15 +94,22 @@ _WELCOME_MESSAGE = """
 |  https://docs.arize.com/phoenix
 |
 |  🚀 Phoenix Server 🚀
-|  Phoenix UI: http://{1}:{2}
-|  Log traces: /v1/traces over HTTP
-|
-"""
+|  Phoenix UI: {{ ui_path }}
+|  Authentication: {{ auth_enabled }}
+|  Websockets: {{ websockets_enabled }}
+|  Log traces:
+|    - gRPC: {{ grpc_path }}
+|    - HTTP: {{ http_path }}
+|  Storage: {{ storage }}
+{% if schema -%}
+|    - Schema: {{ schema }}
+{% endif -%}
+""")
 
 
 def _write_pid_file_when_ready(
     server: Server,
-    wait_up_to_seconds: float = 5,
+    wait_up_to_seconds: float = 60,
 ) -> None:
     """Write PID file after server is started (or when time is up)."""
     time_limit = time() + wait_up_to_seconds
@@ -88,161 +128,310 @@ def _get_pid_file() -> Path:
     return get_pids_path() / str(os.getpid())
 
 
-_Item = TypeVar("_Item", contravariant=True)
-
-
-class _SupportsPut(Protocol[_Item]):
-    def put(self, item: _Item) -> None: ...
-
-
-def _load_items(
-    queue: _SupportsPut[_Item],
-    items: Iterable[_Item],
-    simulate_streaming: Optional[bool] = False,
-) -> None:
-    for item in items:
-        if simulate_streaming:
-            sleep(random())
-        queue.put(item)
+DEFAULT_UMAP_PARAMS_STR = f"{DEFAULT_MIN_DIST},{DEFAULT_N_NEIGHBORS},{DEFAULT_N_SAMPLES}"
 
 
-DEFAULT_UMAP_PARAMS_STR = f"{DEFAULT_MIN_DIST},{DEFAULT_N_NEIGHBORS},{DEFAULT_N_SAMPLES}"
+def main() -> None:
+    initialize_settings()
+    setup_logging()
 
-if __name__ == "__main__":
-    primary_dataset_name: str
-    reference_dataset_name: Optional[str]
+    primary_inferences_name: str
+    reference_inferences_name: Optional[str]
     trace_dataset_name: Optional[str] = None
-    simulate_streaming: Optional[bool] = None
 
-    primary_dataset: Dataset = EMPTY_DATASET
-    reference_dataset: Optional[Dataset] = None
-    corpus_dataset: Optional[Dataset] = None
+    primary_inferences: Inferences = EMPTY_INFERENCES
+    reference_inferences: Optional[Inferences] = None
+    corpus_inferences: Optional[Inferences] = None
 
-    # automatically remove the pid file when the process is being gracefully terminated
     atexit.register(_remove_pid_file)
 
-    parser = ArgumentParser()
-    parser.add_argument("--export_path")
-    parser.add_argument("--host", type=str, required=False)
-    parser.add_argument("--port", type=int, required=False)
-    parser.add_argument("--read-only", type=bool, default=False)
-    parser.add_argument("--no-internet", action="store_true")
-    parser.add_argument("--umap_params", type=str, required=False, default=DEFAULT_UMAP_PARAMS_STR)
-    parser.add_argument("--debug", action="store_false")
-    subparsers = parser.add_subparsers(dest="command", required=True)
+    parser = ArgumentParser(usage="phoenix serve", add_help=False)
+    parser.add_argument(
+        "-h",
+        "--help",
+        action="help",
+        help=SUPPRESS,
+    )
+    parser.add_argument("--database-url", required=False, help=SUPPRESS)
+    parser.add_argument("--export_path", help=SUPPRESS)
+    parser.add_argument("--host", type=str, required=False, help=SUPPRESS)
+    parser.add_argument("--port", type=int, required=False, help=SUPPRESS)
+    parser.add_argument("--read-only", action="store_true", required=False, help=SUPPRESS)
+    parser.add_argument("--no-internet", action="store_true", help=SUPPRESS)
+    parser.add_argument(
+        "--umap_params", type=str, required=False, default=DEFAULT_UMAP_PARAMS_STR, help=SUPPRESS
+    )
+    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")
+    serve_parser.add_argument(
+        "--with-fixture",
+        type=str,
+        required=False,
+        default="",
+        help=("Name of an inference fixture. Example: 'fixture1'"),
+    )
+    serve_parser.add_argument(
+        "--with-trace-fixtures",
+        type=str,
+        required=False,
+        default="",
+        help=(
+            "Comma separated list of tracing fixture names (spaces are ignored). "
+            "Example: 'fixture1, fixture2'"
+        ),
+    )
+    serve_parser.add_argument(
+        "--with-projects",
+        type=str,
+        required=False,
+        default="",
+        help=(
+            "Comma separated list of project names (spaces are ignored). "
+            "Example: 'project1, project2'"
+        ),
+    )
+    serve_parser.add_argument(
+        "--force-fixture-ingestion",
+        action="store_true",
+        required=False,
+        help=(
+            "Whether or not to check the database age before adding the fixtures. "
+            "Default is False, i.e., fixtures will only be added if the "
+            "database is new."
+        ),
+    )
+    serve_parser.add_argument(
+        "--scaffold-datasets",
+        action="store_true",
+        required=False,
+        help=(
+            "Whether or not to add any datasets defined in "
+            "the inputted project or trace fixture. "
+            "Default is False. "
+        ),
+    )
+
     datasets_parser = subparsers.add_parser("datasets")
     datasets_parser.add_argument("--primary", type=str, required=True)
     datasets_parser.add_argument("--reference", type=str, required=False)
     datasets_parser.add_argument("--corpus", type=str, required=False)
     datasets_parser.add_argument("--trace", type=str, required=False)
+
     fixture_parser = subparsers.add_parser("fixture")
     fixture_parser.add_argument("fixture", type=str, choices=[fixture.name for fixture in FIXTURES])
-    fixture_parser.add_argument("--primary-only", type=bool)
+    fixture_parser.add_argument("--primary-only", action="store_true")
+
     trace_fixture_parser = subparsers.add_parser("trace-fixture")
     trace_fixture_parser.add_argument(
         "fixture", type=str, choices=[fixture.name for fixture in TRACES_FIXTURES]
     )
-    trace_fixture_parser.add_argument("--simulate-streaming", type=bool)
+    trace_fixture_parser.add_argument("--simulate-streaming", action="store_true")
+
     demo_parser = subparsers.add_parser("demo")
     demo_parser.add_argument("fixture", type=str, choices=[fixture.name for fixture in FIXTURES])
     demo_parser.add_argument(
         "trace_fixture", type=str, choices=[fixture.name for fixture in TRACES_FIXTURES]
     )
     demo_parser.add_argument("--simulate-streaming", action="store_true")
+
     args = parser.parse_args()
-    export_path = Path(args.export_path) if args.export_path else EXPORT_DIR
-    span_store: Optional[SpanStore] = None
+    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 Path(EXPORT_DIR)
+
+    force_fixture_ingestion = False
+    scaffold_datasets = False
+    tracing_fixture_names = set()
     if args.command == "datasets":
-        primary_dataset_name = args.primary
-        reference_dataset_name = args.reference
-        corpus_dataset_name = args.corpus
-        primary_dataset = Dataset.from_name(primary_dataset_name)
-        reference_dataset = (
-            Dataset.from_name(reference_dataset_name)
-            if reference_dataset_name is not None
+        primary_inferences_name = args.primary
+        reference_inferences_name = args.reference
+        corpus_inferences_name = args.corpus
+        primary_inferences = Inferences.from_name(primary_inferences_name)
+        reference_inferences = (
+            Inferences.from_name(reference_inferences_name)
+            if reference_inferences_name is not None
             else None
         )
-        corpus_dataset = (
-            None if corpus_dataset_name is None else Dataset.from_name(corpus_dataset_name)
+        corpus_inferences = (
+            None if corpus_inferences_name is None else Inferences.from_name(corpus_inferences_name)
         )
     elif args.command == "fixture":
         fixture_name = args.fixture
         primary_only = args.primary_only
-        primary_dataset, reference_dataset, corpus_dataset = get_datasets(
+        primary_inferences, reference_inferences, corpus_inferences = get_inferences(
             fixture_name,
             args.no_internet,
         )
         if primary_only:
-            reference_dataset_name = None
-            reference_dataset = None
+            reference_inferences_name = None
+            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_dataset, reference_dataset, corpus_dataset = get_datasets(
+        primary_inferences, reference_inferences, corpus_inferences = get_inferences(
             fixture_name,
             args.no_internet,
         )
         trace_dataset_name = args.trace_fixture
-        simulate_streaming = args.simulate_streaming
+    elif args.command == "serve":
+        if args.with_fixture:
+            primary_inferences, reference_inferences, corpus_inferences = get_inferences(
+                str(args.with_fixture),
+                args.no_internet,
+            )
+        if args.with_trace_fixtures:
+            tracing_fixture_names.update(
+                [name.strip() for name in args.with_trace_fixtures.split(",")]
+            )
+        if args.with_projects:
+            project_names = [name.strip() for name in args.with_projects.split(",")]
+            tracing_fixture_names.update(
+                fixture.name
+                for name in project_names
+                for fixture in get_trace_fixtures_by_project_name(name)
+            )
+        force_fixture_ingestion = args.force_fixture_ingestion
+        scaffold_datasets = args.scaffold_datasets
+    host: Optional[str] = args.host or get_env_host()
+    if host == "::":
+        host = None
+
+    port = args.port or get_env_port()
+    host_root_path = get_env_host_root_path()
+    read_only = args.read_only
 
-    model = create_model_from_datasets(
-        primary_dataset,
-        reference_dataset,
+    model = create_model_from_inferences(
+        primary_inferences,
+        reference_inferences,
     )
-    traces = Traces()
-    if span_store := get_span_store():
-        Thread(target=load_traces_data_from_store, args=(traces, span_store), daemon=True).start()
+
+    authentication_enabled, secret = get_env_auth_settings()
+
+    fixture_spans: list[Span] = []
+    fixture_evals: list[pb.Evaluation] = []
     if trace_dataset_name is not None:
-        fixture_spans = list(
-            # Apply `encode` here because legacy jsonl files contains UUIDs as strings.
-            # `encode` removes the hyphens in the UUIDs.
-            decode(encode(json_string_to_span(json_span)))
-            for json_span in _download_traces_fixture(
-                _get_trace_fixture_by_name(trace_dataset_name)
-            )
+        fixture_spans, fixture_evals = reset_fixture_span_ids_and_timestamps(
+            (
+                # Apply `encode` here because legacy jsonl files contains UUIDs as strings.
+                # `encode` removes the hyphens in the UUIDs.
+                decode_otlp_span(encode_span_to_otlp(span))
+                for span in load_example_traces(trace_dataset_name).to_spans()
+            ),
+            get_evals_from_fixture(trace_dataset_name),
         )
-        Thread(
-            target=_load_items,
-            args=(traces, fixture_spans, simulate_streaming),
-            daemon=True,
-        ).start()
-        fixture_evals = list(get_evals_from_fixture(trace_dataset_name))
-        Thread(
-            target=_load_items,
-            args=(traces, fixture_evals, simulate_streaming),
-            daemon=True,
-        ).start()
+        dataset_fixtures = list(get_dataset_fixtures(trace_dataset_name))
+        if not read_only:
+            Thread(
+                target=send_dataset_fixtures,
+                args=(f"http://{host}:{port}", dataset_fixtures),
+            ).start()
     umap_params_list = args.umap_params.split(",")
     umap_params = UMAPParameters(
         min_dist=float(umap_params_list[0]),
         n_neighbors=int(umap_params_list[1]),
         n_samples=int(umap_params_list[2]),
     )
-    read_only = args.read_only
-    logger.info(f"Server umap params: {umap_params}")
+
+    if enable_prometheus := get_env_enable_prometheus():
+        from phoenix.server.prometheus import start_prometheus
+
+        start_prometheus()
+
+    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)
+    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(
+        version=phoenix_version,
+        ui_path=root_path,
+        grpc_path=f"http://{host}:{get_env_grpc_port()}",
+        http_path=urljoin(root_path, "v1/traces"),
+        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()
+    scaffolder_config = ScaffolderConfig(
+        db=factory,
+        tracing_fixture_names=tracing_fixture_names,
+        force_fixture_ingestion=force_fixture_ingestion,
+        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 (sender_email := get_env_smtp_mail_from()), "SMTP mail_from is required"
+        email_sender = SimpleEmailSender(
+            smtp_server=mail_sever,
+            smtp_port=get_env_smtp_port(),
+            username=mail_username,
+            password=mail_password,
+            sender_email=sender_email,
+            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,
-        traces=traces,
-        corpus=None if corpus_dataset is None else create_model_from_datasets(corpus_dataset),
+        corpus=corpus_model,
         debug=args.debug,
+        dev=args.dev,
+        serve_ui=not args.no_ui,
         read_only=read_only,
-        span_store=span_store,
+        enable_prometheus=enable_prometheus,
+        initial_spans=fixture_spans,
+        initial_evaluations=fixture_evals,
+        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(),
     )
-    host = args.host or get_env_host()
-    port = args.port or get_env_port()
-    server = Server(config=Config(app, host=host, port=port))
+    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()
 
-    # Print information about the server
-    phoenix_version = pkg_resources.get_distribution("arize-phoenix").version
-    print(
-        _WELCOME_MESSAGE.format(phoenix_version, host if host != "0.0.0.0" else "localhost", port)
-    )
-
     # 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__":
+    main()

phoenix/server/oauth2.py

@@ -0,0 +1,52 @@
+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
+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