arize-phoenix 4.15.0__tar.gz → 4.16.0__tar.gz

{arize_phoenix-4.15.0 → arize_phoenix-4.16.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arize-phoenix
-Version: 4.15.0
+Version: 4.16.0
 Summary: AI Observability and Evaluation
 Project-URL: Documentation, https://docs.arize.com/phoenix/
 Project-URL: Issues, https://github.com/Arize-ai/phoenix/issues
@@ -72,6 +72,7 @@ Requires-Dist: anthropic; extra == 'dev'
 Requires-Dist: arize[autoembeddings,llm-evaluation]; extra == 'dev'
 Requires-Dist: asgi-lifespan; extra == 'dev'
 Requires-Dist: asyncpg; extra == 'dev'
+Requires-Dist: faker>=26.0.0; extra == 'dev'
 Requires-Dist: gcsfs; extra == 'dev'
 Requires-Dist: google-cloud-aiplatform>=1.3; extra == 'dev'
 Requires-Dist: hatch; extra == 'dev'

{arize_phoenix-4.15.0 → arize_phoenix-4.16.0}/pyproject.toml

@@ -93,6 +93,7 @@ dev = [
   "anthropic",
   "prometheus_client",
   "asgi-lifespan",
+  "Faker>=26.0.0",
 ]
 evals = []
 experimental = []
@@ -171,6 +172,7 @@ dependencies = [
   "nest-asyncio", # for executor testing
   "astunparse; python_version<'3.9'",  # `ast.unparse(...)` is only available starting with Python 3.9
   "asgi-lifespan",
+  "Faker>=26.0.0",
 ]
 
 [tool.hatch.envs.type]

{arize_phoenix-4.15.0 → arize_phoenix-4.16.0}/src/phoenix/db/bulk_inserter.py

@@ -1,26 +1,36 @@
 import asyncio
 import logging
-from asyncio import Queue
+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,
     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,
@@ -28,6 +38,9 @@ from phoenix.db.insertion.evaluation import (
 )
 from phoenix.db.insertion.helpers import DataManipulation, DataManipulationEvent
 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.trace.schemas import Span
@@ -55,6 +68,8 @@ class BulkInserter:
         max_ops_per_transaction: int = 1000,
         max_queue_size: int = 1000,
         enable_prometheus: bool = False,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
     ) -> None:
         """
         :param db: A function to initiate a new database session.
@@ -81,6 +96,9 @@ class BulkInserter:
         self._last_updated_at_by_project: LRUCache[ProjectRowId, datetime] = LRUCache(maxsize=100)
         self._cache_for_dataloaders = cache_for_dataloaders
         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):
@@ -90,6 +108,7 @@ class BulkInserter:
     async def __aenter__(
         self,
     ) -> Tuple[
+        Callable[[Any], Awaitable[None]],
         Callable[[Span, str], Awaitable[None]],
         Callable[[pb.Evaluation], Awaitable[None]],
         Callable[[DataManipulation], None],
@@ -98,6 +117,7 @@ class BulkInserter:
         self._operations = Queue(maxsize=self._max_queue_size)
         self._task = asyncio.create_task(self._bulk_insert())
         return (
+            self._enqueue,
             self._queue_span,
             self._queue_evaluation,
             self._enqueue_operation,
@@ -109,6 +129,9 @@ class BulkInserter:
             self._task.cancel()
             self._task = None
 
+    async def _enqueue(self, *items: Any) -> None:
+        await self._queue_inserters.enqueue(*items)
+
     def _enqueue_operation(self, operation: DataManipulation) -> None:
         cast("Queue[DataManipulation]", self._operations).put_nowait(operation)
 
@@ -124,7 +147,17 @@ class BulkInserter:
         assert isinstance(self._operations, Queue)
         spans_buffer, evaluations_buffer = None, None
         # start first insert immediately if the inserter has not run recently
-        while self._running or not self._operations.empty() or self._spans or self._evaluations:
+        while (
+            self._running
+            or not self._queue_inserters.empty
+            or not self._operations.empty()
+            or self._spans
+            or self._evaluations
+        ):
+            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)
             if self._operations.empty() and not (self._spans or self._evaluations):
                 await asyncio.sleep(self._sleep)
                 continue
@@ -244,3 +277,97 @@ 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:
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
+    ) -> None:
+        self._db = db
+        args = (db, retry_delay_sec, retry_allowance)
+        self._span_annotations = SpanAnnotationQueueInserter(*args)
+        self._trace_annotations = TraceAnnotationQueueInserter(*args)
+        self._document_annotations = DocumentAnnotationQueueInserter(*args)
+        self._queues = (
+            self._span_annotations,
+            self._trace_annotations,
+            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
+
+    @property
+    def empty(self) -> bool:
+        return all(q.empty for q in self._queues)
+
+    async def enqueue(self, *items: Any) -> None:
+        for item in items:
+            await self._enqueue(item)
+
+    @singledispatchmethod
+    async def _enqueue(self, item: Any) -> None: ...
+
+    @_enqueue.register(Precursors.SpanAnnotation)
+    @_enqueue.register(Insertables.SpanAnnotation)
+    async def _(self, item: Precursors.SpanAnnotation) -> None:
+        await self._span_annotations.enqueue(item)
+
+    @_enqueue.register(Precursors.TraceAnnotation)
+    @_enqueue.register(Insertables.TraceAnnotation)
+    async def _(self, item: Precursors.TraceAnnotation) -> None:
+        await self._trace_annotations.enqueue(item)
+
+    @_enqueue.register(Precursors.DocumentAnnotation)
+    @_enqueue.register(Insertables.DocumentAnnotation)
+    async def _(self, item: Precursors.DocumentAnnotation) -> None:
+        await self._document_annotations.enqueue(item)

{arize_phoenix-4.15.0 → arize_phoenix-4.16.0}/src/phoenix/db/helpers.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Any, Optional, Tuple
+from typing import Any, Callable, Hashable, Iterable, List, Optional, Set, Tuple, TypeVar
 
 from openinference.semconv.trace import (
     OpenInferenceSpanKindValues,
@@ -80,3 +80,25 @@ def get_project_names_for_experiments(*experiment_ids: int) -> Select[Tuple[Opti
         .where(models.Experiment.id.in_(set(experiment_ids)))
         .where(models.Experiment.project_name.isnot(None))
     )
+
+
+_AnyT = TypeVar("_AnyT")
+_KeyT = TypeVar("_KeyT", bound=Hashable)
+
+
+def dedup(
+    items: Iterable[_AnyT],
+    key: Callable[[_AnyT], _KeyT],
+) -> List[_AnyT]:
+    """
+    Discard subsequent duplicates after the first appearance in `items`.
+    """
+    ans = []
+    seen: Set[_KeyT] = set()
+    for item in items:
+        if (k := key(item)) in seen:
+            continue
+        else:
+            ans.append(item)
+            seen.add(k)
+    return ans

arize_phoenix-4.16.0/src/phoenix/db/insertion/constants.py

@@ -0,0 +1,2 @@
+DEFAULT_RETRY_DELAY_SEC: float = 60
+DEFAULT_RETRY_ALLOWANCE: int = 10

arize_phoenix-4.16.0/src/phoenix/db/insertion/document_annotation.py

@@ -0,0 +1,157 @@
+from datetime import datetime
+from typing import Any, List, Mapping, NamedTuple, Optional, Tuple
+
+from sqlalchemy import Row, Select, and_, select, tuple_
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.helpers import dedup, num_docs_col
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+
+_Name: TypeAlias = str
+_SpanId: TypeAlias = str
+_SpanRowId: TypeAlias = int
+_DocumentPosition: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+_NumDocs: TypeAlias = int
+
+_Key: TypeAlias = Tuple[_Name, _SpanId, _DocumentPosition]
+_UniqueBy: TypeAlias = Tuple[_Name, _SpanRowId, _DocumentPosition]
+_Existing: TypeAlias = Tuple[
+    _SpanRowId,
+    _SpanId,
+    _NumDocs,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[_DocumentPosition],
+    Optional[datetime],
+]
+
+
+class DocumentAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.DocumentAnnotation,
+        Insertables.DocumentAnnotation,
+        models.DocumentAnnotation,
+    ],
+    table=models.DocumentAnnotation,
+    unique_by=("name", "span_rowid", "document_position"),
+):
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.DocumentAnnotation],
+    ) -> Tuple[
+        List[Received[Insertables.DocumentAnnotation]],
+        List[Postponed[Precursors.DocumentAnnotation]],
+        List[Received[Precursors.DocumentAnnotation]],
+    ]:
+        to_insert: List[Received[Insertables.DocumentAnnotation]] = []
+        to_postpone: List[Postponed[Precursors.DocumentAnnotation]] = []
+        to_discard: List[Received[Precursors.DocumentAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: List[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_spans: Mapping[str, _SpanAttr] = {
+            e.span_id: _SpanAttr(e.span_rowid, e.num_docs) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            (e.name, e.span_id, e.document_position): _AnnoAttr(e.span_rowid, e.id, e.updated_at)
+            for e in existing
+            if e.id is not None
+            and e.name is not None
+            and e.document_position is not None
+            and e.updated_at is not None
+        }
+
+        for p in parcels:
+            if (anno := existing_annos.get(_key(p))) is not None:
+                if p.received_at <= anno.updated_at:
+                    to_discard.append(p)
+                else:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                span_rowid=anno.span_rowid,
+                                id_=anno.id_,
+                            ),
+                        )
+                    )
+            elif (span := existing_spans.get(p.item.span_id)) is not None:
+                if 0 <= p.item.document_position < span.num_docs:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                span_rowid=span.span_rowid,
+                            ),
+                        )
+                    )
+                else:
+                    to_discard.append(p)
+            elif isinstance(p, Postponed):
+                if p.retries_left > 1:
+                    to_postpone.append(p.postpone(p.retries_left - 1))
+                else:
+                    to_discard.append(p)
+            elif isinstance(p, Received):
+                to_postpone.append(p.postpone(self._retry_allowance))
+            else:
+                to_discard.append(p)
+
+        assert len(to_insert) + len(to_postpone) + len(to_discard) == len(parcels)
+        to_insert = dedup(sorted(to_insert, key=_time, reverse=True), _unique_by)[::-1]
+        return to_insert, to_postpone, to_discard
+
+    def _select_existing(self, *keys: _Key) -> Select[_Existing]:
+        anno = self.table
+        span = (
+            select(models.Span.id, models.Span.span_id, num_docs_col(self._db.dialect))
+            .where(models.Span.span_id.in_({span_id for _, span_id, *_ in keys}))
+            .cte()
+        )
+        onclause = and_(
+            span.c.id == anno.span_rowid,
+            anno.name.in_({name for name, *_ in keys}),
+            tuple_(anno.name, span.c.span_id, anno.document_position).in_(keys),
+        )
+        return select(
+            span.c.id.label("span_rowid"),
+            span.c.span_id,
+            span.c.num_docs,
+            anno.id,
+            anno.name,
+            anno.document_position,
+            anno.updated_at,
+        ).outerjoin_from(span, anno, onclause)
+
+
+class _SpanAttr(NamedTuple):
+    span_rowid: _SpanRowId
+    num_docs: _NumDocs
+
+
+class _AnnoAttr(NamedTuple):
+    span_rowid: _SpanRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.DocumentAnnotation]) -> _Key:
+    return p.item.obj.name, p.item.span_id, p.item.document_position
+
+
+def _unique_by(p: Received[Insertables.DocumentAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.span_rowid, p.item.document_position
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at

{arize_phoenix-4.15.0 → arize_phoenix-4.16.0}/src/phoenix/db/insertion/helpers.py

@@ -20,6 +20,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.sql.elements import KeyedColumnElement
 from typing_extensions import TypeAlias, assert_never
 
+from phoenix.db import models
 from phoenix.db.helpers import SupportedSQLDialect
 from phoenix.db.models import Base
 
@@ -93,3 +94,15 @@ def _clean(
             yield "metadata", v
         else:
             yield k, v
+
+
+def as_kv(obj: models.Base) -> Iterator[Tuple[str, Any]]:
+    for k, c in obj.__table__.c.items():
+        if k in ["created_at", "updated_at"]:
+            continue
+        k = "metadata_" if k == "metadata" else k
+        v = getattr(obj, k, None)
+        if c.primary_key and v is None:
+            # postgresql disallows None for primary key
+            continue
+        yield k, v

arize_phoenix-4.16.0/src/phoenix/db/insertion/span_annotation.py

@@ -0,0 +1,144 @@
+from datetime import datetime
+from typing import Any, List, Mapping, NamedTuple, Optional, Tuple
+
+from sqlalchemy import Row, Select, and_, select, tuple_
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.helpers import dedup
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+
+_Name: TypeAlias = str
+_SpanId: TypeAlias = str
+_SpanRowId: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+
+_Key: TypeAlias = Tuple[_Name, _SpanId]
+_UniqueBy: TypeAlias = Tuple[_Name, _SpanRowId]
+_Existing: TypeAlias = Tuple[
+    _SpanRowId,
+    _SpanId,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[datetime],
+]
+
+
+class SpanAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.SpanAnnotation,
+        Insertables.SpanAnnotation,
+        models.SpanAnnotation,
+    ],
+    table=models.SpanAnnotation,
+    unique_by=("name", "span_rowid"),
+):
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.SpanAnnotation],
+    ) -> Tuple[
+        List[Received[Insertables.SpanAnnotation]],
+        List[Postponed[Precursors.SpanAnnotation]],
+        List[Received[Precursors.SpanAnnotation]],
+    ]:
+        to_insert: List[Received[Insertables.SpanAnnotation]] = []
+        to_postpone: List[Postponed[Precursors.SpanAnnotation]] = []
+        to_discard: List[Received[Precursors.SpanAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: List[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_spans: Mapping[str, _SpanAttr] = {
+            e.span_id: _SpanAttr(e.span_rowid) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            (e.name, e.span_id): _AnnoAttr(e.span_rowid, e.id, e.updated_at)
+            for e in existing
+            if e.id is not None and e.name is not None and e.updated_at is not None
+        }
+
+        for p in parcels:
+            if (anno := existing_annos.get(_key(p))) is not None:
+                if p.received_at <= anno.updated_at:
+                    to_discard.append(p)
+                else:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                span_rowid=anno.span_rowid,
+                                id_=anno.id_,
+                            ),
+                        )
+                    )
+            elif (span := existing_spans.get(p.item.span_id)) is not None:
+                to_insert.append(
+                    Received(
+                        received_at=p.received_at,
+                        item=p.item.as_insertable(
+                            span_rowid=span.span_rowid,
+                        ),
+                    )
+                )
+            elif isinstance(p, Postponed):
+                if p.retries_left > 1:
+                    to_postpone.append(p.postpone(p.retries_left - 1))
+                else:
+                    to_discard.append(p)
+            elif isinstance(p, Received):
+                to_postpone.append(p.postpone(self._retry_allowance))
+            else:
+                to_discard.append(p)
+
+        assert len(to_insert) + len(to_postpone) + len(to_discard) == len(parcels)
+        to_insert = dedup(sorted(to_insert, key=_time, reverse=True), _unique_by)[::-1]
+        return to_insert, to_postpone, to_discard
+
+    def _select_existing(self, *keys: _Key) -> Select[_Existing]:
+        anno = self.table
+        span = (
+            select(models.Span.id, models.Span.span_id)
+            .where(models.Span.span_id.in_({span_id for _, span_id in keys}))
+            .cte()
+        )
+        onclause = and_(
+            span.c.id == anno.span_rowid,
+            anno.name.in_({name for name, _ in keys}),
+            tuple_(anno.name, span.c.span_id).in_(keys),
+        )
+        return select(
+            span.c.id.label("span_rowid"),
+            span.c.span_id,
+            anno.id,
+            anno.name,
+            anno.updated_at,
+        ).outerjoin_from(span, anno, onclause)
+
+
+class _SpanAttr(NamedTuple):
+    span_rowid: _SpanRowId
+
+
+class _AnnoAttr(NamedTuple):
+    span_rowid: _SpanRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.SpanAnnotation]) -> _Key:
+    return p.item.obj.name, p.item.span_id
+
+
+def _unique_by(p: Received[Insertables.SpanAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.span_rowid
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at