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

phoenix/server/api/dataloaders/experiment_repeated_run_group_annotation_summaries.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from typing import Optional
+
+from sqlalchemy import func, select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentID: TypeAlias = int
+DatasetExampleID: TypeAlias = int
+AnnotationName: TypeAlias = str
+MeanAnnotationScore: TypeAlias = float
+
+
+@dataclass
+class AnnotationSummary:
+    annotation_name: AnnotationName
+    mean_score: Optional[MeanAnnotationScore]
+
+
+Key: TypeAlias = tuple[ExperimentID, DatasetExampleID]
+Result: TypeAlias = list[AnnotationSummary]
+
+
+class ExperimentRepeatedRunGroupAnnotationSummariesDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: DbSessionFactory,
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        annotation_summaries_query = (
+            select(
+                models.ExperimentRun.experiment_id.label("experiment_id"),
+                models.ExperimentRun.dataset_example_id.label("dataset_example_id"),
+                models.ExperimentRunAnnotation.name.label("annotation_name"),
+                func.avg(models.ExperimentRunAnnotation.score).label("mean_score"),
+            )
+            .select_from(models.ExperimentRunAnnotation)
+            .join(
+                models.ExperimentRun,
+                models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
+            )
+            .where(
+                tuple_(
+                    models.ExperimentRun.experiment_id, models.ExperimentRun.dataset_example_id
+                ).in_(set(keys))
+            )
+            .group_by(
+                models.ExperimentRun.experiment_id,
+                models.ExperimentRun.dataset_example_id,
+                models.ExperimentRunAnnotation.name,
+            )
+        )
+        async with self._db() as session:
+            annotation_summaries = (await session.execute(annotation_summaries_query)).all()
+        annotation_summaries_by_key: dict[Key, list[AnnotationSummary]] = {}
+        for summary in annotation_summaries:
+            key = (summary.experiment_id, summary.dataset_example_id)
+            gql_summary = AnnotationSummary(
+                annotation_name=summary.annotation_name,
+                mean_score=summary.mean_score,
+            )
+            if key not in annotation_summaries_by_key:
+                annotation_summaries_by_key[key] = []
+            annotation_summaries_by_key[key].append(gql_summary)
+        return [
+            sorted(
+                annotation_summaries_by_key.get(key, []),
+                key=lambda summary: summary.annotation_name,
+            )
+            for key in keys
+        ]

phoenix/server/api/dataloaders/experiment_repeated_run_groups.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+
+from sqlalchemy import select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentID: TypeAlias = int
+DatasetExampleID: TypeAlias = int
+Key: TypeAlias = tuple[ExperimentID, DatasetExampleID]
+
+
+@dataclass
+class ExperimentRepeatedRunGroup:
+    experiment_rowid: int
+    dataset_example_rowid: int
+    runs: list[models.ExperimentRun]
+
+
+Result: TypeAlias = ExperimentRepeatedRunGroup
+
+
+class ExperimentRepeatedRunGroupsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        repeated_run_groups_query = (
+            select(models.ExperimentRun)
+            .where(
+                tuple_(
+                    models.ExperimentRun.experiment_id,
+                    models.ExperimentRun.dataset_example_id,
+                ).in_(set(keys))
+            )
+            .order_by(models.ExperimentRun.repetition_number)
+        )
+
+        async with self._db() as session:
+            runs_by_key: dict[Key, list[models.ExperimentRun]] = {}
+            for run in (await session.scalars(repeated_run_groups_query)).all():
+                key = (run.experiment_id, run.dataset_example_id)
+                if key not in runs_by_key:
+                    runs_by_key[key] = []
+                runs_by_key[key].append(run)
+
+        return [
+            ExperimentRepeatedRunGroup(
+                experiment_rowid=experiment_id,
+                dataset_example_rowid=dataset_example_id,
+                runs=runs_by_key.get((experiment_id, dataset_example_id), []),
+            )
+            for (experiment_id, dataset_example_id) in keys
+        ]

phoenix/server/api/dataloaders/experiment_runs_by_experiment_and_example.py

@@ -0,0 +1,44 @@
+from collections import defaultdict
+
+from sqlalchemy import select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+Key: TypeAlias = tuple[ExperimentId, DatasetExampleId]
+Result: TypeAlias = list[models.ExperimentRun]
+
+
+class ExperimentRunsByExperimentAndExampleDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        runs_by_key: defaultdict[Key, Result] = defaultdict(list)
+
+        async with self._db() as session:
+            stmt = (
+                select(models.ExperimentRun)
+                .where(
+                    tuple_(
+                        models.ExperimentRun.experiment_id,
+                        models.ExperimentRun.dataset_example_id,
+                    ).in_(keys)
+                )
+                .order_by(
+                    models.ExperimentRun.experiment_id,
+                    models.ExperimentRun.dataset_example_id,
+                    models.ExperimentRun.repetition_number,
+                )
+            )
+            result = await session.stream_scalars(stmt)
+            async for run in result:
+                key = (run.experiment_id, run.dataset_example_id)
+                runs_by_key[key].append(run)
+
+        return [runs_by_key[key] for key in keys]

