arize-phoenix 3.16.0__py3-none-any.whl → 7.7.0__py3-none-any.whl

phoenix/db/insertion/dataset.py

@@ -0,0 +1,227 @@
+import logging
+from collections.abc import Awaitable, Iterable, Iterator, Mapping
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from enum import Enum
+from itertools import chain
+from typing import Any, Optional, Union, cast
+
+from sqlalchemy import insert, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import TypeAlias
+
+from phoenix.db import models
+from phoenix.db.insertion.helpers import DataManipulationEvent
+
+logger = logging.getLogger(__name__)
+
+DatasetId: TypeAlias = int
+DatasetVersionId: TypeAlias = int
+DatasetExampleId: TypeAlias = int
+DatasetExampleRevisionId: TypeAlias = int
+SpanRowId: TypeAlias = int
+
+
+@dataclass(frozen=True)
+class ExampleContent:
+    input: dict[str, Any] = field(default_factory=dict)
+    output: dict[str, Any] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+Examples: TypeAlias = Iterable[ExampleContent]
+
+
+@dataclass(frozen=True)
+class DatasetExampleAdditionEvent(DataManipulationEvent):
+    dataset_id: DatasetId
+
+
+async def insert_dataset(
+    session: AsyncSession,
+    name: str,
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    created_at: Optional[datetime] = None,
+) -> DatasetId:
+    id_ = await session.scalar(
+        insert(models.Dataset)
+        .values(
+            name=name,
+            description=description,
+            metadata_=metadata,
+            created_at=created_at,
+        )
+        .returning(models.Dataset.id)
+    )
+    return cast(DatasetId, id_)
+
+
+async def insert_dataset_version(
+    session: AsyncSession,
+    dataset_id: DatasetId,
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    created_at: Optional[datetime] = None,
+) -> DatasetVersionId:
+    id_ = await session.scalar(
+        insert(models.DatasetVersion)
+        .values(
+            dataset_id=dataset_id,
+            description=description,
+            metadata_=metadata,
+            created_at=created_at,
+        )
+        .returning(models.DatasetVersion.id)
+    )
+    return cast(DatasetVersionId, id_)
+
+
+async def insert_dataset_example(
+    session: AsyncSession,
+    dataset_id: DatasetId,
+    span_rowid: Optional[SpanRowId] = None,
+    created_at: Optional[datetime] = None,
+) -> DatasetExampleId:
+    id_ = await session.scalar(
+        insert(models.DatasetExample)
+        .values(
+            dataset_id=dataset_id,
+            span_rowid=span_rowid,
+            created_at=created_at,
+        )
+        .returning(models.DatasetExample.id)
+    )
+    return cast(DatasetExampleId, id_)
+
+
+class RevisionKind(Enum):
+    CREATE = "CREATE"
+    PATCH = "PATCH"
+    DELETE = "DELETE"
+
+    @classmethod
+    def _missing_(cls, v: Any) -> "RevisionKind":
+        if isinstance(v, str) and v and v.isascii() and not v.isupper():
+            return cls(v.upper())
+        raise ValueError(f"Invalid revision kind: {v}")
+
+
+async def insert_dataset_example_revision(
+    session: AsyncSession,
+    dataset_version_id: DatasetVersionId,
+    dataset_example_id: DatasetExampleId,
+    input: Mapping[str, Any],
+    output: Mapping[str, Any],
+    metadata: Optional[Mapping[str, Any]] = None,
+    revision_kind: RevisionKind = RevisionKind.CREATE,
+    created_at: Optional[datetime] = None,
+) -> DatasetExampleRevisionId:
+    id_ = await session.scalar(
+        insert(models.DatasetExampleRevision)
+        .values(
+            dataset_version_id=dataset_version_id,
+            dataset_example_id=dataset_example_id,
+            input=input,
+            output=output,
+            metadata_=metadata,
+            revision_kind=revision_kind.value,
+            created_at=created_at,
+        )
+        .returning(models.DatasetExampleRevision.id)
+    )
+    return cast(DatasetExampleRevisionId, id_)
+
+
+class DatasetAction(Enum):
+    CREATE = "create"
+    APPEND = "append"
+
+    @classmethod
+    def _missing_(cls, v: Any) -> "DatasetAction":
+        if isinstance(v, str) and v and v.isascii() and not v.islower():
+            return cls(v.lower())
+        raise ValueError(f"Invalid dateset action: {v}")
+
+
+async def add_dataset_examples(
+    session: AsyncSession,
+    name: str,
+    examples: Union[Examples, Awaitable[Examples]],
+    description: Optional[str] = None,
+    metadata: Optional[Mapping[str, Any]] = None,
+    action: DatasetAction = DatasetAction.CREATE,
+) -> Optional[DatasetExampleAdditionEvent]:
+    created_at = datetime.now(timezone.utc)
+    dataset_id: Optional[DatasetId] = None
+    if action is DatasetAction.APPEND and name:
+        dataset_id = await session.scalar(
+            select(models.Dataset.id).where(models.Dataset.name == name)
+        )
+    if action is DatasetAction.CREATE or dataset_id is None:
+        try:
+            dataset_id = await insert_dataset(
+                session=session,
+                name=name,
+                description=description,
+                metadata=metadata,
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(f"Failed to insert dataset: {name=}")
+            raise
+    try:
+        dataset_version_id = await insert_dataset_version(
+            session=session,
+            dataset_id=dataset_id,
+            created_at=created_at,
+        )
+    except Exception:
+        logger.exception(f"Failed to insert dataset version for {dataset_id=}")
+        raise
+    for example in (await examples) if isinstance(examples, Awaitable) else examples:
+        try:
+            dataset_example_id = await insert_dataset_example(
+                session=session,
+                dataset_id=dataset_id,
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(f"Failed to insert dataset example for {dataset_id=}")
+            raise
+        try:
+            await insert_dataset_example_revision(
+                session=session,
+                dataset_version_id=dataset_version_id,
+                dataset_example_id=dataset_example_id,
+                input=example.input,
+                output=example.output,
+                metadata=example.metadata,
+                created_at=created_at,
+            )
+        except Exception:
+            logger.exception(
+                f"Failed to insert dataset example revision for {dataset_version_id=}, "
+                f"{dataset_example_id=}"
+            )
+            raise
+    return DatasetExampleAdditionEvent(dataset_id=dataset_id)
+
+
+@dataclass(frozen=True)
+class DatasetKeys:
+    input: frozenset[str]
+    output: frozenset[str]
+    metadata: frozenset[str]
+
+    def __iter__(self) -> Iterator[str]:
+        yield from sorted(set(chain(self.input, self.output, self.metadata)))
+
+    def check_differences(self, column_headers_set: frozenset[str]) -> None:
+        for category, keys in (
+            ("input", self.input),
+            ("output", self.output),
+            ("metadata", self.metadata),
+        ):
+            if diff := keys.difference(column_headers_set):
+                raise ValueError(f"{category} keys not found in table column headers: {diff}")

phoenix/db/insertion/document_annotation.py

@@ -0,0 +1,171 @@
+from collections.abc import Mapping
+from datetime import datetime
+from typing import Any, NamedTuple, Optional
+
+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.helpers import as_kv
+from phoenix.db.insertion.types import (
+    Insertables,
+    Postponed,
+    Precursors,
+    QueueInserter,
+    Received,
+)
+from phoenix.server.dml_event import DocumentAnnotationDmlEvent
+
+_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,
+        DocumentAnnotationDmlEvent,
+    ],
+    table=models.DocumentAnnotation,
+    unique_by=("name", "span_rowid", "document_position"),
+):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.DocumentAnnotation,
+    ) -> list[DocumentAnnotationDmlEvent]:
+        records = [dict(as_kv(ins.row)) for ins in insertions]
+        stmt = self._insert_on_conflict(*records).returning(self.table.id)
+        ids = tuple([_ async for _ in await session.stream_scalars(stmt)])
+        return [DocumentAnnotationDmlEvent(ids)]
+
+    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/evaluation.py

