arize-phoenix 4.19.0__tar.gz → 4.20.1__tar.gz

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 4.19.0
+Version: 4.20.1
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -96,6 +96,7 @@ Requires-Dist: ruff==0.5.4; extra == 'dev'
 Requires-Dist: strawberry-graphql[debug-server,opentelemetry]==0.236.0; extra == 'dev'
 Requires-Dist: tabulate; extra == 'dev'
 Requires-Dist: types-tabulate; extra == 'dev'
+Requires-Dist: uvloop; (platform_system != 'Windows') and extra == 'dev'
 Provides-Extra: evals
 Provides-Extra: experimental
 Provides-Extra: llama-index

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/pyproject.toml

@@ -94,6 +94,7 @@ dev = [
   "prometheus_client",
   "asgi-lifespan",
   "Faker>=26.0.0",
+  "uvloop; platform_system != 'Windows'",
 ]
 evals = []
 experimental = []
@@ -173,6 +174,7 @@ dependencies = [
   "astunparse; python_version<'3.9'",  # `ast.unparse(...)` is only available starting with Python 3.9
   "asgi-lifespan",
   "Faker>=26.0.0",
+  "uvloop; platform_system != 'Windows'",
 ]
 
 [tool.hatch.envs.type]

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/db/bulk_inserter.py

@@ -1,38 +1,29 @@
 import asyncio
 import logging
 from asyncio import Queue, as_completed
-from collections import defaultdict
 from dataclasses import dataclass, field
-from datetime import datetime, timezone
 from functools import singledispatchmethod
 from itertools import islice
 from time import perf_counter
 from typing import (
     Any,
+    AsyncIterator,
     Awaitable,
     Callable,
-    DefaultDict,
-    Dict,
     Iterable,
     List,
-    Mapping,
     Optional,
     Set,
     Tuple,
-    Type,
     cast,
 )
 
-from cachetools import LRUCache
-from sqlalchemy import Select, select
 from typing_extensions import TypeAlias
 
 import phoenix.trace.v1 as pb
-from phoenix.db import models
 from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
 from phoenix.db.insertion.document_annotation import DocumentAnnotationQueueInserter
 from phoenix.db.insertion.evaluation import (
-    EvaluationInsertionEvent,
     InsertEvaluationError,
     insert_evaluation,
 )
@@ -41,8 +32,8 @@ from phoenix.db.insertion.span import SpanInsertionEvent, insert_span
 from phoenix.db.insertion.span_annotation import SpanAnnotationQueueInserter
 from phoenix.db.insertion.trace_annotation import TraceAnnotationQueueInserter
 from phoenix.db.insertion.types import Insertables, Precursors
-from phoenix.server.api.dataloaders import CacheForDataLoaders
-from phoenix.server.types import DbSessionFactory
+from phoenix.server.dml_event import DmlEvent, SpanInsertEvent
+from phoenix.server.types import CanPutItem, DbSessionFactory
 from phoenix.trace.schemas import Span
 
 logger = logging.getLogger(__name__)
@@ -60,8 +51,7 @@ class BulkInserter:
         self,
         db: DbSessionFactory,
         *,
-        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
-        initial_batch_of_operations: Iterable[DataManipulation] = (),
+        event_queue: CanPutItem[DmlEvent],
         initial_batch_of_spans: Optional[Iterable[Tuple[Span, str]]] = None,
         initial_batch_of_evaluations: Optional[Iterable[pb.Evaluation]] = None,
         sleep: float = 0.1,
@@ -93,18 +83,12 @@ class BulkInserter:
             [] if initial_batch_of_evaluations is None else list(initial_batch_of_evaluations)
         )
         self._task: Optional[asyncio.Task[None]] = None
-        self._last_updated_at_by_project: LRUCache[ProjectRowId, datetime] = LRUCache(maxsize=100)
-        self._cache_for_dataloaders = cache_for_dataloaders
+        self._event_queue = event_queue
         self._enable_prometheus = enable_prometheus
         self._retry_delay_sec = retry_delay_sec
         self._retry_allowance = retry_allowance
         self._queue_inserters = _QueueInserters(db, self._retry_delay_sec, self._retry_allowance)
 
-    def last_updated_at(self, project_rowid: Optional[ProjectRowId] = None) -> Optional[datetime]:
-        if isinstance(project_rowid, ProjectRowId):
-            return self._last_updated_at_by_project.get(project_rowid)
-        return max(self._last_updated_at_by_project.values(), default=None)
-
     async def __aenter__(
         self,
     ) -> Tuple[
@@ -162,21 +146,20 @@ class BulkInserter:
             ):
                 await asyncio.sleep(self._sleep)
                 continue
-            ops_remaining, events = self._max_ops_per_transaction, []
+            ops_remaining = self._max_ops_per_transaction
             async with self._db() as session:
                 while ops_remaining and not self._operations.empty():
                     ops_remaining -= 1
                     op = await self._operations.get()
                     try:
                         async with session.begin_nested():
-                            events.append(await op(session))
+                            await op(session)
                     except Exception as e:
                         if self._enable_prometheus:
                             from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
 
                             BULK_LOADER_EXCEPTIONS.inc()
                         logger.exception(str(e))
-            await self._process_events(events)
             # It's important to grab the buffers at the same time so there's
             # no race condition, since an eval insertion will fail if the span
             # it references doesn't exist. Grabbing the eval buffer later may
@@ -190,25 +173,18 @@ class BulkInserter:
                 self._evaluations = []
             # Spans should be inserted before the evaluations, since an evaluation
             # insertion will fail if the span it references doesn't exist.
-            transaction_result = TransactionResult()
             if spans_buffer:
-                result = await self._insert_spans(spans_buffer)
-                transaction_result.updated_project_rowids.update(result.updated_project_rowids)
+                await self._insert_spans(spans_buffer)
                 spans_buffer = None
             if evaluations_buffer:
-                result = await self._insert_evaluations(evaluations_buffer)
-                transaction_result.updated_project_rowids.update(result.updated_project_rowids)
+                await self._insert_evaluations(evaluations_buffer)
                 evaluations_buffer = None
-            for project_rowid in transaction_result.updated_project_rowids:
-                self._last_updated_at_by_project[project_rowid] = datetime.now(timezone.utc)
-            if not self._queue_inserters.empty:
-                if inserted_ids := await self._queue_inserters.insert():
-                    for project_rowid in await self._get_project_rowids(inserted_ids):
-                        self._last_updated_at_by_project[project_rowid] = datetime.now(timezone.utc)
+            async for event in self._queue_inserters.insert():
+                self._event_queue.put(event)
             await asyncio.sleep(self._sleep)
 
-    async def _insert_spans(self, spans: List[Tuple[Span, str]]) -> TransactionResult:
-        transaction_result = TransactionResult()
+    async def _insert_spans(self, spans: List[Tuple[Span, str]]) -> None:
+        project_ids = set()
         for i in range(0, len(spans), self._max_ops_per_transaction):
             try:
                 start = perf_counter()
@@ -231,9 +207,7 @@ class BulkInserter:
                                 f"Failed to insert span with span_id={span.context.span_id}"
                             )
                         if result is not None:
-                            transaction_result.updated_project_rowids.add(result.project_rowid)
-                            if (cache := self._cache_for_dataloaders) is not None:
-                                cache.invalidate(result)
+                            project_ids.add(result.project_rowid)
                 if self._enable_prometheus:
                     from phoenix.server.prometheus import BULK_LOADER_INSERTION_TIME
 
@@ -244,10 +218,9 @@ class BulkInserter:
 
                     BULK_LOADER_EXCEPTIONS.inc()
                 logger.exception("Failed to insert spans")
-        return transaction_result
+        self._event_queue.put(SpanInsertEvent(tuple(project_ids)))
 
-    async def _insert_evaluations(self, evaluations: List[pb.Evaluation]) -> TransactionResult:
-        transaction_result = TransactionResult()
+    async def _insert_evaluations(self, evaluations: List[pb.Evaluation]) -> None:
         for i in range(0, len(evaluations), self._max_ops_per_transaction):
             try:
                 start = perf_counter()
@@ -257,20 +230,15 @@ class BulkInserter:
                             from phoenix.server.prometheus import BULK_LOADER_EVALUATION_INSERTIONS
 
                             BULK_LOADER_EVALUATION_INSERTIONS.inc()
-                        result: Optional[EvaluationInsertionEvent] = None
                         try:
                             async with session.begin_nested():
-                                result = await insert_evaluation(session, evaluation)
+                                await insert_evaluation(session, evaluation)
                         except InsertEvaluationError as error:
                             if self._enable_prometheus:
                                 from phoenix.server.prometheus import BULK_LOADER_EXCEPTIONS
 
                                 BULK_LOADER_EXCEPTIONS.inc()
                             logger.exception(f"Failed to insert evaluation: {str(error)}")
