arize-phoenix 11.38.0__py3-none-any.whl → 12.2.0__py3-none-any.whl

phoenix/db/helpers.py

@@ -1,7 +1,7 @@
 from collections.abc import Callable, Hashable, Iterable
 from datetime import datetime
 from enum import Enum
-from typing import Any, Literal, Optional, TypeVar, Union
+from typing import Any, Literal, Optional, Sequence, TypeVar, Union
 
 import sqlalchemy as sa
 from openinference.semconv.trace import (
@@ -10,16 +10,24 @@ from openinference.semconv.trace import (
     SpanAttributes,
 )
 from sqlalchemy import (
+    Insert,
     Integer,
     Select,
     SQLColumnExpression,
     and_,
     case,
     distinct,
+    exists,
     func,
+    insert,
+    literal,
+    or_,
     select,
+    util,
 )
+from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import QueryableAttribute
+from sqlalchemy.sql.roles import InElementRole
 from typing_extensions import assert_never
 
 from phoenix.config import PLAYGROUND_PROJECT_NAME
@@ -118,51 +126,205 @@ def dedup(
     return ans
 
 
-def get_dataset_example_revisions(
+def _build_ranked_revisions_query(
     dataset_version_id: int,
-) -> Select[tuple[models.DatasetExampleRevision]]:
-    version = (
+    /,
+    *,
+    dataset_id: Optional[int] = None,
+    example_ids: Optional[Union[Sequence[int], InElementRole]] = None,
+) -> Select[tuple[int]]:
+    """
+    Build a query that ranks revisions per example within a dataset version.
+
+    This performs the core ranking logic using ROW_NUMBER() to find the latest
+    revision for each example within the specified dataset version.
+
+    Args:
+        dataset_version_id: Maximum dataset version to consider
+        dataset_id: Optional dataset ID - if provided, avoids subquery lookup
+
+    Returns:
+        SQLAlchemy SELECT query with revision ranking and basic dataset filtering
+    """
+    stmt = (
         select(
-            models.DatasetVersion.id,
-            models.DatasetVersion.dataset_id,
+            func.row_number()
+            .over(
+                partition_by=models.DatasetExampleRevision.dataset_example_id,
+                order_by=models.DatasetExampleRevision.dataset_version_id.desc(),
+            )
+            .label("rn"),
         )
-        .filter_by(id=dataset_version_id)
-        .subquery()
+        .join(models.DatasetExample)
+        .where(models.DatasetExampleRevision.dataset_version_id <= dataset_version_id)
     )
-    table = models.DatasetExampleRevision
-    revision = (
-        select(
-            table.dataset_example_id,
-            func.max(table.dataset_version_id).label("dataset_version_id"),
-        )
-        .join_from(
-            table,
-            models.DatasetExample,
-            table.dataset_example_id == models.DatasetExample.id,
+
+    if dataset_id is None:
+        version_subquery = (
+            select(models.DatasetVersion.dataset_id)
+            .filter_by(id=dataset_version_id)
+            .scalar_subquery()
         )
-        .join_from(
-            models.DatasetExample,
-            version,
-            models.DatasetExample.dataset_id == version.c.dataset_id,
+        stmt = stmt.where(models.DatasetExample.dataset_id == version_subquery)
+    else:
+        stmt = stmt.where(models.DatasetExample.dataset_id == dataset_id)
+
+    if example_ids is not None:
+        stmt = stmt.where(models.DatasetExampleRevision.dataset_example_id.in_(example_ids))
+
+    return stmt
+
+
+def get_dataset_example_revisions(
+    dataset_version_id: int,
+    /,
+    *,
+    dataset_id: Optional[int] = None,
+    example_ids: Optional[Union[Sequence[int], InElementRole]] = None,
+    split_ids: Optional[Union[Sequence[int], InElementRole]] = None,
+    split_names: Optional[Union[Sequence[str], InElementRole]] = None,
+) -> Select[tuple[models.DatasetExampleRevision]]:
+    """
+    Get the latest revisions for all dataset examples within a specific dataset version.
+
+    Excludes examples where the latest revision is a DELETE.
+
+    Args:
+        dataset_version_id: The dataset version to get revisions for
+        dataset_id: Optional dataset ID - if provided, avoids extra subquery lookup
+        example_ids: Optional filter by specific example IDs (subquery or list of IDs).
+            - None = no filtering
+            - Empty sequences/subqueries = no matches (strict filtering)
+        split_ids: Optional filter by split IDs (subquery or list of split IDs).
+            - None = no filtering
+            - Empty sequences/subqueries = no matches (strict filtering)
+        split_names: Optional filter by split names (subquery or list of split names).
+            - None = no filtering
+            - Empty sequences/subqueries = no matches (strict filtering)
+
+    Note:
+        - split_ids and split_names are mutually exclusive
+        - Use split_ids for better performance when IDs are available (avoids JOIN)
+        - Empty filters use strict behavior: empty inputs return zero results
+    """
+    if split_ids is not None and split_names is not None:
+        raise ValueError(
+            "Cannot specify both split_ids and split_names - they are mutually exclusive"
         )
-        .where(models.DatasetExample.dataset_id == version.c.dataset_id)
-        .where(table.dataset_version_id <= version.c.id)
-        .group_by(table.dataset_example_id)
-        .subquery()
+
+    stmt = _build_ranked_revisions_query(
+        dataset_version_id,
+        dataset_id=dataset_id,
+        example_ids=example_ids,
+    ).add_columns(
+        models.DatasetExampleRevision.id,
+        models.DatasetExampleRevision.revision_kind,
     )
+
+    if split_ids is not None or split_names is not None:
+        if split_names is not None:
+            split_example_ids_subquery = (
+                select(models.DatasetSplitDatasetExample.dataset_example_id)
+                .join(
+                    models.DatasetSplit,
+                    models.DatasetSplit.id == models.DatasetSplitDatasetExample.dataset_split_id,
+                )
+                .where(models.DatasetSplit.name.in_(split_names))
+            )
+            stmt = stmt.where(models.DatasetExample.id.in_(split_example_ids_subquery))
+        else:
+            assert split_ids is not None
+            split_example_ids_subquery = select(
+                models.DatasetSplitDatasetExample.dataset_example_id
+            ).where(models.DatasetSplitDatasetExample.dataset_split_id.in_(split_ids))
+            stmt = stmt.where(models.DatasetExample.id.in_(split_example_ids_subquery))
+
+    ranked_subquery = stmt.subquery()
     return (
-        select(table)
-        .where(table.revision_kind != "DELETE")
+        select(models.DatasetExampleRevision)
         .join(
-            revision,
-            onclause=and_(
-                revision.c.dataset_example_id == table.dataset_example_id,
-                revision.c.dataset_version_id == table.dataset_version_id,
-            ),
+            ranked_subquery,
+            models.DatasetExampleRevision.id == ranked_subquery.c.id,
+        )
+        .where(
+            ranked_subquery.c.rn == 1,
+            ranked_subquery.c.revision_kind != "DELETE",
         )
     )
 
 
+def create_experiment_examples_snapshot_insert(
+    experiment: models.Experiment,
+) -> Insert:
+    """
+    Create an INSERT statement to snapshot dataset examples for an experiment.
+
+    This captures which examples belong to the experiment at the time of creation,
+    respecting any dataset splits assigned to the experiment.
+
+    Args:
+        experiment: The experiment to create the snapshot for
+
+    Returns:
+        SQLAlchemy INSERT statement ready for execution
+    """
+    stmt = _build_ranked_revisions_query(
+        experiment.dataset_version_id,
+        dataset_id=experiment.dataset_id,
+    ).add_columns(
+        models.DatasetExampleRevision.id,
+        models.DatasetExampleRevision.dataset_example_id,
+        models.DatasetExampleRevision.revision_kind,
+    )
+
+    experiment_splits_subquery = select(models.ExperimentDatasetSplit.dataset_split_id).where(
+        models.ExperimentDatasetSplit.experiment_id == experiment.id
+    )
+    has_splits_condition = exists(experiment_splits_subquery)
+    split_filtered_example_ids = select(models.DatasetSplitDatasetExample.dataset_example_id).where(
+        models.DatasetSplitDatasetExample.dataset_split_id.in_(experiment_splits_subquery)
+    )
+
+    stmt = stmt.where(
+        or_(
+            ~has_splits_condition,  # No splits = include all examples
+            models.DatasetExampleRevision.dataset_example_id.in_(
+                split_filtered_example_ids
+            ),  # Has splits = filter by splits
+        )
+    )
+
+    ranked_subquery = stmt.subquery()
+    return insert(models.ExperimentDatasetExample).from_select(
+        [
+            models.ExperimentDatasetExample.experiment_id,
+            models.ExperimentDatasetExample.dataset_example_id,
+            models.ExperimentDatasetExample.dataset_example_revision_id,
+        ],
+        select(
+            literal(experiment.id),
+            ranked_subquery.c.dataset_example_id,
+            ranked_subquery.c.id,
+        ).where(
+            ranked_subquery.c.rn == 1,
+            ranked_subquery.c.revision_kind != "DELETE",
+        ),
+    )
+
+
+async def insert_experiment_with_examples_snapshot(
+    session: AsyncSession,
+    experiment: models.Experiment,
+) -> None:
+    """
+    Insert an experiment with its snapshot of dataset examples.
+    """
+    session.add(experiment)
+    await session.flush()
+    insert_stmt = create_experiment_examples_snapshot_insert(experiment)
+    await session.execute(insert_stmt)
+
+
 _AnyTuple = TypeVar("_AnyTuple", bound=tuple[Any, ...])
 
 
@@ -355,3 +517,10 @@ def get_ancestor_span_rowids(parent_id: str) -> Select[tuple[int]]:
         )
     )
     return select(ancestors.c.id)
+
+
+def truncate_name(name: str, max_len: int = 63) -> str:
+    # https://github.com/sqlalchemy/sqlalchemy/blob/e263825e3c5060bf4f47eed0e833c6660a31658e/lib/sqlalchemy/sql/compiler.py#L7844-L7845
+    if len(name) > max_len:
+        return name[0 : max_len - 8] + "_" + util.md5_hex(name)[-4:]
+    return name

phoenix/db/insertion/dataset.py

@@ -44,6 +44,7 @@ async def insert_dataset(
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     created_at: Optional[datetime] = None,
+    user_id: Optional[int] = None,
 ) -> DatasetId:
     id_ = await session.scalar(
         insert(models.Dataset)
@@ -52,6 +53,7 @@ async def insert_dataset(
             description=description,
             metadata_=metadata,
             created_at=created_at,
+            user_id=user_id,
         )
         .returning(models.Dataset.id)
     )
@@ -64,6 +66,7 @@ async def insert_dataset_version(
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     created_at: Optional[datetime] = None,
+    user_id: Optional[int] = None,
 ) -> DatasetVersionId:
     id_ = await session.scalar(
         insert(models.DatasetVersion)
@@ -72,6 +75,7 @@ async def insert_dataset_version(
             description=description,
             metadata_=metadata,
             created_at=created_at,
+            user_id=user_id,
         )
         .returning(models.DatasetVersion.id)
     )
@@ -152,6 +156,7 @@ async def add_dataset_examples(
     description: Optional[str] = None,
     metadata: Optional[Mapping[str, Any]] = None,
     action: DatasetAction = DatasetAction.CREATE,
+    user_id: Optional[int] = None,
 ) -> Optional[DatasetExampleAdditionEvent]:
     created_at = datetime.now(timezone.utc)
     dataset_id: Optional[DatasetId] = None
@@ -167,6 +172,7 @@ async def add_dataset_examples(
                 description=description,
                 metadata=metadata,
                 created_at=created_at,
+                user_id=user_id,
             )
         except Exception:
             logger.exception(f"Failed to insert dataset: {name=}")
@@ -176,6 +182,7 @@ async def add_dataset_examples(
             session=session,
             dataset_id=dataset_id,
             created_at=created_at,
+            user_id=user_id,
         )
     except Exception:
         logger.exception(f"Failed to insert dataset version for {dataset_id=}")

phoenix/db/insertion/document_annotation.py

@@ -181,7 +181,7 @@ def _key(p: Received[Precursors.DocumentAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.DocumentAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid, p.item.document_position, p.item.identifier
+    return p.item.obj.name, p.item.span_rowid, p.item.document_position, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/helpers.py

@@ -12,7 +12,7 @@ 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.helpers import SupportedSQLDialect, truncate_name
 from phoenix.db.models import Base
 from phoenix.trace.attributes import get_attribute_value
 
@@ -53,7 +53,7 @@ def insert_on_conflict(
             unique_records.append(v)
             seen.add(k)
         records = tuple(reversed(unique_records))
-    constraint = constraint_name or "_".join(("uq", table.__tablename__, *unique_by))
+    constraint = constraint_name or truncate_name("_".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:

phoenix/db/insertion/session_annotation.py

@@ -0,0 +1,176 @@
+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
+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 ProjectSessionAnnotationDmlEvent
+
+# Type alias for consistency with other annotation patterns
+SessionAnnotationDmlEvent = ProjectSessionAnnotationDmlEvent
+
+_Name: TypeAlias = str
+_SessionId: TypeAlias = str
+_SessionRowId: TypeAlias = int
+_AnnoRowId: TypeAlias = int
+_Identifier: TypeAlias = str
+
+
+class _Key(NamedTuple):
+    annotation_name: _Name
+    annotation_identifier: _Identifier
+    session_id: _SessionId
+
+
+_UniqueBy: TypeAlias = tuple[_Name, _SessionRowId, _Identifier]
+_Existing: TypeAlias = tuple[
+    _SessionRowId,
+    _SessionId,
+    Optional[_AnnoRowId],
+    Optional[_Name],
+    Optional[datetime],
+]
+
+
+class SessionAnnotationQueueInserter(
+    QueueInserter[
+        Precursors.SessionAnnotation,
+        Insertables.SessionAnnotation,
+        models.ProjectSessionAnnotation,
+        SessionAnnotationDmlEvent,
+    ],
+    table=models.ProjectSessionAnnotation,
+    unique_by=("name", "project_session_id", "identifier"),
+):
+    async def _events(
+        self,
+        session: AsyncSession,
+        *insertions: Insertables.SessionAnnotation,
+    ) -> list[SessionAnnotationDmlEvent]:
+        records = [{**dict(as_kv(ins.row)), "updated_at": ins.row.updated_at} 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 [SessionAnnotationDmlEvent(ids)]
+
+    async def _partition(
+        self,
+        session: AsyncSession,
+        *parcels: Received[Precursors.SessionAnnotation],
+    ) -> tuple[
+        list[Received[Insertables.SessionAnnotation]],
+        list[Postponed[Precursors.SessionAnnotation]],
+        list[Received[Precursors.SessionAnnotation]],
+    ]:
+        to_insert: list[Received[Insertables.SessionAnnotation]] = []
+        to_postpone: list[Postponed[Precursors.SessionAnnotation]] = []
+        to_discard: list[Received[Precursors.SessionAnnotation]] = []
+
+        stmt = self._select_existing(*map(_key, parcels))
+        existing: list[Row[_Existing]] = [_ async for _ in await session.stream(stmt)]
+        existing_sessions: Mapping[str, _SessionAttr] = {
+            e.session_id: _SessionAttr(e.session_rowid) for e in existing
+        }
+        existing_annos: Mapping[_Key, _AnnoAttr] = {
+            _Key(
+                annotation_name=e.name,
+                annotation_identifier=e.identifier,
+                session_id=e.session_id,
+            ): _AnnoAttr(e.session_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.item.updated_at <= anno.updated_at:
+                    to_discard.append(p)
+                else:
+                    to_insert.append(
+                        Received(
+                            received_at=p.received_at,
+                            item=p.item.as_insertable(
+                                project_session_rowid=anno.session_rowid,
+                            ),
+                        )
+                    )
+            elif (existing_session := existing_sessions.get(p.item.session_id)) is not None:
+                to_insert.append(
+                    Received(
+                        received_at=p.received_at,
+                        item=p.item.as_insertable(
+                            project_session_rowid=existing_session.session_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
+        session = (
+            select(models.ProjectSession.id, models.ProjectSession.session_id)
+            .where(models.ProjectSession.session_id.in_({k.session_id for k in keys}))
+            .cte()
+        )
+        onclause = and_(
+            session.c.id == anno.project_session_id,
+            anno.name.in_({k.annotation_name for k in keys}),
+            tuple_(anno.name, anno.identifier, session.c.session_id).in_(keys),
+        )
+        return select(
+            session.c.id.label("session_rowid"),
+            session.c.session_id,
+            anno.id,
+            anno.name,
+            anno.identifier,
+            anno.updated_at,
+        ).outerjoin_from(session, anno, onclause)
+
+
+class _SessionAttr(NamedTuple):
+    session_rowid: _SessionRowId
+
+
+class _AnnoAttr(NamedTuple):
+    session_rowid: _SessionRowId
+    id_: _AnnoRowId
+    updated_at: datetime
+
+
+def _key(p: Received[Precursors.SessionAnnotation]) -> _Key:
+    return _Key(
+        annotation_name=p.item.obj.name,
+        annotation_identifier=p.item.obj.identifier,
+        session_id=p.item.session_id,
+    )
+
+
+def _unique_by(p: Received[Insertables.SessionAnnotation]) -> _UniqueBy:
+    return p.item.obj.name, p.item.project_session_rowid, p.item.obj.identifier
+
+
+def _time(p: Received[Any]) -> datetime:
+    return p.received_at

phoenix/db/insertion/span_annotation.py

@@ -167,7 +167,7 @@ def _key(p: Received[Precursors.SpanAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.SpanAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.span_rowid, p.item.identifier
+    return p.item.obj.name, p.item.span_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/trace_annotation.py

@@ -166,7 +166,7 @@ def _key(p: Received[Precursors.TraceAnnotation]) -> _Key:
 
 
 def _unique_by(p: Received[Insertables.TraceAnnotation]) -> _UniqueBy:
-    return p.item.obj.name, p.item.trace_rowid, p.item.identifier
+    return p.item.obj.name, p.item.trace_rowid, p.item.obj.identifier
 
 
 def _time(p: Received[Any]) -> datetime:

phoenix/db/insertion/types.py

@@ -225,13 +225,29 @@ class Precursors(ABC):
                 span_rowid=span_rowid,
             )
 
+    @dataclass(frozen=True)
+    class SessionAnnotation:
+        updated_at: datetime
+        session_id: str
+        obj: models.ProjectSessionAnnotation
+
+        def as_insertable(
+            self,
+            project_session_rowid: int,
+        ) -> Insertables.SessionAnnotation:
+            return Insertables.SessionAnnotation(
+                updated_at=self.updated_at,
+                session_id=self.session_id,
+                obj=self.obj,
+                project_session_rowid=project_session_rowid,
+            )
+
 
 class Insertables(ABC):
     @dataclass(frozen=True)
     class SpanAnnotation(Precursors.SpanAnnotation):
         updated_at: datetime
         span_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.SpanAnnotation:
@@ -244,7 +260,6 @@ class Insertables(ABC):
     class TraceAnnotation(Precursors.TraceAnnotation):
         updated_at: datetime
         trace_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.TraceAnnotation:
@@ -257,7 +272,6 @@ class Insertables(ABC):
     class DocumentAnnotation(Precursors.DocumentAnnotation):
         updated_at: datetime
         span_rowid: int
-        identifier: str = ""
 
         @property
         def row(self) -> models.DocumentAnnotation:
@@ -265,3 +279,15 @@ class Insertables(ABC):
             obj.span_rowid = self.span_rowid
             obj.updated_at = self.updated_at
             return obj
+
+    @dataclass(frozen=True)
+    class SessionAnnotation(Precursors.SessionAnnotation):
+        updated_at: datetime
+        project_session_rowid: int
+
+        @property
+        def row(self) -> models.ProjectSessionAnnotation:
+            obj = copy(self.obj)
+            obj.project_session_id = self.project_session_rowid
+            obj.updated_at = self.updated_at
+            return obj

phoenix/db/migrations/versions/01a8342c9cdf_add_user_id_on_datasets.py

@@ -0,0 +1,40 @@
+"""add user_id on datasets
+
+Revision ID: 01a8342c9cdf
+Revises: 0df286449799
+Create Date: 2025-09-25 16:08:51.254947
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "01a8342c9cdf"
+down_revision: Union[str, None] = "0df286449799"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+_Integer = sa.Integer().with_variant(
+    sa.BigInteger(),
+    "postgresql",
+)
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("datasets") as batch_op:
+        batch_op.add_column(
+            sa.Column(
+                "user_id",
+                _Integer,
+                sa.ForeignKey("users.id", ondelete="SET NULL"),
+                nullable=True,
+            ),
+        )
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("datasets") as batch_op:
+        batch_op.drop_column("user_id")