arize-phoenix 10.15.0__py3-none-any.whl → 11.1.0__py3-none-any.whl

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)))

phoenix/server/api/dataloaders/span_cost_detail_summary_entries_by_trace.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
+
+TraceRowId: TypeAlias = int
+Key: TypeAlias = TraceRowId
+Result: TypeAlias = list[SpanCostDetailSummaryEntry]
+
+
+class SpanCostDetailSummaryEntriesByTraceDataLoader(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.trace_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)
+            .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_details_by_span_cost.py

@@ -0,0 +1,27 @@
+from collections import defaultdict
+
+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
+
+SpanCostId: TypeAlias = int
+Key: TypeAlias = SpanCostId
+Result: TypeAlias = list[models.SpanCostDetail]
+
+
+class SpanCostDetailsBySpanCostDataLoader(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]:
+        entity = models.SpanCostDetail
+        stmt = select(entity).where(entity.span_cost_id.in_(keys))
+        result: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for obj in await session.stream_scalars(stmt):
+                result[obj.span_cost_id].append(obj)
+        return list(map(result.__getitem__, keys))

phoenix/server/api/dataloaders/span_cost_summary_by_experiment.py

@@ -0,0 +1,58 @@
+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, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+Key: TypeAlias = ExperimentId
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByExperimentDataLoader(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,
+                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"),
+                coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
+                coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
+                coalesce(func.sum(models.SpanCost.total_tokens), 0).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(models.ExperimentRun.experiment_id.in_(keys))
+            .group_by(models.ExperimentRun.experiment_id)
+        )
+
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                experiment_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] = summary
+        return list(map(results.__getitem__, keys))

phoenix/server/api/dataloaders/span_cost_summary_by_experiment_run.py

@@ -0,0 +1,58 @@
+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, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+ExperimentRunId: TypeAlias = int
+Key: TypeAlias = ExperimentRunId
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByExperimentRunDataLoader(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.id,
+                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"),
+                coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
+                coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
+                coalesce(func.sum(models.SpanCost.total_tokens), 0).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(models.ExperimentRun.id.in_(keys))
+            .group_by(models.ExperimentRun.id)
+        )
+
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                run_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[run_id] = summary
+        return list(map(results.__getitem__, keys))

phoenix/server/api/dataloaders/span_cost_summary_by_generative_model.py

@@ -0,0 +1,55 @@
+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, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+GenerativeModelId: TypeAlias = int
+Key: TypeAlias = GenerativeModelId
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByGenerativeModelDataLoader(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,
+                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"),
+                coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
+                coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
+                coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
+            )
+            .where(pk.in_(keys))
+            .group_by(pk)
+        )
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                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[id_] = summary
+        return list(map(results.__getitem__, keys))

phoenix/server/api/dataloaders/span_cost_summary_by_project.py

