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

phoenix/server/telemetry.py

@@ -0,0 +1,68 @@
+import logging
+import os
+from typing import TYPE_CHECKING
+
+from phoenix.config import (
+    ENV_PHOENIX_SERVER_INSTRUMENTATION_OTLP_TRACE_COLLECTOR_GRPC_ENDPOINT,
+    ENV_PHOENIX_SERVER_INSTRUMENTATION_OTLP_TRACE_COLLECTOR_HTTP_ENDPOINT,
+)
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import TracerProvider
+
+logger = logging.getLogger(__name__)
+
+
+def normalize_http_collector_endpoint(endpoint: str) -> str:
+    normalized_endpoint = endpoint
+    if not normalized_endpoint.startswith("http://") and not normalized_endpoint.startswith(
+        "https://"
+    ):
+        logger.warning(
+            "HTTP collector endpoint should include the protocol (http:// or https://)."
+            "Assuming http."
+        )
+        # assume http if no protocol is provided
+        normalized_endpoint = f"http://{endpoint}"
+    if normalized_endpoint.endswith("/v1/traces"):
+        logger.warning(
+            "HTTP collector endpoint should not include the /v1/traces path. Removing it."
+        )
+        # remove the /v1/traces path
+        normalized_endpoint = normalized_endpoint[: -len("/v1/traces")]
+    # remove trailing slashes
+    normalized_endpoint = normalized_endpoint.rstrip("/")
+    return normalized_endpoint
+
+
+def initialize_opentelemetry_tracer_provider() -> "TracerProvider":
+    logger.info("Initializing OpenTelemetry tracer provider")
+    from opentelemetry.sdk import trace as trace_sdk
+    from opentelemetry.sdk.resources import Resource
+    from opentelemetry.sdk.trace.export import BatchSpanProcessor
+    from opentelemetry.semconv.resource import ResourceAttributes
+
+    tracer_provider = trace_sdk.TracerProvider(
+        resource=Resource(attributes={ResourceAttributes.SERVICE_NAME: "arize-phoenix-server"})
+    )
+    if http_endpoint := os.getenv(
+        ENV_PHOENIX_SERVER_INSTRUMENTATION_OTLP_TRACE_COLLECTOR_HTTP_ENDPOINT
+    ):
+        logger.info(f"Using HTTP collector endpoint: {http_endpoint}")
+        http_endpoint = normalize_http_collector_endpoint(http_endpoint) + "/v1/traces"
+        from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+            OTLPSpanExporter as HttpExporter,
+        )
+
+        tracer_provider.add_span_processor(BatchSpanProcessor(HttpExporter(http_endpoint)))
+    if grpc_endpoint := os.getenv(
+        ENV_PHOENIX_SERVER_INSTRUMENTATION_OTLP_TRACE_COLLECTOR_GRPC_ENDPOINT
+    ):
+        logger.info(f"Using gRPC collector endpoint: {grpc_endpoint}")
+        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
+            OTLPSpanExporter as GrpcExporter,
+        )
+
+        tracer_provider.add_span_processor(BatchSpanProcessor(GrpcExporter(grpc_endpoint)))
+    logger.info("🔭 OpenTelemetry tracer provider initialized")
+    return tracer_provider

phoenix/server/templates/index.html

