arize-phoenix 3.16.1__py3-none-any.whl → 7.7.1__py3-none-any.whl

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

@@ -1,9 +1,10 @@
 from collections import defaultdict
-from typing import Any, Dict, List, Optional
+from typing import Any, 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
 
@@ -122,7 +123,7 @@ class Dimension(Node):
             " Missing values are excluded. Non-categorical dimensions return an empty list."
         )
     )  # type: ignore  # https://github.com/strawberry-graphql/strawberry/issues/1929
-    def categories(self) -> List[str]:
+    def categories(self) -> list[str]:
         return list(self.dimension.categories)
 
     @strawberry.field(
@@ -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:

phoenix/server/api/types/DocumentEvaluationSummary.py

@@ -1,6 +1,7 @@
 import math
+from collections.abc import Iterable
 from functools import cached_property
-from typing import Any, Dict, Iterable, Optional, Tuple
+from typing import Any, Optional
 
 import pandas as pd
 import strawberry
@@ -24,8 +25,8 @@ class DocumentEvaluationSummary:
     ) -> None:
         self.evaluation_name = evaluation_name
         self.metrics_collection = pd.Series(metrics_collection, dtype=object)
-        self._cached_average_ndcg_results: Dict[Optional[int], Tuple[float, int]] = {}
-        self._cached_average_precision_results: Dict[Optional[int], Tuple[float, int]] = {}
+        self._cached_average_ndcg_results: dict[Optional[int], tuple[float, int]] = {}
+        self._cached_average_precision_results: dict[Optional[int], tuple[float, int]] = {}
 
     @strawberry.field
     def average_ndcg(self, k: Optional[int] = UNSET) -> Optional[float]:
@@ -67,7 +68,7 @@ class DocumentEvaluationSummary:
         _, count = self._average_hit
         return count
 
-    def _average_ndcg(self, k: Optional[int] = None) -> Tuple[float, int]:
+    def _average_ndcg(self, k: Optional[int] = None) -> tuple[float, int]:
         if (result := self._cached_average_ndcg_results.get(k)) is not None:
             return result
         values = self.metrics_collection.apply(lambda m: m.ndcg(k))
@@ -75,20 +76,20 @@ class DocumentEvaluationSummary:
         self._cached_average_ndcg_results[k] = result
         return result
 
-    def _average_precision(self, k: Optional[int] = None) -> Tuple[float, int]:
+    def _average_precision(self, k: Optional[int] = None) -> tuple[float, int]:
         if (result := self._cached_average_precision_results.get(k)) is not None:
             return result
         values = self.metrics_collection.apply(lambda m: m.precision(k))
         result = (values.mean(), values.count())
-        self._cached_average_ndcg_results[k] = result
+        self._cached_average_precision_results[k] = result
         return result
 
     @cached_property
-    def _average_reciprocal_rank(self) -> Tuple[float, int]:
+    def _average_reciprocal_rank(self) -> tuple[float, int]:
         values = self.metrics_collection.apply(lambda m: m.reciprocal_rank())
         return values.mean(), values.count()
 
     @cached_property
-    def _average_hit(self) -> Tuple[float, int]:
+    def _average_hit(self) -> tuple[float, int]:
         values = self.metrics_collection.apply(lambda m: m.hit())
         return values.mean(), values.count()

phoenix/server/api/types/EmbeddingDimension.py

@@ -1,13 +1,15 @@
 from collections import defaultdict
+from collections.abc import Iterable, Iterator
 from datetime import timedelta
 from itertools import chain, repeat
-from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union, cast
+from typing import Any, Optional, Union, cast
 
 import numpy as np
 import numpy.typing as npt
 import pandas as pd
 import strawberry
 from strawberry import UNSET
+from strawberry.relay import GlobalID, Node, NodeID
 from strawberry.scalars import ID
 from strawberry.types import Info
 from typing_extensions import Annotated
@@ -22,7 +24,7 @@ from phoenix.core.model_schema import (
     PRIMARY,
     PROMPT,
     REFERENCE,
-    Dataset,
+    Inferences,
 )
 from phoenix.metrics.timeseries import row_interval_from_sorted_time_index
 from phoenix.pointcloud.clustering import Hdbscan
@@ -31,7 +33,7 @@ from phoenix.pointcloud.projectors import Umap
 from phoenix.server.api.context import Context
 from phoenix.server.api.input_types.TimeRange import TimeRange
 from phoenix.server.api.types.Cluster import to_gql_clusters
-from phoenix.server.api.types.DatasetRole import AncillaryDatasetRole, DatasetRole
+from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
 from phoenix.server.api.types.VectorDriftMetricEnum import VectorDriftMetric
 
 from ..input_types.Granularity import Granularity
