arize-phoenix 4.14.1__py3-none-any.whl → 4.16.0__py3-none-any.whl

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

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

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

phoenix/db/insertion/trace_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
+_TraceId: TypeAlias = str
+_TraceRowId: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+
+_Key: TypeAlias = Tuple[_Name, _TraceId]
+_UniqueBy: TypeAlias = Tuple[_Name, _TraceRowId]
+_Existing: TypeAlias = Tuple[
+    _TraceRowId,
+    _TraceId,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[datetime],
+]
+
+
+class TraceAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.TraceAnnotation,
+        Insertables.TraceAnnotation,
+        models.TraceAnnotation,
+    ],
+    table=models.TraceAnnotation,
+    unique_by=("name", "trace_rowid"),
+):
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.TraceAnnotation],
+    ) -> Tuple[
+        List[Received[Insertables.TraceAnnotation]],
+        List[Postponed[Precursors.TraceAnnotation]],
+        List[Received[Precursors.TraceAnnotation]],
+    ]:
+        to_insert: List[Received[Insertables.TraceAnnotation]] = []
+        to_postpone: List[Postponed[Precursors.TraceAnnotation]] = []
+        to_discard: List[Received[Precursors.TraceAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: List[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_traces: Mapping[str, _TraceAttr] = {
+            e.trace_id: _TraceAttr(e.trace_rowid) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            (e.name, e.trace_id): _AnnoAttr(e.trace_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(
+                                trace_rowid=anno.trace_rowid,
+                                id_=anno.id_,
+                            ),
+                        )
+                    )
+            elif (trace := existing_traces.get(p.item.trace_id)) is not None:
+                to_insert.append(
+                    Received(
+                        received_at=p.received_at,
+                        item=p.item.as_insertable(
+                            trace_rowid=trace.trace_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
+        trace = (
+            select(models.Trace.id, models.Trace.trace_id)
+            .where(models.Trace.trace_id.in_({trace_id for _, trace_id in keys}))
+            .cte()
+        )
+        onclause = and_(
+            trace.c.id == anno.trace_rowid,
+            anno.name.in_({name for name, _ in keys}),
+            tuple_(anno.name, trace.c.trace_id).in_(keys),
+        )
+        return select(
+            trace.c.id.label("trace_rowid"),
+            trace.c.trace_id,
+            anno.id,
+            anno.name,
+            anno.updated_at,
+        ).outerjoin_from(trace, anno, onclause)
+
+
+class _TraceAttr(NamedTuple):
+    trace_rowid: _TraceRowId
+
+
+class _AnnoAttr(NamedTuple):
+    trace_rowid: _TraceRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
+    return p.item.obj.name, p.item.trace_id
+
+
+def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.trace_rowid
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at