arize-phoenix 11.7.0__py3-none-any.whl → 11.9.0__py3-none-any.whl

phoenix/db/insertion/span.py

@@ -28,15 +28,6 @@ async def insert_span(
     project_name: str,
 ) -> Optional[SpanInsertionEvent]:
     dialect = SupportedSQLDialect(session.bind.dialect.name)
-    if (
-        project_rowid := await session.scalar(
-            select(models.Project.id).filter_by(name=project_name)
-        )
-    ) is None:
-        project_rowid = await session.scalar(
-            insert(models.Project).values(name=project_name).returning(models.Project.id)
-        )
-    assert project_rowid is not None
 
     trace_id = span.context.trace_id
     trace: models.Trace = await session.scalar(
@@ -44,16 +35,27 @@ async def insert_span(
     ) or models.Trace(trace_id=trace_id)
 
     if trace.id is not None:
+        # We use the existing project_rowid on the trace because we allow users to transfer traces
+        # between projects, so the project_name parameter is ignored for existing traces.
+        project_rowid = trace.project_rowid
         # Trace record may need to be updated.
         if trace.end_time < span.end_time:
             trace.end_time = span.end_time
-            trace.project_rowid = project_rowid
         if span.start_time < trace.start_time:
             trace.start_time = span.start_time
     else:
         # Trace record needs to be persisted for the first time.
         trace.start_time = span.start_time
         trace.end_time = span.end_time
+        if (
+            project_rowid := await session.scalar(
+                select(models.Project.id).filter_by(name=project_name)
+            )
+        ) is None:
+            project_rowid = await session.scalar(
+                insert(models.Project).values(name=project_name).returning(models.Project.id)
+            )
+            assert project_rowid is not None
         trace.project_rowid = project_rowid
         session.add(trace)
 

phoenix/db/insertion/types.py

@@ -94,7 +94,10 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
     async def insert(self) -> Optional[list[_DmlEventT]]:
         if not self._queue:
             return None
-        self._queue, parcels = [], self._queue
+        parcels = self._queue.copy()
+        # IMPORTANT: Use .clear() instead of reassignment, i.e. self._queue = [], to
+        # avoid potential race conditions when appending postponed items to the queue.
+        self._queue.clear()
         events: list[_DmlEventT] = []
         async with self._db() as session:
             to_insert, to_postpone, _ = await self._partition(session, *parcels)
@@ -104,9 +107,13 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
                     to_postpone.extend(to_retry)
         if to_postpone:
             loop = asyncio.get_running_loop()
-            loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
+            loop.call_later(self._retry_delay_sec, self._add_postponed_to_queue, to_postpone)
         return events
 
+    def _add_postponed_to_queue(self, items: list[Postponed[_PrecursorT]]) -> None:
+        """Add postponed items back to the queue for retry."""
+        self._queue.extend(items)
+
     def _insert_on_conflict(self, *records: Mapping[str, Any]) -> Insert:
         return insert_on_conflict(
             *records,

phoenix/server/api/auth.py

@@ -3,8 +3,10 @@ from typing import Any
 
 from strawberry import Info
 from strawberry.permission import BasePermission
+from typing_extensions import override
 
-from phoenix.server.api.exceptions import Unauthorized
+from phoenix.config import get_env_support_email
+from phoenix.server.api.exceptions import InsufficientStorage, Unauthorized
 from phoenix.server.bearer_auth import PhoenixUser
 
 
@@ -20,15 +22,35 @@ class IsNotReadOnly(Authorization):
         return not info.context.read_only
 
 
-class IsLocked(Authorization):
-    """
-    Disables mutations and subscriptions that create or update data but allows
-    queries and delete mutations.
+class IsLocked(BasePermission):
     """
+    Permission class that restricts data-modifying operations when insufficient storage.
+
+    When database storage capacity is exceeded, this permission blocks mutations and
+    subscriptions that create or update data, while allowing queries and delete mutations
+    to continue. This prevents database overflow while maintaining read access and the
+    ability to free up space through deletions.
 
-    message = "Operations that write or modify data are locked"
+    Raises:
+        InsufficientStorage: When storage capacity is exceeded and data operations
+            are temporarily disabled. The error includes guidance for resolution
+            and support contact information if configured.
+    """
 
+    @override
+    def on_unauthorized(self) -> None:
+        """Create user-friendly error message when storage operations are blocked."""
+        message = (
+            "Database operations are disabled due to insufficient storage. "
+            "Please delete old data or increase storage."
+        )
+        if support_email := get_env_support_email():
+            message += f" Need help? Contact us at {support_email}"
+        raise InsufficientStorage(message)
+
+    @override
     def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        """Check if database operations are allowed based on storage capacity and lock status."""
         return not (info.context.db.should_not_insert_or_update or info.context.locked)
 
 

phoenix/server/api/dataloaders/span_cost_summary_by_experiment.py

@@ -1,7 +1,6 @@
 from collections import defaultdict
 
 from sqlalchemy import func, select
-from sqlalchemy.sql.functions import coalesce
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -23,12 +22,12 @@ class SpanCostSummaryByExperimentDataLoader(DataLoader[Key, Result]):
         stmt = (
             select(
                 models.ExperimentRun.experiment_id,
-                coalesce(func.sum(models.SpanCost.prompt_cost), 0).label("prompt_cost"),
-                coalesce(func.sum(models.SpanCost.completion_cost), 0).label("completion_cost"),
-                coalesce(func.sum(models.SpanCost.total_cost), 0).label("total_cost"),
-                coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
-                coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
-                coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
+                func.sum(models.SpanCost.prompt_cost).label("prompt_cost"),
+                func.sum(models.SpanCost.completion_cost).label("completion_cost"),
+                func.sum(models.SpanCost.total_cost).label("total_cost"),
+                func.sum(models.SpanCost.prompt_tokens).label("prompt_tokens"),
+                func.sum(models.SpanCost.completion_tokens).label("completion_tokens"),
+                func.sum(models.SpanCost.total_tokens).label("total_tokens"),
             )
             .select_from(models.ExperimentRun)
             .join(models.Trace, models.ExperimentRun.trace_id == models.Trace.trace_id)

phoenix/server/api/exceptions.py

@@ -27,6 +27,12 @@ class Unauthorized(CustomGraphQLError):
     """
 
 
+class InsufficientStorage(CustomGraphQLError):
+    """
+    An error raised when the database has insufficient storage to complete a request.
+    """
+
+
 class Conflict(CustomGraphQLError):
     """
     An error raised when a mutation cannot be completed due to a conflict with

phoenix/server/api/input_types/CreateProjectInput.py

@@ -0,0 +1,27 @@
+import re
+from typing import Optional
+
+import strawberry
+from strawberry import UNSET
+
+from phoenix.server.api.exceptions import BadRequest
+
+
+@strawberry.input
+class CreateProjectInput:
+    name: str
+    description: Optional[str] = UNSET
+    gradient_start_color: Optional[str] = UNSET
+    gradient_end_color: Optional[str] = UNSET
+
+    def __post_init__(self) -> None:
+        if not self.name.strip():
+            raise BadRequest("Name cannot be empty")
+        if self.gradient_start_color and not re.match(
+            r"^#([0-9a-fA-F]{6})$", self.gradient_start_color
+        ):
+            raise BadRequest("Gradient start color must be a valid hex color")
+        if self.gradient_end_color and not re.match(
+            r"^#([0-9a-fA-F]{6})$", self.gradient_end_color
+        ):
+            raise BadRequest("Gradient end color must be a valid hex color")

phoenix/server/api/input_types/TimeBinConfig.py

@@ -0,0 +1,23 @@
+from enum import Enum
+
+import strawberry
+
+
+@strawberry.enum
+class TimeBinScale(Enum):
+    MINUTE = "minute"
+    HOUR = "hour"
+    DAY = "day"
+    WEEK = "week"
+    MONTH = "month"
+    YEAR = "year"
+
+
+@strawberry.input
+class TimeBinConfig:
+    scale: TimeBinScale = strawberry.field(
+        default=TimeBinScale.HOUR, description="The scale of time bins for aggregation."
+    )
+    utc_offset_minutes: int = strawberry.field(
+        default=0, description="Offset in minutes from UTC for local time binning."
+    )

phoenix/server/api/mutations/project_mutations.py

@@ -1,6 +1,8 @@
 import strawberry
 from sqlalchemy import delete, select
+from sqlalchemy.exc import IntegrityError as PostgreSQLIntegrityError
 from sqlalchemy.orm import load_only
+from sqlean.dbapi2 import IntegrityError as SQLiteIntegrityError  # type: ignore[import-untyped]
 from strawberry.relay import GlobalID
 from strawberry.types import Info
 
@@ -8,14 +10,48 @@ from phoenix.config import DEFAULT_PROJECT_NAME
 from phoenix.db import models
 from phoenix.server.api.auth import IsNotReadOnly
 from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest, Conflict
 from phoenix.server.api.input_types.ClearProjectInput import ClearProjectInput
+from phoenix.server.api.input_types.CreateProjectInput import CreateProjectInput
 from phoenix.server.api.queries import Query
 from phoenix.server.api.types.node import from_global_id_with_expected_type
-from phoenix.server.dml_event import ProjectDeleteEvent, SpanDeleteEvent
+from phoenix.server.api.types.Project import Project, to_gql_project
+from phoenix.server.dml_event import ProjectDeleteEvent, ProjectInsertEvent, SpanDeleteEvent
+
+
+@strawberry.type
+class ProjectMutationPayload:
+    project: Project
+    query: Query
 
 
 @strawberry.type
 class ProjectMutationMixin:
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def create_project(
+        self,
+        info: Info[Context, None],
+        input: CreateProjectInput,
+    ) -> ProjectMutationPayload:
+        if not (name := input.name.strip()):
+            raise BadRequest("Name cannot be empty")
+        description = (input.description or "").strip() or None
+        gradient_start_color = (input.gradient_start_color or "").strip() or None
+        gradient_end_color = (input.gradient_end_color or "").strip() or None
+        project = models.Project(
+            name=name,
+            description=description,
+            gradient_start_color=gradient_start_color,
+            gradient_end_color=gradient_end_color,
+        )
+        try:
+            async with info.context.db() as session:
+                session.add(project)
+        except (PostgreSQLIntegrityError, SQLiteIntegrityError):
+            raise Conflict(f"Project with name '{name}' already exists")
+        info.context.event_queue.put(ProjectInsertEvent((project.id,)))
+        return ProjectMutationPayload(project=to_gql_project(project), query=Query())
+
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
     async def delete_project(self, info: Info[Context, None], id: GlobalID) -> Query:
         project_id = from_global_id_with_expected_type(global_id=id, expected_type_name="Project")

phoenix/server/api/mutations/trace_mutations.py

@@ -1,5 +1,5 @@
 import strawberry
-from sqlalchemy import and_, delete, not_, select
+from sqlalchemy import and_, delete, not_, select, update
 from sqlalchemy.orm import load_only
 from sqlalchemy.sql import literal
 from strawberry.relay import GlobalID
@@ -72,3 +72,47 @@ class TraceMutationMixin:
                 )
             info.context.event_queue.put(SpanDeleteEvent(project_ids))
         return Query()
+
+    @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
+    async def transfer_traces_to_project(
+        self,
+        info: Info[Context, None],
+        trace_ids: list[GlobalID],
+        project_id: GlobalID,
+    ) -> Query:
+        if not trace_ids:
+            raise BadRequest("Must provide at least one trace ID to transfer")
+        trace_ids = list(set(trace_ids))
+        try:
+            trace_rowids = [
+                from_global_id_with_expected_type(global_id=id, expected_type_name="Trace")
+                for id in trace_ids
+            ]
+            dest_project_rowid = from_global_id_with_expected_type(
+                global_id=project_id, expected_type_name="Project"
+            )
+        except ValueError as error:
+            raise BadRequest(str(error))
+
+        async with info.context.db() as session:
+            dest_project = await session.get(models.Project, dest_project_rowid)
+            if dest_project is None:
+                raise BadRequest("Destination project does not exist")
+
+            traces = (
+                await session.scalars(select(models.Trace).where(models.Trace.id.in_(trace_rowids)))
+            ).all()
+            if len(traces) < len(trace_rowids):
+                raise BadRequest("Invalid trace IDs provided")
+
+            source_project_ids = set(trace.project_rowid for trace in traces)
+            if len(source_project_ids) > 1:
+                raise BadRequest("Cannot transfer traces from multiple projects")
+
+            await session.execute(
+                update(models.Trace)
+                .where(models.Trace.id.in_(trace_rowids))
+                .values(project_rowid=dest_project_rowid)
+            )
+
+        return Query()

phoenix/server/api/routers/oauth2.py

@@ -169,7 +169,11 @@ async def create_tokens(
             error=f"OAuth2 IDP {idp_name} does not appear to support OpenID Connect.",
         )
     user_info = await oauth2_client.parse_id_token(token_data, nonce=stored_nonce)
-    user_info = _parse_user_info(user_info)
+    try:
+        user_info = _parse_user_info(user_info)
+    except MissingEmailScope as error:
+        return _redirect_to_login(request=request, error=str(error))
+
     try:
         async with request.app.state.db() as session:
             user = await _process_oauth2_user(
@@ -237,7 +241,12 @@ def _parse_user_info(user_info: dict[str, Any]) -> UserInfo:
     """
     assert isinstance(subject := user_info.get("sub"), (str, int))
     idp_user_id = str(subject)
-    assert isinstance(email := user_info.get("email"), str)
+    email = user_info.get("email")
+    if not isinstance(email, str):
+        raise MissingEmailScope(
+            "Please ensure your OIDC provider is configured to use the 'email' scope."
+        )
+
     assert isinstance(username := user_info.get("name"), str) or username is None
     assert (
         isinstance(profile_picture_url := user_info.get("picture"), str)
@@ -541,6 +550,14 @@ class NotInvited(Exception):
     pass
 
 
+class MissingEmailScope(Exception):
+    """
+    Raised when the OIDC provider does not return the email scope.
+    """
+
+    pass
+
+
 def _redirect_to_login(*, request: Request, error: str) -> RedirectResponse:
     """
     Creates a RedirectResponse to the login page to display an error message.

phoenix/server/api/types/CostBreakdown.py

@@ -5,11 +5,8 @@ import strawberry
 
 @strawberry.type
 class CostBreakdown:
-    tokens: Optional[float] = None
+    tokens: Optional[float] = strawberry.field(
+        default=None,
+        description="Total number of tokens, including tokens for which no cost was computed.",
+    )
     cost: Optional[float] = None
-
-    @strawberry.field
-    def cost_per_token(self) -> Optional[float]:
-        if self.tokens and self.cost:
-            return self.cost / self.tokens
-        return None