@@ -39,7 +41,6 @@ from .DataQualityMetric import DataQualityMetric
 from .EmbeddingMetadata import EmbeddingMetadata
 from .Event import create_event_id, unpack_event_id
 from .EventMetadata import EventMetadata
-from .node import GlobalID, Node
 from .Retrieval import Retrieval
 from .TimeSeries import (
     DataQualityTimeSeries,
@@ -70,6 +71,7 @@ CORPUS = "CORPUS"
 class EmbeddingDimension(Node):
     """A embedding dimension of a model. Represents unstructured data"""
 
+    id_attr: NodeID[int]
     name: str
     dimension: strawberry.Private[ms.EmbeddingDimension]
 
@@ -155,16 +157,16 @@ class EmbeddingDimension(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",
             ),
-        ] = 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
+        dataset = info.context.model[inferences_role.value]
         time_range, granularity = ensure_timeseries_parameters(
             dataset,
             time_range,
@@ -176,7 +178,7 @@ class EmbeddingDimension(Node):
                 metric,
                 time_range,
                 granularity,
-                dataset_role,
+                inferences_role,
             )
         )
 
@@ -312,18 +314,18 @@ class EmbeddingDimension(Node):
         ] = DEFAULT_CLUSTER_SELECTION_EPSILON,
     ) -> UMAPPoints:
         model = info.context.model
-        data: Dict[ID, npt.NDArray[np.float64]] = {}
-        retrievals: List[Tuple[ID, Any, Any]] = []
-        for dataset in model[Dataset]:
-            dataset_id = dataset.role
-            row_id_start, row_id_stop = 0, len(dataset)
-            if dataset_id is PRIMARY:
+        data: dict[ID, npt.NDArray[np.float64]] = {}
+        retrievals: list[tuple[ID, Any, Any]] = []
+        for inferences in model[Inferences]:
+            inferences_id = inferences.role
+            row_id_start, row_id_stop = 0, len(inferences)
+            if inferences_id is PRIMARY:
                 row_id_start, row_id_stop = row_interval_from_sorted_time_index(
-                    time_index=cast(pd.DatetimeIndex, dataset.index),
+                    time_index=cast(pd.DatetimeIndex, inferences.index),
                     time_start=time_range.start,
                     time_stop=time_range.end,
                 )
-            vector_column = self.dimension[dataset_id]
+            vector_column = self.dimension[inferences_id]
             samples_collected = 0
             for row_id in _row_indices(
                 row_id_start,
@@ -337,7 +339,7 @@ class EmbeddingDimension(Node):
                 # of dunder method __len__.
                 if not hasattr(embedding_vector, "__len__"):
                     continue
-                event_id = create_event_id(row_id, dataset_id)
+                event_id = create_event_id(row_id, inferences_id)
                 data[event_id] = embedding_vector
                 samples_collected += 1
                 if isinstance(
@@ -347,23 +349,23 @@ class EmbeddingDimension(Node):
                     retrievals.append(
                         (
                             event_id,
-                            self.dimension.context_retrieval_ids(dataset).iloc[row_id],
-                            self.dimension.context_retrieval_scores(dataset).iloc[row_id],
+                            self.dimension.context_retrieval_ids(inferences).iloc[row_id],
+                            self.dimension.context_retrieval_scores(inferences).iloc[row_id],
                         )
                     )
 
-        context_retrievals: List[Retrieval] = []
+        context_retrievals: list[Retrieval] = []
         if isinstance(
             self.dimension,
             ms.RetrievalEmbeddingDimension,
         ) and (corpus := info.context.corpus):
-            corpus_dataset = corpus[PRIMARY]
-            for row_id, document_embedding_vector in enumerate(corpus_dataset[PROMPT]):
+            corpus_inferences = corpus[PRIMARY]
+            for row_id, document_embedding_vector in enumerate(corpus_inferences[PROMPT]):
                 if not hasattr(document_embedding_vector, "__len__"):
                     continue
-                event_id = create_event_id(row_id, AncillaryDatasetRole.corpus)
+                event_id = create_event_id(row_id, AncillaryInferencesRole.corpus)
                 data[event_id] = document_embedding_vector
-            corpus_primary_key = corpus_dataset.primary_key
+            corpus_primary_key = corpus_inferences.primary_key
             for event_id, retrieval_ids, retrieval_scores in retrievals:
                 if not isinstance(retrieval_ids, Iterable):
                     continue
@@ -385,7 +387,7 @@ class EmbeddingDimension(Node):
                         )
                     except KeyError:
                         continue
-                    document_embedding_vector = corpus_dataset[PROMPT].iloc[document_row_id]
+                    document_embedding_vector = corpus_inferences[PROMPT].iloc[document_row_id]
                     if not hasattr(document_embedding_vector, "__len__"):
                         continue
                     context_retrievals.append(
@@ -393,7 +395,7 @@ class EmbeddingDimension(Node):
                             query_id=event_id,
                             document_id=create_event_id(
                                 document_row_id,
-                                AncillaryDatasetRole.corpus,
+                                AncillaryInferencesRole.corpus,
                             ),
                             relevance=document_score,
                         )
@@ -413,48 +415,53 @@ class EmbeddingDimension(Node):
             ),
         ).generate(data, n_components=n_components)
 
