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

phoenix/server/api/queries.py

@@ -1,14 +1,14 @@
 import re
 from collections import defaultdict
 from datetime import datetime
-from typing import Any, Iterable, Iterator, Optional, Union
+from typing import Any, Iterable, Iterator, Literal, Optional, Union
 from typing import cast as type_cast
 
 import numpy as np
 import numpy.typing as npt
 import strawberry
 from sqlalchemy import ColumnElement, String, and_, case, cast, func, select, text
-from sqlalchemy.orm import aliased, joinedload, load_only
+from sqlalchemy.orm import joinedload, load_only
 from starlette.authentication import UnauthenticatedUser
 from strawberry import ID, UNSET
 from strawberry.relay import Connection, GlobalID, Node
@@ -22,7 +22,10 @@ from phoenix.config import (
 )
 from phoenix.db import models
 from phoenix.db.constants import DEFAULT_PROJECT_TRACE_RETENTION_POLICY_ID
-from phoenix.db.helpers import SupportedSQLDialect, exclude_experiment_projects
+from phoenix.db.helpers import (
+    SupportedSQLDialect,
+    exclude_experiment_projects,
+)
 from phoenix.db.models import LatencyMs
 from phoenix.pointcloud.clustering import Hdbscan
 from phoenix.server.api.auth import MSG_ADMIN_ONLY, IsAdmin
@@ -46,8 +49,10 @@ from phoenix.server.api.input_types.ProjectSort import ProjectColumn, ProjectSor
 from phoenix.server.api.input_types.PromptFilter import PromptFilter
 from phoenix.server.api.types.AnnotationConfig import AnnotationConfig, to_gql_annotation_config
 from phoenix.server.api.types.Cluster import Cluster, to_gql_clusters
-from phoenix.server.api.types.Dataset import Dataset, to_gql_dataset
+from phoenix.server.api.types.Dataset import Dataset
 from phoenix.server.api.types.DatasetExample import DatasetExample
+from phoenix.server.api.types.DatasetLabel import DatasetLabel
+from phoenix.server.api.types.DatasetSplit import DatasetSplit
 from phoenix.server.api.types.Dimension import to_gql_dimension
 from phoenix.server.api.types.EmbeddingDimension import (
     DEFAULT_CLUSTER_SELECTION_EPSILON,
@@ -57,14 +62,24 @@ from phoenix.server.api.types.EmbeddingDimension import (
 )
 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.ExperimentRun import ExperimentRun, to_gql_experiment_run
+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
 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.GenerativeModel import GenerativeModel
 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,
@@ -74,21 +89,21 @@ from phoenix.server.api.types.pagination import (
 )
 from phoenix.server.api.types.PlaygroundModel import PlaygroundModel
 from phoenix.server.api.types.Project import Project
-from phoenix.server.api.types.ProjectSession import ProjectSession, to_gql_project_session
+from phoenix.server.api.types.ProjectSession import ProjectSession
 from phoenix.server.api.types.ProjectTraceRetentionPolicy import ProjectTraceRetentionPolicy
-from phoenix.server.api.types.Prompt import Prompt, to_gql_prompt_from_orm
-from phoenix.server.api.types.PromptLabel import PromptLabel, to_gql_prompt_label
+from phoenix.server.api.types.Prompt import Prompt
+from phoenix.server.api.types.PromptLabel import PromptLabel
 from phoenix.server.api.types.PromptVersion import PromptVersion, to_gql_prompt_version
-from phoenix.server.api.types.PromptVersionTag import PromptVersionTag, to_gql_prompt_version_tag
+from phoenix.server.api.types.PromptVersionTag import PromptVersionTag
 from phoenix.server.api.types.ServerStatus import ServerStatus
 from phoenix.server.api.types.SortDir import SortDir
 from phoenix.server.api.types.Span import Span
-from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.server.api.types.SpanAnnotation import SpanAnnotation
 from phoenix.server.api.types.SystemApiKey import SystemApiKey
 from phoenix.server.api.types.Trace import Trace
-from phoenix.server.api.types.TraceAnnotation import TraceAnnotation, to_gql_trace_annotation
-from phoenix.server.api.types.User import User, to_gql_user
-from phoenix.server.api.types.UserApiKey import UserApiKey, to_gql_api_key
+from phoenix.server.api.types.TraceAnnotation import TraceAnnotation
+from phoenix.server.api.types.User import User
+from phoenix.server.api.types.UserApiKey import UserApiKey
 from phoenix.server.api.types.UserRole import UserRole
 from phoenix.server.api.types.ValidationResult import ValidationResult
 
@@ -108,29 +123,52 @@ class DbTableStats:
 
 
 @strawberry.type
-class MetricCounts:
-    num_increases: int
-    num_decreases: int
-    num_equal: int
-
+class ExperimentRunMetricComparison:
+    num_runs_improved: int = strawberry.field(
+        description=(
+            "The number of runs in which the base experiment improved "
+            "on the best run in any compare experiment."
+        )
+    )
+    num_runs_regressed: int = strawberry.field(
+        description=(
+            "The number of runs in which the base experiment regressed "
+            "on the best run in any compare experiment."
+        )
+    )
+    num_runs_equal: int = strawberry.field(
+        description=(
+            "The number of runs in which the base experiment is equal to the best run "
+            "in any compare experiment."
+        )
+    )
+    num_total_runs: strawberry.Private[int]
 
-@strawberry.type
-class CompareExperimentRunMetricCounts:
-    compare_experiment_id: GlobalID
-    latency: MetricCounts
-    prompt_token_count: MetricCounts
-    completion_token_count: MetricCounts
-    total_token_count: MetricCounts
-    total_cost: MetricCounts
+    @strawberry.field(
+        description=(
+            "The number of runs in the base experiment that could not be compared, either because "
+            "the base experiment run was missing a value or because all compare experiment runs "
+            "were missing values."
+        )
+    )  # type: ignore[misc]
+    def num_runs_without_comparison(self) -> int:
+        return (
+            self.num_total_runs
+            - self.num_runs_improved
+            - self.num_runs_regressed
+            - self.num_runs_equal
+        )
 
 
 @strawberry.type
-class CompareExperimentRunAnnotationMetricCounts:
-    annotation_name: str
-    compare_experiment_id: GlobalID
-    num_increases: int
-    num_decreases: int
-    num_equal: int
+class ExperimentRunMetricComparisons:
+    latency: ExperimentRunMetricComparison
+    total_token_count: ExperimentRunMetricComparison
+    prompt_token_count: ExperimentRunMetricComparison
+    completion_token_count: ExperimentRunMetricComparison
+    total_cost: ExperimentRunMetricComparison
+    prompt_cost: ExperimentRunMetricComparison
+    completion_cost: ExperimentRunMetricComparison
 
 
 @strawberry.type
@@ -150,7 +188,17 @@ class Query:
     async def generative_models(
         self,
         info: Info[Context, None],
-    ) -> list[GenerativeModel]:
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[GenerativeModel]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
         async with info.context.db() as session:
             result = await session.scalars(
                 select(models.GenerativeModel)
@@ -160,17 +208,16 @@ class Query:
                     models.GenerativeModel.provider.nullslast(),
                     models.GenerativeModel.name,
                 )
-                .options(joinedload(models.GenerativeModel.token_prices))
             )
