arize-phoenix 10.0.4__py3-none-any.whl → 12.28.1__py3-none-any.whl

phoenix/server/daemons/generative_model_store.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import logging
+from asyncio import sleep
+from datetime import datetime, timedelta, timezone
+from typing import Any, Mapping, Optional
+
+import sqlalchemy as sa
+from sqlalchemy.orm import joinedload
+
+from phoenix.db import models
+from phoenix.server.cost_tracking.cost_model_lookup import CostModelLookup
+from phoenix.server.types import DaemonTask, DbSessionFactory
+
+logger = logging.getLogger(__name__)
+
+
+class GenerativeModelStore(DaemonTask):
+    """A daemon that periodically fetches generative models and maintains an in-memory cache.
+
+    This daemon periodically fetches generative models and their token prices from the
+    database and maintains an in-memory cache for fast lookups. It uses an incremental
+    fetch strategy to minimize database egress costs. Instead of fetching all models on
+    every refresh, we track the last fetch time and only query for models that have
+    changed since then (using updated_at/deleted_at).
+
+    Rationale: Database egress is expensive in cloud environments (especially managed
+    databases), and generative models change infrequently (mostly static reference data).
+    The cost calculation daemon queries this store frequently (once per span), so trading
+    memory for reduced database egress provides significant cost savings.
+
+    Note:
+        This strategy relies on GenerativeModel.updated_at being properly maintained. Any
+        code that modifies GenerativeModel or TokenPrice records MUST ensure updated_at
+        is explicitly set (see model_mutations.py). Relying solely on SQLAlchemy's
+        onupdate=func.now() is insufficient because SQLAlchemy may skip the UPDATE if it
+        detects no "real" changes to scalar fields (even if child records like TokenPrice
+        are modified).
+    """
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        refresh_interval_seconds: int = 5,
+    ) -> None:
+        super().__init__()
+        self._db = db
+        self._lookup = CostModelLookup()
+        self._last_fetch_time: Optional[datetime] = None
+        self._last_fetch_id: Optional[int] = None
+        self._refresh_interval_seconds = refresh_interval_seconds
+
+    def find_model(
+        self,
+        start_time: datetime,
+        attributes: Mapping[str, Any],
+    ) -> Optional[models.GenerativeModel]:
+        return self._lookup.find_model(start_time, attributes)
+
+    async def _run(self) -> None:
+        while self._running:
+            # Capture time before query with 2-second buffer for clock skew tolerance
+            fetch_start_time = datetime.now(timezone.utc) - timedelta(seconds=2)
+            try:
+                await self._fetch_models()
+            except Exception:
+                logger.exception("Failed to refresh generative models")
+            else:
+                self._last_fetch_time = fetch_start_time
+            await sleep(self._refresh_interval_seconds)
+
+    async def _fetch_models(self) -> None:
+        """
+        Fetch generative models from the database using an incremental strategy.
+
+        On the first run, fetches all models. On subsequent runs, only fetches models
+        where updated_at or deleted_at is at or after the last fetch time (with a 2-second
+        buffer). Some models may be refetched, but .merge() handles duplicates idempotently.
+        """
+        stmt = sa.select(models.GenerativeModel).options(
+            joinedload(models.GenerativeModel.token_prices)
+        )
+        if self._last_fetch_time:
+            # Incremental fetch: get models changed since last fetch.
+            # Use >= for updated_at/deleted_at to catch models from the buffer window.
+            # Include id check as redundant safety check.
+            stmt = stmt.where(
+                sa.or_(
+                    models.GenerativeModel.id > self._last_fetch_id,
+                    models.GenerativeModel.updated_at >= self._last_fetch_time,
+                    models.GenerativeModel.deleted_at >= self._last_fetch_time,
+                )
+            )
+        async with self._db() as session:
+            generative_models = (await session.scalars(stmt)).unique().all()
+
+        if not generative_models:
+            return
+
+        self._lookup.merge(generative_models)
+
+        # Track max id for redundant safety check.
+        self._last_fetch_id = max(model.id for model in generative_models)

