arize-phoenix 4.4.2__py3-none-any.whl → 4.4.4rc0__py3-none-any.whl

phoenix/server/api/types/Dataset.py

@@ -1,80 +1,299 @@
 from datetime import datetime
-from typing import Iterable, List, Optional, Set, Union
+from typing import AsyncIterable, List, Optional, Tuple, cast
 
 import strawberry
-from strawberry.scalars import ID
-from strawberry.unset import UNSET
+from sqlalchemy import and_, func, select
+from sqlalchemy.sql.functions import count
+from strawberry import UNSET
+from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry.scalars import JSON
+from strawberry.types import Info
 
-import phoenix.core.model_schema as ms
-from phoenix.core.model_schema import FEATURE, TAG, ScalarDimension
-
-from ..input_types.DimensionInput import DimensionInput
-from .DatasetRole import AncillaryDatasetRole, DatasetRole
-from .Dimension import Dimension, to_gql_dimension
-from .Event import Event, create_event, create_event_id, parse_event_ids_by_dataset_role
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.DatasetVersionSort import DatasetVersionSort
+from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetVersion import DatasetVersion
+from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
+from phoenix.server.api.types.ExperimentAnnotationSummary import ExperimentAnnotationSummary
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.pagination import (
+    ConnectionArgs,
+    CursorString,
+    connection_from_list,
+)
+from phoenix.server.api.types.SortDir import SortDir
 
 
 @strawberry.type
-class Dataset:
-    start_time: datetime = strawberry.field(description="The start bookend of the data")
-    end_time: datetime = strawberry.field(description="The end bookend of the data")
-    record_count: int = strawberry.field(description="The record count of the data")
-    dataset: strawberry.Private[ms.Dataset]
-    dataset_role: strawberry.Private[Union[DatasetRole, AncillaryDatasetRole]]
-    model: strawberry.Private[ms.Model]
-
-    # type ignored here to get around the following: https://github.com/strawberry-graphql/strawberry/issues/1929
-    @strawberry.field(description="Returns a human friendly name for the dataset.")  # type: ignore
-    def name(self) -> str:
-        return self.dataset.display_name
+class Dataset(Node):
+    id_attr: NodeID[int]
+    name: str
+    description: Optional[str]
+    metadata: JSON
+    created_at: datetime
+    updated_at: datetime
 
     @strawberry.field
-    def events(
+    async def versions(
         self,
-        event_ids: List[ID],
-        dimensions: Optional[List[DimensionInput]] = UNSET,
-    ) -> List[Event]:
-        """
-        Returns events for specific event IDs and dimensions. If no input
-        dimensions are provided, returns all features and tags.
-        """
-        if not event_ids:
-            return []
-        row_ids = parse_event_ids_by_dataset_role(event_ids)
-        if len(row_ids) > 1 or self.dataset_role not in row_ids:
-            raise ValueError("eventIds contains IDs from incorrect dataset.")
-        events = self.dataset[row_ids[self.dataset_role]]
-        requested_gql_dimensions = _get_requested_features_and_tags(
-            core_dimensions=self.model.scalar_dimensions,
-            requested_dimension_names=set(dim.name for dim in dimensions)
-            if isinstance(dimensions, list)
-            else None,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+        sort: Optional[DatasetVersionSort] = UNSET,
+    ) -> Connection[DatasetVersion]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
         )