-
-        return [to_gql_generative_model(model) for model in result.unique()]
+            data = [GenerativeModel(id=model.id, db_record=model) for model in result.unique()]
+        return connection_from_list(data=data, args=args)
 
     @strawberry.field
     async def playground_models(self, input: Optional[ModelsInput] = None) -> list[PlaygroundModel]:
         if input is not None and input.provider_key is not None:
             supported_model_names = PLAYGROUND_CLIENT_REGISTRY.list_models(input.provider_key)
             supported_models = [
-                PlaygroundModel(name=model_name, provider_key=input.provider_key)
+                PlaygroundModel(name_value=model_name, provider_key_value=input.provider_key)
                 for model_name in supported_model_names
             ]
             return supported_models
@@ -179,7 +226,9 @@ class Query:
         all_models: list[PlaygroundModel] = []
         for provider_key, model_name in registered_models:
             if model_name is not None and provider_key is not None:
-                all_models.append(PlaygroundModel(name=model_name, provider_key=provider_key))
+                all_models.append(
+                    PlaygroundModel(name_value=model_name, provider_key_value=provider_key)
+                )
         return all_models
 
     @strawberry.field
@@ -223,7 +272,7 @@ class Query:
         )
         async with info.context.db() as session:
             users = await session.stream_scalars(stmt)
-            data = [to_gql_user(user) async for user in users]
+            data = [User(id=user.id, db_record=user) async for user in users]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field
@@ -253,7 +302,7 @@ class Query:
         )
         async with info.context.db() as session:
             api_keys = await session.scalars(stmt)
-        return [to_gql_api_key(api_key) for api_key in api_keys]
+        return [UserApiKey(id=api_key.id, db_record=api_key) for api_key in api_keys]
 
     @strawberry.field(permission_classes=[IsAdmin])  # type: ignore
     async def system_api_keys(self, info: Info[Context, None]) -> list[SystemApiKey]:
@@ -265,16 +314,7 @@ class Query:
         )
         async with info.context.db() as session:
             api_keys = await session.scalars(stmt)
-        return [
-            SystemApiKey(
-                id_attr=api_key.id,
-                name=api_key.name,
-                description=api_key.description,
-                created_at=api_key.created_at,
-                expires_at=api_key.expires_at,
-            )
-            for api_key in api_keys
-        ]
+        return [SystemApiKey(id=api_key.id, db_record=api_key) for api_key in api_keys]
 
     @strawberry.field
     async def projects(
@@ -315,13 +355,7 @@ class Query:
         stmt = exclude_experiment_projects(stmt)
         async with info.context.db() as session:
             projects = await session.stream_scalars(stmt)
-            data = [
-                Project(
-                    project_rowid=project.id,
-                    db_project=project,
-                )
-                async for project in projects
-            ]
+            data = [Project(id=project.id, db_record=project) async for project in projects]
         return connection_from_list(data=data, args=args)
 
     @strawberry.field
@@ -350,11 +384,39 @@ class Query:
             sort_col = getattr(models.Dataset, sort.col.value)
             stmt = stmt.order_by(sort_col.desc() if sort.dir is SortDir.desc else sort_col.asc())
         if filter:
-            stmt = stmt.where(getattr(models.Dataset, filter.col.value).ilike(f"%{filter.value}%"))
+            # Apply name filter
+            if filter.col and filter.value:
+                stmt = stmt.where(
+                    getattr(models.Dataset, filter.col.value).ilike(f"%{filter.value}%")
+                )
+
+            # Apply label filter
+            if filter.filter_labels and filter.filter_labels is not UNSET:
+                label_rowids = []
+                for label_id in filter.filter_labels:
+                    try:
+                        label_rowid = from_global_id_with_expected_type(
+                            global_id=GlobalID.from_id(label_id),
+                            expected_type_name="DatasetLabel",
+                        )
+                        label_rowids.append(label_rowid)
+                    except ValueError:
+                        continue  # Skip invalid label IDs
+
+                if label_rowids:
+                    # Join with the junction table to filter by labels
+                    stmt = (
+                        stmt.join(
+                            models.DatasetsDatasetLabel,
+                            models.Dataset.id == models.DatasetsDatasetLabel.dataset_id,
+                        )
+                        .where(models.DatasetsDatasetLabel.dataset_label_id.in_(label_rowids))
+                        .distinct()
+                    )
         async with info.context.db() as session:
             datasets = await session.scalars(stmt)
         return connection_from_list(
-            data=[to_gql_dataset(dataset) for dataset in datasets], args=args
+            data=[Dataset(id=dataset.id, db_record=dataset) for dataset in datasets], args=args
         )
 
     @strawberry.field
@@ -413,6 +475,7 @@ class Query:
                     )
                 )
             ).all()
+
             if not experiments or len(experiments) < len(experiment_rowids):
                 raise NotFound("Unable to resolve one or more experiment IDs.")
             num_datasets = len(set(experiment.dataset_id for experiment in experiments))