@@ -1,3 +1,47 @@
+{%- set rendered_chunks = [] -%}
+{%- set css_links = [] -%}
+{%- set js_scripts = [] -%}
+{%- macro collect_assets(chunk_name, level=0) -%}
+    {% if chunk_name not in rendered_chunks %}
+        {% set _ = rendered_chunks.append(chunk_name) %}
+        {% set chunk = manifest[chunk_name] %}
+        {% if chunk.css %}
+            {% for css_file in chunk.css %}
+                {% if css_file not in css_links %}
+                    {% set _ = css_links.append((basename, css_file)) %}
+                {% endif %}
+            {% endfor %}
+        {% endif %}
+        {% if chunk.imports %}
+            {% for import in chunk.imports %}
+                {% set _ = collect_assets(import, level + 1) %}
+            {% endfor %}
+        {% endif %}
+        {% if chunk.file.endswith('.js') %}
+            {% if chunk.file not in js_scripts %}
+                {% set _ = js_scripts.append((basename, chunk.file, level == 0)) %}
+            {% endif %}
+        {% elif chunk.file.endswith('.css') %}
+            {% if chunk.file not in css_links %}
+                {% set _ = css_links.append((basename, chunk.file)) %}
+            {% endif %}
+        {% endif %}
+    {% endif %}
+{%- endmacro -%}
+{%- macro render_css() -%}
+  {%- for basename, css_file in css_links -%}
+    <link rel="stylesheet" href="{{ basename }}/{{ css_file }}">
+  {% endfor -%}
+{%- endmacro -%}
+{%- macro render_js() -%}
+  {%- for basename, js_file, is_entry in js_scripts -%}
+    {%- if is_entry -%}
+      <script type="module" src="{{ basename }}/{{ js_file }}"></script>
+    {% else -%}
+      <link rel="modulepreload" href="{{ basename }}/{{ js_file }}">
+    {% endif -%}
+  {%- endfor -%}
+{%- endmacro -%}
 <!DOCTYPE html>
 <html>
   <head>
@@ -11,37 +55,52 @@
     <meta property="og:description" content="AI Observability & Evaluation" />
     <meta property="og:image" content="https://raw.githubusercontent.com/Arize-ai/phoenix-assets/main/images/socal/social-preview-horizontal.jpg" />
     <meta name="theme-color" content="#ffffff" />
-    <link rel="stylesheet" src="{{basename}}/index.css"></link>
     <link
       rel="stylesheet"
       href="https://fonts.googleapis.com/css2?family=Roboto:wght@100;300;400;500;700;900&display=swap"
     />
+    {% if not is_development -%}
+      {% set _ = collect_assets('index.tsx') -%}
+      {{- render_css() -}}
+    {%- endif -%}
     <script src="{{basename}}/modernizr.js"></script>
   </head>
   <body>
     <div id="root"></div>
     <script>(function() {
-        // Request Info:
-        // URL: {{request.url}}
-        // Base URL: {{request.base_url}}
-        // Root Path: {{root_path}}
-        Object.defineProperty(window, "Config", {
-            // Place any server-side injected config here
-            // E.g. default UMAP parameters etc. that needs to be
-            // injected into the client before React runs
-            value: Object.freeze({
-                basename: "{{basename}}",
-                hasInferences: Boolean("{{has_inferences}}" == "True"),
-                hasCorpus: Boolean("{{has_corpus}}" == "True"),
-                UMAP: {
-                  minDist: parseFloat("{{min_dist}}"),
-                  nNeighbors: parseInt("{{n_neighbors}}"),
-                  nSamples: parseInt("{{n_samples}}"),
-                }
-            }),
-            writable: false
-        });
-    })()</script>
-    <script type="module" src="{{basename}}/index.js"></script>
+      Object.defineProperty(window, "Config", {
+          // Place any server-side injected config here
+          // E.g. default UMAP parameters etc. that needs to be
+          // injected into the client before React runs
+          value: Object.freeze({
+              basename: "{{basename}}",
+              platformVersion: "{{platform_version}}",
+              hasInferences: Boolean("{{has_inferences}}" == "True"),
+              hasCorpus: Boolean("{{has_corpus}}" == "True"),
+              UMAP: {
+                minDist: parseFloat("{{min_dist}}"),
+                nNeighbors: parseInt("{{n_neighbors}}"),
+                nSamples: parseInt("{{n_samples}}"),
+              },
+              authenticationEnabled: Boolean("{{authentication_enabled}}" == "True"),
+              oAuth2Idps: {{ oauth2_idps | tojson }},
+              websocketsEnabled: Boolean("{{websockets_enabled}}" == "True"),
+          }),
+          writable: false
+      });
+  })()</script>
+    {% if is_development -%}
+      <script type="module">
+        import RefreshRuntime from 'http://localhost:5173/@react-refresh'
+        RefreshRuntime.injectIntoGlobalHook(window)
+        window.$RefreshReg$ = () => {}
+        window.$RefreshSig$ = () => (type) => type
+        window.__vite_plugin_react_preamble_installed__ = true
+      </script>
+      <script type="module" src="http://localhost:5173/@vite/client"></script>
+      <script type="module" src="http://localhost:5173/index.tsx"></script>
+    {%- else -%}
+      {{- render_js() -}}
+    {%- endif -%}
   </body>
 </html>