-        return [
-            create_event(
-                event_id=create_event_id(event.id.row_id, self.dataset_role),
-                event=event,
-                dimensions=requested_gql_dimensions,
-                is_document_record=self.dataset_role is AncillaryDatasetRole.corpus,
+        async with info.context.db() as session:
+            stmt = select(models.DatasetVersion).filter_by(dataset_id=self.id_attr)
+            if sort:
+                # For now assume the the column names match 1:1 with the enum values
+                sort_col = getattr(models.DatasetVersion, sort.col.value)
+                if sort.dir is SortDir.desc:
+                    stmt = stmt.order_by(sort_col.desc(), models.DatasetVersion.id.desc())
+                else:
+                    stmt = stmt.order_by(sort_col.asc(), models.DatasetVersion.id.asc())
+            else:
+                stmt = stmt.order_by(models.DatasetVersion.created_at.desc())
+            versions = await session.scalars(stmt)
+        data = [
+            DatasetVersion(
+                id_attr=version.id,
+                description=version.description,
+                metadata=version.metadata_,
+                created_at=version.created_at,
             )
-            for event in events
+            for version in versions
         ]
+        return connection_from_list(data=data, args=args)
+
+    @strawberry.field(
+        description="Number of examples in a specific version if version is specified, or in the "
+        "latest version if version is not specified."
+    )  # type: ignore
+    async def example_count(
+        self,
+        info: Info[Context, None],
+        dataset_version_id: Optional[GlobalID] = UNSET,
+    ) -> int:
+        dataset_id = self.id_attr
+        version_id = (
+            from_global_id_with_expected_type(
+                global_id=dataset_version_id,
+                expected_type_name=DatasetVersion.__name__,
+            )
+            if dataset_version_id
+            else None
+        )
+        revision_ids = (
+            select(func.max(models.DatasetExampleRevision.id))
+            .join(models.DatasetExample)
+            .where(models.DatasetExample.dataset_id == dataset_id)
+            .group_by(models.DatasetExampleRevision.dataset_example_id)
+        )
+        if version_id:
+            version_id_subquery = (
+                select(models.DatasetVersion.id)
+                .where(models.DatasetVersion.dataset_id == dataset_id)
+                .where(models.DatasetVersion.id == version_id)
+                .scalar_subquery()
+            )
+            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")
+        )
+        async with info.context.db() as session:
+            return (await session.scalar(stmt)) or 0
+
+    @strawberry.field
+    async def examples(
+        self,
+        info: Info[Context, None],
+        dataset_version_id: Optional[GlobalID] = UNSET,
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[DatasetExample]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        dataset_id = self.id_attr
+        version_id = (
+            from_global_id_with_expected_type(
+                global_id=dataset_version_id, expected_type_name=DatasetVersion.__name__
+            )
+            if dataset_version_id
+            else None
+        )
+        revision_ids = (
+            select(func.max(models.DatasetExampleRevision.id))
+            .join(models.DatasetExample)
+            .where(models.DatasetExample.dataset_id == dataset_id)
+            .group_by(models.DatasetExampleRevision.dataset_example_id)
+        )
+        if version_id:
+            version_id_subquery = (
+                select(models.DatasetVersion.id)
+                .where(models.DatasetVersion.dataset_id == dataset_id)
+                .where(models.DatasetVersion.id == version_id)
+                .scalar_subquery()
+            )
+            revision_ids = revision_ids.where(
+                models.DatasetExampleRevision.dataset_version_id <= version_id_subquery
+            )
+        query = (
+            select(models.DatasetExample)
+            .join(
+                models.DatasetExampleRevision,
+                onclause=models.DatasetExample.id
+                == models.DatasetExampleRevision.dataset_example_id,
+            )
+            .where(
+                and_(
+                    models.DatasetExampleRevision.id.in_(revision_ids),
+                    models.DatasetExampleRevision.revision_kind != "DELETE",
+                )
+            )
+            .order_by(models.DatasetExampleRevision.dataset_example_id.desc())
+        )
+        async with info.context.db() as session:
+            dataset_examples = [
+                DatasetExample(
+                    id_attr=example.id,
+                    version_id=version_id,
+                    created_at=example.created_at,
+                )
+                async for example in await session.stream_scalars(query)
+            ]
+        return connection_from_list(data=dataset_examples, args=args)
+
+    @strawberry.field(
+        description="Number of experiments for a specific version if version is specified, "
+        "or for all versions if version is not specified."
+    )  # type: ignore
+    async def experiment_count(
+        self,
+        info: Info[Context, None],
+        dataset_version_id: Optional[GlobalID] = UNSET,
+    ) -> int:
+        stmt = select(count(models.Experiment.id)).where(
+            models.Experiment.dataset_id == self.id_attr
+        )
+        version_id = (
+            from_global_id_with_expected_type(
+                global_id=dataset_version_id,
+                expected_type_name=DatasetVersion.__name__,
+            )
+            if dataset_version_id
+            else None
+        )
+        if version_id is not None:
+            stmt = stmt.where(models.Experiment.dataset_version_id == version_id)
+        async with info.context.db() as session:
+            return (await session.scalar(stmt)) or 0
+
+    @strawberry.field
+    async def experiments(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[Experiment]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        dataset_id = self.id_attr
+        row_number = func.row_number().over(order_by=models.Experiment.id).label("row_number")
+        query = (
+            select(models.Experiment, row_number)
+            .where(models.Experiment.dataset_id == dataset_id)
+            .order_by(models.Experiment.id.desc())
+        )
+        async with info.context.db() as session:
+            experiments = [
+                to_gql_experiment(experiment, sequence_number)
+                async for experiment, sequence_number in cast(
+                    AsyncIterable[Tuple[models.Experiment, int]],
+                    await session.stream(query),
+                )
+            ]
+        return connection_from_list(data=experiments, args=args)
+
+    @strawberry.field
+    async def experiment_annotation_summaries(
+        self, info: Info[Context, None]
+    ) -> List[ExperimentAnnotationSummary]:
+        dataset_id = self.id_attr
+        query = (
+            select(
+                models.ExperimentRunAnnotation.name,
+                func.min(models.ExperimentRunAnnotation.score),
+                func.max(models.ExperimentRunAnnotation.score),
+                func.avg(models.ExperimentRunAnnotation.score),
+                func.count(),
+                func.count(models.ExperimentRunAnnotation.error),
+            )
+            .join(
+                models.ExperimentRun,
+                models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
+            )
+            .join(
+                models.Experiment,
+                models.ExperimentRun.experiment_id == models.Experiment.id,
+            )
+            .where(models.Experiment.dataset_id == dataset_id)
+            .group_by(models.ExperimentRunAnnotation.name)
+            .order_by(models.ExperimentRunAnnotation.name)
+        )
+        async with info.context.db() as session:
+            return [
+                ExperimentAnnotationSummary(
+                    annotation_name=annotation_name,
+                    min_score=min_score,
+                    max_score=max_score,
+                    mean_score=mean_score,
+                    count=count_,
+                    error_count=error_count,
+                )
+                async for (
+                    annotation_name,
+                    min_score,
+                    max_score,
+                    mean_score,
+                    count_,
+                    error_count,
+                ) in await session.stream(query)
+            ]
 
 
-def _get_requested_features_and_tags(
-    core_dimensions: Iterable[ScalarDimension],
-    requested_dimension_names: Optional[Set[str]] = UNSET,
-) -> List[Dimension]:
+def to_gql_dataset(dataset: models.Dataset) -> Dataset:
     """
-    Returns requested features and tags as a list of strawberry Datasets. If no
-    dimensions are explicitly requested, returns all features and tags.
+    Converts an ORM dataset to a GraphQL dataset.
     """
-    requested_features_and_tags: List[Dimension] = []
-    for id, dim in enumerate(core_dimensions):
-        is_requested = (
-            not isinstance(requested_dimension_names, Set)
-        ) or dim.name in requested_dimension_names
-        is_feature_or_tag = dim.role in (FEATURE, TAG)
-        if is_requested and is_feature_or_tag:
-            requested_features_and_tags.append(to_gql_dimension(id_attr=id, dimension=dim))
-    return requested_features_and_tags
+    return Dataset(
+        id_attr=dataset.id,
+        name=dataset.name,
+        description=dataset.description,
+        metadata=dataset.metadata_,
+        created_at=dataset.created_at,
+        updated_at=dataset.updated_at,
+    )

phoenix/server/api/types/DatasetExample.py

@@ -0,0 +1,85 @@
+from datetime import datetime
+from typing import Optional
+
+import strawberry
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload
+from strawberry import UNSET
+from strawberry.relay.types import Connection, GlobalID, Node, NodeID
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
+from phoenix.server.api.types.DatasetVersion import DatasetVersion
+from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
+from phoenix.server.api.types.node import from_global_id_with_expected_type
+from phoenix.server.api.types.pagination import (
+    ConnectionArgs,
+    CursorString,
+    connection_from_list,
+)
+from phoenix.server.api.types.Span import Span, to_gql_span
+
+
+@strawberry.type
+class DatasetExample(Node):
+    id_attr: NodeID[int]
+    created_at: datetime
+    version_id: strawberry.Private[Optional[int]] = None
+
+    @strawberry.field
+    async def revision(
+        self,
+        info: Info[Context, None],
+        dataset_version_id: Optional[GlobalID] = UNSET,
+    ) -> DatasetExampleRevision:
+        example_id = self.id_attr
+        version_id: Optional[int] = None
+        if dataset_version_id:
+            version_id = from_global_id_with_expected_type(
+                global_id=dataset_version_id, expected_type_name=DatasetVersion.__name__
+            )
+        elif self.version_id is not None:
+            version_id = self.version_id
+        return await info.context.data_loaders.dataset_example_revisions.load(
+            (example_id, version_id)
+        )
+
+    @strawberry.field
+    async def span(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[Span]:
+        return (
+            to_gql_span(span)
+            if (span := await info.context.data_loaders.dataset_example_spans.load(self.id_attr))
+            else None
+        )
+
+    @strawberry.field
+    async def experiment_runs(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[ExperimentRun]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        example_id = self.id_attr
+        query = (
+            select(models.ExperimentRun)
+            .options(joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id))
+            .join(models.Experiment, models.Experiment.id == models.ExperimentRun.experiment_id)
+            .where(models.ExperimentRun.dataset_example_id == example_id)
+            .order_by(models.Experiment.id.desc())
+        )
+        async with info.context.db() as session:
+            runs = (await session.scalars(query)).all()
+        return connection_from_list([to_gql_experiment_run(run) for run in runs], args)

phoenix/server/api/types/DatasetExampleRevision.py

@@ -0,0 +1,34 @@
+from datetime import datetime
+from enum import Enum
+
+import strawberry
+
+from phoenix.db import models
+from phoenix.server.api.types.ExampleRevisionInterface import ExampleRevision
+
+
+@strawberry.enum
+class RevisionKind(Enum):
+    CREATE = "CREATE"
+    PATCH = "PATCH"
+    DELETE = "DELETE"
+
+
+@strawberry.type
+class DatasetExampleRevision(ExampleRevision):
+    """
+    Represents a revision (i.e., update or alteration) of a dataset example.
+    """
+
+    revision_kind: RevisionKind
+    created_at: datetime
+
+    @classmethod
+    def from_orm_revision(cls, revision: models.DatasetExampleRevision) -> "DatasetExampleRevision":
+        return cls(
+            input=revision.input,
+            output=revision.output,
+            metadata=revision.metadata_,
+            revision_kind=RevisionKind(revision.revision_kind),
+            created_at=revision.created_at,
+        )

phoenix/server/api/types/DatasetVersion.py

@@ -0,0 +1,14 @@
+from datetime import datetime
+from typing import Optional
+
+import strawberry
+from strawberry.relay import Node, NodeID
+from strawberry.scalars import JSON
+
+
+@strawberry.type
+class DatasetVersion(Node):
+    id_attr: NodeID[int]
+    description: Optional[str]
+    metadata: JSON
+    created_at: datetime

phoenix/server/api/types/Dimension.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, List, Optional
 import pandas as pd
 import strawberry
 from strawberry import UNSET
+from strawberry.relay import Node, NodeID
 from strawberry.types import Info
 from typing_extensions import Annotated
 
@@ -17,12 +18,11 @@ from ..context import Context
 from ..input_types.Granularity import Granularity
 from ..input_types.TimeRange import TimeRange
 from .DataQualityMetric import DataQualityMetric
-from .DatasetRole import DatasetRole
 from .DatasetValues import DatasetValues
 from .DimensionDataType import DimensionDataType
 from .DimensionShape import DimensionShape
 from .DimensionType import DimensionType
-from .node import Node
+from .InferencesRole import InferencesRole
 from .ScalarDriftMetricEnum import ScalarDriftMetric
 from .Segments import (
     GqlBinFactory,
@@ -40,6 +40,7 @@ from .TimeSeries import (
 
 @strawberry.type
 class Dimension(Node):
+    id_attr: NodeID[int]
     name: str = strawberry.field(description="The name of the dimension (a.k.a. the column name)")
     type: DimensionType = strawberry.field(
         description="Whether the dimension represents a feature, tag, prediction, or actual."
@@ -62,16 +63,16 @@ class Dimension(Node):
         """
         Computes a drift metric between all reference data and the primary data
         belonging to the input time range (inclusive of the time range start and
-        exclusive of the time range end). Returns None if no reference dataset
-        exists, if no primary data exists in the input time range, or if the
+        exclusive of the time range end). Returns None if no reference inferences
+        exist, if no primary data exists in the input time range, or if the
         input time range is invalid.
         """
         model = info.context.model
         if model[REFERENCE].empty:
             return None
-        dataset = model[PRIMARY]
+        inferences = model[PRIMARY]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
         )
         data = get_drift_timeseries_data(
@@ -92,18 +93,18 @@ class Dimension(Node):
         info: Info[Context, None],
         metric: DataQualityMetric,
         time_range: Optional[TimeRange] = UNSET,
-        dataset_role: Annotated[
-            Optional[DatasetRole],
+        inferences_role: Annotated[
+            Optional[InferencesRole],
             strawberry.argument(
-                description="The dataset (primary or reference) to query",
+                description="The inferences (primary or reference) to query",
             ),
-        ] = DatasetRole.primary,
+        ] = InferencesRole.primary,
     ) -> Optional[float]:
-        if not isinstance(dataset_role, DatasetRole):
-            dataset_role = DatasetRole.primary
-        dataset = info.context.model[dataset_role.value]
+        if not isinstance(inferences_role, InferencesRole):
+            inferences_role = InferencesRole.primary
+        inferences = info.context.model[inferences_role.value]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
         )
         data = get_data_quality_timeseries_data(
@@ -111,7 +112,7 @@ class Dimension(Node):
             metric,
             time_range,
             granularity,
-            dataset_role,
+            inferences_role,
         )
         return data[0].value if len(data) else None
 
@@ -139,18 +140,18 @@ class Dimension(Node):
         metric: DataQualityMetric,
         time_range: TimeRange,
         granularity: Granularity,
-        dataset_role: Annotated[
-            Optional[DatasetRole],
+        inferences_role: Annotated[
+            Optional[InferencesRole],
             strawberry.argument(
-                description="The dataset (primary or reference) to query",
+                description="The inferences (primary or reference) to query",
             ),
-        ] = DatasetRole.primary,
+        ] = InferencesRole.primary,
     ) -> DataQualityTimeSeries:
-        if not isinstance(dataset_role, DatasetRole):
-            dataset_role = DatasetRole.primary
-        dataset = info.context.model[dataset_role.value]
+        if not isinstance(inferences_role, InferencesRole):
+            inferences_role = InferencesRole.primary
+        inferences = info.context.model[inferences_role.value]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
             granularity,
         )
@@ -160,7 +161,7 @@ class Dimension(Node):
                 metric,
                 time_range,
                 granularity,
-                dataset_role,
+                inferences_role,
             )
         )
 
@@ -182,9 +183,9 @@ class Dimension(Node):
         model = info.context.model
         if model[REFERENCE].empty:
             return DriftTimeSeries(data=[])
-        dataset = model[PRIMARY]
+        inferences = model[PRIMARY]
         time_range, granularity = ensure_timeseries_parameters(
-            dataset,
+            inferences,
             time_range,
             granularity,
         )
@@ -202,7 +203,7 @@ class Dimension(Node):
         )
 
     @strawberry.field(
-        description="Returns the segments across both datasets and returns the counts per segment",
+        description="The segments across both inference sets and returns the counts per segment",
     )  # type: ignore
     def segments_comparison(
         self,
@@ -249,8 +250,8 @@ class Dimension(Node):
         if isinstance(binning_method, binning.IntervalBinning) and binning_method.bins is not None:
             all_bins = all_bins.union(binning_method.bins)
         for bin in all_bins:
-            values: Dict[ms.DatasetRole, Any] = defaultdict(lambda: None)
-            for role in ms.DatasetRole:
+            values: Dict[ms.InferencesRole, Any] = defaultdict(lambda: None)
+            for role in ms.InferencesRole:
                 if model[role].empty:
                     continue
                 try: