arize-phoenix 3.25.0__py3-none-any.whl → 4.0.0__py3-none-any.whl

phoenix/server/api/dataloaders/latency_ms_quantile.py

@@ -0,0 +1,198 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import (
+    Any,
+    AsyncContextManager,
+    AsyncIterator,
+    Callable,
+    DefaultDict,
+    List,
+    Literal,
+    Mapping,
+    Optional,
+    Tuple,
+    cast,
+)
+
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import (
+    ARRAY,
+    Float,
+    Integer,
+    Select,
+    SQLColumnExpression,
+    Values,
+    column,
+    func,
+    select,
+    values,
+)
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.functions import percentile_cont
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias, assert_never
+
+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.trace.dsl import SpanFilter
+
+Kind: TypeAlias = Literal["span", "trace"]
+ProjectRowId: TypeAlias = int
+TimeInterval: TypeAlias = Tuple[Optional[datetime], Optional[datetime]]
+FilterCondition: TypeAlias = Optional[str]
+Probability: TypeAlias = float
+QuantileValue: TypeAlias = float
+
+Segment: TypeAlias = Tuple[Kind, TimeInterval, FilterCondition]
+Param: TypeAlias = Tuple[ProjectRowId, Probability]
+
+Key: TypeAlias = Tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition, Probability]
+Result: TypeAlias = Optional[QuantileValue]
+ResultPosition: TypeAlias = int
+DEFAULT_VALUE: Result = None
+
+FloatCol: TypeAlias = SQLColumnExpression[Float[float]]
+
+
+def _cache_key_fn(key: Key) -> Tuple[Segment, Param]:
+    kind, project_rowid, time_range, 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)
+
+
+_Section: TypeAlias = ProjectRowId
+_SubKey: TypeAlias = Tuple[TimeInterval, FilterCondition, Kind, Probability]
+
+
+class LatencyMsQuantileCache(
+    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 * 2 * 16),
+        )
+
+    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)
+
+
+class LatencyMsQuantileDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+        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:
+            dialect = SupportedSQLDialect(session.bind.dialect.name)
+            for segment, params in arguments.items():
+                async for position, quantile_value in _get_results(
+                    dialect, session, segment, params
+                ):
+                    results[position] = quantile_value
+        return results
+
+
+async def _get_results(
+    dialect: SupportedSQLDialect,
+    session: AsyncSession,
+    segment: Segment,
+    params: Mapping[Param, List[ResultPosition]],
+) -> AsyncIterator[Tuple[ResultPosition, QuantileValue]]:
+    kind, (start_time, end_time), 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
+    elif kind == "span":
+        latency_column = cast(FloatCol, models.Span.latency_ms)
+        time_column = models.Span.start_time
+        stmt = stmt.join(models.Span)
+        if filter_condition:
+            sf = SpanFilter(filter_condition)
+            stmt = sf(stmt)
+    else:
+        assert_never(kind)
+    if start_time:
+        stmt = stmt.where(start_time <= time_column)
+    if end_time:
+        stmt = stmt.where(time_column < end_time)
+    if dialect is SupportedSQLDialect.POSTGRESQL:
+        results = _get_results_postgresql(session, stmt, latency_column, params)
+    elif dialect is SupportedSQLDialect.SQLITE:
+        results = _get_results_sqlite(session, stmt, latency_column, params)
+    else:
+        assert_never(dialect)
+    async for position, quantile_value in results:
+        yield position, quantile_value
+
+
+async def _get_results_sqlite(
+    session: AsyncSession,
+    base_stmt: Select[Any],
+    latency_column: FloatCol,
+    params: Mapping[Param, List[ResultPosition]],
+) -> AsyncIterator[Tuple[ResultPosition, QuantileValue]]:
+    projects_per_prob: DefaultDict[Probability, List[ProjectRowId]] = defaultdict(list)
+    for project_rowid, probability in params.keys():
+        projects_per_prob[probability].append(project_rowid)
+    pid = models.Trace.project_rowid
+    for probability, project_rowids in projects_per_prob.items():
+        pctl: FloatCol = func.percentile(latency_column, probability * 100)
+        stmt = base_stmt.add_columns(pctl)
+        stmt = stmt.where(pid.in_(project_rowids))
+        stmt = stmt.group_by(pid)
+        data = await session.stream(stmt)
+        async for project_rowid, quantile_value in data:
+            for position in params[(project_rowid, probability)]:
+                yield position, quantile_value
+
+
+async def _get_results_postgresql(
+    session: AsyncSession,
+    base_stmt: Select[Any],
+    latency_column: FloatCol,
+    params: Mapping[Param, List[ResultPosition]],
+) -> AsyncIterator[Tuple[ResultPosition, QuantileValue]]:
+    probs_per_project: DefaultDict[ProjectRowId, List[Probability]] = defaultdict(list)
+    for project_rowid, probability in params.keys():
+        probs_per_project[project_rowid].append(probability)
+    pp: Values = values(
+        column("project_rowid", Integer),
+        column("probabilities", ARRAY(Float[float])),
+        name="project_probabilities",
+    ).data(probs_per_project.items())  # type: ignore
+    pid = models.Trace.project_rowid
+    pctl: FloatCol = percentile_cont(pp.c.probabilities).within_group(latency_column)
+    stmt = base_stmt.add_columns(pp.c.probabilities, pctl)
+    stmt = stmt.join(pp, pid == pp.c.project_rowid)
+    stmt = stmt.group_by(pid, pp.c.probabilities)
+    data = await session.stream(stmt)
+    async for project_rowid, probabilities, quantile_values in data:
+        for probability, quantile_value in zip(probabilities, quantile_values):
+            for position in params[(project_rowid, probability)]:
+                yield position, quantile_value

