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

phoenix/server/api/dataloaders/experiment_error_rates.py

@@ -1,6 +1,6 @@
 from typing import Optional
 
-from sqlalchemy import case, func, select
+from sqlalchemy import func, select
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -23,36 +23,29 @@ class ExperimentErrorRatesDataLoader(DataLoader[Key, Result]):
 
     async def _load_fn(self, keys: list[Key]) -> list[Result]:
         experiment_ids = keys
-        resolved_experiment_ids = (
-            select(models.Experiment.id)
-            .where(models.Experiment.id.in_(set(experiment_ids)))
-            .subquery()
-        )
-        query = (
+        average_repetition_error_rates_subquery = (
             select(
-                resolved_experiment_ids.c.id,
-                case(
-                    (
-                        func.count(models.ExperimentRun.id) != 0,
-                        func.count(models.ExperimentRun.error)
-                        / func.count(models.ExperimentRun.id),
-                    ),
-                    else_=None,
-                ),
+                models.ExperimentRun.experiment_id.label("experiment_id"),
+                (
+                    func.count(models.ExperimentRun.error) / func.count(models.ExperimentRun.id)
+                ).label("average_repetition_error_rate"),
             )
-            .outerjoin_from(
-                from_=resolved_experiment_ids,
-                target=models.ExperimentRun,
-                onclause=resolved_experiment_ids.c.id == models.ExperimentRun.experiment_id,
-            )
-            .group_by(resolved_experiment_ids.c.id)
+            .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+            .group_by(models.ExperimentRun.dataset_example_id, models.ExperimentRun.experiment_id)
+            .subquery()
+            .alias("average_repetition_error_rates")
         )
+        average_run_error_rates_query = select(
+            average_repetition_error_rates_subquery.c.experiment_id,
+            func.avg(average_repetition_error_rates_subquery.c.average_repetition_error_rate).label(
+                "average_run_error_rates"
+            ),
+        ).group_by(average_repetition_error_rates_subquery.c.experiment_id)
         async with self._db() as session:
-            error_rates = {
+            average_run_error_rates = {
                 experiment_id: error_rate
-                async for experiment_id, error_rate in await session.stream(query)
+                async for experiment_id, error_rate in await session.stream(
+                    average_run_error_rates_query
+                )
             }
-        return [
-            error_rates.get(experiment_id, ValueError(f"Unknown experiment ID: {experiment_id}"))
-            for experiment_id in keys
-        ]
+        return [average_run_error_rates.get(experiment_id) for experiment_id in experiment_ids]

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/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_summary_by_experiment_repeated_run_group.py

@@ -0,0 +1,64 @@
+from collections import defaultdict
+
+from sqlalchemy import func, select, tuple_
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.dataloaders.types import CostBreakdown, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+Key: TypeAlias = tuple[ExperimentId, DatasetExampleId]
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByExperimentRepeatedRunGroupDataLoader(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.ExperimentRun.experiment_id,
+                models.ExperimentRun.dataset_example_id,
+                func.sum(models.SpanCost.prompt_cost).label("prompt_cost"),
+                func.sum(models.SpanCost.completion_cost).label("completion_cost"),
+                func.sum(models.SpanCost.total_cost).label("total_cost"),
+                func.sum(models.SpanCost.prompt_tokens).label("prompt_tokens"),
+                func.sum(models.SpanCost.completion_tokens).label("completion_tokens"),
+                func.sum(models.SpanCost.total_tokens).label("total_tokens"),
+            )
+            .select_from(models.ExperimentRun)
+            .join(models.Trace, models.ExperimentRun.trace_id == models.Trace.trace_id)
+            .join(models.SpanCost, models.SpanCost.trace_rowid == models.Trace.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)
+        )
+
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                experiment_id,
+                dataset_example_id,
+                prompt_cost,
+                completion_cost,
+                total_cost,
+                prompt_tokens,
+                completion_tokens,
+                total_tokens,
+            ) in data:
+                summary = SpanCostSummary(
+                    prompt=CostBreakdown(tokens=prompt_tokens, cost=prompt_cost),
+                    completion=CostBreakdown(tokens=completion_tokens, cost=completion_cost),
+                    total=CostBreakdown(tokens=total_tokens, cost=total_cost),
+                )
+                results[(experiment_id, dataset_example_id)] = summary
+        return [results.get(key, SpanCostSummary()) for key in keys]

phoenix/server/api/dataloaders/span_cost_summary_by_project.py

@@ -12,32 +12,38 @@ from phoenix.db import models
 from phoenix.server.api.dataloaders.cache import TwoTierCache
 from phoenix.server.api.dataloaders.types import CostBreakdown, SpanCostSummary
 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
 
 ProjectRowId: TypeAlias = int
 TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
 FilterCondition: TypeAlias = Optional[str]
+SessionFilterCondition: TypeAlias = Optional[str]
 
-Segment: TypeAlias = tuple[TimeInterval, FilterCondition]
+Segment: TypeAlias = tuple[
+    TimeInterval,
+    FilterCondition,
+    SessionFilterCondition,
+]
 Param: TypeAlias = ProjectRowId
 
-Key: TypeAlias = tuple[ProjectRowId, Optional[TimeRange], FilterCondition]
+Key: TypeAlias = tuple[ProjectRowId, Optional[TimeRange], FilterCondition, SessionFilterCondition]
 Result: TypeAlias = SpanCostSummary
 ResultPosition: TypeAlias = int
 DEFAULT_VALUE: Result = SpanCostSummary()
 
 
 def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
-    project_rowid, time_range, filter_condition = key
+    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 (interval, filter_condition), project_rowid
+    return (interval, filter_condition, session_filter_condition), project_rowid
 
 
 _Section: TypeAlias = ProjectRowId
-_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition]
 
 
 class SpanCostSummaryCache(
@@ -53,8 +59,8 @@ class SpanCostSummaryCache(
         )
 
     def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
-        (interval, filter_condition), project_rowid = _cache_key_fn(key)
-        return project_rowid, (interval, filter_condition)
+        (interval, filter_condition, session_filter_condition), project_rowid = _cache_key_fn(key)
+        return project_rowid, (interval, filter_condition, session_filter_condition)
 
 
 class SpanCostSummaryByProjectDataLoader(DataLoader[Key, Result]):
@@ -106,12 +112,12 @@ def _get_stmt(
     segment: Segment,
     *params: Param,
 ) -> Select[Any]:
-    (start_time, end_time), filter_condition = segment
-    pid = models.Trace.project_rowid
+    project_rowids = params
+    (start_time, end_time), filter_condition, session_filter_condition = segment
 
     stmt: Select[Any] = (
         select(
-            pid,
+            models.Trace.project_rowid,
             coalesce(func.sum(models.SpanCost.prompt_cost), 0).label("prompt_cost"),
             coalesce(func.sum(models.SpanCost.completion_cost), 0).label("completion_cost"),
             coalesce(func.sum(models.SpanCost.total_cost), 0).label("total_cost"),
@@ -119,8 +125,10 @@ def _get_stmt(
             coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
             coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
         )
-        .join_from(models.SpanCost, models.Trace)
-        .group_by(pid)
+        .select_from(models.Trace)
+        .join(models.SpanCost, models.Trace.id == models.SpanCost.trace_rowid)
+        .where(models.Trace.project_rowid.in_(project_rowids))
+        .group_by(models.Trace.project_rowid)
     )
 
     if start_time:
@@ -132,7 +140,13 @@ def _get_stmt(
         sf = SpanFilter(filter_condition)
         stmt = sf(stmt.join_from(models.SpanCost, models.Span))
 
-    project_ids = [rowid for rowid in params]
-    stmt = stmt.where(pid.in_(project_ids))
+    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))
 
     return stmt

phoenix/server/api/dataloaders/span_costs.py

@@ -1,7 +1,6 @@
 from typing import Optional
 
 from sqlalchemy import select
-from sqlalchemy.orm import joinedload, load_only
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
@@ -22,14 +21,9 @@ class SpanCostsDataLoader(DataLoader[Key, Result]):
         span_ids = list(set(keys))
         async with self._db() as session:
             costs = {
-                span.id: span.span_cost
-                async for span in await session.stream_scalars(
-                    select(models.Span)
-                    .where(models.Span.id.in_(span_ids))
-                    .options(
-                        load_only(models.Span.id),
-                        joinedload(models.Span.span_cost),
-                    )
+                span_cost.span_rowid: span_cost
+                async for span_cost in await session.stream_scalars(
+                    select(models.SpanCost).where(models.SpanCost.span_rowid.in_(span_ids))
                 )
             }
         return [costs.get(span_id) for span_id in keys]