arize-phoenix 4.20.0__py3-none-any.whl → 4.20.1__py3-none-any.whl

phoenix/server/api/types/Experiment.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Optional
+from typing import ClassVar, List, Optional, Type
 
 import strawberry
 from sqlalchemy import select
@@ -23,6 +23,7 @@ from phoenix.server.api.types.Project import Project
 
 @strawberry.type
 class Experiment(Node):
+    _table: ClassVar[Type[models.Base]] = models.Experiment
     cached_sequence_number: Private[Optional[int]] = None
     id_attr: NodeID[int]
     name: str
@@ -127,6 +128,10 @@ class Experiment(Node):
             gradient_end_color=db_project.gradient_end_color,
         )
 
+    @strawberry.field
+    def last_updated_at(self, info: Info[Context, None]) -> Optional[datetime]:
+        return info.context.last_updated_at.get(self._table, self.id_attr)
+
 
 def to_gql_experiment(
     experiment: models.Experiment,

phoenix/server/api/types/Project.py

@@ -2,8 +2,10 @@ import operator
 from datetime import datetime
 from typing import (
     Any,
+    ClassVar,
     List,
     Optional,
+    Type,
 )
 
 import strawberry
@@ -38,6 +40,7 @@ from phoenix.trace.dsl import SpanFilter
 
 @strawberry.type
 class Project(Node):
+    _table: ClassVar[Type[models.Base]] = models.Project
     id_attr: NodeID[int]
     name: str
     gradient_start_color: str
@@ -397,7 +400,7 @@ class Project(Node):
         self,
         info: Info[Context, None],
     ) -> Optional[datetime]:
-        return info.context.streaming_last_updated_at(self.id_attr)
+        return info.context.last_updated_at.get(self._table, self.id_attr)
 
     @strawberry.field
     async def validate_span_filter_condition(self, condition: str) -> ValidationResult:

phoenix/server/api/types/Span.py

@@ -194,8 +194,8 @@ class Span(Node):
     ) -> List[SpanAnnotation]:
         span_id = self.id_attr
         annotations = await info.context.data_loaders.span_annotations.load(span_id)
-        sort_key = SpanAnnotationColumn.createdAt.value
-        sort_descending = True
+        sort_key = SpanAnnotationColumn.name.value
+        sort_descending = False
         if sort:
             sort_key = sort.col.value
             sort_descending = sort.dir is SortDir.desc

phoenix/server/app.py

@@ -2,7 +2,6 @@ import asyncio
 import contextlib
 import json
 import logging