phoenix/server/api/dataloaders/min_start_or_max_end_times.py

@@ -0,0 +1,93 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import (
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+)
+
+from cachetools import LFUCache
+from sqlalchemy import func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+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
+
+Kind: TypeAlias = Literal["start", "end"]
+ProjectRowId: TypeAlias = int
+
+Segment: TypeAlias = ProjectRowId
+Param: TypeAlias = Kind
+
+Key: TypeAlias = Tuple[ProjectRowId, Kind]
+Result: TypeAlias = Optional[datetime]
+ResultPosition: TypeAlias = int
+DEFAULT_VALUE: Result = None
+
+_Section = ProjectRowId
+_SubKey = Kind
+
+
+class MinStartOrMaxEndTimeCache(
+    TwoTierCache[Key, Result, _Section, _SubKey],
+):
+    def __init__(self) -> None:
+        super().__init__(
+            main_cache=LFUCache(maxsize=64),
+            sub_cache_factory=lambda: LFUCache(maxsize=2),
+        )
+
+    def _cache_key(self, key: Key) -> Tuple[_Section, _SubKey]:
+        return key
+
+
+class MinStartOrMaxEndTimeDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+        cache_map: Optional[AbstractCache[Key, Result]] = None,
+    ) -> None:
+        super().__init__(
+            load_fn=self._load_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 = key
+            arguments[segment][param].append(position)
+        pid = models.Trace.project_rowid
+        stmt = (
+            select(
+                pid,
+                func.min(models.Trace.start_time).label("min_start"),
+                func.max(models.Trace.end_time).label("max_end"),
+            )
+            .where(pid.in_(arguments.keys()))
+            .group_by(pid)
+        )
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for project_rowid, min_start, max_end in data:
+                for kind, positions in arguments[project_rowid].items():
+                    if kind == "start":
+                        for position in positions:
+                            results[position] = min_start
+                    elif kind == "end":
+                        for position in positions:
+                            results[position] = max_end
+                    else:
+                        assert_never(kind)
+        return results

phoenix/server/api/dataloaders/record_counts.py

@@ -0,0 +1,125 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import (
+    Any,
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+)
+
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import Select, func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+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.trace.dsl import SpanFilter
+
+Kind: TypeAlias = Literal["span", "trace"]
+ProjectRowId: TypeAlias = int
+TimeInterval: TypeAlias = Tuple[Optional[datetime], Optional[datetime]]
+FilterCondition: TypeAlias = Optional[str]
+SpanCount: TypeAlias = int
+
+Segment: TypeAlias = Tuple[Kind, TimeInterval, FilterCondition]
+Param: TypeAlias = ProjectRowId
+
+Key: TypeAlias = Tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+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
+    interval = (
+        (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
+    )
+    return (kind, interval, filter_condition), project_rowid
+
+
+_Section: TypeAlias = ProjectRowId
+_SubKey: TypeAlias = Tuple[TimeInterval, FilterCondition, Kind]
+
+
+class RecordCountCache(
+    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 * 2),
+        )
+
+    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)
+
+
+class RecordCountDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+        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 project_rowid, count in data:
+                    for position in params[project_rowid]:
+                        results[position] = count
+        return results
+
+
+def _get_stmt(
+    segment: Segment,
+    *project_rowids: Param,
+) -> Select[Any]:
+    kind, (start_time, end_time), filter_condition = segment
+    pid = models.Trace.project_rowid
+    stmt = select(pid)
+    if kind == "span":
+        time_column = models.Span.start_time
+        stmt = stmt.join(models.Span)
+        if filter_condition:
+            sf = SpanFilter(filter_condition)
+            stmt = sf(stmt)
+    elif kind == "trace":
+        time_column = models.Trace.start_time
+    else:
+        assert_never(kind)
+    stmt = stmt.add_columns(func.count().label("count"))
+    stmt = stmt.where(pid.in_(project_rowids))
+    stmt = stmt.group_by(pid)
+    if start_time:
+        stmt = stmt.where(start_time <= time_column)
+    if end_time:
+        stmt = stmt.where(time_column < end_time)
+    return stmt

