arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/server/api/dataloaders/experiment_run_counts.py

@@ -0,0 +1,49 @@
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentID: TypeAlias = int
+RunCount: TypeAlias = int
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = RunCount
+
+
+class ExperimentRunCountsDataLoader(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]:
+        experiment_ids = keys
+        resolved_experiment_ids = (
+            select(models.Experiment.id)
+            .where(models.Experiment.id.in_(set(experiment_ids)))
+            .subquery()
+        )
+        query = (
+            select(
+                resolved_experiment_ids.c.id,
+                func.count(models.ExperimentRun.experiment_id),
+            )
+            .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)
+        )
+        async with self._db() as session:
+            run_counts = {
+                experiment_id: run_count
+                async for experiment_id, run_count in await session.stream(query)
+            }
+        return [
+            run_counts.get(experiment_id, ValueError(f"Unknown experiment: {experiment_id}"))
+            for experiment_id in keys
+        ]

phoenix/server/api/dataloaders/experiment_sequence_number.py

@@ -0,0 +1,44 @@
+from typing import Optional
+
+from sqlalchemy import distinct, func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExperimentId: TypeAlias = int
+Key: TypeAlias = ExperimentId
+Result: TypeAlias = Optional[int]
+
+
+class ExperimentSequenceNumberDataLoader(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]:
+        experiment_ids = keys
+        dataset_ids = (
+            select(distinct(models.Experiment.dataset_id))
+            .where(models.Experiment.id.in_(experiment_ids))
+            .scalar_subquery()
+        )
+        row_number = (
+            func.row_number().over(
+                partition_by=models.Experiment.dataset_id,
+                order_by=models.Experiment.id,
+            )
+        ).label("row_number")
+        subq = (
+            select(models.Experiment.id, row_number)
+            .where(models.Experiment.dataset_id.in_(dataset_ids))
+            .subquery()
+        )
+        stmt = select(subq).where(subq.c.id.in_(experiment_ids))
+        async with self._db() as session:
+            result = {
+                experiment_id: sequence_number
+                async for experiment_id, sequence_number in await session.stream(stmt)
+            }
+        return [result.get(experiment_id) for experiment_id in keys]

phoenix/server/api/dataloaders/latency_ms_quantile.py

@@ -0,0 +1,188 @@
+from collections import defaultdict
+from collections.abc import AsyncIterator, Mapping
+from datetime import datetime
+from typing import Any, Literal, Optional, 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.server.types import DbSessionFactory
+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: 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:
+            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,85 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import Literal, Optional
+
+from cachetools import LFUCache
+from sqlalchemy import 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.types import DbSessionFactory
+
+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: DbSessionFactory,
+        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/project_by_name.py

@@ -0,0 +1,31 @@
+from collections import defaultdict
+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
+
+ProjectName: TypeAlias = str
+Key: TypeAlias = ProjectName
+Result: TypeAlias = Optional[models.Project]
+
+
+class ProjectByNameDataLoader(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]:
+        project_names = list(set(keys))
+        projects_by_name: defaultdict[Key, Result] = defaultdict(None)
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(models.Project).where(models.Project.name.in_(project_names))
+            )
+            async for project in data:
+                projects_by_name[project.name] = project
+
+        return [projects_by_name.get(project_name) for project_name in keys]

phoenix/server/api/dataloaders/record_counts.py

@@ -0,0 +1,116 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Literal, Optional
+
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import Select, 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.types import DbSessionFactory
+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: 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 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/session_io.py

