arize-phoenix 4.15.0__py3-none-any.whl → 4.16.1__py3-none-any.whl

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

phoenix/db/insertion/types.py

@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from abc import ABC, abstractmethod
+from copy import copy
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import (
+    Any,
+    Generic,
+    List,
+    Mapping,
+    Optional,
+    Protocol,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+    cast,
+)
+
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.sql.dml import ReturningInsert
+
+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.server.types import DbSessionFactory
+
+logger = logging.getLogger("__name__")
+
+
+class Insertable(Protocol):
+    @property
+    def row(self) -> models.Base: ...
+
+
+_AnyT = TypeVar("_AnyT")
+_PrecursorT = TypeVar("_PrecursorT")
+_InsertableT = TypeVar("_InsertableT", bound=Insertable)
+_RowT = TypeVar("_RowT", bound=models.Base)
+
+
+@dataclass(frozen=True)
+class Received(Generic[_AnyT]):
+    item: _AnyT
+    received_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def postpone(self, retries_left: int = DEFAULT_RETRY_ALLOWANCE) -> Postponed[_AnyT]:
+        return Postponed(item=self.item, received_at=self.received_at, retries_left=retries_left)
+
+
+@dataclass(frozen=True)
+class Postponed(Received[_AnyT]):
+    retries_left: int = field(default=DEFAULT_RETRY_ALLOWANCE)
+
+
+class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT]):
+    table: Type[_RowT]
+    unique_by: Sequence[str]
+
+    def __init_subclass__(
+        cls,
+        table: Type[_RowT],
+        unique_by: Sequence[str],
+    ) -> None:
+        cls.table = table
+        cls.unique_by = unique_by
+
+    def __init__(
+        self,
+        db: DbSessionFactory,
+        retry_delay_sec: float = DEFAULT_RETRY_DELAY_SEC,
+        retry_allowance: int = DEFAULT_RETRY_ALLOWANCE,
+    ) -> None:
+        self._queue: List[Received[_PrecursorT]] = []
+        self._db = db
+        self._retry_delay_sec = retry_delay_sec
+        self._retry_allowance = retry_allowance
+
+    @property
+    def empty(self) -> bool:
+        return not bool(self._queue)
+
+    async def enqueue(self, *items: _PrecursorT) -> None:
+        self._queue.extend([Received(item) for item in items])
+
+    @abstractmethod
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[_PrecursorT],
+    ) -> Tuple[
+        List[Received[_InsertableT]],
+        List[Postponed[_PrecursorT]],
+        List[Received[_PrecursorT]],
+    ]: ...
+
+    async def insert(self) -> Tuple[Type[_RowT], List[int]]:
+        if not self._queue:
+            return self.table, []
+        parcels = self._queue
+        self._queue = []
+        inserted_ids: List[int] = []
+        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)
+        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
+
+    def _stmt(self, *records: Mapping[str, Any]) -> ReturningInsert[Tuple[int]]:
+        pk = next(c for c in self.table.__table__.c if c.primary_key)
+        return insert_on_conflict(
+            *records,
+            table=self.table,
+            unique_by=self.unique_by,
+            dialect=self._db.dialect,
+        ).returning(pk)
+
+    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] = []
+        to_retry: List[Postponed[_PrecursorT]] = []
+        failures: List[Received[_InsertableT]] = []
+        stmt = self._stmt(*records)
+        try:
+            async with session.begin_nested():
+                ids = [id_ async for id_ in await session.stream_scalars(stmt)]
+                inserted_ids.extend(ids)
+        except BaseException:
+            logger.exception(
+                f"Failed to bulk insert for {self.table.__name__}. "
+                f"Will try to insert ({len(records)} records) individually instead."
+            )
+            for i, record in enumerate(records):
+                stmt = self._stmt(record)
+                try:
+                    async with session.begin_nested():
+                        ids = [id_ async for id_ in await session.stream_scalars(stmt)]
+                        inserted_ids.extend(ids)
+                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:
+                        to_retry.append(
+                            Postponed(
+                                item=cast(_PrecursorT, p.item),
+                                received_at=p.received_at,
+                                retries_left=(p.retries_left - 1)
+                                if isinstance(p, Postponed)
+                                else self._retry_allowance,
+                            )
+                        )
+        return inserted_ids, to_retry, failures
+
+
+class Precursors(ABC):
+    @dataclass(frozen=True)
+    class SpanAnnotation:
+        span_id: str
+        obj: models.SpanAnnotation
+
+        def as_insertable(
+            self,
+            span_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.SpanAnnotation:
+            return Insertables.SpanAnnotation(
+                span_id=self.span_id,
+                obj=self.obj,
+                span_rowid=span_rowid,
+                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class TraceAnnotation:
+        trace_id: str
+        obj: models.TraceAnnotation
+
+        def as_insertable(
+            self,
+            trace_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.TraceAnnotation:
+            return Insertables.TraceAnnotation(
+                trace_id=self.trace_id,
+                obj=self.obj,
+                trace_rowid=trace_rowid,
+                id_=id_,
+            )
+
+    @dataclass(frozen=True)
+    class DocumentAnnotation:
+        span_id: str
+        document_position: int
+        obj: models.DocumentAnnotation
+
+        def as_insertable(
+            self,
+            span_rowid: int,
+            id_: Optional[int] = None,
+        ) -> Insertables.DocumentAnnotation:
+            return Insertables.DocumentAnnotation(
+                span_id=self.span_id,
+                document_position=self.document_position,
+                obj=self.obj,
+                span_rowid=span_rowid,
+                id_=id_,
+            )
+
+
+class Insertables(ABC):
+    @dataclass(frozen=True)
+    class SpanAnnotation(Precursors.SpanAnnotation):
+        span_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.SpanAnnotation:
+            obj = copy(self.obj)
+            obj.span_rowid = self.span_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj
+
+    @dataclass(frozen=True)
+    class TraceAnnotation(Precursors.TraceAnnotation):
+        trace_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.TraceAnnotation:
+            obj = copy(self.obj)
+            obj.trace_rowid = self.trace_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj
+
+    @dataclass(frozen=True)
+    class DocumentAnnotation(Precursors.DocumentAnnotation):
+        span_rowid: int
+        id_: Optional[int] = None
+
+        @property
+        def row(self) -> models.DocumentAnnotation:
+            obj = copy(self.obj)
+            obj.span_rowid = self.span_rowid
+            if self.id_ is not None:
+                obj.id = self.id_
+            return obj

phoenix/experiments/types.py

@@ -226,7 +226,7 @@ class ExperimentRun:
 
     def __post_init__(self) -> None:
         if bool(self.output) == bool(self.error):
-            ValueError("Must specify exactly one of experiment_run_output or error")
+            raise ValueError("Must specify exactly one of experiment_run_output or error")
 
 
 @dataclass(frozen=True)
@@ -249,7 +249,7 @@ class EvaluationResult:
 
     def __post_init__(self) -> None:
         if self.score is None and not self.label:
-            ValueError("Must specify score or label, or both")
+            raise ValueError("Must specify score or label, or both")
         if self.score is None and not self.label:
             object.__setattr__(self, "score", 0)
         for k in ("label", "explanation"):
@@ -285,7 +285,7 @@ class ExperimentEvaluationRun:
 
     def __post_init__(self) -> None:
         if bool(self.result) == bool(self.error):
-            ValueError("Must specify either result or error")
+            raise ValueError("Must specify either result or error")
 
 
 ExperimentTask: TypeAlias = Union[

phoenix/server/api/input_types/SpanAnnotationSort.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class SpanAnnotationColumn(Enum):
+    createdAt = "created_at"
+    name = "name"
+
+
+@strawberry.input(description="The sort key and direction for SpanAnnotation connections")
+class SpanAnnotationSort:
+    col: SpanAnnotationColumn
+    dir: SortDir

phoenix/server/api/input_types/TraceAnnotationSort.py

@@ -0,0 +1,17 @@
+from enum import Enum
+
+import strawberry
+
+from phoenix.server.api.types.SortDir import SortDir
+
+
+@strawberry.enum
+class TraceAnnotationColumn(Enum):
+    createdAt = "created_at"
+    name = "name"
+
+
+@strawberry.input(description="The sort key and direction for TraceAnnotation connections")
+class TraceAnnotationSort:
+    col: TraceAnnotationColumn
+    dir: SortDir