phoenix/server/thread_server.py

@@ -1,10 +1,10 @@
 import asyncio
 import logging
+from collections.abc import Generator
 from threading import Thread
 from time import sleep, time
-from typing import Generator
 
-from starlette.applications import Starlette
+from fastapi import FastAPI
 from uvicorn import Config, Server
 from uvicorn.config import LoopSetupType
 
@@ -24,7 +24,7 @@ class ThreadServer(Server):
 
     def __init__(
         self,
-        app: Starlette,
+        app: FastAPI,
         host: str,
         port: int,
         root_path: str,

phoenix/server/types.py

@@ -0,0 +1,275 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from asyncio import Task, create_task, sleep
+from collections import defaultdict
+from collections.abc import Callable, Iterator
+from contextlib import AbstractAsyncContextManager
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import Any, Generic, Optional, Protocol, TypeVar, final
+
+from cachetools import LRUCache
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from phoenix.auth import CanReadToken, ClaimSet, Token, TokenAttributes
+from phoenix.db import enums, models
+from phoenix.db.helpers import SupportedSQLDialect
+
+
+class CanSetLastUpdatedAt(Protocol):
+    def set(self, table: type[models.Base], id_: int) -> None: ...
+
+
+class CanGetLastUpdatedAt(Protocol):
+    def get(self, table: type[models.Base], id_: Optional[int] = None) -> Optional[datetime]: ...
+
+
+class DbSessionFactory:
+    def __init__(
+        self,
+        db: Callable[[], AbstractAsyncContextManager[AsyncSession]],
+        dialect: str,
+    ):
+        self._db = db
+        self.dialect = SupportedSQLDialect(dialect)
+
+    def __call__(self) -> AbstractAsyncContextManager[AsyncSession]:
+        return self._db()
+
+
+_AnyT = TypeVar("_AnyT")
+_ItemT_contra = TypeVar("_ItemT_contra", contravariant=True)
+
+
+class CanPutItem(Protocol[_ItemT_contra]):
+    def put(self, item: _ItemT_contra) -> None: ...
+
+
+class _Batch(CanPutItem[_AnyT], Protocol[_AnyT]):
+    @property
+    def empty(self) -> bool: ...
+    def clear(self) -> None: ...
+    def __iter__(self) -> Iterator[_AnyT]: ...
+
+
+class _HasBatch(Generic[_ItemT_contra], ABC):
+    _batch_factory: Callable[[], _Batch[_ItemT_contra]]
+
+    def __init__(self) -> None:
+        self._batch = self._batch_factory()
+
+    def put(self, item: _ItemT_contra) -> None:
+        self._batch.put(item)
+
+
+class DaemonTask(ABC):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._running = False
+        self._tasks: list[Task[None]] = []
+
+    async def start(self) -> None:
+        self._running = True
+        if not self._tasks:
+            self._tasks.append(create_task(self._run()))
+
+    async def stop(self) -> None:
+        self._running = False
+        for task in reversed(self._tasks):
+            if not task.done():
+                task.cancel()
+        self._tasks.clear()
+
+    async def __aenter__(self) -> None:
+        await self.start()
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        await self.stop()
+
+    @abstractmethod
+    async def _run(self) -> None: ...
+
+
+class BatchedCaller(DaemonTask, _HasBatch[_AnyT], Generic[_AnyT], ABC):
+    def __init__(self, *, sleep_seconds: float = 0.1, **kwargs: Any) -> None:
+        assert sleep_seconds > 0
+        super().__init__(**kwargs)
+        self._seconds = sleep_seconds
+
+    @abstractmethod
+    async def __call__(self) -> None: ...
+
+    async def _run(self) -> None:
+        while self._running:
+            self._tasks.append(create_task(sleep(self._seconds)))
+            await self._tasks[-1]
+            self._tasks.pop()
+            if self._batch.empty:
+                continue
+            self._tasks.append(create_task(self()))
+            await self._tasks[-1]
+            self._tasks.pop()
+            self._batch.clear()
+
+
+class LastUpdatedAt:
+    def __init__(self) -> None:
+        self._cache: defaultdict[
+            type[models.Base],
+            LRUCache[int, datetime],
+        ] = defaultdict(lambda: LRUCache(maxsize=100))
+
+    def get(self, table: type[models.Base], id_: Optional[int] = None) -> Optional[datetime]:
+        if not (cache := self._cache.get(table)):
+            return None
+        if id_ is None:
+            return max(filter(bool, cache.values()), default=None)
+        return cache.get(id_)
+
+    def set(self, table: type[models.Base], id_: int) -> None:
+        self._cache[table][id_] = datetime.now(timezone.utc)
+
+
+class PasswordResetToken(Token): ...
+
+
+class AccessToken(Token): ...
+
+
+class RefreshToken(Token): ...
+
+
+class ApiKey(Token): ...
+
+
+@dataclass(frozen=True)
+class UserTokenAttributes(TokenAttributes):
+    user_role: enums.UserRole
+
+
+@dataclass(frozen=True)
+class RefreshTokenAttributes(UserTokenAttributes): ...
+
+
+@dataclass(frozen=True)
+class PasswordResetTokenAttributes(UserTokenAttributes): ...
+
+
+@dataclass(frozen=True)
+class AccessTokenAttributes(UserTokenAttributes):
+    refresh_token_id: RefreshTokenId
+
+
+@dataclass(frozen=True)
+class ApiKeyAttributes(UserTokenAttributes):
+    name: str
+    description: Optional[str] = None
+
+
+class _DbId(str, ABC):
+    table: type[models.Base]
+
+    def __new__(cls, id_: int) -> _DbId:
+        assert isinstance(id_, int)
+        return super().__new__(cls, f"{cls.table.__name__}:{id_}")
+
+    def __int__(self) -> int:
+        return int(self.split(":")[1])
+
+    def __deepcopy__(self, memo: Any) -> _DbId:
+        return self
+
+
+class TokenId(_DbId, ABC):
+    @classmethod
+    def parse(cls, value: str) -> Optional[TokenId]:
+        table_name, _, id_ = value.partition(":")
+        if not id_.isnumeric():
+            return None
+        for sub in cls.__subclasses__():
+            if sub.table.__name__ == table_name:
+                return sub(int(id_))
+        return None
+
+
+@final
+class PasswordResetTokenId(TokenId):
+    table = models.PasswordResetToken
+
+
+@final
+class AccessTokenId(TokenId):
+    table = models.AccessToken
+
+
+@final
+class RefreshTokenId(TokenId):
+    table = models.RefreshToken
+
+
+@final
+class ApiKeyId(TokenId):
+    table = models.ApiKey
+
+
+@final
+class UserId(_DbId):
+    table = models.User
+
+
+@dataclass(frozen=True)
+class UserClaimSet(ClaimSet):  # type: ignore[override,unused-ignore]
+    subject: Optional[UserId] = None
+    attributes: Optional[UserTokenAttributes] = None
+
+
+@dataclass(frozen=True)
+class PasswordResetTokenClaims(UserClaimSet):  # type: ignore[override,unused-ignore]
+    token_id: Optional[PasswordResetTokenId] = None
+    attributes: Optional[PasswordResetTokenAttributes] = None
+
+
+@dataclass(frozen=True)
+class AccessTokenClaims(UserClaimSet):  # type: ignore[override,unused-ignore]
+    token_id: Optional[AccessTokenId] = None
+    attributes: Optional[AccessTokenAttributes] = None
+
+
+@dataclass(frozen=True)
+class RefreshTokenClaims(UserClaimSet):  # type: ignore[override,unused-ignore]
+    token_id: Optional[RefreshTokenId] = None
+    attributes: Optional[RefreshTokenAttributes] = None
+
+
+@dataclass(frozen=True)
+class ApiKeyClaims(UserClaimSet):  # type: ignore[override,unused-ignore]
+    token_id: Optional[ApiKeyId] = None
+    attributes: Optional[ApiKeyAttributes] = None
+
+
+class CanRevokeTokens(Protocol):
+    async def revoke(self, *token_ids: TokenId) -> None: ...
+
+
+class CanLogOutUser(Protocol):
+    async def log_out(self, user_id: UserId) -> None: ...
+
+
+class TokenStore(CanReadToken, CanRevokeTokens, CanLogOutUser, Protocol):
+    async def create_password_reset_token(
+        self,
+        claims: PasswordResetTokenClaims,
+    ) -> tuple[PasswordResetToken, PasswordResetTokenId]: ...
+    async def create_access_token(
+        self,
+        claims: AccessTokenClaims,
+    ) -> tuple[AccessToken, AccessTokenId]: ...
+    async def create_refresh_token(
+        self,
+        claims: RefreshTokenClaims,
+    ) -> tuple[RefreshToken, RefreshTokenId]: ...
+    async def create_api_key(
+        self,
+        claims: ApiKeyClaims,
+    ) -> tuple[ApiKey, ApiKeyId]: ...

phoenix/services.py

@@ -3,9 +3,10 @@ import os
 import signal
 import subprocess
 import sys
+from collections.abc import Callable
 from pathlib import Path
 from time import sleep, time
-from typing import Callable, List, Optional
+from typing import Optional
 
 import psutil
 
@@ -31,7 +32,7 @@ class Service:
         )
 
     @property
