arize-phoenix 12.3.0__py3-none-any.whl → 12.5.0__py3-none-any.whl

phoenix/server/api/subscriptions.py

@@ -27,7 +27,7 @@ from phoenix.config import PLAYGROUND_PROJECT_NAME
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
 from phoenix.db.helpers import insert_experiment_with_examples_snapshot
-from phoenix.server.api.auth import IsLocked, IsNotReadOnly
+from phoenix.server.api.auth import IsLocked, IsNotReadOnly, IsNotViewer
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
 from phoenix.server.api.helpers.playground_clients import (
@@ -94,7 +94,7 @@ ChatStream: TypeAlias = AsyncGenerator[ChatCompletionSubscriptionPayload, None]
 
 @strawberry.type
 class Subscription:
-    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def chat_completion(
         self, info: Info[Context, None], input: ChatCompletionInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:
@@ -193,7 +193,7 @@ class Subscription:
         info.context.event_queue.put(SpanInsertEvent(ids=(playground_project_id,)))
         yield ChatCompletionSubscriptionResult(span=Span(span_rowid=db_span.id, db_span=db_span))
 
-    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
+    @strawberry.subscription(permission_classes=[IsNotReadOnly, IsNotViewer, IsLocked])  # type: ignore
     async def chat_completion_over_dataset(
         self, info: Info[Context, None], input: ChatCompletionOverDatasetInput
     ) -> AsyncIterator[ChatCompletionSubscriptionPayload]:

phoenix/server/api/types/Dataset.py

@@ -3,7 +3,7 @@ from datetime import datetime
 from typing import ClassVar, Optional, cast
 
 import strawberry
-from sqlalchemy import and_, func, or_, select
+from sqlalchemy import Text, and_, func, or_, select
 from sqlalchemy.sql.functions import count
 from strawberry import UNSET
 from strawberry.relay import Connection, GlobalID, Node, NodeID
@@ -19,6 +19,7 @@ from phoenix.server.api.types.DatasetExperimentAnnotationSummary import (
     DatasetExperimentAnnotationSummary,
 )
 from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
+from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
 from phoenix.server.api.types.node import from_global_id_with_expected_type
@@ -87,6 +88,7 @@ class Dataset(Node):
         self,
         info: Info[Context, None],
         dataset_version_id: Optional[GlobalID] = UNSET,
+        split_ids: Optional[list[GlobalID]] = UNSET,
     ) -> int:
         dataset_id = self.id_attr
         version_id = (
@@ -97,6 +99,20 @@ class Dataset(Node):
             if dataset_version_id
             else None
         )
+
+        # Parse split IDs if provided
+        split_rowids: Optional[list[int]] = None
+        if split_ids:
+            split_rowids = []
+            for split_id in split_ids:
+                try:
+                    split_rowid = from_global_id_with_expected_type(
+                        global_id=split_id, expected_type_name=models.DatasetSplit.__name__
+                    )
+                    split_rowids.append(split_rowid)
+                except Exception:
+                    raise BadRequest(f"Invalid split ID: {split_id}")
+
         revision_ids = (
             select(func.max(models.DatasetExampleRevision.id))
             .join(models.DatasetExample)
@@ -113,11 +129,36 @@ class Dataset(Node):
             revision_ids = revision_ids.where(
                 models.DatasetExampleRevision.dataset_version_id <= version_id_subquery
             )
-        stmt = (
-            select(count(models.DatasetExampleRevision.id))
-            .where(models.DatasetExampleRevision.id.in_(revision_ids))
-            .where(models.DatasetExampleRevision.revision_kind != "DELETE")
-        )
+
+        # Build the count query
+        if split_rowids:
+            # When filtering by splits, count distinct examples that belong to those splits
+            stmt = (
+                select(count(models.DatasetExample.id.distinct()))
+                .join(
+                    models.DatasetExampleRevision,
+                    onclause=(
+                        models.DatasetExample.id == models.DatasetExampleRevision.dataset_example_id
+                    ),
+                )
+                .join(
+                    models.DatasetSplitDatasetExample,
+                    onclause=(
+                        models.DatasetExample.id
+                        == models.DatasetSplitDatasetExample.dataset_example_id
+                    ),
+                )
+                .where(models.DatasetExampleRevision.id.in_(revision_ids))
+                .where(models.DatasetExampleRevision.revision_kind != "DELETE")
+                .where(models.DatasetSplitDatasetExample.dataset_split_id.in_(split_rowids))
+            )
+        else:
+            stmt = (
+                select(count(models.DatasetExampleRevision.id))
+                .where(models.DatasetExampleRevision.id.in_(revision_ids))
+                .where(models.DatasetExampleRevision.revision_kind != "DELETE")
+            )
+
         async with info.context.db() as session:
             return (await session.scalar(stmt)) or 0
 
@@ -126,10 +167,12 @@ class Dataset(Node):
         self,
         info: Info[Context, None],
         dataset_version_id: Optional[GlobalID] = UNSET,
+        split_ids: Optional[list[GlobalID]] = UNSET,
         first: Optional[int] = 50,
         last: Optional[int] = UNSET,
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
+        filter: Optional[str] = UNSET,
     ) -> Connection[DatasetExample]:
         args = ConnectionArgs(
             first=first,
@@ -145,6 +188,20 @@ class Dataset(Node):
             if dataset_version_id
             else None
         )
+
+        # Parse split IDs if provided
+        split_rowids: Optional[list[int]] = None
+        if split_ids:
+            split_rowids = []
+            for split_id in split_ids:
+                try:
+                    split_rowid = from_global_id_with_expected_type(
+                        global_id=split_id, expected_type_name=models.DatasetSplit.__name__
+                    )
+                    split_rowids.append(split_rowid)
+                except Exception:
+                    raise BadRequest(f"Invalid split ID: {split_id}")
+
         revision_ids = (
             select(func.max(models.DatasetExampleRevision.id))
             .join(models.DatasetExample)
@@ -176,6 +233,31 @@ class Dataset(Node):
             )
             .order_by(models.DatasetExampleRevision.dataset_example_id.desc())
         )
