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

phoenix/db/pg_config.py

@@ -10,12 +10,14 @@ from typing_extensions import assert_never
 def get_pg_config(
     url: URL,
     driver: Literal["psycopg", "asyncpg"],
+    enforce_ssl: bool = False,
 ) -> tuple[URL, dict[str, Any]]:
     """Convert SQLAlchemy URL to driver-specific configuration.
 
     Args:
         url: SQLAlchemy URL
         driver: "psycopg" or "asyncpg"
+        enforce_ssl: If True, ensure SSL is enabled (required for AWS RDS IAM auth)
 
     Returns:
         Tuple of (base_url, connect_args):
@@ -26,6 +28,14 @@ def get_pg_config(
     query = url.query
     ssl_args = _get_ssl_args(query)
 
+    if enforce_ssl and not ssl_args:
+        ssl_args = {"sslmode": "require"}
+    elif enforce_ssl and ssl_args.get("sslmode") == "disable":
+        raise ValueError(
+            "SSL cannot be disabled when using AWS RDS IAM authentication. "
+            "Remove 'sslmode=disable' from the connection string."
+        )
+
     # Create base URL without SSL parameters
     base_url = url.set(
         drivername=f"postgresql+{driver}",

phoenix/db/types/trace_retention.py

@@ -7,6 +7,7 @@ import sqlalchemy as sa
 from pydantic import AfterValidator, BaseModel, Field, RootModel
 from sqlalchemy import func
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.roles import InElementRole
 
 from phoenix.utilities import hour_of_week
 
@@ -28,7 +29,7 @@ class _MaxCount(BaseModel):
 
     def max_count_filter(
         self,
-        project_rowids: Union[Iterable[int], sa.ScalarSelect[int]],
+        project_rowids: Union[Iterable[int], InElementRole],
     ) -> sa.ColumnElement[bool]:
         if self.max_count <= 0:
             return sa.literal(False)
@@ -56,7 +57,7 @@ class MaxDaysRule(_MaxDays, BaseModel):
     async def delete_traces(
         self,
         session: AsyncSession,
-        project_rowids: Union[Iterable[int], sa.ScalarSelect[int]],
+        project_rowids: Union[Iterable[int], InElementRole],
     ) -> set[int]:
         if self.max_days <= 0:
             return set()
@@ -80,7 +81,7 @@ class MaxCountRule(_MaxCount, BaseModel):
     async def delete_traces(
         self,
         session: AsyncSession,
-        project_rowids: Union[Iterable[int], sa.ScalarSelect[int]],
+        project_rowids: Union[Iterable[int], InElementRole],
     ) -> set[int]:
         if self.max_count <= 0:
             return set()
@@ -104,7 +105,7 @@ class MaxDaysOrCountRule(_MaxDays, _MaxCount, BaseModel):
     async def delete_traces(
         self,
         session: AsyncSession,
-        project_rowids: Union[Iterable[int], sa.ScalarSelect[int]],
+        project_rowids: Union[Iterable[int], InElementRole],
     ) -> set[int]:
         if self.max_days <= 0 and self.max_count <= 0:
             return set()
@@ -130,7 +131,7 @@ class TraceRetentionRule(RootModel[Union[MaxDaysRule, MaxCountRule, MaxDaysOrCou
     async def delete_traces(
         self,
         session: AsyncSession,
-        project_rowids: Union[Iterable[int], sa.ScalarSelect[int]],
+        project_rowids: Union[Iterable[int], InElementRole],
     ) -> set[int]:
         return await self.root.delete_traces(session, project_rowids)
 
@@ -199,7 +200,7 @@ class TraceRetentionCronExpression(RootModel[str]):
 
 def _parse_field(field: str, min_val: int, max_val: int) -> set[int]:
     """
-    Parse a cron field and return the set of matching values.
+    Parses a cron field and returns the set of matching values.
 
     Args:
         field (str): The cron field to parse

phoenix/experiments/functions.py

@@ -109,6 +109,64 @@ def _phoenix_clients() -> tuple[httpx.Client, httpx.AsyncClient]:
     )
 
 