-                        if result is not None:
-                            transaction_result.updated_project_rowids.add(result.project_rowid)
-                            if (cache := self._cache_for_dataloaders) is not None:
-                                cache.invalidate(result)
                 if self._enable_prometheus:
                     from phoenix.server.prometheus import BULK_LOADER_INSERTION_TIME
 
@@ -281,48 +249,6 @@ class BulkInserter:
 
                     BULK_LOADER_EXCEPTIONS.inc()
                 logger.exception("Failed to insert evaluations")
-        return transaction_result
-
-    async def _get_project_rowids(
-        self,
-        inserted_ids: Mapping[Type[models.Base], List[int]],
-    ) -> Set[int]:
-        ans: Set[int] = set()
-        if not inserted_ids:
-            return ans
-        stmt: Select[Tuple[int]]
-        for table, ids in inserted_ids.items():
-            if not ids:
-                continue
-            if issubclass(table, models.SpanAnnotation):
-                stmt = (
-                    select(models.Project.id)
-                    .join(models.Trace)
-                    .join_from(models.Trace, models.Span)
-                    .join_from(models.Span, models.SpanAnnotation)
-                    .where(models.SpanAnnotation.id.in_(ids))
-                )
-            elif issubclass(table, models.DocumentAnnotation):
-                stmt = (
-                    select(models.Project.id)
-                    .join(models.Trace)
-                    .join_from(models.Trace, models.Span)
-                    .join_from(models.Span, models.DocumentAnnotation)
-                    .where(models.DocumentAnnotation.id.in_(ids))
-                )
-            elif issubclass(table, models.TraceAnnotation):
-                stmt = (
-                    select(models.Project.id)
-                    .join(models.Trace)
-                    .join_from(models.Trace, models.TraceAnnotation)
-                    .where(models.TraceAnnotation.id.in_(ids))
-                )
-            else:
-                continue
-            async with self._db() as session:
-                project_rowids = [_ async for _ in await session.stream_scalars(stmt)]
-                ans.update(project_rowids)
-        return ans
 
 
 class _QueueInserters:
@@ -343,13 +269,13 @@ class _QueueInserters:
             self._document_annotations,
         )
 
-    async def insert(self) -> Dict[Type[models.Base], List[int]]:
-        ans: DefaultDict[Type[models.Base], List[int]] = defaultdict(list)
-        for coro in as_completed([q.insert() for q in self._queues]):
-            table, inserted_ids = await coro
-            if inserted_ids:
-                ans[table].extend(inserted_ids)
-        return ans
+    async def insert(self) -> AsyncIterator[DmlEvent]:
+        if self.empty:
+            return
+        for coro in as_completed([q.insert() for q in self._queues if not q.empty]):
+            if events := cast(Optional[List[DmlEvent]], await coro):
+                for event in events:
+                    yield event
 
     @property
     def empty(self) -> bool:

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/db/insertion/document_annotation.py

@@ -7,6 +7,7 @@ from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.db.helpers import dedup, num_docs_col
+from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import (
     Insertables,
     Postponed,
@@ -14,6 +15,7 @@ from phoenix.db.insertion.types import (
     QueueInserter,
     Received,
 )
+from phoenix.server.dml_event import DocumentAnnotationDmlEvent
 
 _Name: TypeAlias = str
 _SpanId: TypeAlias = str
@@ -40,10 +42,21 @@ class DocumentAnnotationQueueInserter(
         Precursors.DocumentAnnotation,
         Insertables.DocumentAnnotation,
         models.DocumentAnnotation,
+        DocumentAnnotationDmlEvent,
     ],
     table=models.DocumentAnnotation,
     unique_by=("name", "span_rowid", "document_position"),
 ):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.DocumentAnnotation,
