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

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,57 @@
+from collections import defaultdict
+
+from sqlalchemy import func, select
+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,
+                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(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_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_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,152 @@
+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.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,
+    SessionFilterCondition,
+]
+Param: TypeAlias = ProjectRowId
+
+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, session_filter_condition = key
+    interval = (
+        (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
+    )
+    return (interval, filter_condition, session_filter_condition), project_rowid
+
+
+_Section: TypeAlias = ProjectRowId
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition, SessionFilterCondition]
+
+
+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, session_filter_condition), project_rowid = _cache_key_fn(key)
+        return project_rowid, (interval, filter_condition, session_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]:
+    project_rowids = params
+    (start_time, end_time), filter_condition, session_filter_condition = segment
+
+    stmt: Select[Any] = (
+        select(
+            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"),
+            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.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:
+        stmt = stmt.where(start_time <= models.Trace.start_time)
+    if end_time:
+        stmt = stmt.where(models.Trace.start_time < end_time)
+
+    if filter_condition:
+        sf = SpanFilter(filter_condition)
+        stmt = sf(stmt.join_from(models.SpanCost, models.Span))
+
+    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_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))

phoenix/server/api/dataloaders/span_costs.py

@@ -0,0 +1,29 @@
+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
+
+SpanID: TypeAlias = int
+Key: TypeAlias = SpanID
+Result: TypeAlias = Optional[models.SpanCost]
+
+
+class SpanCostsDataLoader(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]:
+        span_ids = list(set(keys))
+        async with self._db() as session:
+            costs = {
+                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]

phoenix/server/api/dataloaders/table_fields.py

@@ -18,7 +18,7 @@ _AttrStrIdentifier: TypeAlias = str
 
 
 class TableFieldsDataLoader(DataLoader[Key, Result]):
-    def __init__(self, db: DbSessionFactory, table: type[models.Base]) -> None:
+    def __init__(self, db: DbSessionFactory, table: type[models.HasId]) -> None:
         super().__init__(load_fn=self._load_fn)
         self._db = db
         self._table = table
@@ -37,7 +37,7 @@ class TableFieldsDataLoader(DataLoader[Key, Result]):
 
 def _get_stmt(
     keys: Iterable[tuple[RowId, QueryableAttribute[Any]]],
-    table: type[models.Base],
+    table: type[models.HasId],
 ) -> tuple[
     Select[Any],
     dict[_ResultColumnPosition, _AttrStrIdentifier],

phoenix/server/api/dataloaders/token_prices_by_model.py

@@ -0,0 +1,30 @@
+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
+
+ModelId: TypeAlias = int
+Key: TypeAlias = ModelId
+Result: TypeAlias = list[models.TokenPrice]
+
+
+class TokenPricesByModelDataLoader(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]:
+        model_ids = keys
+        token_prices: defaultdict[Key, Result] = defaultdict(list)
+
+        async with self._db() as session:
+            async for token_price in await session.stream_scalars(
+                select(models.TokenPrice).where(models.TokenPrice.model_id.in_(model_ids))
+            ):
+                token_prices[token_price.model_id].append(token_price)
+
+        return [token_prices[model_id] for model_id in keys]