arize-phoenix 4.4.4rc5__py3-none-any.whl → 4.5.0__py3-none-any.whl

phoenix/db/migrations/versions/cf03bd6bae1d_init.py

@@ -6,11 +6,13 @@ Create Date: 2024-04-03 19:41:48.871555
 
 """
 
-from typing import Sequence, Union
+from typing import Any, Sequence, Union
 
 import sqlalchemy as sa
 from alembic import op
-from phoenix.db.migrations.types import JSON_
+from sqlalchemy import JSON
+from sqlalchemy.dialects import postgresql
+from sqlalchemy.ext.compiler import compiles
 
 # revision identifiers, used by Alembic.
 revision: str = "cf03bd6bae1d"
@@ -19,6 +21,30 @@ branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
 
+class JSONB(JSON):
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    __visit_name__ = "JSONB"
+
+
+@compiles(JSONB, "sqlite")  # type: ignore
+def _(*args: Any, **kwargs: Any) -> str:
+    # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
+    return "JSONB"
+
+
+JSON_ = (
+    JSON()
+    .with_variant(
+        postgresql.JSONB(),  # type: ignore
+        "postgresql",
+    )
+    .with_variant(
+        JSONB(),
+        "sqlite",
+    )
+)
+
+
 def upgrade() -> None:
     projects_table = op.create_table(
         "projects",

phoenix/db/models.py

@@ -1,5 +1,5 @@
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Optional
 
 from sqlalchemy import (
     JSON,
@@ -15,14 +15,12 @@ from sqlalchemy import (
     String,
     TypeDecorator,
     UniqueConstraint,
-    case,
     func,
     insert,
-    select,
     text,
 )
 from sqlalchemy.dialects import postgresql
-from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
+from sqlalchemy.ext.asyncio import AsyncEngine
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import (
@@ -61,24 +59,6 @@ JSON_ = (
 )
 
 
-class JsonDict(TypeDecorator[Dict[str, Any]]):
-    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
-    cache_ok = True
-    impl = JSON_
-
-    def process_bind_param(self, value: Optional[Dict[str, Any]], _: Dialect) -> Dict[str, Any]:
-        return value if isinstance(value, dict) else {}
-
-
-class JsonList(TypeDecorator[List[Any]]):
-    # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
-    cache_ok = True
-    impl = JSON_
-
-    def process_bind_param(self, value: Optional[List[Any]], _: Dialect) -> List[Any]:
-        return value if isinstance(value, list) else []
-
-
 class UtcTimeStamp(TypeDecorator[datetime]):
     # See # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
     cache_ok = True
@@ -91,10 +71,6 @@ class UtcTimeStamp(TypeDecorator[datetime]):
         return normalize_datetime(value, timezone.utc)
 
 
-class ExperimentResult(TypedDict, total=False):
-    result: Dict[str, Any]
-
-
 class Base(DeclarativeBase):
     # Enforce best practices for naming constraints
     # https://alembic.sqlalchemy.org/en/latest/naming.html#integration-of-naming-conventions-into-operations-autogenerate
@@ -108,9 +84,8 @@ class Base(DeclarativeBase):
         }
     )
     type_annotation_map = {
-        Dict[str, Any]: JsonDict,
-        List[Dict[str, Any]]: JsonList,
-        ExperimentResult: JsonDict,
+        Dict[str, Any]: JSON_,
+        List[Dict[str, Any]]: JSON_,
     }
 
 
@@ -179,10 +154,6 @@ class Trace(Base):
         cascade="all, delete-orphan",
         uselist=True,
     )
-    experiment_runs: Mapped[List["ExperimentRun"]] = relationship(
-        primaryjoin="foreign(ExperimentRun.trace_id) == Trace.trace_id",
-        back_populates="trace",
-    )
     __table_args__ = (
         UniqueConstraint(
             "trace_id",
@@ -232,7 +203,6 @@ class Span(Base):
 
     trace: Mapped["Trace"] = relationship("Trace", back_populates="spans")
     document_annotations: Mapped[List["DocumentAnnotation"]] = relationship(back_populates="span")
-    dataset_examples: Mapped[List["DatasetExample"]] = relationship(back_populates="span")
 
     __table_args__ = (
         UniqueConstraint(
@@ -406,205 +376,3 @@ class DocumentAnnotation(Base):
             "document_position",
         ),
     )
-
-
-class Dataset(Base):
-    __tablename__ = "datasets"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    name: Mapped[str] = mapped_column(unique=True)
-    description: Mapped[Optional[str]]
-    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
-    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-    updated_at: Mapped[datetime] = mapped_column(
-        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
-    )
-
-    @hybrid_property
-    def example_count(self) -> Optional[int]:
-        if hasattr(self, "_example_count_value"):
-            assert isinstance(self._example_count_value, int)
-            return self._example_count_value
-        return None
-
-    @example_count.inplace.expression
-    def _example_count(cls) -> ColumnElement[int]:
-        return (
-            select(
-                func.sum(
-                    case(
-                        (DatasetExampleRevision.revision_kind == "CREATE", 1),
-                        (DatasetExampleRevision.revision_kind == "DELETE", -1),
-                        else_=0,
-                    )
-                )
-            )
-            .select_from(DatasetExampleRevision)
-            .join(
-                DatasetExample,
-                onclause=DatasetExample.id == DatasetExampleRevision.dataset_example_id,
-            )
-            .filter(DatasetExample.dataset_id == cls.id)
-            .label("example_count")
-        )
-
-    async def load_example_count(self, session: AsyncSession) -> None:
-        if not hasattr(self, "_example_count_value"):
-            self._example_count_value = await session.scalar(
-                select(
-                    func.sum(
-                        case(
-                            (DatasetExampleRevision.revision_kind == "CREATE", 1),
-                            (DatasetExampleRevision.revision_kind == "DELETE", -1),
-                            else_=0,
-                        )
-                    )
-                )
-                .select_from(DatasetExampleRevision)
-                .join(
-                    DatasetExample,
-                    onclause=DatasetExample.id == DatasetExampleRevision.dataset_example_id,
-                )
-                .filter(DatasetExample.dataset_id == self.id)
-            )
-
-
-class DatasetVersion(Base):
-    __tablename__ = "dataset_versions"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    dataset_id: Mapped[int] = mapped_column(
-        ForeignKey("datasets.id", ondelete="CASCADE"),
-        index=True,
-    )
-    description: Mapped[Optional[str]]
-    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
-    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-
-
-class DatasetExample(Base):
-    __tablename__ = "dataset_examples"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    dataset_id: Mapped[int] = mapped_column(
-        ForeignKey("datasets.id", ondelete="CASCADE"),
-        index=True,
-    )
-    span_rowid: Mapped[Optional[int]] = mapped_column(
-        ForeignKey("spans.id", ondelete="SET NULL"),
-        index=True,
-        nullable=True,
-    )
-    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-
-    span: Mapped[Optional[Span]] = relationship(back_populates="dataset_examples")
-
-
-class DatasetExampleRevision(Base):
-    __tablename__ = "dataset_example_revisions"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    dataset_example_id: Mapped[int] = mapped_column(
-        ForeignKey("dataset_examples.id", ondelete="CASCADE"),
-        index=True,
-    )
-    dataset_version_id: Mapped[int] = mapped_column(
-        ForeignKey("dataset_versions.id", ondelete="CASCADE"),
-        index=True,
-    )
-    input: Mapped[Dict[str, Any]]
-    output: Mapped[Dict[str, Any]]
-    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
-    revision_kind: Mapped[str] = mapped_column(
-        CheckConstraint(
-            "revision_kind IN ('CREATE', 'PATCH', 'DELETE')", name="valid_revision_kind"
-        ),
-    )
-    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-
-    __table_args__ = (
-        UniqueConstraint(
-            "dataset_example_id",
-            "dataset_version_id",
-        ),
-    )
-
-
-class Experiment(Base):
-    __tablename__ = "experiments"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    dataset_id: Mapped[int] = mapped_column(
-        ForeignKey("datasets.id", ondelete="CASCADE"),
-        index=True,
-    )
-    dataset_version_id: Mapped[int] = mapped_column(
-        ForeignKey("dataset_versions.id", ondelete="CASCADE"),
-        index=True,
-    )
-    name: Mapped[str]
-    description: Mapped[Optional[str]]
-    repetitions: Mapped[int]
-    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
-    project_name: Mapped[Optional[str]] = mapped_column(index=True)
-    created_at: Mapped[datetime] = mapped_column(UtcTimeStamp, server_default=func.now())
-    updated_at: Mapped[datetime] = mapped_column(
-        UtcTimeStamp, server_default=func.now(), onupdate=func.now()
-    )
-
-
-class ExperimentRun(Base):
-    __tablename__ = "experiment_runs"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    experiment_id: Mapped[int] = mapped_column(
-        ForeignKey("experiments.id", ondelete="CASCADE"),
-        index=True,
-    )
-    dataset_example_id: Mapped[int] = mapped_column(
-        ForeignKey("dataset_examples.id", ondelete="CASCADE"),
-        index=True,
-    )
-    repetition_number: Mapped[int]
-    trace_id: Mapped[Optional[str]]
-    output: Mapped[ExperimentResult]
-    start_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
-    end_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
-    prompt_token_count: Mapped[Optional[int]]
-    completion_token_count: Mapped[Optional[int]]
-    error: Mapped[Optional[str]]
-
-    trace: Mapped["Trace"] = relationship(
-        primaryjoin="foreign(ExperimentRun.trace_id) == Trace.trace_id",
-        back_populates="experiment_runs",
-    )
-
-    __table_args__ = (
-        UniqueConstraint(
-            "experiment_id",
-            "dataset_example_id",
-            "repetition_number",
-        ),
-    )
-
-
-class ExperimentRunAnnotation(Base):
-    __tablename__ = "experiment_run_annotations"
-    id: Mapped[int] = mapped_column(primary_key=True)
-    experiment_run_id: Mapped[int] = mapped_column(
-        ForeignKey("experiment_runs.id", ondelete="CASCADE"),
-        index=True,
-    )
-    name: Mapped[str]
-    annotator_kind: Mapped[str] = mapped_column(
-        CheckConstraint("annotator_kind IN ('LLM', 'CODE', 'HUMAN')", name="valid_annotator_kind"),
-    )
-    label: Mapped[Optional[str]]
-    score: Mapped[Optional[float]]
-    explanation: Mapped[Optional[str]]
-    trace_id: Mapped[Optional[str]]
-    error: Mapped[Optional[str]]
-    metadata_: Mapped[Dict[str, Any]] = mapped_column("metadata")
-    start_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
-    end_time: Mapped[datetime] = mapped_column(UtcTimeStamp)
-
-    __table_args__ = (
-        UniqueConstraint(
-            "experiment_run_id",
-            "name",
-        ),
-    )

phoenix/inferences/fixtures.py

@@ -9,7 +9,7 @@ from urllib.parse import quote, urljoin
 
 from pandas import read_parquet
 
-from phoenix.config import INFERENCES_DIR
+from phoenix.config import DATASET_DIR
 from phoenix.inferences.inferences import Inferences
 from phoenix.inferences.schema import (
     EmbeddingColumnNames,
@@ -20,7 +20,7 @@ from phoenix.inferences.schema import (
 logger = logging.getLogger(__name__)
 
 
-class InferencesRole(Enum):
+class DatasetRole(Enum):
     PRIMARY = auto()
     REFERENCE = auto()
     CORPUS = auto()
@@ -39,11 +39,11 @@ class Fixture:
     corpus_file_name: Optional[str] = None
     corpus_schema: Optional[Schema] = None
 
-    def paths(self) -> Iterator[Tuple[InferencesRole, Path]]:
+    def paths(self) -> Iterator[Tuple[DatasetRole, Path]]:
         return (
             (role, Path(self.prefix) / name)
             for role, name in zip(
-                InferencesRole,
+                DatasetRole,
                 (
                     self.primary_file_name,
                     self.reference_file_name,
@@ -413,41 +413,41 @@ FIXTURES: Tuple[Fixture, ...] = (
 NAME_TO_FIXTURE = {fixture.name: fixture for fixture in FIXTURES}
 
 
-def get_inferences(
+def get_datasets(
     fixture_name: str,
     no_internet: bool = False,
 ) -> Tuple[Inferences, Optional[Inferences], Optional[Inferences]]:
     """
-    Downloads primary and reference inferences for a fixture if they are not found
+    Downloads primary and reference datasets for a fixture if they are not found
     locally.
     """
     fixture = _get_fixture_by_name(fixture_name=fixture_name)
     if no_internet:
-        paths = {role: INFERENCES_DIR / path for role, path in fixture.paths()}
+        paths = {role: DATASET_DIR / path for role, path in fixture.paths()}
     else:
-        paths = dict(_download(fixture, INFERENCES_DIR))
-    primary_inferences = Inferences(
-        read_parquet(paths[InferencesRole.PRIMARY]),
+        paths = dict(_download(fixture, DATASET_DIR))
+    primary_dataset = Inferences(
+        read_parquet(paths[DatasetRole.PRIMARY]),
         fixture.primary_schema,
         "production",
     )
-    reference_inferences = None
+    reference_dataset = None
     if fixture.reference_file_name is not None:
-        reference_inferences = Inferences(
-            read_parquet(paths[InferencesRole.REFERENCE]),
+        reference_dataset = Inferences(
+            read_parquet(paths[DatasetRole.REFERENCE]),
             fixture.reference_schema
             if fixture.reference_schema is not None
             else fixture.primary_schema,
             "training",
         )
-    corpus_inferences = None
+    corpus_dataset = None
     if fixture.corpus_file_name is not None:
-        corpus_inferences = Inferences(
-            read_parquet(paths[InferencesRole.CORPUS]),
+        corpus_dataset = Inferences(
+            read_parquet(paths[DatasetRole.CORPUS]),
             fixture.corpus_schema,
             "knowledge_base",
         )
-    return primary_inferences, reference_inferences, corpus_inferences
+    return primary_dataset, reference_dataset, corpus_dataset
 
 
 def _get_fixture_by_name(fixture_name: str) -> Fixture:
@@ -496,14 +496,14 @@ def load_example(use_case: str) -> ExampleInferences:
 
     """
     fixture = _get_fixture_by_name(use_case)
-    primary_inferences, reference_inferences, corpus_inferences = get_inferences(use_case)
+    primary_dataset, reference_dataset, corpus_dataset = get_datasets(use_case)
     print(f"📥 Loaded {use_case} example datasets.")
     print("ℹ️ About this use-case:")
     print(fixture.description)
     return ExampleInferences(
-        primary=primary_inferences,
-        reference=reference_inferences,
-        corpus=corpus_inferences,
+        primary=primary_dataset,
+        reference=reference_dataset,
+        corpus=corpus_dataset,
     )
 
 
@@ -544,7 +544,7 @@ class GCSAssets(NamedTuple):
         )
 
 
-def _download(fixture: Fixture, location: Path) -> Iterator[Tuple[InferencesRole, Path]]:
+def _download(fixture: Fixture, location: Path) -> Iterator[Tuple[DatasetRole, Path]]:
     for role, path in fixture.paths():
         yield role, GCSAssets().metadata(path).save_artifact(location)
 
@@ -556,5 +556,5 @@ if __name__ == "__main__":
     for fixture in FIXTURES:
         start_time = time.time()
         print(f"getting {fixture.name}", end="...")
-        dict(_download(fixture, INFERENCES_DIR))
+        dict(_download(fixture, DATASET_DIR))
         print(f"done ({time.time() - start_time:.2f}s)")

