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

phoenix/server/bearer_auth.py

@@ -0,0 +1,163 @@
+from abc import ABC
+from collections.abc import Awaitable, Callable
+from datetime import datetime, timedelta, timezone
+from functools import cached_property
+from typing import Any, Optional, cast
+
+import grpc
+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
+from starlette.requests import HTTPConnection
+from starlette.status import HTTP_401_UNAUTHORIZED
+
+from phoenix.auth import (
+    PHOENIX_ACCESS_TOKEN_COOKIE_NAME,
+    CanReadToken,
+    ClaimSetStatus,
+    Token,
+)
+from phoenix.db import enums
+from phoenix.db.enums import UserRole
+from phoenix.db.models import User as OrmUser
+from phoenix.server.types import (
+    AccessToken,
+    AccessTokenAttributes,
+    AccessTokenClaims,
+    ApiKeyClaims,
+    RefreshToken,
+    RefreshTokenAttributes,
+    RefreshTokenClaims,
+    TokenStore,
+    UserClaimSet,
+    UserId,
+)
+
+
+class HasTokenStore(ABC):
+    def __init__(self, token_store: CanReadToken) -> None:
+        super().__init__()
+        self._token_store = token_store
+
+
+class BearerTokenAuthBackend(HasTokenStore, AuthenticationBackend):
+    async def authenticate(
+        self,
+        conn: HTTPConnection,
+    ) -> Optional[tuple[AuthCredentials, BaseUser]]:
+        if header := conn.headers.get("Authorization"):
+            scheme, _, token = header.partition(" ")
+            if scheme.lower() != "bearer" or not token:
+                return None
+        elif access_token := conn.cookies.get(PHOENIX_ACCESS_TOKEN_COOKIE_NAME):
+            token = access_token
+        else:
+            return None
+        claims = await self._token_store.read(Token(token))
+        if not (isinstance(claims, UserClaimSet) and isinstance(claims.subject, UserId)):
+            return None
+        if not isinstance(claims, (ApiKeyClaims, AccessTokenClaims)):
+            return None
+        return AuthCredentials(), PhoenixUser(claims.subject, claims)
+
+
+class PhoenixUser(BaseUser):
+    def __init__(self, user_id: UserId, claims: UserClaimSet) -> None:
+        self._user_id = user_id
+        self.claims = claims
+        assert claims.attributes
+        self._is_admin = (
+            claims.status is ClaimSetStatus.VALID
+            and claims.attributes.user_role == enums.UserRole.ADMIN
+        )
+
+    @cached_property
+    def is_admin(self) -> bool:
+        return self._is_admin
+
+    @cached_property
+    def identity(self) -> UserId:
+        return self._user_id
+
+    @cached_property
+    def is_authenticated(self) -> bool:
+        return True
+
+
+class ApiKeyInterceptor(HasTokenStore, AsyncServerInterceptor):
+    async def intercept(
+        self,
+        method: Callable[[Any, grpc.ServicerContext], Awaitable[Any]],
+        request_or_iterator: Any,
+        context: grpc.ServicerContext,
+        method_name: str,
+    ) -> Any:
+        for datum in context.invocation_metadata():
+            if datum.key.lower() == "authorization":
+                scheme, _, token = datum.value.partition(" ")
+                if scheme.lower() != "bearer" or not token:
+                    break
+                claims = await self._token_store.read(Token(token))
+                if not (isinstance(claims, UserClaimSet) and isinstance(claims.subject, UserId)):
+                    break
+                if not isinstance(claims, (ApiKeyClaims, AccessTokenClaims)):
+                    raise Unauthenticated(details="Invalid token")
+                if claims.status is ClaimSetStatus.EXPIRED:
+                    raise Unauthenticated(details="Expired token")
+                if claims.status is ClaimSetStatus.VALID:
+                    return await method(request_or_iterator, context)
+                raise Unauthenticated()
+        raise Unauthenticated()
+
+
+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 or websocket connection is not authenticated.
+    """
+    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")
+    if claims.status is not ClaimSetStatus.VALID:
+        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+
+
+async def create_access_and_refresh_tokens(
+    *,
+    token_store: TokenStore,
+    user: OrmUser,
+    access_token_expiry: timedelta,
+    refresh_token_expiry: timedelta,
+) -> tuple[AccessToken, RefreshToken]:
+    issued_at = datetime.now(timezone.utc)
+    user_id = UserId(user.id)
+    user_role = UserRole(user.role.name)
+    refresh_token_claims = RefreshTokenClaims(
+        subject=user_id,
+        issued_at=issued_at,
+        expiration_time=issued_at + refresh_token_expiry,
+        attributes=RefreshTokenAttributes(
+            user_role=user_role,
+        ),
+    )
+    refresh_token, refresh_token_id = await token_store.create_refresh_token(refresh_token_claims)
+    access_token_claims = AccessTokenClaims(
+        subject=user_id,
+        issued_at=issued_at,
+        expiration_time=issued_at + access_token_expiry,
+        attributes=AccessTokenAttributes(
+            user_role=user_role,
+            refresh_token_id=refresh_token_id,
+        ),
+    )
+    access_token, _ = await token_store.create_access_token(access_token_claims)
+    return access_token, refresh_token

phoenix/server/dml_event.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import ClassVar
+
+from phoenix.db import models
+
+
+@dataclass(frozen=True)
+class DmlEvent(ABC):
+    """
+    Event corresponding to a Data Manipulation Language (DML)
+    operation, e.g. insertion, update, or deletion.
+    """
+
+    table: ClassVar[type[models.Base]]
+    ids: tuple[int, ...] = field(default_factory=tuple)
+
+    def __bool__(self) -> bool:
+        return bool(self.ids)
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+@dataclass(frozen=True)
+class ProjectDmlEvent(DmlEvent):
+    table = models.Project
+
+
+@dataclass(frozen=True)
+class ProjectDeleteEvent(ProjectDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanDmlEvent(ProjectDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanInsertEvent(SpanDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanDeleteEvent(SpanDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DatasetDmlEvent(DmlEvent):
+    table = models.Dataset
+
+
+@dataclass(frozen=True)
+class DatasetInsertEvent(DatasetDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DatasetDeleteEvent(DatasetDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentDmlEvent(DmlEvent):
+    table = models.Experiment
+
+
+@dataclass(frozen=True)
+class ExperimentInsertEvent(ExperimentDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentDeleteEvent(ExperimentDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunDmlEvent(DmlEvent):
+    table = models.ExperimentRun
+
+
+@dataclass(frozen=True)
+class ExperimentRunInsertEvent(ExperimentRunDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunDeleteEvent(ExperimentRunDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationDmlEvent(DmlEvent):
+    table = models.ExperimentRunAnnotation
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationInsertEvent(ExperimentRunAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationDeleteEvent(ExperimentRunAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanAnnotationDmlEvent(DmlEvent):
+    table = models.SpanAnnotation
+
+
+@dataclass(frozen=True)
+class SpanAnnotationInsertEvent(SpanAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanAnnotationDeleteEvent(SpanAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class TraceAnnotationDmlEvent(DmlEvent):
+    table = models.TraceAnnotation
+
+
+@dataclass(frozen=True)
+class TraceAnnotationInsertEvent(TraceAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class TraceAnnotationDeleteEvent(TraceAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationDmlEvent(DmlEvent):
+    table = models.DocumentAnnotation
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationInsertEvent(DocumentAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationDeleteEvent(DocumentAnnotationDmlEvent): ...

phoenix/server/dml_event_handler.py

@@ -0,0 +1,256 @@
+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, Generic, Optional, TypedDict, TypeVar, Union, cast
+
+from sqlalchemy import Select, select
+from typing_extensions import TypeAlias, Unpack
+
+from phoenix.db.models import (
+    Base,
+    DocumentAnnotation,
+    Project,
+    Span,
+    SpanAnnotation,
+    Trace,
+    TraceAnnotation,
+)
+from phoenix.server.api.dataloaders import CacheForDataLoaders
+from phoenix.server.dml_event import (
+    DmlEvent,
+    DocumentAnnotationDmlEvent,
+    SpanAnnotationDmlEvent,
+    SpanDeleteEvent,
+    SpanDmlEvent,
+    TraceAnnotationDmlEvent,
+)
+from phoenix.server.types import (
+    BatchedCaller,
+    CanSetLastUpdatedAt,
+    DbSessionFactory,
+)
+
+_DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)
+
+
+class _DmlEventQueue(Generic[_DmlEventT]):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._events: set[_DmlEventT] = set()
+
+    @property
+    def empty(self) -> bool:
+        return not self._events
+
+    def put(self, event: _DmlEventT) -> None:
+        self._events.add(event)
+
+    def clear(self) -> None:
+        self._events.clear()
+
+    def __iter__(self) -> Iterator[_DmlEventT]:
+        yield from self._events
+
+
+class _HandlerParams(TypedDict):
+    db: DbSessionFactory
+    last_updated_at: CanSetLastUpdatedAt
+    cache_for_dataloaders: Optional[CacheForDataLoaders]
+    sleep_seconds: float
+
+
+class _HasLastUpdatedAt(ABC):
+    def __init__(
+        self,
+        last_updated_at: CanSetLastUpdatedAt,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._last_updated_at = last_updated_at
+
+
+class _HasCacheForDataLoaders(ABC):
+    def __init__(
+        self,
+        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._cache_for_dataloaders = cache_for_dataloaders
+
+
+class _DmlEventHandler(
+    _HasLastUpdatedAt,
+    _HasCacheForDataLoaders,
+    BatchedCaller[_DmlEventT],
+    Generic[_DmlEventT],
+    ABC,
+):
+    _batch_factory = cast(Callable[[], _DmlEventQueue[_DmlEventT]], _DmlEventQueue)
+
+    def __init__(self, *, db: DbSessionFactory, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._db = db
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+class _GenericDmlEventHandler(_DmlEventHandler[DmlEvent]):
+    async def __call__(self) -> None:
+        for e in self._batch:
+            for id_ in e.ids:
+                self._update(e.table, id_)
+
+    def _update(self, table: type[Base], id_: int) -> None:
+        self._last_updated_at.set(table, id_)
+
+
+class _SpanDmlEventHandler(_DmlEventHandler[SpanDmlEvent]):
+    async def __call__(self) -> None:
+        if cache := self._cache_for_dataloaders:
+            for id_ in set(chain.from_iterable(e.ids for e in self._batch)):
+                self._clear(cache, id_)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int) -> None:
+        cache.latency_ms_quantile.invalidate(project_id)
+        cache.token_count.invalidate(project_id)
+        cache.record_count.invalidate(project_id)
+        cache.min_start_or_max_end_time.invalidate(project_id)
+
+
+class _SpanDeleteEventHandler(_SpanDmlEventHandler):
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int) -> None:
+        cache.annotation_summary.invalidate_project(project_id)
+        cache.document_evaluation_summary.invalidate_project(project_id)
+
+
+_AnnotationTable: TypeAlias = Union[
+    type[SpanAnnotation],
+    type[TraceAnnotation],
+    type[DocumentAnnotation],
+]
+
+_AnnotationDmlEventT = TypeVar(
+    "_AnnotationDmlEventT",
+    SpanAnnotationDmlEvent,
+    TraceAnnotationDmlEvent,
+    DocumentAnnotationDmlEvent,
+)
+
+
+class _AnnotationDmlEventHandler(
+    _DmlEventHandler[_AnnotationDmlEventT],
+    Generic[_AnnotationDmlEventT],
+    ABC,
+):
+    _table: _AnnotationTable
+    _base_stmt: Union[Select[tuple[int, str]], Select[tuple[int]]] = (
+        select(Project.id).join_from(Project, Trace).distinct()
+    )
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._base_stmt
+        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]]]:
+        ids = set(chain.from_iterable(e.ids for e in self._batch))
+        return self._stmt.where(self._table.id.in_(ids))
+
+    @staticmethod
+    @abstractmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None: ...
+
+    async def __call__(self) -> None:
+        async with self._db() as session:
+            async for row in await session.stream(self._get_stmt()):
+                self._last_updated_at.set(Project, row.id)
+                if cache := self._cache_for_dataloaders:
+                    self._clear(cache, row.id, row.name)
+
+
+class _SpanAnnotationDmlEventHandler(_AnnotationDmlEventHandler[SpanAnnotationDmlEvent]):
+    _table = SpanAnnotation
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._stmt.join_from(Trace, Span).join_from(Span, self._table)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None:
+        cache.annotation_summary.invalidate((project_id, name, "span"))
+
+
+class _TraceAnnotationDmlEventHandler(_AnnotationDmlEventHandler[TraceAnnotationDmlEvent]):
+    _table = TraceAnnotation
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._stmt.join_from(Trace, self._table)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None:
+        cache.annotation_summary.invalidate((project_id, name, "trace"))
+
+
+class _DocumentAnnotationDmlEventHandler(_AnnotationDmlEventHandler[DocumentAnnotationDmlEvent]):
+    _table = DocumentAnnotation
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._stmt.join_from(Trace, Span).join_from(Span, self._table)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None:
+        cache.document_evaluation_summary.invalidate((project_id, name))
+
+
+class DmlEventHandler:
+    def __init__(
+        self,
+        *,
+        db: DbSessionFactory,
+        last_updated_at: CanSetLastUpdatedAt,
+        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+        sleep_seconds: float = 0.1,
+    ) -> None:
+        kwargs = _HandlerParams(
+            db=db,
+            last_updated_at=last_updated_at,
+            cache_for_dataloaders=cache_for_dataloaders,
+            sleep_seconds=sleep_seconds,
+        )
+        self._handlers: Mapping[type[DmlEvent], Iterable[_DmlEventHandler[Any]]] = {
+            DmlEvent: [_GenericDmlEventHandler(**kwargs)],
+            SpanDmlEvent: [_SpanDmlEventHandler(**kwargs)],
+            SpanDeleteEvent: [_SpanDeleteEventHandler(**kwargs)],
+            SpanAnnotationDmlEvent: [_SpanAnnotationDmlEventHandler(**kwargs)],
+            TraceAnnotationDmlEvent: [_TraceAnnotationDmlEventHandler(**kwargs)],
+            DocumentAnnotationDmlEvent: [_DocumentAnnotationDmlEventHandler(**kwargs)],
+        }
+        self._all_handlers = frozenset(chain.from_iterable(self._handlers.values()))
+
+    async def __aenter__(self) -> None:
+        await gather(*(h.start() for h in self._all_handlers))
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        await gather(*(h.stop() for h in self._all_handlers))
+
+    def put(self, event: DmlEvent) -> None:
+        if not (isinstance(event, DmlEvent) and event):
+            return
+        for cls in getmro(type(event)):
+            if not (issubclass(cls, DmlEvent) and (handlers := self._handlers.get(cls))):
+                continue
+            for h in handlers:
+                h.put(event)
+            if cls is DmlEvent:
+                break

phoenix/server/email/sender.py

@@ -0,0 +1,97 @@
+import asyncio
+import smtplib
+import ssl
+from email.message import EmailMessage
+from pathlib import Path
+from typing import Literal
+
+from jinja2 import Environment, FileSystemLoader, select_autoescape
+
+EMAIL_TEMPLATE_FOLDER = Path(__file__).parent / "templates"
+
+
+class SimpleEmailSender:
+    def __init__(
+        self,
+        smtp_server: str,
+        smtp_port: int,
+        username: str,
+        password: str,
+        sender_email: str,
+        connection_method: Literal["STARTTLS", "SSL", "PLAIN"] = "STARTTLS",
+        validate_certs: bool = True,
+    ) -> None:
+        self.smtp_server = smtp_server
+        self.smtp_port = smtp_port
+        self.username = username
+        self.password = password
+        self.sender_email = sender_email
+        self.connection_method = connection_method.upper()
+        self.validate_certs = validate_certs
+
+        self.env = Environment(
+            loader=FileSystemLoader(EMAIL_TEMPLATE_FOLDER),
+            autoescape=select_autoescape(["html", "xml"]),
+        )
+
+    async def send_password_reset_email(
+        self,
+        email: str,
+        reset_url: str,
+    ) -> None:
+        subject = "[Phoenix] Password Reset Request"
+        template_name = "password_reset.html"
+
+        template = self.env.get_template(template_name)
+        html_content = template.render(reset_url=reset_url)
+
+        msg = EmailMessage()
+        msg["Subject"] = subject
+        msg["From"] = self.sender_email
+        msg["To"] = email
+        msg.set_content(html_content, subtype="html")
+
+        def send_email() -> None:
+            context: ssl.SSLContext
+            if self.validate_certs:
+                context = ssl.create_default_context()
+            else:
+                context = ssl._create_unverified_context()
+
+            methods_to_try = [self.connection_method]
+            # add secure method fallbacks
+            if self.connection_method != "PLAIN":
+                if self.connection_method != "STARTTLS":
+                    methods_to_try.append("STARTTLS")
+                elif self.connection_method != "SSL":
+                    methods_to_try.append("SSL")
+
+            for method in methods_to_try:
+                try:
+                    if method == "STARTTLS":
+                        server = smtplib.SMTP(self.smtp_server, self.smtp_port)
+                        server.ehlo()
+                        server.starttls(context=context)
+                        server.ehlo()
+                    elif method == "SSL":
+                        server = smtplib.SMTP_SSL(self.smtp_server, self.smtp_port, context=context)
+                        server.ehlo()
+                    elif method == "PLAIN":
+                        server = smtplib.SMTP(self.smtp_server, self.smtp_port)
+                        server.ehlo()
+                    else:
+                        continue  # Unsupported method
+
+                    if self.username and self.password:
+                        server.login(self.username, self.password)
+
+                    server.send_message(msg)
+                    server.quit()
+                    break  # Success
+                except Exception as e:
+                    print(f"Failed to send email using {method}: {e}")
+                    continue
+            else:
+                raise Exception("All connection methods failed")
+
+        await asyncio.to_thread(send_email)

phoenix/server/email/templates/password_reset.html

@@ -0,0 +1,19 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="UTF-8" />
+    <title>Password Reset</title>
+  </head>
+  <body>
+    <p>Hello.</p>
+    <p>
+      You have requested a password reset. Please click on the link below to
+      reset your password:
+    </p>
+    <p>
+      <a id="reset-url" href="{{ reset_url }}">Reset Password</a
+      >
+    </p>
+    <p>If you did not make this request, please contact your administrator.</p>
+  </body>
+</html>

phoenix/server/email/types.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class EmailSender(Protocol):
+    async def send_password_reset_email(
+        self,
+        email: str,
+        reset_url: str,
+    ) -> None: ...