arize-phoenix 3.16.1__py3-none-any.whl → 7.7.1__py3-none-any.whl

phoenix/server/api/dataloaders/cache/two_tier_cache.py

@@ -0,0 +1,68 @@
+"""
+The primary intent of a two-tier system is to make cache invalidation more efficient,
+because the cache keys are typically tuples such as (project_id, time_interval, ...),
+but we need to invalidate subsets of keys, e.g. all those associated with a
+specific project, very frequently (i.e. essentially at each span insertion). In a
+single-tier system we would need to check all the keys to see if they are in the
+subset that we want to invalidate.
+"""
+
+from abc import ABC, abstractmethod
+from asyncio import Future
+from collections.abc import Callable
+from typing import Any, Generic, Optional, TypeVar
+
+from cachetools import Cache
+from strawberry.dataloader import AbstractCache
+
+_Key = TypeVar("_Key")
+_Result = TypeVar("_Result")
+
+_Section = TypeVar("_Section")
+_SubKey = TypeVar("_SubKey")
+
+
+class TwoTierCache(
+    AbstractCache[_Key, _Result],
+    Generic[_Key, _Result, _Section, _SubKey],
+    ABC,
+):
+    def __init__(
+        self,
+        main_cache: "Cache[_Section, Cache[_SubKey, Future[_Result]]]",
+        sub_cache_factory: Callable[[], "Cache[_SubKey, Future[_Result]]"],
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(*args, **kwargs)
+        self._cache = main_cache
+        self._sub_cache_factory = sub_cache_factory
+
+    @abstractmethod
+    def _cache_key(self, key: _Key) -> tuple[_Section, _SubKey]: ...
+
+    def invalidate(self, section: _Section) -> None:
+        if sub_cache := self._cache.get(section):
+            sub_cache.clear()
+
+    def get(self, key: _Key) -> Optional["Future[_Result]"]:
+        section, sub_key = self._cache_key(key)
+        if not (sub_cache := self._cache.get(section)):
+            return None
+        return sub_cache.get(sub_key)
+
+    def set(self, key: _Key, value: "Future[_Result]") -> None:
+        section, sub_key = self._cache_key(key)
+        if (sub_cache := self._cache.get(section)) is None:
+            self._cache[section] = sub_cache = self._sub_cache_factory()
+        sub_cache[sub_key] = value
+
+    def delete(self, key: _Key) -> None:
+        section, sub_key = self._cache_key(key)
+        if sub_cache := self._cache.get(section):
+            del sub_cache[sub_key]
+            if not sub_cache:
+                del self._cache[section]
+
+    def clear(self) -> None:
+        self._cache.clear()

phoenix/server/api/dataloaders/dataset_example_revisions.py

@@ -0,0 +1,131 @@
+from typing import Optional, Union
+
+from sqlalchemy import and_, case, func, null, or_, select
+from sqlalchemy.sql.expression import literal
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.api.exceptions import NotFound
+from phoenix.server.api.types.DatasetExampleRevision import DatasetExampleRevision
+from phoenix.server.types import DbSessionFactory
+
+ExampleID: TypeAlias = int
+VersionID: TypeAlias = Optional[int]
+Key: TypeAlias = tuple[ExampleID, Optional[VersionID]]
+Result: TypeAlias = DatasetExampleRevision
+
+
+class DatasetExampleRevisionsDataLoader(DataLoader[Key, Result]):
+    def __init__(self, db: DbSessionFactory) -> None:
+        super().__init__(
+            load_fn=self._load_fn,
+            max_batch_size=200,  # needed to prevent the size of the query from getting too large
+        )
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Union[Result, NotFound]]:
+        example_and_version_ids = tuple(
+            set(
+                (example_id, version_id)
+                for example_id, version_id in keys
+                if version_id is not None
+            )
+        )
+        versionless_example_ids = tuple(
+            set(example_id for example_id, version_id in keys if version_id is None)
+        )
+        resolved_example_and_version_ids = (
+            (
+                select(
+                    models.DatasetExample.id.label("example_id"),
+                    models.DatasetVersion.id.label("version_id"),
+                )
+                .select_from(models.DatasetExample)
+                .join(
+                    models.DatasetVersion,
+                    onclause=literal(True),  # cross join
+                )
+                .where(
+                    or_(
+                        *(
+                            and_(
+                                models.DatasetExample.id == example_id,
+                                models.DatasetVersion.id == version_id,
+                            )
+                            for example_id, version_id in example_and_version_ids
+                        )
+                    )
+                )
+            )
+            .union(
+                select(
+                    models.DatasetExample.id.label("example_id"), null().label("version_id")
+                ).where(models.DatasetExample.id.in_(versionless_example_ids))
+            )
+            .subquery()
+        )
+        revision_ids = (
+            select(
+                resolved_example_and_version_ids.c.example_id,
+                resolved_example_and_version_ids.c.version_id,
+                func.max(models.DatasetExampleRevision.id).label("revision_id"),
+            )
+            .select_from(resolved_example_and_version_ids)
+            .join(
+                models.DatasetExampleRevision,
+                onclause=resolved_example_and_version_ids.c.example_id
+                == models.DatasetExampleRevision.dataset_example_id,
+            )
+            .where(
+                or_(
+                    resolved_example_and_version_ids.c.version_id.is_(None),
+                    models.DatasetExampleRevision.dataset_version_id
+                    <= resolved_example_and_version_ids.c.version_id,
+                )
+            )
+            .group_by(
+                resolved_example_and_version_ids.c.example_id,
+                resolved_example_and_version_ids.c.version_id,
+            )
+        ).subquery()
+        query = (
+            select(
+                revision_ids.c.example_id,
+                revision_ids.c.version_id,
+                case(
+                    (
+                        or_(
+                            revision_ids.c.version_id.is_(None),
+                            models.DatasetVersion.id.is_not(None),
+                        ),
+                        True,
+                    ),
+                    else_=False,
+                ).label("is_valid_version"),  # check that non-null versions exist
+                models.DatasetExampleRevision,
+            )
+            .select_from(revision_ids)
+            .join(
+                models.DatasetExampleRevision,
+                onclause=revision_ids.c.revision_id == models.DatasetExampleRevision.id,
+            )
+            .join(
+                models.DatasetVersion,
+                onclause=revision_ids.c.version_id == models.DatasetVersion.id,
+                isouter=True,  # keep rows where the version id is null
+            )
+            .where(models.DatasetExampleRevision.revision_kind != "DELETE")
+        )
+        async with self._db() as session:
+            results = {
+                (example_id, version_id): DatasetExampleRevision.from_orm_revision(revision)
+                async for (
+                    example_id,
+                    version_id,
+                    is_valid_version,
+                    revision,
+                ) in await session.stream(query)
+                if is_valid_version
+            }
+        return [results.get(key, NotFound("Could not find revision.")) for key in keys]

