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

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/session/client.py

@@ -324,7 +324,7 @@ class Client(TraceDataExtractor):
              Dataset: The dataset object.
         """
         response = self._client.get(
-            urljoin(self._base_url, "/v1/datasets"),
+            urljoin(self._base_url, "v1/datasets"),
             params={"name": name},
         )
         response.raise_for_status()
@@ -366,7 +366,7 @@ class Client(TraceDataExtractor):
             raise ValueError("Dataset id or name must be provided.")
 
         response = self._client.get(
-            urljoin(self._base_url, f"/v1/datasets/{quote(id)}/examples"),
+            urljoin(self._base_url, f"v1/datasets/{quote(id)}/examples"),
             params={"version_id": version_id} if version_id else None,
         )
         response.raise_for_status()

phoenix/version.py

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