@@ -0,0 +1,191 @@
+from typing import NamedTuple, Optional
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from typing_extensions import assert_never
+
+from phoenix.db import models
+from phoenix.db.helpers import SupportedSQLDialect, num_docs_col
+from phoenix.db.insertion.helpers import insert_on_conflict
+from phoenix.exceptions import PhoenixException
+from phoenix.trace import v1 as pb
+
+
+class InsertEvaluationError(PhoenixException):
+    pass
+
+
+class EvaluationInsertionEvent(NamedTuple):
+    project_rowid: int
+    evaluation_name: str
+
+
+class SpanEvaluationInsertionEvent(EvaluationInsertionEvent): ...
+
+
+class TraceEvaluationInsertionEvent(EvaluationInsertionEvent): ...
+
+
+class DocumentEvaluationInsertionEvent(EvaluationInsertionEvent): ...
+
+
+async def insert_evaluation(
+    session: AsyncSession,
+    evaluation: pb.Evaluation,
+) -> Optional[EvaluationInsertionEvent]:
+    evaluation_name = evaluation.name
+    result = evaluation.result
+    label = result.label.value if result.HasField("label") else None
+    score = result.score.value if result.HasField("score") else None
+    explanation = result.explanation.value if result.HasField("explanation") else None
+    if (evaluation_kind := evaluation.subject_id.WhichOneof("kind")) is None:
+        raise InsertEvaluationError("Cannot insert an evaluation that has no evaluation kind")
+    elif evaluation_kind == "trace_id":
+        trace_id = evaluation.subject_id.trace_id
+        return await _insert_trace_evaluation(
+            session, trace_id, evaluation_name, label, score, explanation
+        )
+    elif evaluation_kind == "span_id":
+        span_id = evaluation.subject_id.span_id
+        return await _insert_span_evaluation(
+            session, span_id, evaluation_name, label, score, explanation
+        )
+    elif evaluation_kind == "document_retrieval_id":
+        span_id = evaluation.subject_id.document_retrieval_id.span_id
+        document_position = evaluation.subject_id.document_retrieval_id.document_position
+        return await _insert_document_evaluation(
+            session, span_id, document_position, evaluation_name, label, score, explanation
+        )
+    else:
+        assert_never(evaluation_kind)
+
+
+async def _insert_trace_evaluation(
+    session: AsyncSession,
+    trace_id: str,
+    evaluation_name: str,
+    label: Optional[str],
+    score: Optional[float],
+    explanation: Optional[str],
+) -> TraceEvaluationInsertionEvent:
+    stmt = select(
+        models.Trace.project_rowid,
+        models.Trace.id,
+    ).where(models.Trace.trace_id == trace_id)
+    if not (row := (await session.execute(stmt)).first()):
+        raise InsertEvaluationError(
+            f"Cannot insert a trace evaluation for a missing trace: {evaluation_name=}, {trace_id=}"
+        )
+    project_rowid, trace_rowid = row
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    values = dict(
+        trace_rowid=trace_rowid,
+        name=evaluation_name,
+        label=label,
+        score=score,
+        explanation=explanation,
+        metadata_={},  # `metadata_` must match ORM
+        annotator_kind="LLM",
+    )
+    await session.execute(
+        insert_on_conflict(
+            values,
+            dialect=dialect,
+            table=models.TraceAnnotation,
+            unique_by=("name", "trace_rowid"),
+        )
+    )
+    return TraceEvaluationInsertionEvent(project_rowid, evaluation_name)
+
+
+async def _insert_span_evaluation(
+    session: AsyncSession,
+    span_id: str,
+    evaluation_name: str,
+    label: Optional[str],
+    score: Optional[float],
+    explanation: Optional[str],
+) -> SpanEvaluationInsertionEvent:
+    stmt = (
+        select(
+            models.Trace.project_rowid,
+            models.Span.id,
+        )
+        .join_from(models.Span, models.Trace)
+        .where(models.Span.span_id == span_id)
+    )
+    if not (row := (await session.execute(stmt)).first()):
+        raise InsertEvaluationError(
+            f"Cannot insert a span evaluation for a missing span: {evaluation_name=}, {span_id=}"
+        )
+    project_rowid, span_rowid = row
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    values = dict(
+        span_rowid=span_rowid,
+        name=evaluation_name,
+        label=label,
+        score=score,
+        explanation=explanation,
+        metadata_={},  # `metadata_` must match ORM
+        annotator_kind="LLM",
+    )
+    await session.execute(
+        insert_on_conflict(
+            values,
+            dialect=dialect,
+            table=models.SpanAnnotation,
+            unique_by=("name", "span_rowid"),
+        )
+    )
+    return SpanEvaluationInsertionEvent(project_rowid, evaluation_name)
+
+
+async def _insert_document_evaluation(
+    session: AsyncSession,
+    span_id: str,
+    document_position: int,
+    evaluation_name: str,
+    label: Optional[str],
+    score: Optional[float],
+    explanation: Optional[str],
+) -> EvaluationInsertionEvent:
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    stmt = (
+        select(
+            models.Trace.project_rowid,
+            models.Span.id,
+            num_docs_col(dialect),
+        )
+        .join_from(models.Span, models.Trace)
+        .where(models.Span.span_id == span_id)
+    )
+    if not (row := (await session.execute(stmt)).first()):
+        raise InsertEvaluationError(
+            f"Cannot insert a document evaluation for a missing span: {span_id=}"
+        )
+    project_rowid, span_rowid, num_docs = row
+    if num_docs is None or num_docs <= document_position:
+        raise InsertEvaluationError(
+            f"Cannot insert a document evaluation for a non-existent "
+            f"document position: {evaluation_name=}, {span_id=}, {document_position=}"
+        )
+    dialect = SupportedSQLDialect(session.bind.dialect.name)
+    values = dict(
+        span_rowid=span_rowid,
+        document_position=document_position,
+        name=evaluation_name,
+        label=label,
+        score=score,
+        explanation=explanation,
+        metadata_={},  # `metadata_` must match ORM
+        annotator_kind="LLM",
+    )
+    await session.execute(
+        insert_on_conflict(
+            values,
+            dialect=dialect,
+            table=models.DocumentAnnotation,
+            unique_by=("name", "span_rowid", "document_position"),
+        )
+    )
+    return DocumentEvaluationInsertionEvent(project_rowid, evaluation_name)