-from datetime import datetime
 from functools import cached_property
 from pathlib import Path
 from typing import (
@@ -87,9 +86,16 @@ from phoenix.server.api.dataloaders import (
 from phoenix.server.api.routers.v1 import REST_API_VERSION
 from phoenix.server.api.routers.v1 import router as v1_router
 from phoenix.server.api.schema import schema
+from phoenix.server.dml_event import DmlEvent
+from phoenix.server.dml_event_handler import DmlEventHandler
 from phoenix.server.grpc_server import GrpcServer
 from phoenix.server.telemetry import initialize_opentelemetry_tracer_provider
-from phoenix.server.types import DbSessionFactory
+from phoenix.server.types import (
+    CanGetLastUpdatedAt,
+    CanPutItem,
+    DbSessionFactory,
+    LastUpdatedAt,
+)
 from phoenix.trace.schemas import Span
 from phoenix.utilities.client import PHOENIX_SERVER_VERSION_HEADER
 
@@ -220,6 +226,7 @@ def _lifespan(
     *,
     dialect: SupportedSQLDialect,
     bulk_inserter: BulkInserter,
+    dml_event_handler: DmlEventHandler,
     tracer_provider: Optional["TracerProvider"] = None,
     enable_prometheus: bool = False,
     clean_ups: Iterable[Callable[[], None]] = (),
@@ -239,8 +246,9 @@ def _lifespan(
             disabled=read_only,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
-        ):
+        ), dml_event_handler:
             yield {
+                "event_queue": dml_event_handler,
                 "enqueue": enqueue,
                 "queue_span_for_bulk_insert": queue_span,
                 "queue_evaluation_for_bulk_insert": queue_evaluation,
@@ -263,9 +271,10 @@ def create_graphql_router(
     db: DbSessionFactory,
     model: Model,
     export_path: Path,
+    last_updated_at: CanGetLastUpdatedAt,
     corpus: Optional[Model] = None,
-    streaming_last_updated_at: Callable[[ProjectRowId], Optional[datetime]] = lambda _: None,
     cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+    event_queue: CanPutItem[DmlEvent],
     read_only: bool = False,
 ) -> GraphQLRouter:  # type: ignore[type-arg]
     def get_context() -> Context:
@@ -274,7 +283,8 @@ def create_graphql_router(
             model=model,
             corpus=corpus,
             export_path=export_path,
-            streaming_last_updated_at=streaming_last_updated_at,
+            last_updated_at=last_updated_at,
+            event_queue=event_queue,
             data_loaders=DataLoaders(
                 average_experiment_run_latency=AverageExperimentRunLatencyDataLoader(db),
                 dataset_example_revisions=DatasetExampleRevisionsDataLoader(db),
@@ -420,11 +430,16 @@ def create_app(
     cache_for_dataloaders = (
         CacheForDataLoaders() if db.dialect is SupportedSQLDialect.SQLITE else None
     )
-
+    last_updated_at = LastUpdatedAt()
+    dml_event_handler = DmlEventHandler(
+        db=db,
+        cache_for_dataloaders=cache_for_dataloaders,
+        last_updated_at=last_updated_at,
+    )
     bulk_inserter = BulkInserter(
         db,
         enable_prometheus=enable_prometheus,
-        cache_for_dataloaders=cache_for_dataloaders,
+        event_queue=dml_event_handler,
         initial_batch_of_spans=initial_batch_of_spans,
         initial_batch_of_evaluations=initial_batch_of_evaluations,
     )
@@ -460,7 +475,8 @@ def create_app(
         model=model,
         corpus=corpus,
         export_path=export_path,
-        streaming_last_updated_at=bulk_inserter.last_updated_at,
+        last_updated_at=last_updated_at,
+        event_queue=dml_event_handler,
         cache_for_dataloaders=cache_for_dataloaders,
         read_only=read_only,
     )
@@ -477,6 +493,7 @@ def create_app(
             dialect=db.dialect,
             read_only=read_only,
             bulk_inserter=bulk_inserter,
+            dml_event_handler=dml_event_handler,
             tracer_provider=tracer_provider,
             enable_prometheus=enable_prometheus,
             clean_ups=clean_ups,

phoenix/server/dml_event.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from abc import ABC
+from dataclasses import dataclass, field
+from typing import ClassVar, Tuple, Type
+
+from phoenix.db import models
+
+
+@dataclass(frozen=True)
+class DmlEvent(ABC):
+    """
+    Event corresponding to a Data Manipulation Language (DML)
+    operation, e.g. insertion, update, or deletion.
+    """
+
+    table: ClassVar[Type[models.Base]]
+    ids: Tuple[int, ...] = field(default_factory=tuple)
+
+    def __bool__(self) -> bool:
+        return bool(self.ids)
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+@dataclass(frozen=True)
+class ProjectDmlEvent(DmlEvent):
+    table = models.Project
+
+
+@dataclass(frozen=True)
+class ProjectDeleteEvent(ProjectDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanDmlEvent(ProjectDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanInsertEvent(SpanDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanDeleteEvent(SpanDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DatasetDmlEvent(DmlEvent):
+    table = models.Dataset
+
+
+@dataclass(frozen=True)
+class DatasetInsertEvent(DatasetDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DatasetDeleteEvent(DatasetDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentDmlEvent(DmlEvent):
+    table = models.Experiment
+
+
+@dataclass(frozen=True)
+class ExperimentInsertEvent(ExperimentDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentDeleteEvent(ExperimentDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunDmlEvent(DmlEvent):
+    table = models.ExperimentRun
+
+
+@dataclass(frozen=True)
+class ExperimentRunInsertEvent(ExperimentRunDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunDeleteEvent(ExperimentRunDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationDmlEvent(DmlEvent):
+    table = models.ExperimentRunAnnotation
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationInsertEvent(ExperimentRunAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class ExperimentRunAnnotationDeleteEvent(ExperimentRunAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanAnnotationDmlEvent(DmlEvent):
+    table = models.SpanAnnotation
+
+
+@dataclass(frozen=True)
+class SpanAnnotationInsertEvent(SpanAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class SpanAnnotationDeleteEvent(SpanAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class TraceAnnotationDmlEvent(DmlEvent):
+    table = models.TraceAnnotation
+
+
+@dataclass(frozen=True)
+class TraceAnnotationInsertEvent(TraceAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class TraceAnnotationDeleteEvent(TraceAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationDmlEvent(DmlEvent):
+    table = models.DocumentAnnotation
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationInsertEvent(DocumentAnnotationDmlEvent): ...
+
+
+@dataclass(frozen=True)
+class DocumentAnnotationDeleteEvent(DocumentAnnotationDmlEvent): ...

phoenix/server/dml_event_handler.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from asyncio import gather
+from inspect import getmro
+from itertools import chain
+from typing import (
+    Any,
+    Callable,
+    Generic,
+    Iterable,
+    Iterator,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    TypedDict,
+    TypeVar,
+    Union,
+    cast,
+)
+
+from sqlalchemy import Select, select
+from typing_extensions import TypeAlias, Unpack
+
+from phoenix.db.models import (
+    Base,
+    DocumentAnnotation,
+    Project,
+    Span,
+    SpanAnnotation,
+    Trace,
+    TraceAnnotation,
+)
+from phoenix.server.api.dataloaders import CacheForDataLoaders
+from phoenix.server.dml_event import (
+    DmlEvent,
+    DocumentAnnotationDmlEvent,
+    SpanAnnotationDmlEvent,
+    SpanDeleteEvent,
+    SpanDmlEvent,
+    TraceAnnotationDmlEvent,
+)
+from phoenix.server.types import (
+    BatchedCaller,
+    CanSetLastUpdatedAt,
+    DbSessionFactory,
+)
+
+_DmlEventT = TypeVar("_DmlEventT", bound=DmlEvent)
+
+
+class _DmlEventQueue(Generic[_DmlEventT]):
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._events: Set[_DmlEventT] = set()
+
+    @property
+    def empty(self) -> bool:
+        return not self._events
+
+    def put(self, event: _DmlEventT) -> None:
+        self._events.add(event)
+
+    def clear(self) -> None:
+        self._events.clear()
+
+    def __iter__(self) -> Iterator[_DmlEventT]:
+        yield from self._events
+
+
+class _HandlerParams(TypedDict):
+    db: DbSessionFactory
+    last_updated_at: CanSetLastUpdatedAt
+    cache_for_dataloaders: Optional[CacheForDataLoaders]
+    sleep_seconds: float
+
+
+class _HasLastUpdatedAt(ABC):
+    def __init__(
+        self,
+        last_updated_at: CanSetLastUpdatedAt,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._last_updated_at = last_updated_at
+
+
+class _HasCacheForDataLoaders(ABC):
+    def __init__(
+        self,
+        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._cache_for_dataloaders = cache_for_dataloaders
+
+
+class _DmlEventHandler(
+    _HasLastUpdatedAt,
+    _HasCacheForDataLoaders,
+    BatchedCaller[_DmlEventT],
+    Generic[_DmlEventT],
+    ABC,
+):
+    _batch_factory = cast(Callable[[], _DmlEventQueue[_DmlEventT]], _DmlEventQueue)
+
+    def __init__(self, *, db: DbSessionFactory, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._db = db
+
+    def __hash__(self) -> int:
+        return id(self)
+
+
+class _GenericDmlEventHandler(_DmlEventHandler[DmlEvent]):
+    async def __call__(self) -> None:
+        for e in self._batch:
+            for id_ in e.ids:
+                self._update(e.table, id_)
+
+    def _update(self, table: Type[Base], id_: int) -> None:
+        self._last_updated_at.set(table, id_)
+
+
+class _SpanDmlEventHandler(_DmlEventHandler[SpanDmlEvent]):
+    async def __call__(self) -> None:
+        if cache := self._cache_for_dataloaders:
+            for id_ in set(chain.from_iterable(e.ids for e in self._batch)):
+                self._clear(cache, id_)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int) -> None:
+        cache.latency_ms_quantile.invalidate(project_id)
+        cache.token_count.invalidate(project_id)
+        cache.record_count.invalidate(project_id)
+        cache.min_start_or_max_end_time.invalidate(project_id)
+
+
+class _SpanDeleteEventHandler(_SpanDmlEventHandler):
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int) -> None:
+        cache.annotation_summary.invalidate_project(project_id)
+        cache.evaluation_summary.invalidate_project(project_id)
+        cache.document_evaluation_summary.invalidate_project(project_id)
+
+
+_AnnotationTable: TypeAlias = Union[
+    Type[SpanAnnotation],
+    Type[TraceAnnotation],
+    Type[DocumentAnnotation],
+]
+
+_AnnotationDmlEventT = TypeVar(
+    "_AnnotationDmlEventT",
+    SpanAnnotationDmlEvent,
+    TraceAnnotationDmlEvent,
+    DocumentAnnotationDmlEvent,
+)
+
+
+class _AnnotationDmlEventHandler(
+    _DmlEventHandler[_AnnotationDmlEventT],
+    Generic[_AnnotationDmlEventT],
+    ABC,
+):
+    _table: _AnnotationTable
+    _base_stmt: Union[Select[Tuple[int, str]], Select[Tuple[int]]] = (
+        select(Project.id).join_from(Project, Trace).distinct()
+    )
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._base_stmt
+        if self._cache_for_dataloaders:
+            self._stmt = self._stmt.add_columns(self._table.name)
+
+    def _get_stmt(self) -> Union[Select[Tuple[int, str]], Select[Tuple[int]]]:
+        ids = set(chain.from_iterable(e.ids for e in self._batch))
+        return self._stmt.where(self._table.id.in_(ids))
+
+    @staticmethod
+    @abstractmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None: ...
+
+    async def __call__(self) -> None:
+        async with self._db() as session:
+            async for row in await session.stream(self._get_stmt()):
+                if cache := self._cache_for_dataloaders:
+                    self._clear(cache, row.id, row.name)
+
+
+class _SpanAnnotationDmlEventHandler(_AnnotationDmlEventHandler[SpanAnnotationDmlEvent]):
+    _table = SpanAnnotation
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._stmt.join_from(Trace, Span).join_from(Span, self._table)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None:
+        cache.annotation_summary.invalidate((project_id, name, "span"))
+        cache.evaluation_summary.invalidate((project_id, name, "span"))
+
+
+class _TraceAnnotationDmlEventHandler(_AnnotationDmlEventHandler[TraceAnnotationDmlEvent]):
+    _table = TraceAnnotation
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._stmt.join_from(Trace, self._table)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None:
+        cache.annotation_summary.invalidate((project_id, name, "trace"))
+        cache.evaluation_summary.invalidate((project_id, name, "trace"))
+
+
+class _DocumentAnnotationDmlEventHandler(_AnnotationDmlEventHandler[DocumentAnnotationDmlEvent]):
+    _table = DocumentAnnotation
+
+    def __init__(self, **kwargs: Unpack[_HandlerParams]) -> None:
+        super().__init__(**kwargs)
+        self._stmt = self._stmt.join_from(Trace, Span).join_from(Span, self._table)
+
+    @staticmethod
+    def _clear(cache: CacheForDataLoaders, project_id: int, name: str) -> None:
+        cache.document_evaluation_summary.invalidate((project_id, name))
+
+
+class DmlEventHandler:
+    def __init__(
+        self,
+        *,
+        db: DbSessionFactory,
+        last_updated_at: CanSetLastUpdatedAt,
+        cache_for_dataloaders: Optional[CacheForDataLoaders] = None,
+        sleep_seconds: float = 0.1,
+    ) -> None:
+        kwargs = _HandlerParams(
+            db=db,
+            last_updated_at=last_updated_at,
+            cache_for_dataloaders=cache_for_dataloaders,
+            sleep_seconds=sleep_seconds,
+        )
+        self._handlers: Mapping[Type[DmlEvent], Iterable[_DmlEventHandler[Any]]] = {
+            DmlEvent: [_GenericDmlEventHandler(**kwargs)],
+            SpanDmlEvent: [_SpanDmlEventHandler(**kwargs)],
+            SpanDeleteEvent: [_SpanDeleteEventHandler(**kwargs)],
+            SpanAnnotationDmlEvent: [_SpanAnnotationDmlEventHandler(**kwargs)],
+            TraceAnnotationDmlEvent: [_TraceAnnotationDmlEventHandler(**kwargs)],
+            DocumentAnnotationDmlEvent: [_DocumentAnnotationDmlEventHandler(**kwargs)],
+        }
+        self._all_handlers = frozenset(chain.from_iterable(self._handlers.values()))
+
+    async def __aenter__(self) -> None:
+        await gather(*(h.start() for h in self._all_handlers))
+
+    async def __aexit__(self, *args: Any, **kwargs: Any) -> None:
+        await gather(*(h.stop() for h in self._all_handlers))
+
+    def put(self, event: DmlEvent) -> None:
+        if not (isinstance(event, DmlEvent) and event):
+            return
+        for cls in getmro(type(event)):
+            if not (issubclass(cls, DmlEvent) and (handlers := self._handlers.get(cls))):
+                continue
+            for h in handlers:
+                h.put(event)
+            if cls is DmlEvent:
+                break

phoenix/server/types.py

@@ -1,10 +1,36 @@
-from typing import AsyncContextManager, Callable
+from abc import ABC, abstractmethod
+from asyncio import Task, create_task, sleep
+from collections import defaultdict
+from datetime import datetime, timezone
+from typing import (
+    Any,
+    AsyncContextManager,
+    Callable,
+    DefaultDict,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Protocol,
+    Type,
+    TypeVar,
+)
 
+from cachetools import LRUCache
 from sqlalchemy.ext.asyncio import AsyncSession
 
+from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 
 
+class CanSetLastUpdatedAt(Protocol):
+    def set(self, table: Type[models.Base], id_: int) -> None: ...
+
+
+class CanGetLastUpdatedAt(Protocol):
+    def get(self, table: Type[models.Base], id_: Optional[int] = None) -> Optional[datetime]: ...
+
+
 class DbSessionFactory:
     def __init__(
         self,
@@ -16,3 +42,82 @@ class DbSessionFactory:
 
     def __call__(self) -> AsyncContextManager[AsyncSession]:
         return self._db()
+
+
+_AnyT = TypeVar("_AnyT")
+_ItemT_contra = TypeVar("_ItemT_contra", contravariant=True)
+
+
+class CanPutItem(Protocol[_ItemT_contra]):
+    def put(self, item: _ItemT_contra) -> None: ...
+
+
+class _Batch(CanPutItem[_AnyT], Protocol[_AnyT]):
+    @property
+    def empty(self) -> bool: ...
+    def clear(self) -> None: ...
+    def __iter__(self) -> Iterator[_AnyT]: ...
+
+
+class _HasBatch(Generic[_ItemT_contra], ABC):
+    _batch_factory: Callable[[], _Batch[_ItemT_contra]]
+
+    def __init__(self) -> None:
+        self._batch = self._batch_factory()
+
+    def put(self, item: _ItemT_contra) -> None:
+        self._batch.put(item)
+
+
+class BatchedCaller(_HasBatch[_AnyT], Generic[_AnyT], ABC):
+    def __init__(self, *, sleep_seconds: float = 0.1, **kwargs: Any) -> None:
+        assert sleep_seconds > 0
+        super().__init__(**kwargs)
+        self._running = False
+        self._seconds = sleep_seconds
+        self._tasks: List[Task[None]] = []
+
+    async def start(self) -> None:
+        self._running = True
+        if not self._tasks:
+            self._tasks.append(create_task(self._run()))
+
+    async def stop(self) -> None:
+        self._running = False
+        for task in reversed(self._tasks):
+            if not task.done():
+                task.cancel()
+        self._tasks.clear()
+
+    @abstractmethod
+    async def __call__(self) -> None: ...
+
+    async def _run(self) -> None:
+        while self._running:
+            self._tasks.append(create_task(sleep(self._seconds)))
+            await self._tasks[-1]
+            self._tasks.pop()
+            if self._batch.empty:
+                continue
+            self._tasks.append(create_task(self()))
+            await self._tasks[-1]
+            self._tasks.pop()
+            self._batch.clear()
+
+
+class LastUpdatedAt:
+    def __init__(self) -> None:
+        self._cache: DefaultDict[
+            Type[models.Base],
+            LRUCache[int, datetime],
+        ] = defaultdict(lambda: LRUCache(maxsize=100))
+
+    def get(self, table: Type[models.Base], id_: Optional[int] = None) -> Optional[datetime]:
+        if not (cache := self._cache.get(table)):
+            return None
+        if id_ is None:
+            return max(filter(bool, cache.values()), default=None)
+        return cache.get(id_)
+
+    def set(self, table: Type[models.Base], id_: int) -> None:
+        self._cache[table][id_] = datetime.now(timezone.utc)

phoenix/version.py

@@ -1 +1 @@
-__version__ = "4.20.0"
+__version__ = "4.20.1"