@@ -0,0 +1,140 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Optional
+
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import Select, func, select
+from sqlalchemy.sql.functions import coalesce
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias
+
+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.types import DbSessionFactory
+from phoenix.trace.dsl import SpanFilter
+
+ProjectRowId: TypeAlias = int
+TimeInterval: TypeAlias = tuple[Optional[datetime], Optional[datetime]]
+FilterCondition: TypeAlias = Optional[str]
+
+Segment: TypeAlias = tuple[TimeInterval, FilterCondition]
+Param: TypeAlias = ProjectRowId
+
+Key: TypeAlias = tuple[ProjectRowId, Optional[TimeRange], FilterCondition]
+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
+    interval = (
+        (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
+    )
+    return (interval, filter_condition), project_rowid
+
+
+_Section: TypeAlias = ProjectRowId
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+
+
+class SpanCostSummaryCache(
+    TwoTierCache[Key, Result, _Section, _SubKey],
+):
+    def __init__(self) -> None:
+        super().__init__(
+            # TTL=3600 (1-hour) because time intervals are always moving forward, but
+            # interval endpoints are rounded down to the hour by the UI, so anything
+            # older than an hour most likely won't be a cache-hit anyway.
+            main_cache=TTLCache(maxsize=64, ttl=3600),
+            sub_cache_factory=lambda: LFUCache(maxsize=2 * 2 * 3),
+        )
+
+    def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
+        (interval, filter_condition), project_rowid = _cache_key_fn(key)
+        return project_rowid, (interval, filter_condition)
+
+
+class SpanCostSummaryByProjectDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        cache_map: Optional[AbstractCache[Key, Result]] = None,
+    ) -> None:
+        super().__init__(
+            load_fn=self._load_fn,
+            cache_key_fn=_cache_key_fn,
+            cache_map=cache_map,
+        )
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        results: list[Result] = [DEFAULT_VALUE] * len(keys)
+        arguments: defaultdict[
+            Segment,
+            defaultdict[Param, list[ResultPosition]],
+        ] = defaultdict(lambda: defaultdict(list))
+        for position, key in enumerate(keys):
+            segment, param = _cache_key_fn(key)
+            arguments[segment][param].append(position)
+        async with self._db() as session:
+            for segment, params in arguments.items():
+                stmt = _get_stmt(segment, *params.keys())
+                data = await session.stream(stmt)
+                async for (
+                    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),
+                    )
+                    for position in params.get(id_, []):
+                        results[position] = summary
+        return results
+
+
+def _get_stmt(
+    segment: Segment,
+    *params: Param,
+) -> Select[Any]:
+    (start_time, end_time), filter_condition = segment
+    pid = models.Trace.project_rowid
+
+    stmt: Select[Any] = (
+        select(
+            pid,
+            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"),
+            coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
+            coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
+            coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
+        )
+        .select_from(models.Trace)
+        .join(models.Span, models.Span.trace_rowid == models.Trace.id)
+        .join(models.SpanCost, models.Span.id == models.SpanCost.span_rowid)
+        .group_by(pid)
+    )
+
+    if start_time:
+        stmt = stmt.where(start_time <= models.Span.start_time)
+    if end_time:
+        stmt = stmt.where(models.Span.start_time < end_time)
+
+    if filter_condition:
+        sf = SpanFilter(filter_condition)
+        stmt = sf(stmt)
+
+    project_ids = [rowid for rowid in params]
+    stmt = stmt.where(pid.in_(project_ids))
+
+    return stmt

phoenix/server/api/dataloaders/span_cost_summary_by_project_session.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, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+ProjectSessionRowId: TypeAlias = int
+Key: TypeAlias = ProjectSessionRowId
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByProjectSessionDataLoader(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,
+                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"),
+                coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
+                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)
+            .where(pk.in_(keys))
+            .group_by(pk)
+        )
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                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[id_] = summary
+        return list(map(results.__getitem__, keys))

phoenix/server/api/dataloaders/span_cost_summary_by_trace.py

@@ -0,0 +1,55 @@
+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, SpanCostSummary
+from phoenix.server.types import DbSessionFactory
+
+TraceRowId: TypeAlias = int
+Key: TypeAlias = TraceRowId
+Result: TypeAlias = SpanCostSummary
+
+
+class SpanCostSummaryByTraceDataLoader(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.trace_rowid
+        stmt = (
+            select(
+                pk,
+                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"),
+                coalesce(func.sum(models.SpanCost.prompt_tokens), 0).label("prompt_tokens"),
+                coalesce(func.sum(models.SpanCost.completion_tokens), 0).label("completion_tokens"),
+                coalesce(func.sum(models.SpanCost.total_tokens), 0).label("total_tokens"),
+            )
+            .where(pk.in_(keys))
+            .group_by(pk)
+        )
+        results: defaultdict[Key, Result] = defaultdict(SpanCostSummary)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (
+                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[id_] = summary
+        return list(map(results.__getitem__, keys))