+    ) -> List[DocumentAnnotationDmlEvent]:
+        records = [dict(as_kv(ins.row)) for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [DocumentAnnotationDmlEvent(ids)]
+
     async def _partition(
         self,
         session: AsyncSession,

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/db/insertion/span_annotation.py

@@ -7,6 +7,7 @@ from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.db.helpers import dedup
+from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import (
     Insertables,
     Postponed,
@@ -14,6 +15,7 @@ from phoenix.db.insertion.types import (
     QueueInserter,
     Received,
 )
+from phoenix.server.dml_event import SpanAnnotationDmlEvent
 
 _Name: TypeAlias = str
 _SpanId: TypeAlias = str
@@ -36,10 +38,21 @@ class SpanAnnotationQueueInserter(
         Precursors.SpanAnnotation,
         Insertables.SpanAnnotation,
         models.SpanAnnotation,
+        SpanAnnotationDmlEvent,
     ],
     table=models.SpanAnnotation,
     unique_by=("name", "span_rowid"),
 ):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.SpanAnnotation,
+    ) -> List[SpanAnnotationDmlEvent]:
+        records = [dict(as_kv(ins.row)) for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [SpanAnnotationDmlEvent(ids)]
+
     async def _partition(
         self,
         session: AsyncSession,

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/db/insertion/trace_annotation.py

@@ -7,6 +7,7 @@ from typing_extensions import TypeAlias
 
 from phoenix.db import models
 from phoenix.db.helpers import dedup
+from phoenix.db.insertion.helpers import as_kv
 from phoenix.db.insertion.types import (
     Insertables,
     Postponed,
@@ -14,6 +15,7 @@ from phoenix.db.insertion.types import (
     QueueInserter,
     Received,
 )
+from phoenix.server.dml_event import TraceAnnotationDmlEvent
 
 _Name: TypeAlias = str
 _TraceId: TypeAlias = str
@@ -36,10 +38,21 @@ class TraceAnnotationQueueInserter(
         Precursors.TraceAnnotation,
         Insertables.TraceAnnotation,
         models.TraceAnnotation,
+        TraceAnnotationDmlEvent,
     ],
     table=models.TraceAnnotation,
     unique_by=("name", "trace_rowid"),
 ):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.TraceAnnotation,
+    ) -> List[TraceAnnotationDmlEvent]:
+        records = [dict(as_kv(ins.row)) for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [TraceAnnotationDmlEvent(ids)]
+
     async def _partition(
         self,
         session: AsyncSession,

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/db/insertion/types.py

@@ -21,11 +21,12 @@ from typing import (
 )
 
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.sql.dml import ReturningInsert
+from sqlalchemy.sql.dml import Insert
 
 from phoenix.db import models
 from phoenix.db.insertion.constants import DEFAULT_RETRY_ALLOWANCE, DEFAULT_RETRY_DELAY_SEC
-from phoenix.db.insertion.helpers import as_kv, insert_on_conflict
+from phoenix.db.insertion.helpers import insert_on_conflict
+from phoenix.server.dml_event import DmlEvent
 from phoenix.server.types import DbSessionFactory
 
 logger = logging.getLogger("__name__")
@@ -40,6 +41,7 @@ _AnyT = TypeVar("_AnyT")
 _PrecursorT = TypeVar("_PrecursorT")
 _InsertableT = TypeVar("_InsertableT", bound=Insertable)
 _RowT = TypeVar("_RowT", bound=models.Base)
+_DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)
 
 
 @dataclass(frozen=True)
@@ -56,7 +58,7 @@ class Postponed(Received[_AnyT]):
     retries_left: int = field(default=DEFAULT_RETRY_ALLOWANCE)
 
 
-class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
+class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
     table: Type[_RowT]
     unique_by: Sequence[str]
 
@@ -97,59 +99,63 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
         List[Received[_PrecursorT]],
     ]: ...
 
-    async def insert(self) -> Tuple[Type[_RowT], List[int]]:
+    async def insert(self) -> Optional[List[_DmlEventT]]:
         if not self._queue:
-            return self.table, []
-        parcels = self._queue
-        self._queue = []
-        inserted_ids: List[int] = []
+            return None
+        self._queue, parcels = [], self._queue
+        events: List[_DmlEventT] = []
         async with self._db() as session:
             to_insert, to_postpone, _ = await self._partition(session, *parcels)
             if to_insert:
-                inserted_ids, to_retry, _ = await self._insert(session, *to_insert)
-                to_postpone.extend(to_retry)
+                events, to_retry, _ = await self._insert(session, *to_insert)
+                if to_retry:
+                    to_postpone.extend(to_retry)
         if to_postpone:
             loop = asyncio.get_running_loop()
             loop.call_later(self._retry_delay_sec, self._queue.extend, to_postpone)
-        return self.table, inserted_ids
+        return events
 
-    def _stmt(self, *records: Mapping[str, Any]) -> ReturningInsert[Tuple[int]]:
-        pk = next(c for c in self.table.__table__.c if c.primary_key)
+    def _insert_on_conflict(self, *records: Mapping[str, Any]) -> Insert:
         return insert_on_conflict(
             *records,
             table=self.table,
             unique_by=self.unique_by,
             dialect=self._db.dialect,
-        ).returning(pk)
+        )
+
+    @abstractmethod
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: _InsertableT,
+    ) -> List[_DmlEventT]: ...
 
     async def _insert(
         self,
         session: AsyncSession,
-        *insertions: Received[_InsertableT],
-    ) -> Tuple[List[int], List[Postponed[_PrecursorT]], List[Received[_InsertableT]]]:
-        records = [dict(as_kv(ins.item.row)) for ins in insertions]
-        inserted_ids: List[int] = []
+        *parcels: Received[_InsertableT],
+    ) -> Tuple[
+        List[_DmlEventT],
+        List[Postponed[_PrecursorT]],
+        List[Received[_InsertableT]],
+    ]:
         to_retry: List[Postponed[_PrecursorT]] = []
         failures: List[Received[_InsertableT]] = []