phoenix/db/insertion/helpers.py

@@ -0,0 +1,98 @@
+from abc import ABC
+from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence
+from enum import Enum, auto
+from typing import Any, Optional
+
+from sqlalchemy import Insert
+from sqlalchemy.dialects.postgresql import insert as insert_postgresql
+from sqlalchemy.dialects.sqlite import insert as insert_sqlite
+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
+
+
+class DataManipulationEvent(ABC):
+    """
+    Execution of DML (Data Manipulation Language) statements.
+    """
+
+
+DataManipulation: TypeAlias = Callable[[AsyncSession], Awaitable[Optional[DataManipulationEvent]]]
+
+
+class OnConflict(Enum):
+    DO_NOTHING = auto()
+    DO_UPDATE = auto()
+
+
+def insert_on_conflict(
+    *records: Mapping[str, Any],
+    table: type[Base],
+    dialect: SupportedSQLDialect,
+    unique_by: Sequence[str],
+    on_conflict: OnConflict = OnConflict.DO_UPDATE,
+    set_: Optional[Mapping[str, Any]] = None,
+) -> Insert:
+    """
+    Dialect specific insertion statement using ON CONFLICT DO syntax.
+    """
+    if on_conflict is OnConflict.DO_UPDATE:
+        # postegresql rejects duplicate updates for the same record
+        seen = set()
+        unique_records = []
+        for v in reversed(records):
+            if (k := tuple(v.get(name) for name in unique_by)) in seen:
+                continue
+            unique_records.append(v)
+            seen.add(k)
+        records = tuple(reversed(unique_records))
+    constraint = "_".join(("uq", table.__tablename__, *unique_by))
+    if dialect is SupportedSQLDialect.POSTGRESQL:
+        stmt_postgresql = insert_postgresql(table).values(records)
+        if on_conflict is OnConflict.DO_NOTHING:
+            return stmt_postgresql.on_conflict_do_nothing(constraint=constraint)
+        if on_conflict is OnConflict.DO_UPDATE:
+            return stmt_postgresql.on_conflict_do_update(
+                constraint=constraint,
+                set_=set_ if set_ else dict(_clean(stmt_postgresql.excluded.items())),
+            )
+        assert_never(on_conflict)
+    if dialect is SupportedSQLDialect.SQLITE:
+        stmt_sqlite = insert_sqlite(table).values(records)
+        if on_conflict is OnConflict.DO_NOTHING:
+            return stmt_sqlite.on_conflict_do_nothing(unique_by)
+        if on_conflict is OnConflict.DO_UPDATE:
+            return stmt_sqlite.on_conflict_do_update(
+                unique_by,
+                set_=set_ if set_ else dict(_clean(stmt_sqlite.excluded.items())),
+            )
+        assert_never(on_conflict)
+    assert_never(dialect)
+
+
+def _clean(
+    kv: Iterable[tuple[str, KeyedColumnElement[Any]]],
+) -> Iterator[tuple[str, KeyedColumnElement[Any]]]:
+    for k, v in kv:
+        if v.primary_key or v.foreign_keys or k == "created_at":
+            continue
+        if k == "metadata_":
+            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