@@ -0,0 +1,79 @@
+from functools import cached_property
+from typing import Literal, Optional, cast
+
+from openinference.semconv.trace import SpanAttributes
+from sqlalchemy import Select, func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias, assert_never
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+from phoenix.trace.schemas import MimeType, SpanIOValue
+
+Key: TypeAlias = int
+Result: TypeAlias = Optional[SpanIOValue]
+
+Kind = Literal["first_input", "last_output"]
+
+
+class SessionIODataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory, kind: Kind) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+        self._kind = kind
+
+    @cached_property
+    def _subq(self) -> Select[tuple[Optional[int], str, str, int]]:
+        stmt = (
+            select(models.Trace.project_session_rowid.label("id_"))
+            .join_from(models.Span, models.Trace)
+            .where(models.Span.parent_id.is_(None))
+        )
+        if self._kind == "first_input":
+            stmt = stmt.add_columns(
+                models.Span.attributes[INPUT_VALUE].as_string().label("value"),
+                models.Span.attributes[INPUT_MIME_TYPE].as_string().label("mime_type"),
+                func.row_number()
+                .over(
+                    partition_by=models.Trace.project_session_rowid,
+                    order_by=[models.Trace.start_time.asc(), models.Trace.id.asc()],
+                )
+                .label("rank"),
+            )
+        elif self._kind == "last_output":
+            stmt = stmt.add_columns(
+                models.Span.attributes[OUTPUT_VALUE].as_string().label("value"),
+                models.Span.attributes[OUTPUT_MIME_TYPE].as_string().label("mime_type"),
+                func.row_number()
+                .over(
+                    partition_by=models.Trace.project_session_rowid,
+                    order_by=[models.Trace.start_time.desc(), models.Trace.id.desc()],
+                )
+                .label("rank"),
+            )
+        else:
+            assert_never(self._kind)
+        return cast(Select[tuple[Optional[int], str, str, int]], stmt)
+
+    def _stmt(self, *keys: Key) -> Select[tuple[int, str, str]]:
+        subq = self._subq.where(models.Trace.project_session_rowid.in_(keys)).subquery()
+        return (
+            select(subq.c.id_, subq.c.value, subq.c.mime_type)
+            .filter_by(rank=1)
+            .where(subq.c.value.isnot(None))
+        )
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        async with self._db() as session:
+            result: dict[Key, SpanIOValue] = {
+                id_: SpanIOValue(value=value, mime_type=MimeType(mime_type))
+                async for id_, value, mime_type in await session.stream(self._stmt(*keys))
+                if id_ is not None
+            }
+        return [result.get(key) for key in keys]
+
+
+INPUT_VALUE = SpanAttributes.INPUT_VALUE.split(".")
+INPUT_MIME_TYPE = SpanAttributes.INPUT_MIME_TYPE.split(".")
+OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE.split(".")
+OUTPUT_MIME_TYPE = SpanAttributes.OUTPUT_MIME_TYPE.split(".")

phoenix/server/api/dataloaders/session_num_traces.py

@@ -0,0 +1,30 @@
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = int
+
+
+class SessionNumTracesDataLoader(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.Trace.project_session_rowid.label("id_"),
+                func.count(models.Trace.id).label("value"),
+            )
+            .group_by(models.Trace.project_session_rowid)
+            .where(models.Trace.project_session_rowid.in_(keys))
+        )
+        async with self._db() as session:
+            result: dict[Key, int] = {
+                id_: value async for id_, value in await session.stream(stmt) if id_ is not None
+            }
+        return [result.get(key, 0) for key in keys]

phoenix/server/api/dataloaders/session_num_traces_with_error.py

@@ -0,0 +1,32 @@
+from sqlalchemy import func, select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = int
+
+
+class SessionNumTracesWithErrorDataLoader(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.Trace.project_session_rowid.label("id_"),
+                func.count(models.Trace.id).label("value"),
+            )
+            .join(models.Span)
+            .group_by(models.Trace.project_session_rowid)
+            .where(models.Span.cumulative_error_count > 0)
+            .where(models.Trace.project_session_rowid.in_(keys))
+        )
+        async with self._db() as session:
+            result: dict[Key, int] = {
+                id_: value async for id_, value in await session.stream(stmt) if id_ is not None
+            }
+        return [result.get(key, 0) for key in keys]