-        stmt = self._stmt(*records)
+        events: List[_DmlEventT] = []
         try:
             async with session.begin_nested():
-                ids = [id_ async for id_ in await session.stream_scalars(stmt)]
-                inserted_ids.extend(ids)
+                events.extend(await self._events(session, *(p.item for p in parcels)))
         except BaseException:
             logger.exception(
                 f"Failed to bulk insert for {self.table.__name__}. "
-                f"Will try to insert ({len(records)} records) individually instead."
+                f"Will try to insert ({len(parcels)} records) individually instead."
             )
-            for i, record in enumerate(records):
-                stmt = self._stmt(record)
+            for p in parcels:
                 try:
                     async with session.begin_nested():
-                        ids = [id_ async for id_ in await session.stream_scalars(stmt)]
-                        inserted_ids.extend(ids)
+                        events.extend(await self._events(session, p.item))
                 except BaseException:
                     logger.exception(f"Failed to insert for {self.table.__name__}.")
-                    p = insertions[i]
                     if isinstance(p, Postponed) and p.retries_left == 1:
                         failures.append(p)
                     else:
@@ -162,7 +168,7 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
                                 else self._retry_allowance,
                             )
                         )
-        return inserted_ids, to_retry, failures
+        return events, to_retry, failures
 
 
 class Precursors(ABC):

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/server/api/context.py

@@ -1,10 +1,8 @@
 from dataclasses import dataclass
-from datetime import datetime
 from pathlib import Path
-from typing import Callable, Optional
+from typing import Any, Optional
 
 from strawberry.fastapi import BaseContext