phoenix/server/daemons/span_cost_calculator.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import logging
+from asyncio import sleep
+from collections import deque
+from datetime import datetime
+from typing import Any, Mapping, NamedTuple, Optional
+
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.cost_tracking.cost_details_calculator import SpanCostDetailsCalculator
+from phoenix.server.daemons.generative_model_store import GenerativeModelStore
+from phoenix.server.types import DaemonTask, DbSessionFactory
+
+logger = logging.getLogger(__name__)
+
+_GenerativeModelId: TypeAlias = int
+
+
+class SpanCostCalculatorQueueItem(NamedTuple):
+    span_rowid: int
+    trace_rowid: int
+    attributes: Mapping[str, Any]
+    span_start_time: datetime
+
+
+class SpanCostCalculator(DaemonTask):
+    _SLEEP_INTERVAL = 5  # seconds
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        model_store: GenerativeModelStore,
+    ) -> None:
+        super().__init__()
+        self._db = db
+        self._model_store = model_store
+        self._queue: deque[SpanCostCalculatorQueueItem] = deque()
+        self._max_items_per_transaction = 1000
+
+    async def _run(self) -> None:
+        while self._running:
+            num_items_to_insert = min(self._max_items_per_transaction, len(self._queue))
+            try:
+                await self._insert_costs(num_items_to_insert)
+            except Exception as e:
+                logger.exception(f"Failed to insert costs: {e}")
+            await sleep(self._SLEEP_INTERVAL)
+
+    async def _insert_costs(self, num_items_to_insert: int) -> None:
+        if not num_items_to_insert or not self._queue:
+            return
+        costs: list[models.SpanCost] = []
+        while num_items_to_insert > 0:
+            num_items_to_insert -= 1
+            item = self._queue.popleft()
+            try:
+                cost = self.calculate_cost(item.span_start_time, item.attributes)
+            except Exception as e:
+                logger.exception(f"Failed to calculate cost for span {item.span_rowid}: {e}")
+                continue
+            if not cost:
+                continue
+            cost.span_rowid = item.span_rowid
+            cost.trace_rowid = item.trace_rowid
+            costs.append(cost)
+        try:
+            async with self._db() as session:
+                session.add_all(costs)
+        except Exception as e:
+            logger.exception(f"Failed to insert costs: {e}")
+
+    def put_nowait(self, item: SpanCostCalculatorQueueItem) -> None:
+        self._queue.append(item)
+
+    def calculate_cost(
+        self,
+        start_time: datetime,
+        attributes: Mapping[str, Any],
+    ) -> Optional[models.SpanCost]:
+        if not attributes:
+            return None
+        cost_model = self._model_store.find_model(
+            start_time=start_time,
+            attributes=attributes,
+        )
+        calculator = SpanCostDetailsCalculator(cost_model.token_prices if cost_model else [])
+        details = calculator.calculate_details(attributes)
+        if not details:
+            return None
+
+        cost = models.SpanCost(
+            model_id=cost_model.id if cost_model else None,
+            span_start_time=start_time,
+        )
+        for detail in details:
+            cost.append_detail(detail)
+        return cost

phoenix/server/dml_event.py

@@ -33,6 +33,10 @@ class ProjectDmlEvent(DmlEvent):
 class ProjectDeleteEvent(ProjectDmlEvent): ...
 
 
+@dataclass(frozen=True)
+class ProjectInsertEvent(ProjectDmlEvent): ...
+
+
 @dataclass(frozen=True)
 class SpanDmlEvent(ProjectDmlEvent): ...
 
@@ -123,6 +127,19 @@ class TraceAnnotationInsertEvent(TraceAnnotationDmlEvent): ...
 class TraceAnnotationDeleteEvent(TraceAnnotationDmlEvent): ...
 
 
+@dataclass(frozen=True)
+class ProjectSessionAnnotationDmlEvent(DmlEvent):
+    table = models.ProjectSessionAnnotation
+
+
+@dataclass(frozen=True)
+class ProjectSessionAnnotationInsertEvent(ProjectSessionAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ProjectSessionAnnotationDeleteEvent(ProjectSessionAnnotationDmlEvent): ...
+
+
 @dataclass(frozen=True)
 class DocumentAnnotationDmlEvent(DmlEvent):
     table = models.DocumentAnnotation

phoenix/server/dml_event_handler.py

@@ -120,6 +120,7 @@ class _SpanDmlEventHandler(_DmlEventHandler[SpanDmlEvent]):
     def _clear(cache: CacheForDataLoaders, project_id: int) -> None:
         cache.latency_ms_quantile.invalidate(project_id)
         cache.token_count.invalidate(project_id)
+        cache.token_cost.invalidate(project_id)
         cache.record_count.invalidate(project_id)
         cache.min_start_or_max_end_time.invalidate(project_id)
 
@@ -127,6 +128,10 @@ class _SpanDmlEventHandler(_DmlEventHandler[SpanDmlEvent]):
 class _SpanDeleteEventHandler(_SpanDmlEventHandler):
     @staticmethod
     def _clear(cache: CacheForDataLoaders, project_id: int) -> None:
+        # Call parent's cache invalidation first (core span caches)
+        _SpanDmlEventHandler._clear(cache, project_id)
+
+        # Then invalidate annotation-specific caches
         cache.annotation_summary.invalidate_project(project_id)
         cache.document_evaluation_summary.invalidate_project(project_id)
 

phoenix/server/email/sender.py

@@ -1,15 +1,18 @@
+import logging
 import smtplib
 import ssl
 from email.message import EmailMessage
 from pathlib import Path
 from typing import Literal
-from urllib.parse import urljoin
 
 from anyio import to_thread
+from email_validator import EmailNotValidError, validate_email
 from jinja2 import Environment, FileSystemLoader, select_autoescape
 from typing_extensions import TypeAlias
 
-from phoenix.config import get_env_root_url
+from phoenix.config import get_env_root_url, get_env_support_email
+
+logger = logging.getLogger(__name__)
 
 EMAIL_TEMPLATE_FOLDER = Path(__file__).parent / "templates"
 
@@ -45,13 +48,20 @@ class SimpleEmailSender:
         email: str,
         name: str,
     ) -> None:
+        try:
+            email = validate_email(email, check_deliverability=False).normalized
+        except EmailNotValidError:
+            logger.warning("Skipping welcome email for user with invalid email address")
+            return
+
         subject = "[Phoenix] Welcome to Arize Phoenix"
         template_name = "welcome.html"
 
         template = self.env.get_template(template_name)
+
         html_content = template.render(
             name=name,
-            welcome_url=urljoin(str(get_env_root_url()), "forgot-password"),
+            welcome_url=str(get_env_root_url()),
         )
 
         msg = EmailMessage()
@@ -67,6 +77,12 @@ class SimpleEmailSender:
         email: str,
         reset_url: str,
     ) -> None:
+        try:
+            email = validate_email(email, check_deliverability=False).normalized
+        except EmailNotValidError:
+            logger.warning("Skipping password reset email for user with invalid email address")
+            return
+
         subject = "[Phoenix] Password Reset Request"
         template_name = "password_reset.html"
 
@@ -81,6 +97,43 @@ class SimpleEmailSender:
 
         await to_thread.run_sync(self._send_email, msg)
 
+    async def send_db_usage_warning_email(
+        self,
+        email: str,
+        current_usage_gibibytes: float,
+        allocated_storage_gibibytes: float,
+        notification_threshold_percentage: float,
+    ) -> None:
+        try:
+            email = validate_email(email, check_deliverability=False).normalized
+        except EmailNotValidError:
+            logger.warning(
+                "Skipping database usage warning email for user with invalid email address"
+            )
+            return
+
+        subject = "[Phoenix] Database Disk Space Usage Threshold Exceeded"
+        template_name = "db_disk_usage_notification.html"
+
+        support_email = get_env_support_email()
+        template = self.env.get_template(template_name)
+        html_content = template.render(
+            current_usage_gibibytes=current_usage_gibibytes,
+            allocated_storage_gibibytes=allocated_storage_gibibytes,
+            notification_threshold_percentage=notification_threshold_percentage,
+            support_email=support_email,
+        )
+
+        msg = EmailMessage()
+        msg["Subject"] = subject
+        msg["From"] = self.sender_email
+        msg["To"] = email
+        if support_email:
+            msg["Cc"] = support_email
+        msg.set_content(html_content, subtype="html")
+
+        await to_thread.run_sync(self._send_email, msg)
+
     def _send_email(self, msg: EmailMessage) -> None:
         context: ssl.SSLContext
         if self.validate_certs:

phoenix/server/email/templates/db_disk_usage_notification.html

@@ -0,0 +1,19 @@
+<!DOCTYPE html>
+<html>
+<head>
+  <meta charset="UTF-8" />
+  <title>Database Usage Notification</title>
+</head>
+<body>
+<h1>Database Usage Notification</h1>
+<p>Your Phoenix database usage has exceeded the notification threshold.</p>
+<p><strong>Current Usage:</strong> {{ current_usage_gibibytes|round(1) }} GiB</p>
+<p><strong>Allocated Storage:</strong> {{ allocated_storage_gibibytes|round(1) }} GiB</p>
+<p><strong>Usage Percentage:</strong> {{ ((current_usage_gibibytes / allocated_storage_gibibytes) * 100)|round(1) }}%</p>
+<p><strong>Notification Threshold:</strong> {{ notification_threshold_percentage }}%</p>
+<p>Please consider removing old data or increasing your storage allocation to prevent interruption.</p>
+{% if support_email %}
+<p>If you need assistance, please contact support at <a id="support-email" href="mailto:{{ support_email }}">{{ support_email }}</a>.</p>
+{% endif %}
+</body>
+</html>

