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

phoenix/server/api/dataloaders/annotation_summaries.py

@@ -13,6 +13,7 @@ from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
 from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
+from phoenix.server.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -20,27 +21,41 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 AnnotationName: TypeAlias = str
 
-Segment: TypeAlias = tuple[Kind, ProjectRowId, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[
+    Kind,
+    ProjectRowId,
+    TimeInterval,
+    FilterCondition,
+    SessionFilterCondition,
+]
 Param: TypeAlias = AnnotationName
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, AnnotationName]
+Key: TypeAlias = tuple[
+    Kind,
+    ProjectRowId,
+    Optional[TimeRange],
+    FilterCondition,
+    SessionFilterCondition,
+    AnnotationName,
+]
 Result: TypeAlias = Optional[AnnotationSummary]
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = None
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    kind, project_rowid, time_range, filter_condition, eval_name = key
+    kind, project_rowid, time_range, filter_condition, session_filter_condition, eval_name = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (kind, project_rowid, interval, filter_condition), eval_name
+    return (kind, project_rowid, interval, filter_condition, session_filter_condition), eval_name
 
 
 _Section: TypeAlias = tuple[ProjectRowId, AnnotationName, Kind]
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition]
 
 
 class AnnotationSummaryCache(
@@ -61,8 +76,21 @@ class AnnotationSummaryCache(
                 del self._cache[section]
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (kind, project_rowid, interval, filter_condition), annotation_name = _cache_key_fn(key)
-        return (project_rowid, annotation_name, kind), (interval, filter_condition)
+        (
+            (
+                kind,
+                project_rowid,
+                interval,
+                filter_condition,
+                session_filter_condition,
+            ),
+            annotation_name,
+        ) = _cache_key_fn(key)
+        return (project_rowid, annotation_name, kind), (
+            interval,
+            filter_condition,
+            session_filter_condition,
+        )
 
 
 class AnnotationSummaryDataLoader(DataLoader[Key, Result]):
@@ -102,7 +130,9 @@ def _get_stmt(
     segment: Segment,
     *annotation_names: Param,
 ) -> Select[Any]:
-    kind, project_rowid, (start_time, end_time), filter_condition = segment
+    kind, project_rowid, (start_time, end_time), filter_condition, session_filter_condition = (
+        segment
+    )
 
     annotation_model: Union[Type[models.SpanAnnotation], Type[models.TraceAnnotation]]
     entity_model: Union[Type[models.Span], Type[models.Trace]]
@@ -144,6 +174,19 @@ def _get_stmt(
         entity_count_query = entity_count_query.where(
             cast(Type[models.Trace], entity_model).project_rowid == project_rowid
         )
+    else:
+        assert_never(kind)
+
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=[project_rowid],
+            start_time=start_time,
+            end_time=end_time,
+        )
+        entity_count_query = entity_count_query.where(
+            models.Trace.project_session_rowid.in_(filtered_session_rowids)
+        )
 
     entity_count_query = entity_count_query.where(
         or_(score_column.is_not(None), label_column.is_not(None))
@@ -186,6 +229,15 @@ def _get_stmt(
     else:
         assert_never(kind)
 
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=[project_rowid],
+            start_time=start_time,
+            end_time=end_time,
+        )
+        base_stmt = base_stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
+
     base_stmt = base_stmt.where(or_(score_column.is_not(None), label_column.is_not(None)))
     base_stmt = base_stmt.where(name_column.in_(annotation_names))
 

phoenix/server/api/dataloaders/average_experiment_repeated_run_group_latency.py

@@ -0,0 +1,50 @@
+from typing import Optional
+
+from sqlalchemy import func, select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentID: TypeAlias = int
+DatasetExampleID: TypeAlias = int
+RunLatency: TypeAlias = float
+Key: TypeAlias = tuple[ExperimentID, DatasetExampleID]
+Result: TypeAlias = Optional[RunLatency]
+
+
+class AverageExperimentRepeatedRunGroupLatencyDataLoader(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]:
+        average_latency_query = (
+            select(
+                models.ExperimentRun.experiment_id.label("experiment_id"),
+                models.ExperimentRun.dataset_example_id.label("example_id"),
+                func.avg(models.ExperimentRun.latency_ms).label("average_repetition_latency_ms"),
+            )
+            .select_from(models.ExperimentRun)
+            .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)
+        )
+        async with self._db() as session:
+            average_run_latencies_ms = {
+                (experiment_id, example_id): average_run_latency_ms
+                async for experiment_id, example_id, average_run_latency_ms in await session.stream(
+                    average_latency_query
+                )
+            }
+        return [
+            average_run_latencies_ms.get((experiment_id, example_id))
+            for experiment_id, example_id in keys
+        ]

phoenix/server/api/dataloaders/average_experiment_run_latency.py

@@ -23,32 +23,25 @@ class AverageExperimentRunLatencyDataLoader(DataLoader[Key, Result]):
 
     async def _load_fn(self, keys: list[Key]) -> list[Result]:
         experiment_ids = keys
-        resolved_experiment_ids = (
-            select(models.Experiment.id)
-            .where(models.Experiment.id.in_(set(experiment_ids)))
-            .subquery()
-        )
-        query = (
+        average_repetition_latency_ms = (
             select(
-                resolved_experiment_ids.c.id,
-                func.avg(
-                    func.extract("epoch", models.ExperimentRun.end_time)
-                    - func.extract("epoch", models.ExperimentRun.start_time)
-                ),
+                models.ExperimentRun.experiment_id.label("experiment_id"),
+                func.avg(models.ExperimentRun.latency_ms).label("average_repetition_latency_ms"),
             )
-            .outerjoin_from(
-                from_=resolved_experiment_ids,
-                target=models.ExperimentRun,
-                onclause=resolved_experiment_ids.c.id == models.ExperimentRun.experiment_id,
-            )
-            .group_by(resolved_experiment_ids.c.id)
+            .select_from(models.ExperimentRun)
+            .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+            .group_by(models.ExperimentRun.dataset_example_id, models.ExperimentRun.experiment_id)
+            .subquery()
         )
+        query = select(
+            average_repetition_latency_ms.c.experiment_id,
+            func.avg(average_repetition_latency_ms.c.average_repetition_latency_ms).label(
+                "average_run_latency_ms"
+            ),
+        ).group_by(average_repetition_latency_ms.c.experiment_id)
         async with self._db() as session:
-            avg_latencies = {
-                experiment_id: avg_latency
-                async for experiment_id, avg_latency in await session.stream(query)
+            average_run_latencies_ms = {
+                experiment_id: average_run_latency_ms
+                async for experiment_id, average_run_latency_ms in await session.stream(query)
             }
-        return [
-            avg_latencies.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))
-            for experiment_id in keys
-        ]
+        return [average_run_latencies_ms.get(experiment_id) for experiment_id in keys]

phoenix/server/api/dataloaders/cache/two_tier_cache.py

@@ -7,7 +7,7 @@ single-tier system we would need to check all the keys to see if they are in the
 subset that we want to invalidate.
 """
 
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from asyncio import Future
 from collections.abc import Callable
 from typing import Any, Generic, Optional, TypeVar
@@ -25,7 +25,6 @@ _SubKey = TypeVar("_SubKey")
 class TwoTierCache(
     AbstractCache[_Key, _Result],
     Generic[_Key, _Result, _Section, _SubKey],
-    ABC,
 ):
     def __init__(
         self,

phoenix/server/api/dataloaders/dataset_dataset_splits.py

@@ -0,0 +1,52 @@
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+DatasetID: TypeAlias = int
+Key: TypeAlias = DatasetID
+Result: TypeAlias = list[models.DatasetSplit]
+
+
+class DatasetDatasetSplitsDataLoader(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]:
+        dataset_ids = keys
+        async with self._db() as session:
+            splits: dict[DatasetID, dict[int, models.DatasetSplit]] = {
+                dataset_id: {} for dataset_id in dataset_ids
+            }
+
+            async for dataset_id, split in await session.stream(
+                select(models.DatasetExample.dataset_id, models.DatasetSplit)
+                .select_from(models.DatasetSplit)
+                .join(
+                    models.DatasetSplitDatasetExample,
+                    onclause=(
+                        models.DatasetSplit.id == models.DatasetSplitDatasetExample.dataset_split_id
+                    ),
+                )
+                .join(
+                    models.DatasetExample,
+                    onclause=(
+                        models.DatasetSplitDatasetExample.dataset_example_id
+                        == models.DatasetExample.id
+                    ),
+                )
+                .where(models.DatasetExample.dataset_id.in_(dataset_ids))
+            ):
+                # Use dict to deduplicate splits by split.id
+                if dataset_id in splits:
+                    splits[dataset_id][split.id] = split
+
+            return [
+                sorted(splits.get(dataset_id, {}).values(), key=lambda x: x.name)
+                for dataset_id in keys
+            ]

phoenix/server/api/dataloaders/dataset_example_revisions.py

@@ -91,7 +91,6 @@ class DatasetExampleRevisionsDataLoader(DataLoader[Key, Result]):
                 onclause=revision_ids.c.version_id == models.DatasetVersion.id,
                 isouter=True,  # keep rows where the version id is null
             )
-            .where(models.DatasetExampleRevision.revision_kind != "DELETE")
         )
         async with self._db() as session:
             results = {

phoenix/server/api/dataloaders/dataset_example_splits.py

@@ -0,0 +1,40 @@
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExampleID: TypeAlias = int
+Key: TypeAlias = ExampleID
+Result: TypeAlias = list[models.DatasetSplit]
+
+
+class DatasetExampleSplitsDataLoader(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]:
+        example_ids = keys
+        async with self._db() as session:
+            splits: dict[ExampleID, list[models.DatasetSplit]] = {}
+
+            async for example_id, split in await session.stream(
+                select(models.DatasetSplitDatasetExample.dataset_example_id, models.DatasetSplit)
+                .select_from(models.DatasetSplit)
+                .join(
+                    models.DatasetSplitDatasetExample,
+                    onclause=(
+                        models.DatasetSplit.id == models.DatasetSplitDatasetExample.dataset_split_id
+                    ),
+                )
+                .where(models.DatasetSplitDatasetExample.dataset_example_id.in_(example_ids))
+            ):
+                if example_id not in splits:
+                    splits[example_id] = []
+                splits[example_id].append(split)
+
+            return [sorted(splits.get(example_id, []), key=lambda x: x.name) for example_id in keys]

phoenix/server/api/dataloaders/dataset_examples_and_versions_by_experiment_run.py

@@ -0,0 +1,47 @@
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentRunID: TypeAlias = int
+DatasetExampleID: TypeAlias = int
+DatasetVersionID: TypeAlias = int
+Key: TypeAlias = ExperimentRunID
+Result: TypeAlias = tuple[models.DatasetExample, DatasetVersionID]
+
+
+class DatasetExamplesAndVersionsByExperimentRunDataLoader(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_run_ids = set(keys)
+        examples_and_versions_query = (
+            select(
+                models.ExperimentRun.id.label("experiment_run_id"),
+                models.DatasetExample,
+                models.Experiment.dataset_version_id.label("dataset_version_id"),
+            )
+            .select_from(models.ExperimentRun)
+            .join(
+                models.DatasetExample,
+                models.DatasetExample.id == models.ExperimentRun.dataset_example_id,
+            )
+            .join(
+                models.Experiment,
+                models.Experiment.id == models.ExperimentRun.experiment_id,
+            )
+            .where(models.ExperimentRun.id.in_(experiment_run_ids))
+        )
+        async with self._db() as session:
+            examples_and_versions = {
+                experiment_run_id: (example, version_id)
+                for experiment_run_id, example, version_id in (
+                    await session.execute(examples_and_versions_query)
+                ).all()
+            }
+
+        return [examples_and_versions[key] for key in keys]

phoenix/server/api/dataloaders/dataset_labels.py

@@ -0,0 +1,36 @@
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+DatasetID: TypeAlias = int
+Key: TypeAlias = DatasetID
+Result: TypeAlias = list[models.DatasetLabel]
+
+
+class DatasetLabelsDataLoader(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]:
+        dataset_ids = keys
+        async with self._db() as session:
+            labels: dict[Key, Result] = {}
+            for dataset_id, label in await session.execute(
+                select(models.DatasetsDatasetLabel.dataset_id, models.DatasetLabel)
+                .select_from(models.DatasetLabel)
+                .join(
+                    models.DatasetsDatasetLabel,
+                    models.DatasetLabel.id == models.DatasetsDatasetLabel.dataset_label_id,
+                )
+                .where(models.DatasetsDatasetLabel.dataset_id.in_(dataset_ids))
+            ):
+                if dataset_id not in labels:
+                    labels[dataset_id] = []
+                labels[dataset_id].append(label)
+        return [
+            sorted(labels.get(dataset_id, []), key=lambda label: label.name) for dataset_id in keys
+        ]

phoenix/server/api/dataloaders/document_evaluation_summaries.py

@@ -10,7 +10,7 @@ from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
-from phoenix.db.helpers import SupportedSQLDialect, num_docs_col
+from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.metrics.retrieval_metrics import RetrievalMetrics
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
@@ -122,7 +122,7 @@ def _get_stmt(
         select(
             mda.name,
             models.Span.id,
-            num_docs_col(dialect),
+            models.Span.num_documents.label("num_docs"),
             mda.score,
             mda.document_position,
         )

phoenix/server/api/dataloaders/document_evaluations.py

@@ -5,11 +5,10 @@ from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
 from phoenix.db import models
-from phoenix.server.api.types.Evaluation import DocumentEvaluation
 from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
-Result: TypeAlias = list[DocumentEvaluation]
+Result: TypeAlias = list[models.DocumentAnnotation]
 
 
 class DocumentEvaluationsDataLoader(DataLoader[Key, Result]):
@@ -18,14 +17,12 @@ class DocumentEvaluationsDataLoader(DataLoader[Key, Result]):
         self._db = db
 
     async def _load_fn(self, keys: list[Key]) -> list[Result]:
-        document_evaluations_by_id: defaultdict[Key, Result] = defaultdict(list)
+        document_annotations_by_id: defaultdict[Key, Result] = defaultdict(list)
         mda = models.DocumentAnnotation
         async with self._db() as session:
-            data = await session.stream_scalars(
-                select(mda).where(mda.span_rowid.in_(keys)).where(mda.annotator_kind == "LLM")
-            )
+            data = await session.stream_scalars(select(mda).where(mda.span_rowid.in_(keys)))
             async for document_evaluation in data:
-                document_evaluations_by_id[document_evaluation.span_rowid].append(
-                    DocumentEvaluation.from_sql_document_annotation(document_evaluation)
+                document_annotations_by_id[document_evaluation.span_rowid].append(
+                    document_evaluation
                 )
-        return [document_evaluations_by_id[key] for key in keys]
+        return [document_annotations_by_id[key] for key in keys]

phoenix/server/api/dataloaders/experiment_annotation_summaries.py

@@ -2,7 +2,7 @@ from collections import defaultdict
 from dataclasses import dataclass
 from typing import Optional
 
-from sqlalchemy import func, select
+from sqlalchemy import and_, func, select
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias
 
@@ -37,43 +37,97 @@ class ExperimentAnnotationSummaryDataLoader(DataLoader[Key, Result]):
     async def _load_fn(self, keys: list[Key]) -> list[Result]:
         experiment_ids = keys
         summaries: defaultdict[ExperimentID, Result] = defaultdict(list)
+        repetition_mean_scores_by_example_subquery = (
+            select(
+                models.ExperimentRun.experiment_id.label("experiment_id"),
+                models.ExperimentRunAnnotation.name.label("annotation_name"),
+                func.avg(models.ExperimentRunAnnotation.score).label("mean_repetition_score"),
+            )
+            .select_from(models.ExperimentRunAnnotation)
+            .join(
+                models.ExperimentRun,
+                models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
+            )
+            .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+            .group_by(
+                models.ExperimentRun.experiment_id,
+                models.ExperimentRun.dataset_example_id,
+                models.ExperimentRunAnnotation.name,
+            )
+            .subquery()
+            .alias("repetition_mean_scores_by_example")
+        )
+        repetition_mean_scores_subquery = (
+            select(
+                repetition_mean_scores_by_example_subquery.c.experiment_id.label("experiment_id"),
+                repetition_mean_scores_by_example_subquery.c.annotation_name.label(
+                    "annotation_name"
+                ),
+                func.avg(repetition_mean_scores_by_example_subquery.c.mean_repetition_score).label(
+                    "mean_score"
+                ),
+            )
+            .select_from(repetition_mean_scores_by_example_subquery)
+            .group_by(
+                repetition_mean_scores_by_example_subquery.c.experiment_id,
+                repetition_mean_scores_by_example_subquery.c.annotation_name,
+            )
+            .subquery()
+            .alias("repetition_mean_scores")
+        )
+        repetitions_subquery = (
+            select(
+                models.ExperimentRun.experiment_id.label("experiment_id"),
+                models.ExperimentRunAnnotation.name.label("annotation_name"),
+                func.min(models.ExperimentRunAnnotation.score).label("min_score"),
+                func.max(models.ExperimentRunAnnotation.score).label("max_score"),
+                func.count().label("count"),
+                func.count(models.ExperimentRunAnnotation.error).label("error_count"),
+            )
+            .select_from(models.ExperimentRunAnnotation)
+            .join(
+                models.ExperimentRun,
+                models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
+            )
+            .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+            .group_by(models.ExperimentRun.experiment_id, models.ExperimentRunAnnotation.name)
+            .subquery()
+        )
+        run_scores_query = (
+            select(
+                repetition_mean_scores_subquery.c.experiment_id.label("experiment_id"),
+                repetition_mean_scores_subquery.c.annotation_name.label("annotation_name"),
+                repetition_mean_scores_subquery.c.mean_score.label("mean_score"),
+                repetitions_subquery.c.min_score.label("min_score"),
+                repetitions_subquery.c.max_score.label("max_score"),
+                repetitions_subquery.c.count.label("count_"),
+                repetitions_subquery.c.error_count.label("error_count"),
+            )
+            .select_from(repetition_mean_scores_subquery)
+            .join(
+                repetitions_subquery,
+                and_(
+                    repetitions_subquery.c.experiment_id
+                    == repetition_mean_scores_subquery.c.experiment_id,
+                    repetitions_subquery.c.annotation_name
+                    == repetition_mean_scores_subquery.c.annotation_name,
+                ),
+            )
+            .order_by(repetition_mean_scores_subquery.c.annotation_name)
+        )
         async with self._db() as session:
-            async for (
-                experiment_id,
-                annotation_name,
-                min_score,
-                max_score,
-                mean_score,
-                count,
-                error_count,
-            ) in await session.stream(
-                select(
-                    models.ExperimentRun.experiment_id,
-                    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,
-                )
-                .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
-                .group_by(models.ExperimentRun.experiment_id, models.ExperimentRunAnnotation.name)
-            ):
-                summaries[experiment_id].append(
+            async for scores_tuple in await session.stream(run_scores_query):
+                summaries[scores_tuple.experiment_id].append(
                     ExperimentAnnotationSummary(
-                        annotation_name=annotation_name,
-                        min_score=min_score,
-                        max_score=max_score,
-                        mean_score=mean_score,
-                        count=count,
-                        error_count=error_count,
+                        annotation_name=scores_tuple.annotation_name,
+                        min_score=scores_tuple.min_score,
+                        max_score=scores_tuple.max_score,
+                        mean_score=scores_tuple.mean_score,
+                        count=scores_tuple.count_,
+                        error_count=scores_tuple.error_count,
                     )
                 )
         return [
             sorted(summaries[experiment_id], key=lambda summary: summary.annotation_name)
-            for experiment_id in keys
+            for experiment_id in experiment_ids
         ]

phoenix/server/api/dataloaders/experiment_dataset_splits.py

@@ -0,0 +1,43 @@
+from sqlalchemy import 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
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = list[models.DatasetSplit]
+
+
+class ExperimentDatasetSplitsDataLoader(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
+        async with self._db() as session:
+            splits: dict[ExperimentID, list[models.DatasetSplit]] = {}
+
+            async for experiment_id, split in await session.stream(
+                select(models.ExperimentDatasetSplit.experiment_id, models.DatasetSplit)
+                .select_from(models.DatasetSplit)
+                .join(
+                    models.ExperimentDatasetSplit,
+                    onclause=(
+                        models.DatasetSplit.id == models.ExperimentDatasetSplit.dataset_split_id
+                    ),
+                )
+                .where(models.ExperimentDatasetSplit.experiment_id.in_(experiment_ids))
+            ):
+                if experiment_id not in splits:
+                    splits[experiment_id] = []
+                splits[experiment_id].append(split)
+
+            return [
+                sorted(splits.get(experiment_id, []), key=lambda x: x.name)
+                for experiment_id in keys
+            ]