-        points: Dict[Union[DatasetRole, AncillaryDatasetRole], List[UMAPPoint]] = defaultdict(list)
+        points: dict[Union[InferencesRole, AncillaryInferencesRole], list[UMAPPoint]] = defaultdict(
+            list
+        )
         for event_id, vector in vectors.items():
-            row_id, dataset_role = unpack_event_id(event_id)
-            if isinstance(dataset_role, DatasetRole):
-                dataset = model[dataset_role.value]
+            row_id, inferences_role = unpack_event_id(event_id)
+            if isinstance(inferences_role, InferencesRole):
+                dataset = model[inferences_role.value]
                 embedding_metadata = EmbeddingMetadata(
-                    prediction_id=dataset[PREDICTION_ID][row_id],
-                    link_to_data=dataset[self.dimension.link_to_data][row_id],
-                    raw_data=dataset[self.dimension.raw_data][row_id],
+                    prediction_id=dataset[PREDICTION_ID].iloc[row_id],
+                    link_to_data=dataset[self.dimension.link_to_data].iloc[row_id],
+                    raw_data=dataset[self.dimension.raw_data].iloc[row_id],
                 )
             elif (corpus := info.context.corpus) is not None:
                 dataset = corpus[PRIMARY]
                 dimension = cast(ms.EmbeddingDimension, corpus[PROMPT])
                 embedding_metadata = EmbeddingMetadata(
-                    prediction_id=dataset[PREDICTION_ID][row_id],
-                    link_to_data=dataset[dimension.link_to_data][row_id],
-                    raw_data=dataset[dimension.raw_data][row_id],
+                    prediction_id=dataset[PREDICTION_ID].iloc[row_id],
+                    link_to_data=dataset[dimension.link_to_data].iloc[row_id],
+                    raw_data=dataset[dimension.raw_data].iloc[row_id],
                 )
             else:
                 continue
-            points[dataset_role].append(
+            points[inferences_role].append(
                 UMAPPoint(
-                    id=GlobalID(f"{type(self).__name__}:{str(dataset_role)}", row_id),
+                    id=GlobalID(
+                        type_name=f"{type(self).__name__}:{str(inferences_role)}",
+                        node_id=str(row_id),
+                    ),
                     event_id=event_id,
                     coordinates=to_gql_coordinates(vector),
                     event_metadata=EventMetadata(
-                        prediction_label=dataset[PREDICTION_LABEL][row_id],
-                        prediction_score=dataset[PREDICTION_SCORE][row_id],
-                        actual_label=dataset[ACTUAL_LABEL][row_id],
-                        actual_score=dataset[ACTUAL_SCORE][row_id],
+                        prediction_label=dataset[PREDICTION_LABEL].iloc[row_id],
+                        prediction_score=dataset[PREDICTION_SCORE].iloc[row_id],
+                        actual_label=dataset[ACTUAL_LABEL].iloc[row_id],
+                        actual_score=dataset[ACTUAL_SCORE].iloc[row_id],
                     ),
                     embedding_metadata=embedding_metadata,
                 )
             )
 
         return UMAPPoints(
-            data=points[DatasetRole.primary],
-            reference_data=points[DatasetRole.reference],
+            data=points[InferencesRole.primary],
+            reference_data=points[InferencesRole.reference],
             clusters=to_gql_clusters(
                 clustered_events=clustered_events,
             ),
-            corpus_data=points[AncillaryDatasetRole.corpus],
+            corpus_data=points[AncillaryInferencesRole.corpus],
             context_retrievals=context_retrievals,
         )
 
@@ -470,7 +477,7 @@ def _row_indices(
         return
     shuffled_indices = np.arange(start, stop)
     np.random.shuffle(shuffled_indices)
-    yield from shuffled_indices
+    yield from shuffled_indices  # type: ignore[misc,unused-ignore]
 
 
 def to_gql_embedding_dimension(