phoenix/server/api/dataloaders/span_descendants.py

@@ -0,0 +1,64 @@
+from random import randint
+from typing import (
+    AsyncContextManager,
+    Callable,
+    Dict,
+    List,
+)
+
+from aioitertools.itertools import groupby
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import contains_eager
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+
+SpanId: TypeAlias = str
+
+Key: TypeAlias = SpanId
+Result: TypeAlias = List[models.Span]
+
+
+class SpanDescendantsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        root_ids = set(keys)
+        root_id_label = f"root_id_{randint(0, 10**6):06}"
+        descendant_ids = (
+            select(
+                models.Span.id,
+                models.Span.span_id,
+                models.Span.parent_id.label(root_id_label),
+            )
+            .where(models.Span.parent_id.in_(root_ids))
+            .cte(recursive=True)
+        )
+        parent_ids = descendant_ids.alias()
+        descendant_ids = descendant_ids.union_all(
+            select(
+                models.Span.id,
+                models.Span.span_id,
+                parent_ids.c[root_id_label],
+            ).join(
+                parent_ids,
+                models.Span.parent_id == parent_ids.c.span_id,
+            )
+        )
+        stmt = (
+            select(descendant_ids.c[root_id_label], models.Span)
+            .join(descendant_ids, models.Span.id == descendant_ids.c.id)
+            .join(models.Trace)
+            .options(contains_eager(models.Span.trace))
+            .order_by(descendant_ids.c[root_id_label])
+        )
+        results: Dict[SpanId, Result] = {key: [] for key in keys}
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for root_id, group in groupby(data, key=lambda d: d[0]):
+                results[root_id].extend(span for _, span in group)
+        return [results[key].copy() for key in keys]

phoenix/server/api/dataloaders/span_evaluations.py

@@ -0,0 +1,37 @@
+from collections import defaultdict
+from typing import (
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+)
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.types.Evaluation import SpanEvaluation
+
+Key: TypeAlias = int
+Result: TypeAlias = List[SpanEvaluation]
+
+
+class SpanEvaluationsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        span_evaluations_by_id: DefaultDict[Key, Result] = defaultdict(list)
+        msa = models.SpanAnnotation
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(msa).where(msa.span_rowid.in_(keys)).where(msa.annotator_kind == "LLM")
+            )
+            async for span_evaluation in data:
+                span_evaluations_by_id[span_evaluation.span_rowid].append(
+                    SpanEvaluation.from_sql_span_annotation(span_evaluation)
+                )
+        return [span_evaluations_by_id[key] for key in keys]

