arize-phoenix 11.38.0__py3-none-any.whl → 12.2.0__py3-none-any.whl

phoenix/server/api/subscriptions.py

@@ -26,6 +26,7 @@ from typing_extensions import TypeAlias, assert_never
 from phoenix.config import PLAYGROUND_PROJECT_NAME
 from phoenix.datetime_utils import local_now, normalize_datetime
 from phoenix.db import models
+from phoenix.db.helpers import insert_experiment_with_examples_snapshot
 from phoenix.server.api.auth import IsLocked, IsNotReadOnly
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest, CustomGraphQLError, NotFound
@@ -43,6 +44,7 @@ from phoenix.server.api.helpers.playground_spans import (
     get_db_trace,
     streaming_llm_span,
 )
+from phoenix.server.api.helpers.playground_users import get_user
 from phoenix.server.api.helpers.prompts.models import PromptTemplateFormat
 from phoenix.server.api.input_types.ChatCompletionInput import (
     ChatCompletionInput,
@@ -302,6 +304,7 @@ class Subscription:
                         description="Traces from prompt playground",
                     )
                 )
+            user_id = get_user(info)
             experiment = models.Experiment(
                 dataset_id=from_global_id_with_expected_type(input.dataset_id, Dataset.__name__),
                 dataset_version_id=resolved_version_id,
@@ -311,9 +314,9 @@ class Subscription:
                 repetitions=input.repetitions,
                 metadata_=input.experiment_metadata or dict(),
                 project_name=project_name,
+                user_id=user_id,
             )
-            session.add(experiment)
-            await session.flush()
+            await insert_experiment_with_examples_snapshot(session, experiment)
         yield ChatCompletionSubscriptionExperiment(
             experiment=to_gql_experiment(experiment)
         )  # eagerly yields experiment so it can be linked by consumers of the subscription

phoenix/server/api/types/Dataset.py

@@ -18,6 +18,7 @@ from phoenix.server.api.types.DatasetExample import DatasetExample
 from phoenix.server.api.types.DatasetExperimentAnnotationSummary import (
     DatasetExperimentAnnotationSummary,
 )
+from phoenix.server.api.types.DatasetLabel import DatasetLabel, to_gql_dataset_label
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.Experiment import Experiment, to_gql_experiment
 from phoenix.server.api.types.node import from_global_id_with_expected_type
@@ -303,6 +304,13 @@ class Dataset(Node):
                 async for scores_tuple in await session.stream(query)
             ]
 
+    @strawberry.field
+    async def labels(self, info: Info[Context, None]) -> list[DatasetLabel]:
+        return [
+            to_gql_dataset_label(label)
+            for label in await info.context.data_loaders.dataset_labels.load(self.id_attr)
+        ]
+
     @strawberry.field
     def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
         return info.context.last_updated_at.get(self._table, self.id_attr)

phoenix/server/api/types/DatasetExample.py

@@ -12,6 +12,7 @@ from phoenix.db import models
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import BadRequest
 from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
+from phoenix.server.api.types.DatasetSplit import DatasetSplit, to_gql_dataset_split
 from phoenix.server.api.types.DatasetVersion import DatasetVersion
 from phoenix.server.api.types.ExperimentRepeatedRunGroup import (
     ExperimentRepeatedRunGroup,
@@ -131,3 +132,20 @@ class DatasetExample(Node):
             )
             for group in repeated_run_groups
         ]
+
+    @strawberry.field
+    async def dataset_splits(
+        self,
+        info: Info[Context, None],
+    ) -> list[DatasetSplit]:
+        return [
+            to_gql_dataset_split(split)
+            for split in await info.context.data_loaders.dataset_example_splits.load(self.id_attr)
+        ]
+
+
+def to_gql_dataset_example(example: models.DatasetExample) -> DatasetExample:
+    return DatasetExample(
+        id_attr=example.id,
+        created_at=example.created_at,
+    )

phoenix/server/api/types/DatasetLabel.py

@@ -0,0 +1,23 @@
+from typing import Optional
+
+import strawberry
+from strawberry.relay import Node, NodeID
+
+from phoenix.db import models
+
+
+@strawberry.type
+class DatasetLabel(Node):
+    id_attr: NodeID[int]
+    name: str
+    description: Optional[str]
+    color: str
+
+
+def to_gql_dataset_label(dataset_label: models.DatasetLabel) -> DatasetLabel:
+    return DatasetLabel(
+        id_attr=dataset_label.id,
+        name=dataset_label.name,
+        description=dataset_label.description,
+        color=dataset_label.color,
+    )