@@ -421,37 +484,19 @@ class Query:
             base_experiment = next(
                 experiment for experiment in experiments if experiment.id == base_experiment_rowid
             )
-            revision_ids = (
-                select(func.max(models.DatasetExampleRevision.id))
-                .join(
-                    models.DatasetExample,
-                    models.DatasetExample.id == models.DatasetExampleRevision.dataset_example_id,
-                )
-                .where(
-                    and_(
-                        models.DatasetExampleRevision.dataset_version_id
-                        <= base_experiment.dataset_version_id,
-                        models.DatasetExample.dataset_id == base_experiment.dataset_id,
-                    )
-                )
-                .group_by(models.DatasetExampleRevision.dataset_example_id)
-                .scalar_subquery()
-            )
+
+            # Use ExperimentDatasetExample to pull down examples.
+            # Splits are mutable and should not be used for comparison.
+            # The comparison should only occur against examples which were assigned to the same
+            # splits at the time of execution of the ExperimentRun.
             examples_query = (
                 select(models.DatasetExample)
-                .distinct(models.DatasetExample.id)
-                .join(
-                    models.DatasetExampleRevision,
-                    onclause=and_(
-                        models.DatasetExample.id
-                        == models.DatasetExampleRevision.dataset_example_id,
-                        models.DatasetExampleRevision.id.in_(revision_ids),
-                        models.DatasetExampleRevision.revision_kind != "DELETE",
-                    ),
-                )
+                .join(models.ExperimentDatasetExample)
+                .where(models.ExperimentDatasetExample.experiment_id == base_experiment_rowid)
                 .order_by(models.DatasetExample.id.desc())
                 .limit(page_size + 1)
             )
+
             if cursor is not None:
                 examples_query = examples_query.where(models.DatasetExample.id < cursor.rowid)
 