phoenix/server/api/dataloaders/token_counts.py

@@ -0,0 +1,138 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import (
+    Any,
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+)
+
+from cachetools import LFUCache, TTLCache
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import Select, func, select
+from sqlalchemy.ext.asyncio import AsyncSession
+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.input_types.TimeRange import TimeRange
+from phoenix.trace.dsl import SpanFilter
+
+Kind: TypeAlias = Literal["prompt", "completion", "total"]
+ProjectRowId: TypeAlias = int
+TimeInterval: TypeAlias = Tuple[Optional[datetime], Optional[datetime]]
+FilterCondition: TypeAlias = Optional[str]
+TokenCount: TypeAlias = int
+
+Segment: TypeAlias = Tuple[TimeInterval, FilterCondition]
+Param: TypeAlias = Tuple[ProjectRowId, Kind]
+
+Key: TypeAlias = Tuple[Kind, ProjectRowId, Optional[TimeRange], FilterCondition]
+Result: TypeAlias = TokenCount
+ResultPosition: TypeAlias = int
+DEFAULT_VALUE: Result = 0
+
+
+def _cache_key_fn(key: Key) -> Tuple[Segment, Param]:
+    kind, 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, kind)
+
+
+_Section: TypeAlias = ProjectRowId
+_SubKey: TypeAlias = Tuple[TimeInterval, FilterCondition, Kind]
+
+
+class TokenCountCache(
+    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, kind) = _cache_key_fn(key)
+        return project_rowid, (interval, filter_condition, kind)
+
+
+class TokenCountDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: Callable[[], AsyncContextManager[AsyncSession]],
+        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 project_rowid, prompt, completion, total in data:
+                    for position in params[(project_rowid, "prompt")]:
+                        results[position] = prompt
+                    for position in params[(project_rowid, "completion")]:
+                        results[position] = completion
+                    for position in params[(project_rowid, "total")]:
+                        results[position] = total
+        return results
+
+
+def _get_stmt(
+    segment: Segment,
+    *params: Param,
+) -> Select[Any]:
+    (start_time, end_time), filter_condition = segment
+    prompt = func.sum(models.Span.attributes[_LLM_TOKEN_COUNT_PROMPT].as_float())
+    completion = func.sum(models.Span.attributes[_LLM_TOKEN_COUNT_COMPLETION].as_float())
+    total = coalesce(prompt, 0) + coalesce(completion, 0)
+    pid = models.Trace.project_rowid
+    stmt: Select[Any] = (
+        select(
+            pid,
+            prompt.label("prompt"),
+            completion.label("completion"),
+            total.label("total"),
+        )
+        .join_from(models.Trace, models.Span)
+        .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)
+    stmt = stmt.where(pid.in_([rowid for rowid, _ in params]))
+    return stmt
+
+
+_LLM_TOKEN_COUNT_PROMPT = SpanAttributes.LLM_TOKEN_COUNT_PROMPT.split(".")
+_LLM_TOKEN_COUNT_COMPLETION = SpanAttributes.LLM_TOKEN_COUNT_COMPLETION.split(".")

phoenix/server/api/dataloaders/trace_evaluations.py

@@ -0,0 +1,37 @@
+from collections import defaultdict
+from typing import (
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    List,
+)
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.types.Evaluation import TraceEvaluation
+
+Key: TypeAlias = int
+Result: TypeAlias = List[TraceEvaluation]
+
+
+class TraceEvaluationsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: Callable[[], AsyncContextManager[AsyncSession]]) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: List[Key]) -> List[Result]:
+        trace_evaluations_by_id: DefaultDict[Key, Result] = defaultdict(list)
+        mta = models.TraceAnnotation
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(mta).where(mta.trace_rowid.in_(keys)).where(mta.annotator_kind == "LLM")
+            )
+            async for trace_evaluation in data:
+                trace_evaluations_by_id[trace_evaluation.trace_rowid].append(
+                    TraceEvaluation.from_sql_trace_annotation(trace_evaluation)
+                )
+        return [trace_evaluations_by_id[key] for key in keys]