phoenix/server/api/types/DatasetSplit.py

@@ -0,0 +1,32 @@
+from datetime import datetime
+from typing import ClassVar, Optional
+
+import strawberry
+from strawberry.relay import Node, NodeID
+from strawberry.scalars import JSON
+
+from phoenix.db import models
+
+
+@strawberry.type
+class DatasetSplit(Node):
+    _table: ClassVar[type[models.Base]] = models.DatasetSplit
+    id_attr: NodeID[int]
+    name: str
+    description: Optional[str]
+    metadata: JSON
+    color: str
+    created_at: datetime
+    updated_at: datetime
+
+
+def to_gql_dataset_split(dataset_split: models.DatasetSplit) -> DatasetSplit:
+    return DatasetSplit(
+        id_attr=dataset_split.id,
+        name=dataset_split.name,
+        description=dataset_split.description,
+        color=dataset_split.color or "#ffffff",
+        metadata=dataset_split.metadata_,
+        created_at=dataset_split.created_at,
+        updated_at=dataset_split.updated_at,
+    )

phoenix/server/api/types/Experiment.py

@@ -193,10 +193,6 @@ class Experiment(Node):
                 async for token_type, is_prompt, cost, tokens in data
             ]
 
-    @strawberry.field
-    async def repetition_count(self, info: Info[Context, None]) -> int:
-        return await info.context.data_loaders.experiment_repetition_counts.load(self.id_attr)
-
 
 def to_gql_experiment(
     experiment: models.Experiment,

phoenix/server/api/types/Project.py

@@ -588,6 +588,22 @@ class Project(Node):
         async with info.context.db() as session:
             return list(await session.scalars(stmt))
 
+    @strawberry.field(
+        description="Names of all available annotations for sessions. "
+        "(The list contains no duplicates.)"
+    )  # type: ignore
+    async def session_annotation_names(
+        self,
+        info: Info[Context, None],
+    ) -> list[str]:
+        stmt = (
+            select(distinct(models.ProjectSessionAnnotation.name))
+            .join(models.ProjectSession)
+            .where(models.ProjectSession.project_id == self.project_rowid)
+        )
+        async with info.context.db() as session:
+            return list(await session.scalars(stmt))
+
     @strawberry.field(
         description="Names of available document evaluations.",
     )  # type: ignore

phoenix/server/api/types/ProjectSession.py

@@ -1,14 +1,19 @@
+from collections import defaultdict
+from dataclasses import asdict, dataclass
 from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, ClassVar, Optional, Type
 
+import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import select
 from strawberry import UNSET, Info, Private, lazy
-from strawberry.relay import Connection, GlobalID, Node, NodeID
+from strawberry.relay import Connection, Node, NodeID
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.AnnotationFilter import AnnotationFilter, satisfies_filter
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.MimeType import MimeType
 from phoenix.server.api.types.pagination import ConnectionArgs, CursorString, connection_from_list
@@ -18,6 +23,8 @@ from phoenix.server.api.types.SpanIOValue import SpanIOValue
 from phoenix.server.api.types.TokenUsage import TokenUsage
 
 if TYPE_CHECKING:
+    from phoenix.server.api.types.Project import Project
+    from phoenix.server.api.types.ProjectSessionAnnotation import ProjectSessionAnnotation
     from phoenix.server.api.types.Trace import Trace
 
 
@@ -31,10 +38,13 @@ class ProjectSession(Node):
     end_time: datetime
 
     @strawberry.field
-    async def project_id(self) -> GlobalID:
+    async def project(
+        self,
+        info: Info[Context, None],
+    ) -> Annotated["Project", lazy(".Project")]:
         from phoenix.server.api.types.Project import Project
 
-        return GlobalID(type_name=Project.__name__, node_id=str(self.project_rowid))
+        return Project(project_rowid=self.project_rowid)
 
     @strawberry.field
     async def num_traces(
@@ -165,6 +175,81 @@ class ProjectSession(Node):
             for entry in summary
         ]
 
+    @strawberry.field
+    async def session_annotations(
+        self,
+        info: Info[Context, None],
+    ) -> list[Annotated["ProjectSessionAnnotation", lazy(".ProjectSessionAnnotation")]]:
+        """Get all annotations for this session."""
+        from phoenix.server.api.types.ProjectSessionAnnotation import (
+            to_gql_project_session_annotation,
+        )
+
+        stmt = select(models.ProjectSessionAnnotation).filter_by(project_session_id=self.id_attr)
+        async with info.context.db() as session:
+            annotations = await session.stream_scalars(stmt)
+            return [
+                to_gql_project_session_annotation(annotation) async for annotation in annotations
+            ]
+
+    @strawberry.field(
+        description="Summarizes each annotation (by name) associated with the session"
+    )  # type: ignore
+    async def session_annotation_summaries(
+        self,
+        info: Info[Context, None],
+        filter: Optional[AnnotationFilter] = None,
+    ) -> list[AnnotationSummary]:
+        """
+        Retrieves and summarizes annotations associated with this span.
+
+        This method aggregates annotation data by name and label, calculating metrics
+        such as count of occurrences and sum of scores. The results are organized
+        into a structured format that can be easily converted to a DataFrame.
+
+        Args:
+            info: GraphQL context information
+            filter: Optional filter to apply to annotations before processing
+
+        Returns:
+            A list of AnnotationSummary objects, each containing:
+            - name: The name of the annotation
+            - data: A list of dictionaries with label statistics
+        """
+        # Load all annotations for this span from the data loader
+        annotations = await info.context.data_loaders.session_annotations_by_session.load(
+            self.id_attr
+        )
+
+        # Apply filter if provided to narrow down the annotations
+        if filter:
+            annotations = [
+                annotation for annotation in annotations if satisfies_filter(annotation, filter)
+            ]
+
+        @dataclass
+        class Metrics:
+            record_count: int = 0
+            label_count: int = 0
+            score_sum: float = 0
+            score_count: int = 0
+
+        summaries: defaultdict[str, defaultdict[Optional[str], Metrics]] = defaultdict(
+            lambda: defaultdict(Metrics)
+        )
+        for annotation in annotations:
+            metrics = summaries[annotation.name][annotation.label]
+            metrics.record_count += 1
+            metrics.label_count += int(annotation.label is not None)
+            metrics.score_sum += annotation.score or 0
+            metrics.score_count += int(annotation.score is not None)
+
+        result: list[AnnotationSummary] = []
+        for name, label_metrics in summaries.items():
+            rows = [{"label": label, **asdict(metrics)} for label, metrics in label_metrics.items()]
+            result.append(AnnotationSummary(name=name, df=pd.DataFrame(rows), simple_avg=True))
+        return result
+
 
 def to_gql_project_session(project_session: models.ProjectSession) -> ProjectSession:
     return ProjectSession(

phoenix/server/api/types/ProjectSessionAnnotation.py

@@ -0,0 +1,68 @@
+from typing import Optional
+
+import strawberry
+from strawberry import Private
+from strawberry.relay import GlobalID, Node, NodeID
+from strawberry.scalars import JSON
+from strawberry.types import Info
+
+from phoenix.db import models
+from phoenix.server.api.context import Context
+from phoenix.server.api.types.AnnotatorKind import AnnotatorKind
+
+from .AnnotationSource import AnnotationSource
+from .User import User, to_gql_user
+
+
+@strawberry.type
+class ProjectSessionAnnotation(Node):
+    id_attr: NodeID[int]
+    user_id: Private[Optional[int]]
+    name: str
+    annotator_kind: AnnotatorKind
+    label: Optional[str]
+    score: Optional[float]
+    explanation: Optional[str]
+    metadata: JSON
+    _project_session_id: Private[Optional[int]]
+    identifier: str
+    source: AnnotationSource
+
+    @strawberry.field
+    async def project_session_id(self) -> GlobalID:
+        from phoenix.server.api.types.ProjectSession import ProjectSession
+
+        return GlobalID(type_name=ProjectSession.__name__, node_id=str(self._project_session_id))
+
+    @strawberry.field
+    async def user(
+        self,
+        info: Info[Context, None],
+    ) -> Optional[User]:
+        if self.user_id is None:
+            return None
+        user = await info.context.data_loaders.users.load(self.user_id)
+        if user is None:
+            return None
+        return to_gql_user(user)
+
+
+def to_gql_project_session_annotation(
+    annotation: models.ProjectSessionAnnotation,
+) -> ProjectSessionAnnotation:
+    """
+    Converts an ORM projectSession annotation to a GraphQL ProjectSessionAnnotation.
+    """
+    return ProjectSessionAnnotation(
+        id_attr=annotation.id,
+        user_id=annotation.user_id,
+        _project_session_id=annotation.project_session_id,
+        name=annotation.name,
+        annotator_kind=AnnotatorKind(annotation.annotator_kind),
+        label=annotation.label,
+        score=annotation.score,
+        explanation=annotation.explanation,
+        metadata=JSON(annotation.metadata_),
+        identifier=annotation.identifier,
+        source=AnnotationSource(annotation.source),
+    )

phoenix/server/api/types/Prompt.py

@@ -9,6 +9,7 @@ from strawberry.relay import Connection, GlobalID, Node, NodeID
 from strawberry.types import Info
 
 from phoenix.db import models
+from phoenix.db.types.identifier import Identifier as IdentifierModel
 from phoenix.server.api.context import Context
 from phoenix.server.api.exceptions import NotFound
 from phoenix.server.api.types.Identifier import Identifier
@@ -37,7 +38,10 @@ class Prompt(Node):
 
     @strawberry.field
     async def version(
-        self, info: Info[Context, None], version_id: Optional[GlobalID] = None
+        self,
+        info: Info[Context, None],
+        version_id: Optional[GlobalID] = None,
+        tag_name: Optional[Identifier] = None,
     ) -> PromptVersion:
         async with info.context.db() as session:
             if version_id:
@@ -50,6 +54,19 @@ class Prompt(Node):
                 )
                 if not version:
                     raise NotFound(f"Prompt version not found: {version_id}")
+            elif tag_name:
+                try:
+                    name = IdentifierModel(tag_name)
+                except ValueError:
+                    raise NotFound(f"Prompt version tag not found: {tag_name}")
+                version = await session.scalar(
+                    select(models.PromptVersion)
+                    .where(models.PromptVersion.prompt_id == self.id_attr)
+                    .join_from(models.PromptVersion, models.PromptVersionTag)
+                    .where(models.PromptVersionTag.name == name)
+                )
+                if not version:
+                    raise NotFound(f"This prompt has no associated versions by tag {tag_name}")
             else:
                 stmt = (
                     select(models.PromptVersion)

phoenix/server/api/types/Span.py

@@ -23,11 +23,11 @@ from phoenix.server.api.helpers.dataset_helpers import (
     get_dataset_example_input,
     get_dataset_example_output,
 )
-from phoenix.server.api.input_types.InvocationParameters import InvocationParameter
-from phoenix.server.api.input_types.SpanAnnotationFilter import (
-    SpanAnnotationFilter,
+from phoenix.server.api.input_types.AnnotationFilter import (
+    AnnotationFilter,
     satisfies_filter,
 )
+from phoenix.server.api.input_types.InvocationParameters import InvocationParameter
 from phoenix.server.api.input_types.SpanAnnotationSort import (
     SpanAnnotationColumn,
     SpanAnnotationSort,
@@ -547,7 +547,7 @@ class Span(Node):
         self,
         info: Info[Context, None],
         sort: Optional[SpanAnnotationSort] = UNSET,
-        filter: Optional[SpanAnnotationFilter] = None,
+        filter: Optional[AnnotationFilter] = None,
     ) -> list[SpanAnnotation]:
         span_id = self.span_rowid
         annotations = await info.context.data_loaders.span_annotations.load(span_id)
@@ -580,7 +580,7 @@ class Span(Node):
     async def span_annotation_summaries(
         self,
         info: Info[Context, None],
-        filter: Optional[SpanAnnotationFilter] = None,
+        filter: Optional[AnnotationFilter] = None,
     ) -> list[AnnotationSummary]:
         """
         Retrieves and summarizes annotations associated with this span.

phoenix/server/api/types/Trace.py

@@ -1,8 +1,11 @@
 from __future__ import annotations
 
+from collections import defaultdict
+from dataclasses import asdict, dataclass
 from datetime import datetime
 from typing import TYPE_CHECKING, Annotated, Optional, Union
 
+import pandas as pd
 import strawberry
 from openinference.semconv.trace import SpanAttributes
 from sqlalchemy import desc, select
@@ -13,7 +16,9 @@ from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.server.api.context import Context
+from phoenix.server.api.input_types.AnnotationFilter import AnnotationFilter, satisfies_filter
 from phoenix.server.api.input_types.TraceAnnotationSort import TraceAnnotationSort
+from phoenix.server.api.types.AnnotationSummary import AnnotationSummary
 from phoenix.server.api.types.CostBreakdown import CostBreakdown
 from phoenix.server.api.types.pagination import (
     ConnectionArgs,
@@ -229,6 +234,62 @@ class Trace(Node):
             annotations = await session.scalars(stmt)
         return [to_gql_trace_annotation(annotation) for annotation in annotations]
 
+    @strawberry.field(description="Summarizes each annotation (by name) associated with the trace")  # type: ignore
+    async def trace_annotation_summaries(
+        self,
+        info: Info[Context, None],
+        filter: Optional[AnnotationFilter] = None,
+    ) -> list[AnnotationSummary]:
+        """
+        Retrieves and summarizes annotations associated with this span.
+
+        This method aggregates annotation data by name and label, calculating metrics
+        such as count of occurrences and sum of scores. The results are organized
+        into a structured format that can be easily converted to a DataFrame.
+
+        Args:
+            info: GraphQL context information
+            filter: Optional filter to apply to annotations before processing
+
+        Returns:
+            A list of AnnotationSummary objects, each containing:
+            - name: The name of the annotation
+            - data: A list of dictionaries with label statistics
+        """
+        # Load all annotations for this span from the data loader
+        annotations = await info.context.data_loaders.trace_annotations_by_trace.load(
+            self.trace_rowid
+        )
+
+        # Apply filter if provided to narrow down the annotations
+        if filter:
+            annotations = [
+                annotation for annotation in annotations if satisfies_filter(annotation, filter)
+            ]
+
+        @dataclass
+        class Metrics:
+            record_count: int = 0
+            label_count: int = 0
+            score_sum: float = 0
+            score_count: int = 0
+
+        summaries: defaultdict[str, defaultdict[Optional[str], Metrics]] = defaultdict(
+            lambda: defaultdict(Metrics)
+        )
+        for annotation in annotations:
+            metrics = summaries[annotation.name][annotation.label]
+            metrics.record_count += 1
+            metrics.label_count += int(annotation.label is not None)
+            metrics.score_sum += annotation.score or 0
+            metrics.score_count += int(annotation.score is not None)
+
+        result: list[AnnotationSummary] = []
+        for name, label_metrics in summaries.items():
+            rows = [{"label": label, **asdict(metrics)} for label, metrics in label_metrics.items()]
+            result.append(AnnotationSummary(name=name, df=pd.DataFrame(rows), simple_avg=True))
+        return result
+
     @strawberry.field
     async def cost_summary(
         self,

phoenix/server/app.py

@@ -67,7 +67,6 @@ from phoenix.config import (
     get_env_gql_extension_paths,
     get_env_grpc_interceptor_paths,
     get_env_host,
-    get_env_host_root_path,
     get_env_max_spans_queue_size,
     get_env_port,
     get_env_support_email,
@@ -92,6 +91,7 @@ from phoenix.server.api.dataloaders import (
     DatasetExampleRevisionsDataLoader,
     DatasetExamplesAndVersionsByExperimentRunDataLoader,
     DatasetExampleSpansDataLoader,
+    DatasetExampleSplitsDataLoader,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
@@ -99,7 +99,6 @@ from phoenix.server.api.dataloaders import (
     ExperimentErrorRatesDataLoader,
     ExperimentRepeatedRunGroupAnnotationSummariesDataLoader,
     ExperimentRepeatedRunGroupsDataLoader,
-    ExperimentRepetitionCountsDataLoader,
     ExperimentRunAnnotations,
     ExperimentRunCountsDataLoader,
     ExperimentSequenceNumberDataLoader,
@@ -112,6 +111,7 @@ from phoenix.server.api.dataloaders import (
     ProjectIdsByTraceRetentionPolicyIdDataLoader,
     PromptVersionSequenceNumberDataLoader,
     RecordCountDataLoader,
+    SessionAnnotationsBySessionDataLoader,
     SessionIODataLoader,
     SessionNumTracesDataLoader,
     SessionNumTracesWithErrorDataLoader,
@@ -137,12 +137,14 @@ from phoenix.server.api.dataloaders import (
     SpanProjectsDataLoader,
     TableFieldsDataLoader,
     TokenCountDataLoader,
+    TraceAnnotationsByTraceDataLoader,
     TraceByTraceIdsDataLoader,
     TraceRetentionPolicyIdByProjectIdDataLoader,
     TraceRootSpansDataLoader,
     UserRolesDataLoader,
     UsersDataLoader,
 )
+from phoenix.server.api.dataloaders.dataset_labels import DatasetLabelsDataLoader
 from phoenix.server.api.routers import (
     auth_router,
     create_embeddings_router,
@@ -173,6 +175,7 @@ from phoenix.server.types import (
     LastUpdatedAt,
     TokenStore,
 )
+from phoenix.server.utils import get_root_path, prepend_root_path
 from phoenix.settings import Settings
 from phoenix.trace.fixtures import (
     TracesFixture,
@@ -281,9 +284,6 @@ class Static(StaticFiles):
                 return {}
             raise e
 
-    def _sanitize_basename(self, basename: str) -> str:
-        return basename[:-1] if basename.endswith("/") else basename
-
     async def get_response(self, path: str, scope: Scope) -> Response:
         # Redirect to the oauth2 login page if basic auth is disabled and auto_login is enabled
         # TODO: this needs to be refactored to be cleaner
@@ -292,14 +292,10 @@ class Static(StaticFiles):
             and self._app_config.basic_auth_disabled
             and self._app_config.auto_login_idp_name
         ):
-            request = Request(scope)
-            url = URL(
-                str(
-                    Path(get_env_host_root_path())
-                    / f"oauth2/{self._app_config.auto_login_idp_name}/login"
-                )
+            redirect_path = prepend_root_path(
+                scope, f"oauth2/{self._app_config.auto_login_idp_name}/login"
             )
-            url = url.include_query_params(**request.query_params)
+            url = URL(redirect_path).include_query_params(**Request(scope).query_params)
             return RedirectResponse(url=url)
         try:
             response = await super().get_response(path, scope)
@@ -316,7 +312,7 @@ class Static(StaticFiles):
                     "min_dist": self._app_config.min_dist,
                     "n_neighbors": self._app_config.n_neighbors,
                     "n_samples": self._app_config.n_samples,
-                    "basename": self._sanitize_basename(request.scope.get("root_path", "")),
+                    "basename": get_root_path(scope),
                     "platform_version": phoenix_version,
                     "request": request,
                     "is_development": self._app_config.is_development,
@@ -715,6 +711,8 @@ def create_graphql_router(
                 dataset_examples_and_versions_by_experiment_run=DatasetExamplesAndVersionsByExperimentRunDataLoader(
                     db
                 ),
+                dataset_example_splits=DatasetExampleSplitsDataLoader(db),
+                dataset_labels=DatasetLabelsDataLoader(db),
                 document_evaluation_summaries=DocumentEvaluationSummaryDataLoader(
                     db,
                     cache_map=(
@@ -737,7 +735,6 @@ def create_graphql_router(
                     db
                 ),
                 experiment_repeated_run_groups=ExperimentRepeatedRunGroupsDataLoader(db),
-                experiment_repetition_counts=ExperimentRepetitionCountsDataLoader(db),
                 experiment_run_annotations=ExperimentRunAnnotations(db),
                 experiment_run_counts=ExperimentRunCountsDataLoader(db),
                 experiment_sequence_number=ExperimentSequenceNumberDataLoader(db),
@@ -769,6 +766,7 @@ def create_graphql_router(
                     db,
                     cache_map=cache_for_dataloaders.record_count if cache_for_dataloaders else None,
                 ),
+                session_annotations_by_session=SessionAnnotationsBySessionDataLoader(db),
                 session_first_inputs=SessionIODataLoader(db, "first_input"),
                 session_last_outputs=SessionIODataLoader(db, "last_output"),
                 session_num_traces=SessionNumTracesDataLoader(db),
@@ -815,6 +813,7 @@ def create_graphql_router(
                     db,
                     cache_map=cache_for_dataloaders.token_count if cache_for_dataloaders else None,
                 ),
+                trace_annotations_by_trace=TraceAnnotationsByTraceDataLoader(db),
                 trace_by_trace_ids=TraceByTraceIdsDataLoader(db),
                 trace_fields=TableFieldsDataLoader(db, models.Trace),
                 trace_retention_policy_id_by_project_id=TraceRetentionPolicyIdByProjectIdDataLoader(