phoenix/server/api/dataloaders/last_used_times_by_generative_model_id.py

@@ -0,0 +1,35 @@
+from datetime import datetime
+from typing import Optional
+
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+GenerativeModelID: TypeAlias = int
+Key: TypeAlias = GenerativeModelID
+Result: TypeAlias = Optional[datetime]
+
+
+class LastUsedTimesByGenerativeModelIdDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        async with self._db() as session:
+            last_used_times_by_model_id: dict[Key, Result] = {
+                model_id: last_used_time
+                async for model_id, last_used_time in await session.stream(
+                    select(
+                        models.SpanCost.model_id,
+                        func.max(models.SpanCost.span_start_time).label("last_used_time"),
+                    )
+                    .select_from(models.SpanCost)
+                    .where(models.SpanCost.model_id.in_(keys))
+                    .group_by(models.SpanCost.model_id)
+                )
+            }
+        return [last_used_times_by_model_id.get(model_id) for model_id in keys]

phoenix/server/api/dataloaders/latency_ms_quantile.py

@@ -25,6 +25,7 @@ from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -32,13 +33,16 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 Probability: TypeAlias = float
 QuantileValue: TypeAlias = float
 
-Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition, SessionFilterCondition]
 Param: TypeAlias = tuple[ProjectRowId, Probability]
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, Probability]
+Key: TypeAlias = tuple[
+    Kind, ProjectRowId, Optional[TimeRange], FilterCondition, SessionFilterCondition, Probability
+]
 Result: TypeAlias = Optional[QuantileValue]
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = None
@@ -47,15 +51,18 @@ FloatCol: TypeAlias = SQLColumnExpression[Float[float]]
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    kind, project_rowid, time_range, filter_condition, probability = key
+    kind, project_rowid, time_range, filter_condition, session_filter_condition, probability = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (kind, interval, filter_condition), (project_rowid, probability)
+    return (kind, interval, filter_condition, session_filter_condition), (
+        project_rowid,
+        probability,
+    )
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind, Probability]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition, Kind, Probability]
 
 
 class LatencyMsQuantileCache(
@@ -71,8 +78,17 @@ class LatencyMsQuantileCache(
         )
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (kind, interval, filter_condition), (project_rowid, probability) = _cache_key_fn(key)
-        return project_rowid, (interval, filter_condition, kind, probability)
+        (
+            (kind, interval, filter_condition, session_filter_condition),
+            (project_rowid, probability),
+        ) = _cache_key_fn(key)
+        return project_rowid, (
+            interval,
+            filter_condition,
+            session_filter_condition,
+            kind,
+            probability,
+        )
 
 
 class LatencyMsQuantileDataLoader(DataLoader[Key, Result]):
@@ -113,11 +129,18 @@ async def _get_results(
     segment: Segment,
     params: Mapping[Param, list[ResultPosition]],
 ) -> AsyncIterator[tuple[ResultPosition, QuantileValue]]:
-    kind, (start_time, end_time), filter_condition = segment
+    kind, (start_time, end_time), filter_condition, session_filter_condition = segment
     stmt = select(models.Trace.project_rowid)
     if kind == "trace":
         latency_column = cast(FloatCol, models.Trace.latency_ms)
         time_column = models.Trace.start_time
+        if filter_condition:
+            sf = SpanFilter(filter_condition)
+            stmt = stmt.where(
+                models.Trace.id.in_(
+                    sf(select(models.Span.trace_rowid).distinct()).scalar_subquery()
+                )
+            )
     elif kind == "span":
         latency_column = cast(FloatCol, models.Span.latency_ms)
         time_column = models.Span.start_time
@@ -127,6 +150,15 @@ async def _get_results(
             stmt = sf(stmt)
     else:
         assert_never(kind)
+    if session_filter_condition:
+        project_rowids = [project_rowid for project_rowid, _ in params]
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=project_rowids,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        stmt = stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
     if start_time:
         stmt = stmt.where(start_time <= time_column)
     if end_time:

phoenix/server/api/dataloaders/record_counts.py

@@ -3,13 +3,14 @@ from datetime import datetime
 from typing import Any, Literal, Optional
 
 from cachetools import LFUCache, TTLCache
-from sqlalchemy import Select, func, select
+from sqlalchemy import Select, distinct, func, select
 from strawberry.dataloader import AbstractCache, DataLoader
 from typing_extensions import TypeAlias, assert_never
 
 from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.session_filters import get_filtered_session_rowids_subquery
 from phoenix.server.types import DbSessionFactory
 from phoenix.trace.dsl import SpanFilter
 
@@ -17,27 +18,35 @@ Kind: TypeAlias = Literal["span", "trace"]
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 SpanCount: TypeAlias = int
 
-Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[Kind, TimeInterval, FilterCondition, SessionFilterCondition]
 Param: TypeAlias = ProjectRowId
 
-Key: TypeAlias = tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+Key: TypeAlias = tuple[
+    Kind, ProjectRowId, Optional[TimeRange], FilterCondition, SessionFilterCondition
+]
 Result: TypeAlias = SpanCount
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = 0
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    kind, project_rowid, time_range, filter_condition = key
+    kind, project_rowid, time_range, filter_condition, session_filter_condition = key
     interval = (
         (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
     )
-    return (kind, interval, filter_condition), project_rowid
+    return (
+        kind,
+        interval,
+        filter_condition,
+        session_filter_condition,
+    ), project_rowid
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, Kind]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition, Kind]
 
 
 class RecordCountCache(
@@ -53,8 +62,10 @@ class RecordCountCache(
         )
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (kind, interval, filter_condition), project_rowid = _cache_key_fn(key)
-        return project_rowid, (interval, filter_condition, kind)
+        (kind, interval, filter_condition, session_filter_condition), project_rowid = _cache_key_fn(
+            key
+        )
+        return project_rowid, (interval, filter_condition, session_filter_condition, kind)
 
 
 class RecordCountDataLoader(DataLoader[Key, Result]):
@@ -93,7 +104,7 @@ def _get_stmt(
     segment: Segment,
     *project_rowids: Param,
 ) -> Select[Any]:
-    kind, (start_time, end_time), filter_condition = segment
+    kind, (start_time, end_time), filter_condition, session_filter_condition = segment
     pid = models.Trace.project_rowid
     stmt = select(pid)
     if kind == "span":
@@ -102,12 +113,28 @@ def _get_stmt(
         if filter_condition:
             sf = SpanFilter(filter_condition)
             stmt = sf(stmt)
+        stmt = stmt.add_columns(func.count().label("count"))
     elif kind == "trace":
         time_column = models.Trace.start_time
+        if filter_condition:
+            stmt = stmt.join(models.Span, models.Trace.id == models.Span.trace_rowid)
+            stmt = stmt.add_columns(func.count(distinct(models.Trace.id)).label("count"))
+            sf = SpanFilter(filter_condition)
+            stmt = sf(stmt)
+        else:
+            stmt = stmt.add_columns(func.count().label("count"))
     else:
         assert_never(kind)
-    stmt = stmt.add_columns(func.count().label("count"))
     stmt = stmt.where(pid.in_(project_rowids))
+
+    if session_filter_condition:
+        filtered_session_rowids = get_filtered_session_rowids_subquery(
+            session_filter_condition=session_filter_condition,
+            project_rowids=project_rowids,
+            start_time=start_time,
+            end_time=end_time,
+        )
+        stmt = stmt.where(models.Trace.project_session_rowid.in_(filtered_session_rowids))
     stmt = stmt.group_by(pid)
     if start_time:
         stmt = stmt.where(start_time <= time_column)

phoenix/server/api/dataloaders/session_annotations_by_session.py

@@ -0,0 +1,29 @@
+from collections import defaultdict
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db.models import ProjectSessionAnnotation
+from phoenix.server.types import DbSessionFactory
+
+ProjectSessionId: TypeAlias = int
+Key: TypeAlias = ProjectSessionId
+Result: TypeAlias = list[ProjectSessionAnnotation]
+
+
+class SessionAnnotationsBySessionDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        annotations_by_id: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for annotation in await session.stream_scalars(
+                select(ProjectSessionAnnotation).where(
+                    ProjectSessionAnnotation.project_session_id.in_(keys)
+                )
+            ):
+                annotations_by_id[annotation.project_session_id].append(annotation)
+        return [annotations_by_id[key] for key in keys]

phoenix/server/api/dataloaders/span_cost_by_span.py

@@ -0,0 +1,24 @@
+from typing import Optional
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+SpanRowId: TypeAlias = int
+Key: TypeAlias = SpanRowId
+Result: TypeAlias = Optional[models.SpanCost]
+
+
+class SpanCostBySpanDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        stmt = select(models.SpanCost).where(models.SpanCost.span_rowid.in_(keys))
+        async with self._db() as session:
+            result = {sc.span_rowid: sc async for sc in await session.stream_scalars(stmt)}
+        return list(map(result.get, keys))

phoenix/server/api/dataloaders/span_cost_detail_summary_entries_by_generative_model.py

@@ -0,0 +1,56 @@
+from collections import defaultdict
+
+from sqlalchemy import func, select
+from sqlalchemy.sql.functions import coalesce
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.types import (
+    CostBreakdown,
+    SpanCostDetailSummaryEntry,
+)
+from phoenix.server.types import DbSessionFactory
+
+GenerativeModelId: TypeAlias = int
+Key: TypeAlias = GenerativeModelId
+Result: TypeAlias = list[SpanCostDetailSummaryEntry]
+
+
+class SpanCostDetailSummaryEntriesByGenerativeModelDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        pk = models.SpanCost.model_id
+        stmt = (
+            select(
+                pk,
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                coalesce(func.sum(models.SpanCostDetail.cost), 0).label("cost"),
+                coalesce(func.sum(models.SpanCostDetail.tokens), 0).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .where(pk.in_(keys))
+            .group_by(pk, models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
+        )
+        results: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                id_,
+                token_type,
+                is_prompt,
+                cost,
+                tokens,
+            ) in data:
+                entry = SpanCostDetailSummaryEntry(
+                    token_type=token_type,
+                    is_prompt=is_prompt,
+                    value=CostBreakdown(tokens=tokens, cost=cost),
+                )
+                results[id_].append(entry)
+        return list(map(list, map(results.__getitem__, keys)))

phoenix/server/api/dataloaders/span_cost_detail_summary_entries_by_project_session.py

@@ -0,0 +1,57 @@
+from collections import defaultdict
+
+from sqlalchemy import func, select
+from sqlalchemy.sql.functions import coalesce
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.types import (
+    CostBreakdown,
+    SpanCostDetailSummaryEntry,
+)
+from phoenix.server.types import DbSessionFactory
+
+ProjectSessionRowId: TypeAlias = int
+Key: TypeAlias = ProjectSessionRowId
+Result: TypeAlias = list[SpanCostDetailSummaryEntry]
+
+
+class SpanCostDetailSummaryEntriesByProjectSessionDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        pk = models.Trace.project_session_rowid
+        stmt = (
+            select(
+                pk,
+                models.SpanCostDetail.token_type,
+                models.SpanCostDetail.is_prompt,
+                coalesce(func.sum(models.SpanCostDetail.cost), 0).label("cost"),
+                coalesce(func.sum(models.SpanCostDetail.tokens), 0).label("tokens"),
+            )
+            .select_from(models.SpanCostDetail)
+            .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+            .join(models.Trace, models.SpanCost.trace_rowid == models.Trace.id)
+            .where(pk.in_(keys))
+            .group_by(pk, models.SpanCostDetail.token_type, models.SpanCostDetail.is_prompt)
+        )
+        results: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                id_,
+                token_type,
+                is_prompt,
+                cost,
+                tokens,
+            ) in data:
+                entry = SpanCostDetailSummaryEntry(
+                    token_type=token_type,
+                    is_prompt=is_prompt,
+                    value=CostBreakdown(tokens=tokens, cost=cost),
+                )
+                results[id_].append(entry)
+        return list(map(list, map(results.__getitem__, keys)))

phoenix/server/api/dataloaders/span_cost_detail_summary_entries_by_span.py

@@ -0,0 +1,43 @@
+from collections import defaultdict
+
+from sqlalchemy import select
+from sqlalchemy.orm import contains_eager
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.types import (
+    CostBreakdown,
+    SpanCostDetailSummaryEntry,
+)
+from phoenix.server.types import DbSessionFactory
+
+SpanRowID: TypeAlias = int
+Key: TypeAlias = SpanRowID
+Result: TypeAlias = list[SpanCostDetailSummaryEntry]
+
+
+class SpanCostDetailSummaryEntriesBySpanDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for span_cost_detail in await session.stream_scalars(
+                select(models.SpanCostDetail)
+                .join(models.SpanCost, models.SpanCostDetail.span_cost_id == models.SpanCost.id)
+                .where(models.SpanCost.span_rowid.in_(keys))
+                .options(contains_eager(models.SpanCostDetail.span_cost))
+            ):
+                entry = SpanCostDetailSummaryEntry(
+                    token_type=span_cost_detail.token_type,
+                    is_prompt=span_cost_detail.is_prompt,
+                    value=CostBreakdown(
+                        tokens=span_cost_detail.tokens,
+                        cost=span_cost_detail.cost,
+                    ),
+                )
+                results[span_cost_detail.span_cost.span_rowid].append(entry)
+        return list(map(list, map(results.__getitem__, keys)))