phoenix/server/email/types.py

@@ -19,8 +19,19 @@ class PasswordResetEmailSender(Protocol):
     ) -> None: ...
 
 
+class DbUsageWarningEmailSender(Protocol):
+    async def send_db_usage_warning_email(
+        self,
+        email: str,
+        current_usage_gibibytes: float,
+        allocated_storage_gibibytes: float,
+        notification_threshold_percentage: float,
+    ) -> None: ...
+
+
 class EmailSender(
     WelcomeEmailSender,
     PasswordResetEmailSender,
+    DbUsageWarningEmailSender,
     Protocol,
 ): ...

phoenix/server/experiments/utils.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+import re
+from secrets import token_hex
+
+EXPERIMENT_PROJECT_NAME_PATTERN = re.compile(r"^Experiment-[0-9a-f]{24}$")
+
+
+def generate_experiment_project_name() -> str:
+    return f"Experiment-{token_hex(12)}"
+
+
+def is_experiment_project_name(name: str) -> bool:
+    return bool(EXPERIMENT_PROJECT_NAME_PATTERN.match(name))

phoenix/server/grpc_server.py

@@ -1,5 +1,4 @@
-from collections.abc import Awaitable, Callable
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Iterable, Optional
 
 import grpc
 from grpc.aio import RpcContext, Server, ServerInterceptor
@@ -11,6 +10,7 @@ from opentelemetry.proto.collector.trace.v1.trace_service_pb2_grpc import (
     TraceServiceServicer,
     add_TraceServiceServicer_to_server,
 )
+from starlette.concurrency import run_in_threadpool
 from typing_extensions import TypeAlias
 
 from phoenix.auth import CanReadToken
@@ -34,10 +34,10 @@ ProjectName: TypeAlias = str
 class Servicer(TraceServiceServicer):  # type: ignore[misc,unused-ignore]
     def __init__(
         self,
-        callback: Callable[[Span, ProjectName], Awaitable[None]],
+        enqueue_span: Callable[[Span, ProjectName], Awaitable[None]],
     ) -> None:
         super().__init__()
-        self._callback = callback
+        self._enqueue_span = enqueue_span
 
     async def Export(
         self,
@@ -48,28 +48,28 @@ class Servicer(TraceServiceServicer):  # type: ignore[misc,unused-ignore]
             project_name = get_project_name(resource_spans.resource.attributes)
             for scope_span in resource_spans.scope_spans:
                 for otlp_span in scope_span.spans:
-                    span = decode_otlp_span(otlp_span)
-                    await self._callback(span, project_name)
+                    span = await run_in_threadpool(decode_otlp_span, otlp_span)
+                    await self._enqueue_span(span, project_name)
         return ExportTraceServiceResponse()
 
 
 class GrpcServer:
     def __init__(
         self,
-        callback: Callable[[Span, ProjectName], Awaitable[None]],
+        enqueue_span: Callable[[Span, ProjectName], Awaitable[None]],
         tracer_provider: Optional["TracerProvider"] = None,
         enable_prometheus: bool = False,
         disabled: bool = False,
         token_store: Optional[CanReadToken] = None,
-        interceptors: list[ServerInterceptor] = [],
+        interceptors: Iterable[ServerInterceptor] = (),
     ) -> None:
-        self._callback = callback
+        self._enqueue_span = enqueue_span
         self._server: Optional[Server] = None
         self._tracer_provider = tracer_provider
         self._enable_prometheus = enable_prometheus
         self._disabled = disabled
         self._token_store = token_store
-        self._interceptors = interceptors
+        self._interceptors = list(interceptors)
 
     async def __aenter__(self) -> None:
         interceptors = self._interceptors
@@ -106,7 +106,7 @@ class GrpcServer:
             server.add_secure_port(f"[::]:{get_env_grpc_port()}", server_credentials)
         else:
             server.add_insecure_port(f"[::]:{get_env_grpc_port()}")
-        add_TraceServiceServicer_to_server(Servicer(self._callback), server)  # type: ignore[no-untyped-call,unused-ignore]
+        add_TraceServiceServicer_to_server(Servicer(self._enqueue_span), server)  # type: ignore[no-untyped-call,unused-ignore]
         await server.start()
         self._server = server
 

phoenix/server/jwt_store.py

@@ -20,7 +20,7 @@ from phoenix.auth import (
 )
 from phoenix.config import get_env_enable_prometheus
 from phoenix.db import models
-from phoenix.db.enums import UserRole
+from phoenix.db.models import UserRoleName
 from phoenix.server.types import (
     AccessToken,
     AccessTokenAttributes,
@@ -164,7 +164,7 @@ class JwtStore:
         for token_id in token_ids:
             if isinstance(token_id, PasswordResetTokenId):
                 password_reset_token_ids.append(token_id)
-            if isinstance(token_id, AccessTokenId):
+            elif isinstance(token_id, AccessTokenId):
                 access_token_ids.append(token_id)
             elif isinstance(token_id, RefreshTokenId):
                 refresh_token_ids.append(token_id)
@@ -182,10 +182,10 @@ class JwtStore:
         await gather(*coroutines)
 
     async def log_out(self, user_id: UserId) -> None:
-        for cls in (AccessTokenId, RefreshTokenId):
-            table = cls.table
-            stmt = delete(table).where(table.user_id == int(user_id)).returning(table.id)
-            async with self._db() as session:
+        async with self._db() as session:
+            for cls in (AccessTokenId, RefreshTokenId):
+                table = cls.table
+                stmt = delete(table).where(table.user_id == int(user_id)).returning(table.id)
                 async for id_ in await session.stream_scalars(stmt):
                     await self._evict(cls(id_))
 
@@ -260,7 +260,7 @@ class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC)
         if not record:
             return None
         token, role = record
-        _, claims = self._from_db(token, UserRole(role))
+        _, claims = self._from_db(token, role)
         self._claims[token_id] = claims
         return claims
 
@@ -277,7 +277,7 @@ 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: UserRoleName) -> tuple[_TokenIdT, _ClaimSetT]: ...
 
     @abstractmethod
     def _to_db(self, claims: _ClaimSetT) -> _RecordT: ...