-from typing_extensions import TypeAlias
 
 from phoenix.core.model_schema import Model
 from phoenix.server.api.dataloaders import (
@@ -34,7 +32,8 @@ from phoenix.server.api.dataloaders import (
     TraceEvaluationsDataLoader,
     TraceRowIdsDataLoader,
 )
-from phoenix.server.types import DbSessionFactory
+from phoenix.server.dml_event import DmlEvent
+from phoenix.server.types import CanGetLastUpdatedAt, CanPutItem, DbSessionFactory
 
 
 @dataclass
@@ -54,6 +53,7 @@ class DataLoaders:
     latency_ms_quantile: LatencyMsQuantileDataLoader
     min_start_or_max_end_times: MinStartOrMaxEndTimeDataLoader
     record_counts: RecordCountDataLoader
+    span_annotations: SpanAnnotationsDataLoader
     span_dataset_examples: SpanDatasetExamplesDataLoader
     span_descendants: SpanDescendantsDataLoader
     span_evaluations: SpanEvaluationsDataLoader
@@ -62,10 +62,11 @@ class DataLoaders:
     trace_evaluations: TraceEvaluationsDataLoader
     trace_row_ids: TraceRowIdsDataLoader
     project_by_name: ProjectByNameDataLoader
-    span_annotations: SpanAnnotationsDataLoader
 
 
-ProjectRowId: TypeAlias = int
+class _NoOp:
+    def get(self, *args: Any, **kwargs: Any) -> Any: ...
+    def put(self, *args: Any, **kwargs: Any) -> Any: ...
 
 
 @dataclass
@@ -75,6 +76,7 @@ class Context(BaseContext):
     cache_for_dataloaders: Optional[CacheForDataLoaders]
     model: Model
     export_path: Path
+    last_updated_at: CanGetLastUpdatedAt = _NoOp()
+    event_queue: CanPutItem[DmlEvent] = _NoOp()
     corpus: Optional[Model] = None
-    streaming_last_updated_at: Callable[[ProjectRowId], Optional[datetime]] = lambda _: None
     read_only: bool = False

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/server/api/dataloaders/__init__.py

@@ -1,12 +1,4 @@
 from dataclasses import dataclass, field
-from functools import singledispatchmethod
-
-from phoenix.db.insertion.evaluation import (
-    DocumentEvaluationInsertionEvent,
-    SpanEvaluationInsertionEvent,
-    TraceEvaluationInsertionEvent,
-)
-from phoenix.db.insertion.span import ClearProjectSpansEvent, SpanInsertionEvent
 
 from .annotation_summaries import AnnotationSummaryCache, AnnotationSummaryDataLoader
 from .average_experiment_run_latency import AverageExperimentRunLatencyDataLoader
@@ -88,42 +80,3 @@ class CacheForDataLoaders:
     token_count: TokenCountCache = field(
         default_factory=TokenCountCache,
     )
-
-    def _update_spans(self, project_rowid: int) -> None:
-        self.latency_ms_quantile.invalidate(project_rowid)
-        self.token_count.invalidate(project_rowid)
-        self.record_count.invalidate(project_rowid)
-        self.min_start_or_max_end_time.invalidate(project_rowid)
-
-    def _clear_spans(self, project_rowid: int) -> None:
-        self._update_spans(project_rowid)
-        self.annotation_summary.invalidate_project(project_rowid)
-        self.evaluation_summary.invalidate_project(project_rowid)
-        self.document_evaluation_summary.invalidate_project(project_rowid)
-
-    @singledispatchmethod
-    def invalidate(self, event: SpanInsertionEvent) -> None:
-        project_rowid, *_ = event
-        self._update_spans(project_rowid)
-
-    @invalidate.register
-    def _(self, event: ClearProjectSpansEvent) -> None:
-        project_rowid, *_ = event
-        self._clear_spans(project_rowid)
-
-    @invalidate.register
-    def _(self, event: DocumentEvaluationInsertionEvent) -> None:
-        project_rowid, evaluation_name = event
-        self.document_evaluation_summary.invalidate((project_rowid, evaluation_name))
-
-    @invalidate.register
-    def _(self, event: SpanEvaluationInsertionEvent) -> None:
-        project_rowid, evaluation_name = event
-        self.annotation_summary.invalidate((project_rowid, evaluation_name, "span"))
-        self.evaluation_summary.invalidate((project_rowid, evaluation_name, "span"))
-
-    @invalidate.register
-    def _(self, event: TraceEvaluationInsertionEvent) -> None:
-        project_rowid, evaluation_name = event
-        self.annotation_summary.invalidate((project_rowid, evaluation_name, "trace"))
-        self.evaluation_summary.invalidate((project_rowid, evaluation_name, "trace"))

{arize_phoenix-4.19.0 → arize_phoenix-4.20.1}/src/phoenix/server/api/dataloaders/span_annotations.py

@@ -8,12 +8,11 @@ from sqlalchemy import select
 from strawberry.dataloader import DataLoader
 from typing_extensions import TypeAlias
 
-from phoenix.db import models
-from phoenix.server.api.types.SpanAnnotation import SpanAnnotation, to_gql_span_annotation
+from phoenix.db.models import SpanAnnotation as ORMSpanAnnotation
 from phoenix.server.types import DbSessionFactory
 
 Key: TypeAlias = int
-Result: TypeAlias = List[SpanAnnotation]
+Result: TypeAlias = List[ORMSpanAnnotation]
 
 
 class SpanAnnotationsDataLoader(DataLoader[Key, Result]):
@@ -23,11 +22,9 @@ class SpanAnnotationsDataLoader(DataLoader[Key, Result]):
 
     async def _load_fn(self, keys: List[Key]) -> List[Result]:
         span_annotations_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)))
-            async for span_annotation in data:
-                span_annotations_by_id[span_annotation.span_rowid].append(
-                    to_gql_span_annotation(span_annotation)
-                )
+            async for span_annotation in await session.stream_scalars(
+                select(ORMSpanAnnotation).where(ORMSpanAnnotation.span_rowid.in_(keys))
+            ):
+                span_annotations_by_id[span_annotation.span_rowid].append(span_annotation)
         return [span_annotations_by_id[key] for key in keys]