phoenix/server/api/dataloaders/dataset_example_spans.py

@@ -0,0 +1,38 @@
+from typing import Optional
+
+from sqlalchemy import select
+from sqlalchemy.orm import joinedload
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+ExampleID: TypeAlias = int
+Key: TypeAlias = ExampleID
+Result: TypeAlias = Optional[models.Span]
+
+
+class DatasetExampleSpansDataLoader(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]:
+        example_ids = keys
+        async with self._db() as session:
+            spans = {
+                example_id: span
+                async for example_id, span in await session.stream(
+                    select(models.DatasetExample.id, models.Span)
+                    .select_from(models.DatasetExample)
+                    .join(models.Span, models.DatasetExample.span_rowid == models.Span.id)
+                    .where(models.DatasetExample.id.in_(example_ids))
+                    .options(
+                        joinedload(models.Span.trace, innerjoin=True).load_only(
+                            models.Trace.trace_id
+                        )
+                    )
+                )
+            }
+        return [spans.get(example_id) for example_id in keys]

phoenix/server/api/dataloaders/document_evaluation_summaries.py

@@ -0,0 +1,144 @@
+from collections import defaultdict
+from datetime import datetime
+from typing import Any, Optional
+
+import numpy as np
+from aioitertools.itertools import groupby
+from cachetools import LFUCache, TTLCache
+from sqlalchemy import Select, select
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect, num_docs_col
+from phoenix.metrics.retrieval_metrics import RetrievalMetrics
+from phoenix.server.api.dataloaders.cache import TwoTierCache
+from phoenix.server.api.input_types.TimeRange import TimeRange
+from phoenix.server.api.types.DocumentEvaluationSummary import DocumentEvaluationSummary
+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]
+EvalName: TypeAlias = str
+
+Segment: TypeAlias = tuple[ProjectRowId, TimeInterval, FilterCondition]
+Param: TypeAlias = EvalName
+
+Key: TypeAlias = tuple[ProjectRowId, Optional[TimeRange], FilterCondition, EvalName]
+Result: TypeAlias = Optional[DocumentEvaluationSummary]
+ResultPosition: TypeAlias = int
+DEFAULT_VALUE: Result = None
+
+
+def _cache_key_fn(key: Key) -> tuple[Segment, Param]:
+    project_rowid, time_range, filter_condition, eval_name = key
+    interval = (
+        (time_range.start, time_range.end) if isinstance(time_range, TimeRange) else (None, None)
+    )
+    return (project_rowid, interval, filter_condition), eval_name
+
+
+_Section: TypeAlias = tuple[ProjectRowId, EvalName]
+_SubKey: TypeAlias = tuple[TimeInterval, FilterCondition]
+
+
+class DocumentEvaluationSummaryCache(
+    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 * 32, ttl=3600),
+            sub_cache_factory=lambda: LFUCache(maxsize=2 * 2),
+        )
+
+    def invalidate_project(self, project_rowid: ProjectRowId) -> None:
+        for section in self._cache.keys():
+            if section[0] == project_rowid:
+                del self._cache[section]
+
+    def _cache_key(self, key: Key) -> tuple[_Section, _SubKey]:
+        (project_rowid, interval, filter_condition), eval_name = _cache_key_fn(key)
+        return (project_rowid, eval_name), (interval, filter_condition)
+
+
+class DocumentEvaluationSummaryDataLoader(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)
+        for segment, params in arguments.items():
+            async with self._db() as session:
+                dialect = SupportedSQLDialect(session.bind.dialect.name)
+                stmt = _get_stmt(dialect, segment, *params.keys())
+                data = await session.stream(stmt)
+                async for eval_name, group in groupby(data, lambda d: d.name):
+                    metrics_collection = []
+                    async for (_, num_docs), subgroup in groupby(
+                        group, lambda g: (g.id, g.num_docs)
+                    ):
+                        scores = [np.nan] * num_docs
+                        for row in subgroup:
+                            scores[row.document_position] = row.score
+                        metrics_collection.append(RetrievalMetrics(scores))
+                    summary = DocumentEvaluationSummary(
+                        evaluation_name=eval_name,
+                        metrics_collection=metrics_collection,
+                    )
+                    for position in params[eval_name]:
+                        results[position] = summary
+        return results
+
+
+def _get_stmt(
+    dialect: SupportedSQLDialect,
+    segment: Segment,
+    *eval_names: Param,
+) -> Select[Any]:
+    project_rowid, (start_time, end_time), filter_condition = segment
+    mda = models.DocumentAnnotation
+    stmt = (
+        select(
+            mda.name,
+            models.Span.id,
+            num_docs_col(dialect),
+            mda.score,
+            mda.document_position,
+        )
+        .join(models.Trace)
+        .where(models.Trace.project_rowid == project_rowid)
+        .join(mda)
+        .where(mda.name.in_(eval_names))
+        .where(mda.annotator_kind == "LLM")
+        .where(mda.score.is_not(None))
+        .order_by(mda.name, models.Span.id)
+    )
+    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:
+        span_filter = SpanFilter(condition=filter_condition)
+        stmt = span_filter(stmt)
+    return stmt

phoenix/server/api/dataloaders/document_evaluations.py

@@ -0,0 +1,31 @@
+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.api.types.Evaluation import DocumentEvaluation
+from phoenix.server.types import DbSessionFactory
+
+Key: TypeAlias = int
+Result: TypeAlias = list[DocumentEvaluation]
+
+
+class DocumentEvaluationsDataLoader(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]:
+        document_evaluations_by_id: defaultdict[Key, Result] = defaultdict(list)
+        mda = models.DocumentAnnotation
+        async with self._db() as session:
+            data = await session.stream_scalars(
+                select(mda).where(mda.span_rowid.in_(keys)).where(mda.annotator_kind == "LLM")
+            )
+            async for document_evaluation in data:
+                document_evaluations_by_id[document_evaluation.span_rowid].append(
+                    DocumentEvaluation.from_sql_document_annotation(document_evaluation)
+                )
+        return [document_evaluations_by_id[key] for key in keys]

phoenix/server/api/dataloaders/document_retrieval_metrics.py

@@ -0,0 +1,89 @@
+from collections import defaultdict
+from typing import Optional
+
+import numpy as np
+from aioitertools.itertools import groupby
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.metrics.retrieval_metrics import RetrievalMetrics
+from phoenix.server.api.types.DocumentRetrievalMetrics import DocumentRetrievalMetrics
+from phoenix.server.types import DbSessionFactory
+
+RowId: TypeAlias = int
+NumDocs: TypeAlias = int
+EvalName: TypeAlias = Optional[str]
+
+Key: TypeAlias = tuple[RowId, EvalName, NumDocs]
+Result: TypeAlias = list[DocumentRetrievalMetrics]
+
+
+class DocumentRetrievalMetricsDataLoader(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]:
+        mda = models.DocumentAnnotation
+        stmt = (
+            select(
+                mda.span_rowid,
+                mda.name,
+                mda.score,
+                mda.document_position,
+            )
+            .where(mda.score != None)  # noqa: E711
+            .where(mda.annotator_kind == "LLM")
+            .where(mda.document_position >= 0)
+            .order_by(mda.span_rowid, mda.name)
+        )
+        # Using CTE with VALUES clause is possible in SQLite, but not in
+        # SQLAlchemy v2.0.29, hence the workaround below with over-fetching.
+        # We could use CTE with VALUES for postgresql, but for now we'll keep
+        # it simple and just use one approach for all backends.
+        all_row_ids = {row_id for row_id, _, _ in keys}
+        stmt = stmt.where(mda.span_rowid.in_(all_row_ids))
+        all_eval_names = {eval_name for _, eval_name, _ in keys}
+        if None not in all_eval_names:
+            stmt = stmt.where(mda.name.in_(all_eval_names))
+        max_position = max(num_docs for _, _, num_docs in keys)
+        stmt = stmt.where(mda.document_position < max_position)
+        results: dict[Key, Result] = {key: [] for key in keys}
+        requested_num_docs: defaultdict[tuple[RowId, EvalName], set[NumDocs]] = defaultdict(set)
+        for row_id, eval_name, num_docs in results.keys():
+            requested_num_docs[(row_id, eval_name)].add(num_docs)
+        async with self._db() as session:
+            data = await session.stream(stmt)
+            async for (span_rowid, name), group in groupby(data, lambda r: (r.span_rowid, r.name)):
+                # We need to fulfill two types of potential requests: 1. when it
+                # specifies an evaluation name, and 2. when it doesn't care about
+                # the evaluation name by specifying None.
+                max_requested_num_docs = max(
+                    (
+                        num_docs
+                        for eval_name in (name, None)
+                        for num_docs in (requested_num_docs.get((span_rowid, eval_name)) or ())
+                    ),
+                    default=0,
+                )
+                if max_requested_num_docs <= 0:
+                    # We have over-fetched. Skip this group.
+                    continue
+                scores = [np.nan] * max_requested_num_docs
+                for row in group:
+                    # Length check is necessary due to over-fetching.
+                    if row.document_position < len(scores):
+                        scores[row.document_position] = row.score
+                for eval_name in (name, None):
+                    for num_docs in requested_num_docs.get((span_rowid, eval_name)) or ():
+                        metrics = RetrievalMetrics(scores[:num_docs])
+                        doc_metrics = DocumentRetrievalMetrics(
+                            evaluation_name=name, metrics=metrics
+                        )
+                        key = (span_rowid, eval_name, num_docs)
+                        results[key].append(doc_metrics)
+        # Make sure to copy the result, so we don't return the same list
+        # object to two different requesters.
+        return [results[key].copy() for key in keys]

