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

phoenix/server/api/types/Experiment.py

@@ -1,41 +1,155 @@
 from datetime import datetime
-from typing import ClassVar, Optional
+from typing import TYPE_CHECKING, Annotated, Optional
 
 import strawberry
 from sqlalchemy import func, select
-from sqlalchemy.orm import joinedload
-from sqlalchemy.sql.functions import coalesce
 from strawberry import UNSET, Private
-from strawberry.relay import Connection, Node, NodeID
+from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.scalars import JSON
 from strawberry.types import Info
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.exceptions import BadRequest
+from phoenix.server.api.input_types.ExperimentRunSort import (
+    ExperimentRunSort,
+    add_order_by_and_page_start_to_query,
+    get_experiment_run_cursor,
+)
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
+from phoenix.server.api.types.DatasetSplit import DatasetSplit
+from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.ExperimentAnnotationSummary import ExperimentAnnotationSummary
-from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
+from phoenix.server.api.types.ExperimentRun import ExperimentRun
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
+    Cursor,
     CursorString,
+    connection_from_cursors_and_nodes,
     connection_from_list,
 )
-from phoenix.server.api.types.Project import Project
 from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
 from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
 
+_DEFAULT_EXPERIMENT_RUNS_PAGE_SIZE = 50
+
+if TYPE_CHECKING:
+    from .Project import Project
+
 
 @strawberry.type
 class Experiment(Node):
-    _table: ClassVar[type[models.Base]] = models.Experiment
+    id: NodeID[int]
+    db_record: strawberry.Private[Optional[models.Experiment]] = None
     cached_sequence_number: Private[Optional[int]] = None
-    id_attr: NodeID[int]
-    name: str
-    project_name: Optional[str]
-    description: Optional[str]
-    metadata: JSON
-    created_at: datetime
-    updated_at: datetime
+
+    def __post_init__(self) -> None:
+        if self.db_record and self.id != self.db_record.id:
+            raise ValueError("Experiment ID mismatch")
+
+    @strawberry.field
+    async def name(
+        self,
+        info: Info[Context, None],
+    ) -> str:
+        if self.db_record:
+            val = self.db_record.name
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.name),
+            )
+        return val
+
+    @strawberry.field
+    async def project_name(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.project_name
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.project_name),
+            )
+        return val
+
+    @strawberry.field
+    async def description(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[str]:
+        if self.db_record:
+            val = self.db_record.description
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.description),
+            )
+        return val
+
+    @strawberry.field
+    async def repetitions(
+        self,
+        info: Info[Context, None],
+    ) -> int:
+        if self.db_record:
+            val = self.db_record.repetitions
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.repetitions),
+            )
+        return val
+
+    @strawberry.field
+    async def dataset_version_id(
+        self,
+        info: Info[Context, None],
+    ) -> GlobalID:
+        if self.db_record:
+            version_id = self.db_record.dataset_version_id
+        else:
+            version_id = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.dataset_version_id),
+            )
+        return GlobalID(DatasetVersion.__name__, str(version_id))
+
+    @strawberry.field
+    async def metadata(
+        self,
+        info: Info[Context, None],
+    ) -> JSON:
+        if self.db_record:
+            val = self.db_record.metadata_
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.metadata_),
+            )
+        return val
+
+    @strawberry.field
+    async def created_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.created_at
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.created_at),
+            )
+        return val
+
+    @strawberry.field
+    async def updated_at(
+        self,
+        info: Info[Context, None],
+    ) -> datetime:
+        if self.db_record:
+            val = self.db_record.updated_at
+        else:
+            val = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.updated_at),
+            )
+        return val
 
     @strawberry.field(
         description="Sequence number (1-based) of experiments belonging to the same dataset"
@@ -45,9 +159,9 @@ class Experiment(Node):
         info: Info[Context, None],
     ) -> int:
         if self.cached_sequence_number is None:
-            seq_num = await info.context.data_loaders.experiment_sequence_number.load(self.id_attr)
+            seq_num = await info.context.data_loaders.experiment_sequence_number.load(self.id)
             if seq_num is None:
-                raise ValueError(f"invalid experiment: id={self.id_attr}")
+                raise ValueError(f"invalid experiment: id={self.id}")
             self.cached_sequence_number = seq_num
         return self.cached_sequence_number
 
@@ -55,41 +169,68 @@ class Experiment(Node):
     async def runs(
         self,
         info: Info[Context, None],
-        first: Optional[int] = 50,
-        last: Optional[int] = UNSET,
+        first: Optional[int] = _DEFAULT_EXPERIMENT_RUNS_PAGE_SIZE,
         after: Optional[CursorString] = UNSET,
-        before: Optional[CursorString] = UNSET,
+        sort: Optional[ExperimentRunSort] = 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,
+        if first is not None and first <= 0:
+            raise BadRequest("first must be a positive integer if set")
+        page_size = first or _DEFAULT_EXPERIMENT_RUNS_PAGE_SIZE
+        experiment_runs_query = (
+            select(models.ExperimentRun)
+            .where(models.ExperimentRun.experiment_id == self.id)
+            .limit(page_size + 1)
+        )
+
+        after_experiment_run_rowid = None
+        after_sort_column_value = None
+        if after:
+            cursor = Cursor.from_string(after)
+            after_experiment_run_rowid = cursor.rowid
+            if cursor.sort_column is not None:
+                after_sort_column_value = cursor.sort_column.value
+
+        experiment_runs_query = add_order_by_and_page_start_to_query(
+            query=experiment_runs_query,
+            sort=sort,
+            experiment_rowid=self.id,
+            after_experiment_run_rowid=after_experiment_run_rowid,
+            after_sort_column_value=after_sort_column_value,
         )
-        experiment_id = self.id_attr
+
         async with info.context.db() as session:
-            runs = (
-                await session.scalars(
-                    select(models.ExperimentRun)
-                    .where(models.ExperimentRun.experiment_id == experiment_id)
-                    .order_by(models.ExperimentRun.id.desc())
-                    .options(
-                        joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id)
-                    )
-                )
-            ).all()
-        return connection_from_list([to_gql_experiment_run(run) for run in runs], args)
+            results = (await session.execute(experiment_runs_query)).all()
+
+        has_next_page = False
+        if len(results) > page_size:
+            results = results[:page_size]
+            has_next_page = True
+
+        cursors_and_nodes = []
+        for result in results:
+            run = result[0]
+            annotation_score = result[1] if len(result) > 1 else None
+            gql_run = ExperimentRun(id=run.id, db_record=run)
+            cursor = get_experiment_run_cursor(
+                run=run, annotation_score=annotation_score, sort=sort
+            )
+            cursors_and_nodes.append((cursor, gql_run))
+
+        return connection_from_cursors_and_nodes(
+            cursors_and_nodes=cursors_and_nodes,
+            has_previous_page=False,  # set to false since we are only doing forward pagination (https://relay.dev/graphql/connections.htm#sec-undefined.PageInfo.Fields) # noqa: E501
+            has_next_page=has_next_page,
+        )
 
     @strawberry.field
     async def run_count(self, info: Info[Context, None]) -> int:
-        experiment_id = self.id_attr
-        return await info.context.data_loaders.experiment_run_counts.load(experiment_id)
+        return await info.context.data_loaders.experiment_run_counts.load(self.id)
 
     @strawberry.field
     async def annotation_summaries(
         self, info: Info[Context, None]
     ) -> list[ExperimentAnnotationSummary]:
-        experiment_id = self.id_attr
+        experiment_id = self.id
         return [
             ExperimentAnnotationSummary(
                 annotation_name=summary.annotation_name,
@@ -106,40 +247,42 @@ class Experiment(Node):
 
     @strawberry.field
     async def error_rate(self, info: Info[Context, None]) -> Optional[float]:
-        return await info.context.data_loaders.experiment_error_rates.load(self.id_attr)
+        return await info.context.data_loaders.experiment_error_rates.load(self.id)
 
     @strawberry.field
     async def average_run_latency_ms(self, info: Info[Context, None]) -> Optional[float]:
-        latency_seconds = await info.context.data_loaders.average_experiment_run_latency.load(
-            self.id_attr
-        )
-        return latency_seconds * 1000 if latency_seconds is not None else None
+        latency_ms = await info.context.data_loaders.average_experiment_run_latency.load(self.id)
+        return latency_ms
 
     @strawberry.field
-    async def project(self, info: Info[Context, None]) -> Optional[Project]:
-        if self.project_name is None:
+    async def project(
+        self, info: Info[Context, None]
+    ) -> Optional[Annotated["Project", strawberry.lazy(".Project")]]:
+        if self.db_record:
+            project_name = self.db_record.project_name
+        else:
+            project_name = await info.context.data_loaders.experiment_fields.load(
+                (self.id, models.Experiment.project_name),
+            )
+
+        if project_name is None:
             return None
 
-        db_project = await info.context.data_loaders.project_by_name.load(self.project_name)
+        db_project = await info.context.data_loaders.project_by_name.load(project_name)
 
         if db_project is None:
             return None
+        from .Project import Project
 
-        return Project(
-            project_rowid=db_project.id,
-            db_project=db_project,
-        )
+        return Project(id=db_project.id, db_record=db_project)
 
     @strawberry.field
     def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
-        return info.context.last_updated_at.get(self._table, self.id_attr)
+        return info.context.last_updated_at.get(models.Experiment, self.id)
 
     @strawberry.field
     async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
-        experiment_id = self.id_attr
-        summary = await info.context.data_loaders.span_cost_summary_by_experiment.load(
-            experiment_id
-        )
+        summary = await info.context.data_loaders.span_cost_summary_by_experiment.load(self.id)
         return SpanCostSummary(
             prompt=CostBreakdown(
                 tokens=summary.prompt.tokens,
@@ -159,21 +302,19 @@ class Experiment(Node):
     async def cost_detail_summary_entries(
         self, info: Info[Context, None]
     ) -> list[SpanCostDetailSummaryEntry]:
-        experiment_id = self.id_attr
-
         stmt = (
             select(
                 models.SpanCostDetail.token_type,
                 models.SpanCostDetail.is_prompt,
-                coalesce(func.sum(models.SpanCostDetail.cost), 0).label("cost"),
-                coalesce(func.sum(models.SpanCostDetail.tokens), 0).label("tokens"),
+                func.sum(models.SpanCostDetail.cost).label("cost"),
+                func.sum(models.SpanCostDetail.tokens).label("tokens"),
             )
             .select_from(models.SpanCostDetail)
             .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
             .join(models.Span, models.SpanCost.span_rowid == models.Span.id)
             .join(models.Trace, models.Span.trace_rowid == models.Trace.id)
             .join(models.ExperimentRun, models.ExperimentRun.trace_id == models.Trace.trace_id)
-            .where(models.ExperimentRun.experiment_id == experiment_id)
+            .where(models.ExperimentRun.experiment_id == self.id)
             .group_by(models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
         )
 
@@ -188,6 +329,17 @@ class Experiment(Node):
                 async for token_type, is_prompt, cost, tokens in data
             ]
 
+    @strawberry.field
+    async def dataset_splits(
+        self,
+        info: Info[Context, None],
+    ) -> Connection[DatasetSplit]:
+        """Returns the dataset splits associated with this experiment."""
+        splits = await info.context.data_loaders.experiment_dataset_splits.load(self.id)
+        return connection_from_list(
+            [DatasetSplit(id=split.id, db_record=split) for split in splits], ConnectionArgs()
+        )
+
 
 def to_gql_experiment(
     experiment: models.Experiment,
@@ -197,12 +349,7 @@ def to_gql_experiment(
     Converts an ORM experiment to a GraphQL Experiment.
     """
     return Experiment(
+        id=experiment.id,
+        db_record=experiment,
         cached_sequence_number=sequence_number,
-        id_attr=experiment.id,
-        name=experiment.name,
-        project_name=experiment.project_name,
-        description=experiment.description,
-        metadata=experiment.metadata_,
-        created_at=experiment.created_at,
-        updated_at=experiment.updated_at,
     )

phoenix/server/api/types/ExperimentComparison.py

@@ -1,18 +1,12 @@
 import strawberry
-from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.relay import Node, NodeID
 
 from phoenix.server.api.types.DatasetExample import DatasetExample
-from phoenix.server.api.types.ExperimentRun import ExperimentRun
-
-
-@strawberry.type
-class RunComparisonItem:
-    experiment_id: GlobalID
-    runs: list[ExperimentRun]
+from phoenix.server.api.types.ExperimentRepeatedRunGroup import ExperimentRepeatedRunGroup
 
 
 @strawberry.type
 class ExperimentComparison(Node):
     id_attr: NodeID[int]
     example: DatasetExample
-    run_comparison_items: list[RunComparisonItem]
+    repeated_run_groups: list[ExperimentRepeatedRunGroup]

phoenix/server/api/types/ExperimentRepeatedRunGroup.py

@@ -0,0 +1,155 @@
+import re
+from base64 import b64decode
+from typing import Optional
+
+import strawberry
+from sqlalchemy import func, select
+from strawberry.relay import GlobalID, Node
+from strawberry.types import Info
+from typing_extensions import Self, TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.CostBreakdown import CostBreakdown
+from phoenix.server.api.types.ExperimentRepeatedRunGroupAnnotationSummary import (
+    ExperimentRepeatedRunGroupAnnotationSummary,
+)
+from phoenix.server.api.types.ExperimentRun import ExperimentRun
+from phoenix.server.api.types.SpanCostDetailSummaryEntry import SpanCostDetailSummaryEntry
+from phoenix.server.api.types.SpanCostSummary import SpanCostSummary
+
+ExperimentRowId: TypeAlias = int
+DatasetExampleRowId: TypeAlias = int
+
+
+@strawberry.type
+class ExperimentRepeatedRunGroup(Node):
+    experiment_rowid: strawberry.Private[ExperimentRowId]
+    dataset_example_rowid: strawberry.Private[DatasetExampleRowId]
+    cached_runs: strawberry.Private[Optional[list[ExperimentRun]]] = None
+
+    @strawberry.field
+    async def runs(self, info: Info[Context, None]) -> list[ExperimentRun]:
+        if self.cached_runs is not None:
+            return self.cached_runs
+        runs = await info.context.data_loaders.experiment_runs_by_experiment_and_example.load(
+            (self.experiment_rowid, self.dataset_example_rowid)
+        )
+        return [ExperimentRun(id=run.id, db_record=run) for run in runs]
+
+    @classmethod
+    def resolve_id(
+        cls,
+        root: Self,
+        *,
+        info: Info,
+    ) -> str:
+        return (
+            f"experiment_id={root.experiment_rowid}:dataset_example_id={root.dataset_example_rowid}"
+        )
+
+    @strawberry.field
+    def experiment_id(self) -> strawberry.ID:
+        from phoenix.server.api.types.Experiment import Experiment
+
+        return strawberry.ID(str(GlobalID(Experiment.__name__, str(self.experiment_rowid))))
+
+    @strawberry.field
+    async def average_latency_ms(self, info: Info[Context, None]) -> Optional[float]:
+        return await info.context.data_loaders.average_experiment_repeated_run_group_latency.load(
+            (self.experiment_rowid, self.dataset_example_rowid)
+        )
+
+    @strawberry.field
+    async def cost_summary(self, info: Info[Context, None]) -> SpanCostSummary:
+        experiment_id = self.experiment_rowid
+        example_id = self.dataset_example_rowid
+        summary = (
+            await info.context.data_loaders.span_cost_summary_by_experiment_repeated_run_group.load(
+                (experiment_id, example_id)
+            )
+        )
+        return SpanCostSummary(
+            prompt=CostBreakdown(
+                tokens=summary.prompt.tokens,
+                cost=summary.prompt.cost,
+            ),
+            completion=CostBreakdown(
+                tokens=summary.completion.tokens,
+                cost=summary.completion.cost,
+            ),
+            total=CostBreakdown(
+                tokens=summary.total.tokens,
+                cost=summary.total.cost,
+            ),
+        )
+
+    @strawberry.field
+    async def cost_detail_summary_entries(
+        self, info: Info[Context, None]
+    ) -> list[SpanCostDetailSummaryEntry]:
+        experiment_id = self.experiment_rowid
+        example_id = self.dataset_example_rowid
+        stmt = (
+            select(
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                func.sum(models.SpanCostDetail.cost).label("cost"),
+                func.sum(models.SpanCostDetail.tokens).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .join(models.Trace, models.SpanCost.trace_rowid == models.Trace.id)
+            .join(models.ExperimentRun, models.ExperimentRun.trace_id == models.Trace.trace_id)
+            .where(models.ExperimentRun.experiment_id == experiment_id)
+            .where(models.ExperimentRun.dataset_example_id == example_id)
+            .group_by(models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
+        )
+
+        async with info.context.db() as session:
+            data = await session.stream(stmt)
+            return [
+                SpanCostDetailSummaryEntry(
+                    token_type=token_type,
+                    is_prompt=is_prompt,
+                    value=CostBreakdown(tokens=tokens, cost=cost),
+                )
+                async for token_type, is_prompt, cost, tokens in data
+            ]
+
+    @strawberry.field
+    async def annotation_summaries(
+        self,
+        info: Info[Context, None],
+    ) -> list[ExperimentRepeatedRunGroupAnnotationSummary]:
+        loader = info.context.data_loaders.experiment_repeated_run_group_annotation_summaries
+        summaries = await loader.load((self.experiment_rowid, self.dataset_example_rowid))
+        return [
+            ExperimentRepeatedRunGroupAnnotationSummary(
+                annotation_name=summary.annotation_name,
+                mean_score=summary.mean_score,
+            )
+            for summary in summaries
+        ]
+
+
+_EXPERIMENT_REPEATED_RUN_GROUP_NODE_ID_PATTERN = re.compile(
+    r"ExperimentRepeatedRunGroup:experiment_id=(\d+):dataset_example_id=(\d+)"
+)
+
+
+def parse_experiment_repeated_run_group_node_id(
+    node_id: str,
+) -> tuple[ExperimentRowId, DatasetExampleRowId]:
+    decoded_node_id = _base64_decode(node_id)
+    match = re.match(_EXPERIMENT_REPEATED_RUN_GROUP_NODE_ID_PATTERN, decoded_node_id)
+    if not match:
+        raise ValueError(f"Invalid node ID format: {node_id}")
+
+    experiment_id = int(match.group(1))
+    dataset_example_id = int(match.group(2))
+    return experiment_id, dataset_example_id
+
+
+def _base64_decode(string: str) -> str:
+    return b64decode(string.encode()).decode()

phoenix/server/api/types/ExperimentRepeatedRunGroupAnnotationSummary.py

@@ -0,0 +1,9 @@
+from typing import Optional
+
+import strawberry
+
+
+@strawberry.type
+class ExperimentRepeatedRunGroupAnnotationSummary:
+    annotation_name: str
+    mean_score: Optional[float]