+
+        # Filter by split IDs if provided
+        if split_rowids:
+            query = (
+                query.join(
+                    models.DatasetSplitDatasetExample,
+                    onclause=(
+                        models.DatasetExample.id
+                        == models.DatasetSplitDatasetExample.dataset_example_id
+                    ),
+                )
+                .where(models.DatasetSplitDatasetExample.dataset_split_id.in_(split_rowids))
+                .distinct()
+            )
+        # Apply filter if provided - search through JSON fields (input, output, metadata)
+        if filter is not UNSET and filter:
+            # Create a filter that searches for the filter string in JSON fields
+            # Using PostgreSQL's JSON operators for case-insensitive text search
+            filter_condition = or_(
+                func.cast(models.DatasetExampleRevision.input, Text).ilike(f"%{filter}%"),
+                func.cast(models.DatasetExampleRevision.output, Text).ilike(f"%{filter}%"),
+                func.cast(models.DatasetExampleRevision.metadata_, Text).ilike(f"%{filter}%"),
+            )
+            query = query.where(filter_condition)
+
         async with info.context.db() as session:
             dataset_examples = [
                 DatasetExample(
@@ -187,6 +269,13 @@ class Dataset(Node):
             ]
         return connection_from_list(data=dataset_examples, args=args)
 
+    @strawberry.field
+    async def splits(self, info: Info[Context, None]) -> list[DatasetSplit]:
+        return [
+            to_gql_dataset_split(split)
+            for split in await info.context.data_loaders.dataset_dataset_splits.load(self.id_attr)
+        ]
+
     @strawberry.field(
         description="Number of experiments for a specific version if version is specified, "
         "or for all versions if version is not specified."

phoenix/server/api/types/Project.py

@@ -7,7 +7,6 @@ from aioitertools.itertools import groupby, islice
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import and_, case, desc, distinct, exists, func, or_, select
 from sqlalchemy.dialects import postgresql, sqlite
-from sqlalchemy.sql.elements import ColumnElement
 from sqlalchemy.sql.expression import tuple_
 from sqlalchemy.sql.functions import percentile_cont
 from strawberry import ID, UNSET, Private, lazy
@@ -21,8 +20,8 @@ from phoenix.db.helpers import SupportedSQLDialect, date_trunc
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.input_types.ProjectSessionSort import (
-    ProjectSessionColumn,
     ProjectSessionSort,
+    ProjectSessionSortConfig,
 )
 from phoenix.server.api.input_types.SpanSort import SpanColumn, SpanSort, SpanSortConfig
 from phoenix.server.api.input_types.TimeBinConfig import TimeBinConfig, TimeBinScale
@@ -459,74 +458,31 @@ class Project(Node):
                 end_time=time_range.end if time_range else None,
             )
             stmt = stmt.where(table.id.in_(filtered_session_rowids))
+        sort_config: Optional[ProjectSessionSortConfig] = None
+        cursor_rowid_column: Any = table.id
         if sort:
-            key: ColumnElement[Any]
-            if sort.col is ProjectSessionColumn.startTime:
-                key = table.start_time.label("key")
-            elif sort.col is ProjectSessionColumn.endTime:
-                key = table.end_time.label("key")
-            elif (
-                sort.col is ProjectSessionColumn.tokenCountTotal
-                or sort.col is ProjectSessionColumn.numTraces
-            ):
-                if sort.col is ProjectSessionColumn.tokenCountTotal:
-                    sort_subq = (
-                        select(
-                            models.Trace.project_session_rowid.label("id"),
-                            func.sum(models.Span.cumulative_llm_token_count_total).label("key"),
-                        )
-                        .join_from(models.Trace, models.Span)
-                        .where(models.Span.parent_id.is_(None))
-                        .group_by(models.Trace.project_session_rowid)
-                    ).subquery()
-                elif sort.col is ProjectSessionColumn.numTraces:
-                    sort_subq = (
-                        select(
-                            models.Trace.project_session_rowid.label("id"),
-                            func.count(models.Trace.id).label("key"),
-                        ).group_by(models.Trace.project_session_rowid)
-                    ).subquery()
+            sort_config = sort.update_orm_expr(stmt)
+            stmt = sort_config.stmt
+            if sort_config.dir is SortDir.desc:
+                cursor_rowid_column = desc(cursor_rowid_column)
+        if after:
+            cursor = Cursor.from_string(after)
+            if sort_config and cursor.sort_column:
+                sort_column = cursor.sort_column
+                compare = operator.lt if sort_config.dir is SortDir.desc else operator.gt
+                if sort_column.type is CursorSortColumnDataType.NULL:
+                    stmt = stmt.where(sort_config.orm_expression.is_(None))
+                    stmt = stmt.where(compare(table.id, cursor.rowid))
                 else:
-                    assert_never(sort.col)
-                key = sort_subq.c.key
-                stmt = stmt.join(sort_subq, table.id == sort_subq.c.id)
-            elif sort.col is ProjectSessionColumn.costTotal:
-                sort_subq = (
-                    select(
-                        models.Trace.project_session_rowid.label("id"),
-                        func.sum(models.SpanCost.total_cost).label("key"),
-                    )
-                    .join_from(
-                        models.Trace,
-                        models.SpanCost,
-                        models.Trace.id == models.SpanCost.trace_rowid,
+                    stmt = stmt.where(
+                        compare(
+                            tuple_(sort_config.orm_expression, table.id),
+                            (sort_column.value, cursor.rowid),
+                        )
                     )
-                    .group_by(models.Trace.project_session_rowid)
-                ).subquery()
-                key = sort_subq.c.key
-                stmt = stmt.join(sort_subq, table.id == sort_subq.c.id)
-            else:
-                assert_never(sort.col)
-            stmt = stmt.add_columns(key)
-            if sort.dir is SortDir.asc:
-                stmt = stmt.order_by(key.asc(), table.id.asc())
             else:
-                stmt = stmt.order_by(key.desc(), table.id.desc())
-            if after:
-                cursor = Cursor.from_string(after)
-                assert cursor.sort_column is not None
-                compare = operator.lt if sort.dir is SortDir.desc else operator.gt
-                stmt = stmt.where(
-                    compare(
-                        tuple_(key, table.id),
-                        (cursor.sort_column.value, cursor.rowid),
-                    )
-                )
-        else:
-            stmt = stmt.order_by(table.id.desc())
-            if after:
-                cursor = Cursor.from_string(after)
                 stmt = stmt.where(table.id < cursor.rowid)
+        stmt = stmt.order_by(cursor_rowid_column)
         if first:
             stmt = stmt.limit(
                 first + 1  # over-fetch by one to determine whether there's a next page
@@ -537,10 +493,10 @@ class Project(Node):
             async for record in islice(records, first):
                 project_session = record[0]
                 cursor = Cursor(rowid=project_session.id)
-                if sort:
+                if sort_config:
                     assert len(record) > 1
                     cursor.sort_column = CursorSortColumn(
-                        type=sort.col.data_type,
+                        type=sort_config.column_data_type,
                         value=record[1],
                     )
                 cursors_and_nodes.append((cursor, to_gql_project_session(project_session)))
@@ -724,7 +680,7 @@ class Project(Node):
             stmt = span_filter(select(models.Span))
             dialect = info.context.db.dialect
             if dialect is SupportedSQLDialect.POSTGRESQL:
-                str(stmt.compile(dialect=sqlite.dialect()))  # type: ignore[no-untyped-call]
+                str(stmt.compile(dialect=sqlite.dialect()))
             elif dialect is SupportedSQLDialect.SQLITE:
                 str(stmt.compile(dialect=postgresql.dialect()))  # type: ignore[no-untyped-call]
             else:

phoenix/server/app.py

@@ -45,7 +45,6 @@ from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoin
 from starlette.requests import Request
 from starlette.responses import JSONResponse, PlainTextResponse, RedirectResponse, Response
 from starlette.staticfiles import StaticFiles
-from starlette.status import HTTP_401_UNAUTHORIZED
 from starlette.templating import Jinja2Templates
 from starlette.types import Scope, StatefulLifespan
 from strawberry.extensions import SchemaExtension
@@ -89,6 +88,7 @@ from phoenix.server.api.dataloaders import (
     AverageExperimentRepeatedRunGroupLatencyDataLoader,
     AverageExperimentRunLatencyDataLoader,
     CacheForDataLoaders,
+    DatasetDatasetSplitsDataLoader,
     DatasetExampleRevisionsDataLoader,
     DatasetExamplesAndVersionsByExperimentRunDataLoader,
     DatasetExampleSpansDataLoader,
@@ -352,7 +352,7 @@ class RequestOriginHostnameValidator(BaseHTTPMiddleware):
             if not (url := headers.get(key)):
                 continue
             if urlparse(url).hostname not in self._trusted_hostnames:
-                return Response(f"untrusted {key}", status_code=HTTP_401_UNAUTHORIZED)
+                return Response(f"untrusted {key}", status_code=401)
         return await call_next(request)
 
 
@@ -710,6 +710,7 @@ def create_graphql_router(
                     db
                 ),
                 average_experiment_run_latency=AverageExperimentRunLatencyDataLoader(db),
+                dataset_dataset_splits=DatasetDatasetSplitsDataLoader(db),
                 dataset_example_revisions=DatasetExampleRevisionsDataLoader(db),
                 dataset_example_spans=DatasetExampleSpansDataLoader(db),
                 dataset_examples_and_versions_by_experiment_run=DatasetExamplesAndVersionsByExperimentRunDataLoader(

phoenix/server/authorization.py

@@ -23,7 +23,6 @@ Usage:
 """
 
 from fastapi import HTTPException, Request
-from fastapi import status as fastapi_status
 
 from phoenix.config import get_env_support_email
 from phoenix.server.bearer_auth import PhoenixUser
@@ -43,13 +42,15 @@ def require_admin(request: Request) -> None:
     Behavior:
         - Allows access if the authenticated user is an admin or a system user.
         - Raises HTTP 403 Forbidden if the user is not authorized.
-        - Expects authentication to be enabled and request.user to be set by the authentication.
+        - Allows access if authentication is not enabled.
     """
+    if not request.app.state.authentication_enabled:
+        return
     user = getattr(request, "user", None)
     # System users have all privileges
     if not (isinstance(user, PhoenixUser) and user.is_admin):
         raise HTTPException(
-            status_code=fastapi_status.HTTP_403_FORBIDDEN,
+            status_code=403,
             detail="Only admin or system users can perform this action.",
         )
 
@@ -82,6 +83,6 @@ def is_not_locked(request: Request) -> None:
         if support_email := get_env_support_email():
             detail += f" Need help? Contact us at {support_email}"
         raise HTTPException(
-            status_code=fastapi_status.HTTP_507_INSUFFICIENT_STORAGE,
+            status_code=507,
             detail=detail,
         )

phoenix/server/bearer_auth.py

@@ -9,7 +9,6 @@ from fastapi import HTTPException, Request, WebSocket, WebSocketException
 from grpc_interceptor import AsyncServerInterceptor
 from starlette.authentication import AuthCredentials, AuthenticationBackend, BaseUser
 from starlette.requests import HTTPConnection
-from starlette.status import HTTP_401_UNAUTHORIZED
 from typing_extensions import override
 
 from phoenix import config
@@ -76,11 +75,18 @@ class PhoenixUser(BaseUser):
         self._is_admin = (
             claims.status is ClaimSetStatus.VALID and claims.attributes.user_role == "ADMIN"
         )
+        self._is_viewer = (
+            claims.status is ClaimSetStatus.VALID and claims.attributes.user_role == "VIEWER"
+        )
 
     @cached_property
     def is_admin(self) -> bool:
         return self._is_admin
 
+    @cached_property
+    def is_viewer(self) -> bool:
+        return self._is_viewer
+
     @cached_property
     def identity(self) -> UserId:
         return self._user_id
@@ -93,6 +99,8 @@ class PhoenixUser(BaseUser):
 class PhoenixSystemUser(PhoenixUser):
     def __init__(self, user_id: UserId) -> None:
         self._user_id = user_id
+        self._is_admin = True  # System users have admin privileges
+        self._is_viewer = False  # System users are not viewers
 
     @property
     def is_admin(self) -> bool:
@@ -144,16 +152,16 @@ async def is_authenticated(
     """
     assert request or websocket
     if request and not isinstance((user := request.user), PhoenixUser):
-        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+        raise HTTPException(status_code=401, detail="Invalid token")
     if websocket and not isinstance((user := websocket.user), PhoenixUser):
-        raise WebSocketException(code=HTTP_401_UNAUTHORIZED, reason="Invalid token")
+        raise WebSocketException(code=401, reason="Invalid token")
     if isinstance(user, PhoenixSystemUser):
         return
     claims = user.claims
     if claims.status is ClaimSetStatus.EXPIRED:
-        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Expired token")
+        raise HTTPException(status_code=401, detail="Expired token")
     if claims.status is not ClaimSetStatus.VALID:
-        raise HTTPException(status_code=HTTP_401_UNAUTHORIZED, detail="Invalid token")
+        raise HTTPException(status_code=401, detail="Invalid token")
 
 
 async def create_access_and_refresh_tokens(

phoenix/server/jwt_store.py

@@ -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_))
 
@@ -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: