arize-phoenix 11.32.1__py3-none-any.whl → 11.33.0__py3-none-any.whl

phoenix/server/api/dataloaders/experiment_repetition_counts.py

@@ -0,0 +1,39 @@
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentID: TypeAlias = int
+RepetitionCount: TypeAlias = int
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = RepetitionCount
+
+
+class ExperimentRepetitionCountsDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: DbSessionFactory,
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        experiment_ids = keys
+        repetition_counts_query = (
+            select(
+                models.ExperimentRun.experiment_id,
+                func.max(models.ExperimentRun.repetition_number).label("repetition_count"),
+            )
+            .group_by(models.ExperimentRun.experiment_id)
+            .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+        )
+        async with self._db() as session:
+            repetition_counts = {
+                experiment_id: repetition_count
+                for experiment_id, repetition_count in await session.execute(
+                    repetition_counts_query
+                )
+            }
+        return [repetition_counts.get(experiment_id, 0) for experiment_id in keys]

phoenix/server/api/dataloaders/span_cost_summary_by_experiment_repeated_run_group.py

@@ -0,0 +1,64 @@
+from collections import defaultdict
+
+from sqlalchemy import func, select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.types import CostBreakdown, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+Key: TypeAlias = tuple[ExperimentId, DatasetExampleId]
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByExperimentRepeatedRunGroupDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        stmt = (
+            select(
+                models.ExperimentRun.experiment_id,
+                models.ExperimentRun.dataset_example_id,
+                func.sum(models.SpanCost.prompt_cost).label("prompt_cost"),
+                func.sum(models.SpanCost.completion_cost).label("completion_cost"),
+                func.sum(models.SpanCost.total_cost).label("total_cost"),
+                func.sum(models.SpanCost.prompt_tokens).label("prompt_tokens"),
+                func.sum(models.SpanCost.completion_tokens).label("completion_tokens"),
+                func.sum(models.SpanCost.total_tokens).label("total_tokens"),
+            )
+            .select_from(models.ExperimentRun)
+            .join(models.Trace, models.ExperimentRun.trace_id == models.Trace.trace_id)
+            .join(models.SpanCost, models.SpanCost.trace_rowid == models.Trace.id)
+            .where(
+                tuple_(
+                    models.ExperimentRun.experiment_id, models.ExperimentRun.dataset_example_id
+                ).in_(set(keys))
+            )
+            .group_by(models.ExperimentRun.experiment_id, models.ExperimentRun.dataset_example_id)
+        )
+
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                experiment_id,
+                dataset_example_id,
+                prompt_cost,
+                completion_cost,
+                total_cost,
+                prompt_tokens,
+                completion_tokens,
+                total_tokens,
+            ) in data:
+                summary = SpanCostSummary(
+                    prompt=CostBreakdown(tokens=prompt_tokens, cost=prompt_cost),
+                    completion=CostBreakdown(tokens=completion_tokens, cost=completion_cost),
+                    total=CostBreakdown(tokens=total_tokens, cost=total_cost),
+                )
+                results[(experiment_id, dataset_example_id)] = summary
+        return [results.get(key, SpanCostSummary()) for key in keys]

phoenix/server/api/helpers/playground_clients.py

@@ -1669,7 +1669,11 @@ class AnthropicReasoningStreamingClient(AnthropicStreamingClient):
     provider_key=GenerativeProviderKey.GOOGLE,
     model_names=[
         PROVIDER_DEFAULT,
+        "gemini-2.5-flash",
+        "gemini-2.5-flash-lite",
+        "gemini-2.5-pro",
         "gemini-2.5-pro-preview-03-25",
+        "gemini-2.0-flash",
         "gemini-2.0-flash-lite",
         "gemini-2.0-flash-001",
         "gemini-2.0-flash-thinking-exp-01-21",

phoenix/server/api/mutations/prompt_label_mutations.py

@@ -10,12 +10,10 @@ from strawberry.relay import GlobalID
 from strawberry.types import Info
 
 from phoenix.db import models
-from phoenix.db.types.identifier import Identifier as IdentifierModel
 from phoenix.server.api.auth import IsLocked, IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import Conflict, NotFound
 from phoenix.server.api.queries import Query
-from phoenix.server.api.types.Identifier import Identifier
 from phoenix.server.api.types.node import from_global_id_with_expected_type
 from phoenix.server.api.types.Prompt import Prompt
 from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
@@ -23,37 +21,49 @@ from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_labe
 
 @strawberry.input
 class CreatePromptLabelInput:
-    name: Identifier
+    name: str
     description: Optional[str] = None
+    color: str
 
 
 @strawberry.input
 class PatchPromptLabelInput:
     prompt_label_id: GlobalID
-    name: Optional[Identifier] = None
+    name: Optional[str] = None
     description: Optional[str] = None
 
 
 @strawberry.input
-class DeletePromptLabelInput:
-    prompt_label_id: GlobalID
+class DeletePromptLabelsInput:
+    prompt_label_ids: list[GlobalID]
 
 
 @strawberry.input
-class SetPromptLabelInput:
+class SetPromptLabelsInput:
     prompt_id: GlobalID
-    prompt_label_id: GlobalID
+    prompt_label_ids: list[GlobalID]
 
 
 @strawberry.input
-class UnsetPromptLabelInput:
+class UnsetPromptLabelsInput:
     prompt_id: GlobalID
-    prompt_label_id: GlobalID
+    prompt_label_ids: list[GlobalID]
 
 
 @strawberry.type
 class PromptLabelMutationPayload:
-    prompt_label: Optional["PromptLabel"]
+    prompt_labels: list["PromptLabel"]
+    query: "Query"
+
+
+@strawberry.type
+class PromptLabelDeleteMutationPayload:
+    deleted_prompt_label_ids: list["GlobalID"]
+    query: "Query"
+
+
+@strawberry.type
+class PromptLabelAssociationMutationPayload:
     query: "Query"
 
 
@@ -64,17 +74,18 @@ class PromptLabelMutationMixin:
         self, info: Info[Context, None], input: CreatePromptLabelInput
     ) -> PromptLabelMutationPayload:
         async with info.context.db() as session:
-            name = IdentifierModel.model_validate(str(input.name))
-            label_orm = models.PromptLabel(name=name, description=input.description)
+            label_orm = models.PromptLabel(
+                name=input.name, description=input.description, color=input.color
+            )
             session.add(label_orm)
 
             try:
                 await session.commit()
             except (PostgreSQLIntegrityError, SQLiteIntegrityError):
-                raise Conflict(f"A prompt label named '{name}' already exists.")
+                raise Conflict(f"A prompt label named '{input.name}' already exists.")
 
             return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm),
+                prompt_labels=[to_gql_prompt_label(label_orm)],
                 query=Query(),
             )
 
@@ -82,7 +93,6 @@ class PromptLabelMutationMixin:
     async def patch_prompt_label(
         self, info: Info[Context, None], input: PatchPromptLabelInput
     ) -> PromptLabelMutationPayload:
-        validated_name = IdentifierModel.model_validate(str(input.name)) if input.name else None
         async with info.context.db() as session:
             label_id = from_global_id_with_expected_type(
                 input.prompt_label_id, PromptLabel.__name__
@@ -92,8 +102,8 @@ class PromptLabelMutationMixin:
             if not label_orm:
                 raise NotFound(f"PromptLabel with ID {input.prompt_label_id} not found")
 
-            if validated_name is not None:
-                label_orm.name = validated_name.root
+            if input.name is not None:
+                label_orm.name = input.name
             if input.description is not None:
                 label_orm.description = input.description
 
@@ -103,46 +113,48 @@ class PromptLabelMutationMixin:
                 raise Conflict("Error patching PromptLabel. Possibly a name conflict?")
 
             return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm),
+                prompt_labels=[to_gql_prompt_label(label_orm)],
                 query=Query(),
             )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
-    async def delete_prompt_label(
-        self, info: Info[Context, None], input: DeletePromptLabelInput
-    ) -> PromptLabelMutationPayload:
+    async def delete_prompt_labels(
+        self, info: Info[Context, None], input: DeletePromptLabelsInput
+    ) -> PromptLabelDeleteMutationPayload:
         """
         Deletes a PromptLabel (and any crosswalk references).
         """
         async with info.context.db() as session:
-            label_id = from_global_id_with_expected_type(
-                input.prompt_label_id, PromptLabel.__name__
-            )
-            stmt = delete(models.PromptLabel).where(models.PromptLabel.id == label_id)
-            result = await session.execute(stmt)
-
-            if result.rowcount == 0:
-                raise NotFound(f"PromptLabel with ID {input.prompt_label_id} not found")
+            label_ids = [
+                from_global_id_with_expected_type(prompt_label_id, PromptLabel.__name__)
+                for prompt_label_id in input.prompt_label_ids
+            ]
+            stmt = delete(models.PromptLabel).where(models.PromptLabel.id.in_(label_ids))
+            await session.execute(stmt)
 
             await session.commit()
 
-            return PromptLabelMutationPayload(
-                prompt_label=None,
+            return PromptLabelDeleteMutationPayload(
+                deleted_prompt_label_ids=input.prompt_label_ids,
                 query=Query(),
             )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly, IsLocked])  # type: ignore
-    async def set_prompt_label(
-        self, info: Info[Context, None], input: SetPromptLabelInput
-    ) -> PromptLabelMutationPayload:
+    async def set_prompt_labels(
+        self, info: Info[Context, None], input: SetPromptLabelsInput
+    ) -> PromptLabelAssociationMutationPayload:
         async with info.context.db() as session:
             prompt_id = from_global_id_with_expected_type(input.prompt_id, Prompt.__name__)
-            label_id = from_global_id_with_expected_type(
-                input.prompt_label_id, PromptLabel.__name__
-            )
+            label_ids = [
+                from_global_id_with_expected_type(prompt_label_id, PromptLabel.__name__)
+                for prompt_label_id in input.prompt_label_ids
+            ]
 
-            crosswalk = models.PromptPromptLabel(prompt_id=prompt_id, prompt_label_id=label_id)
-            session.add(crosswalk)
+            crosswalk_items = [
+                models.PromptPromptLabel(prompt_id=prompt_id, prompt_label_id=label_id)
+                for label_id in label_ids
+            ]
+            session.add_all(crosswalk_items)
 
             try:
                 await session.commit()
@@ -152,41 +164,38 @@ class PromptLabelMutationMixin:
                 # - Foreign key violation => prompt_id or label_id doesn't exist
                 raise Conflict("Failed to associate PromptLabel with Prompt.") from e
 
-            label_orm = await session.get(models.PromptLabel, label_id)
-            if not label_orm:
-                raise NotFound(f"PromptLabel with ID {input.prompt_label_id} not found")
-
-            return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm),
+            return PromptLabelAssociationMutationPayload(
                 query=Query(),
             )
 
     @strawberry.mutation(permission_classes=[IsNotReadOnly])  # type: ignore
-    async def unset_prompt_label(
-        self, info: Info[Context, None], input: UnsetPromptLabelInput
-    ) -> PromptLabelMutationPayload:
+    async def unset_prompt_labels(
+        self, info: Info[Context, None], input: UnsetPromptLabelsInput
+    ) -> PromptLabelAssociationMutationPayload:
         """
         Unsets a PromptLabel from a Prompt by removing the row in the crosswalk.
         """
         async with info.context.db() as session:
             prompt_id = from_global_id_with_expected_type(input.prompt_id, Prompt.__name__)
-            label_id = from_global_id_with_expected_type(
-                input.prompt_label_id, PromptLabel.__name__
-            )
+            label_ids = [
+                from_global_id_with_expected_type(prompt_label_id, PromptLabel.__name__)
+                for prompt_label_id in input.prompt_label_ids
+            ]
 
             stmt = delete(models.PromptPromptLabel).where(
                 (models.PromptPromptLabel.prompt_id == prompt_id)
-                & (models.PromptPromptLabel.prompt_label_id == label_id)
+                & (models.PromptPromptLabel.prompt_label_id.in_(label_ids))
             )
             result = await session.execute(stmt)
 
-            if result.rowcount == 0:
-                raise NotFound(f"No association between prompt={prompt_id} and label={label_id}.")
+            if result.rowcount != len(label_ids):
+                label_ids_str = ", ".join(str(i) for i in label_ids)
+                raise NotFound(
+                    f"No association between prompt={prompt_id} and labels={label_ids_str}."
+                )
 
             await session.commit()
 
-            label_orm = await session.get(models.PromptLabel, label_id)
-            return PromptLabelMutationPayload(
-                prompt_label=to_gql_prompt_label(label_orm) if label_orm else None,
+            return PromptLabelAssociationMutationPayload(
                 query=Query(),
             )

phoenix/server/api/queries.py

@@ -56,15 +56,25 @@ from phoenix.server.api.types.EmbeddingDimension import (
     to_gql_embedding_dimension,
 )
 from phoenix.server.api.types.Event import create_event_id, unpack_event_id
-from phoenix.server.api.types.Experiment import Experiment
-from phoenix.server.api.types.ExperimentComparison import ExperimentComparison, RunComparisonItem
+from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
+from phoenix.server.api.types.ExperimentComparison import (
+    ExperimentComparison,
+)
+from phoenix.server.api.types.ExperimentRepeatedRunGroup import (
+    ExperimentRepeatedRunGroup,
+    parse_experiment_repeated_run_group_node_id,
+)
 from phoenix.server.api.types.ExperimentRun import ExperimentRun, to_gql_experiment_run
 from phoenix.server.api.types.Functionality import Functionality
 from phoenix.server.api.types.GenerativeModel import GenerativeModel, to_gql_generative_model
 from phoenix.server.api.types.GenerativeProvider import GenerativeProvider, GenerativeProviderKey
 from phoenix.server.api.types.InferenceModel import InferenceModel
 from phoenix.server.api.types.InferencesRole import AncillaryInferencesRole, InferencesRole
-from phoenix.server.api.types.node import from_global_id, from_global_id_with_expected_type
+from phoenix.server.api.types.node import (
+    from_global_id,
+    from_global_id_with_expected_type,
+    is_global_id,
+)
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
     Cursor,
@@ -513,11 +523,12 @@ class Query:
 
         cursors_and_nodes = []
         for example in examples:
-            run_comparison_items = []
+            repeated_run_groups = []
             for experiment_id in experiment_rowids:
-                run_comparison_items.append(
-                    RunComparisonItem(
-                        experiment_id=GlobalID(Experiment.__name__, str(experiment_id)),
+                repeated_run_groups.append(
+                    ExperimentRepeatedRunGroup(
+                        experiment_rowid=experiment_id,
+                        dataset_example_rowid=example.id,
                         runs=[
                             to_gql_experiment_run(run)
                             for run in sorted(
@@ -533,7 +544,7 @@ class Query:
                     created_at=example.created_at,
                     version_id=base_experiment.dataset_version_id,
                 ),
-                run_comparison_items=run_comparison_items,
+                repeated_run_groups=repeated_run_groups,
             )
             cursors_and_nodes.append((Cursor(rowid=example.id), experiment_comparison))
 
@@ -863,8 +874,37 @@ class Query:
         return InferenceModel()
 
     @strawberry.field
-    async def node(self, id: GlobalID, info: Info[Context, None]) -> Node:
-        type_name, node_id = from_global_id(id)
+    async def node(self, id: strawberry.ID, info: Info[Context, None]) -> Node:
+        if not is_global_id(id):
+            try:
+                experiment_rowid, dataset_example_rowid = (
+                    parse_experiment_repeated_run_group_node_id(id)
+                )
+            except Exception:
+                raise NotFound(f"Unknown node: {id}")
+
+            async with info.context.db() as session:
+                runs = (
+                    await session.scalars(
+                        select(models.ExperimentRun)
+                        .where(models.ExperimentRun.experiment_id == experiment_rowid)
+                        .where(models.ExperimentRun.dataset_example_id == dataset_example_rowid)
+                        .order_by(models.ExperimentRun.repetition_number.asc())
+                        .options(
+                            joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id)
+                        )
+                    )
+                ).all()
+            if not runs:
+                raise NotFound(f"Unknown experiment or dataset example: {id}")
+            return ExperimentRepeatedRunGroup(
+                experiment_rowid=experiment_rowid,
+                dataset_example_rowid=dataset_example_rowid,
+                runs=[to_gql_experiment_run(run) for run in runs],
+            )
+
+        global_id = GlobalID.from_id(id)
+        type_name, node_id = from_global_id(global_id)
         if type_name == "Dimension":
             dimension = info.context.model.scalar_dimensions[node_id]
             return to_gql_dimension(node_id, dimension)
@@ -909,26 +949,9 @@ class Query:
             return to_gql_dataset(dataset)
         elif type_name == DatasetExample.__name__:
             example_id = node_id
-            latest_revision_id = (
-                select(func.max(models.DatasetExampleRevision.id))
-                .where(models.DatasetExampleRevision.dataset_example_id == example_id)
-                .scalar_subquery()
-            )
             async with info.context.db() as session:
                 example = await session.scalar(
-                    select(models.DatasetExample)
-                    .join(
-                        models.DatasetExampleRevision,
-                        onclause=models.DatasetExampleRevision.dataset_example_id
-                        == models.DatasetExample.id,
-                    )
-                    .where(
-                        and_(
-                            models.DatasetExample.id == example_id,
-                            models.DatasetExampleRevision.id == latest_revision_id,
-                            models.DatasetExampleRevision.revision_kind != "DELETE",
-                        )
-                    )
+                    select(models.DatasetExample).where(models.DatasetExample.id == example_id)
                 )
             if not example:
                 raise NotFound(f"Unknown dataset example: {id}")
@@ -943,15 +966,7 @@ class Query:
                 )
             if not experiment:
                 raise NotFound(f"Unknown experiment: {id}")
-            return Experiment(
-                id_attr=experiment.id,
-                name=experiment.name,
-                project_name=experiment.project_name,
-                description=experiment.description,
-                created_at=experiment.created_at,
-                updated_at=experiment.updated_at,
-                metadata=experiment.metadata_,
-            )
+            return to_gql_experiment(experiment)
         elif type_name == ExperimentRun.__name__:
             async with info.context.db() as session:
                 if not (

phoenix/server/api/routers/v1/experiment_runs.py

@@ -27,7 +27,7 @@ class ExperimentRun(V1RoutesBaseModel):
         description="The ID of the dataset example used in the experiment run"
     )
     output: Any = Field(description="The output of the experiment task")
-    repetition_number: int = Field(description="The repetition number of the experiment run")
+    repetition_number: int = Field(description="The repetition number of the experiment run", gt=0)
     start_time: datetime = Field(description="The start time of the experiment run")
     end_time: datetime = Field(description="The end time of the experiment run")
     trace_id: Optional[str] = Field(

phoenix/server/api/routers/v1/experiments.py

@@ -46,7 +46,7 @@ class Experiment(V1RoutesBaseModel):
     dataset_version_id: str = Field(
         description="The ID of the dataset version associated with the experiment"
     )
-    repetitions: int = Field(description="Number of times the experiment is repeated")
+    repetitions: int = Field(description="Number of times the experiment is repeated", gt=0)
     metadata: dict[str, Any] = Field(description="Metadata of the experiment")
     project_name: Optional[str] = Field(
         description="The name of the project associated with the experiment"

phoenix/server/api/types/DatasetExample.py

@@ -10,8 +10,12 @@ 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.types.DatasetExampleRevision import DatasetExampleRevision
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
+from phoenix.server.api.types.ExperimentRepeatedRunGroup import (
+    ExperimentRepeatedRunGroup,
+)
 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 (
@@ -65,6 +69,7 @@ class DatasetExample(Node):
         last: Optional[int] = UNSET,
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
+        experiment_ids: Optional[list[GlobalID]] = UNSET,
     ) -> Connection[ExperimentRun]:
         args = ConnectionArgs(
             first=first,
@@ -78,8 +83,51 @@ class DatasetExample(Node):
             .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())
+            .order_by(
+                models.ExperimentRun.experiment_id.asc(),
+                models.ExperimentRun.repetition_number.asc(),
+            )
         )
+        if experiment_ids:
+            experiment_db_ids = [
+                from_global_id_with_expected_type(
+                    global_id=experiment_id,
+                    expected_type_name=models.Experiment.__name__,
+                )
+                for experiment_id in experiment_ids or []
+            ]
+            query = query.where(models.ExperimentRun.experiment_id.in_(experiment_db_ids))
         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)
+
+    @strawberry.field
+    async def experiment_repeated_run_groups(
+        self,
+        info: Info[Context, None],
+        experiment_ids: list[GlobalID],
+    ) -> list[ExperimentRepeatedRunGroup]:
+        example_rowid = self.id_attr
+        experiment_rowids = []
+        for experiment_id in experiment_ids:
+            try:
+                experiment_rowid = from_global_id_with_expected_type(
+                    global_id=experiment_id,
+                    expected_type_name=models.Experiment.__name__,
+                )
+            except Exception:
+                raise BadRequest(f"Invalid experiment ID: {experiment_id}")
+            experiment_rowids.append(experiment_rowid)
+        repeated_run_groups = (
+            await info.context.data_loaders.experiment_repeated_run_groups.load_many(
+                [(experiment_rowid, example_rowid) for experiment_rowid in experiment_rowids]
+            )
+        )
+        return [
+            ExperimentRepeatedRunGroup(
+                experiment_rowid=group.experiment_rowid,
+                dataset_example_rowid=group.dataset_example_rowid,
+                runs=[to_gql_experiment_run(run) for run in group.runs],
+            )
+            for group in repeated_run_groups
+        ]

phoenix/server/api/types/Experiment.py

@@ -5,13 +5,14 @@ import strawberry
 from sqlalchemy import func, select
 from sqlalchemy.orm import joinedload
 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.types.CostBreakdown import CostBreakdown
+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.pagination import (
@@ -32,6 +33,7 @@ class Experiment(Node):
     name: str
     project_name: Optional[str]
     description: Optional[str]
+    dataset_version_id: GlobalID
     metadata: JSON
     created_at: datetime
     updated_at: datetime
@@ -71,7 +73,10 @@ class Experiment(Node):
                 await session.scalars(
                     select(models.ExperimentRun)
                     .where(models.ExperimentRun.experiment_id == experiment_id)
-                    .order_by(models.ExperimentRun.id.desc())
+                    .order_by(
+                        models.ExperimentRun.dataset_example_id.asc(),
+                        models.ExperimentRun.repetition_number.asc(),
+                    )
                     .options(
                         joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id)
                     )
@@ -187,6 +192,10 @@ class Experiment(Node):
                 async for token_type, is_prompt, cost, tokens in data
             ]
 
+    @strawberry.field
+    async def repetition_count(self, info: Info[Context, None]) -> int:
+        return await info.context.data_loaders.experiment_repetition_counts.load(self.id_attr)
+
 
 def to_gql_experiment(
     experiment: models.Experiment,
@@ -201,6 +210,7 @@ def to_gql_experiment(
         name=experiment.name,
         project_name=experiment.project_name,
         description=experiment.description,
+        dataset_version_id=GlobalID(DatasetVersion.__name__, str(experiment.dataset_version_id)),
         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]