phoenix/server/api/dataloaders/experiment_annotation_summaries.py

@@ -0,0 +1,79 @@
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Optional
+
+from sqlalchemy import func, select
+from strawberry.dataloader import AbstractCache, DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.server.types import DbSessionFactory
+
+
+@dataclass
+class ExperimentAnnotationSummary:
+    annotation_name: str
+    min_score: float
+    max_score: float
+    mean_score: float
+    count: int
+    error_count: int
+
+
+ExperimentID: TypeAlias = int
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = list[ExperimentAnnotationSummary]
+
+
+class ExperimentAnnotationSummaryDataLoader(DataLoader[Key, Result]):
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        cache_map: Optional[AbstractCache[Key, Result]] = None,
+    ) -> None:
+        super().__init__(load_fn=self._load_fn)
+        self._db = db
+
+    async def _load_fn(self, keys: list[Key]) -> list[Result]:
+        experiment_ids = keys
+        summaries: defaultdict[ExperimentID, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for (
+                experiment_id,
+                annotation_name,
+                min_score,
+                max_score,
+                mean_score,
+                count,
+                error_count,
+            ) in await session.stream(
+                select(
+                    models.ExperimentRun.experiment_id,
+                    models.ExperimentRunAnnotation.name,
+                    func.min(models.ExperimentRunAnnotation.score),
+                    func.max(models.ExperimentRunAnnotation.score),
+                    func.avg(models.ExperimentRunAnnotation.score),
+                    func.count(),
+                    func.count(models.ExperimentRunAnnotation.error),
+                )
+                .join(
+                    models.ExperimentRun,
+                    models.ExperimentRunAnnotation.experiment_run_id == models.ExperimentRun.id,
+                )
+                .where(models.ExperimentRun.experiment_id.in_(experiment_ids))
+                .group_by(models.ExperimentRun.experiment_id, models.ExperimentRunAnnotation.name)
+            ):
+                summaries[experiment_id].append(
+                    ExperimentAnnotationSummary(
+                        annotation_name=annotation_name,
+                        min_score=min_score,
+                        max_score=max_score,
+                        mean_score=mean_score,
+                        count=count,
+                        error_count=error_count,
+                    )
+                )
+        return [
+            sorted(summaries[experiment_id], key=lambda summary: summary.annotation_name)
+            for experiment_id in keys
+        ]

phoenix/server/api/dataloaders/experiment_error_rates.py

@@ -0,0 +1,58 @@
+from typing import Optional
+
+from sqlalchemy import case, 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
+ErrorRate: TypeAlias = float
+Key: TypeAlias = ExperimentID
+Result: TypeAlias = Optional[ErrorRate]
+
+
+class ExperimentErrorRatesDataLoader(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,
+                case(
+                    (
+                        func.count(models.ExperimentRun.id) != 0,
+                        func.count(models.ExperimentRun.error)
+                        / func.count(models.ExperimentRun.id),
+                    ),
+                    else_=None,
+                ),
+            )
+            .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:
+            error_rates = {
+                experiment_id: error_rate
+                async for experiment_id, error_rate in await session.stream(query)
+            }
+        return [
+            error_rates.get(experiment_id, ValueError(f"Unknown experiment ID: {experiment_id}"))
+            for experiment_id in keys
+        ]

phoenix/server/api/dataloaders/experiment_run_annotations.py

@@ -0,0 +1,36 @@
+from collections import defaultdict
+
+from sqlalchemy import select
+from strawberry.dataloader import DataLoader
+from typing_extensions import TypeAlias
+
+from phoenix.db.models import ExperimentRunAnnotation as OrmExperimentRunAnnotation
+from phoenix.server.types import DbSessionFactory
+
+ExperimentRunID: TypeAlias = int
+Key: TypeAlias = ExperimentRunID
+Result: TypeAlias = list[OrmExperimentRunAnnotation]
+
+
+class ExperimentRunAnnotations(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]:
+        run_ids = keys
+        annotations: defaultdict[Key, Result] = defaultdict(list)
+        async with self._db() as session:
+            async for run_id, annotation in await session.stream(
+                select(
+                    OrmExperimentRunAnnotation.experiment_run_id, OrmExperimentRunAnnotation
+                ).where(OrmExperimentRunAnnotation.experiment_run_id.in_(run_ids))
+            ):
+                annotations[run_id].append(annotation)
+        return [
+            sorted(annotations[run_id], key=lambda annotation: annotation.name, reverse=True)
+            for run_id in keys
+        ]