@@ -300,12 +300,12 @@ class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC)
                 await self._delete_expired_tokens(session)
             async with session.begin_nested():
                 async for record, role in await session.stream(self._update_stmt):
-                    token_id, claim_set = self._from_db(record, UserRole(role))
+                    token_id, claim_set = self._from_db(record, role)
                     claims[token_id] = claim_set
         self._claims = claims
 
     @cached_property
-    def _update_stmt(self) -> Select[tuple[_RecordT, str]]:
+    def _update_stmt(self) -> Select[tuple[_RecordT, UserRoleName]]:
         return (
             select(self._table, models.UserRole.name)
             .join_from(self._table, models.User)
@@ -314,7 +314,9 @@ class _Store(DaemonTask, Generic[_ClaimSetT, _TokenT, _TokenIdT, _RecordT], ABC)
 
     async def _delete_expired_tokens(self, session: Any) -> None:
         now = datetime.now(timezone.utc)
-        await session.execute(delete(self._table).where(self._table.expires_at < now))
+        # Per JWT RFC 7519 Section 4.1.4, tokens expire "on or after" the expiration time.
+        # Use <= to include tokens expiring at exactly this moment.
+        await session.execute(delete(self._table).where(self._table.expires_at <= now))
 
     async def _run(self) -> None:
         while self._running:
@@ -341,7 +343,7 @@ class _PasswordResetTokenStore(
     def _from_db(
         self,
         record: models.PasswordResetToken,
-        user_role: UserRole,
+        user_role: UserRoleName,
     ) -> tuple[PasswordResetTokenId, PasswordResetTokenClaims]:
         token_id = PasswordResetTokenId(record.id)
         return token_id, PasswordResetTokenClaims(
@@ -380,7 +382,7 @@ class _AccessTokenStore(
     def _from_db(
         self,
         record: models.AccessToken,
-        user_role: UserRole,
+        user_role: UserRoleName,
     ) -> tuple[AccessTokenId, AccessTokenClaims]:
         token_id = AccessTokenId(record.id)
         refresh_token_id = RefreshTokenId(record.refresh_token_id)
@@ -424,7 +426,7 @@ class _RefreshTokenStore(
     def _from_db(
         self,
         record: models.RefreshToken,
-        user_role: UserRole,
+        user_role: UserRoleName,
     ) -> tuple[RefreshTokenId, RefreshTokenClaims]:
         token_id = RefreshTokenId(record.id)
         return token_id, RefreshTokenClaims(
@@ -470,7 +472,7 @@ class _ApiKeyStore(
     def _from_db(
         self,
         record: models.ApiKey,
-        user_role: UserRole,
+        user_role: UserRoleName,
     ) -> tuple[ApiKeyId, ApiKeyClaims]:
         token_id = ApiKeyId(record.id)
         return token_id, ApiKeyClaims(