+def _get_all_experiment_runs(
+    client: httpx.Client,
+    experiment_id: str,
+    page_size: int = 50,
+) -> list[ExperimentRun]:
+    """
+    Fetch all experiment runs using pagination to handle large datasets.
+
+    Args:
+        client: The HTTP client to use for requests.
+        experiment_id: The ID of the experiment.
+        page_size: Number of runs to fetch per page. Defaults to 50.
+
+    Returns:
+        List of all experiment runs as ExperimentRun objects.
+    """
+    all_runs: list[dict[str, Any]] = []
+    cursor = None
+
+    while True:
+        params: dict[str, Any] = {"limit": page_size}
+        if cursor:
+            params["cursor"] = cursor
+
+        try:
+            response = client.get(
+                f"v1/experiments/{experiment_id}/runs",
+                params=params,
+            )
+            response.raise_for_status()
+            data = response.json()
+
+            runs = data["data"]
+            all_runs.extend(runs)
+
+            # Check if there are more pages
+            cursor = data.get("next_cursor")
+            if not cursor:
+                break
+
+        except HTTPStatusError as e:
+            if e.response.status_code == 404:
+                # Experiment doesn't exist - treat as empty result
+                break
+            else:
+                raise
+
+    # Convert dicts to ExperimentRun objects
+    experiment_runs: list[ExperimentRun] = []
+    for run in all_runs:
+        # Parse datetime strings
+        run["start_time"] = datetime.fromisoformat(run["start_time"])
+        run["end_time"] = datetime.fromisoformat(run["end_time"])
+        experiment_runs.append(ExperimentRun.from_dict(run))
+
+    return experiment_runs
+
+
 Evaluators: TypeAlias = Union[
     ExperimentEvaluator,
     Sequence[ExperimentEvaluator],
@@ -231,7 +289,7 @@ def run_experiment(
     }
     if not dry_run:
         experiment_response = sync_client.post(
-            f"/v1/datasets/{normalized_dataset.id}/experiments",
+            f"v1/datasets/{normalized_dataset.id}/experiments",
             json=payload,
         )
         experiment_response.raise_for_status()
@@ -303,7 +361,7 @@ def run_experiment(
                 try:
                     # Try to create the run directly
                     resp = sync_client.post(
-                        f"/v1/experiments/{experiment.id}/runs", json=jsonify(exp_run)
+                        f"v1/experiments/{experiment.id}/runs", json=jsonify(exp_run)
                     )
                     resp.raise_for_status()
                     exp_run = replace(exp_run, id=resp.json()["data"]["id"])
@@ -381,7 +439,7 @@ def run_experiment(
             try:
                 # Try to create the run directly
                 resp = sync_client.post(
-                    f"/v1/experiments/{experiment.id}/runs", json=jsonify(exp_run)
+                    f"v1/experiments/{experiment.id}/runs", json=jsonify(exp_run)
                 )
                 resp.raise_for_status()
                 exp_run = replace(exp_run, id=resp.json()["data"]["id"])
@@ -420,7 +478,7 @@ def run_experiment(
                         None,
                         functools.partial(
                             sync_client.post,
-                            url=f"/v1/experiments/{experiment.id}/runs",
+                            url=f"v1/experiments/{experiment.id}/runs",
                             json=jsonify(exp_run),
                         ),
                     )
@@ -498,7 +556,7 @@ def run_experiment(
                     None,
                     functools.partial(
                         sync_client.post,
-                        url=f"/v1/experiments/{experiment.id}/runs",
+                        url=f"v1/experiments/{experiment.id}/runs",
                         json=jsonify(exp_run),
                     ),
                 )
@@ -548,13 +606,7 @@ def run_experiment(
 
     # Get the final state of runs from the database
     if not dry_run:
-        all_runs = sync_client.get(f"/v1/experiments/{experiment.id}/runs").json()["data"]
-        task_runs = []
-        for run in all_runs:
-            # Parse datetime strings
-            run["start_time"] = datetime.fromisoformat(run["start_time"])
-            run["end_time"] = datetime.fromisoformat(run["end_time"])
-            task_runs.append(ExperimentRun.from_dict(run))
+        task_runs = _get_all_experiment_runs(sync_client, experiment.id)
 
         # Check if we got all expected runs
         expected_runs = len(normalized_dataset.examples) * repetitions
@@ -613,16 +665,14 @@ def evaluate_experiment(
     else:
         dataset = Dataset.from_dict(
             sync_client.get(
-                f"/v1/datasets/{dataset_id}/examples",
+                f"v1/datasets/{dataset_id}/examples",
                 params={"version_id": str(dataset_version_id)},
             ).json()["data"]
         )
         if not dataset.examples:
             raise ValueError(f"Dataset has no examples: {dataset_id=}, {dataset_version_id=}")
-        experiment_runs = {
-            exp_run["id"]: ExperimentRun.from_dict(exp_run)
-            for exp_run in sync_client.get(f"/v1/experiments/{experiment.id}/runs").json()["data"]
-        }
+        all_runs = _get_all_experiment_runs(sync_client, experiment.id)
+        experiment_runs = {exp_run.id: exp_run for exp_run in all_runs}
         if not experiment_runs:
             raise ValueError("Experiment has not been run")
         params = ExperimentParameters(n_examples=len(dataset.examples))
@@ -715,7 +765,7 @@ def evaluate_experiment(
             trace_id=_str_trace_id(span.get_span_context().trace_id),  # type: ignore[no-untyped-call]
         )
         if not dry_run:
-            resp = sync_client.post("/v1/experiment_evaluations", json=jsonify(eval_run))
+            resp = sync_client.post("v1/experiment_evaluations", json=jsonify(eval_run))
             resp.raise_for_status()
             eval_run = replace(eval_run, id=resp.json()["data"]["id"])
         return eval_run
@@ -777,7 +827,7 @@ def evaluate_experiment(
                 None,
                 functools.partial(
                     sync_client.post,
-                    url="/v1/experiment_evaluations",
+                    url="v1/experiment_evaluations",
                     json=jsonify(eval_run),
                 ),
             )

phoenix/inferences/inferences.py

@@ -13,11 +13,10 @@ from pandas import DataFrame, Series, Timestamp, read_parquet
 from pandas.api.types import (
     is_numeric_dtype,
 )
-from typing_extensions import TypeAlias
+from typing_extensions import TypeAlias, deprecated
 
 from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX, INFERENCES_DIR
 from phoenix.datetime_utils import normalize_timestamps
-from phoenix.utilities.deprecation import deprecated
 
 from . import errors as err
 from .schema import (

phoenix/server/api/auth.py

@@ -22,6 +22,15 @@ class IsNotReadOnly(Authorization):
         return not info.context.read_only
 
 
+class IsNotViewer(Authorization):
+    message = "Viewers cannot perform this action"
+
+    def has_permission(self, source: Any, info: Info, **kwargs: Any) -> bool:
+        if not info.context.auth_enabled:
+            return True
+        return isinstance((user := info.context.user), PhoenixUser) and not user.is_viewer
+
+
 class IsLocked(BasePermission):
     """
     Permission class that restricts data-modifying operations when insufficient storage.

phoenix/server/api/auth_messages.py

@@ -0,0 +1,46 @@
+# ruff: noqa: E501
+"""
+Authentication error and success message codes.
+
+These codes are used in authentication flows to safely communicate status
+to users via query parameters. Using codes instead of raw messages prevents
+social engineering and phishing attacks.
+
+The messages are passed to the frontend via window.Config to ensure a single
+source of truth between backend and frontend.
+"""
+
+from types import MappingProxyType
+from typing import Literal, Mapping, get_args
+
+# Error code type - used for type hints in redirect functions
+AuthErrorCode = Literal[
+    "unknown_idp",
+    "auth_failed",
+    "invalid_state",
+    "unsafe_return_url",
+    "oauth_error",
+    "no_oidc_support",
+    "missing_email_scope",
+    "email_in_use",
+    "sign_in_not_allowed",
+]
+
+# Error messages - passed to frontend via window.Config.authErrorMessages
+# Backend generates these codes when redirecting users after OAuth errors
+AUTH_ERROR_MESSAGES: Mapping[AuthErrorCode, str] = MappingProxyType(
+    {
+        "unknown_idp": "Unknown identity provider.",
+        "auth_failed": "Authentication failed. Please contact your administrator.",
+        "invalid_state": "Invalid authentication state. Please try again.",
+        "unsafe_return_url": "Invalid return URL. Please try again.",
+        "oauth_error": "Authentication failed. Please try again.",
+        "no_oidc_support": "Your identity provider does not appear to support OpenID Connect. Please contact your administrator.",
+        "missing_email_scope": "Please ensure your identity provider is configured to use the 'email' scope.",
+        "email_in_use": "An account with this email already exists.",
+        "sign_in_not_allowed": "Sign in is not allowed. Please contact your administrator.",
+    }
+)
+
+# Runtime assertion to ensure AUTH_ERROR_MESSAGES keys match AuthErrorCode Literal values
+assert set(AUTH_ERROR_MESSAGES.keys()) == set(get_args(AuthErrorCode))

phoenix/server/api/context.py

@@ -17,17 +17,25 @@ from phoenix.db import models
 from phoenix.server.api.dataloaders import (
     AnnotationConfigsByProjectDataLoader,
     AnnotationSummaryDataLoader,
+    AverageExperimentRepeatedRunGroupLatencyDataLoader,
     AverageExperimentRunLatencyDataLoader,
     CacheForDataLoaders,
+    DatasetDatasetSplitsDataLoader,
     DatasetExampleRevisionsDataLoader,
+    DatasetExamplesAndVersionsByExperimentRunDataLoader,
     DatasetExampleSpansDataLoader,
+    DatasetExampleSplitsDataLoader,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
     ExperimentAnnotationSummaryDataLoader,
+    ExperimentDatasetSplitsDataLoader,
     ExperimentErrorRatesDataLoader,
+    ExperimentRepeatedRunGroupAnnotationSummariesDataLoader,
+    ExperimentRepeatedRunGroupsDataLoader,
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
+    ExperimentRunsByExperimentAndExampleDataLoader,
     ExperimentSequenceNumberDataLoader,
     LastUsedTimesByGenerativeModelIdDataLoader,
     LatencyMsQuantileDataLoader,
@@ -38,6 +46,7 @@ from phoenix.server.api.dataloaders import (
     ProjectIdsByTraceRetentionPolicyIdDataLoader,
     PromptVersionSequenceNumberDataLoader,
     RecordCountDataLoader,
+    SessionAnnotationsBySessionDataLoader,
     SessionIODataLoader,
     SessionNumTracesDataLoader,
     SessionNumTracesWithErrorDataLoader,
@@ -52,6 +61,7 @@ from phoenix.server.api.dataloaders import (
     SpanCostDetailSummaryEntriesBySpanDataLoader,
     SpanCostDetailSummaryEntriesByTraceDataLoader,
     SpanCostSummaryByExperimentDataLoader,
+    SpanCostSummaryByExperimentRepeatedRunGroupDataLoader,
     SpanCostSummaryByExperimentRunDataLoader,
     SpanCostSummaryByGenerativeModelDataLoader,
     SpanCostSummaryByProjectDataLoader,
@@ -62,12 +72,15 @@ from phoenix.server.api.dataloaders import (
     SpanProjectsDataLoader,
     TableFieldsDataLoader,
     TokenCountDataLoader,
+    TokenPricesByModelDataLoader,
+    TraceAnnotationsByTraceDataLoader,
     TraceByTraceIdsDataLoader,
     TraceRetentionPolicyIdByProjectIdDataLoader,
     TraceRootSpansDataLoader,
     UserRolesDataLoader,
     UsersDataLoader,
 )
+from phoenix.server.api.dataloaders.dataset_labels import DatasetLabelsDataLoader
 from phoenix.server.bearer_auth import PhoenixUser
 from phoenix.server.daemons.span_cost_calculator import SpanCostCalculator
 from phoenix.server.dml_event import DmlEvent
@@ -85,17 +98,42 @@ from phoenix.server.types import (
 class DataLoaders:
     annotation_configs_by_project: AnnotationConfigsByProjectDataLoader
     annotation_summaries: AnnotationSummaryDataLoader
+    average_experiment_repeated_run_group_latency: (
+        AverageExperimentRepeatedRunGroupLatencyDataLoader
+    )
     average_experiment_run_latency: AverageExperimentRunLatencyDataLoader
+    dataset_example_fields: TableFieldsDataLoader
     dataset_example_revisions: DatasetExampleRevisionsDataLoader
     dataset_example_spans: DatasetExampleSpansDataLoader
+    dataset_labels: DatasetLabelsDataLoader
+    dataset_label_fields: TableFieldsDataLoader
+    dataset_dataset_splits: DatasetDatasetSplitsDataLoader
+    dataset_examples_and_versions_by_experiment_run: (
+        DatasetExamplesAndVersionsByExperimentRunDataLoader
+    )
+    dataset_example_splits: DatasetExampleSplitsDataLoader
+    dataset_fields: TableFieldsDataLoader
+    dataset_split_fields: TableFieldsDataLoader
+    dataset_version_fields: TableFieldsDataLoader
+    document_annotation_fields: TableFieldsDataLoader
     document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
     document_evaluations: DocumentEvaluationsDataLoader
     document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
     experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
+    experiment_dataset_splits: ExperimentDatasetSplitsDataLoader
     experiment_error_rates: ExperimentErrorRatesDataLoader
+    experiment_fields: TableFieldsDataLoader
+    experiment_repeated_run_group_annotation_summaries: (
+        ExperimentRepeatedRunGroupAnnotationSummariesDataLoader
+    )
+    experiment_repeated_run_groups: ExperimentRepeatedRunGroupsDataLoader
+    experiment_run_annotation_fields: TableFieldsDataLoader
     experiment_run_annotations: ExperimentRunAnnotations
     experiment_run_counts: ExperimentRunCountsDataLoader
+    experiment_run_fields: TableFieldsDataLoader
+    experiment_runs_by_experiment_and_example: ExperimentRunsByExperimentAndExampleDataLoader
     experiment_sequence_number: ExperimentSequenceNumberDataLoader
+    generative_model_fields: TableFieldsDataLoader
     last_used_times_by_generative_model_id: LastUsedTimesByGenerativeModelIdDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
@@ -105,14 +143,21 @@ class DataLoaders:
     project_fields: TableFieldsDataLoader
     project_trace_retention_policy_fields: TableFieldsDataLoader
     projects_by_trace_retention_policy_id: ProjectIdsByTraceRetentionPolicyIdDataLoader
+    prompt_fields: TableFieldsDataLoader
+    prompt_label_fields: TableFieldsDataLoader
     prompt_version_sequence_number: PromptVersionSequenceNumberDataLoader
+    prompt_version_tag_fields: TableFieldsDataLoader
+    project_session_annotation_fields: TableFieldsDataLoader
+    project_session_fields: TableFieldsDataLoader
     record_counts: RecordCountDataLoader
+    session_annotations_by_session: SessionAnnotationsBySessionDataLoader
     session_first_inputs: SessionIODataLoader
     session_last_outputs: SessionIODataLoader
     session_num_traces: SessionNumTracesDataLoader
     session_num_traces_with_error: SessionNumTracesWithErrorDataLoader
     session_token_usages: SessionTokenUsagesDataLoader
     session_trace_latency_ms_quantile: SessionTraceLatencyMsQuantileDataLoader
+    span_annotation_fields: TableFieldsDataLoader
     span_annotations: SpanAnnotationsDataLoader
     span_by_id: SpanByIdDataLoader
     span_cost_by_span: SpanCostBySpanDataLoader
@@ -128,6 +173,9 @@ class DataLoaders:
     span_cost_details_by_span_cost: SpanCostDetailsBySpanCostDataLoader
     span_cost_fields: TableFieldsDataLoader
     span_cost_summary_by_experiment: SpanCostSummaryByExperimentDataLoader
+    span_cost_summary_by_experiment_repeated_run_group: (
+        SpanCostSummaryByExperimentRepeatedRunGroupDataLoader
+    )
     span_cost_summary_by_experiment_run: SpanCostSummaryByExperimentRunDataLoader
     span_cost_summary_by_generative_model: SpanCostSummaryByGenerativeModelDataLoader
     span_cost_summary_by_project: SpanCostSummaryByProjectDataLoader
@@ -138,11 +186,16 @@ class DataLoaders:
     span_fields: TableFieldsDataLoader
     span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
+    token_prices_by_model: TokenPricesByModelDataLoader
+    trace_annotation_fields: TableFieldsDataLoader
+    trace_annotations_by_trace: TraceAnnotationsByTraceDataLoader
     trace_by_trace_ids: TraceByTraceIdsDataLoader
     trace_fields: TableFieldsDataLoader
     trace_retention_policy_id_by_project_id: TraceRetentionPolicyIdByProjectIdDataLoader
     trace_root_spans: TraceRootSpansDataLoader
     user_roles: UserRolesDataLoader
+    user_api_key_fields: TableFieldsDataLoader
+    user_fields: TableFieldsDataLoader
     users: UsersDataLoader
 
 
@@ -217,3 +270,10 @@ class Context(BaseContext):
     @cached_property
     def user(self) -> PhoenixUser:
         return cast(PhoenixUser, self.get_request().user)
+
+    @cached_property
+    def user_id(self) -> Optional[int]:
+        try:
+            return int(self.user.identity)
+        except Exception:
+            return None

phoenix/server/api/dataloaders/__init__.py

@@ -6,9 +6,18 @@ from phoenix.server.api.dataloaders.span_cost_detail_summary_entries_by_project_
 
 from .annotation_configs_by_project import AnnotationConfigsByProjectDataLoader
 from .annotation_summaries import AnnotationSummaryCache, AnnotationSummaryDataLoader
+from .average_experiment_repeated_run_group_latency import (
+    AverageExperimentRepeatedRunGroupLatencyDataLoader,
+)
 from .average_experiment_run_latency import AverageExperimentRunLatencyDataLoader
+from .dataset_dataset_splits import DatasetDatasetSplitsDataLoader
 from .dataset_example_revisions import DatasetExampleRevisionsDataLoader
 from .dataset_example_spans import DatasetExampleSpansDataLoader
+from .dataset_example_splits import DatasetExampleSplitsDataLoader
+from .dataset_examples_and_versions_by_experiment_run import (
+    DatasetExamplesAndVersionsByExperimentRunDataLoader,
+)
+from .dataset_labels import DatasetLabelsDataLoader
 from .document_evaluation_summaries import (
     DocumentEvaluationSummaryCache,
     DocumentEvaluationSummaryDataLoader,
@@ -16,9 +25,17 @@ from .document_evaluation_summaries import (
 from .document_evaluations import DocumentEvaluationsDataLoader
 from .document_retrieval_metrics import DocumentRetrievalMetricsDataLoader
 from .experiment_annotation_summaries import ExperimentAnnotationSummaryDataLoader
+from .experiment_dataset_splits import ExperimentDatasetSplitsDataLoader
 from .experiment_error_rates import ExperimentErrorRatesDataLoader
+from .experiment_repeated_run_group_annotation_summaries import (
+    ExperimentRepeatedRunGroupAnnotationSummariesDataLoader,
+)
+from .experiment_repeated_run_groups import ExperimentRepeatedRunGroupsDataLoader
 from .experiment_run_annotations import ExperimentRunAnnotations
 from .experiment_run_counts import ExperimentRunCountsDataLoader
+from .experiment_runs_by_experiment_and_example import (
+    ExperimentRunsByExperimentAndExampleDataLoader,
+)
 from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .last_used_times_by_generative_model_id import LastUsedTimesByGenerativeModelIdDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
@@ -29,6 +46,7 @@ from .project_by_name import ProjectByNameDataLoader
 from .project_ids_by_trace_retention_policy_id import ProjectIdsByTraceRetentionPolicyIdDataLoader
 from .prompt_version_sequence_number import PromptVersionSequenceNumberDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
+from .session_annotations_by_session import SessionAnnotationsBySessionDataLoader
 from .session_io import SessionIODataLoader
 from .session_num_traces import SessionNumTracesDataLoader
 from .session_num_traces_with_error import SessionNumTracesWithErrorDataLoader
@@ -44,6 +62,9 @@ from .span_cost_detail_summary_entries_by_span import SpanCostDetailSummaryEntri
 from .span_cost_detail_summary_entries_by_trace import SpanCostDetailSummaryEntriesByTraceDataLoader
 from .span_cost_details_by_span_cost import SpanCostDetailsBySpanCostDataLoader
 from .span_cost_summary_by_experiment import SpanCostSummaryByExperimentDataLoader
+from .span_cost_summary_by_experiment_repeated_run_group import (
+    SpanCostSummaryByExperimentRepeatedRunGroupDataLoader,
+)
 from .span_cost_summary_by_experiment_run import SpanCostSummaryByExperimentRunDataLoader
 from .span_cost_summary_by_generative_model import SpanCostSummaryByGenerativeModelDataLoader
 from .span_cost_summary_by_project import SpanCostSummaryByProjectDataLoader, SpanCostSummaryCache
@@ -55,6 +76,8 @@ from .span_descendants import SpanDescendantsDataLoader
 from .span_projects import SpanProjectsDataLoader
 from .table_fields import TableFieldsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
+from .token_prices_by_model import TokenPricesByModelDataLoader
+from .trace_annotations_by_trace import TraceAnnotationsByTraceDataLoader
 from .trace_by_trace_ids import TraceByTraceIdsDataLoader
 from .trace_retention_policy_id_by_project_id import TraceRetentionPolicyIdByProjectIdDataLoader
 from .trace_root_spans import TraceRootSpansDataLoader
@@ -64,17 +87,26 @@ from .users import UsersDataLoader
 __all__ = [
     "AnnotationConfigsByProjectDataLoader",
     "AnnotationSummaryDataLoader",
+    "AverageExperimentRepeatedRunGroupLatencyDataLoader",
     "AverageExperimentRunLatencyDataLoader",
     "CacheForDataLoaders",
+    "DatasetDatasetSplitsDataLoader",
     "DatasetExampleRevisionsDataLoader",
     "DatasetExampleSpansDataLoader",
+    "DatasetExamplesAndVersionsByExperimentRunDataLoader",
+    "DatasetExampleSplitsDataLoader",
+    "DatasetLabelsDataLoader",
+    "ExperimentDatasetSplitsDataLoader",
     "DocumentEvaluationSummaryDataLoader",
     "DocumentEvaluationsDataLoader",
     "DocumentRetrievalMetricsDataLoader",
     "ExperimentAnnotationSummaryDataLoader",
     "ExperimentErrorRatesDataLoader",
+    "ExperimentRepeatedRunGroupsDataLoader",
+    "ExperimentRepeatedRunGroupAnnotationSummariesDataLoader",
     "ExperimentRunAnnotations",
     "ExperimentRunCountsDataLoader",
+    "ExperimentRunsByExperimentAndExampleDataLoader",
     "ExperimentSequenceNumberDataLoader",
     "LastUsedTimesByGenerativeModelIdDataLoader",
     "LatencyMsQuantileDataLoader",
@@ -85,6 +117,7 @@ __all__ = [
     "ProjectIdsByTraceRetentionPolicyIdDataLoader",
     "PromptVersionSequenceNumberDataLoader",
     "RecordCountDataLoader",
+    "SessionAnnotationsBySessionDataLoader",
     "SessionIODataLoader",
     "SessionNumTracesDataLoader",
     "SessionNumTracesWithErrorDataLoader",
@@ -99,6 +132,7 @@ __all__ = [
     "SpanCostDetailSummaryEntriesByTraceDataLoader",
     "SpanCostDetailsBySpanCostDataLoader",
     "SpanCostSummaryByExperimentDataLoader",
+    "SpanCostSummaryByExperimentRepeatedRunGroupDataLoader",
     "SpanCostSummaryByExperimentRunDataLoader",
     "SpanCostSummaryByGenerativeModelDataLoader",
     "SpanCostSummaryByProjectDataLoader",
@@ -110,6 +144,8 @@ __all__ = [
     "SpanProjectsDataLoader",
     "TableFieldsDataLoader",
     "TokenCountDataLoader",
+    "TokenPricesByModelDataLoader",
+    "TraceAnnotationsByTraceDataLoader",
     "TraceByTraceIdsDataLoader",
     "TraceRetentionPolicyIdByProjectIdDataLoader",
     "TraceRootSpansDataLoader",