@@ -490,15 +535,17 @@ 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)),
-                        runs=[
-                            to_gql_experiment_run(run)
+                repeated_run_groups.append(
+                    ExperimentRepeatedRunGroup(
+                        experiment_rowid=experiment_id,
+                        dataset_example_rowid=example.id,
+                        cached_runs=[
+                            ExperimentRun(id=run.id, db_record=run)
                             for run in sorted(
-                                runs[example.id][experiment_id], key=lambda run: run.id
+                                runs[example.id][experiment_id],
+                                key=lambda run: run.repetition_number,
                             )
                         ],
                     )
@@ -506,11 +553,11 @@ class Query:
             experiment_comparison = ExperimentComparison(
                 id_attr=example.id,
                 example=DatasetExample(
-                    id_attr=example.id,
-                    created_at=example.created_at,
+                    id=example.id,
+                    db_record=example,
                     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))
 
@@ -521,12 +568,12 @@ class Query:
         )
 
     @strawberry.field
-    async def compare_experiment_run_metric_counts(
+    async def experiment_run_metric_comparisons(
         self,
         info: Info[Context, None],
         base_experiment_id: GlobalID,
         compare_experiment_ids: list[GlobalID],
-    ) -> list[CompareExperimentRunMetricCounts]:
+    ) -> ExperimentRunMetricComparisons:
         if base_experiment_id in compare_experiment_ids:
             raise BadRequest("Compare experiment IDs cannot contain the base experiment ID")
         if not compare_experiment_ids:
@@ -553,375 +600,256 @@ class Query:
                 raise BadRequest(f"Invalid compare experiment ID: {compare_experiment_id}")
 
         base_experiment_runs = (
-            select(models.ExperimentRun)
+            select(
+                models.ExperimentRun.dataset_example_id,
+                func.min(models.ExperimentRun.start_time).label("start_time"),
+                func.min(models.ExperimentRun.end_time).label("end_time"),
+                func.sum(models.SpanCost.total_tokens).label("total_tokens"),
+                func.sum(models.SpanCost.prompt_tokens).label("prompt_tokens"),
+                func.sum(models.SpanCost.completion_tokens).label("completion_tokens"),
+                func.sum(models.SpanCost.total_cost).label("total_cost"),
+                func.sum(models.SpanCost.prompt_cost).label("prompt_cost"),
+                func.sum(models.SpanCost.completion_cost).label("completion_cost"),
+            )
+            .select_from(models.ExperimentRun)
+            .join(
+                models.Trace,
+                onclause=models.ExperimentRun.trace_id == models.Trace.trace_id,
+                isouter=True,
+            )
+            .join(
+                models.SpanCost,
+                onclause=models.Trace.id == models.SpanCost.trace_rowid,
+                isouter=True,
+            )
             .where(models.ExperimentRun.experiment_id == base_experiment_rowid)
+            .group_by(models.ExperimentRun.dataset_example_id)
             .subquery()
             .alias("base_experiment_runs")
         )
-        base_experiment_traces = aliased(models.Trace, name="base_experiment_traces")
-        base_experiment_span_costs = (
+        compare_experiment_runs = (
             select(
-                models.SpanCost.trace_rowid,
-                func.coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
-                func.coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
-                func.coalesce(func.sum(models.SpanCost.completion_tokens), 0).label(
-                    "completion_tokens"
-                ),
-                func.coalesce(func.sum(models.SpanCost.total_cost), 0).label("total_cost"),
+                models.ExperimentRun.dataset_example_id,
+                func.min(
+                    LatencyMs(models.ExperimentRun.start_time, models.ExperimentRun.end_time)
+                ).label("min_latency_ms"),
+                func.min(models.SpanCost.total_tokens).label("min_total_tokens"),
+                func.min(models.SpanCost.prompt_tokens).label("min_prompt_tokens"),
+                func.min(models.SpanCost.completion_tokens).label("min_completion_tokens"),
+                func.min(models.SpanCost.total_cost).label("min_total_cost"),
+                func.min(models.SpanCost.prompt_cost).label("min_prompt_cost"),
+                func.min(models.SpanCost.completion_cost).label("min_completion_cost"),
             )
-            .select_from(models.SpanCost)
-            .group_by(
-                models.SpanCost.trace_rowid,
-            )
-            .subquery()
-            .alias("base_experiment_span_costs")
-        )
-
-        query = (
-            select()  # add selected columns below
-            .select_from(base_experiment_runs)
+            .select_from(models.ExperimentRun)
             .join(
-                base_experiment_traces,
-                onclause=base_experiment_runs.c.trace_id == base_experiment_traces.trace_id,
+                models.Trace,
+                onclause=models.ExperimentRun.trace_id == models.Trace.trace_id,
                 isouter=True,
             )
             .join(
-                base_experiment_span_costs,
-                onclause=base_experiment_traces.id == base_experiment_span_costs.c.trace_rowid,
+                models.SpanCost,
+                onclause=models.Trace.id == models.SpanCost.trace_rowid,
                 isouter=True,
             )
+            .where(
+                models.ExperimentRun.experiment_id.in_(compare_experiment_rowids),
+            )
+            .group_by(models.ExperimentRun.dataset_example_id)
+            .subquery()
+            .alias("comp_exp_run_mins")
         )
 
         base_experiment_run_latency = LatencyMs(
             base_experiment_runs.c.start_time, base_experiment_runs.c.end_time
         ).label("base_experiment_run_latency_ms")
-        base_experiment_run_prompt_token_count = base_experiment_span_costs.c.prompt_tokens
-        base_experiment_run_completion_token_count = base_experiment_span_costs.c.completion_tokens
-        base_experiment_run_total_token_count = base_experiment_span_costs.c.total_tokens
-        base_experiment_run_total_cost = base_experiment_span_costs.c.total_cost
-
-        for compare_experiment_index, compare_experiment_rowid in enumerate(
-            compare_experiment_rowids
-        ):
-            compare_experiment_runs = (
-                select(models.ExperimentRun)
-                .where(models.ExperimentRun.experiment_id == compare_experiment_rowid)
-                .subquery()
-                .alias(f"comp_exp_{compare_experiment_index}_runs")
-            )
-            compare_experiment_traces = aliased(
-                models.Trace, name=f"comp_exp_{compare_experiment_index}_traces"
-            )
-            compare_experiment_span_costs = (
-                select(
-                    models.SpanCost.trace_rowid,
-                    func.coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
-                    func.coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label(
-                        "prompt_tokens"
-                    ),
-                    func.coalesce(func.sum(models.SpanCost.completion_tokens), 0).label(
-                        "completion_tokens"
-                    ),
-                    func.coalesce(func.sum(models.SpanCost.total_cost), 0).label("total_cost"),
-                )
-                .select_from(models.SpanCost)
-                .group_by(models.SpanCost.trace_rowid)
-                .subquery()
-                .alias(f"comp_exp_{compare_experiment_index}_span_costs")
-            )
-            compare_experiment_run_latency = LatencyMs(
-                compare_experiment_runs.c.start_time, compare_experiment_runs.c.end_time
-            ).label(f"comp_exp_{compare_experiment_index}_run_latency_ms")
-            compare_experiment_run_prompt_token_count = (
-                compare_experiment_span_costs.c.prompt_tokens
-            )
-            compare_experiment_run_completion_token_count = (
-                compare_experiment_span_costs.c.completion_tokens
-            )
-            compare_experiment_run_total_token_count = compare_experiment_span_costs.c.total_tokens
-            compare_experiment_run_total_cost = compare_experiment_span_costs.c.total_cost
-
-            query = (
-                query.add_columns(
-                    _count_rows(
-                        base_experiment_run_latency < compare_experiment_run_latency,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_increased_latency"),
-                    _count_rows(
-                        base_experiment_run_latency > compare_experiment_run_latency,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_decreased_latency"),
-                    _count_rows(
-                        base_experiment_run_latency == compare_experiment_run_latency,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_equal_latency"),
-                    _count_rows(
-                        base_experiment_run_prompt_token_count
-                        < compare_experiment_run_prompt_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_increased_prompt_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_prompt_token_count
-                        > compare_experiment_run_prompt_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_decreased_prompt_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_prompt_token_count
-                        == compare_experiment_run_prompt_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_equal_prompt_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_completion_token_count
-                        < compare_experiment_run_completion_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_increased_completion_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_completion_token_count
-                        > compare_experiment_run_completion_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_decreased_completion_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_completion_token_count
-                        == compare_experiment_run_completion_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_equal_completion_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_total_token_count
-                        < compare_experiment_run_total_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_increased_total_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_total_token_count
-                        > compare_experiment_run_total_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_decreased_total_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_total_token_count
-                        == compare_experiment_run_total_token_count,
-                    ).label(
-                        f"comp_exp_{compare_experiment_index}_num_runs_equal_total_token_count"
-                    ),
-                    _count_rows(
-                        base_experiment_run_total_cost < compare_experiment_run_total_cost,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_increased_total_cost"),
-                    _count_rows(
-                        base_experiment_run_total_cost > compare_experiment_run_total_cost,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_decreased_total_cost"),
-                    _count_rows(
-                        base_experiment_run_total_cost == compare_experiment_run_total_cost,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_equal_total_cost"),
-                )
-                .join(
-                    compare_experiment_runs,
-                    onclause=base_experiment_runs.c.dataset_example_id
-                    == compare_experiment_runs.c.dataset_example_id,
-                    isouter=True,
-                )
-                .join(
-                    compare_experiment_traces,
-                    onclause=compare_experiment_runs.c.trace_id
-                    == compare_experiment_traces.trace_id,
-                    isouter=True,
-                )
-                .join(
-                    compare_experiment_span_costs,
-                    onclause=compare_experiment_traces.id
-                    == compare_experiment_span_costs.c.trace_rowid,
-                    isouter=True,
-                )
-            )
 
-        async with info.context.db() as session:
-            result = (await session.execute(query)).first()
-        assert result is not None
-
-        num_columns_per_compare_experiment = len(query.columns) // len(compare_experiment_ids)
-        counts = []
-        for compare_experiment_index, compare_experiment_id in enumerate(compare_experiment_ids):
-            start_index = compare_experiment_index * num_columns_per_compare_experiment
-            end_index = start_index + num_columns_per_compare_experiment
-            (
-                num_runs_with_increased_latency,
-                num_runs_with_decreased_latency,
-                num_runs_with_equal_latency,
-                num_runs_with_increased_prompt_token_count,
-                num_runs_with_decreased_prompt_token_count,
-                num_runs_with_equal_prompt_token_count,
-                num_runs_with_increased_completion_token_count,
-                num_runs_with_decreased_completion_token_count,
-                num_runs_with_equal_completion_token_count,
-                num_runs_with_increased_total_token_count,
-                num_runs_with_decreased_total_token_count,
-                num_runs_with_equal_total_token_count,
-                num_runs_with_increased_total_cost,
-                num_runs_with_decreased_total_cost,
-                num_runs_with_equal_total_cost,
-            ) = result[start_index:end_index]
-            counts.append(
-                CompareExperimentRunMetricCounts(
-                    compare_experiment_id=compare_experiment_id,
-                    latency=MetricCounts(
-                        num_increases=num_runs_with_increased_latency,
-                        num_decreases=num_runs_with_decreased_latency,
-                        num_equal=num_runs_with_equal_latency,
-                    ),
-                    prompt_token_count=MetricCounts(
-                        num_increases=num_runs_with_increased_prompt_token_count,
-                        num_decreases=num_runs_with_decreased_prompt_token_count,
-                        num_equal=num_runs_with_equal_prompt_token_count,
-                    ),
-                    completion_token_count=MetricCounts(
-                        num_increases=num_runs_with_increased_completion_token_count,
-                        num_decreases=num_runs_with_decreased_completion_token_count,
-                        num_equal=num_runs_with_equal_completion_token_count,
-                    ),
-                    total_token_count=MetricCounts(
-                        num_increases=num_runs_with_increased_total_token_count,
-                        num_decreases=num_runs_with_decreased_total_token_count,
-                        num_equal=num_runs_with_equal_total_token_count,
-                    ),
-                    total_cost=MetricCounts(
-                        num_increases=num_runs_with_increased_total_cost,
-                        num_decreases=num_runs_with_decreased_total_cost,
-                        num_equal=num_runs_with_equal_total_cost,
-                    ),
-                )
-            )
-        return counts
-
-    @strawberry.field
-    async def compare_experiment_run_annotation_metric_counts(
-        self,
-        info: Info[Context, None],
-        base_experiment_id: GlobalID,
-        compare_experiment_ids: list[GlobalID],
-    ) -> list[CompareExperimentRunAnnotationMetricCounts]:
-        if base_experiment_id in compare_experiment_ids:
-            raise BadRequest("Compare experiment IDs cannot contain the base experiment ID")
-        if not compare_experiment_ids:
-            raise BadRequest("At least one compare experiment ID must be provided")
-        if len(set(compare_experiment_ids)) < len(compare_experiment_ids):
-            raise BadRequest("Compare experiment IDs must be unique")
-
-        try:
-            base_experiment_rowid = from_global_id_with_expected_type(
-                base_experiment_id, models.Experiment.__name__
-            )
-        except ValueError:
-            raise BadRequest(f"Invalid base experiment ID: {base_experiment_id}")
-
-        compare_experiment_rowids = []
-        for compare_experiment_id in compare_experiment_ids:
-            try:
-                compare_experiment_rowids.append(
-                    from_global_id_with_expected_type(
-                        compare_experiment_id, models.Experiment.__name__
-                    )
-                )
-            except ValueError:
-                raise BadRequest(f"Invalid compare experiment ID: {compare_experiment_id}")
-
-        base_experiment_runs = (
-            select(models.ExperimentRun)
-            .where(
-                models.ExperimentRun.experiment_id == base_experiment_rowid,
+        comparisons_query = (
+            select(
+                func.count().label("num_base_experiment_runs"),
+                _comparison_count_expression(
+                    base_column=base_experiment_run_latency,
+                    compare_column=compare_experiment_runs.c.min_latency_ms,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_latency_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_run_latency,
+                    compare_column=compare_experiment_runs.c.min_latency_ms,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_latency_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_run_latency,
+                    compare_column=compare_experiment_runs.c.min_latency_ms,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_latency_is_equal"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.total_tokens,
+                    compare_column=compare_experiment_runs.c.min_total_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_total_token_count_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.total_tokens,
+                    compare_column=compare_experiment_runs.c.min_total_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_total_token_count_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.total_tokens,
+                    compare_column=compare_experiment_runs.c.min_total_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_total_token_count_is_equal"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.prompt_tokens,
+                    compare_column=compare_experiment_runs.c.min_prompt_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_prompt_token_count_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.prompt_tokens,
+                    compare_column=compare_experiment_runs.c.min_prompt_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_prompt_token_count_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.prompt_tokens,
+                    compare_column=compare_experiment_runs.c.min_prompt_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_prompt_token_count_is_equal"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.completion_tokens,
+                    compare_column=compare_experiment_runs.c.min_completion_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_completion_token_count_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.completion_tokens,
+                    compare_column=compare_experiment_runs.c.min_completion_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_completion_token_count_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.completion_tokens,
+                    compare_column=compare_experiment_runs.c.min_completion_tokens,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_completion_token_count_is_equal"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.total_cost,
+                    compare_column=compare_experiment_runs.c.min_total_cost,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_total_cost_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.total_cost,
+                    compare_column=compare_experiment_runs.c.min_total_cost,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_total_cost_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.total_cost,
+                    compare_column=compare_experiment_runs.c.min_total_cost,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_total_cost_is_equal"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.prompt_cost,
+                    compare_column=compare_experiment_runs.c.min_prompt_cost,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_prompt_cost_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.prompt_cost,
+                    compare_column=compare_experiment_runs.c.min_prompt_cost,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_prompt_cost_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.prompt_cost,
+                    compare_column=compare_experiment_runs.c.min_prompt_cost,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_prompt_cost_is_equal"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.completion_cost,
+                    compare_column=compare_experiment_runs.c.min_completion_cost,
+                    optimization_direction="minimize",
+                    comparison_type="improvement",
+                ).label("num_completion_cost_improved"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.completion_cost,
+                    compare_column=compare_experiment_runs.c.min_completion_cost,
+                    optimization_direction="minimize",
+                    comparison_type="regression",
+                ).label("num_completion_cost_regressed"),
+                _comparison_count_expression(
+                    base_column=base_experiment_runs.c.completion_cost,
+                    compare_column=compare_experiment_runs.c.min_completion_cost,
+                    optimization_direction="minimize",
+                    comparison_type="equality",
+                ).label("num_completion_cost_is_equal"),
             )
-            .subquery()
-            .alias("base_experiment_runs")
-        )
-        base_experiment_run_annotations = aliased(
-            models.ExperimentRunAnnotation, name="base_experiment_run_annotations"
-        )
-        query = (
-            select(base_experiment_run_annotations.name)
             .select_from(base_experiment_runs)
             .join(
-                base_experiment_run_annotations,
-                onclause=base_experiment_runs.c.id
-                == base_experiment_run_annotations.experiment_run_id,
+                compare_experiment_runs,
+                onclause=base_experiment_runs.c.dataset_example_id
+                == compare_experiment_runs.c.dataset_example_id,
                 isouter=True,
             )
-            .group_by(base_experiment_run_annotations.name)
-            .order_by(base_experiment_run_annotations.name)
         )
-        for compare_experiment_index, compare_experiment_rowid in enumerate(
-            compare_experiment_rowids
-        ):
-            compare_experiment_runs = (
-                select(models.ExperimentRun)
-                .where(
-                    models.ExperimentRun.experiment_id == compare_experiment_rowid,
-                )
-                .subquery()
-                .alias(f"comp_exp_{compare_experiment_index}_runs")
-            )
-            compare_experiment_run_annotations = aliased(
-                models.ExperimentRunAnnotation,
-                name=f"comp_exp_{compare_experiment_index}_run_annotations",
-            )
-            query = (
-                query.add_columns(
-                    _count_rows(
-                        base_experiment_run_annotations.score
-                        < compare_experiment_run_annotations.score,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_increased_score"),
-                    _count_rows(
-                        base_experiment_run_annotations.score
-                        > compare_experiment_run_annotations.score,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_decreased_score"),
-                    _count_rows(
-                        base_experiment_run_annotations.score
-                        == compare_experiment_run_annotations.score,
-                    ).label(f"comp_exp_{compare_experiment_index}_num_runs_equal_score"),
-                )
-                .join(
-                    compare_experiment_runs,
-                    onclause=base_experiment_runs.c.dataset_example_id
-                    == compare_experiment_runs.c.dataset_example_id,
-                    isouter=True,
-                )
-                .join(
-                    compare_experiment_run_annotations,
-                    onclause=compare_experiment_runs.c.id
-                    == compare_experiment_run_annotations.experiment_run_id,
-                    isouter=True,
-                )
-                .where(
-                    base_experiment_run_annotations.name == compare_experiment_run_annotations.name
-                )
-            )
+
         async with info.context.db() as session:
-            result = (await session.execute(query)).all()
+            result = (await session.execute(comparisons_query)).first()
         assert result is not None
-        num_columns_per_compare_experiment = (len(query.columns) - 1) // len(compare_experiment_ids)
-        metric_counts = []
-        for record in result:
-            annotation_name, *counts = record
-            for compare_experiment_index, compare_experiment_id in enumerate(
-                compare_experiment_ids
-            ):
-                start_index = compare_experiment_index * num_columns_per_compare_experiment
-                end_index = start_index + num_columns_per_compare_experiment
-                (
-                    num_runs_with_increased_score,
-                    num_runs_with_decreased_score,
-                    num_runs_with_equal_score,
-                ) = counts[start_index:end_index]
-                metric_counts.append(
-                    CompareExperimentRunAnnotationMetricCounts(
-                        annotation_name=annotation_name,
-                        compare_experiment_id=compare_experiment_id,
-                        num_increases=num_runs_with_increased_score,
-                        num_decreases=num_runs_with_decreased_score,
-                        num_equal=num_runs_with_equal_score,
-                    )
-                )
-        return metric_counts
+
+        return ExperimentRunMetricComparisons(
+            latency=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_latency_improved,
+                num_runs_regressed=result.num_latency_regressed,
+                num_runs_equal=result.num_latency_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+            total_token_count=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_total_token_count_improved,
+                num_runs_regressed=result.num_total_token_count_regressed,
+                num_runs_equal=result.num_total_token_count_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+            prompt_token_count=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_prompt_token_count_improved,
+                num_runs_regressed=result.num_prompt_token_count_regressed,
+                num_runs_equal=result.num_prompt_token_count_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+            completion_token_count=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_completion_token_count_improved,
+                num_runs_regressed=result.num_completion_token_count_regressed,
+                num_runs_equal=result.num_completion_token_count_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+            total_cost=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_total_cost_improved,
+                num_runs_regressed=result.num_total_cost_regressed,
+                num_runs_equal=result.num_total_cost_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+            prompt_cost=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_prompt_cost_improved,
+                num_runs_regressed=result.num_prompt_cost_regressed,
+                num_runs_equal=result.num_prompt_cost_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+            completion_cost=ExperimentRunMetricComparison(
+                num_runs_improved=result.num_completion_cost_improved,
+                num_runs_regressed=result.num_completion_cost_regressed,
+                num_runs_equal=result.num_completion_cost_is_equal,
+                num_total_runs=result.num_base_experiment_runs,
+            ),
+        )
 
     @strawberry.field
     async def validate_experiment_run_filter_condition(
@@ -959,136 +887,51 @@ 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}")
+            return ExperimentRepeatedRunGroup(
+                experiment_rowid=experiment_rowid,
+                dataset_example_rowid=dataset_example_rowid,
+            )
+
+        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)
         elif type_name == "EmbeddingDimension":
             embedding_dimension = info.context.model.embedding_dimensions[node_id]
             return to_gql_embedding_dimension(node_id, embedding_dimension)
-        elif type_name == "Project":
-            project_stmt = select(models.Project).filter_by(id=node_id)
-            async with info.context.db() as session:
-                project = await session.scalar(project_stmt)
-            if project is None:
-                raise NotFound(f"Unknown project: {id}")
-            return Project(
-                project_rowid=project.id,
-                db_project=project,
-            )
-        elif type_name == "Trace":
-            trace_stmt = select(models.Trace).filter_by(id=node_id)
-            async with info.context.db() as session:
-                trace = await session.scalar(trace_stmt)
-            if trace is None:
-                raise NotFound(f"Unknown trace: {id}")
-            return Trace(trace_rowid=trace.id, db_trace=trace)
+        elif type_name == Project.__name__:
+            return Project(id=node_id)
+        elif type_name == Trace.__name__:
+            return Trace(id=node_id)
         elif type_name == Span.__name__:
-            span_stmt = (
-                select(models.Span)
-                .options(
-                    joinedload(models.Span.trace, innerjoin=True).load_only(models.Trace.trace_id)
-                )
-                .where(models.Span.id == node_id)
-            )
-            async with info.context.db() as session:
-                span = await session.scalar(span_stmt)
-            if span is None:
-                raise NotFound(f"Unknown span: {id}")
-            return Span(span_rowid=span.id, db_span=span)
+            return Span(id=node_id)
         elif type_name == Dataset.__name__:
-            dataset_stmt = select(models.Dataset).where(models.Dataset.id == node_id)
-            async with info.context.db() as session:
-                if (dataset := await session.scalar(dataset_stmt)) is None:
-                    raise NotFound(f"Unknown dataset: {id}")
-            return to_gql_dataset(dataset)
+            return Dataset(id=node_id)
         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",
-                        )
-                    )
-                )
-            if not example:
-                raise NotFound(f"Unknown dataset example: {id}")
-            return DatasetExample(
-                id_attr=example.id,
-                created_at=example.created_at,
-            )
+            return DatasetExample(id=node_id)
+        elif type_name == DatasetSplit.__name__:
+            return DatasetSplit(id=node_id)
         elif type_name == Experiment.__name__:
-            async with info.context.db() as session:
-                experiment = await session.scalar(
-                    select(models.Experiment).where(models.Experiment.id == node_id)
-                )
-            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 Experiment(id=node_id)
         elif type_name == ExperimentRun.__name__:
-            async with info.context.db() as session:
-                if not (
-                    run := await session.scalar(
-                        select(models.ExperimentRun)
-                        .where(models.ExperimentRun.id == node_id)
-                        .options(
-                            joinedload(models.ExperimentRun.trace).load_only(models.Trace.trace_id)
-                        )
-                    )
-                ):
-                    raise NotFound(f"Unknown experiment run: {id}")
-            return to_gql_experiment_run(run)
+            return ExperimentRun(id=node_id)
         elif type_name == User.__name__:
             if int((user := info.context.user).identity) != node_id and not user.is_admin:
                 raise Unauthorized(MSG_ADMIN_ONLY)
-            async with info.context.db() as session:
-                if not (
-                    user := await session.scalar(
-                        select(models.User).where(models.User.id == node_id)
-                    )
-                ):
-                    raise NotFound(f"Unknown user: {id}")
-            return to_gql_user(user)
+            return User(id=node_id)
         elif type_name == ProjectSession.__name__:
-            async with info.context.db() as session:
-                if not (
-                    project_session := await session.scalar(
-                        select(models.ProjectSession).filter_by(id=node_id)
-                    )
-                ):
-                    raise NotFound(f"Unknown user: {id}")
-            return to_gql_project_session(project_session)
+            return ProjectSession(id=node_id)
         elif type_name == Prompt.__name__:
-            async with info.context.db() as session:
-                if orm_prompt := await session.scalar(
-                    select(models.Prompt).where(models.Prompt.id == node_id)
-                ):
-                    return to_gql_prompt_from_orm(orm_prompt)
-                else:
-                    raise NotFound(f"Unknown prompt: {id}")
+            return Prompt(id=node_id)
         elif type_name == PromptVersion.__name__:
             async with info.context.db() as session:
                 if orm_prompt_version := await session.scalar(
@@ -1098,51 +941,17 @@ class Query:
                 else:
                     raise NotFound(f"Unknown prompt version: {id}")
         elif type_name == PromptLabel.__name__:
-            async with info.context.db() as session:
-                if not (
-                    prompt_label := await session.scalar(
-                        select(models.PromptLabel).where(models.PromptLabel.id == node_id)
-                    )
-                ):
-                    raise NotFound(f"Unknown prompt label: {id}")
-            return to_gql_prompt_label(prompt_label)
+            return PromptLabel(id=node_id)
         elif type_name == PromptVersionTag.__name__:
-            async with info.context.db() as session:
-                if not (prompt_version_tag := await session.get(models.PromptVersionTag, node_id)):
-                    raise NotFound(f"Unknown prompt version tag: {id}")
-            return to_gql_prompt_version_tag(prompt_version_tag)
+            return PromptVersionTag(id=node_id)
         elif type_name == ProjectTraceRetentionPolicy.__name__:
-            async with info.context.db() as session:
-                db_policy = await session.scalar(
-                    select(models.ProjectTraceRetentionPolicy).filter_by(id=node_id)
-                )
-                if not db_policy:
-                    raise NotFound(f"Unknown project trace retention policy: {id}")
-            return ProjectTraceRetentionPolicy(id=db_policy.id, db_policy=db_policy)
+            return ProjectTraceRetentionPolicy(id=node_id)
         elif type_name == SpanAnnotation.__name__:
-            async with info.context.db() as session:
-                span_annotation = await session.get(models.SpanAnnotation, node_id)
-                if not span_annotation:
-                    raise NotFound(f"Unknown span annotation: {id}")
-            return to_gql_span_annotation(span_annotation)
+            return SpanAnnotation(id=node_id)
         elif type_name == TraceAnnotation.__name__:
-            async with info.context.db() as session:
-                trace_annotation = await session.get(models.TraceAnnotation, node_id)
-                if not trace_annotation:
-                    raise NotFound(f"Unknown trace annotation: {id}")
-            return to_gql_trace_annotation(trace_annotation)
+            return TraceAnnotation(id=node_id)
         elif type_name == GenerativeModel.__name__:
-            async with info.context.db() as session:
-                stmt = (
-                    select(models.GenerativeModel)
-                    .where(models.GenerativeModel.deleted_at.is_(None))
-                    .where(models.GenerativeModel.id == node_id)
-                    .options(joinedload(models.GenerativeModel.token_prices))
-                )
-                model = await session.scalar(stmt)
-                if not model:
-                    raise NotFound(f"Unknown model: {id}")
-            return to_gql_generative_model(model)
+            return GenerativeModel(id=node_id)
         raise NotFound(f"Unknown node type: {type_name}")
 
     @strawberry.field
@@ -1154,16 +963,7 @@ class Query:
             return None
         if isinstance(user, UnauthenticatedUser):
             return None
-        async with info.context.db() as session:
-            if (
-                user := await session.scalar(
-                    select(models.User)
-                    .where(models.User.id == int(user.identity))
-                    .options(joinedload(models.User.role))
-                )
-            ) is None:
-                return None
-        return to_gql_user(user)
+        return User(id=int(user.identity))
 
     @strawberry.field
     async def prompts(
@@ -1174,6 +974,7 @@ class Query:
         after: Optional[CursorString] = UNSET,
         before: Optional[CursorString] = UNSET,
         filter: Optional[PromptFilter] = UNSET,
+        labelIds: Optional[list[GlobalID]] = UNSET,
     ) -> Connection[Prompt]:
         args = ConnectionArgs(
             first=first,
@@ -1190,9 +991,21 @@ class Query:
             stmt = stmt.where(column.ilike(f"%{filter.value}%")).order_by(
                 models.Prompt.updated_at.desc()
             )
+        if labelIds:
+            stmt = stmt.join(models.PromptPromptLabel).where(
+                models.PromptPromptLabel.prompt_label_id.in_(
+                    from_global_id_with_expected_type(
+                        global_id=label_id, expected_type_name="PromptLabel"
+                    )
+                    for label_id in labelIds
+                )
+            )
+            stmt = stmt.distinct()
         async with info.context.db() as session:
             orm_prompts = await session.stream_scalars(stmt)
-            data = [to_gql_prompt_from_orm(orm_prompt) async for orm_prompt in orm_prompts]
+            data = [
+                Prompt(id=orm_prompt.id, db_record=orm_prompt) async for orm_prompt in orm_prompts
+            ]
             return connection_from_list(
                 data=data,
                 args=args,
@@ -1215,7 +1028,58 @@ class Query:
         )
         async with info.context.db() as session:
             prompt_labels = await session.stream_scalars(select(models.PromptLabel))
-            data = [to_gql_prompt_label(prompt_label) async for prompt_label in prompt_labels]
+            data = [
+                PromptLabel(id=prompt_label.id, db_record=prompt_label)
+                async for prompt_label in prompt_labels
+            ]
+            return connection_from_list(
+                data=data,
+                args=args,
+            )
+
+    @strawberry.field
+    async def dataset_labels(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[DatasetLabel]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        async with info.context.db() as session:
+            dataset_labels = await session.scalars(
+                select(models.DatasetLabel).order_by(models.DatasetLabel.name.asc())
+            )
+        data = [
+            DatasetLabel(id=dataset_label.id, db_record=dataset_label)
+            for dataset_label in dataset_labels
+        ]
+        return connection_from_list(data=data, args=args)
+
+    @strawberry.field
+    async def dataset_splits(
+        self,
+        info: Info[Context, None],
+        first: Optional[int] = 50,
+        last: Optional[int] = UNSET,
+        after: Optional[CursorString] = UNSET,
+        before: Optional[CursorString] = UNSET,
+    ) -> Connection[DatasetSplit]:
+        args = ConnectionArgs(
+            first=first,
+            after=after if isinstance(after, CursorString) else None,
+            last=last,
+            before=before if isinstance(before, CursorString) else None,
+        )
+        async with info.context.db() as session:
+            splits = await session.stream_scalars(select(models.DatasetSplit))
+            data = [DatasetSplit(id=split.id, db_record=split) async for split in splits]
             return connection_from_list(
                 data=data,
                 args=args,
@@ -1486,7 +1350,7 @@ class Query:
         async with info.context.db() as session:
             span_rowid = await session.scalar(stmt)
         if span_rowid:
-            return Span(span_rowid=span_rowid)
+            return Span(id=span_rowid)
         return None
 
     @strawberry.field
@@ -1499,7 +1363,7 @@ class Query:
         async with info.context.db() as session:
             trace_rowid = await session.scalar(stmt)
         if trace_rowid:
-            return Trace(trace_rowid=trace_rowid)
+            return Trace(id=trace_rowid)
         return None
 
     @strawberry.field
@@ -1512,7 +1376,7 @@ class Query:
         async with info.context.db() as session:
             session_row = await session.scalar(stmt)
         if session_row:
-            return to_gql_project_session(session_row)
+            return ProjectSession(id=session_row.id, db_record=session_row)
         return None
 
 
@@ -1550,16 +1414,36 @@ def _longest_matching_prefix(s: str, prefixes: Iterable[str]) -> str:
     return longest
 
 
-def _count_rows(
-    condition: ColumnElement[Any],
-) -> ColumnElement[Any]:
+def _comparison_count_expression(
+    *,
+    base_column: ColumnElement[Any],
+    compare_column: ColumnElement[Any],
+    optimization_direction: Literal["maximize", "minimize"],
+    comparison_type: Literal["improvement", "regression", "equality"],
+) -> ColumnElement[int]:
     """
-    Returns an expression that counts the number of rows satisfying the condition.
+    Given a base and compare column, returns an expression counting the number of
+    improvements, regressions, or equalities given the optimization direction.
     """
+    if optimization_direction == "maximize":
+        raise NotImplementedError
+
+    if comparison_type == "improvement":
+        condition = compare_column > base_column
+    elif comparison_type == "regression":
+        condition = compare_column < base_column
+    elif comparison_type == "equality":
+        condition = compare_column == base_column
+    else:
+        assert_never(comparison_type)
+
     return func.coalesce(
         func.sum(
             case(
-                (condition, 1),
+                (
+                    condition,
+                    1,
+                ),
                 else_=0,
             )
         ),