-    def command(self) -> List[str]:
+    def command(self) -> list[str]:
         raise NotImplementedError(f"{type(self)} must define `command`")
 
     def start(self) -> psutil.Popen:
@@ -99,40 +100,46 @@ class AppService(Service):
 
     working_dir = SERVER_DIR
 
-    # Internal references to the name / directory of the dataset(s)
-    __primary_dataset_name: str
-    __reference_dataset_name: Optional[str]
-    __corpus_dataset_name: Optional[str]
+    # Internal references to the name / directory of the inferences(s)
+    __primary_inferences_name: str
+    __reference_inferences_name: Optional[str]
+    __corpus_inferences_name: Optional[str]
     __trace_dataset_name: Optional[str]
 
     def __init__(
         self,
+        database_url: str,
         export_path: Path,
         host: str,
         port: int,
         root_path: str,
-        primary_dataset_name: str,
+        primary_inferences_name: str,
         umap_params: str,
-        reference_dataset_name: Optional[str],
-        corpus_dataset_name: Optional[str],
+        reference_inferences_name: Optional[str],
+        corpus_inferences_name: Optional[str],
         trace_dataset_name: Optional[str],
+        enable_websockets: bool,
     ):
+        self.database_url = database_url
         self.export_path = export_path
         self.host = host
         self.port = port
         self.root_path = root_path  # TODO(mikeldking): Add support for root_path