phoenix/inferences/inferences.py

@@ -15,7 +15,7 @@ from pandas.api.types import (
 )
 from typing_extensions import TypeAlias
 
-from phoenix.config import GENERATED_INFERENCES_NAME_PREFIX, INFERENCES_DIR
+from phoenix.config import DATASET_DIR, GENERATED_DATASET_NAME_PREFIX
 from phoenix.datetime_utils import normalize_timestamps
 from phoenix.utilities.deprecation import deprecated
 
@@ -31,7 +31,7 @@ from .schema import (
     SchemaFieldName,
     SchemaFieldValue,
 )
-from .validation import validate_inferences_inputs
+from .validation import validate_dataset_inputs
 
 logger = logging.getLogger(__name__)
 
@@ -62,7 +62,7 @@ class Inferences:
 
     Examples
     --------
-    >>> primary_inferences = px.Inferences(
+    >>> primary_dataset = px.Inferences(
     >>>    dataframe=production_dataframe, schema=schema, name="primary"
     >>> )
     """
@@ -81,7 +81,7 @@ class Inferences:
         # allow for schema like objects
         if not isinstance(schema, Schema):
             schema = _get_schema_from_unknown_schema_param(schema)
-        errors = validate_inferences_inputs(
+        errors = validate_dataset_inputs(
             dataframe=dataframe,
             schema=schema,
         )
@@ -95,7 +95,7 @@ class Inferences:
         self.__dataframe: DataFrame = dataframe
         self.__schema: Schema = schema
         self.__name: str = (
-            name if name is not None else f"{GENERATED_INFERENCES_NAME_PREFIX}{str(uuid.uuid4())}"
+            name if name is not None else f"{GENERATED_DATASET_NAME_PREFIX}{str(uuid.uuid4())}"
         )
         self._is_empty = self.dataframe.empty
         logger.info(f"""Dataset: {self.__name} initialized""")
@@ -118,7 +118,7 @@ class Inferences:
     @classmethod
     def from_name(cls, name: str) -> "Inferences":
         """Retrieves a dataset by name from the file system"""
-        directory = INFERENCES_DIR / name
+        directory = DATASET_DIR / name
         df = read_parquet(directory / cls._data_file_name)
         with open(directory / cls._schema_file_name) as schema_file:
             schema_json = schema_file.read()
@@ -127,7 +127,7 @@ class Inferences:
 
     def to_disc(self) -> None:
         """writes the data and schema to disc"""
-        directory = INFERENCES_DIR / self.name
+        directory = DATASET_DIR / self.name
         directory.mkdir(parents=True, exist_ok=True)
         self.dataframe.to_parquet(
             directory / self._data_file_name,

phoenix/inferences/validation.py

@@ -34,7 +34,7 @@ def _check_valid_schema(schema: Schema) -> List[err.ValidationError]:
     return []
 
 
-def validate_inferences_inputs(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
+def validate_dataset_inputs(dataframe: DataFrame, schema: Schema) -> List[err.ValidationError]:
     errors = _check_missing_columns(dataframe, schema)
     if errors:
         return errors

phoenix/server/api/context.py

@@ -12,51 +12,33 @@ from typing_extensions import TypeAlias
 from phoenix.core.model_schema import Model
 from phoenix.server.api.dataloaders import (
     CacheForDataLoaders,
-    DatasetExampleRevisionsDataLoader,
-    DatasetExampleSpansDataLoader,
     DocumentEvaluationsDataLoader,
     DocumentEvaluationSummaryDataLoader,
     DocumentRetrievalMetricsDataLoader,
     EvaluationSummaryDataLoader,
-    ExperimentAnnotationSummaryDataLoader,
-    ExperimentErrorRatesDataLoader,
-    ExperimentRunCountsDataLoader,
-    ExperimentSequenceNumberDataLoader,
     LatencyMsQuantileDataLoader,
     MinStartOrMaxEndTimeDataLoader,
-    ProjectByNameDataLoader,
     RecordCountDataLoader,
     SpanDescendantsDataLoader,
     SpanEvaluationsDataLoader,
-    SpanProjectsDataLoader,
     TokenCountDataLoader,
     TraceEvaluationsDataLoader,
-    TraceRowIdsDataLoader,
 )
 
 
 @dataclass
 class DataLoaders:
-    dataset_example_revisions: DatasetExampleRevisionsDataLoader
-    dataset_example_spans: DatasetExampleSpansDataLoader
     document_evaluation_summaries: DocumentEvaluationSummaryDataLoader
     document_evaluations: DocumentEvaluationsDataLoader
     document_retrieval_metrics: DocumentRetrievalMetricsDataLoader
     evaluation_summaries: EvaluationSummaryDataLoader
-    experiment_annotation_summaries: ExperimentAnnotationSummaryDataLoader
-    experiment_error_rates: ExperimentErrorRatesDataLoader
-    experiment_run_counts: ExperimentRunCountsDataLoader
-    experiment_sequence_number: ExperimentSequenceNumberDataLoader
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
     record_counts: RecordCountDataLoader
     span_descendants: SpanDescendantsDataLoader
     span_evaluations: SpanEvaluationsDataLoader
-    span_projects: SpanProjectsDataLoader
     token_counts: TokenCountDataLoader
     trace_evaluations: TraceEvaluationsDataLoader
-    trace_row_ids: TraceRowIdsDataLoader
-    project_by_name: ProjectByNameDataLoader
 
 
 ProjectRowId: TypeAlias = int

phoenix/server/api/dataloaders/__init__.py

@@ -8,8 +8,6 @@ from phoenix.db.insertion.evaluation import (
 )
 from phoenix.db.insertion.span import ClearProjectSpansEvent, SpanInsertionEvent
 
-from .dataset_example_revisions import DatasetExampleRevisionsDataLoader
-from .dataset_example_spans import DatasetExampleSpansDataLoader
 from .document_evaluation_summaries import (
     DocumentEvaluationSummaryCache,
     DocumentEvaluationSummaryDataLoader,
@@ -17,43 +15,27 @@ from .document_evaluation_summaries import (
 from .document_evaluations import DocumentEvaluationsDataLoader
 from .document_retrieval_metrics import DocumentRetrievalMetricsDataLoader
 from .evaluation_summaries import EvaluationSummaryCache, EvaluationSummaryDataLoader
-from .experiment_annotation_summaries import ExperimentAnnotationSummaryDataLoader
-from .experiment_error_rates import ExperimentErrorRatesDataLoader
-from .experiment_run_counts import ExperimentRunCountsDataLoader
-from .experiment_sequence_number import ExperimentSequenceNumberDataLoader
 from .latency_ms_quantile import LatencyMsQuantileCache, LatencyMsQuantileDataLoader
 from .min_start_or_max_end_times import MinStartOrMaxEndTimeCache, MinStartOrMaxEndTimeDataLoader
-from .project_by_name import ProjectByNameDataLoader
 from .record_counts import RecordCountCache, RecordCountDataLoader
 from .span_descendants import SpanDescendantsDataLoader
 from .span_evaluations import SpanEvaluationsDataLoader
-from .span_projects import SpanProjectsDataLoader
 from .token_counts import TokenCountCache, TokenCountDataLoader
 from .trace_evaluations import TraceEvaluationsDataLoader
-from .trace_row_ids import TraceRowIdsDataLoader
 
 __all__ = [
     "CacheForDataLoaders",
-    "DatasetExampleRevisionsDataLoader",
-    "DatasetExampleSpansDataLoader",
     "DocumentEvaluationSummaryDataLoader",
     "DocumentEvaluationsDataLoader",
     "DocumentRetrievalMetricsDataLoader",
     "EvaluationSummaryDataLoader",
-    "ExperimentAnnotationSummaryDataLoader",
-    "ExperimentErrorRatesDataLoader",
-    "ExperimentRunCountsDataLoader",
-    "ExperimentSequenceNumberDataLoader",
     "LatencyMsQuantileDataLoader",
     "MinStartOrMaxEndTimeDataLoader",
     "RecordCountDataLoader",
     "SpanDescendantsDataLoader",
     "SpanEvaluationsDataLoader",
-    "SpanProjectsDataLoader",
     "TokenCountDataLoader",
     "TraceEvaluationsDataLoader",
-    "TraceRowIdsDataLoader",
-    "ProjectByNameDataLoader",
 ]
 
 

phoenix/server/api/dataloaders/span_descendants.py

@@ -9,7 +9,7 @@ from typing import (
 from aioitertools.itertools import groupby
 from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import contains_eager
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -52,7 +52,8 @@ class SpanDescendantsDataLoader(DataLoader[Key, Result]):
         stmt = (
             select(descendant_ids.c[root_id_label], models.Span)
             .join(descendant_ids, models.Span.id == descendant_ids.c.id)
-            .options(joinedload(models.Span.trace, innerjoin=True).load_only(models.Trace.trace_id))
+            .join(models.Trace)
+            .options(contains_eager(models.Span.trace))
             .order_by(descendant_ids.c[root_id_label])
         )
         results: Dict[SpanId, Result] = {key: [] for key in keys}