-        self.__primary_dataset_name = primary_dataset_name
+        self.__primary_inferences_name = primary_inferences_name
         self.__umap_params = umap_params
-        self.__reference_dataset_name = reference_dataset_name
-        self.__corpus_dataset_name = corpus_dataset_name
+        self.__reference_inferences_name = reference_inferences_name
+        self.__corpus_inferences_name = corpus_inferences_name
         self.__trace_dataset_name = trace_dataset_name
+        self.enable_websockets = enable_websockets
         super().__init__()
 
     @property
-    def command(self) -> List[str]:
+    def command(self) -> list[str]:
         command = [
             sys.executable,
             "main.py",
+            "--database-url",
+            self.database_url,
             "--export_path",
             str(self.export_path),
             "--host",
@@ -143,13 +150,15 @@ class AppService(Service):
             self.__umap_params,
             "datasets",
             "--primary",
-            str(self.__primary_dataset_name),
+            str(self.__primary_inferences_name),
         ]
-        if self.__reference_dataset_name is not None:
-            command.extend(["--reference", str(self.__reference_dataset_name)])
-        if self.__corpus_dataset_name is not None:
-            command.extend(["--corpus", str(self.__corpus_dataset_name)])
+        if self.__reference_inferences_name is not None:
+            command.extend(["--reference", str(self.__reference_inferences_name)])
+        if self.__corpus_inferences_name is not None:
+            command.extend(["--corpus", str(self.__corpus_inferences_name)])
         if self.__trace_dataset_name is not None:
             command.extend(["--trace", str(self.__trace_dataset_name)])
+        if self.enable_websockets:
+            command.append("--enable-websockets")
         logger.info(f"command: {' '.join(command)}")
         return command